tenferro-gpu 0.1.0

CubeCL-backed CUDA and WebGPU provider backends for tenferro tensors.
Documentation 
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

mod classify;
mod codegen;
mod launch;

use classify::classify;

use crate::backend::ElementwiseFusionPlan;
use crate::cubecl::CudaBackend;
use crate::Tensor;

pub(crate) fn execute_elementwise_fusion(
    backend: &CudaBackend,
    inputs: &[&Tensor],
    plan: &ElementwiseFusionPlan,
) -> crate::Result<Option<Vec<Tensor>>> {
    match plan.dtype() {
        crate::DType::F32 => {
            let Some(classified) = classify::<f32>(inputs, plan)? else {
                return Ok(None);
            };
            let outputs = launch::launch(backend.runtime(), classified)?;
            Ok(Some(outputs.into_iter().map(Tensor::F32).collect()))
        }
        crate::DType::F64 => {
            let Some(classified) = classify::<f64>(inputs, plan)? else {
                return Ok(None);
            };
            let outputs = launch::launch(backend.runtime(), classified)?;
            Ok(Some(outputs.into_iter().map(Tensor::F64).collect()))
        }
        crate::DType::C32 => {
            let Some(classified) = classify::<num_complex::Complex32>(inputs, plan)? else {
                return Ok(None);
            };
            let outputs = launch::launch(backend.runtime(), classified)?;
            Ok(Some(outputs.into_iter().map(Tensor::C32).collect()))
        }
        crate::DType::C64 => {
            let Some(classified) = classify::<num_complex::Complex64>(inputs, plan)? else {
                return Ok(None);
            };
            let outputs = launch::launch(backend.runtime(), classified)?;
            Ok(Some(outputs.into_iter().map(Tensor::C64).collect()))
        }
        crate::DType::I32 | crate::DType::I64 | crate::DType::Bool => Ok(None),
    }
}