strided_einsum2/backend.rs
1//! Backend abstraction for batched GEMM dispatch.
2//!
3//! This module defines the [`Backend`] trait, marker structs for each backend,
4//! and the `ActiveBackend` type alias that serves as the single point of
5//! backend selection based on Cargo features.
6
7use strided_kernel::ExecContext;
8use strided_view::ElementOp;
9
10/// Trait for backends that can execute batched GEMM on contiguous operands.
11///
12/// Each backend declares its configuration (conjugation materialization,
13/// stride requirements) and provides a GEMM implementation.
14///
15/// Implementations are provided by each backend module (faer, blas).
16/// External crates can implement this trait for custom scalar types
17/// (e.g., tropical semiring) and pass the backend to [`einsum2_with_backend_into`].
18///
19/// [`einsum2_with_backend_into`]: crate::einsum2_with_backend_into
20pub trait Backend<T: crate::ScalarBase> {
21 /// Whether the backend needs conjugation materialized into the data
22 /// before GEMM (e.g., CBLAS has no conjugation flag for `?gemm`).
23 const MATERIALIZES_CONJ: bool;
24
25 /// Whether the backend requires at least one unit stride per matrix
26 /// dimension (row or column stride must be 1). CBLAS `?gemm` requires
27 /// this; faer does not.
28 const REQUIRES_UNIT_STRIDE: bool;
29
30 /// Execute batched GEMM: `C = alpha * A * B + beta * C` for each batch.
31 ///
32 /// - `c`: mutable output operand (batch x m x n)
33 /// - `a`: input operand (batch x m x k)
34 /// - `b`: input operand (batch x k x n)
35 /// - `batch_dims`: sizes of the batch dimensions
36 /// - `m`, `n`, `k`: fused matrix dimensions
37 /// - `alpha`, `beta`: scaling factors
38 fn bgemm_contiguous_into(
39 c: &mut crate::contiguous::ContiguousOperandMut<T>,
40 a: &crate::contiguous::ContiguousOperand<T>,
41 b: &crate::contiguous::ContiguousOperand<T>,
42 batch_dims: &[usize],
43 m: usize,
44 n: usize,
45 k: usize,
46 alpha: T,
47 beta: T,
48 ) -> strided_view::Result<()>;
49}
50
51/// Private overwrite-only backend contract. The initialized `Backend` trait
52/// remains unchanged for beta-bearing callers.
53#[allow(dead_code)]
54pub(crate) trait OverwriteBackend<T: crate::ScalarBase> {
55 fn bgemm_contiguous_overwrite(
56 c: &mut crate::contiguous::UninitContiguousOperand<'_, '_, T>,
57 a: &crate::contiguous::ContiguousOperand<T>,
58 b: &crate::contiguous::ContiguousOperand<T>,
59 batch_dims: &[usize],
60 m: usize,
61 n: usize,
62 k: usize,
63 alpha: T,
64 ctx: &ExecContext,
65 ) -> strided_view::Result<()>;
66}
67
68// ---------------------------------------------------------------------------
69// Marker structs
70// ---------------------------------------------------------------------------
71
72/// Batched GEMM backend using the [`faer`] library.
73///
74/// `Backend<T>` is implemented in `bgemm_faer.rs`.
75#[cfg(feature = "faer")]
76pub struct FaerBackend;
77
78/// Batched GEMM backend using CBLAS (via `cblas-sys` or `cblas-inject`).
79///
80/// `Backend<T>` is implemented in `bgemm_blas.rs`.
81#[cfg(any(feature = "blas", feature = "blas-inject"))]
82pub struct BlasBackend;
83
84/// Fallback batched GEMM backend using explicit loops (no external library).
85///
86/// This backend is used as `ActiveBackend` when no GEMM feature is enabled.
87/// The GEMM dispatch in `einsum2_into` calls `bgemm_naive` directly rather
88/// than going through the `Backend` trait, so `bgemm_contiguous_into` is
89/// unreachable.
90#[allow(dead_code)]
91pub struct NaiveBackend;
92
93impl<T> Backend<T> for NaiveBackend
94where
95 T: crate::ScalarBase + strided_view::ElementOpApply,
96{
97 const MATERIALIZES_CONJ: bool = false;
98 const REQUIRES_UNIT_STRIDE: bool = false;
99
100 fn bgemm_contiguous_into(
101 c: &mut crate::contiguous::ContiguousOperandMut<T>,
102 a: &crate::contiguous::ContiguousOperand<T>,
103 b: &crate::contiguous::ContiguousOperand<T>,
104 batch_dims: &[usize],
105 m: usize,
106 n: usize,
107 k: usize,
108 alpha: T,
109 beta: T,
110 ) -> strided_view::Result<()> {
111 let a_ptr = a.ptr();
112 let b_ptr = b.ptr();
113 let c_ptr = c.ptr();
114 let a_rs = a.row_stride();
115 let a_cs = a.col_stride();
116 let b_rs = b.row_stride();
117 let b_cs = b.col_stride();
118 let c_rs = c.row_stride();
119 let c_cs = c.col_stride();
120
121 let mut batch_iter = crate::util::MultiIndex::new(batch_dims);
122 while batch_iter.next().is_some() {
123 let a_base = batch_iter.offset(a.batch_strides());
124 let b_base = batch_iter.offset(b.batch_strides());
125 let c_base = batch_iter.offset(c.batch_strides());
126
127 for i in 0..m {
128 for j in 0..n {
129 let mut acc = T::zero();
130 for l in 0..k {
131 let mut a_val = unsafe {
132 *a_ptr.offset(a_base + i as isize * a_rs + l as isize * a_cs)
133 };
134 let mut b_val = unsafe {
135 *b_ptr.offset(b_base + l as isize * b_rs + j as isize * b_cs)
136 };
137 if a.conj() {
138 a_val = strided_view::Conj::apply(a_val);
139 }
140 if b.conj() {
141 b_val = strided_view::Conj::apply(b_val);
142 }
143 acc = acc + a_val * b_val;
144 }
145 unsafe {
146 let c_elem = c_ptr.offset(c_base + i as isize * c_rs + j as isize * c_cs);
147 if beta == T::zero() {
148 *c_elem = alpha * acc;
149 } else {
150 *c_elem = alpha * acc + beta * (*c_elem);
151 }
152 }
153 }
154 }
155 }
156 Ok(())
157 }
158}
159
160#[cfg(not(any(feature = "blas", feature = "blas-inject")))]
161impl<T> OverwriteBackend<T> for NaiveBackend
162where
163 T: crate::ScalarBase + strided_view::ElementOpApply,
164{
165 fn bgemm_contiguous_overwrite(
166 c: &mut crate::contiguous::UninitContiguousOperand<'_, '_, T>,
167 a: &crate::contiguous::ContiguousOperand<T>,
168 b: &crate::contiguous::ContiguousOperand<T>,
169 batch_dims: &[usize],
170 m: usize,
171 n: usize,
172 k: usize,
173 alpha: T,
174 ctx: &ExecContext,
175 ) -> strided_view::Result<()> {
176 crate::uninit::bgemm_contiguous_naive(c, a, b, batch_dims, m, n, k, alpha, ctx)
177 }
178}
179
180// ---------------------------------------------------------------------------
181// ActiveBackend type alias -- the SINGLE point of backend selection
182// ---------------------------------------------------------------------------
183
184/// The active GEMM backend, selected by Cargo features.
185///
186/// - `blas` or `blas-inject` -> `BlasBackend`
187/// - `faer` without BLAS -> `FaerBackend`
188/// - no backend feature -> `NaiveBackend`
189/// - invalid combos -> `NaiveBackend` (placeholder; `compile_error!` fires first)
190#[cfg(any(
191 all(feature = "blas", not(feature = "blas-inject")),
192 all(feature = "blas-inject", not(feature = "blas"))
193))]
194pub type ActiveBackend = BlasBackend;
195
196#[cfg(all(feature = "faer", not(any(feature = "blas", feature = "blas-inject"))))]
197pub type ActiveBackend = FaerBackend;
198
199#[cfg(not(any(feature = "faer", feature = "blas", feature = "blas-inject")))]
200pub type ActiveBackend = NaiveBackend;
201
202/// Placeholder for invalid mutually-exclusive feature combinations.
203///
204/// The crate emits `compile_error!` for these combinations (in `lib.rs`), so this
205/// alias only suppresses cascading type-resolution errors.
206#[cfg(all(feature = "blas", feature = "blas-inject"))]
207pub type ActiveBackend = NaiveBackend;