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
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84

use tenferro_core_ops::{descriptor as primitive_descriptor, DTypePolicy, PrimitiveOpKind};

/// Host-side launch family used by the CubeCL primitive dispatcher.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum GpuLaunchKind {
    BinaryFloatComplexInt,
    BinaryFloatInt,
    UnaryFloatComplex,
    UnaryFloatComplexInt,
    UnaryFloatInt,
    UnaryFloatOnly,
    CompareFloatIntToBool,
    SelectBoolFloatInt,
    ClampFloat,
    Reduction,
}

/// GPU-specific primitive metadata derived from the core operation catalog.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) struct GpuOpDescriptor {
    pub(crate) kind: PrimitiveOpKind,
    pub(crate) name: &'static str,
    pub(crate) dtype_policy: DTypePolicy,
    pub(crate) launch: GpuLaunchKind,
}

pub(crate) fn gpu_descriptor(kind: PrimitiveOpKind) -> Option<GpuOpDescriptor> {
    let launch = match kind {
        PrimitiveOpKind::Add | PrimitiveOpKind::Sub | PrimitiveOpKind::Mul => {
            GpuLaunchKind::BinaryFloatComplexInt
        }
        PrimitiveOpKind::Div => GpuLaunchKind::BinaryFloatComplexInt,
        PrimitiveOpKind::Maximum | PrimitiveOpKind::Minimum | PrimitiveOpKind::Rem => {
            GpuLaunchKind::BinaryFloatInt
        }
        PrimitiveOpKind::Pow => GpuLaunchKind::BinaryFloatInt,
        PrimitiveOpKind::Neg => GpuLaunchKind::UnaryFloatComplexInt,
        PrimitiveOpKind::Conj => GpuLaunchKind::UnaryFloatComplex,
        PrimitiveOpKind::Abs | PrimitiveOpKind::Sign => GpuLaunchKind::UnaryFloatInt,
        PrimitiveOpKind::Exp
        | PrimitiveOpKind::Log
        | PrimitiveOpKind::Sin
        | PrimitiveOpKind::Cos
        | PrimitiveOpKind::Tanh
        | PrimitiveOpKind::Sqrt
        | PrimitiveOpKind::Rsqrt
        | PrimitiveOpKind::Expm1
        | PrimitiveOpKind::Log1p => GpuLaunchKind::UnaryFloatOnly,
        PrimitiveOpKind::Compare => GpuLaunchKind::CompareFloatIntToBool,
        PrimitiveOpKind::Select => GpuLaunchKind::SelectBoolFloatInt,
        PrimitiveOpKind::Clamp => GpuLaunchKind::ClampFloat,
        PrimitiveOpKind::ReduceSum
        | PrimitiveOpKind::ReduceSumSquares
        | PrimitiveOpKind::ReduceProd
        | PrimitiveOpKind::ReduceMax
        | PrimitiveOpKind::ReduceMin => GpuLaunchKind::Reduction,
        _ => return None,
    };

    let descriptor = primitive_descriptor(kind);
    Some(GpuOpDescriptor {
        kind,
        name: descriptor.name,
        dtype_policy: descriptor.dtype_policy,
        launch,
    })
}

pub(crate) fn require_gpu_descriptor(
    kind: PrimitiveOpKind,
    expected: GpuLaunchKind,
) -> crate::Result<GpuOpDescriptor> {
    let descriptor = gpu_descriptor(kind).ok_or_else(|| {
        crate::Error::unsupported(
            "gpu_dispatch",
            format!("primitive {kind:?} is not implemented by CubeCL dispatch"),
        )
    })?;
    debug_assert_eq!(descriptor.launch, expected);
    Ok(descriptor)
}

#[cfg(test)]
mod tests;