tenferro_gpu/cuda/
capability.rs

1use tenferro_core_ops::PrimitiveOpKind;
2use tenferro_tensor::{
3    BackendId, DType, OperationCapability, SupportLevel, TensorBackendCapability,
4};
5
6use super::CudaEngine;
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/// The table states what this backend serves natively; it is written down rather than derived
15/// from the kernel catalogue's `accepts` list, and `tests/seam_capability.rs` holds each entry to
16/// the engine's observed behaviour.
17///
18/// # Examples
19///
20/// ```rust
21/// use tenferro_core_ops::PrimitiveOpKind;
22/// use tenferro_gpu::cuda::cuda_capabilities;
23/// use tenferro_tensor::{DType, SupportLevel};
24///
25/// let add = cuda_capabilities()
26///     .iter()
27///     .find(|entry| entry.op == PrimitiveOpKind::Add && entry.dtype == DType::I32)
28///     .unwrap();
29/// assert_eq!(add.result, SupportLevel::Native);
30/// ```
31#[must_use]
32pub const fn cuda_capabilities() -> &'static [OperationCapability] {
33    CUDA_CAPABILITIES
34}
35
36impl TensorBackendCapability for CudaEngine {
37    fn backend_id(&self) -> BackendId {
38        BackendId::Cuda
39    }
40
41    fn capabilities(&self) -> &'static [OperationCapability] {
42        CUDA_CAPABILITIES
43    }
44}
45
46const fn owned_level(
47    op: PrimitiveOpKind,
48    dtype: DType,
49    output_dtype: DType,
50    level: SupportLevel,
51) -> OperationCapability {
52    OperationCapability {
53        backend: BackendId::Cuda,
54        op,
55        dtype,
56        output_dtype,
57        result: level,
58        read_inputs: if matches!(level, SupportLevel::Native) {
59            F
60        } else {
61            U
62        },
63        write_output: U,
64        strided_output: U,
65        accumulation: U,
66    }
67}
68
69const fn same_dtype(op: PrimitiveOpKind, dtype: DType, level: SupportLevel) -> OperationCapability {
70    owned_level(op, dtype, dtype, level)
71}
72
73const fn dot_native(dtype: DType) -> OperationCapability {
74    OperationCapability {
75        backend: BackendId::Cuda,
76        op: PrimitiveOpKind::DotGeneral,
77        dtype,
78        output_dtype: dtype,
79        result: N,
80        // A borrowed operand reaches cuBLAS through a compact copy: cuBLAS takes no strides.
81        read_inputs: F,
82        write_output: N,
83        // The dotted result is written into an owned output only; a strided destination region
84        // is refused by both `dot_general_read_into` entries.
85        strided_output: U,
86        accumulation: N,
87    }
88}
89
90const CUDA_CAPABILITIES: &[OperationCapability] = &[
91    same_dtype(PrimitiveOpKind::Add, DType::F32, N),
92    same_dtype(PrimitiveOpKind::Add, DType::F64, N),
93    same_dtype(PrimitiveOpKind::Add, DType::I32, N),
94    same_dtype(PrimitiveOpKind::Add, DType::I64, N),
95    same_dtype(PrimitiveOpKind::Add, DType::C32, N),
96    same_dtype(PrimitiveOpKind::Add, DType::C64, N),
97    same_dtype(PrimitiveOpKind::Sub, DType::F32, N),
98    same_dtype(PrimitiveOpKind::Sub, DType::F64, N),
99    same_dtype(PrimitiveOpKind::Sub, DType::I32, N),
100    same_dtype(PrimitiveOpKind::Sub, DType::I64, N),
101    same_dtype(PrimitiveOpKind::Sub, DType::C32, N),
102    same_dtype(PrimitiveOpKind::Sub, DType::C64, N),
103    same_dtype(PrimitiveOpKind::Mul, DType::F32, N),
104    same_dtype(PrimitiveOpKind::Mul, DType::F64, N),
105    same_dtype(PrimitiveOpKind::Mul, DType::I32, N),
106    same_dtype(PrimitiveOpKind::Mul, DType::I64, N),
107    same_dtype(PrimitiveOpKind::Mul, DType::C32, N),
108    same_dtype(PrimitiveOpKind::Mul, DType::C64, N),
109    same_dtype(PrimitiveOpKind::Neg, DType::F32, N),
110    same_dtype(PrimitiveOpKind::Neg, DType::F64, N),
111    same_dtype(PrimitiveOpKind::Neg, DType::I32, N),
112    same_dtype(PrimitiveOpKind::Neg, DType::I64, N),
113    same_dtype(PrimitiveOpKind::Neg, DType::C32, N),
114    same_dtype(PrimitiveOpKind::Neg, DType::C64, N),
115    same_dtype(PrimitiveOpKind::Conj, DType::F32, N),
116    same_dtype(PrimitiveOpKind::Conj, DType::F64, N),
117    same_dtype(PrimitiveOpKind::Conj, DType::C32, N),
118    same_dtype(PrimitiveOpKind::Conj, DType::C64, N),
119    same_dtype(PrimitiveOpKind::Div, DType::F32, N),
120    same_dtype(PrimitiveOpKind::Div, DType::F64, N),
121    same_dtype(PrimitiveOpKind::Div, DType::I32, N),
122    same_dtype(PrimitiveOpKind::Div, DType::I64, N),
123    same_dtype(PrimitiveOpKind::Div, DType::C32, N),
124    same_dtype(PrimitiveOpKind::Div, DType::C64, N),
125    same_dtype(PrimitiveOpKind::Rem, DType::F32, N),
126    same_dtype(PrimitiveOpKind::Rem, DType::F64, N),
127    same_dtype(PrimitiveOpKind::Rem, DType::I32, N),
128    same_dtype(PrimitiveOpKind::Rem, DType::I64, N),
129    same_dtype(PrimitiveOpKind::Abs, DType::F32, N),
130    same_dtype(PrimitiveOpKind::Abs, DType::F64, N),
131    same_dtype(PrimitiveOpKind::Abs, DType::I32, N),
132    same_dtype(PrimitiveOpKind::Abs, DType::I64, N),
133    owned_level(PrimitiveOpKind::Abs, DType::C32, DType::F32, N),
134    owned_level(PrimitiveOpKind::Abs, DType::C64, DType::F64, N),
135    same_dtype(PrimitiveOpKind::Sign, DType::F32, N),
136    same_dtype(PrimitiveOpKind::Sign, DType::F64, N),
137    same_dtype(PrimitiveOpKind::Sign, DType::I32, N),
138    same_dtype(PrimitiveOpKind::Sign, DType::I64, N),
139    same_dtype(PrimitiveOpKind::Sign, DType::C32, N),
140    same_dtype(PrimitiveOpKind::Sign, DType::C64, N),
141    same_dtype(PrimitiveOpKind::Maximum, DType::F32, N),
142    same_dtype(PrimitiveOpKind::Maximum, DType::F64, N),
143    same_dtype(PrimitiveOpKind::Maximum, DType::I32, N),
144    same_dtype(PrimitiveOpKind::Maximum, DType::I64, N),
145    same_dtype(PrimitiveOpKind::Minimum, DType::F32, N),
146    same_dtype(PrimitiveOpKind::Minimum, DType::F64, N),
147    same_dtype(PrimitiveOpKind::Minimum, DType::I32, N),
148    same_dtype(PrimitiveOpKind::Minimum, DType::I64, N),
149    owned_level(PrimitiveOpKind::Compare, DType::F32, DType::Bool, N),
150    owned_level(PrimitiveOpKind::Compare, DType::F64, DType::Bool, N),
151    owned_level(PrimitiveOpKind::Compare, DType::I32, DType::Bool, N),
152    owned_level(PrimitiveOpKind::Compare, DType::I64, DType::Bool, N),
153    owned_level(PrimitiveOpKind::Compare, DType::Bool, DType::Bool, U),
154    same_dtype(PrimitiveOpKind::Select, DType::F32, N),
155    same_dtype(PrimitiveOpKind::Select, DType::F64, N),
156    same_dtype(PrimitiveOpKind::Select, DType::I32, N),
157    same_dtype(PrimitiveOpKind::Select, DType::I64, N),
158    same_dtype(PrimitiveOpKind::Select, DType::Bool, U),
159    same_dtype(PrimitiveOpKind::Select, DType::C32, U),
160    same_dtype(PrimitiveOpKind::Select, DType::C64, U),
161    same_dtype(PrimitiveOpKind::Clamp, DType::F32, N),
162    same_dtype(PrimitiveOpKind::Clamp, DType::F64, N),
163    same_dtype(PrimitiveOpKind::Exp, DType::F32, N),
164    same_dtype(PrimitiveOpKind::Exp, DType::F64, N),
165    same_dtype(PrimitiveOpKind::Exp, DType::C32, U),
166    same_dtype(PrimitiveOpKind::Exp, DType::C64, U),
167    same_dtype(PrimitiveOpKind::Log, DType::F32, N),
168    same_dtype(PrimitiveOpKind::Log, DType::F64, N),
169    same_dtype(PrimitiveOpKind::Log, DType::C32, U),
170    same_dtype(PrimitiveOpKind::Log, DType::C64, U),
171    same_dtype(PrimitiveOpKind::Sin, DType::F32, N),
172    same_dtype(PrimitiveOpKind::Sin, DType::F64, N),
173    same_dtype(PrimitiveOpKind::Sin, DType::C32, U),
174    same_dtype(PrimitiveOpKind::Sin, DType::C64, U),
175    same_dtype(PrimitiveOpKind::Cos, DType::F32, N),
176    same_dtype(PrimitiveOpKind::Cos, DType::F64, N),
177    same_dtype(PrimitiveOpKind::Cos, DType::C32, U),
178    same_dtype(PrimitiveOpKind::Cos, DType::C64, U),
179    same_dtype(PrimitiveOpKind::Tanh, DType::F32, N),
180    same_dtype(PrimitiveOpKind::Tanh, DType::F64, N),
181    same_dtype(PrimitiveOpKind::Tanh, DType::C32, U),
182    same_dtype(PrimitiveOpKind::Tanh, DType::C64, U),
183    same_dtype(PrimitiveOpKind::Sqrt, DType::F32, N),
184    same_dtype(PrimitiveOpKind::Sqrt, DType::F64, N),
185    same_dtype(PrimitiveOpKind::Sqrt, DType::C32, U),
186    same_dtype(PrimitiveOpKind::Sqrt, DType::C64, U),
187    same_dtype(PrimitiveOpKind::Rsqrt, DType::F32, N),
188    same_dtype(PrimitiveOpKind::Rsqrt, DType::F64, N),
189    same_dtype(PrimitiveOpKind::Rsqrt, DType::C32, U),
190    same_dtype(PrimitiveOpKind::Rsqrt, DType::C64, U),
191    same_dtype(PrimitiveOpKind::Pow, DType::F32, N),
192    same_dtype(PrimitiveOpKind::Pow, DType::F64, N),
193    same_dtype(PrimitiveOpKind::Pow, DType::I32, N),
194    same_dtype(PrimitiveOpKind::Pow, DType::I64, N),
195    same_dtype(PrimitiveOpKind::Pow, DType::C32, U),
196    same_dtype(PrimitiveOpKind::Pow, DType::C64, U),
197    same_dtype(PrimitiveOpKind::Expm1, DType::F32, N),
198    same_dtype(PrimitiveOpKind::Expm1, DType::F64, N),
199    same_dtype(PrimitiveOpKind::Expm1, DType::C32, U),
200    same_dtype(PrimitiveOpKind::Expm1, DType::C64, U),
201    same_dtype(PrimitiveOpKind::Log1p, DType::F32, N),
202    same_dtype(PrimitiveOpKind::Log1p, DType::F64, N),
203    same_dtype(PrimitiveOpKind::Log1p, DType::C32, U),
204    same_dtype(PrimitiveOpKind::Log1p, DType::C64, U),
205    same_dtype(PrimitiveOpKind::Erf, DType::F32, N),
206    same_dtype(PrimitiveOpKind::Erf, DType::F64, N),
207    same_dtype(PrimitiveOpKind::ReduceSum, DType::F32, N),
208    same_dtype(PrimitiveOpKind::ReduceSum, DType::F64, N),
209    same_dtype(PrimitiveOpKind::ReduceSum, DType::I32, N),
210    same_dtype(PrimitiveOpKind::ReduceSum, DType::I64, N),
211    same_dtype(PrimitiveOpKind::ReduceSum, DType::C32, N),
212    same_dtype(PrimitiveOpKind::ReduceSum, DType::C64, N),
213    same_dtype(PrimitiveOpKind::ReduceSumSquares, DType::F32, N),
214    same_dtype(PrimitiveOpKind::ReduceSumSquares, DType::F64, N),
215    same_dtype(PrimitiveOpKind::ReduceProd, DType::F32, N),
216    same_dtype(PrimitiveOpKind::ReduceProd, DType::F64, N),
217    same_dtype(PrimitiveOpKind::ReduceProd, DType::I32, N),
218    same_dtype(PrimitiveOpKind::ReduceProd, DType::I64, N),
219    same_dtype(PrimitiveOpKind::ReduceProd, DType::C32, N),
220    same_dtype(PrimitiveOpKind::ReduceProd, DType::C64, N),
221    same_dtype(PrimitiveOpKind::ReduceMax, DType::F32, N),
222    same_dtype(PrimitiveOpKind::ReduceMax, DType::F64, N),
223    same_dtype(PrimitiveOpKind::ReduceMax, DType::I32, N),
224    same_dtype(PrimitiveOpKind::ReduceMax, DType::I64, N),
225    same_dtype(PrimitiveOpKind::ReduceMin, DType::F32, N),
226    same_dtype(PrimitiveOpKind::ReduceMin, DType::F64, N),
227    same_dtype(PrimitiveOpKind::ReduceMin, DType::I32, N),
228    same_dtype(PrimitiveOpKind::ReduceMin, DType::I64, N),
229    dot_native(DType::F32),
230    dot_native(DType::F64),
231    dot_native(DType::C32),
232    dot_native(DType::C64),
233];
tenferro_gpu/cuda/capability.rs

tenferro_gpu/cuda/
capability.rs
