strided_einsum2/
dot_general.rs

1//! Axis-based general dot product API.
2//!
3//! The dimension-numbering config borrows axis slices so callers can reuse
4//! their existing metadata without allocating a second owned config.
5//!
6//! ```
7//! use strided_einsum2::{dot_general_into, DotGeneralConfig};
8//! use strided_view::StridedArray;
9//!
10//! let a = StridedArray::<f64>::from_fn_col_major(&[2, 3], |idx| {
11//!     (idx[0] + 2 * idx[1] + 1) as f64
12//! });
13//! let b = StridedArray::<f64>::from_fn_col_major(&[3, 2], |idx| {
14//!     (idx[0] + 3 * idx[1] + 1) as f64
15//! });
16//! let mut c = StridedArray::<f64>::col_major(&[2, 2]);
17//!
18//! let config = DotGeneralConfig {
19//!     lhs_contracting_dims: &[1],
20//!     rhs_contracting_dims: &[0],
21//!     lhs_batch_dims: &[],
22//!     rhs_batch_dims: &[],
23//! };
24//! dot_general_into(c.view_mut(), &a.view(), &b.view(), &config, 1.0, 0.0).unwrap();
25//! ```
26
27use smallvec::SmallVec;
28use std::mem::MaybeUninit;
29use strided_kernel::ExecContext;
30use strided_view::{RawStridedMut, StridedView, StridedViewMut};
31
32use crate::backend::Backend;
33use crate::{einsum2_dispatch, Einsum2Plan, EinsumError, Result, ScalarBase};
34
35/// DotGeneral dimension configuration.
36///
37/// The output shape is `[lhs_free..., rhs_free..., batch...]`, matching
38/// tenferro's batch-trailing col-major convention.
39#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
40pub struct DotGeneralConfig<'a> {
41    pub lhs_contracting_dims: &'a [usize],
42    pub rhs_contracting_dims: &'a [usize],
43    pub lhs_batch_dims: &'a [usize],
44    pub rhs_batch_dims: &'a [usize],
45}
46
47#[derive(Clone, Debug)]
48pub(crate) struct DotGeneralLabels {
49    pub lhs_labels: SmallVec<[usize; 8]>,
50    pub rhs_labels: SmallVec<[usize; 8]>,
51    pub out_labels: SmallVec<[usize; 8]>,
52}
53
54fn check_no_duplicates(dims: &[usize], label: &str) -> Result<()> {
55    for (i, &dim) in dims.iter().enumerate() {
56        if dims[..i].contains(&dim) {
57            return Err(EinsumError::InvalidDotGeneralConfig(format!(
58                "{label} contains duplicate dim {dim}"
59            )));
60        }
61    }
62    Ok(())
63}
64
65fn check_bounds(dims: &[usize], rank: usize, label: &str) -> Result<()> {
66    for &dim in dims {
67        if dim >= rank {
68            return Err(EinsumError::InvalidDotGeneralConfig(format!(
69                "{label} dim {dim} out of bounds for rank {rank}"
70            )));
71        }
72    }
73    Ok(())
74}
75
76fn free_dims(rank: usize, contracting: &[usize], batch: &[usize]) -> SmallVec<[usize; 8]> {
77    (0..rank)
78        .filter(|dim| !contracting.contains(dim) && !batch.contains(dim))
79        .collect()
80}
81
82impl DotGeneralConfig<'_> {
83    /// Validate dimension indices for explicit operand ranks.
84    pub fn validate_dims_with_ranks(&self, lhs_rank: usize, rhs_rank: usize) -> Result<()> {
85        check_bounds(&self.lhs_contracting_dims, lhs_rank, "lhs_contracting")?;
86        check_bounds(&self.rhs_contracting_dims, rhs_rank, "rhs_contracting")?;
87        check_bounds(&self.lhs_batch_dims, lhs_rank, "lhs_batch")?;
88        check_bounds(&self.rhs_batch_dims, rhs_rank, "rhs_batch")?;
89        check_no_duplicates(&self.lhs_contracting_dims, "lhs_contracting_dims")?;
90        check_no_duplicates(&self.rhs_contracting_dims, "rhs_contracting_dims")?;
91        check_no_duplicates(&self.lhs_batch_dims, "lhs_batch_dims")?;
92        check_no_duplicates(&self.rhs_batch_dims, "rhs_batch_dims")?;
93
94        for &dim in self.lhs_contracting_dims {
95            if self.lhs_batch_dims.contains(&dim) {
96                return Err(EinsumError::InvalidDotGeneralConfig(format!(
97                    "lhs dim {dim} appears in both contracting and batch dims"
98                )));
99            }
100        }
101        for &dim in self.rhs_contracting_dims {
102            if self.rhs_batch_dims.contains(&dim) {
103                return Err(EinsumError::InvalidDotGeneralConfig(format!(
104                    "rhs dim {dim} appears in both contracting and batch dims"
105                )));
106            }
107        }
108        if self.lhs_contracting_dims.len() != self.rhs_contracting_dims.len() {
109            return Err(EinsumError::InvalidDotGeneralConfig(format!(
110                "lhs/rhs contracting dim counts differ ({} vs {})",
111                self.lhs_contracting_dims.len(),
112                self.rhs_contracting_dims.len()
113            )));
114        }
115        if self.lhs_batch_dims.len() != self.rhs_batch_dims.len() {
116            return Err(EinsumError::InvalidDotGeneralConfig(format!(
117                "lhs/rhs batch dim counts differ ({} vs {})",
118                self.lhs_batch_dims.len(),
119                self.rhs_batch_dims.len()
120            )));
121        }
122        Ok(())
123    }
124
125    /// Compute the output shape `[lhs_free..., rhs_free..., batch...]`.
126    pub fn expected_output_shape(
127        &self,
128        lhs_shape: &[usize],
129        rhs_shape: &[usize],
130    ) -> Result<Vec<usize>> {
131        self.validate_dims_with_ranks(lhs_shape.len(), rhs_shape.len())?;
132        self.validate_pair_shapes(lhs_shape, rhs_shape)?;
133
134        let lhs_free = free_dims(
135            lhs_shape.len(),
136            &self.lhs_contracting_dims,
137            &self.lhs_batch_dims,
138        );
139        let rhs_free = free_dims(
140            rhs_shape.len(),
141            &self.rhs_contracting_dims,
142            &self.rhs_batch_dims,
143        );
144
145        let mut out =
146            Vec::with_capacity(lhs_free.len() + rhs_free.len() + self.lhs_batch_dims.len());
147        out.extend(lhs_free.iter().map(|&dim| lhs_shape[dim]));
148        out.extend(rhs_free.iter().map(|&dim| rhs_shape[dim]));
149        out.extend(self.lhs_batch_dims.iter().map(|&dim| lhs_shape[dim]));
150        Ok(out)
151    }
152
153    fn validate_pair_shapes(&self, lhs_shape: &[usize], rhs_shape: &[usize]) -> Result<()> {
154        for (&lhs_dim, &rhs_dim) in self.lhs_batch_dims.iter().zip(self.rhs_batch_dims) {
155            if lhs_shape[lhs_dim] != rhs_shape[rhs_dim] {
156                return Err(EinsumError::DimensionMismatch {
157                    axis: format!("batch lhs {lhs_dim} rhs {rhs_dim}"),
158                    dim_a: lhs_shape[lhs_dim],
159                    dim_b: rhs_shape[rhs_dim],
160                });
161            }
162        }
163        for (&lhs_dim, &rhs_dim) in self
164            .lhs_contracting_dims
165            .iter()
166            .zip(self.rhs_contracting_dims)
167        {
168            if lhs_shape[lhs_dim] != rhs_shape[rhs_dim] {
169                return Err(EinsumError::DimensionMismatch {
170                    axis: format!("contract lhs {lhs_dim} rhs {rhs_dim}"),
171                    dim_a: lhs_shape[lhs_dim],
172                    dim_b: rhs_shape[rhs_dim],
173                });
174            }
175        }
176        Ok(())
177    }
178
179    pub(crate) fn labels_for_shapes(
180        &self,
181        lhs_shape: &[usize],
182        rhs_shape: &[usize],
183        out_shape: &[usize],
184    ) -> Result<DotGeneralLabels> {
185        self.validate_dims_with_ranks(lhs_shape.len(), rhs_shape.len())?;
186        self.validate_pair_shapes(lhs_shape, rhs_shape)?;
187
188        let lhs_free = free_dims(
189            lhs_shape.len(),
190            &self.lhs_contracting_dims,
191            &self.lhs_batch_dims,
192        );
193        let rhs_free = free_dims(
194            rhs_shape.len(),
195            &self.rhs_contracting_dims,
196            &self.rhs_batch_dims,
197        );
198
199        let mut lhs_labels = smallvec::smallvec![usize::MAX; lhs_shape.len()];
200        let mut rhs_labels = smallvec::smallvec![usize::MAX; rhs_shape.len()];
201        let mut next_label = 0usize;
202
203        for (&lhs_dim, &rhs_dim) in self.lhs_batch_dims.iter().zip(self.rhs_batch_dims) {
204            lhs_labels[lhs_dim] = next_label;
205            rhs_labels[rhs_dim] = next_label;
206            next_label += 1;
207        }
208        for &lhs_dim in &lhs_free {
209            lhs_labels[lhs_dim] = next_label;
210            next_label += 1;
211        }
212        for &rhs_dim in &rhs_free {
213            rhs_labels[rhs_dim] = next_label;
214            next_label += 1;
215        }
216        for (&lhs_dim, &rhs_dim) in self
217            .lhs_contracting_dims
218            .iter()
219            .zip(self.rhs_contracting_dims)
220        {
221            lhs_labels[lhs_dim] = next_label;
222            rhs_labels[rhs_dim] = next_label;
223            next_label += 1;
224        }
225
226        let expected_out_shape = self.expected_output_shape(lhs_shape, rhs_shape)?;
227        if expected_out_shape.as_slice() != out_shape {
228            return Err(EinsumError::OutputShapeMismatch {
229                expected: expected_out_shape,
230                got: out_shape.to_vec(),
231            });
232        }
233
234        let mut out_labels = SmallVec::<[usize; 8]>::new();
235        out_labels.extend(lhs_free.iter().map(|&dim| lhs_labels[dim]));
236        out_labels.extend(rhs_free.iter().map(|&dim| rhs_labels[dim]));
237        out_labels.extend(self.lhs_batch_dims.iter().map(|&dim| lhs_labels[dim]));
238
239        Ok(DotGeneralLabels {
240            lhs_labels,
241            rhs_labels,
242            out_labels,
243        })
244    }
245}
246
247/// Compute `C = alpha * dot_general(A, B) + beta * C` with an explicit backend.
248pub fn dot_general_with_backend_into<T, B>(
249    c: StridedViewMut<T>,
250    a: &StridedView<T>,
251    b: &StridedView<T>,
252    config: &DotGeneralConfig<'_>,
253    alpha: T,
254    beta: T,
255) -> Result<()>
256where
257    T: ScalarBase,
258    B: Backend<T>,
259{
260    let labels = config.labels_for_shapes(a.dims(), b.dims(), c.dims())?;
261    let plan = Einsum2Plan::new(
262        labels.lhs_labels.as_slice(),
263        labels.rhs_labels.as_slice(),
264        labels.out_labels.as_slice(),
265    )?;
266    einsum2_dispatch::<T, B, _>(c, a, b, &plan, alpha, beta, false, false, None)
267}
268
269/// Compute `C = alpha * dot_general(A, B) + beta * C` with the active backend.
270#[cfg(any(feature = "faer", feature = "blas", feature = "blas-inject"))]
271pub fn dot_general_into<T: crate::Scalar>(
272    c: StridedViewMut<T>,
273    a: &StridedView<T>,
274    b: &StridedView<T>,
275    config: &DotGeneralConfig<'_>,
276    alpha: T,
277    beta: T,
278) -> Result<()>
279where
280    crate::backend::ActiveBackend: Backend<T>,
281{
282    dot_general_with_backend_into::<T, crate::backend::ActiveBackend>(c, a, b, config, alpha, beta)
283}
284
285/// Compute dot-general into a genuinely uninitialized destination.
286///
287/// The destination is only exposed as initialized after this function returns
288/// `Ok(())`; holes in a strided backing allocation are never touched.
289pub fn dot_general_into_uninit<T: crate::Scalar>(
290    c: &mut RawStridedMut<'_, MaybeUninit<T>>,
291    a: &StridedView<'_, T>,
292    b: &StridedView<'_, T>,
293    config: &DotGeneralConfig<'_>,
294    alpha: T,
295    ctx: &ExecContext,
296) -> Result<()> {
297    let labels = config.labels_for_shapes(a.dims(), b.dims(), c.dims())?;
298    crate::einsum2_into_uninit(
299        c,
300        a,
301        b,
302        labels.out_labels.as_slice(),
303        labels.lhs_labels.as_slice(),
304        labels.rhs_labels.as_slice(),
305        alpha,
306        ctx,
307    )
308}
309
310#[cfg(test)]
311mod tests {
312    use super::*;
313    use std::mem::MaybeUninit;
314
315    #[test]
316    fn uninit_entry_point_routes_through_einsum_validation() {
317        let a = strided_view::StridedArray::from_fn_col_major(&[2, 3], |idx| {
318            (idx[0] + 2 * idx[1] + 1) as f64
319        });
320        let b = strided_view::StridedArray::from_fn_col_major(&[3, 2], |idx| {
321            (idx[0] + 3 * idx[1] + 1) as f64
322        });
323        let mut storage = vec![MaybeUninit::<f64>::uninit(); 4];
324        let mut c = RawStridedMut::new(&mut storage, &[2, 2], &[2, 1], 0).unwrap();
325        let config = DotGeneralConfig {
326            lhs_contracting_dims: &[1],
327            rhs_contracting_dims: &[0],
328            lhs_batch_dims: &[],
329            rhs_batch_dims: &[],
330        };
331
332        let err = dot_general_into_uninit(
333            &mut c,
334            &a.view(),
335            &b.view(),
336            &config,
337            1.0,
338            &ExecContext::serial(),
339        )
340        .unwrap_err();
341        assert!(matches!(err, crate::EinsumError::Unsupported(_)));
342    }
343}
344
345/// Compute `C = alpha * dot_general(A, B) + beta * C` with the naive fallback.
346#[cfg(not(any(feature = "faer", feature = "blas", feature = "blas-inject")))]
347pub fn dot_general_into<T: crate::Scalar>(
348    c: StridedViewMut<T>,
349    a: &StridedView<T>,
350    b: &StridedView<T>,
351    config: &DotGeneralConfig<'_>,
352    alpha: T,
353    beta: T,
354) -> Result<()> {
355    let labels = config.labels_for_shapes(a.dims(), b.dims(), c.dims())?;
356    crate::einsum2_naive_into(
357        c,
358        a,
359        b,
360        labels.out_labels.as_slice(),
361        labels.lhs_labels.as_slice(),
362        labels.rhs_labels.as_slice(),
363        alpha,
364        beta,
365        |x| x,
366        |x| x,
367    )
368}
strided_einsum2/dot_general.rs

strided_einsum2/
dot_general.rs
