candela/tensor/backend/
mod.rs

1mod common;
2mod common_kernels;
3#[cfg(feature = "mkl")]
4mod cpu_mkl;
5mod cpu_pure;
6
7use std::fmt::Debug;
8
9use crate::tensor::mem_formats::layout::Layout;
10use crate::tensor::ops::def_op::OpKind;
11use crate::tensor::storage::TensorData;
12use crate::tensor::traits::Numeric;
13
14/// Scalar element type the tensor framework supports.
15pub trait Dtype: Copy + Numeric {}
16impl Dtype for f32 {}
17impl Dtype for f64 {}
18
19/// Per-`(T, B)` compute dispatch. Splitting kernels across this trait lets each
20/// dtype specialize independently without trait-coherence conflicts; the
21/// [`Backend`] methods delegate here.
22pub trait ComputeFor<B: Backend>: Dtype {
23    fn compute(
24        op: &OpKind<Self>,
25        output_buffer: Vec<Self>,
26        output_layout: &Layout,
27        inputs: &[TensorData<Self>],
28    ) -> TensorData<Self>;
29
30    fn compute_inplace(
31        op: &OpKind<Self>,
32        output_layout: &Layout,
33        inputs: Vec<TensorData<Self>>,
34        output_idx: usize,
35    ) -> TensorData<Self>;
36}
37
38/// Compute strategy used by [`Tensor`](crate::tensor::Tensor) and the planner
39/// to execute graph nodes. A `Backend` impl is a zero-sized policy type; all
40/// state lives in the `TensorData` buffers passed through `compute`.
41///
42/// # Required behaviour
43///
44/// Implementations must accept stride-0 batch axes in
45/// `OpKind::MatMul`. A batched
46/// matmul whose leading axis has stride 0 is computed by re-reading the same
47/// matrix per batch iteration; the op layer relies on this to skip
48/// materializing batch broadcasts.
49pub trait Backend: Sized + Debug {
50    /// `true` if `OpKind::MatMul`
51    /// accepts a 2D input whose strides describe a transposed view
52    /// (`row_stride == 1`, `col_stride > 1`) - i.e. the underlying GEMM is
53    /// invoked with a trans-flag and no copy is required. The fast path is
54    /// deliberately scoped to rank 2; higher-rank tensors are always
55    /// contiguified by the op layer before reaching the kernel.
56    ///
57    /// When `false`, the op layer inserts an `AsContiguous` on any matmul
58    /// input that is not already contiguous.
59    const SUPPORTS_2D_TRANSPOSED_MATMUL: bool = Self::SUPPORTS_NON_CONTIGUOUS_MATMUL;
60    /// `true` if `OpKind::MatMul`
61    /// accepts any memory configuration as long the last 2 axis are contiguous.
62    ///
63    /// When `false`, the op layer inserts an `AsContiguous` on any matmul
64    /// input that is not already contiguous.
65    const SUPPORTS_NON_CONTIGUOUS_MATMUL: bool;
66
67    /// Run `op` over `inputs` into a fresh allocation. `output_buffer` is the
68    /// destination `Vec<T>`; the returned `TensorData` wraps it with
69    /// `output_layout`.
70    fn compute<T>(
71        op: &OpKind<T>,
72        output_buffer: Vec<T>,
73        output_layout: &Layout,
74        inputs: &[TensorData<T>],
75    ) -> TensorData<T>
76    where
77        T: Dtype + ComputeFor<Self>;
78
79    /// Run `op` reusing `inputs[output_idx]`'s buffer as the destination. The
80    /// planner guarantees that buffer is no longer referenced by any live
81    /// node at this point, so the in-place write is sound.
82    fn compute_inplace<T>(
83        op: &OpKind<T>,
84        output_layout: &Layout,
85        inputs: Vec<TensorData<T>>,
86        output_idx: usize,
87    ) -> TensorData<T>
88    where
89        T: Dtype + ComputeFor<Self>;
90}
91
92/// Backend selected when no explicit type parameter is supplied at the
93/// [`Tensor`](crate::tensor::Tensor) construction site. Defaults to the
94/// pure-Rust backend; enabling the `mkl` feature switches it to the Intel MKL
95/// backend.
96///
97/// # Examples
98///
99/// ```
100/// use candela::backend::DefaultBackend;
101/// use candela::Tensor;
102///
103/// // `Tensor<T>` is shorthand for `Tensor<T, DefaultBackend>` - the same type.
104/// let a: Tensor<f64> = Tensor::from_scalar(1.0, &[3]);
105/// let b: Tensor<f64, DefaultBackend> = Tensor::from_scalar(1.0, &[3]);
106/// assert_eq!(a.data(), b.data());
107/// ```
108#[cfg(feature = "mkl")]
109pub type DefaultBackend = cpu_mkl::CpuMkl;
110#[cfg(not(feature = "mkl"))]
111pub type DefaultBackend = cpu_pure::CpuPure;
112
113pub mod implementation {
114    #[cfg(feature = "mkl")]
115    pub use crate::tensor::backend::cpu_mkl::CpuMkl;
116    pub use crate::tensor::backend::cpu_pure::CpuPure;
117}
candela/tensor/backend/mod.rs

candela/tensor/backend/
mod.rs
