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::simd;
14use crate::view::{StridedView, StridedViewMut};
15use crate::{Result, StridedError};
16use core::mem::MaybeUninit;
17use std::ops::Mul;
18use strided_view::{ElementOp, Identity};
19
20use crate::fuse::compute_costs;
21use crate::threading::{
22    for_each_inner_block_with_offsets, is_parallel, mapreduce_threaded, MINTHREADLENGTH,
23};
24use crate::Exec;
25
26type AxisVec<T> = crate::inline::InlineVec<T>;
27
28/// Largest fused inner:row aspect at which the tiled transposed-scalar kernel
29/// is selected.
30///
31/// The cache-tiled kernel wins or ties the generic blocked plan across the
32/// measured shape/dtype table (u4-tiled #210), except when the fused inner run
33/// is more than 8x the row extent: then its strided destination writes lose to
34/// the blocked plan on f64 (`4096x256` ~1.4x, `8192x256` ~1.3x at 1T), so the
35/// detector declines and the generic path runs.
36const TRANSPOSED_SCALAR_TILE_MAX_INNER_ROW_ASPECT: usize = 8;
37
38/// Internal validation marker used by kernel-family execution adapters.
39///
40/// This marker is not bound to a destination. It does not authorize unchecked
41/// replay with a different geometry; callers of the unsafe execution entries
42/// must establish that association themselves.
43#[derive(Clone, Copy, Debug, Default)]
44pub struct ValidatedDestinationLayout(());
45
46/// Validate destination injectivity and return the replay marker.
47///
48/// Runs the exact (allocation-free pairwise, then spilling) injectivity check,
49/// so every injective layout the other kernel entry points accept is accepted
50/// here. Kernel families that need the marker for more than one step (for
51/// example validating shapes and then taking a specialization) should call this
52/// once and pass the marker on instead of re-running the check.
53#[inline]
54pub fn 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 + Send + Sync + 'static {
325    type Slot: Copy + Send + Sync + '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 + Send + Sync + '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 + Send + Sync + '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
818fn strided_offset_for_contiguous_linear_index(
819    dims: &[usize],
820    strides: &[isize],
821    axis_order: &[usize],
822    mut index: usize,
823) -> isize {
824    let mut offset = 0isize;
825    for &axis in axis_order {
826        let dim = dims[axis];
827        if dim == 0 {
828            return 0;
829        }
830        let coord = index % dim;
831        index /= dim;
832        offset += coord as isize * strides[axis];
833    }
834    offset
835}
836
837#[allow(clippy::too_many_arguments)]
838fn try_contiguous_range_mul<
839    O: MulOutput<D>,
840    D: Copy + Send + Sync + 'static,
841    A: Copy + Send + Sync + Mul<B, Output = D> + 'static,
842    B: Copy + Send + Sync + '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    exec: Exec<'_>,
852) -> bool {
853    // An overflowing element count falls back to the general path, which
854    // reports it as `OffsetOverflow`.
855    let Ok(total) = total_len(dims) else {
856        return false;
857    };
858    if total == 0 {
859        return true;
860    }
861    if total <= MINTHREADLENGTH {
862        return false;
863    }
864
865    let Some(plan) = contiguous_mul_range_plan(dims, dst_strides, a_strides, b_strides) else {
866        return false;
867    };
868    if !plan.walks_inputs_without_blocking() {
869        return false;
870    }
871
872    let inner_len = plan.inner_len.max(1);
873    let row_len = plan.row_len.max(1);
874    let block_len = inner_len.saturating_mul(row_len).max(1);
875    let outer_groups = total.div_ceil(block_len);
876
877    let nthreads = exec.budget().get();
878    if nthreads > 1 {
879        use crate::threading::{parallel_for_each, SendPtr};
880
881        let dst = SendPtr(dst_ptr);
882        let a = SendPtr(a_ptr as *mut A);
883        let b = SendPtr(b_ptr as *mut B);
884
885        if outer_groups < nthreads {
886            let chunk_len = total.div_ceil(nthreads);
887            let nchunks = total.div_ceil(chunk_len);
888
889            parallel_for_each(0..nchunks, exec, &|chunks| {
890                for chunk in chunks {
891                    let start = chunk * chunk_len;
892                    let end = (start + chunk_len).min(total);
893                    let mut index = start;
894
895                    while index < end {
896                        let in_inner = index % inner_len;
897                        let len = (inner_len - in_inner).min(end - index);
898                        let a_offset = strided_offset_for_contiguous_linear_index(
899                            dims,
900                            a_strides,
901                            &plan.axis_order,
902                            index,
903                        );
904                        let b_offset = strided_offset_for_contiguous_linear_index(
905                            dims,
906                            b_strides,
907                            &plan.axis_order,
908                            index,
909                        );
910
911                        unsafe {
912                            inner_loop_mul2::<O, D, A, B>(
913                                dst.as_ptr().add(index),
914                                1,
915                                a.as_const().offset(a_offset),
916                                plan.a_fast_stride,
917                                b.as_const().offset(b_offset),
918                                plan.b_fast_stride,
919                                len,
920                            );
921                        }
922                        index += len;
923                    }
924                }
925            });
926
927            return true;
928        }
929
930        let groups_per_chunk = outer_groups.div_ceil(nthreads);
931        let nchunks = outer_groups.div_ceil(groups_per_chunk);
932
933        parallel_for_each(0..nchunks, exec, &|chunks| {
934            for chunk in chunks {
935                let group_start = chunk * groups_per_chunk;
936                let group_end = (group_start + groups_per_chunk).min(outer_groups);
937                let mut cursor =
938                    ContiguousMulOuterCursor::new(dims, a_strides, b_strides, &plan, group_start);
939
940                for group in group_start..group_end {
941                    let index = group * block_len;
942                    unsafe {
943                        run_contiguous_mul_row_block::<O, D, A, B>(
944                            dst.as_ptr(),
945                            a.as_const(),
946                            b.as_const(),
947                            &plan,
948                            index,
949                            total,
950                            cursor.a_offset,
951                            cursor.b_offset,
952                        );
953                    }
954                    cursor.advance();
955                }
956            }
957        });
958
959        true
960    } else {
961        run_contiguous_range_mul_single_thread::<O, D, A, B>(
962            dst_ptr,
963            dims,
964            a_ptr,
965            a_strides,
966            b_ptr,
967            b_strides,
968            &plan,
969            total,
970            block_len,
971            outer_groups,
972        )
973    }
974}
975
976fn run_contiguous_range_mul_single_thread<
977    O: MulOutput<D>,
978    D: Copy + Send + Sync + 'static,
979    A: Copy + Send + Sync + Mul<B, Output = D> + 'static,
980    B: Copy + Send + Sync + 'static,
981>(
982    dst_ptr: *mut O::Slot,
983    dims: &[usize],
984    a_ptr: *const A,
985    a_strides: &[isize],
986    b_ptr: *const B,
987    b_strides: &[isize],
988    plan: &ContiguousMulRangePlan,
989    total: usize,
990    block_len: usize,
991    outer_groups: usize,
992) -> bool {
993    let mut cursor = ContiguousMulOuterCursor::new(dims, a_strides, b_strides, plan, 0);
994    for group in 0..outer_groups {
995        let index = group * block_len;
996        unsafe {
997            run_contiguous_mul_row_block::<O, D, A, B>(
998                dst_ptr,
999                a_ptr,
1000                b_ptr,
1001                plan,
1002                index,
1003                total,
1004                cursor.a_offset,
1005                cursor.b_offset,
1006            );
1007        }
1008        cursor.advance();
1009    }
1010    true
1011}
1012
1013/// Ternary inner loop: `dest[i] = f(a[i], b[i], c[i])`.
1014#[inline(always)]
1015unsafe fn inner_loop_map3<
1016    D: Copy,
1017    A: Copy,
1018    B: Copy,
1019    C: Copy,
1020    OpA: ElementOp<A>,
1021    OpB: ElementOp<B>,
1022    OpC: ElementOp<C>,
1023>(
1024    dp: *mut D,
1025    ds: isize,
1026    ap: *const A,
1027    a_s: isize,
1028    bp: *const B,
1029    b_s: isize,
1030    cp: *const C,
1031    c_s: isize,
1032    len: usize,
1033    f: &impl Fn(A, B, C) -> D,
1034) {
1035    if ds == 1 && a_s == 1 && b_s == 1 && c_s == 1 {
1036        let src_a = std::slice::from_raw_parts(ap, len);
1037        let src_b = std::slice::from_raw_parts(bp, len);
1038        let src_c = std::slice::from_raw_parts(cp, len);
1039        let dst = std::slice::from_raw_parts_mut(dp, len);
1040        simd::dispatch_if_large(len, || {
1041            for (((d, &a), &b), &c) in dst.iter_mut().zip(src_a).zip(src_b).zip(src_c) {
1042                *d = f(OpA::apply(a), OpB::apply(b), OpC::apply(c));
1043            }
1044        });
1045    } else {
1046        let mut dp = dp;
1047        let mut ap = ap;
1048        let mut bp = bp;
1049        let mut cp = cp;
1050        for _ in 0..len {
1051            *dp = f(OpA::apply(*ap), OpB::apply(*bp), OpC::apply(*cp));
1052            dp = dp.wrapping_offset(ds);
1053            ap = ap.wrapping_offset(a_s);
1054            bp = bp.wrapping_offset(b_s);
1055            cp = cp.wrapping_offset(c_s);
1056        }
1057    }
1058}
1059
1060/// Quaternary inner loop: `dest[i] = f(a[i], b[i], c[i], e[i])`.
1061#[inline(always)]
1062unsafe fn inner_loop_map4<
1063    D: Copy,
1064    A: Copy,
1065    B: Copy,
1066    C: Copy,
1067    E: Copy,
1068    OpA: ElementOp<A>,
1069    OpB: ElementOp<B>,
1070    OpC: ElementOp<C>,
1071    OpE: ElementOp<E>,
1072>(
1073    dp: *mut D,
1074    ds: isize,
1075    ap: *const A,
1076    a_s: isize,
1077    bp: *const B,
1078    b_s: isize,
1079    cp: *const C,
1080    c_s: isize,
1081    ep: *const E,
1082    e_s: isize,
1083    len: usize,
1084    f: &impl Fn(A, B, C, E) -> D,
1085) {
1086    if ds == 1 && a_s == 1 && b_s == 1 && c_s == 1 && e_s == 1 {
1087        let src_a = std::slice::from_raw_parts(ap, len);
1088        let src_b = std::slice::from_raw_parts(bp, len);
1089        let src_c = std::slice::from_raw_parts(cp, len);
1090        let src_e = std::slice::from_raw_parts(ep, len);
1091        let dst = std::slice::from_raw_parts_mut(dp, len);
1092        simd::dispatch_if_large(len, || {
1093            for i in 0..len {
1094                dst[i] = f(
1095                    OpA::apply(src_a[i]),
1096                    OpB::apply(src_b[i]),
1097                    OpC::apply(src_c[i]),
1098                    OpE::apply(src_e[i]),
1099                );
1100            }
1101        });
1102    } else {
1103        let mut dp = dp;
1104        let mut ap = ap;
1105        let mut bp = bp;
1106        let mut cp = cp;
1107        let mut ep = ep;
1108        for _ in 0..len {
1109            *dp = f(
1110                OpA::apply(*ap),
1111                OpB::apply(*bp),
1112                OpC::apply(*cp),
1113                OpE::apply(*ep),
1114            );
1115            dp = dp.wrapping_offset(ds);
1116            ap = ap.wrapping_offset(a_s);
1117            bp = bp.wrapping_offset(b_s);
1118            cp = cp.wrapping_offset(c_s);
1119            ep = ep.wrapping_offset(e_s);
1120        }
1121    }
1122}
1123
1124/// Apply a function element-wise from source to destination.
1125///
1126/// The element operation `Op` is applied lazily when reading from `src`.
1127/// Source and destination may have different element types.
1128pub fn map_into<D: Copy + Send + Sync, A: Copy + Send + Sync, Op: ElementOp<A>>(
1129    dest: &mut StridedViewMut<D>,
1130    src: &StridedView<A, Op>,
1131    f: impl Fn(A) -> D + Sync,
1132    exec: Exec<'_>,
1133) -> Result<()> {
1134    map_parts_into::<D, A, Op, NoKernel>(
1135        dest.as_mut_ptr(),
1136        dest.dims(),
1137        dest.strides(),
1138        src.ptr(),
1139        src.dims(),
1140        src.strides(),
1141        f,
1142        &NoKernel,
1143        exec,
1144    )
1145}
1146
1147pub(crate) fn map_into_validated<D: Copy + Send + Sync, A: Copy + Send + Sync, Op: ElementOp<A>>(
1148    dest: &mut StridedViewMut<D>,
1149    src: &StridedView<A, Op>,
1150    f: impl Fn(A) -> D + Sync,
1151    validated: ValidatedDestinationLayout,
1152    exec: Exec<'_>,
1153) -> Result<()> {
1154    ensure_same_shape(dest.dims(), src.dims())?;
1155    map_parts_into_validated::<D, A, Op, NoKernel>(
1156        dest.as_mut_ptr(),
1157        dest.dims(),
1158        dest.strides(),
1159        src.ptr(),
1160        src.strides(),
1161        f,
1162        &NoKernel,
1163        validated,
1164        exec,
1165    )
1166}
1167
1168/// Apply an element-wise function into a potentially uninitialized destination, handing a
1169/// vectorized kernel the flat runs it can take.
1170///
1171/// `kernel` is offered a flat, unit-stride run whenever this crate has one to give: the whole
1172/// buffer when the two views share a contiguous layout and the operation runs sequentially, and
1173/// one block at a time on the parallel path. It must write every element of the run or write
1174/// nothing — returning `false` makes that run fall back to `scalar`, element by element. Runs whose
1175/// strides are not both 1 never reach it. It is called from whichever tasks the budget of
1176/// `exec` selected, so it must be `Sync`.
1177///
1178/// The point of the entry point is that a caller which owns vector kernels — a SIMD math library,
1179/// say — reuses this crate's layout logic and this crate's threading decision instead of
1180/// rebuilding either. `scalar` and `kernel` must agree up to the caller's own error bound; this
1181/// crate only guarantees that every destination element is written exactly once.
1182///
1183/// # Examples
1184/// ```
1185/// use std::mem::MaybeUninit;
1186/// use strided_basic::{map_into_uninit_with_kernel, Exec, StridedView, StridedViewMut};
1187///
1188/// let src = vec![1.0f64, 2.0, 3.0];
1189/// let mut out = vec![MaybeUninit::<f64>::uninit(); 3];
1190/// let mut dest = StridedViewMut::new(&mut out, &[3], &[1], 0).unwrap();
1191/// let source = StridedView::new(&src, &[3], &[1], 0).unwrap();
1192/// let kernel = |src: &[f64], dst: &mut [MaybeUninit<f64>]| {
1193///     for (slot, value) in dst.iter_mut().zip(src) {
1194///         slot.write(value * 2.0);
1195///     }
1196///     true
1197/// };
1198/// map_into_uninit_with_kernel(&mut dest, &source, kernel, |x| x * 2.0, Exec::serial()).unwrap();
1199/// // SAFETY: the kernel wrote every element.
1200/// let values: Vec<f64> = out.into_iter().map(|x| unsafe { x.assume_init() }).collect();
1201/// assert_eq!(values, [2.0, 4.0, 6.0]);
1202/// ```
1203///
1204/// # Errors
1205/// Returns [`StridedError::ShapeMismatch`] when the shapes differ and
1206/// [`StridedError::NonInjectiveOutputLayout`] when the destination layout cannot be proven
1207/// injective.
1208pub fn map_into_uninit_with_kernel<T: Copy + Send + Sync>(
1209    dest: &mut StridedViewMut<MaybeUninit<T>>,
1210    src: &StridedView<T>,
1211    kernel: impl Fn(&[T], &mut [MaybeUninit<T>]) -> bool + Sync,
1212    scalar: impl Fn(T) -> T + Sync,
1213    exec: Exec<'_>,
1214) -> Result<()> {
1215    map_parts_into::<MaybeUninit<T>, T, Identity, _>(
1216        dest.as_mut_ptr(),
1217        dest.dims(),
1218        dest.strides(),
1219        src.ptr(),
1220        src.dims(),
1221        src.strides(),
1222        |value| MaybeUninit::new(scalar(value)),
1223        &SliceKernel(kernel),
1224        exec,
1225    )
1226}
1227
1228/// Map a raw destination from a raw source without preparing a plan.
1229///
1230/// Only the raw-replay unit tests need this entry now; the crate's production
1231/// callers either go through the fused traversal (raw ops, `CopyPlan`) or
1232/// [`map_raw_into_validated`].
1233#[cfg(test)]
1234pub(crate) fn map_raw_into<D: Copy + Send + Sync, A: Copy + Send + Sync, Op: ElementOp<A>>(
1235    dest: &mut crate::RawStridedMut<'_, D>,
1236    src: &crate::RawStridedRef<'_, A>,
1237    f: impl Fn(A) -> D + Sync,
1238    exec: Exec<'_>,
1239) -> Result<()> {
1240    map_parts_into::<D, A, Op, NoKernel>(
1241        dest.as_mut_ptr(),
1242        dest.dims(),
1243        dest.strides(),
1244        src.ptr(),
1245        src.dims(),
1246        src.strides(),
1247        f,
1248        &NoKernel,
1249        exec,
1250    )
1251}
1252
1253pub(crate) fn map_raw_into_validated<
1254    D: Copy + Send + Sync,
1255    A: Copy + Send + Sync,
1256    Op: ElementOp<A>,
1257>(
1258    dest: &mut crate::RawStridedMut<'_, D>,
1259    src: &crate::RawStridedRef<'_, A>,
1260    f: impl Fn(A) -> D + Sync,
1261    validated: ValidatedDestinationLayout,
1262    exec: Exec<'_>,
1263) -> Result<()> {
1264    ensure_same_shape(dest.dims(), src.dims())?;
1265    map_parts_into_validated::<D, A, Op, NoKernel>(
1266        dest.as_mut_ptr(),
1267        dest.dims(),
1268        dest.strides(),
1269        src.ptr(),
1270        src.strides(),
1271        f,
1272        &NoKernel,
1273        validated,
1274        exec,
1275    )
1276}
1277
1278#[allow(clippy::too_many_arguments)]
1279fn map_parts_into<
1280    D: Copy + Send + Sync,
1281    A: Copy + Send + Sync,
1282    Op: ElementOp<A>,
1283    K: MapKernel<A, D> + Sync,
1284>(
1285    dst_ptr: *mut D,
1286    dst_dims: &[usize],
1287    dst_strides: &[isize],
1288    src_ptr: *const A,
1289    src_dims: &[usize],
1290    src_strides: &[isize],
1291    f: impl Fn(A) -> D + Sync,
1292    kernel: &K,
1293    exec: Exec<'_>,
1294) -> Result<()> {
1295    ensure_same_shape(dst_dims, src_dims)?;
1296    let validated = validate_destination_layout(dst_dims, dst_strides)?;
1297    map_parts_into_validated::<D, A, Op, K>(
1298        dst_ptr,
1299        dst_dims,
1300        dst_strides,
1301        src_ptr,
1302        src_strides,
1303        f,
1304        kernel,
1305        validated,
1306        exec,
1307    )
1308}
1309
1310#[allow(clippy::too_many_arguments)]
1311/// Apply `f` (with the element operation) over a fully contiguous pair.
1312///
1313/// Kept out of line, and out of the caller's dispatch frame, so that it cannot
1314/// change the code shape of the blocked kernels below; the operation is moved in
1315/// by value because a dispatched closure that only borrows it stays scalar
1316/// (tensor4all/strided-rs#301).
1317#[inline(never)]
1318fn map_contiguous<D: Copy, A: Copy, Op: ElementOp<A>>(
1319    dst: &mut [D],
1320    src: &[A],
1321    f: impl Fn(A) -> D,
1322) {
1323    simd::dispatch_if_large(dst.len(), move || {
1324        for i in 0..dst.len() {
1325            dst[i] = f(Op::apply(src[i]));
1326        }
1327    });
1328}
1329
1330#[allow(clippy::too_many_arguments)]
1331fn map_parts_into_validated<
1332    D: Copy + Send + Sync,
1333    A: Copy + Send + Sync,
1334    Op: ElementOp<A>,
1335    K: MapKernel<A, D> + Sync,
1336>(
1337    dst_ptr: *mut D,
1338    dst_dims: &[usize],
1339    dst_strides: &[isize],
1340    src_ptr: *const A,
1341    src_strides: &[isize],
1342    f: impl Fn(A) -> D + Sync,
1343    kernel: &K,
1344    _validated: ValidatedDestinationLayout,
1345    exec: Exec<'_>,
1346) -> Result<()> {
1347    if sequential_contiguous_layout(dst_dims, &[dst_strides, src_strides], &exec)?.is_some() {
1348        let len = total_len(dst_dims)?;
1349        // SAFETY: the contiguous-layout check proves that `dst` and `src` address `len` elements in
1350        // the same order starting at their base pointer, so each is exactly a flat run of `len`
1351        // elements.
1352        let dst = unsafe { std::slice::from_raw_parts_mut(dst_ptr, len) };
1353        let src = unsafe { std::slice::from_raw_parts(src_ptr, len) };
1354        if kernel.apply(src, dst) {
1355            return Ok(());
1356        }
1357        map_contiguous::<D, A, Op>(dst, src, f);
1358        return Ok(());
1359    }
1360
1361    let strides_list: [&[isize]; 2] = [dst_strides, src_strides];
1362    let elem_size = std::mem::size_of::<D>().max(std::mem::size_of::<A>());
1363    let total = total_len(dst_dims)?;
1364
1365    // Small tensor fast path: skip compute_order and compute_block_sizes
1366    let (fused_dims, ordered_strides, plan) = if total <= SMALL_TENSOR_THRESHOLD {
1367        build_plan_fused_small(dst_dims, &strides_list)
1368    } else {
1369        build_plan_fused(dst_dims, &strides_list, Some(0), elem_size)
1370    };
1371
1372    let total = total_len(&fused_dims)?;
1373    if is_parallel(total, &exec) {
1374        use crate::threading::SendPtr;
1375        let dst_send = SendPtr(dst_ptr);
1376        let src_send = SendPtr(src_ptr as *mut A);
1377
1378        let costs = compute_costs(&ordered_strides);
1379        let initial_offsets: crate::inline::OperandVec<isize> =
1380            strides_list.iter().map(|_| 0isize).collect();
1381        return mapreduce_threaded(
1382            &fused_dims,
1383            &plan.block,
1384            &ordered_strides,
1385            &initial_offsets,
1386            &costs,
1387            exec,
1388            &|dims, blocks, strides_list, offsets| {
1389                for_each_inner_block_with_offsets(
1390                    dims,
1391                    blocks,
1392                    strides_list,
1393                    offsets,
1394                    |offsets, len, strides| {
1395                        let dp = unsafe { dst_send.as_ptr().offset(offsets[0]) };
1396                        let sp = unsafe { src_send.as_const().offset(offsets[1]) };
1397                        unsafe {
1398                            inner_loop_map1::<D, A, Op, K>(
1399                                dp, strides[0], sp, strides[1], len, &f, kernel,
1400                            )
1401                        };
1402                        Ok(())
1403                    },
1404                )
1405            },
1406        );
1407    }
1408
1409    let initial_offsets: crate::inline::OperandVec<isize> =
1410        ordered_strides.iter().map(|_| 0isize).collect();
1411    for_each_inner_block_preordered(
1412        &fused_dims,
1413        &plan.block,
1414        &ordered_strides,
1415        &initial_offsets,
1416        |offsets, len, strides| {
1417            let dp = unsafe { dst_ptr.offset(offsets[0]) };
1418            let sp = unsafe { src_ptr.offset(offsets[1]) };
1419            unsafe {
1420                inner_loop_map1::<D, A, Op, K>(dp, strides[0], sp, strides[1], len, &f, kernel)
1421            };
1422            Ok(())
1423        },
1424    )
1425}
1426
1427/// Binary element-wise operation: `dest[i] = f(a[i], b[i])`.
1428///
1429/// Source operands `a` and `b` may have different element types from each other
1430/// and from `dest`. The closure `f` handles per-element type conversion.
1431pub fn zip_map2_into<
1432    D: Copy + Send + Sync,
1433    A: Copy + Send + Sync,
1434    B: Copy + Send + Sync,
1435    OpA: ElementOp<A>,
1436    OpB: ElementOp<B>,
1437>(
1438    dest: &mut StridedViewMut<D>,
1439    a: &StridedView<A, OpA>,
1440    b: &StridedView<B, OpB>,
1441    f: impl Fn(A, B) -> D + Sync,
1442    exec: Exec<'_>,
1443) -> Result<()> {
1444    zip_map2_parts_into::<D, A, B, OpA, OpB>(
1445        dest.as_mut_ptr(),
1446        dest.dims(),
1447        dest.strides(),
1448        a.ptr(),
1449        a.dims(),
1450        a.strides(),
1451        b.ptr(),
1452        b.dims(),
1453        b.strides(),
1454        f,
1455        exec,
1456    )
1457}
1458
1459pub(crate) fn zip_map2_into_validated<
1460    D: Copy + Send + Sync,
1461    A: Copy + Send + Sync,
1462    B: Copy + Send + Sync,
1463    OpA: ElementOp<A>,
1464    OpB: ElementOp<B>,
1465>(
1466    dest: &mut StridedViewMut<D>,
1467    a: &StridedView<A, OpA>,
1468    b: &StridedView<B, OpB>,
1469    f: impl Fn(A, B) -> D + Sync,
1470    validated: ValidatedDestinationLayout,
1471    exec: Exec<'_>,
1472) -> Result<()> {
1473    zip_map2_parts_into_validated::<D, A, B, OpA, OpB>(
1474        dest.as_mut_ptr(),
1475        dest.dims(),
1476        dest.strides(),
1477        a.ptr(),
1478        a.strides(),
1479        b.ptr(),
1480        b.strides(),
1481        f,
1482        validated,
1483        exec,
1484    )
1485}
1486
1487/// Runtime comparison selected once before entering the element loop.
1488#[non_exhaustive]
1489#[derive(Clone, Copy, Debug, Eq, PartialEq)]
1490pub enum CompareOp {
1491    Eq,
1492    Lt,
1493    Le,
1494    Gt,
1495    Ge,
1496}
1497
1498/// Compare two ordered views elementwise into a Boolean destination.
1499///
1500/// Unlike embedding a runtime operation match inside a [`zip_map2_into`]
1501/// closure, this entry point selects a fixed comparison before traversal.
1502///
1503/// # Errors
1504///
1505/// Returns [`StridedError::ShapeMismatch`] when the source and destination
1506/// shapes differ, or [`StridedError::NonInjectiveOutputLayout`] when distinct
1507/// logical destination elements may overlap.
1508pub fn compare_into<T, OpA, OpB>(
1509    dest: &mut StridedViewMut<bool>,
1510    a: &StridedView<T, OpA>,
1511    b: &StridedView<T, OpB>,
1512    op: CompareOp,
1513    exec: Exec<'_>,
1514) -> Result<()>
1515where
1516    T: Copy + Send + Sync + PartialOrd,
1517    OpA: ElementOp<T>,
1518    OpB: ElementOp<T>,
1519{
1520    match op {
1521        CompareOp::Eq => zip_map2_into(dest, a, b, |lhs, rhs| lhs == rhs, exec),
1522        CompareOp::Lt => zip_map2_into(dest, a, b, |lhs, rhs| lhs < rhs, exec),
1523        CompareOp::Le => zip_map2_into(dest, a, b, |lhs, rhs| lhs <= rhs, exec),
1524        CompareOp::Gt => zip_map2_into(dest, a, b, |lhs, rhs| lhs > rhs, exec),
1525        CompareOp::Ge => zip_map2_into(dest, a, b, |lhs, rhs| lhs >= rhs, exec),
1526    }
1527}
1528
1529/// Compare two views into a fully overwritten uninitialized Boolean output.
1530///
1531/// Dtype-independent shape, destination-injectivity, and reachable-byte
1532/// overlap validation completes before the first write. Safe Rust borrows
1533/// already prevent input/output aliasing; the explicit overlap check preserves
1534/// the contract for views produced through unsafe constructors.
1535///
1536/// `Ok(())` means every logical destination element is initialized. An error
1537/// occurs before writes. A panic during replay may leave a partially initialized
1538/// destination, which remains safe to drop as `MaybeUninit<bool>`.
1539///
1540/// # Errors
1541///
1542/// Returns [`StridedError::ShapeMismatch`] for unequal shapes,
1543/// [`StridedError::NonInjectiveOutputLayout`] for an overlapping output layout,
1544/// [`StridedError::OverlappingInputOutput`] for aliased storage, or
1545/// [`StridedError::OffsetOverflow`] when a reachable byte range is not
1546/// representable.
1547pub fn compare_into_uninit<T, OpA, OpB>(
1548    dest: &mut StridedViewMut<MaybeUninit<bool>>,
1549    a: &StridedView<T, OpA>,
1550    b: &StridedView<T, OpB>,
1551    op: CompareOp,
1552    exec: Exec<'_>,
1553) -> Result<()>
1554where
1555    T: Copy + Send + Sync + PartialOrd,
1556    OpA: ElementOp<T>,
1557    OpB: ElementOp<T>,
1558{
1559    ensure_same_shape(dest.dims(), a.dims())?;
1560    ensure_same_shape(dest.dims(), b.dims())?;
1561    let validated = validate_destination_layout(dest.dims(), dest.strides())?;
1562    validate_typed_no_overlap(dest, a, 0)?;
1563    validate_typed_no_overlap(dest, b, 1)?;
1564    match op {
1565        CompareOp::Eq => zip_map2_into_validated(
1566            dest,
1567            a,
1568            b,
1569            |lhs, rhs| MaybeUninit::new(lhs == rhs),
1570            validated,
1571            exec,
1572        ),
1573        CompareOp::Lt => zip_map2_into_validated(
1574            dest,
1575            a,
1576            b,
1577            |lhs, rhs| MaybeUninit::new(lhs < rhs),
1578            validated,
1579            exec,
1580        ),
1581        CompareOp::Le => zip_map2_into_validated(
1582            dest,
1583            a,
1584            b,
1585            |lhs, rhs| MaybeUninit::new(lhs <= rhs),
1586            validated,
1587            exec,
1588        ),
1589        CompareOp::Gt => zip_map2_into_validated(
1590            dest,
1591            a,
1592            b,
1593            |lhs, rhs| MaybeUninit::new(lhs > rhs),
1594            validated,
1595            exec,
1596        ),
1597        CompareOp::Ge => zip_map2_into_validated(
1598            dest,
1599            a,
1600            b,
1601            |lhs, rhs| MaybeUninit::new(lhs >= rhs),
1602            validated,
1603            exec,
1604        ),
1605    }
1606}
1607
1608pub(crate) fn zip_map2_raw_into_validated<
1609    D: Copy + Send + Sync,
1610    A: Copy + Send + Sync,
1611    B: Copy + Send + Sync,
1612    OpA: ElementOp<A>,
1613    OpB: ElementOp<B>,
1614>(
1615    dest: &mut crate::RawStridedMut<'_, D>,
1616    a: &crate::RawStridedRef<'_, A>,
1617    b: &crate::RawStridedRef<'_, B>,
1618    f: impl Fn(A, B) -> D + Sync,
1619    validated: ValidatedDestinationLayout,
1620    exec: Exec<'_>,
1621) -> Result<()> {
1622    ensure_same_shape(dest.dims(), a.dims())?;
1623    ensure_same_shape(dest.dims(), b.dims())?;
1624    zip_map2_parts_into_validated::<D, A, B, OpA, OpB>(
1625        dest.as_mut_ptr(),
1626        dest.dims(),
1627        dest.strides(),
1628        a.ptr(),
1629        a.strides(),
1630        b.ptr(),
1631        b.strides(),
1632        f,
1633        validated,
1634        exec,
1635    )
1636}
1637
1638#[allow(clippy::too_many_arguments)]
1639fn zip_map2_parts_into<
1640    D: Copy + Send + Sync,
1641    A: Copy + Send + Sync,
1642    B: Copy + Send + Sync,
1643    OpA: ElementOp<A>,
1644    OpB: ElementOp<B>,
1645>(
1646    dst_ptr: *mut D,
1647    dst_dims: &[usize],
1648    dst_strides: &[isize],
1649    a_ptr: *const A,
1650    a_dims: &[usize],
1651    a_strides: &[isize],
1652    b_ptr: *const B,
1653    b_dims: &[usize],
1654    b_strides: &[isize],
1655    f: impl Fn(A, B) -> D + Sync,
1656    exec: Exec<'_>,
1657) -> Result<()> {
1658    ensure_same_shape(dst_dims, a_dims)?;
1659    ensure_same_shape(dst_dims, b_dims)?;
1660    let validated = validate_destination_layout(dst_dims, dst_strides)?;
1661    zip_map2_parts_into_validated::<D, A, B, OpA, OpB>(
1662        dst_ptr,
1663        dst_dims,
1664        dst_strides,
1665        a_ptr,
1666        a_strides,
1667        b_ptr,
1668        b_strides,
1669        f,
1670        validated,
1671        exec,
1672    )
1673}
1674
1675#[allow(clippy::too_many_arguments)]
1676fn zip_map2_parts_into_validated<
1677    D: Copy + Send + Sync,
1678    A: Copy + Send + Sync,
1679    B: Copy + Send + Sync,
1680    OpA: ElementOp<A>,
1681    OpB: ElementOp<B>,
1682>(
1683    dst_ptr: *mut D,
1684    dst_dims: &[usize],
1685    dst_strides: &[isize],
1686    a_ptr: *const A,
1687    a_strides: &[isize],
1688    b_ptr: *const B,
1689    b_strides: &[isize],
1690    f: impl Fn(A, B) -> D + Sync,
1691    _validated: ValidatedDestinationLayout,
1692    exec: Exec<'_>,
1693) -> Result<()> {
1694    if sequential_contiguous_layout(dst_dims, &[dst_strides, a_strides, b_strides], &exec)?
1695        .is_some()
1696    {
1697        let len = total_len(dst_dims)?;
1698        let dst = unsafe { std::slice::from_raw_parts_mut(dst_ptr, len) };
1699        let sa = unsafe { std::slice::from_raw_parts(a_ptr, len) };
1700        let sb = unsafe { std::slice::from_raw_parts(b_ptr, len) };
1701        simd::dispatch_if_large(len, || {
1702            for i in 0..len {
1703                dst[i] = f(OpA::apply(sa[i]), OpB::apply(sb[i]));
1704            }
1705        });
1706        return Ok(());
1707    }
1708
1709    let strides_list: [&[isize]; 3] = [dst_strides, a_strides, b_strides];
1710    let elem_size = std::mem::size_of::<D>()
1711        .max(std::mem::size_of::<A>())
1712        .max(std::mem::size_of::<B>());
1713    let total = total_len(dst_dims)?;
1714
1715    // Small tensor fast path: skip compute_order and compute_block_sizes
1716    let (fused_dims, ordered_strides, plan) = if total <= SMALL_TENSOR_THRESHOLD {
1717        build_plan_fused_small(dst_dims, &strides_list)
1718    } else {
1719        build_plan_fused(dst_dims, &strides_list, Some(0), elem_size)
1720    };
1721
1722    let total = total_len(&fused_dims)?;
1723    if is_parallel(total, &exec) {
1724        use crate::threading::SendPtr;
1725        let dst_send = SendPtr(dst_ptr);
1726        let a_send = SendPtr(a_ptr as *mut A);
1727        let b_send = SendPtr(b_ptr as *mut B);
1728
1729        let costs = compute_costs(&ordered_strides);
1730        let initial_offsets: crate::inline::OperandVec<isize> =
1731            strides_list.iter().map(|_| 0isize).collect();
1732        return mapreduce_threaded(
1733            &fused_dims,
1734            &plan.block,
1735            &ordered_strides,
1736            &initial_offsets,
1737            &costs,
1738            exec,
1739            &|dims, blocks, strides_list, offsets| {
1740                for_each_inner_block_with_offsets(
1741                    dims,
1742                    blocks,
1743                    strides_list,
1744                    offsets,
1745                    |offsets, len, strides| {
1746                        let dp = unsafe { dst_send.as_ptr().offset(offsets[0]) };
1747                        let ap = unsafe { a_send.as_const().offset(offsets[1]) };
1748                        let bp = unsafe { b_send.as_const().offset(offsets[2]) };
1749                        unsafe {
1750                            inner_loop_map2::<D, A, B, OpA, OpB>(
1751                                dp, strides[0], ap, strides[1], bp, strides[2], len, &f,
1752                            )
1753                        };
1754                        Ok(())
1755                    },
1756                )
1757            },
1758        );
1759    }
1760
1761    let initial_offsets: crate::inline::OperandVec<isize> =
1762        ordered_strides.iter().map(|_| 0isize).collect();
1763    for_each_inner_block_preordered(
1764        &fused_dims,
1765        &plan.block,
1766        &ordered_strides,
1767        &initial_offsets,
1768        |offsets, len, strides| {
1769            let dp = unsafe { dst_ptr.offset(offsets[0]) };
1770            let ap = unsafe { a_ptr.offset(offsets[1]) };
1771            let bp = unsafe { b_ptr.offset(offsets[2]) };
1772            unsafe {
1773                inner_loop_map2::<D, A, B, OpA, OpB>(
1774                    dp, strides[0], ap, strides[1], bp, strides[2], len, &f,
1775                )
1776            };
1777            Ok(())
1778        },
1779    )
1780}
1781
1782#[allow(clippy::too_many_arguments)]
1783fn mul_identity_into_raw<
1784    O: MulOutput<D>,
1785    D: Copy + Send + Sync + 'static,
1786    A: Copy + Send + Sync + Mul<B, Output = D> + 'static,
1787    B: Copy + Send + Sync + 'static,
1788>(
1789    dst_ptr: *mut O::Slot,
1790    dst_dims: &[usize],
1791    dst_strides: &[isize],
1792    a_ptr: *const A,
1793    a_strides: &[isize],
1794    b_ptr: *const B,
1795    b_strides: &[isize],
1796    _validated: ValidatedDestinationLayout,
1797    mut exec: Exec<'_>,
1798) -> Result<()> {
1799    debug_assert_eq!(dst_dims.len(), a_strides.len());
1800    debug_assert_eq!(dst_dims.len(), b_strides.len());
1801
1802    if sequential_contiguous_layout(dst_dims, &[dst_strides, a_strides, b_strides], &exec)?
1803        .is_some()
1804    {
1805        let len = total_len(dst_dims)?;
1806        let sa = unsafe { std::slice::from_raw_parts(a_ptr, len) };
1807        let sb = unsafe { std::slice::from_raw_parts(b_ptr, len) };
1808        if unsafe { O::try_contiguous(dst_ptr, len, sa, sb) } {
1809            return Ok(());
1810        }
1811        for i in 0..len {
1812            unsafe { O::write(dst_ptr.add(i), multiply_value(sa[i], sb[i])) };
1813        }
1814        return Ok(());
1815    }
1816
1817    let strides_list: [&[isize]; 3] = [dst_strides, a_strides, b_strides];
1818    let elem_size = std::mem::size_of::<D>()
1819        .max(std::mem::size_of::<A>())
1820        .max(std::mem::size_of::<B>());
1821    let total = total_len(dst_dims)?;
1822
1823    if try_contiguous_range_mul::<O, D, A, B>(
1824        dst_ptr,
1825        dst_dims,
1826        dst_strides,
1827        a_ptr,
1828        a_strides,
1829        b_ptr,
1830        b_strides,
1831        exec.reborrow(),
1832    ) {
1833        return Ok(());
1834    }
1835
1836    let (fused_dims, ordered_strides, plan) = if total <= SMALL_TENSOR_THRESHOLD {
1837        build_plan_fused_small(dst_dims, &strides_list)
1838    } else {
1839        build_plan_fused(dst_dims, &strides_list, Some(0), elem_size)
1840    };
1841
1842    let total = total_len(&fused_dims)?;
1843    if is_parallel(total, &exec) {
1844        use crate::threading::SendPtr;
1845        let dst_send = SendPtr(dst_ptr);
1846        let a_send = SendPtr(a_ptr as *mut A);
1847        let b_send = SendPtr(b_ptr as *mut B);
1848
1849        let costs = compute_costs(&ordered_strides);
1850        let initial_offsets: crate::inline::OperandVec<isize> =
1851            strides_list.iter().map(|_| 0isize).collect();
1852        return mapreduce_threaded(
1853            &fused_dims,
1854            &plan.block,
1855            &ordered_strides,
1856            &initial_offsets,
1857            &costs,
1858            exec,
1859            &|dims, blocks, strides_list, offsets| {
1860                for_each_inner_block_with_offsets(
1861                    dims,
1862                    blocks,
1863                    strides_list,
1864                    offsets,
1865                    |offsets, len, strides| {
1866                        let dp = unsafe { dst_send.as_ptr().offset(offsets[0]) };
1867                        let ap = unsafe { a_send.as_const().offset(offsets[1]) };
1868                        let bp = unsafe { b_send.as_const().offset(offsets[2]) };
1869                        unsafe {
1870                            inner_loop_mul2::<O, D, A, B>(
1871                                dp, strides[0], ap, strides[1], bp, strides[2], len,
1872                            )
1873                        };
1874                        Ok(())
1875                    },
1876                )
1877            },
1878        );
1879    }
1880
1881    let initial_offsets: crate::inline::OperandVec<isize> =
1882        ordered_strides.iter().map(|_| 0isize).collect();
1883    for_each_inner_block_preordered(
1884        &fused_dims,
1885        &plan.block,
1886        &ordered_strides,
1887        &initial_offsets,
1888        |offsets, len, strides| {
1889            let dp = unsafe { dst_ptr.offset(offsets[0]) };
1890            let ap = unsafe { a_ptr.offset(offsets[1]) };
1891            let bp = unsafe { b_ptr.offset(offsets[2]) };
1892            unsafe {
1893                inner_loop_mul2::<O, D, A, B>(dp, strides[0], ap, strides[1], bp, strides[2], len)
1894            };
1895            Ok(())
1896        },
1897    )
1898}
1899
1900/// Element-wise multiplication: `dest[i] = a[i] * b[i]`.
1901///
1902/// All views must have the same shape. Broadcast operands should be represented
1903/// as stride-0 views before calling this function.
1904pub fn mul_into<
1905    D: Copy + Send + Sync + 'static,
1906    A: Copy + Mul<B, Output = D> + Send + Sync + 'static,
1907    B: Copy + Send + Sync + 'static,
1908    OpA: ElementOp<A>,
1909    OpB: ElementOp<B>,
1910>(
1911    dest: &mut StridedViewMut<D>,
1912    a: &StridedView<A, OpA>,
1913    b: &StridedView<B, OpB>,
1914    exec: Exec<'_>,
1915) -> Result<()> {
1916    ensure_same_shape(dest.dims(), a.dims())?;
1917    ensure_same_shape(dest.dims(), b.dims())?;
1918    let validated = validate_destination_layout(dest.dims(), dest.strides())?;
1919
1920    if OpA::IS_IDENTITY && OpB::IS_IDENTITY {
1921        return mul_identity_into_raw::<InitializedOutput, D, A, B>(
1922            dest.as_mut_ptr(),
1923            dest.dims(),
1924            dest.strides(),
1925            a.ptr(),
1926            a.strides(),
1927            b.ptr(),
1928            b.strides(),
1929            validated,
1930            exec,
1931        );
1932    }
1933
1934    zip_map2_into_validated(dest, a, b, multiply_value, validated, exec)
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 + Send + Sync + 'static,
1957    A: Copy + Mul<B, Output = D> + Send + Sync + 'static,
1958    B: Copy + Send + Sync + '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    exec: Exec<'_>,
1966) -> Result<()> {
1967    ensure_same_shape(dest.dims(), a.dims())?;
1968    ensure_same_shape(dest.dims(), b.dims())?;
1969    let validated = validate_destination_layout(dest.dims(), dest.strides())?;
1970    validate_typed_no_overlap(dest, a, 0)?;
1971    validate_typed_no_overlap(dest, b, 1)?;
1972
1973    if OpA::IS_IDENTITY && OpB::IS_IDENTITY {
1974        return mul_identity_into_raw::<UninitializedOutput, D, A, B>(
1975            dest.as_mut_ptr(),
1976            dest.dims(),
1977            dest.strides(),
1978            a.ptr(),
1979            a.strides(),
1980            b.ptr(),
1981            b.strides(),
1982            validated,
1983            exec,
1984        );
1985    }
1986
1987    zip_map2_into_validated(
1988        dest,
1989        a,
1990        b,
1991        |lhs, rhs| MaybeUninit::new(multiply_value(lhs, rhs)),
1992        validated,
1993        exec,
1994    )
1995}
1996
1997fn broadcast_strides_for_axes(
1998    source_dims: &[usize],
1999    source_strides: &[isize],
2000    target_dims: &[usize],
2001    axes: &[usize],
2002) -> Result<AxisVec<isize>> {
2003    if source_dims.len() != axes.len() {
2004        return Err(StridedError::RankMismatch(source_dims.len(), axes.len()));
2005    }
2006    debug_assert_eq!(source_dims.len(), source_strides.len());
2007
2008    let mut seen = AxisVec::<bool>::new();
2009    seen.resize(target_dims.len(), false);
2010    let mut strides = AxisVec::<isize>::new();
2011    strides.resize(target_dims.len(), 0);
2012    for (src_axis, &dst_axis) in axes.iter().enumerate() {
2013        if dst_axis >= target_dims.len() {
2014            return Err(StridedError::InvalidAxis {
2015                axis: dst_axis,
2016                rank: target_dims.len(),
2017            });
2018        }
2019        if seen[dst_axis] {
2020            return Err(StridedError::InvalidAxis {
2021                axis: dst_axis,
2022                rank: target_dims.len(),
2023            });
2024        }
2025        seen[dst_axis] = true;
2026
2027        let source_dim = source_dims[src_axis];
2028        let target_dim = target_dims[dst_axis];
2029        if source_dim != target_dim && source_dim != 1 {
2030            return Err(StridedError::ShapeMismatch(
2031                source_dims.to_vec(),
2032                target_dims.to_vec(),
2033            ));
2034        }
2035        if source_dim == target_dim {
2036            strides[dst_axis] = source_strides[src_axis];
2037        }
2038    }
2039
2040    Ok(strides)
2041}
2042
2043/// A view of `view`'s data with stride-0 broadcast axes.
2044///
2045/// The metadata is computed for this call only; [`StridedView::new_unchecked`]
2046/// copies it inline, so the view still borrows only the element data.
2047fn broadcast_view_with_strides<'a, T, Op: ElementOp<T>>(
2048    view: &StridedView<'a, T, Op>,
2049    target_dims: &[usize],
2050    strides: &[isize],
2051) -> StridedView<'a, T, Op> {
2052    unsafe { StridedView::new_unchecked(view.data(), target_dims, strides, view.offset()) }
2053}
2054
2055/// Broadcasted element-wise multiplication: `dest[i] = a[i] * b[i]`.
2056///
2057/// `a_axes` and `b_axes` map each source axis to an axis of `dest`. Output axes
2058/// not referenced by a source operand are treated as stride-0 broadcast axes.
2059pub fn broadcast_mul_into<
2060    D: Copy + Send + Sync + 'static,
2061    A: Copy + Mul<B, Output = D> + Send + Sync + 'static,
2062    B: Copy + Send + Sync + 'static,
2063    OpA: ElementOp<A>,
2064    OpB: ElementOp<B>,
2065>(
2066    dest: &mut StridedViewMut<D>,
2067    a: &StridedView<A, OpA>,
2068    a_axes: &[usize],
2069    b: &StridedView<B, OpB>,
2070    b_axes: &[usize],
2071    exec: Exec<'_>,
2072) -> Result<()> {
2073    let validated = validate_destination_layout(dest.dims(), dest.strides())?;
2074    let a_strides = broadcast_strides_for_axes(a.dims(), a.strides(), dest.dims(), a_axes)?;
2075    let b_strides = broadcast_strides_for_axes(b.dims(), b.strides(), dest.dims(), b_axes)?;
2076
2077    if OpA::IS_IDENTITY && OpB::IS_IDENTITY {
2078        return mul_identity_into_raw::<InitializedOutput, D, A, B>(
2079            dest.as_mut_ptr(),
2080            dest.dims(),
2081            dest.strides(),
2082            a.ptr(),
2083            &a_strides,
2084            b.ptr(),
2085            &b_strides,
2086            validated,
2087            exec,
2088        );
2089    }
2090
2091    let a = broadcast_view_with_strides(a, dest.dims(), &a_strides);
2092    let b = broadcast_view_with_strides(b, dest.dims(), &b_strides);
2093    zip_map2_into_validated(dest, &a, &b, multiply_value, validated, exec)
2094}
2095
2096/// Broadcast and multiply into a fully overwritten uninitialized output.
2097///
2098/// Axis mappings, shapes, destination injectivity, and reachable-byte overlap
2099/// are validated before the first write. Safe Rust borrows already prevent
2100/// input/output aliasing; the explicit overlap check preserves the contract for
2101/// views produced through unsafe constructors.
2102///
2103/// `Ok(())` means every logical destination element is initialized. An error
2104/// occurs before writes. A panic during replay may leave a partially initialized
2105/// destination, which remains safe to drop as `MaybeUninit<D>`.
2106///
2107/// # Errors
2108///
2109/// Returns a typed rank, axis, or shape error for an invalid broadcast mapping,
2110/// [`StridedError::NonInjectiveOutputLayout`] for an overlapping output layout,
2111/// [`StridedError::OverlappingInputOutput`] for aliased storage, or
2112/// [`StridedError::OffsetOverflow`] when a reachable byte range is not
2113/// representable.
2114pub fn broadcast_mul_into_uninit<
2115    D: Copy + Send + Sync + 'static,
2116    A: Copy + Mul<B, Output = D> + Send + Sync + 'static,
2117    B: Copy + Send + Sync + 'static,
2118    OpA: ElementOp<A>,
2119    OpB: ElementOp<B>,
2120>(
2121    dest: &mut StridedViewMut<MaybeUninit<D>>,
2122    a: &StridedView<A, OpA>,
2123    a_axes: &[usize],
2124    b: &StridedView<B, OpB>,
2125    b_axes: &[usize],
2126    exec: Exec<'_>,
2127) -> Result<()> {
2128    let validated = validate_destination_layout(dest.dims(), dest.strides())?;
2129    let a_strides = broadcast_strides_for_axes(a.dims(), a.strides(), dest.dims(), a_axes)?;
2130    let b_strides = broadcast_strides_for_axes(b.dims(), b.strides(), dest.dims(), b_axes)?;
2131    validate_typed_no_overlap(dest, a, 0)?;
2132    validate_typed_no_overlap(dest, b, 1)?;
2133
2134    if OpA::IS_IDENTITY && OpB::IS_IDENTITY {
2135        return mul_identity_into_raw::<UninitializedOutput, D, A, B>(
2136            dest.as_mut_ptr(),
2137            dest.dims(),
2138            dest.strides(),
2139            a.ptr(),
2140            &a_strides,
2141            b.ptr(),
2142            &b_strides,
2143            validated,
2144            exec,
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        exec,
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 + Send + Sync,
2163    A: Copy + Send + Sync,
2164    B: Copy + Send + Sync,
2165    C: Copy + Send + Sync,
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 + Sync,
2175    exec: Exec<'_>,
2176) -> Result<()> {
2177    ensure_same_shape(dest.dims(), a.dims())?;
2178    ensure_same_shape(dest.dims(), b.dims())?;
2179    ensure_same_shape(dest.dims(), c.dims())?;
2180    let validated = validate_destination_layout(dest.dims(), dest.strides())?;
2181    zip_map3_into_validated(dest, a, b, c, f, validated, exec)
2182}
2183
2184pub(crate) fn zip_map3_into_validated<
2185    D: Copy + Send + Sync,
2186    A: Copy + Send + Sync,
2187    B: Copy + Send + Sync,
2188    C: Copy + Send + Sync,
2189    OpA: ElementOp<A>,
2190    OpB: ElementOp<B>,
2191    OpC: ElementOp<C>,
2192>(
2193    dest: &mut StridedViewMut<D>,
2194    a: &StridedView<A, OpA>,
2195    b: &StridedView<B, OpB>,
2196    c: &StridedView<C, OpC>,
2197    f: impl Fn(A, B, C) -> D + Sync,
2198    _validated: ValidatedDestinationLayout,
2199    exec: Exec<'_>,
2200) -> Result<()> {
2201    ensure_same_shape(dest.dims(), a.dims())?;
2202    ensure_same_shape(dest.dims(), b.dims())?;
2203    ensure_same_shape(dest.dims(), c.dims())?;
2204    let dst_ptr = dest.as_mut_ptr();
2205    let a_ptr = a.ptr();
2206    let b_ptr = b.ptr();
2207    let c_ptr = c.ptr();
2208
2209    let dst_dims = dest.dims();
2210    let dst_strides = dest.strides();
2211
2212    if sequential_contiguous_layout(
2213        dst_dims,
2214        &[dst_strides, a.strides(), b.strides(), c.strides()],
2215        &exec,
2216    )?
2217    .is_some()
2218    {
2219        let len = total_len(dst_dims)?;
2220        let dst = unsafe { std::slice::from_raw_parts_mut(dst_ptr, len) };
2221        let sa = unsafe { std::slice::from_raw_parts(a_ptr, len) };
2222        let sb = unsafe { std::slice::from_raw_parts(b_ptr, len) };
2223        let sc = unsafe { std::slice::from_raw_parts(c_ptr, len) };
2224        simd::dispatch_if_large(len, || {
2225            for (((d, &a), &b), &c) in dst.iter_mut().zip(sa).zip(sb).zip(sc) {
2226                *d = f(OpA::apply(a), OpB::apply(b), OpC::apply(c));
2227            }
2228        });
2229        return Ok(());
2230    }
2231
2232    let strides_list: [&[isize]; 4] = [dst_strides, a.strides(), b.strides(), c.strides()];
2233    let elem_size = std::mem::size_of::<D>()
2234        .max(std::mem::size_of::<A>())
2235        .max(std::mem::size_of::<B>())
2236        .max(std::mem::size_of::<C>());
2237    let total = total_len(dst_dims)?;
2238
2239    // Small tensor fast path: skip compute_order and compute_block_sizes
2240    let (fused_dims, ordered_strides, plan) = if total <= SMALL_TENSOR_THRESHOLD {
2241        build_plan_fused_small(dst_dims, &strides_list)
2242    } else {
2243        build_plan_fused(dst_dims, &strides_list, Some(0), elem_size)
2244    };
2245
2246    let total = total_len(&fused_dims)?;
2247    if is_parallel(total, &exec) {
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            exec,
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    let initial_offsets: crate::inline::OperandVec<isize> =
2289        ordered_strides.iter().map(|_| 0isize).collect();
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 + Send + Sync,
2313    A: Copy + Send + Sync,
2314    B: Copy + Send + Sync,
2315    C: Copy + Send + Sync,
2316    E: Copy + Send + Sync,
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 + Sync,
2328    exec: Exec<'_>,
2329) -> Result<()> {
2330    ensure_same_shape(dest.dims(), a.dims())?;
2331    ensure_same_shape(dest.dims(), b.dims())?;
2332    ensure_same_shape(dest.dims(), c.dims())?;
2333    ensure_same_shape(dest.dims(), e.dims())?;
2334    let validated = validate_destination_layout(dest.dims(), dest.strides())?;
2335    zip_map4_into_validated(dest, a, b, c, e, f, validated, exec)
2336}
2337
2338#[allow(clippy::too_many_arguments)]
2339pub(crate) fn zip_map4_into_validated<
2340    D: Copy + Send + Sync,
2341    A: Copy + Send + Sync,
2342    B: Copy + Send + Sync,
2343    C: Copy + Send + Sync,
2344    E: Copy + Send + Sync,
2345    OpA: ElementOp<A>,
2346    OpB: ElementOp<B>,
2347    OpC: ElementOp<C>,
2348    OpE: ElementOp<E>,
2349>(
2350    dest: &mut StridedViewMut<D>,
2351    a: &StridedView<A, OpA>,
2352    b: &StridedView<B, OpB>,
2353    c: &StridedView<C, OpC>,
2354    e: &StridedView<E, OpE>,
2355    f: impl Fn(A, B, C, E) -> D + Sync,
2356    _validated: ValidatedDestinationLayout,
2357    exec: Exec<'_>,
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        &exec,
2382    )?
2383    .is_some()
2384    {
2385        let len = total_len(dst_dims)?;
2386        let dst = unsafe { std::slice::from_raw_parts_mut(dst_ptr, len) };
2387        let sa = unsafe { std::slice::from_raw_parts(a_ptr, len) };
2388        let sb = unsafe { std::slice::from_raw_parts(b_ptr, len) };
2389        let sc = unsafe { std::slice::from_raw_parts(c_ptr, len) };
2390        let se = unsafe { std::slice::from_raw_parts(e_ptr, len) };
2391        simd::dispatch_if_large(len, || {
2392            for i in 0..len {
2393                dst[i] = f(
2394                    OpA::apply(sa[i]),
2395                    OpB::apply(sb[i]),
2396                    OpC::apply(sc[i]),
2397                    OpE::apply(se[i]),
2398                );
2399            }
2400        });
2401        return Ok(());
2402    }
2403
2404    let strides_list: [&[isize]; 5] = [
2405        dst_strides,
2406        a.strides(),
2407        b.strides(),
2408        c.strides(),
2409        e.strides(),
2410    ];
2411    let elem_size = std::mem::size_of::<D>()
2412        .max(std::mem::size_of::<A>())
2413        .max(std::mem::size_of::<B>())
2414        .max(std::mem::size_of::<C>())
2415        .max(std::mem::size_of::<E>());
2416    let total = total_len(dst_dims)?;
2417
2418    // Small tensor fast path: skip compute_order and compute_block_sizes
2419    let (fused_dims, ordered_strides, plan) = if total <= SMALL_TENSOR_THRESHOLD {
2420        build_plan_fused_small(dst_dims, &strides_list)
2421    } else {
2422        build_plan_fused(dst_dims, &strides_list, Some(0), elem_size)
2423    };
2424
2425    let total = total_len(&fused_dims)?;
2426    if is_parallel(total, &exec) {
2427        use crate::threading::SendPtr;
2428        let dst_send = SendPtr(dst_ptr);
2429        let a_send = SendPtr(a_ptr as *mut A);
2430        let b_send = SendPtr(b_ptr as *mut B);
2431        let c_send = SendPtr(c_ptr as *mut C);
2432        let e_send = SendPtr(e_ptr as *mut E);
2433
2434        let costs = compute_costs(&ordered_strides);
2435        let initial_offsets: crate::inline::OperandVec<isize> =
2436            strides_list.iter().map(|_| 0isize).collect();
2437        return mapreduce_threaded(
2438            &fused_dims,
2439            &plan.block,
2440            &ordered_strides,
2441            &initial_offsets,
2442            &costs,
2443            exec,
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], ep,
2459                                strides[4], len, &f,
2460                            )
2461                        };
2462                        Ok(())
2463                    },
2464                )
2465            },
2466        );
2467    }
2468
2469    let initial_offsets: crate::inline::OperandVec<isize> =
2470        ordered_strides.iter().map(|_| 0isize).collect();
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(test)]
2498#[path = "map_view/tests/tests.rs"]
2499mod tests;