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