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

strided_basic/
dense_update.rs
