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

strided_basic/
update_view.rs
