tenferro-gpu 0.3.0

CubeCL-backed CUDA and WebGPU provider backends for tenferro tensors.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52

use std::sync::Arc;

use smallvec::SmallVec;
use tenferro_runtime::runtime::{
    EventDomainError, EventDomainId, EventDomainOperation, EventToken,
};
use tenferro_runtime::{Error as RuntimeError, Result};

pub(crate) fn admit_event_token<'a, T: 'static>(
    dependency: &'a dyn EventToken,
    expected: EventDomainId,
    token_type: &'static str,
) -> Result<&'a T> {
    let actual = dependency.origin();
    if actual != expected {
        return Err(RuntimeError::from(
            EventDomainError::DependencyDomainMismatch {
                operation: EventDomainOperation::Enqueue,
                node_index: None,
                expected,
                actual,
            },
        ));
    }

    dependency.as_any().downcast_ref::<T>().ok_or_else(|| {
        RuntimeError::from(EventDomainError::IncompatibleTokenType {
            operation: EventDomainOperation::Enqueue,
            node_index: None,
            expected,
            actual,
            token_type,
        })
    })
}

pub(crate) fn admit_event_tokens<'a, T: 'static, R>(
    dependencies: &'a [Arc<dyn EventToken>],
    expected: EventDomainId,
    token_type: &'static str,
    launch: impl FnOnce(&[&'a T]) -> Result<R>,
) -> Result<R> {
    let mut admitted = SmallVec::<[&'a T; 4]>::new();
    for dependency in dependencies {
        admitted.push(admit_event_token::<T>(
            dependency.as_ref(),
            expected,
            token_type,
        )?);
    }
    launch(&admitted)
}