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

strided_basic/
update_view.rs
