tenferro_cpu/
capability.rs

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

tenferro_cpu/
capability.rs
