tenferro_gpu/cubecl/
capability.rs

1use tenferro_core_ops::PrimitiveOpKind;
2use tenferro_tensor::{
3    BackendId, DType, OperationCapability, SupportLevel, TensorBackendCapability,
4};
5
6use super::CudaBackend;
7
8const U: SupportLevel = SupportLevel::Unsupported;
9const F: SupportLevel = SupportLevel::FallbackCopy;
10const N: SupportLevel = SupportLevel::Native;
11
12/// Return the CUDA backend operation capability descriptor table.
13///
14/// # Examples
15///
16/// ```rust
17/// use tenferro_core_ops::PrimitiveOpKind;
18/// use tenferro_gpu::cuda::cuda_capabilities;
19/// use tenferro_tensor::{DType, SupportLevel};
20///
21/// let add = cuda_capabilities()
22///     .iter()
23///     .find(|entry| entry.op == PrimitiveOpKind::Add && entry.dtype == DType::I32)
24///     .unwrap();
25/// assert_eq!(add.result, SupportLevel::Native);
26/// ```
27#[must_use]
28pub const fn cuda_capabilities() -> &'static [OperationCapability] {
29    CUDA_CAPABILITIES
30}
31
32impl TensorBackendCapability for CudaBackend {
33    fn backend_id(&self) -> BackendId {
34        BackendId::Cuda
35    }
36
37    fn capabilities(&self) -> &'static [OperationCapability] {
38        CUDA_CAPABILITIES
39    }
40}
41
42const fn owned_level(
43    op: PrimitiveOpKind,
44    dtype: DType,
45    output_dtype: DType,
46    level: SupportLevel,
47) -> OperationCapability {
48    OperationCapability {
49        backend: BackendId::Cuda,
50        op,
51        dtype,
52        output_dtype,
53        result: level,
54        read_inputs: if matches!(level, SupportLevel::Native) {
55            F
56        } else {
57            U
58        },
59        write_output: U,
60        strided_output: U,
61        accumulation: U,
62    }
63}
64
65const fn same_dtype(op: PrimitiveOpKind, dtype: DType, level: SupportLevel) -> OperationCapability {
66    owned_level(op, dtype, dtype, level)
67}
68
69const fn dot_native(dtype: DType) -> OperationCapability {
70    OperationCapability {
71        backend: BackendId::Cuda,
72        op: PrimitiveOpKind::DotGeneral,
73        dtype,
74        output_dtype: dtype,
75        result: N,
76        read_inputs: N,
77        write_output: N,
78        strided_output: N,
79        accumulation: N,
80    }
81}
82
83const CUDA_CAPABILITIES: &[OperationCapability] = &[
84    same_dtype(PrimitiveOpKind::Add, DType::F32, N),
85    same_dtype(PrimitiveOpKind::Add, DType::F64, N),
86    same_dtype(PrimitiveOpKind::Add, DType::I32, N),
87    same_dtype(PrimitiveOpKind::Add, DType::I64, N),
88    same_dtype(PrimitiveOpKind::Add, DType::C32, N),
89    same_dtype(PrimitiveOpKind::Add, DType::C64, N),
90    same_dtype(PrimitiveOpKind::Sub, DType::F32, N),
91    same_dtype(PrimitiveOpKind::Sub, DType::F64, N),
92    same_dtype(PrimitiveOpKind::Sub, DType::I32, N),
93    same_dtype(PrimitiveOpKind::Sub, DType::I64, N),
94    same_dtype(PrimitiveOpKind::Sub, DType::C32, N),
95    same_dtype(PrimitiveOpKind::Sub, DType::C64, N),
96    same_dtype(PrimitiveOpKind::Mul, DType::F32, N),
97    same_dtype(PrimitiveOpKind::Mul, DType::F64, N),
98    same_dtype(PrimitiveOpKind::Mul, DType::I32, N),
99    same_dtype(PrimitiveOpKind::Mul, DType::I64, N),
100    same_dtype(PrimitiveOpKind::Mul, DType::C32, N),
101    same_dtype(PrimitiveOpKind::Mul, DType::C64, N),
102    same_dtype(PrimitiveOpKind::Neg, DType::F32, N),
103    same_dtype(PrimitiveOpKind::Neg, DType::F64, N),
104    same_dtype(PrimitiveOpKind::Neg, DType::I32, N),
105    same_dtype(PrimitiveOpKind::Neg, DType::I64, N),
106    same_dtype(PrimitiveOpKind::Neg, DType::C32, N),
107    same_dtype(PrimitiveOpKind::Neg, DType::C64, N),
108    same_dtype(PrimitiveOpKind::Conj, DType::F32, N),
109    same_dtype(PrimitiveOpKind::Conj, DType::F64, N),
110    same_dtype(PrimitiveOpKind::Conj, DType::C32, N),
111    same_dtype(PrimitiveOpKind::Conj, DType::C64, N),
112    same_dtype(PrimitiveOpKind::Div, DType::F32, N),
113    same_dtype(PrimitiveOpKind::Div, DType::F64, N),
114    same_dtype(PrimitiveOpKind::Div, DType::I32, N),
115    same_dtype(PrimitiveOpKind::Div, DType::I64, N),
116    same_dtype(PrimitiveOpKind::Div, DType::C32, N),
117    same_dtype(PrimitiveOpKind::Div, DType::C64, N),
118    same_dtype(PrimitiveOpKind::Rem, DType::F32, N),
119    same_dtype(PrimitiveOpKind::Rem, DType::F64, N),
120    same_dtype(PrimitiveOpKind::Rem, DType::I32, N),
121    same_dtype(PrimitiveOpKind::Rem, DType::I64, N),
122    same_dtype(PrimitiveOpKind::Abs, DType::F32, N),
123    same_dtype(PrimitiveOpKind::Abs, DType::F64, N),
124    same_dtype(PrimitiveOpKind::Abs, DType::I32, N),
125    same_dtype(PrimitiveOpKind::Abs, DType::I64, N),
126    owned_level(PrimitiveOpKind::Abs, DType::C32, DType::F32, N),
127    owned_level(PrimitiveOpKind::Abs, DType::C64, DType::F64, N),
128    same_dtype(PrimitiveOpKind::Sign, DType::F32, N),
129    same_dtype(PrimitiveOpKind::Sign, DType::F64, N),
130    same_dtype(PrimitiveOpKind::Sign, DType::I32, N),
131    same_dtype(PrimitiveOpKind::Sign, DType::I64, N),
132    same_dtype(PrimitiveOpKind::Maximum, DType::F32, N),
133    same_dtype(PrimitiveOpKind::Maximum, DType::F64, N),
134    same_dtype(PrimitiveOpKind::Maximum, DType::I32, N),
135    same_dtype(PrimitiveOpKind::Maximum, DType::I64, N),
136    same_dtype(PrimitiveOpKind::Minimum, DType::F32, N),
137    same_dtype(PrimitiveOpKind::Minimum, DType::F64, N),
138    same_dtype(PrimitiveOpKind::Minimum, DType::I32, N),
139    same_dtype(PrimitiveOpKind::Minimum, DType::I64, N),
140    owned_level(PrimitiveOpKind::Compare, DType::F32, DType::Bool, N),
141    owned_level(PrimitiveOpKind::Compare, DType::F64, DType::Bool, N),
142    owned_level(PrimitiveOpKind::Compare, DType::I32, DType::Bool, N),
143    owned_level(PrimitiveOpKind::Compare, DType::I64, DType::Bool, N),
144    owned_level(PrimitiveOpKind::Compare, DType::Bool, DType::Bool, U),
145    same_dtype(PrimitiveOpKind::Select, DType::F32, N),
146    same_dtype(PrimitiveOpKind::Select, DType::F64, N),
147    same_dtype(PrimitiveOpKind::Select, DType::I32, N),
148    same_dtype(PrimitiveOpKind::Select, DType::I64, N),
149    same_dtype(PrimitiveOpKind::Select, DType::Bool, U),
150    same_dtype(PrimitiveOpKind::Select, DType::C32, U),
151    same_dtype(PrimitiveOpKind::Select, DType::C64, U),
152    same_dtype(PrimitiveOpKind::Clamp, DType::F32, N),
153    same_dtype(PrimitiveOpKind::Clamp, DType::F64, N),
154    same_dtype(PrimitiveOpKind::Exp, DType::F32, N),
155    same_dtype(PrimitiveOpKind::Exp, DType::F64, N),
156    same_dtype(PrimitiveOpKind::Exp, DType::C32, U),
157    same_dtype(PrimitiveOpKind::Exp, DType::C64, U),
158    same_dtype(PrimitiveOpKind::Log, DType::F32, N),
159    same_dtype(PrimitiveOpKind::Log, DType::F64, N),
160    same_dtype(PrimitiveOpKind::Log, DType::C32, U),
161    same_dtype(PrimitiveOpKind::Log, DType::C64, U),
162    same_dtype(PrimitiveOpKind::Sin, DType::F32, N),
163    same_dtype(PrimitiveOpKind::Sin, DType::F64, N),
164    same_dtype(PrimitiveOpKind::Sin, DType::C32, U),
165    same_dtype(PrimitiveOpKind::Sin, DType::C64, U),
166    same_dtype(PrimitiveOpKind::Cos, DType::F32, N),
167    same_dtype(PrimitiveOpKind::Cos, DType::F64, N),
168    same_dtype(PrimitiveOpKind::Cos, DType::C32, U),
169    same_dtype(PrimitiveOpKind::Cos, DType::C64, U),
170    same_dtype(PrimitiveOpKind::Tanh, DType::F32, N),
171    same_dtype(PrimitiveOpKind::Tanh, DType::F64, N),
172    same_dtype(PrimitiveOpKind::Tanh, DType::C32, U),
173    same_dtype(PrimitiveOpKind::Tanh, DType::C64, U),
174    same_dtype(PrimitiveOpKind::Sqrt, DType::F32, N),
175    same_dtype(PrimitiveOpKind::Sqrt, DType::F64, N),
176    same_dtype(PrimitiveOpKind::Sqrt, DType::C32, U),
177    same_dtype(PrimitiveOpKind::Sqrt, DType::C64, U),
178    same_dtype(PrimitiveOpKind::Rsqrt, DType::F32, N),
179    same_dtype(PrimitiveOpKind::Rsqrt, DType::F64, N),
180    same_dtype(PrimitiveOpKind::Rsqrt, DType::C32, U),
181    same_dtype(PrimitiveOpKind::Rsqrt, DType::C64, U),
182    same_dtype(PrimitiveOpKind::Pow, DType::F32, N),
183    same_dtype(PrimitiveOpKind::Pow, DType::F64, N),
184    same_dtype(PrimitiveOpKind::Pow, DType::I32, N),
185    same_dtype(PrimitiveOpKind::Pow, DType::I64, N),
186    same_dtype(PrimitiveOpKind::Pow, DType::C32, U),
187    same_dtype(PrimitiveOpKind::Pow, DType::C64, U),
188    same_dtype(PrimitiveOpKind::Expm1, DType::F32, N),
189    same_dtype(PrimitiveOpKind::Expm1, DType::F64, N),
190    same_dtype(PrimitiveOpKind::Expm1, DType::C32, U),
191    same_dtype(PrimitiveOpKind::Expm1, DType::C64, U),
192    same_dtype(PrimitiveOpKind::Log1p, DType::F32, N),
193    same_dtype(PrimitiveOpKind::Log1p, DType::F64, N),
194    same_dtype(PrimitiveOpKind::Log1p, DType::C32, U),
195    same_dtype(PrimitiveOpKind::Log1p, DType::C64, U),
196    same_dtype(PrimitiveOpKind::ReduceSum, DType::F32, N),
197    same_dtype(PrimitiveOpKind::ReduceSum, DType::F64, N),
198    same_dtype(PrimitiveOpKind::ReduceSum, DType::I32, N),
199    same_dtype(PrimitiveOpKind::ReduceSum, DType::I64, N),
200    same_dtype(PrimitiveOpKind::ReduceSum, DType::C32, N),
201    same_dtype(PrimitiveOpKind::ReduceSum, DType::C64, N),
202    same_dtype(PrimitiveOpKind::ReduceSumSquares, DType::F32, N),
203    same_dtype(PrimitiveOpKind::ReduceSumSquares, DType::F64, N),
204    same_dtype(PrimitiveOpKind::ReduceProd, DType::F32, N),
205    same_dtype(PrimitiveOpKind::ReduceProd, DType::F64, N),
206    same_dtype(PrimitiveOpKind::ReduceProd, DType::I32, N),
207    same_dtype(PrimitiveOpKind::ReduceProd, DType::I64, N),
208    same_dtype(PrimitiveOpKind::ReduceProd, DType::C32, N),
209    same_dtype(PrimitiveOpKind::ReduceProd, DType::C64, N),
210    same_dtype(PrimitiveOpKind::ReduceMax, DType::F32, N),
211    same_dtype(PrimitiveOpKind::ReduceMax, DType::F64, N),
212    same_dtype(PrimitiveOpKind::ReduceMax, DType::I32, N),
213    same_dtype(PrimitiveOpKind::ReduceMax, DType::I64, N),
214    same_dtype(PrimitiveOpKind::ReduceMin, DType::F32, N),
215    same_dtype(PrimitiveOpKind::ReduceMin, DType::F64, N),
216    same_dtype(PrimitiveOpKind::ReduceMin, DType::I32, N),
217    same_dtype(PrimitiveOpKind::ReduceMin, DType::I64, N),
218    dot_native(DType::F32),
219    dot_native(DType::F64),
220    dot_native(DType::C32),
221    dot_native(DType::C64),
222];
tenferro_gpu/cubecl/capability.rs

tenferro_gpu/cubecl/
capability.rs
