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