strided_basic/
dense_update.rs

1//! Dense column-major update and structural kernels.
2//!
3//! AXPBY and triangular masking are adapted from tenferro-rs d8759f43,
4//! `crates/tenferro-cpu/src/{blas1,structural}.rs` (MIT OR Apache-2.0).
5//! Tensor allocation/placement stay with the caller. Diagonal embedding uses
6//! the existing strided map traversal rather than tenferro's coordinate loop.
7
8use core::mem::MaybeUninit;
9use core::ops::{Add, Mul};
10
11use crate::cost::KernelClass;
12use crate::{Exec, Result, StridedError, StridedView, StridedViewMut};
13
14fn same_len(actual: usize, expected: usize) -> Result<()> {
15    if actual == expected {
16        Ok(())
17    } else {
18        Err(StridedError::ShapeMismatch(vec![actual], vec![expected]))
19    }
20}
21
22fn product(shape: &[usize]) -> Result<usize> {
23    if shape.contains(&0) {
24        return Ok(0);
25    }
26    shape.iter().try_fold(1usize, |n, &d| {
27        n.checked_mul(d).ok_or(StridedError::OffsetOverflow)
28    })
29}
30
31fn for_each_chunk<T: Send + Sync, F: Fn(usize, &mut [T]) + Sync>(
32    output: &mut [T],
33    operation: F,
34    exec: Exec<'_>,
35) {
36    let exec = crate::cost::for_len(exec, KernelClass::Map, output.len());
37    if !exec.is_serial() {
38        let ptr = crate::threading::SendPtr(output.as_mut_ptr());
39        crate::threading::parallel_for_each(0..output.len(), exec, &|range| {
40            // SAFETY: the scoped scheduler partitions the exclusive output
41            // borrow into disjoint ranges; all tasks finish before return.
42            let chunk = unsafe {
43                core::slice::from_raw_parts_mut(ptr.as_ptr().add(range.start), range.len())
44            };
45            operation(range.start, chunk);
46        });
47        return;
48    }
49    operation(0, output);
50}
51
52/// Source offset of a linear batch index within `shape[2..]`.
53fn batch_source_offset(mut batch: usize, shape: &[usize], strides: &[isize]) -> isize {
54    let mut offset = 0isize;
55    for (&dim, &stride) in shape.iter().zip(strides) {
56        let coord = batch % dim;
57        batch /= dim;
58        offset += coord as isize * stride;
59    }
60    offset
61}
62
63/// Borrow `data` as a dense column-major view of `shape`.
64fn dense_view<'a, T>(data: &'a [T], shape: &[usize]) -> Result<StridedView<'a, T>> {
65    let len = product(shape)?;
66    same_len(data.len(), len)?;
67    let strides = crate::col_major_strides(shape);
68    StridedView::new(data, shape, &strides, 0)
69}
70
71/// Copy `dst.len()` source elements starting at reachable byte-free offset
72/// `base` with element step `stride`.
73fn copy_strided_run<T: Copy>(dst: &mut [MaybeUninit<T>], data: &[T], base: isize, stride: isize) {
74    if stride == 1 {
75        // INVARIANT: `base` is reachable inside a view-validated buffer, so
76        // the run start plus its length stays in bounds.
77        let start = base as usize;
78        let src = &data[start..start + dst.len()];
79        for (slot, &value) in dst.iter_mut().zip(src) {
80            slot.write(value);
81        }
82    } else {
83        let mut offset = base;
84        for slot in dst.iter_mut() {
85            slot.write(data[offset as usize]);
86            offset += stride;
87        }
88    }
89}
90
91/// Update caller-owned contiguous storage: `y = alpha * x + beta * y`.
92///
93/// This is an accumulation operation: it reads the old `y`, even for beta=0
94/// (including ordinary IEEE NaN propagation). No temporary tensor is allocated.
95/// The update splits within the budget of `exec` above the shared parallel
96/// threshold.
97///
98/// # Examples
99/// ```
100/// let mut y = [3.0, 4.0];
101/// strided_basic::axpby_accum(&mut y, &[1.0, 2.0], 2.0, 3.0, strided_basic::Exec::serial()).unwrap();
102/// assert_eq!(y, [11.0, 16.0]);
103/// ```
104/// # Errors
105/// Returns `ShapeMismatch` for unequal slice lengths, before modifying `y`.
106pub fn axpby_accum<T>(y: &mut [T], x: &[T], alpha: T, beta: T, exec: Exec<'_>) -> Result<()>
107where
108    T: Copy + Send + Sync + Add<Output = T> + Mul<Output = T>,
109{
110    same_len(x.len(), y.len())?;
111    for_each_chunk(
112        y,
113        |start, dst| {
114            let len = dst.len();
115            for (out, &src) in dst.iter_mut().zip(&x[start..start + len]) {
116                *out = alpha * src + beta * *out;
117            }
118        },
119        exec,
120    );
121    Ok(())
122}
123
124/// Copy dense column-major matrices, replacing the masked triangle with `fill`.
125///
126/// The first two dimensions are rows and columns; remaining dimensions are
127/// batches. `upper=true` keeps `row <= column-k`, otherwise the lower triangle.
128/// Every output element is initialized on success; old output is never read.
129///
130/// # Examples
131/// ```
132/// use core::mem::MaybeUninit;
133/// let mut out = [MaybeUninit::uninit(); 4];
134/// strided_basic::triangular_mask_into_uninit(
135///     &mut out, &[1, 2, 3, 4], &[2, 2], 0, false, 0, strided_basic::Exec::serial()).unwrap();
136/// // SAFETY: successful full-overwrite kernel initialized every element.
137/// assert_eq!(out.map(|v| unsafe { v.assume_init() }), [1, 2, 0, 4]);
138/// ```
139/// # Errors
140/// Returns `RankMismatch`, `ShapeMismatch`, or `OffsetOverflow` for invalid
141/// shape/buffer lengths, before writing output.
142pub fn triangular_mask_into_uninit<T: Copy + Send + Sync>(
143    output: &mut [MaybeUninit<T>],
144    input: &[T],
145    shape: &[usize],
146    k: i64,
147    upper: bool,
148    fill: T,
149    exec: Exec<'_>,
150) -> Result<()> {
151    let src = dense_view(input, shape)?;
152    triangular_mask_view_into_uninit(output, &src, k, upper, fill, exec)
153}
154
155/// Mask a triangular region of a strided input into caller-owned contiguous
156/// storage.
157///
158/// The source layout is given by `src`; reads follow its strides and offset, so
159/// a reversed, transposed or otherwise strided source needs no staging copy.
160/// Kept elements come from `src`; masked elements receive `fill`. Every output
161/// element is written exactly once and the masked region of the source is
162/// never read.
163///
164/// # Errors
165/// Returns `RankMismatch` or `ShapeMismatch` for invalid shapes and buffer
166/// lengths, and `OffsetOverflow` for unrepresentable layouts, before writing.
167pub fn triangular_mask_view_into_uninit<T: Copy + Send + Sync>(
168    output: &mut [MaybeUninit<T>],
169    src: &StridedView<'_, T>,
170    k: i64,
171    upper: bool,
172    fill: T,
173    exec: Exec<'_>,
174) -> Result<()> {
175    let shape = src.dims();
176    if shape.len() < 2 {
177        return Err(StridedError::RankMismatch(shape.len(), 2));
178    }
179    let len = product(shape)?;
180    same_len(output.len(), len)?;
181    if len == 0 {
182        return Ok(());
183    }
184    let rows = shape[0];
185    let cols = shape[1];
186    let data = src.data();
187    let strides = src.strides();
188    let offset = src.offset();
189    for_each_chunk(
190        output,
191        |start, chunk| {
192            let mut pos = 0;
193            let mut batch = usize::MAX;
194            let mut batch_offset = 0isize;
195            while pos < chunk.len() {
196                let flat = start + pos;
197                let row = flat % rows;
198                let col = (flat / rows) % cols;
199                // INVARIANT: `product(shape)` succeeded, so `rows * cols` fits.
200                let run_batch = flat / (rows * cols);
201                if run_batch != batch {
202                    batch = run_batch;
203                    batch_offset = batch_source_offset(run_batch, &shape[2..], &strides[2..]);
204                }
205                let count = (rows - row).min(chunk.len() - pos);
206                let boundary = col as i128 - k as i128;
207                let split = if upper {
208                    boundary.saturating_add(1)
209                } else {
210                    boundary
211                }
212                .clamp(row as i128, (row + count) as i128) as usize
213                    - row;
214                let (low, high) = chunk[pos..pos + count].split_at_mut(split);
215                let (kept, first_row, masked) = if upper {
216                    (low, row, high)
217                } else {
218                    (high, row + split, low)
219                };
220                let base = offset
221                    + batch_offset
222                    + (col as isize) * strides[1]
223                    + (first_row as isize) * strides[0];
224                // Full overwrite: never read masked input or copy it only to
225                // overwrite it a second time. Each output element is written once.
226                copy_strided_run(kept, data, base, strides[0]);
227                masked.fill(MaybeUninit::new(fill));
228                pos += count;
229            }
230        },
231        exec,
232    );
233    Ok(())
234}
235
236/// Embed a dense column-major input on a diagonal in a new axis.
237///
238/// Insert an axis of size `shape[axis]` at `insert_axis`. Its coordinate must
239/// equal the original `axis` coordinate; other entries receive `zero`.
240/// Output is fully initialized without reading its previous contents. The
241/// budget of `exec` governs fill and strided diagonal-copy traversal.
242///
243/// # Examples
244/// ```
245/// use core::mem::MaybeUninit;
246/// let mut out = [MaybeUninit::uninit(); 4];
247/// strided_basic::embed_diagonal_into_uninit(&mut out, &[3, 5], &[2], 0, 1, 0, strided_basic::Exec::serial()).unwrap();
248/// // SAFETY: successful full-overwrite kernel initialized every element.
249/// assert_eq!(out.map(|v| unsafe { v.assume_init() }), [3, 0, 0, 5]);
250/// ```
251/// # Errors
252/// Returns `InvalidAxis`, `ShapeMismatch`, or `OffsetOverflow` for invalid
253/// axes, lengths or unrepresentable layouts. Destination layout validation
254/// may return `NonInjectiveOutputLayout` if injectivity cannot be established.
255pub fn embed_diagonal_into_uninit<T: Copy + Send + Sync + 'static>(
256    output: &mut [MaybeUninit<T>],
257    input: &[T],
258    shape: &[usize],
259    axis: usize,
260    insert_axis: usize,
261    zero: T,
262    exec: Exec<'_>,
263) -> Result<()> {
264    if axis >= shape.len() {
265        return Err(StridedError::InvalidAxis {
266            axis,
267            rank: shape.len(),
268        });
269    }
270    if insert_axis > shape.len() {
271        return Err(StridedError::InvalidAxis {
272            axis: insert_axis,
273            rank: shape.len() + 1,
274        });
275    }
276    let len = product(shape)?;
277    same_len(input.len(), len)?;
278    let src = dense_view(input, shape)?;
279    embed_diagonal_view_into_uninit(output, &src, axis, insert_axis, zero, exec)
280}
281
282/// Embed a strided input on a diagonal in a new axis of caller-owned
283/// contiguous storage.
284///
285/// The source layout is given by `src`; the diagonal copy follows its strides
286/// and offset, so a strided source needs no staging copy. Off-diagonal entries
287/// receive `zero`, which is part of the result rather than scratch
288/// initialization.
289///
290/// # Errors
291/// Returns `InvalidAxis`, `ShapeMismatch`, or `OffsetOverflow` for invalid
292/// axes, lengths or unrepresentable layouts. Destination layout validation may
293/// return `NonInjectiveOutputLayout` if injectivity cannot be established.
294pub fn embed_diagonal_view_into_uninit<T: Copy + Send + Sync + 'static>(
295    output: &mut [MaybeUninit<T>],
296    src: &StridedView<'_, T>,
297    axis: usize,
298    insert_axis: usize,
299    zero: T,
300    exec: Exec<'_>,
301) -> Result<()> {
302    let mut exec = exec;
303    let shape = src.dims();
304    if axis >= shape.len() {
305        return Err(StridedError::InvalidAxis {
306            axis,
307            rank: shape.len(),
308        });
309    }
310    if insert_axis > shape.len() {
311        return Err(StridedError::InvalidAxis {
312            axis: insert_axis,
313            rank: shape.len() + 1,
314        });
315    }
316    let len = product(shape)?;
317    let out_len = len
318        .checked_mul(shape[axis])
319        .ok_or(StridedError::OffsetOverflow)?;
320    same_len(output.len(), out_len)?;
321    if out_len == 0 {
322        return Ok(());
323    }
324    isize::try_from(out_len).map_err(|_| StridedError::OffsetOverflow)?;
325    let mut out_shape: crate::InlineVec<usize> = shape.iter().copied().collect();
326    out_shape.insert(insert_axis, shape[axis]);
327    let mut dst_strides = crate::col_major_strides(&out_shape);
328    let inserted_stride = dst_strides.remove(insert_axis);
329    dst_strides[axis] = dst_strides[axis]
330        .checked_add(inserted_stride)
331        .ok_or(StridedError::OffsetOverflow)?;
332    let mut dst = StridedViewMut::new(output, shape, &dst_strides, 0)?;
333    crate::map_view::validate_destination_layout_without_alloc(shape, &dst_strides)?;
334    // INVARIANT: off-diagonal zeros are part of the mathematical result, not
335    // scratch initialization. Only the diagonal subset is overwritten below.
336    for_each_chunk(
337        dst.data_mut(),
338        |_, chunk| chunk.fill(MaybeUninit::new(zero)),
339        exec.reborrow(),
340    );
341    crate::map_view::copy_map_into(&mut dst, src, MaybeUninit::new, exec)
342}
strided_basic/dense_update.rs

strided_basic/
dense_update.rs
