strided_kernel/
outer_product.rs

1//! Semantic outer-product API on dynamic-rank strided views.
2
3use core::mem::MaybeUninit;
4use std::ops::Mul;
5
6#[cfg(feature = "parallel")]
7use smallvec::SmallVec;
8
9use crate::map_view::{broadcast_mul_into, broadcast_mul_into_uninit};
10use crate::maybe_sync::MaybeSendSync;
11use crate::view::{StridedView, StridedViewMut};
12use crate::{ElementOp, Result, StridedError};
13
14#[cfg(feature = "parallel")]
15type AxisVec<T> = SmallVec<[T; 8]>;
16#[cfg(not(feature = "parallel"))]
17type AxisVec<T> = Vec<T>;
18
19/// Compute `dest[lhs_free..., rhs_free..., batch...] =
20/// lhs[lhs_free..., batch...] * rhs[rhs_free..., batch...]`.
21///
22/// This is a semantic convenience wrapper over [`broadcast_mul_into`]. The
23/// broadcast/mul planner owns kernel selection, so explicit outer-product calls
24/// and equivalent broadcasted multiplication use the same implementation path.
25pub fn batched_outer_product_into<D, A, B, OpA, OpB>(
26    dest: &mut StridedViewMut<D>,
27    lhs: &StridedView<A, OpA>,
28    rhs: &StridedView<B, OpB>,
29    lhs_free_ndim: usize,
30    rhs_free_ndim: usize,
31) -> Result<()>
32where
33    D: Copy + MaybeSendSync + 'static,
34    A: Copy + MaybeSendSync + Mul<B, Output = D> + 'static,
35    B: Copy + MaybeSendSync + 'static,
36    OpA: ElementOp<A>,
37    OpB: ElementOp<B>,
38{
39    validate_batched_outer_shape(dest, lhs, rhs, lhs_free_ndim, rhs_free_ndim)?;
40
41    let batch_ndim = lhs.ndim() - lhs_free_ndim;
42    let mut lhs_axes = AxisVec::<usize>::with_capacity(lhs.ndim());
43    let mut rhs_axes = AxisVec::<usize>::with_capacity(rhs.ndim());
44
45    lhs_axes.extend(0..lhs_free_ndim);
46    rhs_axes.extend(lhs_free_ndim..lhs_free_ndim + rhs_free_ndim);
47
48    let batch_axis_start = lhs_free_ndim + rhs_free_ndim;
49    lhs_axes.extend(batch_axis_start..batch_axis_start + batch_ndim);
50    rhs_axes.extend(batch_axis_start..batch_axis_start + batch_ndim);
51
52    broadcast_mul_into(dest, lhs, &lhs_axes, rhs, &rhs_axes)
53}
54
55/// Compute a batched outer product into a fully overwritten uninitialized output.
56///
57/// Rank, shape, destination-injectivity, and reachable-byte overlap validation
58/// completes before the first write. Safe Rust borrows already prevent
59/// input/output aliasing; the explicit overlap check in the shared broadcast
60/// kernel preserves the contract for views produced through unsafe constructors.
61///
62/// `Ok(())` means every logical destination element is initialized. An error
63/// occurs before writes. A panic during replay may leave a partially initialized
64/// destination, which remains safe to drop as `MaybeUninit<D>`.
65///
66/// # Errors
67///
68/// Returns a typed rank or shape error for incompatible free/batch dimensions,
69/// [`StridedError::NonInjectiveOutputLayout`] for an overlapping output layout,
70/// [`StridedError::OverlappingInputOutput`] for aliased storage, or
71/// [`StridedError::OffsetOverflow`] when a reachable byte range is not
72/// representable.
73pub fn batched_outer_product_into_uninit<D, A, B, OpA, OpB>(
74    dest: &mut StridedViewMut<MaybeUninit<D>>,
75    lhs: &StridedView<A, OpA>,
76    rhs: &StridedView<B, OpB>,
77    lhs_free_ndim: usize,
78    rhs_free_ndim: usize,
79) -> Result<()>
80where
81    D: Copy + MaybeSendSync + 'static,
82    A: Copy + MaybeSendSync + Mul<B, Output = D> + 'static,
83    B: Copy + MaybeSendSync + 'static,
84    OpA: ElementOp<A>,
85    OpB: ElementOp<B>,
86{
87    validate_batched_outer_shape(dest, lhs, rhs, lhs_free_ndim, rhs_free_ndim)?;
88    let batch_ndim = lhs.ndim() - lhs_free_ndim;
89    let mut lhs_axes = AxisVec::<usize>::with_capacity(lhs.ndim());
90    let mut rhs_axes = AxisVec::<usize>::with_capacity(rhs.ndim());
91    lhs_axes.extend(0..lhs_free_ndim);
92    rhs_axes.extend(lhs_free_ndim..lhs_free_ndim + rhs_free_ndim);
93    let batch_axis_start = lhs_free_ndim + rhs_free_ndim;
94    lhs_axes.extend(batch_axis_start..batch_axis_start + batch_ndim);
95    rhs_axes.extend(batch_axis_start..batch_axis_start + batch_ndim);
96    broadcast_mul_into_uninit(dest, lhs, &lhs_axes, rhs, &rhs_axes)
97}
98
99fn validate_batched_outer_shape<D, A, OpA, B, OpB>(
100    dest: &StridedViewMut<D>,
101    lhs: &StridedView<A, OpA>,
102    rhs: &StridedView<B, OpB>,
103    lhs_free_ndim: usize,
104    rhs_free_ndim: usize,
105) -> Result<()> {
106    if lhs_free_ndim > lhs.ndim() {
107        return Err(StridedError::RankMismatch(lhs_free_ndim, lhs.ndim()));
108    }
109    if rhs_free_ndim > rhs.ndim() {
110        return Err(StridedError::RankMismatch(rhs_free_ndim, rhs.ndim()));
111    }
112
113    let lhs_batch_ndim = lhs.ndim() - lhs_free_ndim;
114    let rhs_batch_ndim = rhs.ndim() - rhs_free_ndim;
115    if lhs_batch_ndim != rhs_batch_ndim {
116        return Err(StridedError::RankMismatch(lhs_batch_ndim, rhs_batch_ndim));
117    }
118
119    let expected_dest_rank = lhs_free_ndim + rhs_free_ndim + lhs_batch_ndim;
120    if dest.ndim() != expected_dest_rank {
121        return Err(StridedError::RankMismatch(dest.ndim(), expected_dest_rank));
122    }
123
124    ensure_dims(&dest.dims()[..lhs_free_ndim], &lhs.dims()[..lhs_free_ndim])?;
125    ensure_dims(
126        &dest.dims()[lhs_free_ndim..lhs_free_ndim + rhs_free_ndim],
127        &rhs.dims()[..rhs_free_ndim],
128    )?;
129    ensure_dims(
130        &dest.dims()[lhs_free_ndim + rhs_free_ndim..],
131        &lhs.dims()[lhs_free_ndim..],
132    )?;
133    ensure_dims(
134        &dest.dims()[lhs_free_ndim + rhs_free_ndim..],
135        &rhs.dims()[rhs_free_ndim..],
136    )?;
137
138    Ok(())
139}
140
141fn ensure_dims(actual: &[usize], expected: &[usize]) -> Result<()> {
142    if actual == expected {
143        Ok(())
144    } else {
145        Err(StridedError::ShapeMismatch(
146            actual.to_vec(),
147            expected.to_vec(),
148        ))
149    }
150}
strided_kernel/outer_product.rs

strided_kernel/
outer_product.rs
