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
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222

use tenferro_core_ops::PrimitiveOpKind;
use tenferro_tensor::{
    BackendId, DType, OperationCapability, SupportLevel, TensorBackendCapability,
};

use super::CudaBackend;

const U: SupportLevel = SupportLevel::Unsupported;
const F: SupportLevel = SupportLevel::FallbackCopy;
const N: SupportLevel = SupportLevel::Native;

/// Return the CUDA backend operation capability descriptor table.
///
/// # Examples
///
/// ```rust
/// use tenferro_core_ops::PrimitiveOpKind;
/// use tenferro_gpu::cuda::cuda_capabilities;
/// use tenferro_tensor::{DType, SupportLevel};
///
/// let add = cuda_capabilities()
///     .iter()
///     .find(|entry| entry.op == PrimitiveOpKind::Add && entry.dtype == DType::I32)
///     .unwrap();
/// assert_eq!(add.result, SupportLevel::Native);
/// ```
#[must_use]
pub const fn cuda_capabilities() -> &'static [OperationCapability] {
    CUDA_CAPABILITIES
}

impl TensorBackendCapability for CudaBackend {
    fn backend_id(&self) -> BackendId {
        BackendId::Cuda
    }

    fn capabilities(&self) -> &'static [OperationCapability] {
        CUDA_CAPABILITIES
    }
}

const fn owned_level(
    op: PrimitiveOpKind,
    dtype: DType,
    output_dtype: DType,
    level: SupportLevel,
) -> OperationCapability {
    OperationCapability {
        backend: BackendId::Cuda,
        op,
        dtype,
        output_dtype,
        result: level,
        read_inputs: if matches!(level, SupportLevel::Native) {
            F
        } else {
            U
        },
        write_output: U,
        strided_output: U,
        accumulation: U,
    }
}

const fn same_dtype(op: PrimitiveOpKind, dtype: DType, level: SupportLevel) -> OperationCapability {
    owned_level(op, dtype, dtype, level)
}

const fn dot_native(dtype: DType) -> OperationCapability {
    OperationCapability {
        backend: BackendId::Cuda,
        op: PrimitiveOpKind::DotGeneral,
        dtype,
        output_dtype: dtype,
        result: N,
        read_inputs: N,
        write_output: N,
        strided_output: N,
        accumulation: N,
    }
}

const CUDA_CAPABILITIES: &[OperationCapability] = &[
    same_dtype(PrimitiveOpKind::Add, DType::F32, N),
    same_dtype(PrimitiveOpKind::Add, DType::F64, N),
    same_dtype(PrimitiveOpKind::Add, DType::I32, N),
    same_dtype(PrimitiveOpKind::Add, DType::I64, N),
    same_dtype(PrimitiveOpKind::Add, DType::C32, N),
    same_dtype(PrimitiveOpKind::Add, DType::C64, N),
    same_dtype(PrimitiveOpKind::Sub, DType::F32, N),
    same_dtype(PrimitiveOpKind::Sub, DType::F64, N),
    same_dtype(PrimitiveOpKind::Sub, DType::I32, N),
    same_dtype(PrimitiveOpKind::Sub, DType::I64, N),
    same_dtype(PrimitiveOpKind::Sub, DType::C32, N),
    same_dtype(PrimitiveOpKind::Sub, DType::C64, N),
    same_dtype(PrimitiveOpKind::Mul, DType::F32, N),
    same_dtype(PrimitiveOpKind::Mul, DType::F64, N),
    same_dtype(PrimitiveOpKind::Mul, DType::I32, N),
    same_dtype(PrimitiveOpKind::Mul, DType::I64, N),
    same_dtype(PrimitiveOpKind::Mul, DType::C32, N),
    same_dtype(PrimitiveOpKind::Mul, DType::C64, N),
    same_dtype(PrimitiveOpKind::Neg, DType::F32, N),
    same_dtype(PrimitiveOpKind::Neg, DType::F64, N),
    same_dtype(PrimitiveOpKind::Neg, DType::I32, N),
    same_dtype(PrimitiveOpKind::Neg, DType::I64, N),
    same_dtype(PrimitiveOpKind::Neg, DType::C32, N),
    same_dtype(PrimitiveOpKind::Neg, DType::C64, N),
    same_dtype(PrimitiveOpKind::Conj, DType::F32, N),
    same_dtype(PrimitiveOpKind::Conj, DType::F64, N),
    same_dtype(PrimitiveOpKind::Conj, DType::C32, N),
    same_dtype(PrimitiveOpKind::Conj, DType::C64, N),
    same_dtype(PrimitiveOpKind::Div, DType::F32, N),
    same_dtype(PrimitiveOpKind::Div, DType::F64, N),
    same_dtype(PrimitiveOpKind::Div, DType::I32, N),
    same_dtype(PrimitiveOpKind::Div, DType::I64, N),
    same_dtype(PrimitiveOpKind::Div, DType::C32, N),
    same_dtype(PrimitiveOpKind::Div, DType::C64, N),
    same_dtype(PrimitiveOpKind::Rem, DType::F32, N),
    same_dtype(PrimitiveOpKind::Rem, DType::F64, N),
    same_dtype(PrimitiveOpKind::Rem, DType::I32, N),
    same_dtype(PrimitiveOpKind::Rem, DType::I64, N),
    same_dtype(PrimitiveOpKind::Abs, DType::F32, N),
    same_dtype(PrimitiveOpKind::Abs, DType::F64, N),
    same_dtype(PrimitiveOpKind::Abs, DType::I32, N),
    same_dtype(PrimitiveOpKind::Abs, DType::I64, N),
    owned_level(PrimitiveOpKind::Abs, DType::C32, DType::F32, N),
    owned_level(PrimitiveOpKind::Abs, DType::C64, DType::F64, N),
    same_dtype(PrimitiveOpKind::Sign, DType::F32, N),
    same_dtype(PrimitiveOpKind::Sign, DType::F64, N),
    same_dtype(PrimitiveOpKind::Sign, DType::I32, N),
    same_dtype(PrimitiveOpKind::Sign, DType::I64, N),
    same_dtype(PrimitiveOpKind::Maximum, DType::F32, N),
    same_dtype(PrimitiveOpKind::Maximum, DType::F64, N),
    same_dtype(PrimitiveOpKind::Maximum, DType::I32, N),
    same_dtype(PrimitiveOpKind::Maximum, DType::I64, N),
    same_dtype(PrimitiveOpKind::Minimum, DType::F32, N),
    same_dtype(PrimitiveOpKind::Minimum, DType::F64, N),
    same_dtype(PrimitiveOpKind::Minimum, DType::I32, N),
    same_dtype(PrimitiveOpKind::Minimum, DType::I64, N),
    owned_level(PrimitiveOpKind::Compare, DType::F32, DType::Bool, N),
    owned_level(PrimitiveOpKind::Compare, DType::F64, DType::Bool, N),
    owned_level(PrimitiveOpKind::Compare, DType::I32, DType::Bool, N),
    owned_level(PrimitiveOpKind::Compare, DType::I64, DType::Bool, N),
    owned_level(PrimitiveOpKind::Compare, DType::Bool, DType::Bool, U),
    same_dtype(PrimitiveOpKind::Select, DType::F32, N),
    same_dtype(PrimitiveOpKind::Select, DType::F64, N),
    same_dtype(PrimitiveOpKind::Select, DType::I32, N),
    same_dtype(PrimitiveOpKind::Select, DType::I64, N),
    same_dtype(PrimitiveOpKind::Select, DType::Bool, U),
    same_dtype(PrimitiveOpKind::Select, DType::C32, U),
    same_dtype(PrimitiveOpKind::Select, DType::C64, U),
    same_dtype(PrimitiveOpKind::Clamp, DType::F32, N),
    same_dtype(PrimitiveOpKind::Clamp, DType::F64, N),
    same_dtype(PrimitiveOpKind::Exp, DType::F32, N),
    same_dtype(PrimitiveOpKind::Exp, DType::F64, N),
    same_dtype(PrimitiveOpKind::Exp, DType::C32, U),
    same_dtype(PrimitiveOpKind::Exp, DType::C64, U),
    same_dtype(PrimitiveOpKind::Log, DType::F32, N),
    same_dtype(PrimitiveOpKind::Log, DType::F64, N),
    same_dtype(PrimitiveOpKind::Log, DType::C32, U),
    same_dtype(PrimitiveOpKind::Log, DType::C64, U),
    same_dtype(PrimitiveOpKind::Sin, DType::F32, N),
    same_dtype(PrimitiveOpKind::Sin, DType::F64, N),
    same_dtype(PrimitiveOpKind::Sin, DType::C32, U),
    same_dtype(PrimitiveOpKind::Sin, DType::C64, U),
    same_dtype(PrimitiveOpKind::Cos, DType::F32, N),
    same_dtype(PrimitiveOpKind::Cos, DType::F64, N),
    same_dtype(PrimitiveOpKind::Cos, DType::C32, U),
    same_dtype(PrimitiveOpKind::Cos, DType::C64, U),
    same_dtype(PrimitiveOpKind::Tanh, DType::F32, N),
    same_dtype(PrimitiveOpKind::Tanh, DType::F64, N),
    same_dtype(PrimitiveOpKind::Tanh, DType::C32, U),
    same_dtype(PrimitiveOpKind::Tanh, DType::C64, U),
    same_dtype(PrimitiveOpKind::Sqrt, DType::F32, N),
    same_dtype(PrimitiveOpKind::Sqrt, DType::F64, N),
    same_dtype(PrimitiveOpKind::Sqrt, DType::C32, U),
    same_dtype(PrimitiveOpKind::Sqrt, DType::C64, U),
    same_dtype(PrimitiveOpKind::Rsqrt, DType::F32, N),
    same_dtype(PrimitiveOpKind::Rsqrt, DType::F64, N),
    same_dtype(PrimitiveOpKind::Rsqrt, DType::C32, U),
    same_dtype(PrimitiveOpKind::Rsqrt, DType::C64, U),
    same_dtype(PrimitiveOpKind::Pow, DType::F32, N),
    same_dtype(PrimitiveOpKind::Pow, DType::F64, N),
    same_dtype(PrimitiveOpKind::Pow, DType::I32, N),
    same_dtype(PrimitiveOpKind::Pow, DType::I64, N),
    same_dtype(PrimitiveOpKind::Pow, DType::C32, U),
    same_dtype(PrimitiveOpKind::Pow, DType::C64, U),
    same_dtype(PrimitiveOpKind::Expm1, DType::F32, N),
    same_dtype(PrimitiveOpKind::Expm1, DType::F64, N),
    same_dtype(PrimitiveOpKind::Expm1, DType::C32, U),
    same_dtype(PrimitiveOpKind::Expm1, DType::C64, U),
    same_dtype(PrimitiveOpKind::Log1p, DType::F32, N),
    same_dtype(PrimitiveOpKind::Log1p, DType::F64, N),
    same_dtype(PrimitiveOpKind::Log1p, DType::C32, U),
    same_dtype(PrimitiveOpKind::Log1p, DType::C64, U),
    same_dtype(PrimitiveOpKind::ReduceSum, DType::F32, N),
    same_dtype(PrimitiveOpKind::ReduceSum, DType::F64, N),
    same_dtype(PrimitiveOpKind::ReduceSum, DType::I32, N),
    same_dtype(PrimitiveOpKind::ReduceSum, DType::I64, N),
    same_dtype(PrimitiveOpKind::ReduceSum, DType::C32, N),
    same_dtype(PrimitiveOpKind::ReduceSum, DType::C64, N),
    same_dtype(PrimitiveOpKind::ReduceSumSquares, DType::F32, N),
    same_dtype(PrimitiveOpKind::ReduceSumSquares, DType::F64, N),
    same_dtype(PrimitiveOpKind::ReduceProd, DType::F32, N),
    same_dtype(PrimitiveOpKind::ReduceProd, DType::F64, N),
    same_dtype(PrimitiveOpKind::ReduceProd, DType::I32, N),
    same_dtype(PrimitiveOpKind::ReduceProd, DType::I64, N),
    same_dtype(PrimitiveOpKind::ReduceProd, DType::C32, N),
    same_dtype(PrimitiveOpKind::ReduceProd, DType::C64, N),
    same_dtype(PrimitiveOpKind::ReduceMax, DType::F32, N),
    same_dtype(PrimitiveOpKind::ReduceMax, DType::F64, N),
    same_dtype(PrimitiveOpKind::ReduceMax, DType::I32, N),
    same_dtype(PrimitiveOpKind::ReduceMax, DType::I64, N),
    same_dtype(PrimitiveOpKind::ReduceMin, DType::F32, N),
    same_dtype(PrimitiveOpKind::ReduceMin, DType::F64, N),
    same_dtype(PrimitiveOpKind::ReduceMin, DType::I32, N),
    same_dtype(PrimitiveOpKind::ReduceMin, DType::I64, N),
    dot_native(DType::F32),
    dot_native(DType::F64),
    dot_native(DType::C32),
    dot_native(DType::C64),
];