Skip to main content

strided_basic/
map_view.rs

1//! Map operations on dynamic-rank strided views.
2//!
3//! These are the canonical view-based map functions, equivalent to Julia's `Base.map!`.
4//! Mutable destinations must have an injective layout. Validation is
5//! conservative and bounded: layouts that are injective but cannot be proven
6//! by the bounded checker are rejected with
7//! [`StridedError::NonInjectiveOutputLayout`].
8
9use crate::kernel::{
10    build_plan_fused, build_plan_fused_small, ensure_same_shape, for_each_inner_block_preordered,
11    sequential_contiguous_layout, total_len, SMALL_TENSOR_THRESHOLD,
12};
13use crate::maybe_sync::{MaybeSendSync, MaybeSync};
14use crate::simd;
15use crate::view::{StridedView, StridedViewMut};
16use crate::{Result, StridedError};
17use core::mem::MaybeUninit;
18use std::ops::Mul;
19use strided_view::{ElementOp, Identity};
20
21#[cfg(feature = "parallel")]
22use crate::fuse::compute_costs;
23#[cfg(feature = "parallel")]
24use crate::threading::{for_each_inner_block_with_offsets, mapreduce_threaded, MINTHREADLENGTH};
25#[cfg(feature = "parallel")]
26use smallvec::SmallVec;
27
28#[cfg(feature = "parallel")]
29type AxisVec<T> = SmallVec<[T; 8]>;
30#[cfg(not(feature = "parallel"))]
31type AxisVec<T> = Vec<T>;
32
33const CONTIGUOUS_RANGE_MIN_LEN: usize = 1 << 15;
34
35/// Internal validation marker used by kernel-family execution adapters.
36///
37/// This marker is not bound to a destination. It does not authorize unchecked
38/// replay with a different geometry; callers of the unsafe execution entries
39/// must establish that association themselves.
40#[derive(Clone, Copy, Debug)]
41pub struct ValidatedDestinationLayout(());
42
43#[inline]
44fn validate_destination_layout(
45    dims: &[usize],
46    strides: &[isize],
47) -> Result<ValidatedDestinationLayout> {
48    if crate::layout_check::is_injective_layout(dims, strides) {
49        Ok(ValidatedDestinationLayout(()))
50    } else {
51        Err(StridedError::NonInjectiveOutputLayout)
52    }
53}
54
55#[inline]
56/// Validate injectivity without allocating a set of offsets.
57///
58/// # Examples
59///
60/// ```
61/// use strided_basic::execution::validate_destination_layout_without_alloc;
62/// let marker = validate_destination_layout_without_alloc(&[2], &[1]).unwrap();
63/// assert!(validate_destination_layout_without_alloc(&[2], &[0]).is_err());
64/// let _ = marker;
65/// ```
66///
67/// # Errors
68/// Returns `NonInjectiveOutputLayout` for invalid or unsupported output geometry.
69pub fn validate_destination_layout_without_alloc(
70    dims: &[usize],
71    strides: &[isize],
72) -> Result<ValidatedDestinationLayout> {
73    if crate::layout_check::is_injective_layout_without_alloc(dims, strides) {
74        Ok(ValidatedDestinationLayout(()))
75    } else {
76        Err(StridedError::NonInjectiveOutputLayout)
77    }
78}
79
80fn reachable_byte_range(
81    ptr: usize,
82    elem_size: usize,
83    dims: &[usize],
84    strides: &[isize],
85) -> Result<Option<(usize, usize)>> {
86    if elem_size == 0 || dims.contains(&0) {
87        return Ok(None);
88    }
89
90    let mut min_offset = 0isize;
91    let mut max_offset = 0isize;
92    for (&dim, &stride) in dims.iter().zip(strides) {
93        if dim <= 1 {
94            continue;
95        }
96        let extent = isize::try_from(dim - 1).map_err(|_| StridedError::OffsetOverflow)?;
97        let span = stride
98            .checked_mul(extent)
99            .ok_or(StridedError::OffsetOverflow)?;
100        if span < 0 {
101            min_offset = min_offset
102                .checked_add(span)
103                .ok_or(StridedError::OffsetOverflow)?;
104        } else {
105            max_offset = max_offset
106                .checked_add(span)
107                .ok_or(StridedError::OffsetOverflow)?;
108        }
109    }
110
111    let ptr = ptr as i128;
112    let elem_size = elem_size as i128;
113    let start = ptr
114        .checked_add(
115            (min_offset as i128)
116                .checked_mul(elem_size)
117                .ok_or(StridedError::OffsetOverflow)?,
118        )
119        .ok_or(StridedError::OffsetOverflow)?;
120    let end = ptr
121        .checked_add(
122            (max_offset as i128)
123                .checked_mul(elem_size)
124                .and_then(|offset| offset.checked_add(elem_size))
125                .ok_or(StridedError::OffsetOverflow)?,
126        )
127        .ok_or(StridedError::OffsetOverflow)?;
128    if start < 0 || end < 0 || start > usize::MAX as i128 || end > usize::MAX as i128 {
129        return Err(StridedError::OffsetOverflow);
130    }
131    Ok(Some((start as usize, end as usize)))
132}
133
134fn validate_typed_no_overlap<D, A, Op: ElementOp<A>>(
135    dest: &StridedViewMut<MaybeUninit<D>>,
136    input: &StridedView<A, Op>,
137    input_index: usize,
138) -> Result<()> {
139    let dest_range = reachable_byte_range(
140        dest.ptr() as usize,
141        core::mem::size_of::<D>(),
142        dest.dims(),
143        dest.strides(),
144    )?;
145    let input_range = reachable_byte_range(
146        input.ptr() as usize,
147        core::mem::size_of::<A>(),
148        input.dims(),
149        input.strides(),
150    )?;
151    if let (Some((dest_start, dest_end)), Some((input_start, input_end))) =
152        (dest_range, input_range)
153    {
154        if dest_start < input_end && input_start < dest_end {
155            return Err(StridedError::OverlappingInputOutput { input: input_index });
156        }
157    }
158    Ok(())
159}
160
161// ============================================================================
162// Stride-specialized inner loop helpers
163//
164// When all inner strides are 1 (contiguous in the innermost dimension),
165// we use slice-based iteration so LLVM can auto-vectorize effectively.
166// This is the Rust equivalent of Julia's @simd on the innermost loop.
167// ============================================================================
168
169// INVARIANT: validated layouts keep every dereferenced cursor in bounds.
170// Pointwise leaves use wrapping advances: the unused final cursor may lie
171// outside the allocation for negative or gapped strides.
172/// Unary inner loop: `dest[i] = f(Op::apply(src[i]))` for `len` elements.
173///
174/// A vectorized kernel a map operation offers its flat runs to.
175///
176/// [`NoKernel`] is what the entry points without a kernel pass, so the check it introduces is
177/// monomorphized away; [`SliceKernel`] adapts the closure a caller supplied.
178trait MapKernel<A, D> {
179    /// Writes every element of `dst` from `src`, or writes nothing and returns `false`.
180    fn apply(&self, src: &[A], dst: &mut [D]) -> bool;
181}
182
183struct NoKernel;
184
185impl<A, D> MapKernel<A, D> for NoKernel {
186    #[inline(always)]
187    fn apply(&self, _src: &[A], _dst: &mut [D]) -> bool {
188        false
189    }
190}
191
192struct SliceKernel<K>(K);
193
194impl<A, D, K: Fn(&[A], &mut [D]) -> bool> MapKernel<A, D> for SliceKernel<K> {
195    #[inline(always)]
196    fn apply(&self, src: &[A], dst: &mut [D]) -> bool {
197        (self.0)(src, dst)
198    }
199}
200
201unsafe fn inner_loop_map1<D: Copy, A: Copy, Op: ElementOp<A>, K: MapKernel<A, D>>(
202    dp: *mut D,
203    ds: isize,
204    sp: *const A,
205    ss: isize,
206    len: usize,
207    f: &impl Fn(A) -> D,
208    kernel: &K,
209) {
210    if ds == 1 && ss == 1 {
211        let src = std::slice::from_raw_parts(sp, len);
212        let dst = std::slice::from_raw_parts_mut(dp, len);
213        if kernel.apply(src, dst) {
214            return;
215        }
216        simd::dispatch_if_large(len, || {
217            for (d, s) in dst.iter_mut().zip(src.iter()) {
218                *d = f(Op::apply(*s));
219            }
220        });
221    } else {
222        let mut dp = dp;
223        let mut sp = sp;
224        for _ in 0..len {
225            *dp = f(Op::apply(*sp));
226            dp = dp.wrapping_offset(ds);
227            sp = sp.wrapping_offset(ss);
228        }
229    }
230}
231
232/// Binary inner loop: `dest[i] = f(OpA::apply(a[i]), OpB::apply(b[i]))`.
233#[inline(always)]
234unsafe fn inner_loop_map2<D: Copy, A: Copy, B: Copy, OpA: ElementOp<A>, OpB: ElementOp<B>>(
235    dp: *mut D,
236    ds: isize,
237    ap: *const A,
238    a_s: isize,
239    bp: *const B,
240    b_s: isize,
241    len: usize,
242    f: &impl Fn(A, B) -> D,
243) {
244    if ds == 1 && a_s == 1 && b_s == 1 {
245        let src_a = std::slice::from_raw_parts(ap, len);
246        let src_b = std::slice::from_raw_parts(bp, len);
247        let dst = std::slice::from_raw_parts_mut(dp, len);
248        simd::dispatch_if_large(len, || {
249            for i in 0..len {
250                dst[i] = f(OpA::apply(src_a[i]), OpB::apply(src_b[i]));
251            }
252        });
253    } else if ds == 1 && a_s == 1 && b_s == 0 {
254        let src_a = std::slice::from_raw_parts(ap, len);
255        let b = OpB::apply(*bp);
256        let dst = std::slice::from_raw_parts_mut(dp, len);
257        simd::dispatch_if_large(len, || {
258            for i in 0..len {
259                dst[i] = f(OpA::apply(src_a[i]), b);
260            }
261        });
262    } else if ds == 1 && a_s == 0 && b_s == 1 {
263        let a = OpA::apply(*ap);
264        let src_b = std::slice::from_raw_parts(bp, len);
265        let dst = std::slice::from_raw_parts_mut(dp, len);
266        simd::dispatch_if_large(len, || {
267            for i in 0..len {
268                dst[i] = f(a, OpB::apply(src_b[i]));
269            }
270        });
271    } else if ds == 1 && a_s == 0 && b_s == 0 {
272        let a = OpA::apply(*ap);
273        let b = OpB::apply(*bp);
274        let dst = std::slice::from_raw_parts_mut(dp, len);
275        simd::dispatch_if_large(len, || {
276            for d in dst.iter_mut() {
277                *d = f(a, b);
278            }
279        });
280    } else if ds == 1 && b_s == 0 {
281        let b = OpB::apply(*bp);
282        let dst = std::slice::from_raw_parts_mut(dp, len);
283        let mut ap = ap;
284        simd::dispatch_if_large(len, || {
285            for d in dst.iter_mut() {
286                *d = f(OpA::apply(*ap), b);
287                ap = ap.wrapping_offset(a_s);
288            }
289        });
290    } else if ds == 1 && a_s == 0 {
291        let a = OpA::apply(*ap);
292        let dst = std::slice::from_raw_parts_mut(dp, len);
293        let mut bp = bp;
294        simd::dispatch_if_large(len, || {
295            for d in dst.iter_mut() {
296                *d = f(a, OpB::apply(*bp));
297                bp = bp.wrapping_offset(b_s);
298            }
299        });
300    } else {
301        let mut dp = dp;
302        let mut ap = ap;
303        let mut bp = bp;
304        for _ in 0..len {
305            *dp = f(OpA::apply(*ap), OpB::apply(*bp));
306            dp = dp.wrapping_offset(ds);
307            ap = ap.wrapping_offset(a_s);
308            bp = bp.wrapping_offset(b_s);
309        }
310    }
311}
312
313/// Binary multiplication inner loop for identity element ops.
314trait MulOutput<D: Copy + 'static>: Copy + MaybeSendSync + 'static {
315    type Slot: Copy + MaybeSendSync + 'static;
316
317    unsafe fn write(dst: *mut Self::Slot, value: D);
318
319    unsafe fn try_contiguous<A: 'static, B: 'static>(
320        dst: *mut Self::Slot,
321        len: usize,
322        a: &[A],
323        b: &[B],
324    ) -> bool;
325}
326
327#[derive(Clone, Copy)]
328struct InitializedOutput;
329
330impl<D: Copy + MaybeSendSync + 'static> MulOutput<D> for InitializedOutput {
331    type Slot = D;
332
333    #[inline(always)]
334    unsafe fn write(dst: *mut D, value: D) {
335        unsafe { dst.write(value) };
336    }
337
338    #[inline(always)]
339    unsafe fn try_contiguous<A: 'static, B: 'static>(
340        dst: *mut D,
341        len: usize,
342        a: &[A],
343        b: &[B],
344    ) -> bool {
345        unsafe { simd::try_mul_contiguous_ptr(dst, len, a, b) }
346    }
347}
348
349#[derive(Clone, Copy)]
350struct UninitializedOutput;
351
352impl<D: Copy + MaybeSendSync + 'static> MulOutput<D> for UninitializedOutput {
353    type Slot = MaybeUninit<D>;
354
355    #[inline(always)]
356    unsafe fn write(dst: *mut MaybeUninit<D>, value: D) {
357        unsafe { dst.write(MaybeUninit::new(value)) };
358    }
359
360    #[inline(always)]
361    unsafe fn try_contiguous<A: 'static, B: 'static>(
362        dst: *mut MaybeUninit<D>,
363        len: usize,
364        a: &[A],
365        b: &[B],
366    ) -> bool {
367        // SIMD stores write complete values directly through raw pointers; no
368        // initialized reference is formed over the uninitialized destination.
369        unsafe { simd::try_mul_contiguous_ptr(dst.cast::<D>(), len, a, b) }
370    }
371}
372
373#[inline(always)]
374fn multiply_value<A, B, D>(lhs: A, rhs: B) -> D
375where
376    A: Copy + Mul<B, Output = D> + 'static,
377    B: Copy + 'static,
378    D: Copy + 'static,
379{
380    use core::any::TypeId;
381
382    if TypeId::of::<A>() == TypeId::of::<i32>()
383        && TypeId::of::<B>() == TypeId::of::<i32>()
384        && TypeId::of::<D>() == TypeId::of::<i32>()
385    {
386        // The exact TypeId checks prove identical layout and validity.
387        let lhs = unsafe { *(&lhs as *const A).cast::<i32>() };
388        let rhs = unsafe { *(&rhs as *const B).cast::<i32>() };
389        let value = lhs.wrapping_mul(rhs);
390        return unsafe { core::mem::transmute_copy(&value) };
391    }
392    if TypeId::of::<A>() == TypeId::of::<i64>()
393        && TypeId::of::<B>() == TypeId::of::<i64>()
394        && TypeId::of::<D>() == TypeId::of::<i64>()
395    {
396        // The exact TypeId checks prove identical layout and validity.
397        let lhs = unsafe { *(&lhs as *const A).cast::<i64>() };
398        let rhs = unsafe { *(&rhs as *const B).cast::<i64>() };
399        let value = lhs.wrapping_mul(rhs);
400        return unsafe { core::mem::transmute_copy(&value) };
401    }
402    lhs * rhs
403}
404
405#[inline(always)]
406unsafe fn inner_loop_mul2<
407    O: MulOutput<D>,
408    D: Copy + 'static,
409    A: Copy + Mul<B, Output = D> + 'static,
410    B: Copy + 'static,
411>(
412    dp: *mut O::Slot,
413    ds: isize,
414    ap: *const A,
415    a_s: isize,
416    bp: *const B,
417    b_s: isize,
418    len: usize,
419) {
420    if ds == 1 && a_s == 1 && b_s == 1 {
421        let src_a = std::slice::from_raw_parts(ap, len);
422        let src_b = std::slice::from_raw_parts(bp, len);
423        if len >= 64 && O::try_contiguous(dp, len, src_a, src_b) {
424            return;
425        }
426        for i in 0..len {
427            O::write(dp.add(i), multiply_value(src_a[i], src_b[i]));
428        }
429    } else if ds == 1 && a_s == 1 && b_s == 0 {
430        let src_a = std::slice::from_raw_parts(ap, len);
431        let b = *bp;
432        for i in 0..len {
433            O::write(dp.add(i), multiply_value(src_a[i], b));
434        }
435    } else if ds == 1 && a_s == 0 && b_s == 1 {
436        let a = *ap;
437        let src_b = std::slice::from_raw_parts(bp, len);
438        for i in 0..len {
439            O::write(dp.add(i), multiply_value(a, src_b[i]));
440        }
441    } else if ds == 1 && a_s == 0 && b_s == 0 {
442        let a = *ap;
443        let b = *bp;
444        for i in 0..len {
445            O::write(dp.add(i), multiply_value(a, b));
446        }
447    } else if ds == 1 && b_s == 0 {
448        let b = *bp;
449        let mut ap = ap;
450        for i in 0..len {
451            O::write(dp.add(i), multiply_value(*ap, b));
452            ap = ap.wrapping_offset(a_s);
453        }
454    } else if ds == 1 && a_s == 0 {
455        let a = *ap;
456        let mut bp = bp;
457        for i in 0..len {
458            O::write(dp.add(i), multiply_value(a, *bp));
459            bp = bp.wrapping_offset(b_s);
460        }
461    } else {
462        let mut dp = dp;
463        let mut ap = ap;
464        let mut bp = bp;
465        for _ in 0..len {
466            O::write(dp, multiply_value(*ap, *bp));
467            dp = dp.wrapping_offset(ds);
468            ap = ap.wrapping_offset(a_s);
469            bp = bp.wrapping_offset(b_s);
470        }
471    }
472}
473
474#[derive(Clone, Debug, Eq, PartialEq)]
475struct ContiguousMulRangePlan {
476    axis_order: AxisVec<usize>,
477    inner_len: usize,
478    inner_axis_count: usize,
479    row_len: usize,
480    outer_axis_start: usize,
481    fast_axis: usize,
482    a_fast_stride: isize,
483    b_fast_stride: isize,
484    a_row_stride: isize,
485    b_row_stride: isize,
486}
487
488impl ContiguousMulRangePlan {
489    /// Whether running the plan row by row, without cache blocking, is sound
490    /// for performance.
491    ///
492    /// The destination is walked contiguously, so a unit-stride or broadcast
493    /// (stride 0) input streams along with it. An input read at any other
494    /// stride along the fast axis (a transposed operand) is swept across the
495    /// whole fast axis once per row, touching a new cache line per element and
496    /// again on the next row; the blocked kernel tiles that access. The one
497    /// exception is the SIMD transposed-scalar tile, which blocks itself.
498    fn walks_inputs_without_blocking(&self) -> bool {
499        let streams = |stride: isize| stride == 0 || stride == 1;
500        if streams(self.a_fast_stride) && streams(self.b_fast_stride) {
501            return true;
502        }
503        #[cfg(feature = "parallel")]
504        if transposed_scalar_tile_kind(self).is_some() {
505            return true;
506        }
507        false
508    }
509}
510
511#[cfg(feature = "parallel")]
512#[derive(Clone, Copy, Debug, Eq, PartialEq)]
513enum TransposedScalarTileKind {
514    RhsScalar,
515    LhsScalar,
516}
517
518#[cfg(feature = "parallel")]
519fn transposed_scalar_tile_kind(plan: &ContiguousMulRangePlan) -> Option<TransposedScalarTileKind> {
520    let row_len = isize::try_from(plan.row_len).ok()?;
521    if plan.a_fast_stride == row_len
522        && plan.a_row_stride == 1
523        && plan.b_fast_stride == 0
524        && plan.b_row_stride == 0
525    {
526        return Some(TransposedScalarTileKind::RhsScalar);
527    }
528
529    if plan.b_fast_stride == row_len
530        && plan.b_row_stride == 1
531        && plan.a_fast_stride == 0
532        && plan.a_row_stride == 0
533    {
534        return Some(TransposedScalarTileKind::LhsScalar);
535    }
536
537    None
538}
539
540fn compact_axis_order(dims: &[usize], strides: &[isize]) -> Option<AxisVec<usize>> {
541    if dims.len() != strides.len() {
542        return None;
543    }
544
545    let mut active = AxisVec::<usize>::new();
546    let mut inactive = AxisVec::<usize>::new();
547    for (axis, (&dim, &stride)) in dims.iter().zip(strides.iter()).enumerate() {
548        if stride < 0 {
549            return None;
550        }
551        if dim > 1 {
552            active.push(axis);
553        } else {
554            inactive.push(axis);
555        }
556    }
557
558    active.sort_by(|&lhs, &rhs| strides[lhs].cmp(&strides[rhs]).then_with(|| lhs.cmp(&rhs)));
559
560    let mut expected = 1isize;
561    for &axis in &active {
562        if strides[axis] != expected {
563            return None;
564        }
565        expected = expected.saturating_mul(dims[axis] as isize);
566    }
567
568    active.extend(inactive);
569    Some(active)
570}
571
572fn can_fuse_contiguous_range_axis(dim: usize, prev_stride: isize, next_stride: isize) -> bool {
573    dim <= 1 || (prev_stride == 0 && next_stride == 0) || next_stride == prev_stride * dim as isize
574}
575
576fn contiguous_mul_range_plan(
577    dims: &[usize],
578    dst_strides: &[isize],
579    a_strides: &[isize],
580    b_strides: &[isize],
581) -> Option<ContiguousMulRangePlan> {
582    let axis_order = compact_axis_order(dims, dst_strides)?;
583    if dims.is_empty() {
584        return Some(ContiguousMulRangePlan {
585            axis_order,
586            inner_len: 1,
587            inner_axis_count: 0,
588            row_len: 1,
589            outer_axis_start: 0,
590            fast_axis: 0,
591            a_fast_stride: 0,
592            b_fast_stride: 0,
593            a_row_stride: 0,
594            b_row_stride: 0,
595        });
596    }
597
598    let first_pos = axis_order
599        .iter()
600        .position(|&axis| dims[axis] > 1)
601        .unwrap_or(0);
602    let first_axis = axis_order[first_pos];
603    let mut inner_len = dims[first_axis].max(1);
604    let mut inner_axis_count = first_pos + 1;
605    let mut prev_axis = first_axis;
606
607    for &axis in axis_order.iter().skip(first_pos + 1) {
608        if can_fuse_contiguous_range_axis(
609            dims[prev_axis],
610            dst_strides[prev_axis],
611            dst_strides[axis],
612        ) && can_fuse_contiguous_range_axis(
613            dims[prev_axis],
614            a_strides[prev_axis],
615            a_strides[axis],
616        ) && can_fuse_contiguous_range_axis(
617            dims[prev_axis],
618            b_strides[prev_axis],
619            b_strides[axis],
620        ) {
621            inner_len = inner_len.checked_mul(dims[axis].max(1))?;
622            inner_axis_count += 1;
623            prev_axis = axis;
624        } else {
625            break;
626        }
627    }
628
629    let row = axis_order
630        .iter()
631        .enumerate()
632        .skip(inner_axis_count)
633        .find(|&(_, &axis)| dims[axis] > 1);
634    let (row_len, outer_axis_start, a_row_stride, b_row_stride) =
635        if let Some((row_pos, &row_axis)) = row {
636            (
637                dims[row_axis],
638                row_pos + 1,
639                a_strides[row_axis],
640                b_strides[row_axis],
641            )
642        } else {
643            (1, axis_order.len(), 0, 0)
644        };
645
646    Some(ContiguousMulRangePlan {
647        axis_order,
648        inner_len,
649        inner_axis_count,
650        row_len,
651        outer_axis_start,
652        fast_axis: first_axis,
653        a_fast_stride: a_strides[first_axis],
654        b_fast_stride: b_strides[first_axis],
655        a_row_stride,
656        b_row_stride,
657    })
658}
659
660struct ContiguousMulOuterCursor<'a> {
661    dims: &'a [usize],
662    a_strides: &'a [isize],
663    b_strides: &'a [isize],
664    axes: AxisVec<usize>,
665    coords: AxisVec<usize>,
666    a_offset: isize,
667    b_offset: isize,
668}
669
670impl<'a> ContiguousMulOuterCursor<'a> {
671    fn new(
672        dims: &'a [usize],
673        a_strides: &'a [isize],
674        b_strides: &'a [isize],
675        plan: &ContiguousMulRangePlan,
676        outer_group: usize,
677    ) -> Self {
678        let axes: AxisVec<usize> = plan
679            .axis_order
680            .iter()
681            .skip(plan.outer_axis_start)
682            .copied()
683            .collect();
684        let mut coords = AxisVec::<usize>::with_capacity(axes.len());
685        let mut rem = outer_group;
686        let mut a_offset = 0isize;
687        let mut b_offset = 0isize;
688
689        for &axis in &axes {
690            let dim = dims[axis].max(1);
691            let coord = rem % dim;
692            rem /= dim;
693            coords.push(coord);
694            a_offset += coord as isize * a_strides[axis];
695            b_offset += coord as isize * b_strides[axis];
696        }
697
698        Self {
699            dims,
700            a_strides,
701            b_strides,
702            axes,
703            coords,
704            a_offset,
705            b_offset,
706        }
707    }
708
709    fn advance(&mut self) {
710        for (i, &axis) in self.axes.iter().enumerate() {
711            let dim = self.dims[axis].max(1);
712            if dim <= 1 {
713                continue;
714            }
715
716            // Advance only when another position along this axis follows,
717            // and rewind from the last one, so offsets stay reachable.
718            if self.coords[i] + 1 < dim {
719                self.coords[i] += 1;
720                self.a_offset += self.a_strides[axis];
721                self.b_offset += self.b_strides[axis];
722                break;
723            }
724
725            let last = (dim - 1) as isize;
726            self.coords[i] = 0;
727            self.a_offset -= last * self.a_strides[axis];
728            self.b_offset -= last * self.b_strides[axis];
729        }
730    }
731}
732
733#[inline(always)]
734unsafe fn run_contiguous_mul_row_block<
735    O: MulOutput<D>,
736    D: Copy + 'static,
737    A: Copy + Mul<B, Output = D> + 'static,
738    B: Copy + 'static,
739>(
740    dst_ptr: *mut O::Slot,
741    a_ptr: *const A,
742    b_ptr: *const B,
743    plan: &ContiguousMulRangePlan,
744    base_index: usize,
745    total: usize,
746    base_a_offset: isize,
747    base_b_offset: isize,
748) {
749    let inner_len = plan.inner_len.max(1);
750    let row_len = plan.row_len.max(1);
751    #[cfg(feature = "parallel")]
752    let block_len = inner_len.saturating_mul(row_len);
753
754    #[cfg(feature = "parallel")]
755    if total.saturating_sub(base_index) >= block_len {
756        match transposed_scalar_tile_kind(plan) {
757            Some(TransposedScalarTileKind::RhsScalar) => {
758                if simd::try_mul_transposed_scalar_rhs_2d::<D, A, B>(
759                    dst_ptr.add(base_index).cast::<D>(),
760                    a_ptr.offset(base_a_offset),
761                    b_ptr.offset(base_b_offset),
762                    inner_len,
763                    row_len,
764                    plan.a_fast_stride,
765                    plan.a_row_stride,
766                ) {
767                    return;
768                }
769            }
770            Some(TransposedScalarTileKind::LhsScalar) => {
771                if simd::try_mul_transposed_scalar_lhs_2d::<D, A, B>(
772                    dst_ptr.add(base_index).cast::<D>(),
773                    a_ptr.offset(base_a_offset),
774                    b_ptr.offset(base_b_offset),
775                    inner_len,
776                    row_len,
777                    plan.b_fast_stride,
778                    plan.b_row_stride,
779                ) {
780                    return;
781                }
782            }
783            None => {}
784        }
785    }
786
787    let mut index = base_index;
788    let mut a_offset = base_a_offset;
789    let mut b_offset = base_b_offset;
790
791    for _ in 0..row_len {
792        if index >= total {
793            break;
794        }
795        let len = inner_len.min(total - index);
796        inner_loop_mul2::<O, D, A, B>(
797            dst_ptr.add(index),
798            1,
799            a_ptr.offset(a_offset),
800            plan.a_fast_stride,
801            b_ptr.offset(b_offset),
802            plan.b_fast_stride,
803            len,
804        );
805        index += inner_len;
806        a_offset += plan.a_row_stride;
807        b_offset += plan.b_row_stride;
808    }
809}
810
811#[cfg(feature = "parallel")]
812fn strided_offset_for_contiguous_linear_index(
813    dims: &[usize],
814    strides: &[isize],
815    axis_order: &[usize],
816    mut index: usize,
817) -> isize {
818    let mut offset = 0isize;
819    for &axis in axis_order {
820        let dim = dims[axis];
821        if dim == 0 {
822            return 0;
823        }
824        let coord = index % dim;
825        index /= dim;
826        offset += coord as isize * strides[axis];
827    }
828    offset
829}
830
831fn try_contiguous_range_mul<
832    O: MulOutput<D>,
833    D: Copy + MaybeSendSync + 'static,
834    A: Copy + MaybeSendSync + Mul<B, Output = D> + 'static,
835    B: Copy + MaybeSendSync + 'static,
836>(
837    dst_ptr: *mut O::Slot,
838    dims: &[usize],
839    dst_strides: &[isize],
840    a_ptr: *const A,
841    a_strides: &[isize],
842    b_ptr: *const B,
843    b_strides: &[isize],
844) -> bool {
845    // An overflowing element count falls back to the general path, which
846    // reports it as `OffsetOverflow`.
847    let Ok(total) = total_len(dims) else {
848        return false;
849    };
850    if total == 0 {
851        return true;
852    }
853    if total <= CONTIGUOUS_RANGE_MIN_LEN {
854        return false;
855    }
856
857    let Some(plan) = contiguous_mul_range_plan(dims, dst_strides, a_strides, b_strides) else {
858        return false;
859    };
860    if !plan.walks_inputs_without_blocking() {
861        return false;
862    }
863
864    let inner_len = plan.inner_len.max(1);
865    let row_len = plan.row_len.max(1);
866    let block_len = inner_len.saturating_mul(row_len).max(1);
867    let outer_groups = total.div_ceil(block_len);
868
869    #[cfg(feature = "parallel")]
870    {
871        let nthreads = crate::execution_policy::rayon_threads();
872        if nthreads > 1 {
873            use crate::threading::{parallel_for_each, SendPtr};
874
875            let dst = SendPtr(dst_ptr);
876            let a = SendPtr(a_ptr as *mut A);
877            let b = SendPtr(b_ptr as *mut B);
878
879            if outer_groups < nthreads {
880                let chunk_len = total.div_ceil(nthreads);
881                let nchunks = total.div_ceil(chunk_len);
882
883                parallel_for_each(0..nchunks, nthreads, &|chunks| {
884                    for chunk in chunks {
885                        let start = chunk * chunk_len;
886                        let end = (start + chunk_len).min(total);
887                        let mut index = start;
888
889                        while index < end {
890                            let in_inner = index % inner_len;
891                            let len = (inner_len - in_inner).min(end - index);
892                            let a_offset = strided_offset_for_contiguous_linear_index(
893                                dims,
894                                a_strides,
895                                &plan.axis_order,
896                                index,
897                            );
898                            let b_offset = strided_offset_for_contiguous_linear_index(
899                                dims,
900                                b_strides,
901                                &plan.axis_order,
902                                index,
903                            );
904
905                            unsafe {
906                                inner_loop_mul2::<O, D, A, B>(
907                                    dst.as_ptr().add(index),
908                                    1,
909                                    a.as_const().offset(a_offset),
910                                    plan.a_fast_stride,
911                                    b.as_const().offset(b_offset),
912                                    plan.b_fast_stride,
913                                    len,
914                                );
915                            }
916                            index += len;
917                        }
918                    }
919                });
920
921                return true;
922            }
923
924            let groups_per_chunk = outer_groups.div_ceil(nthreads);
925            let nchunks = outer_groups.div_ceil(groups_per_chunk);
926
927            parallel_for_each(0..nchunks, nthreads, &|chunks| {
928                for chunk in chunks {
929                    let group_start = chunk * groups_per_chunk;
930                    let group_end = (group_start + groups_per_chunk).min(outer_groups);
931                    let mut cursor = ContiguousMulOuterCursor::new(
932                        dims,
933                        a_strides,
934                        b_strides,
935                        &plan,
936                        group_start,
937                    );
938
939                    for group in group_start..group_end {
940                        let index = group * block_len;
941                        unsafe {
942                            run_contiguous_mul_row_block::<O, D, A, B>(
943                                dst.as_ptr(),
944                                a.as_const(),
945                                b.as_const(),
946                                &plan,
947                                index,
948                                total,
949                                cursor.a_offset,
950                                cursor.b_offset,
951                            );
952                        }
953                        cursor.advance();
954                    }
955                }
956            });
957
958            true
959        } else {
960            run_contiguous_range_mul_single_thread::<O, D, A, B>(
961                dst_ptr,
962                dims,
963                a_ptr,
964                a_strides,
965                b_ptr,
966                b_strides,
967                &plan,
968                total,
969                block_len,
970                outer_groups,
971            )
972        }
973    }
974
975    #[cfg(not(feature = "parallel"))]
976    {
977        run_contiguous_range_mul_single_thread::<O, D, A, B>(
978            dst_ptr,
979            dims,
980            a_ptr,
981            a_strides,
982            b_ptr,
983            b_strides,
984            &plan,
985            total,
986            block_len,
987            outer_groups,
988        )
989    }
990}
991
992fn run_contiguous_range_mul_single_thread<
993    O: MulOutput<D>,
994    D: Copy + MaybeSendSync + 'static,
995    A: Copy + MaybeSendSync + Mul<B, Output = D> + 'static,
996    B: Copy + MaybeSendSync + 'static,
997>(
998    dst_ptr: *mut O::Slot,
999    dims: &[usize],
1000    a_ptr: *const A,
1001    a_strides: &[isize],
1002    b_ptr: *const B,
1003    b_strides: &[isize],
1004    plan: &ContiguousMulRangePlan,
1005    total: usize,
1006    block_len: usize,
1007    outer_groups: usize,
1008) -> bool {
1009    let mut cursor = ContiguousMulOuterCursor::new(dims, a_strides, b_strides, plan, 0);
1010    for group in 0..outer_groups {
1011        let index = group * block_len;
1012        unsafe {
1013            run_contiguous_mul_row_block::<O, D, A, B>(
1014                dst_ptr,
1015                a_ptr,
1016                b_ptr,
1017                plan,
1018                index,
1019                total,
1020                cursor.a_offset,
1021                cursor.b_offset,
1022            );
1023        }
1024        cursor.advance();
1025    }
1026    true
1027}
1028
1029/// Ternary inner loop: `dest[i] = f(a[i], b[i], c[i])`.
1030#[inline(always)]
1031unsafe fn inner_loop_map3<
1032    D: Copy,
1033    A: Copy,
1034    B: Copy,
1035    C: Copy,
1036    OpA: ElementOp<A>,
1037    OpB: ElementOp<B>,
1038    OpC: ElementOp<C>,
1039>(
1040    dp: *mut D,
1041    ds: isize,
1042    ap: *const A,
1043    a_s: isize,
1044    bp: *const B,
1045    b_s: isize,
1046    cp: *const C,
1047    c_s: isize,
1048    len: usize,
1049    f: &impl Fn(A, B, C) -> D,
1050) {
1051    if ds == 1 && a_s == 1 && b_s == 1 && c_s == 1 {
1052        let src_a = std::slice::from_raw_parts(ap, len);
1053        let src_b = std::slice::from_raw_parts(bp, len);
1054        let src_c = std::slice::from_raw_parts(cp, len);
1055        let dst = std::slice::from_raw_parts_mut(dp, len);
1056        simd::dispatch_if_large(len, || {
1057            for (((d, &a), &b), &c) in dst.iter_mut().zip(src_a).zip(src_b).zip(src_c) {
1058                *d = f(OpA::apply(a), OpB::apply(b), OpC::apply(c));
1059            }
1060        });
1061    } else {
1062        let mut dp = dp;
1063        let mut ap = ap;
1064        let mut bp = bp;
1065        let mut cp = cp;
1066        for _ in 0..len {
1067            *dp = f(OpA::apply(*ap), OpB::apply(*bp), OpC::apply(*cp));
1068            dp = dp.wrapping_offset(ds);
1069            ap = ap.wrapping_offset(a_s);
1070            bp = bp.wrapping_offset(b_s);
1071            cp = cp.wrapping_offset(c_s);
1072        }
1073    }
1074}
1075
1076/// Quaternary inner loop: `dest[i] = f(a[i], b[i], c[i], e[i])`.
1077#[inline(always)]
1078unsafe fn inner_loop_map4<
1079    D: Copy,
1080    A: Copy,
1081    B: Copy,
1082    C: Copy,
1083    E: Copy,
1084    OpA: ElementOp<A>,
1085    OpB: ElementOp<B>,
1086    OpC: ElementOp<C>,
1087    OpE: ElementOp<E>,
1088>(
1089    dp: *mut D,
1090    ds: isize,
1091    ap: *const A,
1092    a_s: isize,
1093    bp: *const B,
1094    b_s: isize,
1095    cp: *const C,
1096    c_s: isize,
1097    ep: *const E,
1098    e_s: isize,
1099    len: usize,
1100    f: &impl Fn(A, B, C, E) -> D,
1101) {
1102    if ds == 1 && a_s == 1 && b_s == 1 && c_s == 1 && e_s == 1 {
1103        let src_a = std::slice::from_raw_parts(ap, len);
1104        let src_b = std::slice::from_raw_parts(bp, len);
1105        let src_c = std::slice::from_raw_parts(cp, len);
1106        let src_e = std::slice::from_raw_parts(ep, len);
1107        let dst = std::slice::from_raw_parts_mut(dp, len);
1108        simd::dispatch_if_large(len, || {
1109            for i in 0..len {
1110                dst[i] = f(
1111                    OpA::apply(src_a[i]),
1112                    OpB::apply(src_b[i]),
1113                    OpC::apply(src_c[i]),
1114                    OpE::apply(src_e[i]),
1115                );
1116            }
1117        });
1118    } else {
1119        let mut dp = dp;
1120        let mut ap = ap;
1121        let mut bp = bp;
1122        let mut cp = cp;
1123        let mut ep = ep;
1124        for _ in 0..len {
1125            *dp = f(
1126                OpA::apply(*ap),
1127                OpB::apply(*bp),
1128                OpC::apply(*cp),
1129                OpE::apply(*ep),
1130            );
1131            dp = dp.wrapping_offset(ds);
1132            ap = ap.wrapping_offset(a_s);
1133            bp = bp.wrapping_offset(b_s);
1134            cp = cp.wrapping_offset(c_s);
1135            ep = ep.wrapping_offset(e_s);
1136        }
1137    }
1138}
1139
1140/// Apply a function element-wise from source to destination.
1141///
1142/// The element operation `Op` is applied lazily when reading from `src`.
1143/// Source and destination may have different element types.
1144pub fn map_into<D: Copy + MaybeSendSync, A: Copy + MaybeSendSync, Op: ElementOp<A>>(
1145    dest: &mut StridedViewMut<D>,
1146    src: &StridedView<A, Op>,
1147    f: impl Fn(A) -> D + MaybeSync,
1148) -> Result<()> {
1149    map_parts_into::<D, A, Op, NoKernel>(
1150        dest.as_mut_ptr(),
1151        dest.dims(),
1152        dest.strides(),
1153        src.ptr(),
1154        src.dims(),
1155        src.strides(),
1156        f,
1157        &NoKernel,
1158    )
1159}
1160
1161pub(crate) fn map_into_validated<
1162    D: Copy + MaybeSendSync,
1163    A: Copy + MaybeSendSync,
1164    Op: ElementOp<A>,
1165>(
1166    dest: &mut StridedViewMut<D>,
1167    src: &StridedView<A, Op>,
1168    f: impl Fn(A) -> D + MaybeSync,
1169    validated: ValidatedDestinationLayout,
1170) -> Result<()> {
1171    ensure_same_shape(dest.dims(), src.dims())?;
1172    map_parts_into_validated::<D, A, Op, NoKernel>(
1173        dest.as_mut_ptr(),
1174        dest.dims(),
1175        dest.strides(),
1176        src.ptr(),
1177        src.strides(),
1178        f,
1179        &NoKernel,
1180        validated,
1181    )
1182}
1183
1184/// Apply an element-wise function into a potentially uninitialized destination, handing a
1185/// vectorized kernel the flat runs it can take.
1186///
1187/// `kernel` is offered a flat, unit-stride run whenever this crate has one to give: the whole
1188/// buffer when the two views share a contiguous layout and the operation runs sequentially, and
1189/// one block at a time on the parallel path. It must write every element of the run or write
1190/// nothing — returning `false` makes that run fall back to `scalar`, element by element. Runs whose
1191/// strides are not both 1 never reach it. It is called from whichever threads this crate's
1192/// execution policy selected, so it must be `Sync`.
1193///
1194/// The point of the entry point is that a caller which owns vector kernels — a SIMD math library,
1195/// say — reuses this crate's layout logic and this crate's threading decision instead of
1196/// rebuilding either. `scalar` and `kernel` must agree up to the caller's own error bound; this
1197/// crate only guarantees that every destination element is written exactly once.
1198///
1199/// # Examples
1200/// ```
1201/// use std::mem::MaybeUninit;
1202/// use strided_basic::{map_into_uninit_with_kernel, StridedView, StridedViewMut};
1203///
1204/// let src = vec![1.0f64, 2.0, 3.0];
1205/// let mut out = vec![MaybeUninit::<f64>::uninit(); 3];
1206/// let mut dest = StridedViewMut::new(&mut out, &[3], &[1], 0).unwrap();
1207/// let source = StridedView::new(&src, &[3], &[1], 0).unwrap();
1208/// let kernel = |src: &[f64], dst: &mut [MaybeUninit<f64>]| {
1209///     for (slot, value) in dst.iter_mut().zip(src) {
1210///         slot.write(value * 2.0);
1211///     }
1212///     true
1213/// };
1214/// map_into_uninit_with_kernel(&mut dest, &source, kernel, |x| x * 2.0).unwrap();
1215/// // SAFETY: the kernel wrote every element.
1216/// let values: Vec<f64> = out.into_iter().map(|x| unsafe { x.assume_init() }).collect();
1217/// assert_eq!(values, [2.0, 4.0, 6.0]);
1218/// ```
1219///
1220/// # Errors
1221/// Returns [`StridedError::ShapeMismatch`] when the shapes differ and
1222/// [`StridedError::NonInjectiveOutputLayout`] when the destination layout cannot be proven
1223/// injective.
1224pub fn map_into_uninit_with_kernel<T: Copy + MaybeSendSync>(
1225    dest: &mut StridedViewMut<MaybeUninit<T>>,
1226    src: &StridedView<T>,
1227    kernel: impl Fn(&[T], &mut [MaybeUninit<T>]) -> bool + Sync,
1228    scalar: impl Fn(T) -> T + MaybeSync,
1229) -> Result<()> {
1230    map_parts_into::<MaybeUninit<T>, T, Identity, _>(
1231        dest.as_mut_ptr(),
1232        dest.dims(),
1233        dest.strides(),
1234        src.ptr(),
1235        src.dims(),
1236        src.strides(),
1237        |value| MaybeUninit::new(scalar(value)),
1238        &SliceKernel(kernel),
1239    )
1240}
1241
1242pub(crate) fn map_raw_into<D: Copy + MaybeSendSync, A: Copy + MaybeSendSync, Op: ElementOp<A>>(
1243    dest: &mut crate::RawStridedMut<'_, D>,
1244    src: &crate::RawStridedRef<'_, A>,
1245    f: impl Fn(A) -> D + MaybeSync,
1246) -> Result<()> {
1247    map_parts_into::<D, A, Op, NoKernel>(
1248        dest.as_mut_ptr(),
1249        dest.dims(),
1250        dest.strides(),
1251        src.ptr(),
1252        src.dims(),
1253        src.strides(),
1254        f,
1255        &NoKernel,
1256    )
1257}
1258
1259pub(crate) fn map_raw_into_validated<
1260    D: Copy + MaybeSendSync,
1261    A: Copy + MaybeSendSync,
1262    Op: ElementOp<A>,
1263>(
1264    dest: &mut crate::RawStridedMut<'_, D>,
1265    src: &crate::RawStridedRef<'_, A>,
1266    f: impl Fn(A) -> D + MaybeSync,
1267    validated: ValidatedDestinationLayout,
1268) -> Result<()> {
1269    ensure_same_shape(dest.dims(), src.dims())?;
1270    map_parts_into_validated::<D, A, Op, NoKernel>(
1271        dest.as_mut_ptr(),
1272        dest.dims(),
1273        dest.strides(),
1274        src.ptr(),
1275        src.strides(),
1276        f,
1277        &NoKernel,
1278        validated,
1279    )
1280}
1281
1282#[allow(clippy::too_many_arguments)]
1283fn map_parts_into<
1284    D: Copy + MaybeSendSync,
1285    A: Copy + MaybeSendSync,
1286    Op: ElementOp<A>,
1287    K: MapKernel<A, D> + Sync,
1288>(
1289    dst_ptr: *mut D,
1290    dst_dims: &[usize],
1291    dst_strides: &[isize],
1292    src_ptr: *const A,
1293    src_dims: &[usize],
1294    src_strides: &[isize],
1295    f: impl Fn(A) -> D + MaybeSync,
1296    kernel: &K,
1297) -> Result<()> {
1298    ensure_same_shape(dst_dims, src_dims)?;
1299    let validated = validate_destination_layout(dst_dims, dst_strides)?;
1300    map_parts_into_validated::<D, A, Op, K>(
1301        dst_ptr,
1302        dst_dims,
1303        dst_strides,
1304        src_ptr,
1305        src_strides,
1306        f,
1307        kernel,
1308        validated,
1309    )
1310}
1311
1312#[allow(clippy::too_many_arguments)]
1313fn map_parts_into_validated<
1314    D: Copy + MaybeSendSync,
1315    A: Copy + MaybeSendSync,
1316    Op: ElementOp<A>,
1317    K: MapKernel<A, D> + Sync,
1318>(
1319    dst_ptr: *mut D,
1320    dst_dims: &[usize],
1321    dst_strides: &[isize],
1322    src_ptr: *const A,
1323    src_strides: &[isize],
1324    f: impl Fn(A) -> D + MaybeSync,
1325    kernel: &K,
1326    _validated: ValidatedDestinationLayout,
1327) -> Result<()> {
1328    if sequential_contiguous_layout(dst_dims, &[dst_strides, src_strides])?.is_some() {
1329        let len = total_len(dst_dims)?;
1330        // SAFETY: the contiguous-layout check proves that `dst` and `src` address `len` elements in
1331        // the same order starting at their base pointer, so each is exactly a flat run of `len`
1332        // elements.
1333        let dst = unsafe { std::slice::from_raw_parts_mut(dst_ptr, len) };
1334        let src = unsafe { std::slice::from_raw_parts(src_ptr, len) };
1335        if kernel.apply(src, dst) {
1336            return Ok(());
1337        }
1338        simd::dispatch_if_large(len, || {
1339            for i in 0..len {
1340                dst[i] = f(Op::apply(src[i]));
1341            }
1342        });
1343        return Ok(());
1344    }
1345
1346    let strides_list: [&[isize]; 2] = [dst_strides, src_strides];
1347    let elem_size = std::mem::size_of::<D>().max(std::mem::size_of::<A>());
1348    let total = total_len(dst_dims)?;
1349
1350    // Small tensor fast path: skip compute_order and compute_block_sizes
1351    let (fused_dims, ordered_strides, plan) = if total <= SMALL_TENSOR_THRESHOLD {
1352        build_plan_fused_small(dst_dims, &strides_list)
1353    } else {
1354        build_plan_fused(dst_dims, &strides_list, Some(0), elem_size)
1355    };
1356
1357    #[cfg(feature = "parallel")]
1358    {
1359        let total = total_len(&fused_dims)?;
1360        let nthreads = crate::execution_policy::rayon_threads();
1361        if total > MINTHREADLENGTH && nthreads > 1 {
1362            use crate::threading::SendPtr;
1363            let dst_send = SendPtr(dst_ptr);
1364            let src_send = SendPtr(src_ptr as *mut A);
1365
1366            let costs = compute_costs(&ordered_strides);
1367            let initial_offsets = vec![0isize; strides_list.len()];
1368            return mapreduce_threaded(
1369                &fused_dims,
1370                &plan.block,
1371                &ordered_strides,
1372                &initial_offsets,
1373                &costs,
1374                nthreads,
1375                0,
1376                1,
1377                &|dims, blocks, strides_list, offsets| {
1378                    for_each_inner_block_with_offsets(
1379                        dims,
1380                        blocks,
1381                        strides_list,
1382                        offsets,
1383                        |offsets, len, strides| {
1384                            let dp = unsafe { dst_send.as_ptr().offset(offsets[0]) };
1385                            let sp = unsafe { src_send.as_const().offset(offsets[1]) };
1386                            unsafe {
1387                                inner_loop_map1::<D, A, Op, K>(
1388                                    dp, strides[0], sp, strides[1], len, &f, kernel,
1389                                )
1390                            };
1391                            Ok(())
1392                        },
1393                    )
1394                },
1395            );
1396        }
1397    }
1398
1399    let initial_offsets = vec![0isize; ordered_strides.len()];
1400    for_each_inner_block_preordered(
1401        &fused_dims,
1402        &plan.block,
1403        &ordered_strides,
1404        &initial_offsets,
1405        |offsets, len, strides| {
1406            let dp = unsafe { dst_ptr.offset(offsets[0]) };
1407            let sp = unsafe { src_ptr.offset(offsets[1]) };
1408            unsafe {
1409                inner_loop_map1::<D, A, Op, K>(dp, strides[0], sp, strides[1], len, &f, kernel)
1410            };
1411            Ok(())
1412        },
1413    )
1414}
1415
1416/// Binary element-wise operation: `dest[i] = f(a[i], b[i])`.
1417///
1418/// Source operands `a` and `b` may have different element types from each other
1419/// and from `dest`. The closure `f` handles per-element type conversion.
1420pub fn zip_map2_into<
1421    D: Copy + MaybeSendSync,
1422    A: Copy + MaybeSendSync,
1423    B: Copy + MaybeSendSync,
1424    OpA: ElementOp<A>,
1425    OpB: ElementOp<B>,
1426>(
1427    dest: &mut StridedViewMut<D>,
1428    a: &StridedView<A, OpA>,
1429    b: &StridedView<B, OpB>,
1430    f: impl Fn(A, B) -> D + MaybeSync,
1431) -> Result<()> {
1432    zip_map2_parts_into::<D, A, B, OpA, OpB>(
1433        dest.as_mut_ptr(),
1434        dest.dims(),
1435        dest.strides(),
1436        a.ptr(),
1437        a.dims(),
1438        a.strides(),
1439        b.ptr(),
1440        b.dims(),
1441        b.strides(),
1442        f,
1443    )
1444}
1445
1446pub(crate) fn zip_map2_into_validated<
1447    D: Copy + MaybeSendSync,
1448    A: Copy + MaybeSendSync,
1449    B: Copy + MaybeSendSync,
1450    OpA: ElementOp<A>,
1451    OpB: ElementOp<B>,
1452>(
1453    dest: &mut StridedViewMut<D>,
1454    a: &StridedView<A, OpA>,
1455    b: &StridedView<B, OpB>,
1456    f: impl Fn(A, B) -> D + MaybeSync,
1457    validated: ValidatedDestinationLayout,
1458) -> Result<()> {
1459    zip_map2_parts_into_validated::<D, A, B, OpA, OpB>(
1460        dest.as_mut_ptr(),
1461        dest.dims(),
1462        dest.strides(),
1463        a.ptr(),
1464        a.strides(),
1465        b.ptr(),
1466        b.strides(),
1467        f,
1468        validated,
1469    )
1470}
1471
1472/// Runtime comparison selected once before entering the element loop.
1473#[non_exhaustive]
1474#[derive(Clone, Copy, Debug, Eq, PartialEq)]
1475pub enum CompareOp {
1476    Eq,
1477    Lt,
1478    Le,
1479    Gt,
1480    Ge,
1481}
1482
1483/// Compare two ordered views elementwise into a Boolean destination.
1484///
1485/// Unlike embedding a runtime operation match inside a [`zip_map2_into`]
1486/// closure, this entry point selects a fixed comparison before traversal.
1487///
1488/// # Errors
1489///
1490/// Returns [`StridedError::ShapeMismatch`] when the source and destination
1491/// shapes differ, or [`StridedError::NonInjectiveOutputLayout`] when distinct
1492/// logical destination elements may overlap.
1493pub fn compare_into<T, OpA, OpB>(
1494    dest: &mut StridedViewMut<bool>,
1495    a: &StridedView<T, OpA>,
1496    b: &StridedView<T, OpB>,
1497    op: CompareOp,
1498) -> Result<()>
1499where
1500    T: Copy + MaybeSendSync + PartialOrd,
1501    OpA: ElementOp<T>,
1502    OpB: ElementOp<T>,
1503{
1504    match op {
1505        CompareOp::Eq => zip_map2_into(dest, a, b, |lhs, rhs| lhs == rhs),
1506        CompareOp::Lt => zip_map2_into(dest, a, b, |lhs, rhs| lhs < rhs),
1507        CompareOp::Le => zip_map2_into(dest, a, b, |lhs, rhs| lhs <= rhs),
1508        CompareOp::Gt => zip_map2_into(dest, a, b, |lhs, rhs| lhs > rhs),
1509        CompareOp::Ge => zip_map2_into(dest, a, b, |lhs, rhs| lhs >= rhs),
1510    }
1511}
1512
1513/// Compare two views into a fully overwritten uninitialized Boolean output.
1514///
1515/// Dtype-independent shape, destination-injectivity, and reachable-byte
1516/// overlap validation completes before the first write. Safe Rust borrows
1517/// already prevent input/output aliasing; the explicit overlap check preserves
1518/// the contract for views produced through unsafe constructors.
1519///
1520/// `Ok(())` means every logical destination element is initialized. An error
1521/// occurs before writes. A panic during replay may leave a partially initialized
1522/// destination, which remains safe to drop as `MaybeUninit<bool>`.
1523///
1524/// # Errors
1525///
1526/// Returns [`StridedError::ShapeMismatch`] for unequal shapes,
1527/// [`StridedError::NonInjectiveOutputLayout`] for an overlapping output layout,
1528/// [`StridedError::OverlappingInputOutput`] for aliased storage, or
1529/// [`StridedError::OffsetOverflow`] when a reachable byte range is not
1530/// representable.
1531pub fn compare_into_uninit<T, OpA, OpB>(
1532    dest: &mut StridedViewMut<MaybeUninit<bool>>,
1533    a: &StridedView<T, OpA>,
1534    b: &StridedView<T, OpB>,
1535    op: CompareOp,
1536) -> Result<()>
1537where
1538    T: Copy + MaybeSendSync + PartialOrd,
1539    OpA: ElementOp<T>,
1540    OpB: ElementOp<T>,
1541{
1542    ensure_same_shape(dest.dims(), a.dims())?;
1543    ensure_same_shape(dest.dims(), b.dims())?;
1544    let validated = validate_destination_layout(dest.dims(), dest.strides())?;
1545    validate_typed_no_overlap(dest, a, 0)?;
1546    validate_typed_no_overlap(dest, b, 1)?;
1547    match op {
1548        CompareOp::Eq => zip_map2_into_validated(
1549            dest,
1550            a,
1551            b,
1552            |lhs, rhs| MaybeUninit::new(lhs == rhs),
1553            validated,
1554        ),
1555        CompareOp::Lt => zip_map2_into_validated(
1556            dest,
1557            a,
1558            b,
1559            |lhs, rhs| MaybeUninit::new(lhs < rhs),
1560            validated,
1561        ),
1562        CompareOp::Le => zip_map2_into_validated(
1563            dest,
1564            a,
1565            b,
1566            |lhs, rhs| MaybeUninit::new(lhs <= rhs),
1567            validated,
1568        ),
1569        CompareOp::Gt => zip_map2_into_validated(
1570            dest,
1571            a,
1572            b,
1573            |lhs, rhs| MaybeUninit::new(lhs > rhs),
1574            validated,
1575        ),
1576        CompareOp::Ge => zip_map2_into_validated(
1577            dest,
1578            a,
1579            b,
1580            |lhs, rhs| MaybeUninit::new(lhs >= rhs),
1581            validated,
1582        ),
1583    }
1584}
1585
1586pub(crate) fn zip_map2_raw_into_validated<
1587    D: Copy + MaybeSendSync,
1588    A: Copy + MaybeSendSync,
1589    B: Copy + MaybeSendSync,
1590    OpA: ElementOp<A>,
1591    OpB: ElementOp<B>,
1592>(
1593    dest: &mut crate::RawStridedMut<'_, D>,
1594    a: &crate::RawStridedRef<'_, A>,
1595    b: &crate::RawStridedRef<'_, B>,
1596    f: impl Fn(A, B) -> D + MaybeSync,
1597    validated: ValidatedDestinationLayout,
1598) -> Result<()> {
1599    ensure_same_shape(dest.dims(), a.dims())?;
1600    ensure_same_shape(dest.dims(), b.dims())?;
1601    zip_map2_parts_into_validated::<D, A, B, OpA, OpB>(
1602        dest.as_mut_ptr(),
1603        dest.dims(),
1604        dest.strides(),
1605        a.ptr(),
1606        a.strides(),
1607        b.ptr(),
1608        b.strides(),
1609        f,
1610        validated,
1611    )
1612}
1613
1614#[allow(clippy::too_many_arguments)]
1615fn zip_map2_parts_into<
1616    D: Copy + MaybeSendSync,
1617    A: Copy + MaybeSendSync,
1618    B: Copy + MaybeSendSync,
1619    OpA: ElementOp<A>,
1620    OpB: ElementOp<B>,
1621>(
1622    dst_ptr: *mut D,
1623    dst_dims: &[usize],
1624    dst_strides: &[isize],
1625    a_ptr: *const A,
1626    a_dims: &[usize],
1627    a_strides: &[isize],
1628    b_ptr: *const B,
1629    b_dims: &[usize],
1630    b_strides: &[isize],
1631    f: impl Fn(A, B) -> D + MaybeSync,
1632) -> Result<()> {
1633    ensure_same_shape(dst_dims, a_dims)?;
1634    ensure_same_shape(dst_dims, b_dims)?;
1635    let validated = validate_destination_layout(dst_dims, dst_strides)?;
1636    zip_map2_parts_into_validated::<D, A, B, OpA, OpB>(
1637        dst_ptr,
1638        dst_dims,
1639        dst_strides,
1640        a_ptr,
1641        a_strides,
1642        b_ptr,
1643        b_strides,
1644        f,
1645        validated,
1646    )
1647}
1648
1649#[allow(clippy::too_many_arguments)]
1650fn zip_map2_parts_into_validated<
1651    D: Copy + MaybeSendSync,
1652    A: Copy + MaybeSendSync,
1653    B: Copy + MaybeSendSync,
1654    OpA: ElementOp<A>,
1655    OpB: ElementOp<B>,
1656>(
1657    dst_ptr: *mut D,
1658    dst_dims: &[usize],
1659    dst_strides: &[isize],
1660    a_ptr: *const A,
1661    a_strides: &[isize],
1662    b_ptr: *const B,
1663    b_strides: &[isize],
1664    f: impl Fn(A, B) -> D + MaybeSync,
1665    _validated: ValidatedDestinationLayout,
1666) -> Result<()> {
1667    if sequential_contiguous_layout(dst_dims, &[dst_strides, a_strides, b_strides])?.is_some() {
1668        let len = total_len(dst_dims)?;
1669        let dst = unsafe { std::slice::from_raw_parts_mut(dst_ptr, len) };
1670        let sa = unsafe { std::slice::from_raw_parts(a_ptr, len) };
1671        let sb = unsafe { std::slice::from_raw_parts(b_ptr, len) };
1672        simd::dispatch_if_large(len, || {
1673            for i in 0..len {
1674                dst[i] = f(OpA::apply(sa[i]), OpB::apply(sb[i]));
1675            }
1676        });
1677        return Ok(());
1678    }
1679
1680    let strides_list: [&[isize]; 3] = [dst_strides, a_strides, b_strides];
1681    let elem_size = std::mem::size_of::<D>()
1682        .max(std::mem::size_of::<A>())
1683        .max(std::mem::size_of::<B>());
1684    let total = total_len(dst_dims)?;
1685
1686    // Small tensor fast path: skip compute_order and compute_block_sizes
1687    let (fused_dims, ordered_strides, plan) = if total <= SMALL_TENSOR_THRESHOLD {
1688        build_plan_fused_small(dst_dims, &strides_list)
1689    } else {
1690        build_plan_fused(dst_dims, &strides_list, Some(0), elem_size)
1691    };
1692
1693    #[cfg(feature = "parallel")]
1694    {
1695        let total = total_len(&fused_dims)?;
1696        let nthreads = crate::execution_policy::rayon_threads();
1697        if total > MINTHREADLENGTH && nthreads > 1 {
1698            use crate::threading::SendPtr;
1699            let dst_send = SendPtr(dst_ptr);
1700            let a_send = SendPtr(a_ptr as *mut A);
1701            let b_send = SendPtr(b_ptr as *mut B);
1702
1703            let costs = compute_costs(&ordered_strides);
1704            let initial_offsets = vec![0isize; strides_list.len()];
1705            return mapreduce_threaded(
1706                &fused_dims,
1707                &plan.block,
1708                &ordered_strides,
1709                &initial_offsets,
1710                &costs,
1711                nthreads,
1712                0,
1713                1,
1714                &|dims, blocks, strides_list, offsets| {
1715                    for_each_inner_block_with_offsets(
1716                        dims,
1717                        blocks,
1718                        strides_list,
1719                        offsets,
1720                        |offsets, len, strides| {
1721                            let dp = unsafe { dst_send.as_ptr().offset(offsets[0]) };
1722                            let ap = unsafe { a_send.as_const().offset(offsets[1]) };
1723                            let bp = unsafe { b_send.as_const().offset(offsets[2]) };
1724                            unsafe {
1725                                inner_loop_map2::<D, A, B, OpA, OpB>(
1726                                    dp, strides[0], ap, strides[1], bp, strides[2], len, &f,
1727                                )
1728                            };
1729                            Ok(())
1730                        },
1731                    )
1732                },
1733            );
1734        }
1735    }
1736
1737    let initial_offsets = vec![0isize; ordered_strides.len()];
1738    for_each_inner_block_preordered(
1739        &fused_dims,
1740        &plan.block,
1741        &ordered_strides,
1742        &initial_offsets,
1743        |offsets, len, strides| {
1744            let dp = unsafe { dst_ptr.offset(offsets[0]) };
1745            let ap = unsafe { a_ptr.offset(offsets[1]) };
1746            let bp = unsafe { b_ptr.offset(offsets[2]) };
1747            unsafe {
1748                inner_loop_map2::<D, A, B, OpA, OpB>(
1749                    dp, strides[0], ap, strides[1], bp, strides[2], len, &f,
1750                )
1751            };
1752            Ok(())
1753        },
1754    )
1755}
1756
1757fn mul_identity_into_raw<
1758    O: MulOutput<D>,
1759    D: Copy + MaybeSendSync + 'static,
1760    A: Copy + MaybeSendSync + Mul<B, Output = D> + 'static,
1761    B: Copy + MaybeSendSync + 'static,
1762>(
1763    dst_ptr: *mut O::Slot,
1764    dst_dims: &[usize],
1765    dst_strides: &[isize],
1766    a_ptr: *const A,
1767    a_strides: &[isize],
1768    b_ptr: *const B,
1769    b_strides: &[isize],
1770    _validated: ValidatedDestinationLayout,
1771) -> Result<()> {
1772    debug_assert_eq!(dst_dims.len(), a_strides.len());
1773    debug_assert_eq!(dst_dims.len(), b_strides.len());
1774
1775    if sequential_contiguous_layout(dst_dims, &[dst_strides, a_strides, b_strides])?.is_some() {
1776        let len = total_len(dst_dims)?;
1777        let sa = unsafe { std::slice::from_raw_parts(a_ptr, len) };
1778        let sb = unsafe { std::slice::from_raw_parts(b_ptr, len) };
1779        if unsafe { O::try_contiguous(dst_ptr, len, sa, sb) } {
1780            return Ok(());
1781        }
1782        for i in 0..len {
1783            unsafe { O::write(dst_ptr.add(i), multiply_value(sa[i], sb[i])) };
1784        }
1785        return Ok(());
1786    }
1787
1788    let strides_list: [&[isize]; 3] = [dst_strides, a_strides, b_strides];
1789    let elem_size = std::mem::size_of::<D>()
1790        .max(std::mem::size_of::<A>())
1791        .max(std::mem::size_of::<B>());
1792    let total = total_len(dst_dims)?;
1793
1794    if try_contiguous_range_mul::<O, D, A, B>(
1795        dst_ptr,
1796        dst_dims,
1797        dst_strides,
1798        a_ptr,
1799        a_strides,
1800        b_ptr,
1801        b_strides,
1802    ) {
1803        return Ok(());
1804    }
1805
1806    let (fused_dims, ordered_strides, plan) = if total <= SMALL_TENSOR_THRESHOLD {
1807        build_plan_fused_small(dst_dims, &strides_list)
1808    } else {
1809        build_plan_fused(dst_dims, &strides_list, Some(0), elem_size)
1810    };
1811
1812    #[cfg(feature = "parallel")]
1813    {
1814        let total = total_len(&fused_dims)?;
1815        let nthreads = crate::execution_policy::rayon_threads();
1816        if total > MINTHREADLENGTH && nthreads > 1 {
1817            use crate::threading::SendPtr;
1818            let dst_send = SendPtr(dst_ptr);
1819            let a_send = SendPtr(a_ptr as *mut A);
1820            let b_send = SendPtr(b_ptr as *mut B);
1821
1822            let costs = compute_costs(&ordered_strides);
1823            let initial_offsets = vec![0isize; strides_list.len()];
1824            return mapreduce_threaded(
1825                &fused_dims,
1826                &plan.block,
1827                &ordered_strides,
1828                &initial_offsets,
1829                &costs,
1830                nthreads,
1831                0,
1832                1,
1833                &|dims, blocks, strides_list, offsets| {
1834                    for_each_inner_block_with_offsets(
1835                        dims,
1836                        blocks,
1837                        strides_list,
1838                        offsets,
1839                        |offsets, len, strides| {
1840                            let dp = unsafe { dst_send.as_ptr().offset(offsets[0]) };
1841                            let ap = unsafe { a_send.as_const().offset(offsets[1]) };
1842                            let bp = unsafe { b_send.as_const().offset(offsets[2]) };
1843                            unsafe {
1844                                inner_loop_mul2::<O, D, A, B>(
1845                                    dp, strides[0], ap, strides[1], bp, strides[2], len,
1846                                )
1847                            };
1848                            Ok(())
1849                        },
1850                    )
1851                },
1852            );
1853        }
1854    }
1855
1856    let initial_offsets = vec![0isize; ordered_strides.len()];
1857    for_each_inner_block_preordered(
1858        &fused_dims,
1859        &plan.block,
1860        &ordered_strides,
1861        &initial_offsets,
1862        |offsets, len, strides| {
1863            let dp = unsafe { dst_ptr.offset(offsets[0]) };
1864            let ap = unsafe { a_ptr.offset(offsets[1]) };
1865            let bp = unsafe { b_ptr.offset(offsets[2]) };
1866            unsafe {
1867                inner_loop_mul2::<O, D, A, B>(dp, strides[0], ap, strides[1], bp, strides[2], len)
1868            };
1869            Ok(())
1870        },
1871    )
1872}
1873
1874/// Element-wise multiplication: `dest[i] = a[i] * b[i]`.
1875///
1876/// All views must have the same shape. Broadcast operands should be represented
1877/// as stride-0 views before calling this function.
1878pub fn mul_into<
1879    D: Copy + MaybeSendSync + 'static,
1880    A: Copy + Mul<B, Output = D> + MaybeSendSync + 'static,
1881    B: Copy + MaybeSendSync + 'static,
1882    OpA: ElementOp<A>,
1883    OpB: ElementOp<B>,
1884>(
1885    dest: &mut StridedViewMut<D>,
1886    a: &StridedView<A, OpA>,
1887    b: &StridedView<B, OpB>,
1888) -> Result<()> {
1889    ensure_same_shape(dest.dims(), a.dims())?;
1890    ensure_same_shape(dest.dims(), b.dims())?;
1891    let validated = validate_destination_layout(dest.dims(), dest.strides())?;
1892
1893    if OpA::IS_IDENTITY && OpB::IS_IDENTITY {
1894        return mul_identity_into_raw::<InitializedOutput, D, A, B>(
1895            dest.as_mut_ptr(),
1896            dest.dims(),
1897            dest.strides(),
1898            a.ptr(),
1899            a.strides(),
1900            b.ptr(),
1901            b.strides(),
1902            validated,
1903        );
1904    }
1905
1906    zip_map2_into_validated(dest, a, b, multiply_value, validated)
1907}
1908
1909/// Multiply two views into a fully overwritten uninitialized output.
1910///
1911/// Shape, destination-injectivity, and reachable-byte overlap validation
1912/// completes before the first write. Safe Rust borrows already prevent
1913/// input/output aliasing; the explicit overlap check preserves the contract for
1914/// views produced through unsafe constructors.
1915///
1916/// `Ok(())` means every logical destination element is initialized. An error
1917/// occurs before writes. A panic during replay may leave a partially initialized
1918/// destination, which remains safe to drop as `MaybeUninit<D>`.
1919///
1920/// # Errors
1921///
1922/// Returns [`StridedError::ShapeMismatch`] for unequal shapes,
1923/// [`StridedError::NonInjectiveOutputLayout`] for an overlapping output layout,
1924/// [`StridedError::OverlappingInputOutput`] for aliased storage, or
1925/// [`StridedError::OffsetOverflow`] when a reachable byte range is not
1926/// representable.
1927pub fn mul_into_uninit<
1928    D: Copy + MaybeSendSync + 'static,
1929    A: Copy + Mul<B, Output = D> + MaybeSendSync + 'static,
1930    B: Copy + MaybeSendSync + 'static,
1931    OpA: ElementOp<A>,
1932    OpB: ElementOp<B>,
1933>(
1934    dest: &mut StridedViewMut<MaybeUninit<D>>,
1935    a: &StridedView<A, OpA>,
1936    b: &StridedView<B, OpB>,
1937) -> Result<()> {
1938    ensure_same_shape(dest.dims(), a.dims())?;
1939    ensure_same_shape(dest.dims(), b.dims())?;
1940    let validated = validate_destination_layout(dest.dims(), dest.strides())?;
1941    validate_typed_no_overlap(dest, a, 0)?;
1942    validate_typed_no_overlap(dest, b, 1)?;
1943
1944    if OpA::IS_IDENTITY && OpB::IS_IDENTITY {
1945        return mul_identity_into_raw::<UninitializedOutput, D, A, B>(
1946            dest.as_mut_ptr(),
1947            dest.dims(),
1948            dest.strides(),
1949            a.ptr(),
1950            a.strides(),
1951            b.ptr(),
1952            b.strides(),
1953            validated,
1954        );
1955    }
1956
1957    zip_map2_into_validated(
1958        dest,
1959        a,
1960        b,
1961        |lhs, rhs| MaybeUninit::new(multiply_value(lhs, rhs)),
1962        validated,
1963    )
1964}
1965
1966fn broadcast_strides_for_axes(
1967    source_dims: &[usize],
1968    source_strides: &[isize],
1969    target_dims: &[usize],
1970    axes: &[usize],
1971) -> Result<AxisVec<isize>> {
1972    if source_dims.len() != axes.len() {
1973        return Err(StridedError::RankMismatch(source_dims.len(), axes.len()));
1974    }
1975    debug_assert_eq!(source_dims.len(), source_strides.len());
1976
1977    let mut seen = AxisVec::<bool>::new();
1978    seen.resize(target_dims.len(), false);
1979    let mut strides = AxisVec::<isize>::new();
1980    strides.resize(target_dims.len(), 0);
1981    for (src_axis, &dst_axis) in axes.iter().enumerate() {
1982        if dst_axis >= target_dims.len() {
1983            return Err(StridedError::InvalidAxis {
1984                axis: dst_axis,
1985                rank: target_dims.len(),
1986            });
1987        }
1988        if seen[dst_axis] {
1989            return Err(StridedError::InvalidAxis {
1990                axis: dst_axis,
1991                rank: target_dims.len(),
1992            });
1993        }
1994        seen[dst_axis] = true;
1995
1996        let source_dim = source_dims[src_axis];
1997        let target_dim = target_dims[dst_axis];
1998        if source_dim != target_dim && source_dim != 1 {
1999            return Err(StridedError::ShapeMismatch(
2000                source_dims.to_vec(),
2001                target_dims.to_vec(),
2002            ));
2003        }
2004        if source_dim == target_dim {
2005            strides[dst_axis] = source_strides[src_axis];
2006        }
2007    }
2008
2009    Ok(strides)
2010}
2011
2012/// A view of `view`'s data with stride-0 broadcast axes, owning the metadata.
2013///
2014/// The metadata is computed for this call only and the caller's `dest` must
2015/// stay mutable while the returned view is used, so the view cannot borrow it.
2016fn broadcast_view_with_strides<'a, T, Op: ElementOp<T>>(
2017    view: &StridedView<'a, T, Op>,
2018    target_dims: &[usize],
2019    strides: &[isize],
2020) -> StridedView<'a, T, Op> {
2021    unsafe {
2022        StridedView::new_owned_unchecked(
2023            view.data(),
2024            target_dims.to_vec(),
2025            strides.to_vec(),
2026            view.offset(),
2027        )
2028    }
2029}
2030
2031/// Broadcasted element-wise multiplication: `dest[i] = a[i] * b[i]`.
2032///
2033/// `a_axes` and `b_axes` map each source axis to an axis of `dest`. Output axes
2034/// not referenced by a source operand are treated as stride-0 broadcast axes.
2035pub fn broadcast_mul_into<
2036    D: Copy + MaybeSendSync + 'static,
2037    A: Copy + Mul<B, Output = D> + MaybeSendSync + 'static,
2038    B: Copy + MaybeSendSync + 'static,
2039    OpA: ElementOp<A>,
2040    OpB: ElementOp<B>,
2041>(
2042    dest: &mut StridedViewMut<D>,
2043    a: &StridedView<A, OpA>,
2044    a_axes: &[usize],
2045    b: &StridedView<B, OpB>,
2046    b_axes: &[usize],
2047) -> Result<()> {
2048    let validated = validate_destination_layout(dest.dims(), dest.strides())?;
2049    let a_strides = broadcast_strides_for_axes(a.dims(), a.strides(), dest.dims(), a_axes)?;
2050    let b_strides = broadcast_strides_for_axes(b.dims(), b.strides(), dest.dims(), b_axes)?;
2051
2052    if OpA::IS_IDENTITY && OpB::IS_IDENTITY {
2053        return mul_identity_into_raw::<InitializedOutput, D, A, B>(
2054            dest.as_mut_ptr(),
2055            dest.dims(),
2056            dest.strides(),
2057            a.ptr(),
2058            &a_strides,
2059            b.ptr(),
2060            &b_strides,
2061            validated,
2062        );
2063    }
2064
2065    let a = broadcast_view_with_strides(a, dest.dims(), &a_strides);
2066    let b = broadcast_view_with_strides(b, dest.dims(), &b_strides);
2067    zip_map2_into_validated(dest, &a, &b, multiply_value, validated)
2068}
2069
2070/// Broadcast and multiply into a fully overwritten uninitialized output.
2071///
2072/// Axis mappings, shapes, destination injectivity, and reachable-byte overlap
2073/// are validated before the first write. Safe Rust borrows already prevent
2074/// input/output aliasing; the explicit overlap check preserves the contract for
2075/// views produced through unsafe constructors.
2076///
2077/// `Ok(())` means every logical destination element is initialized. An error
2078/// occurs before writes. A panic during replay may leave a partially initialized
2079/// destination, which remains safe to drop as `MaybeUninit<D>`.
2080///
2081/// # Errors
2082///
2083/// Returns a typed rank, axis, or shape error for an invalid broadcast mapping,
2084/// [`StridedError::NonInjectiveOutputLayout`] for an overlapping output layout,
2085/// [`StridedError::OverlappingInputOutput`] for aliased storage, or
2086/// [`StridedError::OffsetOverflow`] when a reachable byte range is not
2087/// representable.
2088pub fn broadcast_mul_into_uninit<
2089    D: Copy + MaybeSendSync + 'static,
2090    A: Copy + Mul<B, Output = D> + MaybeSendSync + 'static,
2091    B: Copy + MaybeSendSync + 'static,
2092    OpA: ElementOp<A>,
2093    OpB: ElementOp<B>,
2094>(
2095    dest: &mut StridedViewMut<MaybeUninit<D>>,
2096    a: &StridedView<A, OpA>,
2097    a_axes: &[usize],
2098    b: &StridedView<B, OpB>,
2099    b_axes: &[usize],
2100) -> Result<()> {
2101    let validated = validate_destination_layout(dest.dims(), dest.strides())?;
2102    let a_strides = broadcast_strides_for_axes(a.dims(), a.strides(), dest.dims(), a_axes)?;
2103    let b_strides = broadcast_strides_for_axes(b.dims(), b.strides(), dest.dims(), b_axes)?;
2104    validate_typed_no_overlap(dest, a, 0)?;
2105    validate_typed_no_overlap(dest, b, 1)?;
2106
2107    if OpA::IS_IDENTITY && OpB::IS_IDENTITY {
2108        return mul_identity_into_raw::<UninitializedOutput, D, A, B>(
2109            dest.as_mut_ptr(),
2110            dest.dims(),
2111            dest.strides(),
2112            a.ptr(),
2113            &a_strides,
2114            b.ptr(),
2115            &b_strides,
2116            validated,
2117        );
2118    }
2119
2120    let a = broadcast_view_with_strides(a, dest.dims(), &a_strides);
2121    let b = broadcast_view_with_strides(b, dest.dims(), &b_strides);
2122    zip_map2_into_validated(
2123        dest,
2124        &a,
2125        &b,
2126        |lhs, rhs| MaybeUninit::new(multiply_value(lhs, rhs)),
2127        validated,
2128    )
2129}
2130
2131/// Ternary element-wise operation: `dest[i] = f(a[i], b[i], c[i])`.
2132pub fn zip_map3_into<
2133    D: Copy + MaybeSendSync,
2134    A: Copy + MaybeSendSync,
2135    B: Copy + MaybeSendSync,
2136    C: Copy + MaybeSendSync,
2137    OpA: ElementOp<A>,
2138    OpB: ElementOp<B>,
2139    OpC: ElementOp<C>,
2140>(
2141    dest: &mut StridedViewMut<D>,
2142    a: &StridedView<A, OpA>,
2143    b: &StridedView<B, OpB>,
2144    c: &StridedView<C, OpC>,
2145    f: impl Fn(A, B, C) -> D + MaybeSync,
2146) -> Result<()> {
2147    ensure_same_shape(dest.dims(), a.dims())?;
2148    ensure_same_shape(dest.dims(), b.dims())?;
2149    ensure_same_shape(dest.dims(), c.dims())?;
2150    let validated = validate_destination_layout(dest.dims(), dest.strides())?;
2151    zip_map3_into_validated(dest, a, b, c, f, validated)
2152}
2153
2154pub(crate) fn zip_map3_into_validated<
2155    D: Copy + MaybeSendSync,
2156    A: Copy + MaybeSendSync,
2157    B: Copy + MaybeSendSync,
2158    C: Copy + MaybeSendSync,
2159    OpA: ElementOp<A>,
2160    OpB: ElementOp<B>,
2161    OpC: ElementOp<C>,
2162>(
2163    dest: &mut StridedViewMut<D>,
2164    a: &StridedView<A, OpA>,
2165    b: &StridedView<B, OpB>,
2166    c: &StridedView<C, OpC>,
2167    f: impl Fn(A, B, C) -> D + MaybeSync,
2168    _validated: ValidatedDestinationLayout,
2169) -> Result<()> {
2170    ensure_same_shape(dest.dims(), a.dims())?;
2171    ensure_same_shape(dest.dims(), b.dims())?;
2172    ensure_same_shape(dest.dims(), c.dims())?;
2173    let dst_ptr = dest.as_mut_ptr();
2174    let a_ptr = a.ptr();
2175    let b_ptr = b.ptr();
2176    let c_ptr = c.ptr();
2177
2178    let dst_dims = dest.dims();
2179    let dst_strides = dest.strides();
2180
2181    if sequential_contiguous_layout(
2182        dst_dims,
2183        &[dst_strides, a.strides(), b.strides(), c.strides()],
2184    )?
2185    .is_some()
2186    {
2187        let len = total_len(dst_dims)?;
2188        let dst = unsafe { std::slice::from_raw_parts_mut(dst_ptr, len) };
2189        let sa = unsafe { std::slice::from_raw_parts(a_ptr, len) };
2190        let sb = unsafe { std::slice::from_raw_parts(b_ptr, len) };
2191        let sc = unsafe { std::slice::from_raw_parts(c_ptr, len) };
2192        simd::dispatch_if_large(len, || {
2193            for (((d, &a), &b), &c) in dst.iter_mut().zip(sa).zip(sb).zip(sc) {
2194                *d = f(OpA::apply(a), OpB::apply(b), OpC::apply(c));
2195            }
2196        });
2197        return Ok(());
2198    }
2199
2200    let strides_list: [&[isize]; 4] = [dst_strides, a.strides(), b.strides(), c.strides()];
2201    let elem_size = std::mem::size_of::<D>()
2202        .max(std::mem::size_of::<A>())
2203        .max(std::mem::size_of::<B>())
2204        .max(std::mem::size_of::<C>());
2205    let total = total_len(dst_dims)?;
2206
2207    // Small tensor fast path: skip compute_order and compute_block_sizes
2208    let (fused_dims, ordered_strides, plan) = if total <= SMALL_TENSOR_THRESHOLD {
2209        build_plan_fused_small(dst_dims, &strides_list)
2210    } else {
2211        build_plan_fused(dst_dims, &strides_list, Some(0), elem_size)
2212    };
2213
2214    #[cfg(feature = "parallel")]
2215    {
2216        let total = total_len(&fused_dims)?;
2217        let nthreads = crate::execution_policy::rayon_threads();
2218        if total > MINTHREADLENGTH && nthreads > 1 {
2219            use crate::threading::SendPtr;
2220            let dst_send = SendPtr(dst_ptr);
2221            let a_send = SendPtr(a_ptr as *mut A);
2222            let b_send = SendPtr(b_ptr as *mut B);
2223            let c_send = SendPtr(c_ptr as *mut C);
2224
2225            let costs = compute_costs(&ordered_strides);
2226            let initial_offsets = vec![0isize; strides_list.len()];
2227            return mapreduce_threaded(
2228                &fused_dims,
2229                &plan.block,
2230                &ordered_strides,
2231                &initial_offsets,
2232                &costs,
2233                nthreads,
2234                0,
2235                1,
2236                &|dims, blocks, strides_list, offsets| {
2237                    for_each_inner_block_with_offsets(
2238                        dims,
2239                        blocks,
2240                        strides_list,
2241                        offsets,
2242                        |offsets, len, strides| {
2243                            let dp = unsafe { dst_send.as_ptr().offset(offsets[0]) };
2244                            let ap = unsafe { a_send.as_const().offset(offsets[1]) };
2245                            let bp = unsafe { b_send.as_const().offset(offsets[2]) };
2246                            let cp = unsafe { c_send.as_const().offset(offsets[3]) };
2247                            unsafe {
2248                                inner_loop_map3::<D, A, B, C, OpA, OpB, OpC>(
2249                                    dp, strides[0], ap, strides[1], bp, strides[2], cp, strides[3],
2250                                    len, &f,
2251                                )
2252                            };
2253                            Ok(())
2254                        },
2255                    )
2256                },
2257            );
2258        }
2259    }
2260
2261    let initial_offsets = vec![0isize; ordered_strides.len()];
2262    for_each_inner_block_preordered(
2263        &fused_dims,
2264        &plan.block,
2265        &ordered_strides,
2266        &initial_offsets,
2267        |offsets, len, strides| {
2268            let dp = unsafe { dst_ptr.offset(offsets[0]) };
2269            let ap = unsafe { a_ptr.offset(offsets[1]) };
2270            let bp = unsafe { b_ptr.offset(offsets[2]) };
2271            let cp = unsafe { c_ptr.offset(offsets[3]) };
2272            unsafe {
2273                inner_loop_map3::<D, A, B, C, OpA, OpB, OpC>(
2274                    dp, strides[0], ap, strides[1], bp, strides[2], cp, strides[3], len, &f,
2275                )
2276            };
2277            Ok(())
2278        },
2279    )
2280}
2281
2282/// Quaternary element-wise operation: `dest[i] = f(a[i], b[i], c[i], e[i])`.
2283pub fn zip_map4_into<
2284    D: Copy + MaybeSendSync,
2285    A: Copy + MaybeSendSync,
2286    B: Copy + MaybeSendSync,
2287    C: Copy + MaybeSendSync,
2288    E: Copy + MaybeSendSync,
2289    OpA: ElementOp<A>,
2290    OpB: ElementOp<B>,
2291    OpC: ElementOp<C>,
2292    OpE: ElementOp<E>,
2293>(
2294    dest: &mut StridedViewMut<D>,
2295    a: &StridedView<A, OpA>,
2296    b: &StridedView<B, OpB>,
2297    c: &StridedView<C, OpC>,
2298    e: &StridedView<E, OpE>,
2299    f: impl Fn(A, B, C, E) -> D + MaybeSync,
2300) -> Result<()> {
2301    ensure_same_shape(dest.dims(), a.dims())?;
2302    ensure_same_shape(dest.dims(), b.dims())?;
2303    ensure_same_shape(dest.dims(), c.dims())?;
2304    ensure_same_shape(dest.dims(), e.dims())?;
2305    let validated = validate_destination_layout(dest.dims(), dest.strides())?;
2306    zip_map4_into_validated(dest, a, b, c, e, f, validated)
2307}
2308
2309pub(crate) fn zip_map4_into_validated<
2310    D: Copy + MaybeSendSync,
2311    A: Copy + MaybeSendSync,
2312    B: Copy + MaybeSendSync,
2313    C: Copy + MaybeSendSync,
2314    E: Copy + MaybeSendSync,
2315    OpA: ElementOp<A>,
2316    OpB: ElementOp<B>,
2317    OpC: ElementOp<C>,
2318    OpE: ElementOp<E>,
2319>(
2320    dest: &mut StridedViewMut<D>,
2321    a: &StridedView<A, OpA>,
2322    b: &StridedView<B, OpB>,
2323    c: &StridedView<C, OpC>,
2324    e: &StridedView<E, OpE>,
2325    f: impl Fn(A, B, C, E) -> D + MaybeSync,
2326    _validated: ValidatedDestinationLayout,
2327) -> Result<()> {
2328    ensure_same_shape(dest.dims(), a.dims())?;
2329    ensure_same_shape(dest.dims(), b.dims())?;
2330    ensure_same_shape(dest.dims(), c.dims())?;
2331    ensure_same_shape(dest.dims(), e.dims())?;
2332    let dst_ptr = dest.as_mut_ptr();
2333    let a_ptr = a.ptr();
2334    let b_ptr = b.ptr();
2335    let c_ptr = c.ptr();
2336    let e_ptr = e.ptr();
2337
2338    let dst_dims = dest.dims();
2339    let dst_strides = dest.strides();
2340
2341    if sequential_contiguous_layout(
2342        dst_dims,
2343        &[
2344            dst_strides,
2345            a.strides(),
2346            b.strides(),
2347            c.strides(),
2348            e.strides(),
2349        ],
2350    )?
2351    .is_some()
2352    {
2353        let len = total_len(dst_dims)?;
2354        let dst = unsafe { std::slice::from_raw_parts_mut(dst_ptr, len) };
2355        let sa = unsafe { std::slice::from_raw_parts(a_ptr, len) };
2356        let sb = unsafe { std::slice::from_raw_parts(b_ptr, len) };
2357        let sc = unsafe { std::slice::from_raw_parts(c_ptr, len) };
2358        let se = unsafe { std::slice::from_raw_parts(e_ptr, len) };
2359        simd::dispatch_if_large(len, || {
2360            for i in 0..len {
2361                dst[i] = f(
2362                    OpA::apply(sa[i]),
2363                    OpB::apply(sb[i]),
2364                    OpC::apply(sc[i]),
2365                    OpE::apply(se[i]),
2366                );
2367            }
2368        });
2369        return Ok(());
2370    }
2371
2372    let strides_list: [&[isize]; 5] = [
2373        dst_strides,
2374        a.strides(),
2375        b.strides(),
2376        c.strides(),
2377        e.strides(),
2378    ];
2379    let elem_size = std::mem::size_of::<D>()
2380        .max(std::mem::size_of::<A>())
2381        .max(std::mem::size_of::<B>())
2382        .max(std::mem::size_of::<C>())
2383        .max(std::mem::size_of::<E>());
2384    let total = total_len(dst_dims)?;
2385
2386    // Small tensor fast path: skip compute_order and compute_block_sizes
2387    let (fused_dims, ordered_strides, plan) = if total <= SMALL_TENSOR_THRESHOLD {
2388        build_plan_fused_small(dst_dims, &strides_list)
2389    } else {
2390        build_plan_fused(dst_dims, &strides_list, Some(0), elem_size)
2391    };
2392
2393    #[cfg(feature = "parallel")]
2394    {
2395        let total = total_len(&fused_dims)?;
2396        let nthreads = crate::execution_policy::rayon_threads();
2397        if total > MINTHREADLENGTH && nthreads > 1 {
2398            use crate::threading::SendPtr;
2399            let dst_send = SendPtr(dst_ptr);
2400            let a_send = SendPtr(a_ptr as *mut A);
2401            let b_send = SendPtr(b_ptr as *mut B);
2402            let c_send = SendPtr(c_ptr as *mut C);
2403            let e_send = SendPtr(e_ptr as *mut E);
2404
2405            let costs = compute_costs(&ordered_strides);
2406            let initial_offsets = vec![0isize; strides_list.len()];
2407            return mapreduce_threaded(
2408                &fused_dims,
2409                &plan.block,
2410                &ordered_strides,
2411                &initial_offsets,
2412                &costs,
2413                nthreads,
2414                0,
2415                1,
2416                &|dims, blocks, strides_list, offsets| {
2417                    for_each_inner_block_with_offsets(
2418                        dims,
2419                        blocks,
2420                        strides_list,
2421                        offsets,
2422                        |offsets, len, strides| {
2423                            let dp = unsafe { dst_send.as_ptr().offset(offsets[0]) };
2424                            let ap = unsafe { a_send.as_const().offset(offsets[1]) };
2425                            let bp = unsafe { b_send.as_const().offset(offsets[2]) };
2426                            let cp = unsafe { c_send.as_const().offset(offsets[3]) };
2427                            let ep = unsafe { e_send.as_const().offset(offsets[4]) };
2428                            unsafe {
2429                                inner_loop_map4::<D, A, B, C, E, OpA, OpB, OpC, OpE>(
2430                                    dp, strides[0], ap, strides[1], bp, strides[2], cp, strides[3],
2431                                    ep, strides[4], len, &f,
2432                                )
2433                            };
2434                            Ok(())
2435                        },
2436                    )
2437                },
2438            );
2439        }
2440    }
2441
2442    let initial_offsets = vec![0isize; ordered_strides.len()];
2443    for_each_inner_block_preordered(
2444        &fused_dims,
2445        &plan.block,
2446        &ordered_strides,
2447        &initial_offsets,
2448        |offsets, len, strides| {
2449            let dp = unsafe { dst_ptr.offset(offsets[0]) };
2450            let ap = unsafe { a_ptr.offset(offsets[1]) };
2451            let bp = unsafe { b_ptr.offset(offsets[2]) };
2452            let cp = unsafe { c_ptr.offset(offsets[3]) };
2453            let ep = unsafe { e_ptr.offset(offsets[4]) };
2454            unsafe {
2455                inner_loop_map4::<D, A, B, C, E, OpA, OpB, OpC, OpE>(
2456                    dp, strides[0], ap, strides[1], bp, strides[2], cp, strides[3], ep, strides[4],
2457                    len, &f,
2458                )
2459            };
2460            Ok(())
2461        },
2462    )
2463}
2464
2465#[cfg(test)]
2466#[path = "map_view/tests/scalar_branch_tests.rs"]
2467mod scalar_branch_tests;
2468
2469#[cfg(all(test, feature = "parallel"))]
2470#[path = "map_view/tests/tests.rs"]
2471mod tests;