strided_basic/
outer_product.rs

1//! Semantic outer-product API on dynamic-rank strided views.
2
3use core::mem::MaybeUninit;
4use std::ops::Mul;
5
6use crate::view::{StridedView, StridedViewMut};
7use crate::{broadcast_mul_into, broadcast_mul_into_uninit, Exec};
8use crate::{ElementOp, Result, StridedError};
9
10type AxisVec<T> = crate::inline::InlineVec<T>;
11
12/// Compute `dest[lhs_free..., rhs_free..., batch...] =
13/// lhs[lhs_free..., batch...] * rhs[rhs_free..., batch...]`.
14///
15/// This is a semantic convenience wrapper over [`broadcast_mul_into`]. The
16/// broadcast/mul planner owns kernel selection, so explicit outer-product calls
17/// and equivalent broadcasted multiplication use the same implementation path.
18pub fn batched_outer_product_into<D, A, B, OpA, OpB>(
19    dest: &mut StridedViewMut<D>,
20    lhs: &StridedView<A, OpA>,
21    rhs: &StridedView<B, OpB>,
22    lhs_free_ndim: usize,
23    rhs_free_ndim: usize,
24    exec: Exec<'_>,
25) -> Result<()>
26where
27    D: Copy + Send + Sync + 'static,
28    A: Copy + Send + Sync + Mul<B, Output = D> + 'static,
29    B: Copy + Send + Sync + 'static,
30    OpA: ElementOp<A>,
31    OpB: ElementOp<B>,
32{
33    validate_batched_outer_shape(dest, lhs, rhs, lhs_free_ndim, rhs_free_ndim)?;
34
35    let batch_ndim = lhs.ndim() - lhs_free_ndim;
36    let mut lhs_axes = AxisVec::<usize>::with_capacity(lhs.ndim());
37    let mut rhs_axes = AxisVec::<usize>::with_capacity(rhs.ndim());
38
39    lhs_axes.extend(0..lhs_free_ndim);
40    rhs_axes.extend(lhs_free_ndim..lhs_free_ndim + rhs_free_ndim);
41
42    let batch_axis_start = lhs_free_ndim + rhs_free_ndim;
43    lhs_axes.extend(batch_axis_start..batch_axis_start + batch_ndim);
44    rhs_axes.extend(batch_axis_start..batch_axis_start + batch_ndim);
45
46    broadcast_mul_into(dest, lhs, &lhs_axes, rhs, &rhs_axes, exec)
47}
48
49/// Compute a batched outer product into a fully overwritten uninitialized output.
50///
51/// Rank, shape, destination-injectivity, and reachable-byte overlap validation
52/// completes before the first write. Safe Rust borrows already prevent
53/// input/output aliasing; the explicit overlap check in the shared broadcast
54/// kernel preserves the contract for views produced through unsafe constructors.
55///
56/// `Ok(())` means every logical destination element is initialized. An error
57/// occurs before writes. A panic during replay may leave a partially initialized
58/// destination, which remains safe to drop as `MaybeUninit<D>`.
59///
60/// # Errors
61///
62/// Returns a typed rank or shape error for incompatible free/batch dimensions,
63/// [`StridedError::NonInjectiveOutputLayout`] for an overlapping output layout,
64/// [`StridedError::OverlappingInputOutput`] for aliased storage, or
65/// [`StridedError::OffsetOverflow`] when a reachable byte range is not
66/// representable.
67pub fn batched_outer_product_into_uninit<D, A, B, OpA, OpB>(
68    dest: &mut StridedViewMut<MaybeUninit<D>>,
69    lhs: &StridedView<A, OpA>,
70    rhs: &StridedView<B, OpB>,
71    lhs_free_ndim: usize,
72    rhs_free_ndim: usize,
73    exec: Exec<'_>,
74) -> Result<()>
75where
76    D: Copy + Send + Sync + 'static,
77    A: Copy + Send + Sync + Mul<B, Output = D> + 'static,
78    B: Copy + Send + Sync + 'static,
79    OpA: ElementOp<A>,
80    OpB: ElementOp<B>,
81{
82    validate_batched_outer_shape(dest, lhs, rhs, lhs_free_ndim, rhs_free_ndim)?;
83    let batch_ndim = lhs.ndim() - lhs_free_ndim;
84    let mut lhs_axes = AxisVec::<usize>::with_capacity(lhs.ndim());
85    let mut rhs_axes = AxisVec::<usize>::with_capacity(rhs.ndim());
86    lhs_axes.extend(0..lhs_free_ndim);
87    rhs_axes.extend(lhs_free_ndim..lhs_free_ndim + rhs_free_ndim);
88    let batch_axis_start = lhs_free_ndim + rhs_free_ndim;
89    lhs_axes.extend(batch_axis_start..batch_axis_start + batch_ndim);
90    rhs_axes.extend(batch_axis_start..batch_axis_start + batch_ndim);
91    broadcast_mul_into_uninit(dest, lhs, &lhs_axes, rhs, &rhs_axes, exec)
92}
93
94fn validate_batched_outer_shape<D, A, OpA, B, OpB>(
95    dest: &StridedViewMut<D>,
96    lhs: &StridedView<A, OpA>,
97    rhs: &StridedView<B, OpB>,
98    lhs_free_ndim: usize,
99    rhs_free_ndim: usize,
100) -> Result<()> {
101    if lhs_free_ndim > lhs.ndim() {
102        return Err(StridedError::RankMismatch(lhs_free_ndim, lhs.ndim()));
103    }
104    if rhs_free_ndim > rhs.ndim() {
105        return Err(StridedError::RankMismatch(rhs_free_ndim, rhs.ndim()));
106    }
107
108    let lhs_batch_ndim = lhs.ndim() - lhs_free_ndim;
109    let rhs_batch_ndim = rhs.ndim() - rhs_free_ndim;
110    if lhs_batch_ndim != rhs_batch_ndim {
111        return Err(StridedError::RankMismatch(lhs_batch_ndim, rhs_batch_ndim));
112    }
113
114    let expected_dest_rank = lhs_free_ndim + rhs_free_ndim + lhs_batch_ndim;
115    if dest.ndim() != expected_dest_rank {
116        return Err(StridedError::RankMismatch(dest.ndim(), expected_dest_rank));
117    }
118
119    ensure_dims(&dest.dims()[..lhs_free_ndim], &lhs.dims()[..lhs_free_ndim])?;
120    ensure_dims(
121        &dest.dims()[lhs_free_ndim..lhs_free_ndim + rhs_free_ndim],
122        &rhs.dims()[..rhs_free_ndim],
123    )?;
124    ensure_dims(
125        &dest.dims()[lhs_free_ndim + rhs_free_ndim..],
126        &lhs.dims()[lhs_free_ndim..],
127    )?;
128    ensure_dims(
129        &dest.dims()[lhs_free_ndim + rhs_free_ndim..],
130        &rhs.dims()[rhs_free_ndim..],
131    )?;
132
133    Ok(())
134}
135
136fn ensure_dims(actual: &[usize], expected: &[usize]) -> Result<()> {
137    if actual == expected {
138        Ok(())
139    } else {
140        Err(StridedError::ShapeMismatch(
141            actual.to_vec(),
142            expected.to_vec(),
143        ))
144    }
145}
146
147/// Storage layout for a lazily ordered outer product.
148///
149/// Returned by [`plan_lazy_outer_product`]. The caller allocates a dense
150/// column-major base of [`base_dims`](Self::base_dims), fills it through a
151/// destination descriptor over that base with the logical output dims and
152/// [`output_strides`](Self::output_strides) at offset zero (for example with
153/// [`broadcast_mul_into_uninit`] and the original axis maps), and then exposes
154/// the result as a strided view with those same dims and strides.
155#[derive(Clone, Debug, PartialEq, Eq)]
156pub struct LazyOuterProductLayout {
157    /// Column-major extents of the dense base allocation.
158    pub base_dims: Vec<usize>,
159    /// Strides, in base elements, of the logical output axes over the base.
160    pub output_strides: Vec<isize>,
161}
162
163/// Plan an outer-product output whose memory order follows the inputs'
164/// physical stride order instead of the logical output order.
165///
166/// `lhs_axes[k]` (respectively `rhs_axes[k]`) names the output axis that
167/// operand axis `k` maps to, with the operand extent equal to the output
168/// extent. Output axes mapped only by `lhs` are its free axes, axes mapped
169/// only by `rhs` are its free axes, and axes mapped by both are batch axes.
170///
171/// A layout is returned only when the product splits into free and batch
172/// groups, both free groups have more than one element, the logical output
173/// order is `[lhs_free, rhs_free, batch]` or `[rhs_free, lhs_free, batch]`,
174/// all strides are non-negative, and sorting either operand's free axes by
175/// `(stride, axis)` changes their order. The base then holds the leading
176/// group's free axes in the leading operand's physical order, then the
177/// trailing group's free axes in the trailing operand's physical order, then
178/// the batch axes in output order. Traversing inputs in their physical order
179/// while writing the base contiguously is what makes the layout useful.
180///
181/// # Examples
182///
183/// ```
184/// use strided_basic::plan_lazy_outer_product;
185///
186/// // lhs is a transposed 2 x 3 matrix (row-major strides), rhs is a vector.
187/// let layout = plan_lazy_outer_product(&[2, 3, 4], &[2, 3], &[3, 1], &[0, 1], &[4], &[1], &[2])
188///     .unwrap()
189///     .unwrap();
190/// assert_eq!(layout.base_dims, [3, 2, 4]);
191/// assert_eq!(layout.output_strides, [3, 1, 6]);
192/// ```
193///
194/// # Errors
195///
196/// Returns [`StridedError::RankMismatch`] when an operand's dims, strides,
197/// and axis map lengths disagree, and [`StridedError::OffsetOverflow`] when a
198/// base extent product or stride is not representable. Inputs that are valid
199/// but ineligible return `Ok(None)`.
200pub fn plan_lazy_outer_product(
201    output_dims: &[usize],
202    lhs_dims: &[usize],
203    lhs_strides: &[isize],
204    lhs_axes: &[usize],
205    rhs_dims: &[usize],
206    rhs_strides: &[isize],
207    rhs_axes: &[usize],
208) -> Result<Option<LazyOuterProductLayout>> {
209    for (dims, strides, axes) in [
210        (lhs_dims, lhs_strides, lhs_axes),
211        (rhs_dims, rhs_strides, rhs_axes),
212    ] {
213        if dims.len() != strides.len() {
214            return Err(StridedError::RankMismatch(dims.len(), strides.len()));
215        }
216        if dims.len() != axes.len() {
217            return Err(StridedError::RankMismatch(dims.len(), axes.len()));
218        }
219    }
220    if !extents_match_output(output_dims, lhs_dims, lhs_axes)
221        || !extents_match_output(output_dims, rhs_dims, rhs_axes)
222        || lhs_strides
223            .iter()
224            .chain(rhs_strides)
225            .any(|&stride| stride < 0)
226    {
227        return Ok(None);
228    }
229    let Some(partition) = classify_outer_axes(output_dims.len(), lhs_axes, rhs_axes) else {
230        return Ok(None);
231    };
232    if checked_axes_product(lhs_dims, &partition.lhs_free)? <= 1
233        || checked_axes_product(rhs_dims, &partition.rhs_free)? <= 1
234    {
235        return Ok(None);
236    }
237
238    let output_rank = output_dims.len();
239    let lhs_prefix = partition
240        .lhs_free_out
241        .iter()
242        .chain(&partition.rhs_free_out)
243        .chain(&partition.batch_out)
244        .copied()
245        .eq(0..output_rank);
246    let rhs_prefix = !lhs_prefix
247        && partition
248            .rhs_free_out
249            .iter()
250            .chain(&partition.lhs_free_out)
251            .chain(&partition.batch_out)
252            .copied()
253            .eq(0..output_rank);
254    if !lhs_prefix && !rhs_prefix {
255        return Ok(None);
256    }
257
258    let lhs_physical = axes_by_physical_stride(lhs_strides, &partition.lhs_free);
259    let rhs_physical = axes_by_physical_stride(rhs_strides, &partition.rhs_free);
260    if lhs_physical == partition.lhs_free && rhs_physical == partition.rhs_free {
261        return Ok(None);
262    }
263
264    // Base axis order: leading free, trailing free, then batch in output order.
265    let mut base_out_axes = Vec::with_capacity(output_rank);
266    let (leading, leading_axes, trailing, trailing_axes) = if lhs_prefix {
267        (&lhs_physical, lhs_axes, &rhs_physical, rhs_axes)
268    } else {
269        (&rhs_physical, rhs_axes, &lhs_physical, lhs_axes)
270    };
271    base_out_axes.extend(leading.iter().map(|&axis| leading_axes[axis]));
272    base_out_axes.extend(trailing.iter().map(|&axis| trailing_axes[axis]));
273    base_out_axes.extend(partition.batch_out.iter().copied());
274
275    let base_dims: Vec<usize> = base_out_axes
276        .iter()
277        .map(|&axis| output_dims[axis])
278        .collect();
279    let mut output_strides = vec![0isize; output_rank];
280    let mut stride = 1isize;
281    for (&out_axis, &extent) in base_out_axes.iter().zip(&base_dims) {
282        // INVARIANT: `classify_outer_axes` proved every output axis appears
283        // exactly once in the free/batch groups, so each slot is written once.
284        output_strides[out_axis] = stride;
285        let extent = isize::try_from(extent).map_err(|_| StridedError::OffsetOverflow)?;
286        stride = stride
287            .checked_mul(extent)
288            .ok_or(StridedError::OffsetOverflow)?;
289    }
290    Ok(Some(LazyOuterProductLayout {
291        base_dims,
292        output_strides,
293    }))
294}
295
296struct OuterAxisPartition {
297    lhs_free_out: Vec<usize>,
298    rhs_free_out: Vec<usize>,
299    batch_out: Vec<usize>,
300    lhs_free: Vec<usize>,
301    rhs_free: Vec<usize>,
302}
303
304fn extents_match_output(output_dims: &[usize], dims: &[usize], axes: &[usize]) -> bool {
305    dims.iter()
306        .zip(axes)
307        .all(|(&dim, &axis)| output_dims.get(axis) == Some(&dim))
308}
309
310fn operand_axes_by_output(axes: &[usize], output_rank: usize) -> Option<Vec<Option<usize>>> {
311    let mut by_output = vec![None; output_rank];
312    for (operand_axis, &output_axis) in axes.iter().enumerate() {
313        if by_output
314            .get_mut(output_axis)?
315            .replace(operand_axis)
316            .is_some()
317        {
318            return None;
319        }
320    }
321    Some(by_output)
322}
323
324fn classify_outer_axes(
325    output_rank: usize,
326    lhs_axes: &[usize],
327    rhs_axes: &[usize],
328) -> Option<OuterAxisPartition> {
329    let lhs_by_output = operand_axes_by_output(lhs_axes, output_rank)?;
330    let rhs_by_output = operand_axes_by_output(rhs_axes, output_rank)?;
331    let mut partition = OuterAxisPartition {
332        lhs_free_out: Vec::new(),
333        rhs_free_out: Vec::new(),
334        batch_out: Vec::new(),
335        lhs_free: Vec::new(),
336        rhs_free: Vec::new(),
337    };
338    for output_axis in 0..output_rank {
339        match (lhs_by_output[output_axis], rhs_by_output[output_axis]) {
340            (Some(_), Some(_)) => partition.batch_out.push(output_axis),
341            (Some(lhs_axis), None) => {
342                partition.lhs_free_out.push(output_axis);
343                partition.lhs_free.push(lhs_axis);
344            }
345            (None, Some(rhs_axis)) => {
346                partition.rhs_free_out.push(output_axis);
347                partition.rhs_free.push(rhs_axis);
348            }
349            (None, None) => return None,
350        }
351    }
352    Some(partition)
353}
354
355fn checked_axes_product(dims: &[usize], axes: &[usize]) -> Result<usize> {
356    // A zero extent makes the group empty even when the other extents
357    // overflow.
358    if axes.iter().any(|&axis| dims[axis] == 0) {
359        return Ok(0);
360    }
361    axes.iter().try_fold(1usize, |acc, &axis| {
362        acc.checked_mul(dims[axis])
363            .ok_or(StridedError::OffsetOverflow)
364    })
365}
366
367fn axes_by_physical_stride(strides: &[isize], axes: &[usize]) -> Vec<usize> {
368    let mut sorted = axes.to_vec();
369    sorted.sort_by(|&lhs, &rhs| strides[lhs].cmp(&strides[rhs]).then(lhs.cmp(&rhs)));
370    sorted
371}
strided_basic/outer_product.rs

strided_basic/
outer_product.rs
