strided_basic/
update_view.rs

1//! In-place update operations: the destination is also an input.
2//!
3//! `map_into` and `zip_map*_into` write `dest[i] = f(src...)` and never read
4//! the old destination. These functions compute
5//!
6//! ```text
7//! dest[i] = f(OpD(dest[i]), OpA(a[i]), OpB(b[i]))
8//! ```
9//!
10//! reading the previous `dest[i]` through the destination view itself, so no
11//! second view over the destination's storage is ever formed (two live views
12//! of one allocation, one mutable, would be aliasing). Every element of an
13//! injective destination is read and written exactly once, by one call of `f`.
14//!
15//! The element operation of each input, and of the destination read, is a type
16//! parameter (`Identity`, `Conj`, ...), so a conjugating update is a
17//! monomorphized loop with no per-element flag. Traversal, blocking, contiguous
18//! inner loops and threading are those of the other map kernels (threading
19//! follows the budget of the [`Exec`](crate::Exec) token each entry takes).
20
21use crate::kernel::{
22    build_plan_fused, build_plan_fused_small, ensure_same_shape, for_each_inner_block_preordered,
23    sequential_contiguous_layout, total_len, SMALL_TENSOR_THRESHOLD,
24};
25use crate::layout_check::is_injective_layout;
26use crate::simd;
27use crate::view::{StridedView, StridedViewMut};
28use crate::{Exec, Result, StridedError};
29use strided_view::ElementOp;
30
31use crate::cost::KernelClass;
32use crate::fuse::compute_costs;
33use crate::threading::{for_each_inner_block_with_offsets, is_parallel, mapreduce_threaded};
34
35/// The per-inner-block callback: `(offsets, len, strides)`.
36type Body<'a> = &'a (dyn Fn(&[isize], usize, &[isize]) + Sync);
37
38/// A raw pointer that may cross threads: dereferenced only at the validated,
39/// disjoint (for the destination) offsets the traversal generates.
40struct Raw<T>(*mut T);
41impl<T> Clone for Raw<T> {
42    fn clone(&self) -> Self {
43        *self
44    }
45}
46impl<T> Copy for Raw<T> {}
47// SAFETY: see the type documentation; every user upholds it.
48unsafe impl<T> Send for Raw<T> {}
49unsafe impl<T> Sync for Raw<T> {}
50impl<T> Raw<T> {
51    // A method, so closures capture the whole wrapper rather than its field.
52    fn get(self) -> *mut T {
53        self.0
54    }
55}
56
57fn validate_destination(dims: &[usize], strides: &[isize]) -> Result<()> {
58    if is_injective_layout(dims, strides) {
59        Ok(())
60    } else {
61        Err(StridedError::NonInjectiveOutputLayout)
62    }
63}
64
65/// Walk `dims` with the strided traversal and call `body(offsets, len,
66/// strides)` per inner block; the destination is `strides_list[0]`.
67///
68/// The destination layout must already be validated injective, which makes
69/// the blocks' destination regions disjoint when they run on different
70/// threads.
71///
72/// Not generic: the traversal and threading machinery is compiled once here,
73/// and each operation passes its inner loop as a `dyn` callback. The call is
74/// per inner block (a run of elements), never per element, so it costs nothing
75/// measurable while keeping the machinery out of every caller's monomorphization.
76fn run_update(
77    dims: &[usize],
78    strides_list: &[&[isize]],
79    elem_size: usize,
80    body: Body<'_>,
81    exec: Exec<'_>,
82) -> Result<()> {
83    let total = total_len(dims)?;
84    if total == 0 {
85        return Ok(());
86    }
87    // Contiguous fast path: every operand is a flat run in the same order, so
88    // one unit-stride block replaces planning and blocking. The sequential
89    // guard keeps the parallel branch reachable for large tensors (see
90    // `use_sequential_fast_path`).
91    if sequential_contiguous_layout(KernelClass::Map, dims, strides_list, &exec)?.is_some() {
92        let zeros: crate::inline::OperandVec<isize> = strides_list.iter().map(|_| 0isize).collect();
93        let ones: crate::inline::OperandVec<isize> = strides_list.iter().map(|_| 1isize).collect();
94        body(&zeros, total, &ones);
95        return Ok(());
96    }
97    // One element (rank 0, or every extent one): a single length-1 block.
98    if dims.iter().all(|&d| d == 1) {
99        let zeros: crate::inline::OperandVec<isize> = strides_list.iter().map(|_| 0isize).collect();
100        body(&zeros, 1, &zeros);
101        return Ok(());
102    }
103    // Small tensor fast path: skip compute_order and compute_block_sizes
104    let (fused_dims, ordered_strides, plan) = if total <= SMALL_TENSOR_THRESHOLD {
105        build_plan_fused_small(dims, strides_list)
106    } else {
107        build_plan_fused(dims, strides_list, Some(0), elem_size)
108    };
109
110    {
111        let total = total_len(&fused_dims)?;
112        if is_parallel(KernelClass::Map, total, &exec) {
113            let costs = compute_costs(&ordered_strides);
114            let initial_offsets: crate::inline::OperandVec<isize> =
115                strides_list.iter().map(|_| 0isize).collect();
116            return mapreduce_threaded(
117                &fused_dims,
118                &plan.block,
119                &ordered_strides,
120                &initial_offsets,
121                &costs,
122                KernelClass::Map,
123                exec,
124                &|dims, blocks, strides_list, offsets| {
125                    for_each_inner_block_with_offsets(
126                        dims,
127                        blocks,
128                        strides_list,
129                        offsets,
130                        |offsets, len, strides| {
131                            body(offsets, len, strides);
132                            Ok(())
133                        },
134                    )
135                },
136            );
137        }
138    }
139
140    let initial_offsets: crate::inline::OperandVec<isize> =
141        ordered_strides.iter().map(|_| 0isize).collect();
142    for_each_inner_block_preordered(
143        &fused_dims,
144        &plan.block,
145        &ordered_strides,
146        &initial_offsets,
147        |offsets, len, strides| {
148            body(offsets, len, strides);
149            Ok(())
150        },
151    )
152}
153
154// INVARIANT: validated layouts keep all dereferences in bounds. Wrapping
155// advances permit an unused final cursor outside a reversed/gapped allocation.
156/// Unary inner loop: `d[i] = f(OpD(d[i]))`.
157///
158/// # Safety
159/// `dp` addresses `len` live, exclusively owned elements at stride `ds`.
160#[inline(always)]
161unsafe fn inner_loop_update1<D: Copy, OpD: ElementOp<D>>(
162    dp: *mut D,
163    ds: isize,
164    len: usize,
165    f: &impl Fn(D) -> D,
166) {
167    if ds == 1 {
168        let dst = std::slice::from_raw_parts_mut(dp, len);
169        simd::dispatch_if_large(len, || {
170            for d in dst.iter_mut() {
171                *d = f(OpD::apply(*d));
172            }
173        });
174    } else {
175        let mut dp = dp;
176        for _ in 0..len {
177            *dp = f(OpD::apply(*dp));
178            dp = dp.wrapping_offset(ds);
179        }
180    }
181}
182
183/// Binary inner loop: `d[i] = f(OpD(d[i]), OpA(a[i]))`.
184///
185/// # Safety
186/// As [`inner_loop_update1`]; `ap` addresses `len` live elements that do not
187/// overlap the destination.
188#[inline(always)]
189unsafe fn inner_loop_update2<D: Copy, A: Copy, OpD: ElementOp<D>, OpA: ElementOp<A>>(
190    dp: *mut D,
191    ds: isize,
192    ap: *const A,
193    a_s: isize,
194    len: usize,
195    f: &impl Fn(D, A) -> D,
196) {
197    if ds == 1 && a_s == 1 {
198        let src_a = std::slice::from_raw_parts(ap, len);
199        let dst = std::slice::from_raw_parts_mut(dp, len);
200        simd::dispatch_if_large(len, || {
201            for (d, &a) in dst.iter_mut().zip(src_a) {
202                *d = f(OpD::apply(*d), OpA::apply(a));
203            }
204        });
205    } else {
206        let (mut dp, mut ap) = (dp, ap);
207        for _ in 0..len {
208            *dp = f(OpD::apply(*dp), OpA::apply(*ap));
209            dp = dp.wrapping_offset(ds);
210            ap = ap.wrapping_offset(a_s);
211        }
212    }
213}
214
215/// Ternary inner loop: `d[i] = f(OpD(d[i]), OpA(a[i]), OpB(b[i]))`.
216///
217/// # Safety
218/// As [`inner_loop_update2`], with `bp` likewise.
219#[inline(always)]
220#[allow(clippy::too_many_arguments)] // INVARIANT: three strided operands, length and body.
221unsafe fn inner_loop_update3<
222    D: Copy,
223    A: Copy,
224    B: Copy,
225    OpD: ElementOp<D>,
226    OpA: ElementOp<A>,
227    OpB: ElementOp<B>,
228>(
229    dp: *mut D,
230    ds: isize,
231    ap: *const A,
232    a_s: isize,
233    bp: *const B,
234    b_s: isize,
235    len: usize,
236    f: &impl Fn(D, A, B) -> D,
237) {
238    if ds == 1 && a_s == 1 && b_s == 1 {
239        let src_a = std::slice::from_raw_parts(ap, len);
240        let src_b = std::slice::from_raw_parts(bp, len);
241        let dst = std::slice::from_raw_parts_mut(dp, len);
242        simd::dispatch_if_large(len, || {
243            for ((d, &a), &b) in dst.iter_mut().zip(src_a).zip(src_b) {
244                *d = f(OpD::apply(*d), OpA::apply(a), OpB::apply(b));
245            }
246        });
247    } else {
248        let (mut dp, mut ap, mut bp) = (dp, ap, bp);
249        for _ in 0..len {
250            *dp = f(OpD::apply(*dp), OpA::apply(*ap), OpB::apply(*bp));
251            dp = dp.wrapping_offset(ds);
252            ap = ap.wrapping_offset(a_s);
253            bp = bp.wrapping_offset(b_s);
254        }
255    }
256}
257
258/// Update in place: `dest[i] = f(OpD(dest[i]))`.
259///
260/// `OpD` is applied lazily to the old value before `f` sees it.
261///
262/// # Errors
263///
264/// [`StridedError::NonInjectiveOutputLayout`] when `dest` maps two indices to
265/// one element (checked before any write).
266///
267/// # Examples
268///
269/// ```
270/// use strided_basic::{map_update_into, Exec, Identity, StridedArray};
271/// let mut d = StridedArray::<f64>::from_fn_col_major(&[2, 2], |i| (i[0] + 2 * i[1]) as f64);
272/// map_update_into::<_, Identity>(&mut d.view_mut(), |x| 2.0 * x + 1.0, Exec::serial()).unwrap();
273/// assert_eq!(d.get(&[1, 1]), 7.0);
274/// ```
275pub fn map_update_into<D, OpD>(
276    dest: &mut StridedViewMut<D>,
277    f: impl Fn(D) -> D + Sync,
278    exec: Exec<'_>,
279) -> Result<()>
280where
281    D: Copy + Send + Sync,
282    OpD: ElementOp<D>,
283{
284    validate_destination(dest.dims(), dest.strides())?;
285    let dp = Raw(dest.as_mut_ptr());
286    run_update(
287        dest.dims(),
288        &[dest.strides()],
289        std::mem::size_of::<D>(),
290        &|offsets, len, strides| {
291            // SAFETY: the destination is injective and in bounds; blocks are disjoint.
292            unsafe {
293                inner_loop_update1::<D, OpD>(dp.get().offset(offsets[0]), strides[0], len, &f);
294            }
295        },
296        exec,
297    )
298}
299
300/// Update in place from one input: `dest[i] = f(OpD(dest[i]), OpA(a[i]))`.
301///
302/// # Errors
303///
304/// [`StridedError::ShapeMismatch`] for unequal shapes and
305/// [`StridedError::NonInjectiveOutputLayout`], both before any write.
306///
307/// # Examples
308///
309/// ```
310/// use strided_basic::{zip_update2_into, Exec, Identity, StridedArray};
311/// let a = StridedArray::<f64>::from_fn_col_major(&[3], |i| i[0] as f64);
312/// let mut d = StridedArray::<f64>::from_fn_col_major(&[3], |_| 10.0);
313/// // d = 2 * a + 0.5 * d
314/// zip_update2_into::<_, _, Identity, Identity>(
315///     &mut d.view_mut(), &a.view(), |d, a| 2.0 * a + 0.5 * d, Exec::serial(),
316/// )
317/// .unwrap();
318/// assert_eq!(d.get(&[2]), 9.0);
319/// ```
320pub fn zip_update2_into<D, A, OpD, OpA>(
321    dest: &mut StridedViewMut<D>,
322    a: &StridedView<A, OpA>,
323    f: impl Fn(D, A) -> D + Sync,
324    exec: Exec<'_>,
325) -> Result<()>
326where
327    D: Copy + Send + Sync,
328    A: Copy + Send + Sync,
329    OpD: ElementOp<D>,
330    OpA: ElementOp<A>,
331{
332    ensure_same_shape(dest.dims(), a.dims())?;
333    validate_destination(dest.dims(), dest.strides())?;
334    let dp = Raw(dest.as_mut_ptr());
335    let ap = Raw(a.ptr() as *mut A);
336    run_update(
337        dest.dims(),
338        &[dest.strides(), a.strides()],
339        std::mem::size_of::<D>().max(std::mem::size_of::<A>()),
340        &|offsets, len, strides| {
341            // SAFETY: bounds are the views'; `a` is a distinct borrow from `dest`.
342            unsafe {
343                inner_loop_update2::<D, A, OpD, OpA>(
344                    dp.get().offset(offsets[0]),
345                    strides[0],
346                    ap.get().offset(offsets[1]).cast_const(),
347                    strides[1],
348                    len,
349                    &f,
350                );
351            }
352        },
353        exec,
354    )
355}
356
357/// Update in place from two inputs:
358/// `dest[i] = f(OpD(dest[i]), OpA(a[i]), OpB(b[i]))`.
359///
360/// # Errors
361///
362/// As [`zip_update2_into`].
363///
364/// # Examples
365///
366/// ```
367/// use strided_basic::{zip_update3_into, Conj, Exec, Identity, StridedArray};
368/// use num_complex::Complex64;
369/// let one = Complex64::new(1.0, 1.0);
370/// let a = StridedArray::<Complex64>::from_fn_col_major(&[2], |_| one);
371/// let b = StridedArray::<Complex64>::from_fn_col_major(&[2], |_| one);
372/// let mut d = StridedArray::<Complex64>::from_fn_col_major(&[2], |_| one);
373/// // d = conj(d) + a * b
374/// zip_update3_into::<_, _, _, Conj, Identity, Identity>(
375///     &mut d.view_mut(), &a.view(), &b.view(), |d, a, b| d + a * b, Exec::serial(),
376/// )
377/// .unwrap();
378/// assert_eq!(d.get(&[0]), Complex64::new(1.0, 1.0)); // (1 - i) + 2i
379/// ```
380pub fn zip_update3_into<D, A, B, OpD, OpA, OpB>(
381    dest: &mut StridedViewMut<D>,
382    a: &StridedView<A, OpA>,
383    b: &StridedView<B, OpB>,
384    f: impl Fn(D, A, B) -> D + Sync,
385    exec: Exec<'_>,
386) -> Result<()>
387where
388    D: Copy + Send + Sync,
389    A: Copy + Send + Sync,
390    B: Copy + Send + Sync,
391    OpD: ElementOp<D>,
392    OpA: ElementOp<A>,
393    OpB: ElementOp<B>,
394{
395    ensure_same_shape(dest.dims(), a.dims())?;
396    ensure_same_shape(dest.dims(), b.dims())?;
397    validate_destination(dest.dims(), dest.strides())?;
398    let dp = Raw(dest.as_mut_ptr());
399    let ap = Raw(a.ptr() as *mut A);
400    let bp = Raw(b.ptr() as *mut B);
401    run_update(
402        dest.dims(),
403        &[dest.strides(), a.strides(), b.strides()],
404        std::mem::size_of::<D>()
405            .max(std::mem::size_of::<A>())
406            .max(std::mem::size_of::<B>()),
407        &|offsets, len, strides| {
408            // SAFETY: bounds are the views'; `a`, `b` are distinct borrows from `dest`.
409            unsafe {
410                inner_loop_update3::<D, A, B, OpD, OpA, OpB>(
411                    dp.get().offset(offsets[0]),
412                    strides[0],
413                    ap.get().offset(offsets[1]).cast_const(),
414                    strides[1],
415                    bp.get().offset(offsets[2]).cast_const(),
416                    strides[2],
417                    len,
418                    &f,
419                );
420            }
421        },
422        exec,
423    )
424}
425
426#[cfg(test)]
427#[path = "update_view/tests/tests.rs"]
428mod tests;
strided_basic/update_view.rs

strided_basic/
update_view.rs
