strided_basic/

map_view.rs

//! Map operations on dynamic-rank strided views.
//!
//! These are the canonical view-based map functions, equivalent to Julia's `Base.map!`.
//! Mutable destinations must have an injective layout. Validation is
//! conservative and bounded: layouts that are injective but cannot be proven
//! by the bounded checker are rejected with
//! [`StridedError::NonInjectiveOutputLayout`].

use crate::kernel::{
    build_plan_fused, build_plan_fused_small, ensure_same_shape, for_each_inner_block_preordered,
    sequential_contiguous_layout, total_len, SMALL_TENSOR_THRESHOLD,
};
use crate::simd;
use crate::view::{StridedView, StridedViewMut};
use crate::{Result, StridedError};
use core::mem::MaybeUninit;
use std::ops::Mul;
use strided_view::{ElementOp, Identity};

use crate::cost::{for_len, KernelClass};
use crate::fuse::compute_costs;
use crate::threading::{for_each_inner_block_with_offsets, is_parallel, mapreduce_threaded};
use crate::Exec;

type AxisVec<T> = crate::inline::InlineVec<T>;

/// Largest fused inner:row aspect at which the tiled transposed-scalar kernel
/// is selected.
///
/// The cache-tiled kernel wins or ties the generic blocked plan across the
/// measured shape/dtype table (u4-tiled #210), except when the fused inner run
/// is more than 8x the row extent: then its strided destination writes lose to
/// the blocked plan on f64 (`4096x256` ~1.4x, `8192x256` ~1.3x at 1T), so the
/// detector declines and the generic path runs.
const TRANSPOSED_SCALAR_TILE_MAX_INNER_ROW_ASPECT: usize = 8;

/// Internal validation marker used by kernel-family execution adapters.
///
/// This marker is not bound to a destination. It does not authorize unchecked
/// replay with a different geometry; callers of the unsafe execution entries
/// must establish that association themselves.
#[derive(Clone, Copy, Debug, Default)]
pub struct ValidatedDestinationLayout(());

/// Validate destination injectivity and return the replay marker.
///
/// Runs the exact (allocation-free pairwise, then spilling) injectivity check,
/// so every injective layout the other kernel entry points accept is accepted
/// here. Kernel families that need the marker for more than one step (for
/// example validating shapes and then taking a specialization) should call this
/// once and pass the marker on instead of re-running the check.
#[inline]
pub fn validate_destination_layout(
    dims: &[usize],
    strides: &[isize],
) -> Result<ValidatedDestinationLayout> {
    if crate::layout_check::is_injective_layout(dims, strides) {
        Ok(ValidatedDestinationLayout(()))
    } else {
        Err(StridedError::NonInjectiveOutputLayout)
    }
}

#[inline]
/// Validate injectivity without allocating a set of offsets.
///
/// # Examples
///
/// ```
/// use strided_basic::execution::validate_destination_layout_without_alloc;
/// let marker = validate_destination_layout_without_alloc(&[2], &[1]).unwrap();
/// assert!(validate_destination_layout_without_alloc(&[2], &[0]).is_err());
/// let _ = marker;
/// ```
///
/// # Errors
/// Returns `NonInjectiveOutputLayout` for invalid or unsupported output geometry.
pub fn validate_destination_layout_without_alloc(
    dims: &[usize],
    strides: &[isize],
) -> Result<ValidatedDestinationLayout> {
    if crate::layout_check::is_injective_layout_without_alloc(dims, strides) {
        Ok(ValidatedDestinationLayout(()))
    } else {
        Err(StridedError::NonInjectiveOutputLayout)
    }
}

fn reachable_byte_range(
    ptr: usize,
    elem_size: usize,
    dims: &[usize],
    strides: &[isize],
) -> Result<Option<(usize, usize)>> {
    if elem_size == 0 || dims.contains(&0) {
        return Ok(None);
    }

    let mut min_offset = 0isize;
    let mut max_offset = 0isize;
    for (&dim, &stride) in dims.iter().zip(strides) {
        if dim <= 1 {
            continue;
        }
        let extent = isize::try_from(dim - 1).map_err(|_| StridedError::OffsetOverflow)?;
        let span = stride
            .checked_mul(extent)
            .ok_or(StridedError::OffsetOverflow)?;
        if span < 0 {
            min_offset = min_offset
                .checked_add(span)
                .ok_or(StridedError::OffsetOverflow)?;
        } else {
            max_offset = max_offset
                .checked_add(span)
                .ok_or(StridedError::OffsetOverflow)?;
        }
    }

    let ptr = ptr as i128;
    let elem_size = elem_size as i128;
    let start = ptr
        .checked_add(
            (min_offset as i128)
                .checked_mul(elem_size)
                .ok_or(StridedError::OffsetOverflow)?,
        )
        .ok_or(StridedError::OffsetOverflow)?;
    let end = ptr
        .checked_add(
            (max_offset as i128)
                .checked_mul(elem_size)
                .and_then(|offset| offset.checked_add(elem_size))
                .ok_or(StridedError::OffsetOverflow)?,
        )
        .ok_or(StridedError::OffsetOverflow)?;
    if start < 0 || end < 0 || start > usize::MAX as i128 || end > usize::MAX as i128 {
        return Err(StridedError::OffsetOverflow);
    }
    Ok(Some((start as usize, end as usize)))
}

fn validate_typed_no_overlap<D, A, Op: ElementOp<A>>(
    dest: &StridedViewMut<MaybeUninit<D>>,
    input: &StridedView<A, Op>,
    input_index: usize,
) -> Result<()> {
    let dest_range = reachable_byte_range(
        dest.ptr() as usize,
        core::mem::size_of::<D>(),
        dest.dims(),
        dest.strides(),
    )?;
    let input_range = reachable_byte_range(
        input.ptr() as usize,
        core::mem::size_of::<A>(),
        input.dims(),
        input.strides(),
    )?;
    if let (Some((dest_start, dest_end)), Some((input_start, input_end))) =
        (dest_range, input_range)
    {
        if dest_start < input_end && input_start < dest_end {
            return Err(StridedError::OverlappingInputOutput { input: input_index });
        }
    }
    Ok(())
}

// ============================================================================
// Stride-specialized inner loop helpers
//
// When all inner strides are 1 (contiguous in the innermost dimension),
// we use slice-based iteration so LLVM can auto-vectorize effectively.
// This is the Rust equivalent of Julia's @simd on the innermost loop.
// ============================================================================

// INVARIANT: validated layouts keep every dereferenced cursor in bounds.
// Pointwise leaves use wrapping advances: the unused final cursor may lie
// outside the allocation for negative or gapped strides.
/// Unary inner loop: `dest[i] = f(Op::apply(src[i]))` for `len` elements.
///
/// A vectorized kernel a map operation offers its flat runs to.
///
/// [`NoKernel`] is what the entry points without a kernel pass, so the check it introduces is
/// monomorphized away; [`SliceKernel`] adapts the closure a caller supplied.
trait MapKernel<A, D> {
    /// Writes every element of `dst` from `src`, or writes nothing and returns `false`.
    fn apply(&self, src: &[A], dst: &mut [D]) -> bool;
}

struct NoKernel;

impl<A, D> MapKernel<A, D> for NoKernel {
    #[inline(always)]
    fn apply(&self, _src: &[A], _dst: &mut [D]) -> bool {
        false
    }
}

struct SliceKernel<K>(K);

impl<A, D, K: Fn(&[A], &mut [D]) -> bool> MapKernel<A, D> for SliceKernel<K> {
    #[inline(always)]
    fn apply(&self, src: &[A], dst: &mut [D]) -> bool {
        (self.0)(src, dst)
    }
}

unsafe fn inner_loop_map1<D: Copy, A: Copy, Op: ElementOp<A>, K: MapKernel<A, D>>(
    dp: *mut D,
    ds: isize,
    sp: *const A,
    ss: isize,
    len: usize,
    f: &impl Fn(A) -> D,
    kernel: &K,
) {
    if ds == 1 && ss == 1 {
        let src = std::slice::from_raw_parts(sp, len);
        let dst = std::slice::from_raw_parts_mut(dp, len);
        if kernel.apply(src, dst) {
            return;
        }
        simd::dispatch_if_large(len, || {
            for (d, s) in dst.iter_mut().zip(src.iter()) {
                *d = f(Op::apply(*s));
            }
        });
    } else {
        let mut dp = dp;
        let mut sp = sp;
        for _ in 0..len {
            *dp = f(Op::apply(*sp));
            dp = dp.wrapping_offset(ds);
            sp = sp.wrapping_offset(ss);
        }
    }
}

/// Binary inner loop: `dest[i] = f(OpA::apply(a[i]), OpB::apply(b[i]))`.
#[inline(always)]
unsafe fn inner_loop_map2<D: Copy, A: Copy, B: Copy, OpA: ElementOp<A>, OpB: ElementOp<B>>(
    dp: *mut D,
    ds: isize,
    ap: *const A,
    a_s: isize,
    bp: *const B,
    b_s: isize,
    len: usize,
    f: &impl Fn(A, B) -> D,
) {
    if ds == 1 && a_s == 1 && b_s == 1 {
        let src_a = std::slice::from_raw_parts(ap, len);
        let src_b = std::slice::from_raw_parts(bp, len);
        let dst = std::slice::from_raw_parts_mut(dp, len);
        simd::dispatch_if_large(len, || {
            for i in 0..len {
                dst[i] = f(OpA::apply(src_a[i]), OpB::apply(src_b[i]));
            }
        });
    } else if ds == 1 && a_s == 1 && b_s == 0 {
        let src_a = std::slice::from_raw_parts(ap, len);
        let b = OpB::apply(*bp);
        let dst = std::slice::from_raw_parts_mut(dp, len);
        simd::dispatch_if_large(len, || {
            for i in 0..len {
                dst[i] = f(OpA::apply(src_a[i]), b);
            }
        });
    } else if ds == 1 && a_s == 0 && b_s == 1 {
        let a = OpA::apply(*ap);
        let src_b = std::slice::from_raw_parts(bp, len);
        let dst = std::slice::from_raw_parts_mut(dp, len);
        simd::dispatch_if_large(len, || {
            for i in 0..len {
                dst[i] = f(a, OpB::apply(src_b[i]));
            }
        });
    } else if ds == 1 && a_s == 0 && b_s == 0 {
        let a = OpA::apply(*ap);
        let b = OpB::apply(*bp);
        let dst = std::slice::from_raw_parts_mut(dp, len);
        simd::dispatch_if_large(len, || {
            for d in dst.iter_mut() {
                *d = f(a, b);
            }
        });
    } else if ds == 1 && b_s == 0 {
        let b = OpB::apply(*bp);
        let dst = std::slice::from_raw_parts_mut(dp, len);
        let mut ap = ap;
        simd::dispatch_if_large(len, || {
            for d in dst.iter_mut() {
                *d = f(OpA::apply(*ap), b);
                ap = ap.wrapping_offset(a_s);
            }
        });
    } else if ds == 1 && a_s == 0 {
        let a = OpA::apply(*ap);
        let dst = std::slice::from_raw_parts_mut(dp, len);
        let mut bp = bp;
        simd::dispatch_if_large(len, || {
            for d in dst.iter_mut() {
                *d = f(a, OpB::apply(*bp));
                bp = bp.wrapping_offset(b_s);
            }
        });
    } else {
        let mut dp = dp;
        let mut ap = ap;
        let mut bp = bp;
        for _ in 0..len {
            *dp = f(OpA::apply(*ap), OpB::apply(*bp));
            dp = dp.wrapping_offset(ds);
            ap = ap.wrapping_offset(a_s);
            bp = bp.wrapping_offset(b_s);
        }
    }
}

/// Binary multiplication inner loop for identity element ops.
trait MulOutput<D: Copy + 'static>: Copy + Send + Sync + 'static {
    type Slot: Copy + Send + Sync + 'static;

    unsafe fn write(dst: *mut Self::Slot, value: D);

    unsafe fn try_contiguous<A: 'static, B: 'static>(
        dst: *mut Self::Slot,
        len: usize,
        a: &[A],
        b: &[B],
    ) -> bool;
}

#[derive(Clone, Copy)]
struct InitializedOutput;

impl<D: Copy + Send + Sync + 'static> MulOutput<D> for InitializedOutput {
    type Slot = D;

    #[inline(always)]
    unsafe fn write(dst: *mut D, value: D) {
        unsafe { dst.write(value) };
    }

    #[inline(always)]
    unsafe fn try_contiguous<A: 'static, B: 'static>(
        dst: *mut D,
        len: usize,
        a: &[A],
        b: &[B],
    ) -> bool {
        unsafe { simd::try_mul_contiguous_ptr(dst, len, a, b) }
    }
}

#[derive(Clone, Copy)]
struct UninitializedOutput;

impl<D: Copy + Send + Sync + 'static> MulOutput<D> for UninitializedOutput {
    type Slot = MaybeUninit<D>;

    #[inline(always)]
    unsafe fn write(dst: *mut MaybeUninit<D>, value: D) {
        unsafe { dst.write(MaybeUninit::new(value)) };
    }

    #[inline(always)]
    unsafe fn try_contiguous<A: 'static, B: 'static>(
        dst: *mut MaybeUninit<D>,
        len: usize,
        a: &[A],
        b: &[B],
    ) -> bool {
        // SIMD stores write complete values directly through raw pointers; no
        // initialized reference is formed over the uninitialized destination.
        unsafe { simd::try_mul_contiguous_ptr(dst.cast::<D>(), len, a, b) }
    }
}

#[inline(always)]
fn multiply_value<A, B, D>(lhs: A, rhs: B) -> D
where
    A: Copy + Mul<B, Output = D> + 'static,
    B: Copy + 'static,
    D: Copy + 'static,
{
    use core::any::TypeId;

    if TypeId::of::<A>() == TypeId::of::<i32>()
        && TypeId::of::<B>() == TypeId::of::<i32>()
        && TypeId::of::<D>() == TypeId::of::<i32>()
    {
        // The exact TypeId checks prove identical layout and validity.
        let lhs = unsafe { *(&lhs as *const A).cast::<i32>() };
        let rhs = unsafe { *(&rhs as *const B).cast::<i32>() };
        let value = lhs.wrapping_mul(rhs);
        return unsafe { core::mem::transmute_copy(&value) };
    }
    if TypeId::of::<A>() == TypeId::of::<i64>()
        && TypeId::of::<B>() == TypeId::of::<i64>()
        && TypeId::of::<D>() == TypeId::of::<i64>()
    {
        // The exact TypeId checks prove identical layout and validity.
        let lhs = unsafe { *(&lhs as *const A).cast::<i64>() };
        let rhs = unsafe { *(&rhs as *const B).cast::<i64>() };
        let value = lhs.wrapping_mul(rhs);
        return unsafe { core::mem::transmute_copy(&value) };
    }
    lhs * rhs
}

#[inline(always)]
unsafe fn inner_loop_mul2<
    O: MulOutput<D>,
    D: Copy + 'static,
    A: Copy + Mul<B, Output = D> + 'static,
    B: Copy + 'static,
>(
    dp: *mut O::Slot,
    ds: isize,
    ap: *const A,
    a_s: isize,
    bp: *const B,
    b_s: isize,
    len: usize,
) {
    if ds == 1 && a_s == 1 && b_s == 1 {
        let src_a = std::slice::from_raw_parts(ap, len);
        let src_b = std::slice::from_raw_parts(bp, len);
        if len >= 64 && O::try_contiguous(dp, len, src_a, src_b) {
            return;
        }
        for i in 0..len {
            O::write(dp.add(i), multiply_value(src_a[i], src_b[i]));
        }
    } else if ds == 1 && a_s == 1 && b_s == 0 {
        let src_a = std::slice::from_raw_parts(ap, len);
        let b = *bp;
        for i in 0..len {
            O::write(dp.add(i), multiply_value(src_a[i], b));
        }
    } else if ds == 1 && a_s == 0 && b_s == 1 {
        let a = *ap;
        let src_b = std::slice::from_raw_parts(bp, len);
        for i in 0..len {
            O::write(dp.add(i), multiply_value(a, src_b[i]));
        }
    } else if ds == 1 && a_s == 0 && b_s == 0 {
        let a = *ap;
        let b = *bp;
        for i in 0..len {
            O::write(dp.add(i), multiply_value(a, b));
        }
    } else if ds == 1 && b_s == 0 {
        let b = *bp;
        let mut ap = ap;
        for i in 0..len {
            O::write(dp.add(i), multiply_value(*ap, b));
            ap = ap.wrapping_offset(a_s);
        }
    } else if ds == 1 && a_s == 0 {
        let a = *ap;
        let mut bp = bp;
        for i in 0..len {
            O::write(dp.add(i), multiply_value(a, *bp));
            bp = bp.wrapping_offset(b_s);
        }
    } else {
        let mut dp = dp;
        let mut ap = ap;
        let mut bp = bp;
        for _ in 0..len {
            O::write(dp, multiply_value(*ap, *bp));
            dp = dp.wrapping_offset(ds);
            ap = ap.wrapping_offset(a_s);
            bp = bp.wrapping_offset(b_s);
        }
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
struct ContiguousMulRangePlan {
    axis_order: AxisVec<usize>,
    inner_len: usize,
    inner_axis_count: usize,
    row_len: usize,
    outer_axis_start: usize,
    fast_axis: usize,
    a_fast_stride: isize,
    b_fast_stride: isize,
    a_row_stride: isize,
    b_row_stride: isize,
}

impl ContiguousMulRangePlan {
    /// Whether running the plan row by row, without cache blocking, is sound
    /// for performance.
    ///
    /// The destination is walked contiguously, so a unit-stride or broadcast
    /// (stride 0) input streams along with it. An input read at any other
    /// stride along the fast axis (a transposed operand) is swept across the
    /// whole fast axis once per row, touching a new cache line per element and
    /// again on the next row; the blocked kernel tiles that access. The one
    /// exception is the transposed-scalar tile, which blocks itself.
    fn walks_inputs_without_blocking(&self) -> bool {
        let streams = |stride: isize| stride == 0 || stride == 1;
        if streams(self.a_fast_stride) && streams(self.b_fast_stride) {
            return true;
        }
        if transposed_scalar_tile_kind(self).is_some() {
            return true;
        }
        false
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum TransposedScalarTileKind {
    RhsScalar,
    LhsScalar,
}

fn transposed_scalar_tile_kind(plan: &ContiguousMulRangePlan) -> Option<TransposedScalarTileKind> {
    if plan.inner_len > TRANSPOSED_SCALAR_TILE_MAX_INNER_ROW_ASPECT.saturating_mul(plan.row_len) {
        return None;
    }
    let row_len = isize::try_from(plan.row_len).ok()?;
    if plan.a_fast_stride == row_len
        && plan.a_row_stride == 1
        && plan.b_fast_stride == 0
        && plan.b_row_stride == 0
    {
        return Some(TransposedScalarTileKind::RhsScalar);
    }

    if plan.b_fast_stride == row_len
        && plan.b_row_stride == 1
        && plan.a_fast_stride == 0
        && plan.a_row_stride == 0
    {
        return Some(TransposedScalarTileKind::LhsScalar);
    }

    None
}

fn compact_axis_order(dims: &[usize], strides: &[isize]) -> Option<AxisVec<usize>> {
    if dims.len() != strides.len() {
        return None;
    }

    let mut active = AxisVec::<usize>::new();
    let mut inactive = AxisVec::<usize>::new();
    for (axis, (&dim, &stride)) in dims.iter().zip(strides.iter()).enumerate() {
        if stride < 0 {
            return None;
        }
        if dim > 1 {
            active.push(axis);
        } else {
            inactive.push(axis);
        }
    }

    active.sort_by(|&lhs, &rhs| strides[lhs].cmp(&strides[rhs]).then_with(|| lhs.cmp(&rhs)));

    let mut expected = 1isize;
    for &axis in &active {
        if strides[axis] != expected {
            return None;
        }
        expected = expected.saturating_mul(dims[axis] as isize);
    }

    active.extend(inactive);
    Some(active)
}

fn can_fuse_contiguous_range_axis(dim: usize, prev_stride: isize, next_stride: isize) -> bool {
    dim <= 1 || (prev_stride == 0 && next_stride == 0) || next_stride == prev_stride * dim as isize
}

fn contiguous_mul_range_plan(
    dims: &[usize],
    dst_strides: &[isize],
    a_strides: &[isize],
    b_strides: &[isize],
) -> Option<ContiguousMulRangePlan> {
    let axis_order = compact_axis_order(dims, dst_strides)?;
    if dims.is_empty() {
        return Some(ContiguousMulRangePlan {
            axis_order,
            inner_len: 1,
            inner_axis_count: 0,
            row_len: 1,
            outer_axis_start: 0,
            fast_axis: 0,
            a_fast_stride: 0,
            b_fast_stride: 0,
            a_row_stride: 0,
            b_row_stride: 0,
        });
    }

    let first_pos = axis_order
        .iter()
        .position(|&axis| dims[axis] > 1)
        .unwrap_or(0);
    let first_axis = axis_order[first_pos];
    let mut inner_len = dims[first_axis].max(1);
    let mut inner_axis_count = first_pos + 1;
    let mut prev_axis = first_axis;

    for &axis in axis_order.iter().skip(first_pos + 1) {
        if can_fuse_contiguous_range_axis(
            dims[prev_axis],
            dst_strides[prev_axis],
            dst_strides[axis],
        ) && can_fuse_contiguous_range_axis(
            dims[prev_axis],
            a_strides[prev_axis],
            a_strides[axis],
        ) && can_fuse_contiguous_range_axis(
            dims[prev_axis],
            b_strides[prev_axis],
            b_strides[axis],
        ) {
            inner_len = inner_len.checked_mul(dims[axis].max(1))?;
            inner_axis_count += 1;
            prev_axis = axis;
        } else {
            break;
        }
    }

    let row = axis_order
        .iter()
        .enumerate()
        .skip(inner_axis_count)
        .find(|&(_, &axis)| dims[axis] > 1);
    let (row_len, outer_axis_start, a_row_stride, b_row_stride) =
        if let Some((row_pos, &row_axis)) = row {
            (
                dims[row_axis],
                row_pos + 1,
                a_strides[row_axis],
                b_strides[row_axis],
            )
        } else {
            (1, axis_order.len(), 0, 0)
        };

    Some(ContiguousMulRangePlan {
        axis_order,
        inner_len,
        inner_axis_count,
        row_len,
        outer_axis_start,
        fast_axis: first_axis,
        a_fast_stride: a_strides[first_axis],
        b_fast_stride: b_strides[first_axis],
        a_row_stride,
        b_row_stride,
    })
}

struct ContiguousMulOuterCursor<'a> {
    dims: &'a [usize],
    a_strides: &'a [isize],
    b_strides: &'a [isize],
    axes: AxisVec<usize>,
    coords: AxisVec<usize>,
    a_offset: isize,
    b_offset: isize,
}

impl<'a> ContiguousMulOuterCursor<'a> {
    fn new(
        dims: &'a [usize],
        a_strides: &'a [isize],
        b_strides: &'a [isize],
        plan: &ContiguousMulRangePlan,
        outer_group: usize,
    ) -> Self {
        let axes: AxisVec<usize> = plan
            .axis_order
            .iter()
            .skip(plan.outer_axis_start)
            .copied()
            .collect();
        let mut coords = AxisVec::<usize>::with_capacity(axes.len());
        let mut rem = outer_group;
        let mut a_offset = 0isize;
        let mut b_offset = 0isize;

        for &axis in &axes {
            let dim = dims[axis].max(1);
            let coord = rem % dim;
            rem /= dim;
            coords.push(coord);
            a_offset += coord as isize * a_strides[axis];
            b_offset += coord as isize * b_strides[axis];
        }

        Self {
            dims,
            a_strides,
            b_strides,
            axes,
            coords,
            a_offset,
            b_offset,
        }
    }

    fn advance(&mut self) {
        for (i, &axis) in self.axes.iter().enumerate() {
            let dim = self.dims[axis].max(1);
            if dim <= 1 {
                continue;
            }

            // Advance only when another position along this axis follows,
            // and rewind from the last one, so offsets stay reachable.
            if self.coords[i] + 1 < dim {
                self.coords[i] += 1;
                self.a_offset += self.a_strides[axis];
                self.b_offset += self.b_strides[axis];
                break;
            }

            let last = (dim - 1) as isize;
            self.coords[i] = 0;
            self.a_offset -= last * self.a_strides[axis];
            self.b_offset -= last * self.b_strides[axis];
        }
    }
}

#[inline(always)]
unsafe fn run_contiguous_mul_row_block<
    O: MulOutput<D>,
    D: Copy + 'static,
    A: Copy + Mul<B, Output = D> + 'static,
    B: Copy + 'static,
>(
    dst_ptr: *mut O::Slot,
    a_ptr: *const A,
    b_ptr: *const B,
    plan: &ContiguousMulRangePlan,
    base_index: usize,
    total: usize,
    base_a_offset: isize,
    base_b_offset: isize,
) {
    let inner_len = plan.inner_len.max(1);
    let row_len = plan.row_len.max(1);
    let block_len = inner_len.saturating_mul(row_len);

    if total.saturating_sub(base_index) >= block_len {
        let tiled = match transposed_scalar_tile_kind(plan) {
            Some(TransposedScalarTileKind::RhsScalar) => {
                simd::try_mul_transposed_scalar_rhs_2d::<D, A, B>(
                    dst_ptr.add(base_index).cast::<D>(),
                    a_ptr.offset(base_a_offset),
                    b_ptr.offset(base_b_offset),
                    inner_len,
                    row_len,
                    plan.a_fast_stride,
                    plan.a_row_stride,
                )
            }
            Some(TransposedScalarTileKind::LhsScalar) => {
                simd::try_mul_transposed_scalar_lhs_2d::<D, A, B>(
                    dst_ptr.add(base_index).cast::<D>(),
                    a_ptr.offset(base_a_offset),
                    b_ptr.offset(base_b_offset),
                    inner_len,
                    row_len,
                    plan.b_fast_stride,
                    plan.b_row_stride,
                )
            }
            None => false,
        };
        if tiled {
            return;
        }
    }

    let mut index = base_index;
    let mut a_offset = base_a_offset;
    let mut b_offset = base_b_offset;

    for _ in 0..row_len {
        if index >= total {
            break;
        }
        let len = inner_len.min(total - index);
        inner_loop_mul2::<O, D, A, B>(
            dst_ptr.add(index),
            1,
            a_ptr.offset(a_offset),
            plan.a_fast_stride,
            b_ptr.offset(b_offset),
            plan.b_fast_stride,
            len,
        );
        index += inner_len;
        a_offset += plan.a_row_stride;
        b_offset += plan.b_row_stride;
    }
}

fn strided_offset_for_contiguous_linear_index(
    dims: &[usize],
    strides: &[isize],
    axis_order: &[usize],
    mut index: usize,
) -> isize {
    let mut offset = 0isize;
    for &axis in axis_order {
        let dim = dims[axis];
        if dim == 0 {
            return 0;
        }
        let coord = index % dim;
        index /= dim;
        offset += coord as isize * strides[axis];
    }
    offset
}

#[allow(clippy::too_many_arguments)]
fn try_contiguous_range_mul<
    O: MulOutput<D>,
    D: Copy + Send + Sync + 'static,
    A: Copy + Send + Sync + Mul<B, Output = D> + 'static,
    B: Copy + Send + Sync + 'static,
>(
    dst_ptr: *mut O::Slot,
    dims: &[usize],
    dst_strides: &[isize],
    a_ptr: *const A,
    a_strides: &[isize],
    b_ptr: *const B,
    b_strides: &[isize],
    exec: Exec<'_>,
) -> bool {
    // An overflowing element count falls back to the general path, which
    // reports it as `OffsetOverflow`.
    let Ok(total) = total_len(dims) else {
        return false;
    };
    if total == 0 {
        return true;
    }
    // One flat section over the chunks: the width is the token's decision.
    let exec = for_len(exec, KernelClass::Map, total);
    if exec.is_serial() {
        return false;
    }

    let Some(plan) = contiguous_mul_range_plan(dims, dst_strides, a_strides, b_strides) else {
        return false;
    };
    if !plan.walks_inputs_without_blocking() {
        return false;
    }

    let inner_len = plan.inner_len.max(1);
    let row_len = plan.row_len.max(1);
    let block_len = inner_len.saturating_mul(row_len).max(1);
    let outer_groups = total.div_ceil(block_len);

    let nthreads = exec.budget().get();
    if nthreads > 1 {
        use crate::threading::{parallel_for_each, SendPtr};

        let dst = SendPtr(dst_ptr);
        let a = SendPtr(a_ptr as *mut A);
        let b = SendPtr(b_ptr as *mut B);

        if outer_groups < nthreads {
            let chunk_len = total.div_ceil(nthreads);
            let nchunks = total.div_ceil(chunk_len);

            parallel_for_each(0..nchunks, exec, &|chunks| {
                for chunk in chunks {
                    let start = chunk * chunk_len;
                    let end = (start + chunk_len).min(total);
                    let mut index = start;

                    while index < end {
                        let in_inner = index % inner_len;
                        let len = (inner_len - in_inner).min(end - index);
                        let a_offset = strided_offset_for_contiguous_linear_index(
                            dims,
                            a_strides,
                            &plan.axis_order,
                            index,
                        );
                        let b_offset = strided_offset_for_contiguous_linear_index(
                            dims,
                            b_strides,
                            &plan.axis_order,
                            index,
                        );

                        unsafe {
                            inner_loop_mul2::<O, D, A, B>(
                                dst.as_ptr().add(index),
                                1,
                                a.as_const().offset(a_offset),
                                plan.a_fast_stride,
                                b.as_const().offset(b_offset),
                                plan.b_fast_stride,
                                len,
                            );
                        }
                        index += len;
                    }
                }
            });

            return true;
        }

        let groups_per_chunk = outer_groups.div_ceil(nthreads);
        let nchunks = outer_groups.div_ceil(groups_per_chunk);

        parallel_for_each(0..nchunks, exec, &|chunks| {
            for chunk in chunks {
                let group_start = chunk * groups_per_chunk;
                let group_end = (group_start + groups_per_chunk).min(outer_groups);
                let mut cursor =
                    ContiguousMulOuterCursor::new(dims, a_strides, b_strides, &plan, group_start);

                for group in group_start..group_end {
                    let index = group * block_len;
                    unsafe {
                        run_contiguous_mul_row_block::<O, D, A, B>(
                            dst.as_ptr(),
                            a.as_const(),
                            b.as_const(),
                            &plan,
                            index,
                            total,
                            cursor.a_offset,
                            cursor.b_offset,
                        );
                    }
                    cursor.advance();
                }
            }
        });

        true
    } else {
        run_contiguous_range_mul_single_thread::<O, D, A, B>(
            dst_ptr,
            dims,
            a_ptr,
            a_strides,
            b_ptr,
            b_strides,
            &plan,
            total,
            block_len,
            outer_groups,
        )
    }
}

fn run_contiguous_range_mul_single_thread<
    O: MulOutput<D>,
    D: Copy + Send + Sync + 'static,
    A: Copy + Send + Sync + Mul<B, Output = D> + 'static,
    B: Copy + Send + Sync + 'static,
>(
    dst_ptr: *mut O::Slot,
    dims: &[usize],
    a_ptr: *const A,
    a_strides: &[isize],
    b_ptr: *const B,
    b_strides: &[isize],
    plan: &ContiguousMulRangePlan,
    total: usize,
    block_len: usize,
    outer_groups: usize,
) -> bool {
    let mut cursor = ContiguousMulOuterCursor::new(dims, a_strides, b_strides, plan, 0);
    for group in 0..outer_groups {
        let index = group * block_len;
        unsafe {
            run_contiguous_mul_row_block::<O, D, A, B>(
                dst_ptr,
                a_ptr,
                b_ptr,
                plan,
                index,
                total,
                cursor.a_offset,
                cursor.b_offset,
            );
        }
        cursor.advance();
    }
    true
}

/// Ternary inner loop: `dest[i] = f(a[i], b[i], c[i])`.
#[inline(always)]
unsafe fn inner_loop_map3<
    D: Copy,
    A: Copy,
    B: Copy,
    C: Copy,
    OpA: ElementOp<A>,
    OpB: ElementOp<B>,
    OpC: ElementOp<C>,
>(
    dp: *mut D,
    ds: isize,
    ap: *const A,
    a_s: isize,
    bp: *const B,
    b_s: isize,
    cp: *const C,
    c_s: isize,
    len: usize,
    f: &impl Fn(A, B, C) -> D,
) {
    if ds == 1 && a_s == 1 && b_s == 1 && c_s == 1 {
        let src_a = std::slice::from_raw_parts(ap, len);
        let src_b = std::slice::from_raw_parts(bp, len);
        let src_c = std::slice::from_raw_parts(cp, len);
        let dst = std::slice::from_raw_parts_mut(dp, len);
        simd::dispatch_if_large(len, || {
            for (((d, &a), &b), &c) in dst.iter_mut().zip(src_a).zip(src_b).zip(src_c) {
                *d = f(OpA::apply(a), OpB::apply(b), OpC::apply(c));
            }
        });
    } else {
        let mut dp = dp;
        let mut ap = ap;
        let mut bp = bp;
        let mut cp = cp;
        for _ in 0..len {
            *dp = f(OpA::apply(*ap), OpB::apply(*bp), OpC::apply(*cp));
            dp = dp.wrapping_offset(ds);
            ap = ap.wrapping_offset(a_s);
            bp = bp.wrapping_offset(b_s);
            cp = cp.wrapping_offset(c_s);
        }
    }
}

/// Quaternary inner loop: `dest[i] = f(a[i], b[i], c[i], e[i])`.
#[inline(always)]
unsafe fn inner_loop_map4<
    D: Copy,
    A: Copy,
    B: Copy,
    C: Copy,
    E: Copy,
    OpA: ElementOp<A>,
    OpB: ElementOp<B>,
    OpC: ElementOp<C>,
    OpE: ElementOp<E>,
>(
    dp: *mut D,
    ds: isize,
    ap: *const A,
    a_s: isize,
    bp: *const B,
    b_s: isize,
    cp: *const C,
    c_s: isize,
    ep: *const E,
    e_s: isize,
    len: usize,
    f: &impl Fn(A, B, C, E) -> D,
) {
    if ds == 1 && a_s == 1 && b_s == 1 && c_s == 1 && e_s == 1 {
        let src_a = std::slice::from_raw_parts(ap, len);
        let src_b = std::slice::from_raw_parts(bp, len);
        let src_c = std::slice::from_raw_parts(cp, len);
        let src_e = std::slice::from_raw_parts(ep, len);
        let dst = std::slice::from_raw_parts_mut(dp, len);
        simd::dispatch_if_large(len, || {
            for i in 0..len {
                dst[i] = f(
                    OpA::apply(src_a[i]),
                    OpB::apply(src_b[i]),
                    OpC::apply(src_c[i]),
                    OpE::apply(src_e[i]),
                );
            }
        });
    } else {
        let mut dp = dp;
        let mut ap = ap;
        let mut bp = bp;
        let mut cp = cp;
        let mut ep = ep;
        for _ in 0..len {
            *dp = f(
                OpA::apply(*ap),
                OpB::apply(*bp),
                OpC::apply(*cp),
                OpE::apply(*ep),
            );
            dp = dp.wrapping_offset(ds);
            ap = ap.wrapping_offset(a_s);
            bp = bp.wrapping_offset(b_s);
            cp = cp.wrapping_offset(c_s);
            ep = ep.wrapping_offset(e_s);
        }
    }
}

/// An element function that carries its [`KernelClass`].
///
/// [`map_into`] and its siblings gate their parallel fork on the serial time of the
/// region, which is the element count times the throughput of the function's class.
/// A plain closure or function item is a [`KernelClass::Map`] (a few arithmetic
/// operations per element) through the blanket implementation, so existing callers
/// are unchanged. A caller whose function is expensive declares it with
/// [`Classed`], usually through [`KernelClass::with`], so a region too short to pay
/// for a fork as arithmetic is still forked when each element costs a libm call.
///
/// The trait is implemented for every `Fn(A) -> D + Sync` and for [`Classed`]; it is
/// not meant to be implemented by hand.
///
/// # Examples
/// ```
/// use strided_basic::execution::KernelClass;
/// use strided_basic::{map_into, Exec, StridedView, StridedViewMut};
///
/// let src = vec![0.0f64, 1.0, 2.0];
/// let mut out = vec![0.0f64; 3];
/// let mut dest = StridedViewMut::new(&mut out, &[3], &[1], 0).unwrap();
/// let source = StridedView::<f64>::new(&src, &[3], &[1], 0).unwrap();
/// // `exp` costs about as much as ten arithmetic operations: say so.
/// let exp = KernelClass::Transcendental.with(|x: f64| x.exp());
/// map_into(&mut dest, &source, exp, Exec::serial()).unwrap();
/// assert_eq!(out[0], 1.0);
/// ```
pub trait MapFn<A, D>: Sync {
    /// Apply the function to one element.
    fn call(&self, a: A) -> D;

    /// Serial cost class the parallel gate charges per element; `Map` by default.
    fn class(&self) -> KernelClass {
        KernelClass::Map
    }
}

impl<A, D, F: Fn(A) -> D + Sync> MapFn<A, D> for F {
    #[inline(always)]
    fn call(&self, a: A) -> D {
        self(a)
    }
}

/// An element function together with the [`KernelClass`] it is charged as.
///
/// Build one with [`Classed::new`] or [`KernelClass::with`]; see [`MapFn`].
#[derive(Debug, Clone, Copy)]
pub struct Classed<F> {
    /// The element function.
    pub f: F,
    /// The class the parallel gate charges per element.
    pub class: KernelClass,
}

impl<F> Classed<F> {
    /// Declare `f` as a function of cost class `class`.
    pub const fn new(class: KernelClass, f: F) -> Self {
        Self { f, class }
    }
}

impl KernelClass {
    /// Declare an element function of this class, for [`map_into`] and its siblings.
    pub const fn with<F>(self, f: F) -> Classed<F> {
        Classed::new(self, f)
    }
}

impl<A, D, F: Fn(A) -> D + Sync> MapFn<A, D> for Classed<F> {
    #[inline(always)]
    fn call(&self, a: A) -> D {
        (self.f)(a)
    }

    fn class(&self) -> KernelClass {
        self.class
    }
}

/// Apply a function element-wise from source to destination.
///
/// The element operation `Op` is applied lazily when reading from `src`.
/// Source and destination may have different element types.
///
/// `f` is any `Fn(A) -> D + Sync`, gated as [`KernelClass::Map`], or a function declared
/// with a class of its own (see [`MapFn`]), so an expensive element function forks sooner.
pub fn map_into<D: Copy + Send + Sync, A: Copy + Send + Sync, Op: ElementOp<A>>(
    dest: &mut StridedViewMut<D>,
    src: &StridedView<A, Op>,
    f: impl MapFn<A, D>,
    exec: Exec<'_>,
) -> Result<()> {
    let class = f.class();
    map_parts_into::<D, A, Op, NoKernel>(
        dest.as_mut_ptr(),
        dest.dims(),
        dest.strides(),
        src.ptr(),
        src.dims(),
        src.strides(),
        |a| f.call(a),
        &NoKernel,
        class,
        exec,
    )
}

/// `map_into` for a pure copy: the same traversal, gated as [`KernelClass::Copy`].
///
/// A copy has no arithmetic per element, so a fork over it pays later than over
/// a general element map; asking as `Map` forked it too early.
pub(crate) fn copy_map_into<D: Copy + Send + Sync, A: Copy + Send + Sync, Op: ElementOp<A>>(
    dest: &mut StridedViewMut<D>,
    src: &StridedView<A, Op>,
    f: impl Fn(A) -> D + Sync,
    exec: Exec<'_>,
) -> Result<()> {
    map_parts_into::<D, A, Op, NoKernel>(
        dest.as_mut_ptr(),
        dest.dims(),
        dest.strides(),
        src.ptr(),
        src.dims(),
        src.strides(),
        f,
        &NoKernel,
        KernelClass::Copy,
        exec,
    )
}

pub(crate) fn map_into_validated<D: Copy + Send + Sync, A: Copy + Send + Sync, Op: ElementOp<A>>(
    dest: &mut StridedViewMut<D>,
    src: &StridedView<A, Op>,
    f: impl MapFn<A, D>,
    validated: ValidatedDestinationLayout,
    exec: Exec<'_>,
) -> Result<()> {
    ensure_same_shape(dest.dims(), src.dims())?;
    let class = f.class();
    map_parts_into_validated::<D, A, Op, NoKernel>(
        dest.as_mut_ptr(),
        dest.dims(),
        dest.strides(),
        src.ptr(),
        src.strides(),
        |a| f.call(a),
        &NoKernel,
        validated,
        class,
        exec,
    )
}

/// Apply an element-wise function into a potentially uninitialized destination, handing a
/// vectorized kernel the flat runs it can take.
///
/// `kernel` is offered a flat, unit-stride run whenever this crate has one to give: the whole
/// buffer when the two views share a contiguous layout and the operation runs sequentially, and
/// one block at a time on the parallel path. It must write every element of the run or write
/// nothing — returning `false` makes that run fall back to `scalar`, element by element. Runs whose
/// strides are not both 1 never reach it. It is called from whichever tasks the budget of
/// `exec` selected, so it must be `Sync`.
///
/// The point of the entry point is that a caller which owns vector kernels — a SIMD math library,
/// say — reuses this crate's layout logic and this crate's threading decision instead of
/// rebuilding either. `scalar` and `kernel` must agree up to the caller's own error bound; this
/// crate only guarantees that every destination element is written exactly once.
///
/// `scalar` may carry a [`KernelClass`] (see [`MapFn`]); the gate charges the region as that
/// class, so a transcendental `scalar` forks sooner.
///
/// # Examples
/// ```
/// use std::mem::MaybeUninit;
/// use strided_basic::{map_into_uninit_with_kernel, Exec, StridedView, StridedViewMut};
///
/// let src = vec![1.0f64, 2.0, 3.0];
/// let mut out = vec![MaybeUninit::<f64>::uninit(); 3];
/// let mut dest = StridedViewMut::new(&mut out, &[3], &[1], 0).unwrap();
/// let source = StridedView::new(&src, &[3], &[1], 0).unwrap();
/// let kernel = |src: &[f64], dst: &mut [MaybeUninit<f64>]| {
///     for (slot, value) in dst.iter_mut().zip(src) {
///         slot.write(value * 2.0);
///     }
///     true
/// };
/// map_into_uninit_with_kernel(&mut dest, &source, kernel, |x| x * 2.0, Exec::serial()).unwrap();
/// // SAFETY: the kernel wrote every element.
/// let values: Vec<f64> = out.into_iter().map(|x| unsafe { x.assume_init() }).collect();
/// assert_eq!(values, [2.0, 4.0, 6.0]);
/// ```
///
/// # Errors
/// Returns [`StridedError::ShapeMismatch`] when the shapes differ and
/// [`StridedError::NonInjectiveOutputLayout`] when the destination layout cannot be proven
/// injective.
pub fn map_into_uninit_with_kernel<T: Copy + Send + Sync>(
    dest: &mut StridedViewMut<MaybeUninit<T>>,
    src: &StridedView<T>,
    kernel: impl Fn(&[T], &mut [MaybeUninit<T>]) -> bool + Sync,
    scalar: impl MapFn<T, T>,
    exec: Exec<'_>,
) -> Result<()> {
    let class = scalar.class();
    map_parts_into::<MaybeUninit<T>, T, Identity, _>(
        dest.as_mut_ptr(),
        dest.dims(),
        dest.strides(),
        src.ptr(),
        src.dims(),
        src.strides(),
        |value| MaybeUninit::new(scalar.call(value)),
        &SliceKernel(kernel),
        class,
        exec,
    )
}

/// Map a raw destination from a raw source without preparing a plan.
///
/// Only the raw-replay unit tests need this entry now; the crate's production
/// callers either go through the fused traversal (raw ops, `CopyPlan`) or
/// [`map_raw_into_validated`].
#[cfg(test)]
pub(crate) fn map_raw_into<D: Copy + Send + Sync, A: Copy + Send + Sync, Op: ElementOp<A>>(
    dest: &mut crate::RawStridedMut<'_, D>,
    src: &crate::RawStridedRef<'_, A>,
    f: impl Fn(A) -> D + Sync,
    exec: Exec<'_>,
) -> Result<()> {
    map_parts_into::<D, A, Op, NoKernel>(
        dest.as_mut_ptr(),
        dest.dims(),
        dest.strides(),
        src.ptr(),
        src.dims(),
        src.strides(),
        f,
        &NoKernel,
        KernelClass::Map,
        exec,
    )
}

pub(crate) fn map_raw_into_validated<
    D: Copy + Send + Sync,
    A: Copy + Send + Sync,
    Op: ElementOp<A>,
>(
    dest: &mut crate::RawStridedMut<'_, D>,
    src: &crate::RawStridedRef<'_, A>,
    f: impl Fn(A) -> D + Sync,
    validated: ValidatedDestinationLayout,
    exec: Exec<'_>,
) -> Result<()> {
    ensure_same_shape(dest.dims(), src.dims())?;
    map_parts_into_validated::<D, A, Op, NoKernel>(
        dest.as_mut_ptr(),
        dest.dims(),
        dest.strides(),
        src.ptr(),
        src.strides(),
        f,
        &NoKernel,
        validated,
        KernelClass::Map,
        exec,
    )
}

#[allow(clippy::too_many_arguments)]
fn map_parts_into<
    D: Copy + Send + Sync,
    A: Copy + Send + Sync,
    Op: ElementOp<A>,
    K: MapKernel<A, D> + Sync,
>(
    dst_ptr: *mut D,
    dst_dims: &[usize],
    dst_strides: &[isize],
    src_ptr: *const A,
    src_dims: &[usize],
    src_strides: &[isize],
    f: impl Fn(A) -> D + Sync,
    kernel: &K,
    class: KernelClass,
    exec: Exec<'_>,
) -> Result<()> {
    ensure_same_shape(dst_dims, src_dims)?;
    let validated = validate_destination_layout(dst_dims, dst_strides)?;
    map_parts_into_validated::<D, A, Op, K>(
        dst_ptr,
        dst_dims,
        dst_strides,
        src_ptr,
        src_strides,
        f,
        kernel,
        validated,
        class,
        exec,
    )
}

#[allow(clippy::too_many_arguments)]
/// Apply `f` (with the element operation) over a fully contiguous pair.
///
/// Kept out of line, and out of the caller's dispatch frame, so that it cannot
/// change the code shape of the blocked kernels below; the operation is moved in
/// by value because a dispatched closure that only borrows it stays scalar
/// (tensor4all/strided-rs#301).
#[inline(never)]
fn map_contiguous<D: Copy, A: Copy, Op: ElementOp<A>>(
    dst: &mut [D],
    src: &[A],
    f: impl Fn(A) -> D,
) {
    simd::dispatch_if_large(dst.len(), move || {
        for i in 0..dst.len() {
            dst[i] = f(Op::apply(src[i]));
        }
    });
}

#[allow(clippy::too_many_arguments)]
fn map_parts_into_validated<
    D: Copy + Send + Sync,
    A: Copy + Send + Sync,
    Op: ElementOp<A>,
    K: MapKernel<A, D> + Sync,
>(
    dst_ptr: *mut D,
    dst_dims: &[usize],
    dst_strides: &[isize],
    src_ptr: *const A,
    src_strides: &[isize],
    f: impl Fn(A) -> D + Sync,
    kernel: &K,
    _validated: ValidatedDestinationLayout,
    class: KernelClass,
    exec: Exec<'_>,
) -> Result<()> {
    if sequential_contiguous_layout(class, dst_dims, &[dst_strides, src_strides], &exec)?.is_some()
    {
        let len = total_len(dst_dims)?;
        // SAFETY: the contiguous-layout check proves that `dst` and `src` address `len` elements in
        // the same order starting at their base pointer, so each is exactly a flat run of `len`
        // elements.
        let dst = unsafe { std::slice::from_raw_parts_mut(dst_ptr, len) };
        let src = unsafe { std::slice::from_raw_parts(src_ptr, len) };
        if kernel.apply(src, dst) {
            return Ok(());
        }
        map_contiguous::<D, A, Op>(dst, src, f);
        return Ok(());
    }

    let strides_list: [&[isize]; 2] = [dst_strides, src_strides];
    let elem_size = std::mem::size_of::<D>().max(std::mem::size_of::<A>());
    let total = total_len(dst_dims)?;

    // Small tensor fast path: skip compute_order and compute_block_sizes
    let (fused_dims, ordered_strides, plan) = if total <= SMALL_TENSOR_THRESHOLD {
        build_plan_fused_small(dst_dims, &strides_list)
    } else {
        build_plan_fused(dst_dims, &strides_list, Some(0), elem_size)
    };

    let total = total_len(&fused_dims)?;
    if is_parallel(class, total, &exec) {
        use crate::threading::SendPtr;
        let dst_send = SendPtr(dst_ptr);
        let src_send = SendPtr(src_ptr as *mut A);

        let costs = compute_costs(&ordered_strides);
        let initial_offsets: crate::inline::OperandVec<isize> =
            strides_list.iter().map(|_| 0isize).collect();
        return mapreduce_threaded(
            &fused_dims,
            &plan.block,
            &ordered_strides,
            &initial_offsets,
            &costs,
            class,
            exec,
            &|dims, blocks, strides_list, offsets| {
                for_each_inner_block_with_offsets(
                    dims,
                    blocks,
                    strides_list,
                    offsets,
                    |offsets, len, strides| {
                        let dp = unsafe { dst_send.as_ptr().offset(offsets[0]) };
                        let sp = unsafe { src_send.as_const().offset(offsets[1]) };
                        unsafe {
                            inner_loop_map1::<D, A, Op, K>(
                                dp, strides[0], sp, strides[1], len, &f, kernel,
                            )
                        };
                        Ok(())
                    },
                )
            },
        );
    }

    let initial_offsets: crate::inline::OperandVec<isize> =
        ordered_strides.iter().map(|_| 0isize).collect();
    for_each_inner_block_preordered(
        &fused_dims,
        &plan.block,
        &ordered_strides,
        &initial_offsets,
        |offsets, len, strides| {
            let dp = unsafe { dst_ptr.offset(offsets[0]) };
            let sp = unsafe { src_ptr.offset(offsets[1]) };
            unsafe {
                inner_loop_map1::<D, A, Op, K>(dp, strides[0], sp, strides[1], len, &f, kernel)
            };
            Ok(())
        },
    )
}

/// Binary element-wise operation: `dest[i] = f(a[i], b[i])`.
///
/// Source operands `a` and `b` may have different element types from each other
/// and from `dest`. The closure `f` handles per-element type conversion.
pub fn zip_map2_into<
    D: Copy + Send + Sync,
    A: Copy + Send + Sync,
    B: Copy + Send + Sync,
    OpA: ElementOp<A>,
    OpB: ElementOp<B>,
>(
    dest: &mut StridedViewMut<D>,
    a: &StridedView<A, OpA>,
    b: &StridedView<B, OpB>,
    f: impl Fn(A, B) -> D + Sync,
    exec: Exec<'_>,
) -> Result<()> {
    zip_map2_parts_into::<D, A, B, OpA, OpB>(
        dest.as_mut_ptr(),
        dest.dims(),
        dest.strides(),
        a.ptr(),
        a.dims(),
        a.strides(),
        b.ptr(),
        b.dims(),
        b.strides(),
        f,
        exec,
    )
}

pub(crate) fn zip_map2_into_validated<
    D: Copy + Send + Sync,
    A: Copy + Send + Sync,
    B: Copy + Send + Sync,
    OpA: ElementOp<A>,
    OpB: ElementOp<B>,
>(
    dest: &mut StridedViewMut<D>,
    a: &StridedView<A, OpA>,
    b: &StridedView<B, OpB>,
    f: impl Fn(A, B) -> D + Sync,
    validated: ValidatedDestinationLayout,
    exec: Exec<'_>,
) -> Result<()> {
    zip_map2_parts_into_validated::<D, A, B, OpA, OpB>(
        dest.as_mut_ptr(),
        dest.dims(),
        dest.strides(),
        a.ptr(),
        a.strides(),
        b.ptr(),
        b.strides(),
        f,
        validated,
        exec,
    )
}

/// Runtime comparison selected once before entering the element loop.
#[non_exhaustive]
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum CompareOp {
    Eq,
    Lt,
    Le,
    Gt,
    Ge,
}

/// Compare two ordered views elementwise into a Boolean destination.
///
/// Unlike embedding a runtime operation match inside a [`zip_map2_into`]
/// closure, this entry point selects a fixed comparison before traversal.
///
/// # Errors
///
/// Returns [`StridedError::ShapeMismatch`] when the source and destination
/// shapes differ, or [`StridedError::NonInjectiveOutputLayout`] when distinct
/// logical destination elements may overlap.
pub fn compare_into<T, OpA, OpB>(
    dest: &mut StridedViewMut<bool>,
    a: &StridedView<T, OpA>,
    b: &StridedView<T, OpB>,
    op: CompareOp,
    exec: Exec<'_>,
) -> Result<()>
where
    T: Copy + Send + Sync + PartialOrd,
    OpA: ElementOp<T>,
    OpB: ElementOp<T>,
{
    match op {
        CompareOp::Eq => zip_map2_into(dest, a, b, |lhs, rhs| lhs == rhs, exec),
        CompareOp::Lt => zip_map2_into(dest, a, b, |lhs, rhs| lhs < rhs, exec),
        CompareOp::Le => zip_map2_into(dest, a, b, |lhs, rhs| lhs <= rhs, exec),
        CompareOp::Gt => zip_map2_into(dest, a, b, |lhs, rhs| lhs > rhs, exec),
        CompareOp::Ge => zip_map2_into(dest, a, b, |lhs, rhs| lhs >= rhs, exec),
    }
}

/// Compare two views into a fully overwritten uninitialized Boolean output.
///
/// Dtype-independent shape, destination-injectivity, and reachable-byte
/// overlap validation completes before the first write. Safe Rust borrows
/// already prevent input/output aliasing; the explicit overlap check preserves
/// the contract for views produced through unsafe constructors.
///
/// `Ok(())` means every logical destination element is initialized. An error
/// occurs before writes. A panic during replay may leave a partially initialized
/// destination, which remains safe to drop as `MaybeUninit<bool>`.
///
/// # Errors
///
/// Returns [`StridedError::ShapeMismatch`] for unequal shapes,
/// [`StridedError::NonInjectiveOutputLayout`] for an overlapping output layout,
/// [`StridedError::OverlappingInputOutput`] for aliased storage, or
/// [`StridedError::OffsetOverflow`] when a reachable byte range is not
/// representable.
pub fn compare_into_uninit<T, OpA, OpB>(
    dest: &mut StridedViewMut<MaybeUninit<bool>>,
    a: &StridedView<T, OpA>,
    b: &StridedView<T, OpB>,
    op: CompareOp,
    exec: Exec<'_>,
) -> Result<()>
where
    T: Copy + Send + Sync + PartialOrd,
    OpA: ElementOp<T>,
    OpB: ElementOp<T>,
{
    ensure_same_shape(dest.dims(), a.dims())?;
    ensure_same_shape(dest.dims(), b.dims())?;
    let validated = validate_destination_layout(dest.dims(), dest.strides())?;
    validate_typed_no_overlap(dest, a, 0)?;
    validate_typed_no_overlap(dest, b, 1)?;
    match op {
        CompareOp::Eq => zip_map2_into_validated(
            dest,
            a,
            b,
            |lhs, rhs| MaybeUninit::new(lhs == rhs),
            validated,
            exec,
        ),
        CompareOp::Lt => zip_map2_into_validated(
            dest,
            a,
            b,
            |lhs, rhs| MaybeUninit::new(lhs < rhs),
            validated,
            exec,
        ),
        CompareOp::Le => zip_map2_into_validated(
            dest,
            a,
            b,
            |lhs, rhs| MaybeUninit::new(lhs <= rhs),
            validated,
            exec,
        ),
        CompareOp::Gt => zip_map2_into_validated(
            dest,
            a,
            b,
            |lhs, rhs| MaybeUninit::new(lhs > rhs),
            validated,
            exec,
        ),
        CompareOp::Ge => zip_map2_into_validated(
            dest,
            a,
            b,
            |lhs, rhs| MaybeUninit::new(lhs >= rhs),
            validated,
            exec,
        ),
    }
}

pub(crate) fn zip_map2_raw_into_validated<
    D: Copy + Send + Sync,
    A: Copy + Send + Sync,
    B: Copy + Send + Sync,
    OpA: ElementOp<A>,
    OpB: ElementOp<B>,
>(
    dest: &mut crate::RawStridedMut<'_, D>,
    a: &crate::RawStridedRef<'_, A>,
    b: &crate::RawStridedRef<'_, B>,
    f: impl Fn(A, B) -> D + Sync,
    validated: ValidatedDestinationLayout,
    exec: Exec<'_>,
) -> Result<()> {
    ensure_same_shape(dest.dims(), a.dims())?;
    ensure_same_shape(dest.dims(), b.dims())?;
    zip_map2_parts_into_validated::<D, A, B, OpA, OpB>(
        dest.as_mut_ptr(),
        dest.dims(),
        dest.strides(),
        a.ptr(),
        a.strides(),
        b.ptr(),
        b.strides(),
        f,
        validated,
        exec,
    )
}

#[allow(clippy::too_many_arguments)]
fn zip_map2_parts_into<
    D: Copy + Send + Sync,
    A: Copy + Send + Sync,
    B: Copy + Send + Sync,
    OpA: ElementOp<A>,
    OpB: ElementOp<B>,
>(
    dst_ptr: *mut D,
    dst_dims: &[usize],
    dst_strides: &[isize],
    a_ptr: *const A,
    a_dims: &[usize],
    a_strides: &[isize],
    b_ptr: *const B,
    b_dims: &[usize],
    b_strides: &[isize],
    f: impl Fn(A, B) -> D + Sync,
    exec: Exec<'_>,
) -> Result<()> {
    ensure_same_shape(dst_dims, a_dims)?;
    ensure_same_shape(dst_dims, b_dims)?;
    let validated = validate_destination_layout(dst_dims, dst_strides)?;
    zip_map2_parts_into_validated::<D, A, B, OpA, OpB>(
        dst_ptr,
        dst_dims,
        dst_strides,
        a_ptr,
        a_strides,
        b_ptr,
        b_strides,
        f,
        validated,
        exec,
    )
}

#[allow(clippy::too_many_arguments)]
fn zip_map2_parts_into_validated<
    D: Copy + Send + Sync,
    A: Copy + Send + Sync,
    B: Copy + Send + Sync,
    OpA: ElementOp<A>,
    OpB: ElementOp<B>,
>(
    dst_ptr: *mut D,
    dst_dims: &[usize],
    dst_strides: &[isize],
    a_ptr: *const A,
    a_strides: &[isize],
    b_ptr: *const B,
    b_strides: &[isize],
    f: impl Fn(A, B) -> D + Sync,
    _validated: ValidatedDestinationLayout,
    exec: Exec<'_>,
) -> Result<()> {
    if sequential_contiguous_layout(
        KernelClass::Map,
        dst_dims,
        &[dst_strides, a_strides, b_strides],
        &exec,
    )?
    .is_some()
    {
        let len = total_len(dst_dims)?;
        let dst = unsafe { std::slice::from_raw_parts_mut(dst_ptr, len) };
        let sa = unsafe { std::slice::from_raw_parts(a_ptr, len) };
        let sb = unsafe { std::slice::from_raw_parts(b_ptr, len) };
        simd::dispatch_if_large(len, || {
            for i in 0..len {
                dst[i] = f(OpA::apply(sa[i]), OpB::apply(sb[i]));
            }
        });
        return Ok(());
    }

    let strides_list: [&[isize]; 3] = [dst_strides, a_strides, b_strides];
    let elem_size = std::mem::size_of::<D>()
        .max(std::mem::size_of::<A>())
        .max(std::mem::size_of::<B>());
    let total = total_len(dst_dims)?;

    // Small tensor fast path: skip compute_order and compute_block_sizes
    let (fused_dims, ordered_strides, plan) = if total <= SMALL_TENSOR_THRESHOLD {
        build_plan_fused_small(dst_dims, &strides_list)
    } else {
        build_plan_fused(dst_dims, &strides_list, Some(0), elem_size)
    };

    let total = total_len(&fused_dims)?;
    if is_parallel(KernelClass::Map, total, &exec) {
        use crate::threading::SendPtr;
        let dst_send = SendPtr(dst_ptr);
        let a_send = SendPtr(a_ptr as *mut A);
        let b_send = SendPtr(b_ptr as *mut B);

        let costs = compute_costs(&ordered_strides);
        let initial_offsets: crate::inline::OperandVec<isize> =
            strides_list.iter().map(|_| 0isize).collect();
        return mapreduce_threaded(
            &fused_dims,
            &plan.block,
            &ordered_strides,
            &initial_offsets,
            &costs,
            KernelClass::Map,
            exec,
            &|dims, blocks, strides_list, offsets| {
                for_each_inner_block_with_offsets(
                    dims,
                    blocks,
                    strides_list,
                    offsets,
                    |offsets, len, strides| {
                        let dp = unsafe { dst_send.as_ptr().offset(offsets[0]) };
                        let ap = unsafe { a_send.as_const().offset(offsets[1]) };
                        let bp = unsafe { b_send.as_const().offset(offsets[2]) };
                        unsafe {
                            inner_loop_map2::<D, A, B, OpA, OpB>(
                                dp, strides[0], ap, strides[1], bp, strides[2], len, &f,
                            )
                        };
                        Ok(())
                    },
                )
            },
        );
    }

    let initial_offsets: crate::inline::OperandVec<isize> =
        ordered_strides.iter().map(|_| 0isize).collect();
    for_each_inner_block_preordered(
        &fused_dims,
        &plan.block,
        &ordered_strides,
        &initial_offsets,
        |offsets, len, strides| {
            let dp = unsafe { dst_ptr.offset(offsets[0]) };
            let ap = unsafe { a_ptr.offset(offsets[1]) };
            let bp = unsafe { b_ptr.offset(offsets[2]) };
            unsafe {
                inner_loop_map2::<D, A, B, OpA, OpB>(
                    dp, strides[0], ap, strides[1], bp, strides[2], len, &f,
                )
            };
            Ok(())
        },
    )
}

#[allow(clippy::too_many_arguments)]
fn mul_identity_into_raw<
    O: MulOutput<D>,
    D: Copy + Send + Sync + 'static,
    A: Copy + Send + Sync + Mul<B, Output = D> + 'static,
    B: Copy + Send + Sync + 'static,
>(
    dst_ptr: *mut O::Slot,
    dst_dims: &[usize],
    dst_strides: &[isize],
    a_ptr: *const A,
    a_strides: &[isize],
    b_ptr: *const B,
    b_strides: &[isize],
    _validated: ValidatedDestinationLayout,
    mut exec: Exec<'_>,
) -> Result<()> {
    debug_assert_eq!(dst_dims.len(), a_strides.len());
    debug_assert_eq!(dst_dims.len(), b_strides.len());

    if sequential_contiguous_layout(
        KernelClass::Map,
        dst_dims,
        &[dst_strides, a_strides, b_strides],
        &exec,
    )?
    .is_some()
    {
        let len = total_len(dst_dims)?;
        let sa = unsafe { std::slice::from_raw_parts(a_ptr, len) };
        let sb = unsafe { std::slice::from_raw_parts(b_ptr, len) };
        if unsafe { O::try_contiguous(dst_ptr, len, sa, sb) } {
            return Ok(());
        }
        for i in 0..len {
            unsafe { O::write(dst_ptr.add(i), multiply_value(sa[i], sb[i])) };
        }
        return Ok(());
    }

    let strides_list: [&[isize]; 3] = [dst_strides, a_strides, b_strides];
    let elem_size = std::mem::size_of::<D>()
        .max(std::mem::size_of::<A>())
        .max(std::mem::size_of::<B>());
    let total = total_len(dst_dims)?;

    if try_contiguous_range_mul::<O, D, A, B>(
        dst_ptr,
        dst_dims,
        dst_strides,
        a_ptr,
        a_strides,
        b_ptr,
        b_strides,
        exec.reborrow(),
    ) {
        return Ok(());
    }

    let (fused_dims, ordered_strides, plan) = if total <= SMALL_TENSOR_THRESHOLD {
        build_plan_fused_small(dst_dims, &strides_list)
    } else {
        build_plan_fused(dst_dims, &strides_list, Some(0), elem_size)
    };

    let total = total_len(&fused_dims)?;
    if is_parallel(KernelClass::Map, total, &exec) {
        use crate::threading::SendPtr;
        let dst_send = SendPtr(dst_ptr);
        let a_send = SendPtr(a_ptr as *mut A);
        let b_send = SendPtr(b_ptr as *mut B);

        let costs = compute_costs(&ordered_strides);
        let initial_offsets: crate::inline::OperandVec<isize> =
            strides_list.iter().map(|_| 0isize).collect();
        return mapreduce_threaded(
            &fused_dims,
            &plan.block,
            &ordered_strides,
            &initial_offsets,
            &costs,
            KernelClass::Map,
            exec,
            &|dims, blocks, strides_list, offsets| {
                for_each_inner_block_with_offsets(
                    dims,
                    blocks,
                    strides_list,
                    offsets,
                    |offsets, len, strides| {
                        let dp = unsafe { dst_send.as_ptr().offset(offsets[0]) };
                        let ap = unsafe { a_send.as_const().offset(offsets[1]) };
                        let bp = unsafe { b_send.as_const().offset(offsets[2]) };
                        unsafe {
                            inner_loop_mul2::<O, D, A, B>(
                                dp, strides[0], ap, strides[1], bp, strides[2], len,
                            )
                        };
                        Ok(())
                    },
                )
            },
        );
    }

    let initial_offsets: crate::inline::OperandVec<isize> =
        ordered_strides.iter().map(|_| 0isize).collect();
    for_each_inner_block_preordered(
        &fused_dims,
        &plan.block,
        &ordered_strides,
        &initial_offsets,
        |offsets, len, strides| {
            let dp = unsafe { dst_ptr.offset(offsets[0]) };
            let ap = unsafe { a_ptr.offset(offsets[1]) };
            let bp = unsafe { b_ptr.offset(offsets[2]) };
            unsafe {
                inner_loop_mul2::<O, D, A, B>(dp, strides[0], ap, strides[1], bp, strides[2], len)
            };
            Ok(())
        },
    )
}

/// Element-wise multiplication: `dest[i] = a[i] * b[i]`.
///
/// All views must have the same shape. Broadcast operands should be represented
/// as stride-0 views before calling this function.
pub fn mul_into<
    D: Copy + Send + Sync + 'static,
    A: Copy + Mul<B, Output = D> + Send + Sync + 'static,
    B: Copy + Send + Sync + 'static,
    OpA: ElementOp<A>,
    OpB: ElementOp<B>,
>(
    dest: &mut StridedViewMut<D>,
    a: &StridedView<A, OpA>,
    b: &StridedView<B, OpB>,
    exec: Exec<'_>,
) -> Result<()> {
    ensure_same_shape(dest.dims(), a.dims())?;
    ensure_same_shape(dest.dims(), b.dims())?;
    let validated = validate_destination_layout(dest.dims(), dest.strides())?;

    if OpA::IS_IDENTITY && OpB::IS_IDENTITY {
        return mul_identity_into_raw::<InitializedOutput, D, A, B>(
            dest.as_mut_ptr(),
            dest.dims(),
            dest.strides(),
            a.ptr(),
            a.strides(),
            b.ptr(),
            b.strides(),
            validated,
            exec,
        );
    }

    zip_map2_into_validated(dest, a, b, multiply_value, validated, exec)
}

/// Multiply two views into a fully overwritten uninitialized output.
///
/// Shape, destination-injectivity, and reachable-byte overlap validation
/// completes before the first write. Safe Rust borrows already prevent
/// input/output aliasing; the explicit overlap check preserves the contract for
/// views produced through unsafe constructors.
///
/// `Ok(())` means every logical destination element is initialized. An error
/// occurs before writes. A panic during replay may leave a partially initialized
/// destination, which remains safe to drop as `MaybeUninit<D>`.
///
/// # Errors
///
/// Returns [`StridedError::ShapeMismatch`] for unequal shapes,
/// [`StridedError::NonInjectiveOutputLayout`] for an overlapping output layout,
/// [`StridedError::OverlappingInputOutput`] for aliased storage, or
/// [`StridedError::OffsetOverflow`] when a reachable byte range is not
/// representable.
pub fn mul_into_uninit<
    D: Copy + Send + Sync + 'static,
    A: Copy + Mul<B, Output = D> + Send + Sync + 'static,
    B: Copy + Send + Sync + 'static,
    OpA: ElementOp<A>,
    OpB: ElementOp<B>,
>(
    dest: &mut StridedViewMut<MaybeUninit<D>>,
    a: &StridedView<A, OpA>,
    b: &StridedView<B, OpB>,
    exec: Exec<'_>,
) -> Result<()> {
    ensure_same_shape(dest.dims(), a.dims())?;
    ensure_same_shape(dest.dims(), b.dims())?;
    let validated = validate_destination_layout(dest.dims(), dest.strides())?;
    validate_typed_no_overlap(dest, a, 0)?;
    validate_typed_no_overlap(dest, b, 1)?;

    if OpA::IS_IDENTITY && OpB::IS_IDENTITY {
        return mul_identity_into_raw::<UninitializedOutput, D, A, B>(
            dest.as_mut_ptr(),
            dest.dims(),
            dest.strides(),
            a.ptr(),
            a.strides(),
            b.ptr(),
            b.strides(),
            validated,
            exec,
        );
    }

    zip_map2_into_validated(
        dest,
        a,
        b,
        |lhs, rhs| MaybeUninit::new(multiply_value(lhs, rhs)),
        validated,
        exec,
    )
}

fn broadcast_strides_for_axes(
    source_dims: &[usize],
    source_strides: &[isize],
    target_dims: &[usize],
    axes: &[usize],
) -> Result<AxisVec<isize>> {
    if source_dims.len() != axes.len() {
        return Err(StridedError::RankMismatch(source_dims.len(), axes.len()));
    }
    debug_assert_eq!(source_dims.len(), source_strides.len());

    let mut seen = AxisVec::<bool>::new();
    seen.resize(target_dims.len(), false);
    let mut strides = AxisVec::<isize>::new();
    strides.resize(target_dims.len(), 0);
    for (src_axis, &dst_axis) in axes.iter().enumerate() {
        if dst_axis >= target_dims.len() {
            return Err(StridedError::InvalidAxis {
                axis: dst_axis,
                rank: target_dims.len(),
            });
        }
        if seen[dst_axis] {
            return Err(StridedError::InvalidAxis {
                axis: dst_axis,
                rank: target_dims.len(),
            });
        }
        seen[dst_axis] = true;

        let source_dim = source_dims[src_axis];
        let target_dim = target_dims[dst_axis];
        if source_dim != target_dim && source_dim != 1 {
            return Err(StridedError::ShapeMismatch(
                source_dims.to_vec(),
                target_dims.to_vec(),
            ));
        }
        if source_dim == target_dim {
            strides[dst_axis] = source_strides[src_axis];
        }
    }

    Ok(strides)
}

/// A view of `view`'s data with stride-0 broadcast axes.
///
/// The metadata is computed for this call only; [`StridedView::new_unchecked`]
/// copies it inline, so the view still borrows only the element data.
fn broadcast_view_with_strides<'a, T, Op: ElementOp<T>>(
    view: &StridedView<'a, T, Op>,
    target_dims: &[usize],
    strides: &[isize],
) -> StridedView<'a, T, Op> {
    unsafe { StridedView::new_unchecked(view.data(), target_dims, strides, view.offset()) }
}

/// Broadcasted element-wise multiplication: `dest[i] = a[i] * b[i]`.
///
/// `a_axes` and `b_axes` map each source axis to an axis of `dest`. Output axes
/// not referenced by a source operand are treated as stride-0 broadcast axes.
pub fn broadcast_mul_into<
    D: Copy + Send + Sync + 'static,
    A: Copy + Mul<B, Output = D> + Send + Sync + 'static,
    B: Copy + Send + Sync + 'static,
    OpA: ElementOp<A>,
    OpB: ElementOp<B>,
>(
    dest: &mut StridedViewMut<D>,
    a: &StridedView<A, OpA>,
    a_axes: &[usize],
    b: &StridedView<B, OpB>,
    b_axes: &[usize],
    exec: Exec<'_>,
) -> Result<()> {
    let validated = validate_destination_layout(dest.dims(), dest.strides())?;
    let a_strides = broadcast_strides_for_axes(a.dims(), a.strides(), dest.dims(), a_axes)?;
    let b_strides = broadcast_strides_for_axes(b.dims(), b.strides(), dest.dims(), b_axes)?;

    if OpA::IS_IDENTITY && OpB::IS_IDENTITY {
        return mul_identity_into_raw::<InitializedOutput, D, A, B>(
            dest.as_mut_ptr(),
            dest.dims(),
            dest.strides(),
            a.ptr(),
            &a_strides,
            b.ptr(),
            &b_strides,
            validated,
            exec,
        );
    }

    let a = broadcast_view_with_strides(a, dest.dims(), &a_strides);
    let b = broadcast_view_with_strides(b, dest.dims(), &b_strides);
    zip_map2_into_validated(dest, &a, &b, multiply_value, validated, exec)
}

/// Broadcast and multiply into a fully overwritten uninitialized output.
///
/// Axis mappings, shapes, destination injectivity, and reachable-byte overlap
/// are validated before the first write. Safe Rust borrows already prevent
/// input/output aliasing; the explicit overlap check preserves the contract for
/// views produced through unsafe constructors.
///
/// `Ok(())` means every logical destination element is initialized. An error
/// occurs before writes. A panic during replay may leave a partially initialized
/// destination, which remains safe to drop as `MaybeUninit<D>`.
///
/// # Errors
///
/// Returns a typed rank, axis, or shape error for an invalid broadcast mapping,
/// [`StridedError::NonInjectiveOutputLayout`] for an overlapping output layout,
/// [`StridedError::OverlappingInputOutput`] for aliased storage, or
/// [`StridedError::OffsetOverflow`] when a reachable byte range is not
/// representable.
pub fn broadcast_mul_into_uninit<
    D: Copy + Send + Sync + 'static,
    A: Copy + Mul<B, Output = D> + Send + Sync + 'static,
    B: Copy + Send + Sync + 'static,
    OpA: ElementOp<A>,
    OpB: ElementOp<B>,
>(
    dest: &mut StridedViewMut<MaybeUninit<D>>,
    a: &StridedView<A, OpA>,
    a_axes: &[usize],
    b: &StridedView<B, OpB>,
    b_axes: &[usize],
    exec: Exec<'_>,
) -> Result<()> {
    let validated = validate_destination_layout(dest.dims(), dest.strides())?;
    let a_strides = broadcast_strides_for_axes(a.dims(), a.strides(), dest.dims(), a_axes)?;
    let b_strides = broadcast_strides_for_axes(b.dims(), b.strides(), dest.dims(), b_axes)?;
    validate_typed_no_overlap(dest, a, 0)?;
    validate_typed_no_overlap(dest, b, 1)?;

    if OpA::IS_IDENTITY && OpB::IS_IDENTITY {
        return mul_identity_into_raw::<UninitializedOutput, D, A, B>(
            dest.as_mut_ptr(),
            dest.dims(),
            dest.strides(),
            a.ptr(),
            &a_strides,
            b.ptr(),
            &b_strides,
            validated,
            exec,
        );
    }

    let a = broadcast_view_with_strides(a, dest.dims(), &a_strides);
    let b = broadcast_view_with_strides(b, dest.dims(), &b_strides);
    zip_map2_into_validated(
        dest,
        &a,
        &b,
        |lhs, rhs| MaybeUninit::new(multiply_value(lhs, rhs)),
        validated,
        exec,
    )
}

/// Ternary element-wise operation: `dest[i] = f(a[i], b[i], c[i])`.
pub fn zip_map3_into<
    D: Copy + Send + Sync,
    A: Copy + Send + Sync,
    B: Copy + Send + Sync,
    C: Copy + Send + Sync,
    OpA: ElementOp<A>,
    OpB: ElementOp<B>,
    OpC: ElementOp<C>,
>(
    dest: &mut StridedViewMut<D>,
    a: &StridedView<A, OpA>,
    b: &StridedView<B, OpB>,
    c: &StridedView<C, OpC>,
    f: impl Fn(A, B, C) -> D + Sync,
    exec: Exec<'_>,
) -> Result<()> {
    ensure_same_shape(dest.dims(), a.dims())?;
    ensure_same_shape(dest.dims(), b.dims())?;
    ensure_same_shape(dest.dims(), c.dims())?;
    let validated = validate_destination_layout(dest.dims(), dest.strides())?;
    zip_map3_into_validated(dest, a, b, c, f, validated, exec)
}

pub(crate) fn zip_map3_into_validated<
    D: Copy + Send + Sync,
    A: Copy + Send + Sync,
    B: Copy + Send + Sync,
    C: Copy + Send + Sync,
    OpA: ElementOp<A>,
    OpB: ElementOp<B>,
    OpC: ElementOp<C>,
>(
    dest: &mut StridedViewMut<D>,
    a: &StridedView<A, OpA>,
    b: &StridedView<B, OpB>,
    c: &StridedView<C, OpC>,
    f: impl Fn(A, B, C) -> D + Sync,
    _validated: ValidatedDestinationLayout,
    exec: Exec<'_>,
) -> Result<()> {
    ensure_same_shape(dest.dims(), a.dims())?;
    ensure_same_shape(dest.dims(), b.dims())?;
    ensure_same_shape(dest.dims(), c.dims())?;
    let dst_ptr = dest.as_mut_ptr();
    let a_ptr = a.ptr();
    let b_ptr = b.ptr();
    let c_ptr = c.ptr();

    let dst_dims = dest.dims();
    let dst_strides = dest.strides();

    if sequential_contiguous_layout(
        KernelClass::Map,
        dst_dims,
        &[dst_strides, a.strides(), b.strides(), c.strides()],
        &exec,
    )?
    .is_some()
    {
        let len = total_len(dst_dims)?;
        let dst = unsafe { std::slice::from_raw_parts_mut(dst_ptr, len) };
        let sa = unsafe { std::slice::from_raw_parts(a_ptr, len) };
        let sb = unsafe { std::slice::from_raw_parts(b_ptr, len) };
        let sc = unsafe { std::slice::from_raw_parts(c_ptr, len) };
        simd::dispatch_if_large(len, || {
            for (((d, &a), &b), &c) in dst.iter_mut().zip(sa).zip(sb).zip(sc) {
                *d = f(OpA::apply(a), OpB::apply(b), OpC::apply(c));
            }
        });
        return Ok(());
    }

    let strides_list: [&[isize]; 4] = [dst_strides, a.strides(), b.strides(), c.strides()];
    let elem_size = std::mem::size_of::<D>()
        .max(std::mem::size_of::<A>())
        .max(std::mem::size_of::<B>())
        .max(std::mem::size_of::<C>());
    let total = total_len(dst_dims)?;

    // Small tensor fast path: skip compute_order and compute_block_sizes
    let (fused_dims, ordered_strides, plan) = if total <= SMALL_TENSOR_THRESHOLD {
        build_plan_fused_small(dst_dims, &strides_list)
    } else {
        build_plan_fused(dst_dims, &strides_list, Some(0), elem_size)
    };

    let total = total_len(&fused_dims)?;
    if is_parallel(KernelClass::Map, total, &exec) {
        use crate::threading::SendPtr;
        let dst_send = SendPtr(dst_ptr);
        let a_send = SendPtr(a_ptr as *mut A);
        let b_send = SendPtr(b_ptr as *mut B);
        let c_send = SendPtr(c_ptr as *mut C);

        let costs = compute_costs(&ordered_strides);
        let initial_offsets: crate::inline::OperandVec<isize> =
            strides_list.iter().map(|_| 0isize).collect();
        return mapreduce_threaded(
            &fused_dims,
            &plan.block,
            &ordered_strides,
            &initial_offsets,
            &costs,
            KernelClass::Map,
            exec,
            &|dims, blocks, strides_list, offsets| {
                for_each_inner_block_with_offsets(
                    dims,
                    blocks,
                    strides_list,
                    offsets,
                    |offsets, len, strides| {
                        let dp = unsafe { dst_send.as_ptr().offset(offsets[0]) };
                        let ap = unsafe { a_send.as_const().offset(offsets[1]) };
                        let bp = unsafe { b_send.as_const().offset(offsets[2]) };
                        let cp = unsafe { c_send.as_const().offset(offsets[3]) };
                        unsafe {
                            inner_loop_map3::<D, A, B, C, OpA, OpB, OpC>(
                                dp, strides[0], ap, strides[1], bp, strides[2], cp, strides[3],
                                len, &f,
                            )
                        };
                        Ok(())
                    },
                )
            },
        );
    }

    let initial_offsets: crate::inline::OperandVec<isize> =
        ordered_strides.iter().map(|_| 0isize).collect();
    for_each_inner_block_preordered(
        &fused_dims,
        &plan.block,
        &ordered_strides,
        &initial_offsets,
        |offsets, len, strides| {
            let dp = unsafe { dst_ptr.offset(offsets[0]) };
            let ap = unsafe { a_ptr.offset(offsets[1]) };
            let bp = unsafe { b_ptr.offset(offsets[2]) };
            let cp = unsafe { c_ptr.offset(offsets[3]) };
            unsafe {
                inner_loop_map3::<D, A, B, C, OpA, OpB, OpC>(
                    dp, strides[0], ap, strides[1], bp, strides[2], cp, strides[3], len, &f,
                )
            };
            Ok(())
        },
    )
}

/// Quaternary element-wise operation: `dest[i] = f(a[i], b[i], c[i], e[i])`.
pub fn zip_map4_into<
    D: Copy + Send + Sync,
    A: Copy + Send + Sync,
    B: Copy + Send + Sync,
    C: Copy + Send + Sync,
    E: Copy + Send + Sync,
    OpA: ElementOp<A>,
    OpB: ElementOp<B>,
    OpC: ElementOp<C>,
    OpE: ElementOp<E>,
>(
    dest: &mut StridedViewMut<D>,
    a: &StridedView<A, OpA>,
    b: &StridedView<B, OpB>,
    c: &StridedView<C, OpC>,
    e: &StridedView<E, OpE>,
    f: impl Fn(A, B, C, E) -> D + Sync,
    exec: Exec<'_>,
) -> Result<()> {
    ensure_same_shape(dest.dims(), a.dims())?;
    ensure_same_shape(dest.dims(), b.dims())?;
    ensure_same_shape(dest.dims(), c.dims())?;
    ensure_same_shape(dest.dims(), e.dims())?;
    let validated = validate_destination_layout(dest.dims(), dest.strides())?;
    zip_map4_into_validated(dest, a, b, c, e, f, validated, exec)
}

#[allow(clippy::too_many_arguments)]
pub(crate) fn zip_map4_into_validated<
    D: Copy + Send + Sync,
    A: Copy + Send + Sync,
    B: Copy + Send + Sync,
    C: Copy + Send + Sync,
    E: Copy + Send + Sync,
    OpA: ElementOp<A>,
    OpB: ElementOp<B>,
    OpC: ElementOp<C>,
    OpE: ElementOp<E>,
>(
    dest: &mut StridedViewMut<D>,
    a: &StridedView<A, OpA>,
    b: &StridedView<B, OpB>,
    c: &StridedView<C, OpC>,
    e: &StridedView<E, OpE>,
    f: impl Fn(A, B, C, E) -> D + Sync,
    _validated: ValidatedDestinationLayout,
    exec: Exec<'_>,
) -> Result<()> {
    ensure_same_shape(dest.dims(), a.dims())?;
    ensure_same_shape(dest.dims(), b.dims())?;
    ensure_same_shape(dest.dims(), c.dims())?;
    ensure_same_shape(dest.dims(), e.dims())?;
    let dst_ptr = dest.as_mut_ptr();
    let a_ptr = a.ptr();
    let b_ptr = b.ptr();
    let c_ptr = c.ptr();
    let e_ptr = e.ptr();

    let dst_dims = dest.dims();
    let dst_strides = dest.strides();

    if sequential_contiguous_layout(
        KernelClass::Map,
        dst_dims,
        &[
            dst_strides,
            a.strides(),
            b.strides(),
            c.strides(),
            e.strides(),
        ],
        &exec,
    )?
    .is_some()
    {
        let len = total_len(dst_dims)?;
        let dst = unsafe { std::slice::from_raw_parts_mut(dst_ptr, len) };
        let sa = unsafe { std::slice::from_raw_parts(a_ptr, len) };
        let sb = unsafe { std::slice::from_raw_parts(b_ptr, len) };
        let sc = unsafe { std::slice::from_raw_parts(c_ptr, len) };
        let se = unsafe { std::slice::from_raw_parts(e_ptr, len) };
        simd::dispatch_if_large(len, || {
            for i in 0..len {
                dst[i] = f(
                    OpA::apply(sa[i]),
                    OpB::apply(sb[i]),
                    OpC::apply(sc[i]),
                    OpE::apply(se[i]),
                );
            }
        });
        return Ok(());
    }

    let strides_list: [&[isize]; 5] = [
        dst_strides,
        a.strides(),
        b.strides(),
        c.strides(),
        e.strides(),
    ];
    let elem_size = std::mem::size_of::<D>()
        .max(std::mem::size_of::<A>())
        .max(std::mem::size_of::<B>())
        .max(std::mem::size_of::<C>())
        .max(std::mem::size_of::<E>());
    let total = total_len(dst_dims)?;

    // Small tensor fast path: skip compute_order and compute_block_sizes
    let (fused_dims, ordered_strides, plan) = if total <= SMALL_TENSOR_THRESHOLD {
        build_plan_fused_small(dst_dims, &strides_list)
    } else {
        build_plan_fused(dst_dims, &strides_list, Some(0), elem_size)
    };

    let total = total_len(&fused_dims)?;
    if is_parallel(KernelClass::Map, total, &exec) {
        use crate::threading::SendPtr;
        let dst_send = SendPtr(dst_ptr);
        let a_send = SendPtr(a_ptr as *mut A);
        let b_send = SendPtr(b_ptr as *mut B);
        let c_send = SendPtr(c_ptr as *mut C);
        let e_send = SendPtr(e_ptr as *mut E);

        let costs = compute_costs(&ordered_strides);
        let initial_offsets: crate::inline::OperandVec<isize> =
            strides_list.iter().map(|_| 0isize).collect();
        return mapreduce_threaded(
            &fused_dims,
            &plan.block,
            &ordered_strides,
            &initial_offsets,
            &costs,
            KernelClass::Map,
            exec,
            &|dims, blocks, strides_list, offsets| {
                for_each_inner_block_with_offsets(
                    dims,
                    blocks,
                    strides_list,
                    offsets,
                    |offsets, len, strides| {
                        let dp = unsafe { dst_send.as_ptr().offset(offsets[0]) };
                        let ap = unsafe { a_send.as_const().offset(offsets[1]) };
                        let bp = unsafe { b_send.as_const().offset(offsets[2]) };
                        let cp = unsafe { c_send.as_const().offset(offsets[3]) };
                        let ep = unsafe { e_send.as_const().offset(offsets[4]) };
                        unsafe {
                            inner_loop_map4::<D, A, B, C, E, OpA, OpB, OpC, OpE>(
                                dp, strides[0], ap, strides[1], bp, strides[2], cp, strides[3], ep,
                                strides[4], len, &f,
                            )
                        };
                        Ok(())
                    },
                )
            },
        );
    }

    let initial_offsets: crate::inline::OperandVec<isize> =
        ordered_strides.iter().map(|_| 0isize).collect();
    for_each_inner_block_preordered(
        &fused_dims,
        &plan.block,
        &ordered_strides,
        &initial_offsets,
        |offsets, len, strides| {
            let dp = unsafe { dst_ptr.offset(offsets[0]) };
            let ap = unsafe { a_ptr.offset(offsets[1]) };
            let bp = unsafe { b_ptr.offset(offsets[2]) };
            let cp = unsafe { c_ptr.offset(offsets[3]) };
            let ep = unsafe { e_ptr.offset(offsets[4]) };
            unsafe {
                inner_loop_map4::<D, A, B, C, E, OpA, OpB, OpC, OpE>(
                    dp, strides[0], ap, strides[1], bp, strides[2], cp, strides[3], ep, strides[4],
                    len, &f,
                )
            };
            Ok(())
        },
    )
}

#[cfg(test)]
#[path = "map_view/tests/scalar_branch_tests.rs"]
mod scalar_branch_tests;

#[cfg(test)]
#[path = "map_view/tests/tests.rs"]
mod tests;