tenferro_tensor/
capability.rs

1//! Backend operation capability descriptors.
2//!
3//! # Examples
4//!
5//! ```rust
6//! use tenferro_core_ops::PrimitiveOpKind;
7//! use tenferro_tensor::{capability_output_dtype, DType};
8//!
9//! assert_eq!(
10//!     capability_output_dtype(PrimitiveOpKind::Compare, DType::F64),
11//!     Some(DType::Bool)
12//! );
13//! ```
14
15use std::fmt;
16
17use tenferro_core_ops::{descriptor, DTypePolicy, PrimitiveOpKind};
18
19use crate::DType;
20
21/// Stable backend identifier used by capability descriptors.
22///
23/// # Examples
24///
25/// ```rust
26/// use tenferro_tensor::BackendId;
27///
28/// assert_eq!(BackendId::Cuda.as_str(), "cuda");
29/// ```
30#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
31pub enum BackendId {
32    Cpu,
33    Cuda,
34    WebGpu,
35    Other(&'static str),
36}
37
38impl BackendId {
39    /// Return the stable backend name used in diagnostics and generated docs.
40    ///
41    /// # Examples
42    ///
43    /// ```rust
44    /// use tenferro_tensor::BackendId;
45    ///
46    /// assert_eq!(BackendId::Cpu.as_str(), "cpu");
47    /// ```
48    #[must_use]
49    pub const fn as_str(self) -> &'static str {
50        match self {
51            Self::Cpu => "cpu",
52            Self::Cuda => "cuda",
53            Self::WebGpu => "webgpu",
54            Self::Other(name) => name,
55        }
56    }
57}
58
59impl fmt::Display for BackendId {
60    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
61        f.write_str(self.as_str())
62    }
63}
64
65/// Three-valued support level for a backend capability axis.
66///
67/// `FallbackCopy` is intentionally distinct from `Native`: the operation is
68/// accepted, but only by materializing through a default/copy path.
69///
70/// # Examples
71///
72/// ```rust
73/// use tenferro_tensor::SupportLevel;
74///
75/// assert!(SupportLevel::Native > SupportLevel::FallbackCopy);
76/// assert!(SupportLevel::FallbackCopy.is_supported());
77/// assert!(!SupportLevel::Unsupported.is_supported());
78/// ```
79#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
80pub enum SupportLevel {
81    Unsupported,
82    FallbackCopy,
83    Native,
84}
85
86impl SupportLevel {
87    /// Return whether this level represents any usable implementation.
88    ///
89    /// # Examples
90    ///
91    /// ```rust
92    /// use tenferro_tensor::SupportLevel;
93    ///
94    /// assert!(SupportLevel::Native.is_supported());
95    /// assert!(!SupportLevel::Unsupported.is_supported());
96    /// ```
97    #[must_use]
98    pub const fn is_supported(self) -> bool {
99        !matches!(self, Self::Unsupported)
100    }
101}
102
103/// Axis within an operation capability entry.
104///
105/// # Examples
106///
107/// ```rust
108/// use tenferro_tensor::CapabilityAxis;
109///
110/// let axis = CapabilityAxis::ReadInputs;
111/// assert_eq!(format!("{axis:?}"), "ReadInputs");
112/// ```
113#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
114pub enum CapabilityAxis {
115    OwnedResult,
116    ReadInputs,
117    WriteOutput,
118    StridedOutput,
119    Accumulation,
120}
121
122/// Query key for a backend capability lookup.
123///
124/// # Examples
125///
126/// ```rust
127/// use tenferro_core_ops::PrimitiveOpKind;
128/// use tenferro_tensor::{CapabilityQuery, DType};
129///
130/// let query = CapabilityQuery::new(PrimitiveOpKind::Add, DType::F32);
131/// assert_eq!(query.dtype, DType::F32);
132/// ```
133#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
134pub struct CapabilityQuery {
135    pub op: PrimitiveOpKind,
136    pub dtype: DType,
137}
138
139impl CapabilityQuery {
140    /// Build a capability lookup key.
141    ///
142    /// # Examples
143    ///
144    /// ```rust
145    /// use tenferro_core_ops::PrimitiveOpKind;
146    /// use tenferro_tensor::{CapabilityQuery, DType};
147    ///
148    /// assert_eq!(
149    ///     CapabilityQuery::new(PrimitiveOpKind::Mul, DType::I64).op,
150    ///     PrimitiveOpKind::Mul
151    /// );
152    /// ```
153    #[must_use]
154    pub const fn new(op: PrimitiveOpKind, dtype: DType) -> Self {
155        Self { op, dtype }
156    }
157}
158
159/// One backend capability entry for a primitive op and input dtype.
160///
161/// # Examples
162///
163/// ```rust
164/// use tenferro_core_ops::PrimitiveOpKind;
165/// use tenferro_tensor::{
166///     BackendId, CapabilityAxis, DType, OperationCapability, SupportLevel,
167/// };
168///
169/// let entry = OperationCapability {
170///     backend: BackendId::Cpu,
171///     op: PrimitiveOpKind::Add,
172///     dtype: DType::F64,
173///     output_dtype: DType::F64,
174///     result: SupportLevel::Native,
175///     read_inputs: SupportLevel::Native,
176///     write_output: SupportLevel::Native,
177///     strided_output: SupportLevel::Native,
178///     accumulation: SupportLevel::Unsupported,
179/// };
180/// assert_eq!(entry.axis(CapabilityAxis::OwnedResult), SupportLevel::Native);
181/// ```
182#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
183pub struct OperationCapability {
184    pub backend: BackendId,
185    pub op: PrimitiveOpKind,
186    pub dtype: DType,
187    pub output_dtype: DType,
188    pub result: SupportLevel,
189    pub read_inputs: SupportLevel,
190    pub write_output: SupportLevel,
191    pub strided_output: SupportLevel,
192    pub accumulation: SupportLevel,
193}
194
195impl OperationCapability {
196    /// Return the support level for one capability axis.
197    ///
198    /// # Examples
199    ///
200    /// ```rust
201    /// use tenferro_core_ops::PrimitiveOpKind;
202    /// use tenferro_tensor::{
203    ///     BackendId, CapabilityAxis, DType, OperationCapability, SupportLevel,
204    /// };
205    ///
206    /// let entry = OperationCapability {
207    ///     backend: BackendId::Cuda,
208    ///     op: PrimitiveOpKind::ReduceSum,
209    ///     dtype: DType::I32,
210    ///     output_dtype: DType::I32,
211    ///     result: SupportLevel::Native,
212    ///     read_inputs: SupportLevel::FallbackCopy,
213    ///     write_output: SupportLevel::Unsupported,
214    ///     strided_output: SupportLevel::Unsupported,
215    ///     accumulation: SupportLevel::Unsupported,
216    /// };
217    /// assert_eq!(entry.axis(CapabilityAxis::ReadInputs), SupportLevel::FallbackCopy);
218    /// ```
219    #[must_use]
220    pub const fn axis(&self, axis: CapabilityAxis) -> SupportLevel {
221        match axis {
222            CapabilityAxis::OwnedResult => self.result,
223            CapabilityAxis::ReadInputs => self.read_inputs,
224            CapabilityAxis::WriteOutput => self.write_output,
225            CapabilityAxis::StridedOutput => self.strided_output,
226            CapabilityAxis::Accumulation => self.accumulation,
227        }
228    }
229}
230
231/// Backend capability query surface.
232///
233/// # Examples
234///
235/// ```rust
236/// use tenferro_core_ops::PrimitiveOpKind;
237/// use tenferro_tensor::{
238///     BackendId, CapabilityQuery, DType, OperationCapability, SupportLevel,
239///     TensorBackendCapability,
240/// };
241///
242/// struct Backend;
243///
244/// const ENTRIES: &[OperationCapability] = &[OperationCapability {
245///     backend: BackendId::Cpu,
246///     op: PrimitiveOpKind::Add,
247///     dtype: DType::F32,
248///     output_dtype: DType::F32,
249///     result: SupportLevel::Native,
250///     read_inputs: SupportLevel::Native,
251///     write_output: SupportLevel::Native,
252///     strided_output: SupportLevel::Native,
253///     accumulation: SupportLevel::Unsupported,
254/// }];
255///
256/// impl TensorBackendCapability for Backend {
257///     fn backend_id(&self) -> BackendId { BackendId::Cpu }
258///     fn capabilities(&self) -> &'static [OperationCapability] { ENTRIES }
259/// }
260///
261/// assert!(Backend
262///     .capability(CapabilityQuery::new(PrimitiveOpKind::Add, DType::F32))
263///     .is_some());
264/// ```
265pub trait TensorBackendCapability {
266    fn backend_id(&self) -> BackendId;
267    fn capabilities(&self) -> &'static [OperationCapability];
268
269    /// Look up one operation/dtype capability for this backend.
270    ///
271    /// # Examples
272    ///
273    /// ```rust
274    /// use tenferro_core_ops::PrimitiveOpKind;
275    /// use tenferro_tensor::{
276    ///     BackendId, CapabilityQuery, DType, OperationCapability, SupportLevel,
277    ///     TensorBackendCapability,
278    /// };
279    ///
280    /// struct Backend;
281    /// const ENTRIES: &[OperationCapability] = &[OperationCapability {
282    ///     backend: BackendId::Cpu,
283    ///     op: PrimitiveOpKind::Mul,
284    ///     dtype: DType::I64,
285    ///     output_dtype: DType::I64,
286    ///     result: SupportLevel::Native,
287    ///     read_inputs: SupportLevel::Native,
288    ///     write_output: SupportLevel::Unsupported,
289    ///     strided_output: SupportLevel::Unsupported,
290    ///     accumulation: SupportLevel::Unsupported,
291    /// }];
292    /// impl TensorBackendCapability for Backend {
293    ///     fn backend_id(&self) -> BackendId { BackendId::Cpu }
294    ///     fn capabilities(&self) -> &'static [OperationCapability] { ENTRIES }
295    /// }
296    ///
297    /// let entry = Backend
298    ///     .capability(CapabilityQuery::new(PrimitiveOpKind::Mul, DType::I64))
299    ///     .unwrap();
300    /// assert_eq!(entry.result, SupportLevel::Native);
301    /// ```
302    #[must_use]
303    fn capability(&self, query: CapabilityQuery) -> Option<OperationCapability> {
304        self.capabilities().iter().copied().find(|entry| {
305            entry.backend == self.backend_id() && entry.op == query.op && entry.dtype == query.dtype
306        })
307    }
308
309    /// Require support for one operation/dtype/axis, returning a structured
310    /// unsupported error otherwise.
311    ///
312    /// # Examples
313    ///
314    /// ```rust
315    /// use tenferro_core_ops::PrimitiveOpKind;
316    /// use tenferro_tensor::{
317    ///     BackendId, CapabilityAxis, CapabilityQuery, DType, Error, OperationCapability,
318    ///     SupportLevel, TensorBackendCapability,
319    /// };
320    ///
321    /// struct Backend;
322    /// const ENTRIES: &[OperationCapability] = &[OperationCapability {
323    ///     backend: BackendId::Cuda,
324    ///     op: PrimitiveOpKind::Neg,
325    ///     dtype: DType::I32,
326    ///     output_dtype: DType::I32,
327    ///     result: SupportLevel::Unsupported,
328    ///     read_inputs: SupportLevel::Unsupported,
329    ///     write_output: SupportLevel::Unsupported,
330    ///     strided_output: SupportLevel::Unsupported,
331    ///     accumulation: SupportLevel::Unsupported,
332    /// }];
333    /// impl TensorBackendCapability for Backend {
334    ///     fn backend_id(&self) -> BackendId { BackendId::Cuda }
335    ///     fn capabilities(&self) -> &'static [OperationCapability] { ENTRIES }
336    /// }
337    ///
338    /// let err = Backend
339    ///     .require_capability(
340    ///         CapabilityQuery::new(PrimitiveOpKind::Neg, DType::I32),
341    ///         CapabilityAxis::OwnedResult,
342    ///     )
343    ///     .unwrap_err();
344    /// assert!(matches!(err, Error::UnsupportedOpDType { backend: BackendId::Cuda, .. }));
345    /// ```
346    fn require_capability(
347        &self,
348        query: CapabilityQuery,
349        axis: CapabilityAxis,
350    ) -> crate::Result<OperationCapability> {
351        let entry = self.capability(query).ok_or_else(|| {
352            crate::Error::unsupported_op_dtype(
353                descriptor(query.op).name,
354                query.dtype,
355                self.backend_id(),
356            )
357        })?;
358        if entry.axis(axis).is_supported() {
359            Ok(entry)
360        } else {
361            Err(crate::Error::unsupported_op_dtype(
362                descriptor(query.op).name,
363                query.dtype,
364                self.backend_id(),
365            ))
366        }
367    }
368}
369
370/// Return the output dtype allowed by the core op catalog policy for a unary
371/// dtype representative.
372///
373/// `None` means the catalog policy does not admit the queried dtype. Backend
374/// descriptors add implementation support on top of this semantic policy.
375///
376/// # Examples
377///
378/// ```rust
379/// use tenferro_core_ops::PrimitiveOpKind;
380/// use tenferro_tensor::{capability_output_dtype, DType};
381///
382/// assert_eq!(
383///     capability_output_dtype(PrimitiveOpKind::Abs, DType::C64),
384///     Some(DType::F64)
385/// );
386/// assert_eq!(
387///     capability_output_dtype(PrimitiveOpKind::Pow, DType::I32),
388///     Some(DType::I32)
389/// );
390/// ```
391#[must_use]
392pub fn capability_output_dtype(op: PrimitiveOpKind, dtype: DType) -> Option<DType> {
393    let policy = descriptor(op).dtype_policy;
394    match policy {
395        DTypePolicy::SameAny => Some(dtype),
396        DTypePolicy::SameNumeric => numeric_dtype(dtype).then_some(dtype),
397        DTypePolicy::SameFloat => float_dtype(dtype).then_some(dtype),
398        DTypePolicy::AbsToReal => match dtype {
399            DType::F32 => Some(DType::F32),
400            DType::F64 => Some(DType::F64),
401            DType::I32 => Some(DType::I32),
402            DType::I64 => Some(DType::I64),
403            DType::C32 => Some(DType::F32),
404            DType::C64 => Some(DType::F64),
405            DType::Bool => None,
406        },
407        DTypePolicy::SameFloatOrComplex => float_or_complex_dtype(dtype).then_some(dtype),
408        DTypePolicy::CompareToBool => comparable_dtype(dtype).then_some(DType::Bool),
409        DTypePolicy::BoolSelect => Some(dtype),
410        DTypePolicy::Convert | DTypePolicy::Shape | DTypePolicy::Constant => Some(dtype),
411    }
412}
413
414const fn numeric_dtype(dtype: DType) -> bool {
415    matches!(
416        dtype,
417        DType::F32 | DType::F64 | DType::I32 | DType::I64 | DType::C32 | DType::C64
418    )
419}
420
421const fn float_dtype(dtype: DType) -> bool {
422    matches!(dtype, DType::F32 | DType::F64)
423}
424
425const fn float_or_complex_dtype(dtype: DType) -> bool {
426    matches!(dtype, DType::F32 | DType::F64 | DType::C32 | DType::C64)
427}
428
429const fn comparable_dtype(dtype: DType) -> bool {
430    matches!(
431        dtype,
432        DType::F32 | DType::F64 | DType::I32 | DType::I64 | DType::Bool
433    )
434}
tenferro_tensor/capability.rs

tenferro_tensor/
capability.rs
