strided_kernel/

gather_plan.rs

//! Prepared indexed plans over raw strided value and index layouts.
//!
//! This module owns the generic gather, dynamic-slice/update, and scatter
//! traversals used by the erased replay layer. It models the XLA/tenferro
//! indexed shape vocabulary, but keeps tensor allocation, dtype promotion, and
//! frontend error policy outside `strided-kernel`.

use core::{mem::MaybeUninit, ops::Add};

use crate::copy_plan::{CopyPlan, OverwriteWriter, ReadModifyWrite};
use crate::{
    MaybeSendSync, RawStridedMut, RawStridedRef, Result, StridedError, RAW_FUSED_RANK_LIMIT,
};

#[cfg(feature = "parallel")]
type AxisVec<T> = smallvec::SmallVec<[T; RAW_FUSED_RANK_LIMIT]>;
#[cfg(not(feature = "parallel"))]
type AxisVec<T> = Vec<T>;

/// Gather configuration shared by generic and erased replay.
///
/// The fields follow the usual gather vocabulary:
///
/// - `start_index_map[component]` names the operand axis controlled by a
///   component in the index vector;
/// - `collapsed_slice_dims` names operand axes whose slice size is one and
///   which do not appear as output window axes;
/// - `offset_dims` names output axes that represent window offsets;
/// - output axes not in `offset_dims` are batch axes from `start_indices`;
/// - `index_vector_dim == start_indices_rank` represents scalar index vectors.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct GatherSpec {
    pub offset_dims: Vec<usize>,
    pub collapsed_slice_dims: Vec<usize>,
    pub start_index_map: Vec<usize>,
    pub index_vector_dim: usize,
    pub slice_sizes: Vec<usize>,
}

/// Index scalar types accepted by [`GatherPlan`].
pub trait GatherIndex: Copy + MaybeSendSync {
    fn to_i64(self) -> i64;
}

impl GatherIndex for i32 {
    #[inline]
    fn to_i64(self) -> i64 {
        i64::from(self)
    }
}

impl GatherIndex for i64 {
    #[inline]
    fn to_i64(self) -> i64 {
        self
    }
}

/// A compiled gather traversal for one value layout, index layout, and output layout.
#[derive(Clone, Debug)]
pub struct GatherPlan {
    operand_dims: AxisVec<usize>,
    operand_strides: AxisVec<isize>,
    index_dims: AxisVec<usize>,
    index_strides: AxisVec<isize>,
    dest_dims: AxisVec<usize>,
    dest_strides: AxisVec<isize>,
    spec: GatherSpec,
    batch_shape: AxisVec<usize>,
    out_axis_to_operand_dim: AxisVec<Option<usize>>,
    total: usize,
}

/// Scatter configuration shared by generic and erased replay.
///
/// `ScatterPlan` implements tenferro's current additive scatter semantics:
/// every update value is added to the selected output slot, so overlapping
/// windows accumulate in deterministic column-major replay order.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ScatterSpec {
    pub update_window_dims: Vec<usize>,
    pub inserted_window_dims: Vec<usize>,
    pub scatter_dims_to_operand_dims: Vec<usize>,
    pub index_vector_dim: usize,
}

/// A compiled fixed-window dynamic-slice traversal.
#[derive(Clone, Debug)]
pub struct DynamicSlicePlan {
    operand_dims: AxisVec<usize>,
    operand_strides: AxisVec<isize>,
    start_dims: AxisVec<usize>,
    start_strides: AxisVec<isize>,
    dest_dims: AxisVec<usize>,
    dest_strides: AxisVec<isize>,
    slice_sizes: AxisVec<usize>,
    total: usize,
}

/// A compiled dynamic-update-slice traversal.
///
/// Execution first copies `operand` into `dest`, then overwrites the clamped
/// update window. The plan performs no allocation for ranks at most
/// [`RAW_FUSED_RANK_LIMIT`].
#[derive(Clone, Debug)]
pub struct DynamicUpdateSlicePlan {
    operand_dims: AxisVec<usize>,
    operand_strides: AxisVec<isize>,
    start_dims: AxisVec<usize>,
    start_strides: AxisVec<isize>,
    update_dims: AxisVec<usize>,
    update_strides: AxisVec<isize>,
    dest_dims: AxisVec<usize>,
    dest_strides: AxisVec<isize>,
    total: usize,
    copy_plan: CopyPlan,
}

/// A compiled additive scatter traversal.
///
/// Execution first copies `operand` into `dest`, then applies additive updates
/// in deterministic column-major order. Boolean values are intentionally not
/// supported because additive scatter has no bool semantics.
#[derive(Clone, Debug)]
pub struct ScatterPlan {
    operand_dims: AxisVec<usize>,
    operand_strides: AxisVec<isize>,
    index_dims: AxisVec<usize>,
    index_strides: AxisVec<isize>,
    update_dims: AxisVec<usize>,
    update_strides: AxisVec<isize>,
    dest_dims: AxisVec<usize>,
    dest_strides: AxisVec<isize>,
    spec: ScatterSpec,
    batch_shape: AxisVec<usize>,
    window_dims: AxisVec<usize>,
    window_shape: AxisVec<usize>,
    window_shape_updates: AxisVec<usize>,
    is_update_window_dim: AxisVec<bool>,
    batch_elems: usize,
    window_elems: usize,
    copy_plan: CopyPlan,
}

impl GatherPlan {
    /// Compile a gather plan for fixed operand, index, and destination layouts.
    pub fn compile(
        operand_dims: &[usize],
        operand_strides: &[isize],
        index_dims: &[usize],
        index_strides: &[isize],
        dest_dims: &[usize],
        dest_strides: &[isize],
        spec: GatherSpec,
    ) -> Result<Self> {
        if operand_dims.len() != operand_strides.len()
            || index_dims.len() != index_strides.len()
            || dest_dims.len() != dest_strides.len()
        {
            return Err(StridedError::StrideLengthMismatch);
        }
        checked_total_len(operand_dims)?;
        checked_total_len(index_dims)?;
        let total = checked_total_len(dest_dims)?;
        if !crate::fused::is_injective_layout(dest_dims, dest_strides) {
            return Err(StridedError::NonInjectiveOutputLayout);
        }

        let operand_rank = operand_dims.len();
        if spec.slice_sizes.len() != operand_rank {
            return Err(StridedError::RankMismatch(
                spec.slice_sizes.len(),
                operand_rank,
            ));
        }
        validate_unique_axes(&spec.collapsed_slice_dims, operand_rank)?;
        validate_unique_axes(&spec.start_index_map, operand_rank)?;
        if spec.index_vector_dim > index_dims.len() {
            return Err(StridedError::InvalidAxis {
                axis: spec.index_vector_dim,
                rank: index_dims.len() + 1,
            });
        }

        for (axis, (&window, &dim)) in spec.slice_sizes.iter().zip(operand_dims.iter()).enumerate()
        {
            if window > dim {
                return Err(StridedError::InvalidAxis {
                    axis,
                    rank: operand_rank,
                });
            }
        }
        for &axis in &spec.collapsed_slice_dims {
            if spec.slice_sizes[axis] != 1 {
                return Err(StridedError::InvalidAxis {
                    axis,
                    rank: operand_rank,
                });
            }
        }

        let index_vector_size = if spec.index_vector_dim == index_dims.len() {
            1
        } else {
            index_dims[spec.index_vector_dim]
        };
        if index_vector_size != spec.start_index_map.len() {
            return Err(StridedError::RankMismatch(
                index_vector_size,
                spec.start_index_map.len(),
            ));
        }

        let window_dims = operand_window_dims(operand_rank, &spec.collapsed_slice_dims);
        if spec.offset_dims.len() != window_dims.len() {
            return Err(StridedError::RankMismatch(
                spec.offset_dims.len(),
                window_dims.len(),
            ));
        }

        let batch_shape = index_batch_shape(index_dims, spec.index_vector_dim);
        let out_rank = batch_shape.len() + spec.offset_dims.len();
        validate_unique_axes(&spec.offset_dims, out_rank)?;

        let mut out_axis_to_operand_dim: AxisVec<Option<usize>> =
            (0..out_rank).map(|_| None).collect();
        for (offset_axis, &out_axis) in spec.offset_dims.iter().enumerate() {
            out_axis_to_operand_dim[out_axis] = Some(window_dims[offset_axis]);
        }

        let mut expected_dest_dims: AxisVec<usize> = AxisVec::with_capacity(out_rank);
        let mut batch_axis = 0usize;
        for &operand_dim in &out_axis_to_operand_dim {
            match operand_dim {
                Some(axis) => expected_dest_dims.push(spec.slice_sizes[axis]),
                None => {
                    expected_dest_dims.push(batch_shape[batch_axis]);
                    batch_axis += 1;
                }
            }
        }
        if dest_dims != &expected_dest_dims[..] {
            return Err(StridedError::ShapeMismatch(
                dest_dims.to_vec(),
                expected_dest_dims.to_vec(),
            ));
        }

        Ok(Self {
            operand_dims: operand_dims.into(),
            operand_strides: operand_strides.into(),
            index_dims: index_dims.into(),
            index_strides: index_strides.into(),
            dest_dims: dest_dims.into(),
            dest_strides: dest_strides.into(),
            spec,
            batch_shape,
            out_axis_to_operand_dim,
            total,
        })
    }

    #[inline]
    pub fn spec(&self) -> &GatherSpec {
        &self.spec
    }

    #[inline]
    pub fn dest_dims(&self) -> &[usize] {
        &self.dest_dims
    }

    /// Execute the prepared gather traversal.
    pub fn execute<T, I>(
        &self,
        dest: &mut RawStridedMut<'_, T>,
        operand: &RawStridedRef<'_, T>,
        start_indices: &RawStridedRef<'_, I>,
    ) -> Result<()>
    where
        T: Copy + MaybeSendSync,
        I: GatherIndex,
    {
        self.execute_with_writer(dest, operand, start_indices)
    }

    /// Execute the prepared gather into a destination whose reachable slots
    /// may be uninitialized. Every logical destination slot is written.
    pub(crate) fn execute_uninit<T, I>(
        &self,
        dest: &mut RawStridedMut<'_, MaybeUninit<T>>,
        operand: &RawStridedRef<'_, T>,
        start_indices: &RawStridedRef<'_, I>,
    ) -> Result<()>
    where
        T: Copy + MaybeSendSync,
        I: GatherIndex,
    {
        self.execute_with_writer(dest, operand, start_indices)
    }

    fn execute_with_writer<T, I, W>(
        &self,
        dest: &mut W,
        operand: &RawStridedRef<'_, T>,
        start_indices: &RawStridedRef<'_, I>,
    ) -> Result<()>
    where
        T: Copy + MaybeSendSync,
        I: GatherIndex,
        W: OverwriteWriter<T>,
    {
        self.check_call(dest, operand, start_indices)?;
        if self.total == 0 {
            return Ok(());
        }
        #[cfg(feature = "parallel")]
        {
            let nthreads = crate::threading::parallel_threads_for_len(self.total);
            if nthreads > 1 {
                return self.execute_parallel(dest, operand, start_indices, nthreads);
            }
        }

        let mut out_idx_storage = CoordScratch::new(self.dest_dims.len());
        let mut batch_idx_storage = CoordScratch::new(self.batch_shape.len());
        let mut operand_idx_storage = CoordScratch::new(self.operand_dims.len());
        let mut window_offsets_storage = CoordScratch::new(self.operand_dims.len());
        let out_idx = out_idx_storage.as_mut_slice();
        let batch_idx = batch_idx_storage.as_mut_slice();
        let operand_idx = operand_idx_storage.as_mut_slice();
        let window_offsets = window_offsets_storage.as_mut_slice();

        let dest_offset_base = dest.offset();
        let operand_offset_base = operand.offset();
        let operand_strides = operand.strides();
        let index_offset_base = start_indices.offset();
        let index_strides = start_indices.strides();
        let operand_data = operand.data();
        let index_data = start_indices.data();

        for _ in 0..self.total {
            window_offsets.fill(0);
            let mut batch_axis = 0usize;
            for (out_axis, &operand_dim) in self.out_axis_to_operand_dim.iter().enumerate() {
                match operand_dim {
                    Some(axis) => window_offsets[axis] = out_idx[out_axis],
                    None => {
                        batch_idx[batch_axis] = out_idx[out_axis];
                        batch_axis += 1;
                    }
                }
            }

            operand_idx.fill(0);
            for (component, &operand_dim) in self.spec.start_index_map.iter().enumerate() {
                let start = self.index_component(
                    start_indices.dims(),
                    index_strides,
                    index_offset_base,
                    index_data,
                    &batch_idx,
                    component,
                )?;
                operand_idx[operand_dim] = self.clamp_window_start(start, operand_dim);
            }
            for axis in 0..operand_idx.len() {
                operand_idx[axis] += window_offsets[axis];
            }

            let dest_offset = checked_strided_offset(dest_offset_base, dest.strides(), &out_idx)?;
            let operand_offset =
                checked_strided_offset(operand_offset_base, operand_strides, &operand_idx)?;
            // SAFETY: the validated operand layout proves this source offset.
            let value = unsafe { *operand_data.as_ptr().offset(operand_offset) };
            // SAFETY: the validated plan proves this logical offset is in-bounds.
            unsafe { dest.write_at(dest_offset, value) };
            advance_col_major_index(out_idx, &self.dest_dims);
        }
        Ok(())
    }

    #[cfg(feature = "parallel")]
    fn execute_parallel<T, I, W>(
        &self,
        dest: &mut W,
        operand: &RawStridedRef<'_, T>,
        start_indices: &RawStridedRef<'_, I>,
        nthreads: usize,
    ) -> Result<()>
    where
        T: Copy + MaybeSendSync,
        I: GatherIndex,
        W: OverwriteWriter<T>,
    {
        let dest_offset_base = dest.offset();
        let operand_offset_base = operand.offset();
        let index_offset_base = start_indices.offset();
        // SAFETY: the validated writer owns the destination allocation.
        let dest_ptr = crate::threading::SendPtr(unsafe { dest.data_ptr() });
        let operand_ptr = crate::threading::SendPtr(operand.data().as_ptr() as *mut T);
        let index_ptr = crate::threading::SendPtr(start_indices.data().as_ptr() as *mut I);

        crate::threading::parallel_map_reduce(
            0..self.total,
            nthreads,
            &|range| {
                let mut out_idx_storage = CoordScratch::new(self.dest_dims.len());
                let mut batch_idx_storage = CoordScratch::new(self.batch_shape.len());
                let mut operand_idx_storage = CoordScratch::new(self.operand_dims.len());
                let mut window_offsets_storage = CoordScratch::new(self.operand_dims.len());
                let out_idx = out_idx_storage.as_mut_slice();
                let batch_idx = batch_idx_storage.as_mut_slice();
                let operand_idx = operand_idx_storage.as_mut_slice();
                let window_offsets = window_offsets_storage.as_mut_slice();
                fill_col_major_index(range.start, &self.dest_dims, out_idx);
                let dest_ptr = dest_ptr.as_ptr();
                let operand_ptr = operand_ptr.as_const();
                let index_ptr = index_ptr.as_const();

                for _ in range {
                    window_offsets.fill(0);
                    let mut batch_axis = 0usize;
                    for (out_axis, &operand_dim) in self.out_axis_to_operand_dim.iter().enumerate()
                    {
                        match operand_dim {
                            Some(axis) => window_offsets[axis] = out_idx[out_axis],
                            None => {
                                batch_idx[batch_axis] = out_idx[out_axis];
                                batch_axis += 1;
                            }
                        }
                    }

                    operand_idx.fill(0);
                    for (component, &operand_dim) in self.spec.start_index_map.iter().enumerate() {
                        let start = self.index_component_ptr(
                            start_indices.dims(),
                            index_offset_base,
                            index_ptr,
                            batch_idx,
                            component,
                        )?;
                        operand_idx[operand_dim] = self.clamp_window_start(start, operand_dim);
                    }
                    for axis in 0..operand_idx.len() {
                        operand_idx[axis] += window_offsets[axis];
                    }

                    let dest_offset =
                        checked_strided_offset(dest_offset_base, &self.dest_strides, out_idx)?;
                    let operand_offset = checked_strided_offset(
                        operand_offset_base,
                        &self.operand_strides,
                        operand_idx,
                    )?;
                    unsafe {
                        // SAFETY: gather writes one value per logical output,
                        // and compile rejected non-injective destination
                        // layouts.
                        dest_ptr
                            .offset(dest_offset)
                            .write(operand_ptr.offset(operand_offset).read());
                    }
                    advance_col_major_index(out_idx, &self.dest_dims);
                }
                Ok(())
            },
            &|left, right| left.and(right),
        )
    }

    fn check_call<T, I, W>(
        &self,
        dest: &W,
        operand: &RawStridedRef<'_, T>,
        start_indices: &RawStridedRef<'_, I>,
    ) -> Result<()>
    where
        W: OverwriteWriter<T>,
    {
        if dest.dims() != &self.dest_dims[..]
            || dest.strides() != &self.dest_strides[..]
            || operand.dims() != &self.operand_dims[..]
            || operand.strides() != &self.operand_strides[..]
            || start_indices.dims() != &self.index_dims[..]
            || start_indices.strides() != &self.index_strides[..]
        {
            return Err(StridedError::PlanLayoutMismatch);
        }
        Ok(())
    }

    fn index_component<I>(
        &self,
        index_dims: &[usize],
        index_strides: &[isize],
        index_offset_base: isize,
        index_data: &[I],
        batch_idx: &[usize],
        component: usize,
    ) -> Result<i64>
    where
        I: GatherIndex,
    {
        let mut offset = index_offset_base;
        let mut batch_axis = 0usize;
        for axis in 0..index_dims.len() {
            let coord = if axis == self.spec.index_vector_dim {
                component
            } else {
                let coord = batch_idx[batch_axis];
                batch_axis += 1;
                coord
            };
            offset = checked_offset_add(offset, index_strides[axis], coord)?;
        }
        Ok(unsafe { *index_data.as_ptr().offset(offset) }.to_i64())
    }

    #[cfg(feature = "parallel")]
    fn index_component_ptr<I>(
        &self,
        index_dims: &[usize],
        index_offset_base: isize,
        index_ptr: *const I,
        batch_idx: &[usize],
        component: usize,
    ) -> Result<i64>
    where
        I: GatherIndex,
    {
        let mut offset = index_offset_base;
        let mut batch_axis = 0usize;
        for axis in 0..index_dims.len() {
            let coord = if axis == self.spec.index_vector_dim {
                component
            } else {
                let coord = batch_idx[batch_axis];
                batch_axis += 1;
                coord
            };
            offset = checked_offset_add(offset, self.index_strides[axis], coord)?;
        }
        Ok(unsafe { *index_ptr.offset(offset) }.to_i64())
    }

    #[inline]
    fn clamp_window_start(&self, start: i64, operand_dim: usize) -> usize {
        let dim_size = self.operand_dims[operand_dim];
        let window_size = self.spec.slice_sizes[operand_dim];
        let max_start = dim_size.saturating_sub(window_size) as i64;
        start.clamp(0, max_start) as usize
    }
}

impl DynamicSlicePlan {
    /// Compile a fixed-window dynamic-slice traversal.
    ///
    /// `start_*` describes a rank-1 index vector whose length equals the
    /// operand rank. `dest_dims` must equal `slice_sizes`.
    pub fn compile(
        operand_dims: &[usize],
        operand_strides: &[isize],
        start_dims: &[usize],
        start_strides: &[isize],
        dest_dims: &[usize],
        dest_strides: &[isize],
        slice_sizes: &[usize],
    ) -> Result<Self> {
        if operand_dims.len() != operand_strides.len()
            || start_dims.len() != start_strides.len()
            || dest_dims.len() != dest_strides.len()
        {
            return Err(StridedError::StrideLengthMismatch);
        }
        if slice_sizes.len() != operand_dims.len() {
            return Err(StridedError::RankMismatch(
                slice_sizes.len(),
                operand_dims.len(),
            ));
        }
        validate_start_vector(start_dims, operand_dims.len())?;
        checked_total_len(operand_dims)?;
        checked_total_len(start_dims)?;
        let total = checked_total_len(dest_dims)?;
        if dest_dims != slice_sizes {
            return Err(StridedError::ShapeMismatch(
                dest_dims.to_vec(),
                slice_sizes.to_vec(),
            ));
        }
        if !crate::fused::is_injective_layout(dest_dims, dest_strides) {
            return Err(StridedError::NonInjectiveOutputLayout);
        }
        validate_window_sizes(operand_dims, slice_sizes)?;

        Ok(Self {
            operand_dims: operand_dims.into(),
            operand_strides: operand_strides.into(),
            start_dims: start_dims.into(),
            start_strides: start_strides.into(),
            dest_dims: dest_dims.into(),
            dest_strides: dest_strides.into(),
            slice_sizes: slice_sizes.into(),
            total,
        })
    }

    /// Execute the prepared dynamic-slice traversal.
    pub fn execute<T, I>(
        &self,
        dest: &mut RawStridedMut<'_, T>,
        operand: &RawStridedRef<'_, T>,
        starts: &RawStridedRef<'_, I>,
    ) -> Result<()>
    where
        T: Copy + MaybeSendSync,
        I: GatherIndex,
    {
        self.execute_with_writer(dest, operand, starts)
    }

    pub(crate) fn execute_uninit<T, I>(
        &self,
        dest: &mut RawStridedMut<'_, MaybeUninit<T>>,
        operand: &RawStridedRef<'_, T>,
        starts: &RawStridedRef<'_, I>,
    ) -> Result<()>
    where
        T: Copy + MaybeSendSync,
        I: GatherIndex,
    {
        self.execute_with_writer(dest, operand, starts)
    }

    fn execute_with_writer<T, I, W>(
        &self,
        dest: &mut W,
        operand: &RawStridedRef<'_, T>,
        starts: &RawStridedRef<'_, I>,
    ) -> Result<()>
    where
        T: Copy + MaybeSendSync,
        I: GatherIndex,
        W: OverwriteWriter<T>,
    {
        self.check_call(dest, operand, starts)?;
        if self.total == 0 {
            return Ok(());
        }
        if self.uses_rank_one_contiguous_path() {
            return self.execute_rank_one_contiguous(dest, operand, starts);
        }
        #[cfg(feature = "parallel")]
        {
            let nthreads = crate::threading::parallel_threads_for_len(self.total);
            if nthreads > 1 {
                return self.execute_parallel(dest, operand, starts, nthreads);
            }
        }

        let mut starts_storage = CoordScratch::new(self.operand_dims.len());
        let mut dest_idx_storage = CoordScratch::new(self.dest_dims.len());
        let mut operand_idx_storage = CoordScratch::new(self.operand_dims.len());
        let clamped_starts = starts_storage.as_mut_slice();
        let dest_idx = dest_idx_storage.as_mut_slice();
        let operand_idx = operand_idx_storage.as_mut_slice();
        read_clamped_starts(
            starts,
            &self.operand_dims,
            &self.slice_sizes,
            clamped_starts,
        )?;

        let operand_offset_base = operand.offset();
        let operand_strides = operand.strides();
        let dest_offset_base = dest.offset();
        let operand_data = operand.data();

        for _ in 0..self.total {
            for axis in 0..operand_idx.len() {
                operand_idx[axis] = clamped_starts[axis] + dest_idx[axis];
            }
            let operand_offset =
                checked_strided_offset(operand_offset_base, operand_strides, operand_idx)?;
            let dest_offset = checked_strided_offset(dest_offset_base, dest.strides(), dest_idx)?;
            // SAFETY: the validated plan proves both offsets.
            let value = unsafe { *operand_data.as_ptr().offset(operand_offset) };
            // SAFETY: the validated plan proves this logical offset is in-bounds.
            unsafe { dest.write_at(dest_offset, value) };
            advance_col_major_index(dest_idx, &self.dest_dims);
        }
        Ok(())
    }

    #[inline]
    fn uses_rank_one_contiguous_path(&self) -> bool {
        self.operand_dims.len() == 1 && self.operand_strides[0] == 1 && self.dest_strides[0] == 1
    }

    fn execute_rank_one_contiguous<T, I, W>(
        &self,
        dest: &mut W,
        operand: &RawStridedRef<'_, T>,
        starts: &RawStridedRef<'_, I>,
    ) -> Result<()>
    where
        T: Copy,
        I: GatherIndex,
        W: OverwriteWriter<T>,
    {
        let mut clamped_starts = [0usize; 1];
        read_clamped_starts(
            starts,
            &self.operand_dims,
            &self.slice_sizes,
            &mut clamped_starts,
        )?;
        let source_start = checked_offset_add(operand.offset(), 1, clamped_starts[0])?;
        let source_start =
            usize::try_from(source_start).map_err(|_| StridedError::OffsetOverflow)?;
        let dest_start =
            usize::try_from(dest.offset()).map_err(|_| StridedError::OffsetOverflow)?;
        let source_end = source_start
            .checked_add(self.total)
            .ok_or(StridedError::OffsetOverflow)?;
        let source = operand
            .data()
            .get(source_start..source_end)
            .ok_or(StridedError::OffsetOverflow)?;
        // SAFETY: the validated writer owns the destination allocation.
        let dest_ptr = unsafe { dest.data_ptr() };
        // SAFETY: bounds were checked above and the writer owns the logical
        // destination storage.
        unsafe {
            core::ptr::copy_nonoverlapping(source.as_ptr(), dest_ptr.add(dest_start), self.total);
        }
        Ok(())
    }

    #[cfg(feature = "parallel")]
    fn execute_parallel<T, I, W>(
        &self,
        dest: &mut W,
        operand: &RawStridedRef<'_, T>,
        starts: &RawStridedRef<'_, I>,
        nthreads: usize,
    ) -> Result<()>
    where
        T: Copy + MaybeSendSync,
        I: GatherIndex,
        W: OverwriteWriter<T>,
    {
        let mut clamped_starts: AxisVec<usize> = (0..self.operand_dims.len()).map(|_| 0).collect();
        read_clamped_starts(
            starts,
            &self.operand_dims,
            &self.slice_sizes,
            &mut clamped_starts,
        )?;

        let operand_offset_base = operand.offset();
        let dest_offset_base = dest.offset();
        let operand_ptr = crate::threading::SendPtr(operand.data().as_ptr() as *mut T);
        // SAFETY: the validated writer owns the destination allocation.
        let dest_ptr = crate::threading::SendPtr(unsafe { dest.data_ptr() });

        crate::threading::parallel_map_reduce(
            0..self.total,
            nthreads,
            &|range| {
                let mut dest_idx_storage = CoordScratch::new(self.dest_dims.len());
                let mut operand_idx_storage = CoordScratch::new(self.operand_dims.len());
                let dest_idx = dest_idx_storage.as_mut_slice();
                let operand_idx = operand_idx_storage.as_mut_slice();
                fill_col_major_index(range.start, &self.dest_dims, dest_idx);
                let operand_ptr = operand_ptr.as_const();
                let dest_ptr = dest_ptr.as_ptr();

                for _ in range {
                    for axis in 0..operand_idx.len() {
                        operand_idx[axis] = clamped_starts[axis] + dest_idx[axis];
                    }
                    let operand_offset = checked_strided_offset(
                        operand_offset_base,
                        &self.operand_strides,
                        operand_idx,
                    )?;
                    let dest_offset =
                        checked_strided_offset(dest_offset_base, &self.dest_strides, dest_idx)?;
                    unsafe {
                        // SAFETY: dynamic slice writes one value per logical
                        // output, and compile rejected non-injective
                        // destination layouts.
                        dest_ptr
                            .offset(dest_offset)
                            .write(operand_ptr.offset(operand_offset).read());
                    }
                    advance_col_major_index(dest_idx, &self.dest_dims);
                }
                Ok(())
            },
            &|left, right| left.and(right),
        )
    }

    fn check_call<T, I, W>(
        &self,
        dest: &W,
        operand: &RawStridedRef<'_, T>,
        starts: &RawStridedRef<'_, I>,
    ) -> Result<()>
    where
        W: OverwriteWriter<T>,
    {
        if dest.dims() != &self.dest_dims[..]
            || dest.strides() != &self.dest_strides[..]
            || operand.dims() != &self.operand_dims[..]
            || operand.strides() != &self.operand_strides[..]
            || starts.dims() != &self.start_dims[..]
            || starts.strides() != &self.start_strides[..]
        {
            return Err(StridedError::PlanLayoutMismatch);
        }
        Ok(())
    }
}

impl DynamicUpdateSlicePlan {
    /// Compile a dynamic-update-slice traversal.
    ///
    /// `dest_dims` must match `operand_dims`; execution materializes
    /// `dest = operand` and then overwrites the clamped update window.
    #[allow(clippy::too_many_arguments)]
    pub fn compile(
        operand_dims: &[usize],
        operand_strides: &[isize],
        start_dims: &[usize],
        start_strides: &[isize],
        update_dims: &[usize],
        update_strides: &[isize],
        dest_dims: &[usize],
        dest_strides: &[isize],
    ) -> Result<Self> {
        if operand_dims.len() != operand_strides.len()
            || start_dims.len() != start_strides.len()
            || update_dims.len() != update_strides.len()
            || dest_dims.len() != dest_strides.len()
        {
            return Err(StridedError::StrideLengthMismatch);
        }
        if update_dims.len() != operand_dims.len() {
            return Err(StridedError::RankMismatch(
                update_dims.len(),
                operand_dims.len(),
            ));
        }
        validate_start_vector(start_dims, operand_dims.len())?;
        checked_total_len(operand_dims)?;
        checked_total_len(start_dims)?;
        let total = checked_total_len(update_dims)?;
        if dest_dims != operand_dims {
            return Err(StridedError::ShapeMismatch(
                dest_dims.to_vec(),
                operand_dims.to_vec(),
            ));
        }
        if !crate::fused::is_injective_layout(dest_dims, dest_strides) {
            return Err(StridedError::NonInjectiveOutputLayout);
        }
        validate_window_sizes(operand_dims, update_dims)?;
        let copy_plan = CopyPlan::compile(operand_dims, dest_strides, operand_strides)?;

        Ok(Self {
            operand_dims: operand_dims.into(),
            operand_strides: operand_strides.into(),
            start_dims: start_dims.into(),
            start_strides: start_strides.into(),
            update_dims: update_dims.into(),
            update_strides: update_strides.into(),
            dest_dims: dest_dims.into(),
            dest_strides: dest_strides.into(),
            total,
            copy_plan,
        })
    }

    /// Execute the prepared dynamic-update-slice traversal.
    pub fn execute<T, I>(
        &self,
        dest: &mut RawStridedMut<'_, T>,
        operand: &RawStridedRef<'_, T>,
        update: &RawStridedRef<'_, T>,
        starts: &RawStridedRef<'_, I>,
    ) -> Result<()>
    where
        T: Copy + MaybeSendSync,
        I: GatherIndex,
    {
        self.check_call(dest, operand, update, starts)?;
        self.copy_plan.execute(dest, operand)?;
        self.execute_update_with_writer(dest, update, starts)
    }

    /// Execute dynamic update into a destination whose reachable slots may be
    /// uninitialized. The copy completes before any read-modify-write access.
    pub(crate) fn execute_uninit<'a, T, I>(
        &self,
        dest: &'a mut RawStridedMut<'a, MaybeUninit<T>>,
        operand: &RawStridedRef<'_, T>,
        update: &RawStridedRef<'_, T>,
        starts: &RawStridedRef<'_, I>,
    ) -> Result<()>
    where
        T: Copy + MaybeSendSync,
        I: GatherIndex,
    {
        self.check_call(dest, operand, update, starts)?;
        self.copy_plan
            .execute_uninit_then(dest, operand, |mut receipt| {
                self.execute_update_with_writer(&mut receipt, update, starts)
            })?
    }

    fn execute_update_with_writer<T, I, W>(
        &self,
        dest: &mut W,
        update: &RawStridedRef<'_, T>,
        starts: &RawStridedRef<'_, I>,
    ) -> Result<()>
    where
        T: Copy + MaybeSendSync,
        I: GatherIndex,
        W: OverwriteWriter<T>,
    {
        if self.total == 0 {
            return Ok(());
        }
        if self.uses_rank_one_contiguous_path() {
            return self.execute_rank_one_contiguous(dest, update, starts);
        }
        #[cfg(feature = "parallel")]
        {
            let nthreads = crate::threading::parallel_threads_for_len(self.total);
            if nthreads > 1 {
                return self.execute_update_parallel(dest, update, starts, nthreads);
            }
        }

        let mut starts_storage = CoordScratch::new(self.operand_dims.len());
        let mut update_idx_storage = CoordScratch::new(self.update_dims.len());
        let mut dest_idx_storage = CoordScratch::new(self.operand_dims.len());
        let clamped_starts = starts_storage.as_mut_slice();
        let update_idx = update_idx_storage.as_mut_slice();
        let dest_idx = dest_idx_storage.as_mut_slice();
        read_clamped_starts(
            starts,
            &self.operand_dims,
            &self.update_dims,
            clamped_starts,
        )?;

        let update_offset_base = update.offset();
        let update_strides = update.strides();
        let dest_offset_base = dest.offset();
        let update_data = update.data();

        for _ in 0..self.total {
            for axis in 0..dest_idx.len() {
                dest_idx[axis] = clamped_starts[axis] + update_idx[axis];
            }
            let update_offset =
                checked_strided_offset(update_offset_base, update_strides, update_idx)?;
            let dest_offset = checked_strided_offset(dest_offset_base, dest.strides(), dest_idx)?;
            let value = unsafe { *update_data.as_ptr().offset(update_offset) };
            // SAFETY: the validated plan proves this logical offset is in-bounds.
            unsafe { dest.write_at(dest_offset, value) };
            advance_col_major_index(update_idx, &self.update_dims);
        }
        Ok(())
    }

    #[inline]
    fn uses_rank_one_contiguous_path(&self) -> bool {
        self.operand_dims.len() == 1
            && self.operand_strides[0] == 1
            && self.update_strides[0] == 1
            && self.dest_strides[0] == 1
    }

    fn execute_rank_one_contiguous<T, I, W>(
        &self,
        dest: &mut W,
        update: &RawStridedRef<'_, T>,
        starts: &RawStridedRef<'_, I>,
    ) -> Result<()>
    where
        T: Copy,
        I: GatherIndex,
        W: OverwriteWriter<T>,
    {
        let mut clamped_starts = [0usize; 1];
        read_clamped_starts(
            starts,
            &self.operand_dims,
            &self.update_dims,
            &mut clamped_starts,
        )?;
        let update_start =
            usize::try_from(update.offset()).map_err(|_| StridedError::OffsetOverflow)?;
        let dest_start = checked_offset_add(dest.offset(), 1, clamped_starts[0])?;
        let dest_start = usize::try_from(dest_start).map_err(|_| StridedError::OffsetOverflow)?;
        let update_end = update_start
            .checked_add(self.total)
            .ok_or(StridedError::OffsetOverflow)?;
        let update = update
            .data()
            .get(update_start..update_end)
            .ok_or(StridedError::OffsetOverflow)?;
        // SAFETY: the validated writer owns the destination allocation.
        let dest_ptr = unsafe { dest.data_ptr() };
        // SAFETY: the checked ranges are inside the destination allocation.
        unsafe {
            core::ptr::copy_nonoverlapping(update.as_ptr(), dest_ptr.add(dest_start), self.total);
        }
        Ok(())
    }

    #[cfg(feature = "parallel")]
    fn execute_update_parallel<T, I, W>(
        &self,
        dest: &mut W,
        update: &RawStridedRef<'_, T>,
        starts: &RawStridedRef<'_, I>,
        nthreads: usize,
    ) -> Result<()>
    where
        T: Copy + MaybeSendSync,
        I: GatherIndex,
        W: OverwriteWriter<T>,
    {
        let mut clamped_starts: AxisVec<usize> = (0..self.operand_dims.len()).map(|_| 0).collect();
        read_clamped_starts(
            starts,
            &self.operand_dims,
            &self.update_dims,
            &mut clamped_starts,
        )?;

        let update_offset_base = update.offset();
        let dest_offset_base = dest.offset();
        let update_ptr = crate::threading::SendPtr(update.data().as_ptr() as *mut T);
        // SAFETY: the validated writer owns the destination allocation.
        let dest_ptr = crate::threading::SendPtr(unsafe { dest.data_ptr() });

        crate::threading::parallel_map_reduce(
            0..self.total,
            nthreads,
            &|range| {
                let mut update_idx_storage = CoordScratch::new(self.update_dims.len());
                let mut dest_idx_storage = CoordScratch::new(self.operand_dims.len());
                let update_idx = update_idx_storage.as_mut_slice();
                let dest_idx = dest_idx_storage.as_mut_slice();
                fill_col_major_index(range.start, &self.update_dims, update_idx);
                let update_ptr = update_ptr.as_const();
                let dest_ptr = dest_ptr.as_ptr();

                for _ in range {
                    for axis in 0..dest_idx.len() {
                        dest_idx[axis] = clamped_starts[axis] + update_idx[axis];
                    }
                    let update_offset = checked_strided_offset(
                        update_offset_base,
                        &self.update_strides,
                        update_idx,
                    )?;
                    let dest_offset =
                        checked_strided_offset(dest_offset_base, &self.dest_strides, dest_idx)?;
                    unsafe {
                        // SAFETY: each update-domain logical index maps to a
                        // distinct destination position for a fixed window, and
                        // compile rejected non-injective destination layouts.
                        dest_ptr
                            .offset(dest_offset)
                            .write(update_ptr.offset(update_offset).read());
                    }
                    advance_col_major_index(update_idx, &self.update_dims);
                }
                Ok(())
            },
            &|left, right| left.and(right),
        )
    }

    fn check_call<T, I, W>(
        &self,
        dest: &W,
        operand: &RawStridedRef<'_, T>,
        update: &RawStridedRef<'_, T>,
        starts: &RawStridedRef<'_, I>,
    ) -> Result<()>
    where
        W: OverwriteWriter<T>,
    {
        if dest.dims() != &self.dest_dims[..]
            || dest.strides() != &self.dest_strides[..]
            || operand.dims() != &self.operand_dims[..]
            || operand.strides() != &self.operand_strides[..]
            || update.dims() != &self.update_dims[..]
            || update.strides() != &self.update_strides[..]
            || starts.dims() != &self.start_dims[..]
            || starts.strides() != &self.start_strides[..]
        {
            return Err(StridedError::PlanLayoutMismatch);
        }
        Ok(())
    }
}

impl ScatterPlan {
    /// Compile an additive scatter traversal.
    #[allow(clippy::too_many_arguments)]
    pub fn compile(
        operand_dims: &[usize],
        operand_strides: &[isize],
        index_dims: &[usize],
        index_strides: &[isize],
        update_dims: &[usize],
        update_strides: &[isize],
        dest_dims: &[usize],
        dest_strides: &[isize],
        spec: ScatterSpec,
    ) -> Result<Self> {
        if operand_dims.len() != operand_strides.len()
            || index_dims.len() != index_strides.len()
            || update_dims.len() != update_strides.len()
            || dest_dims.len() != dest_strides.len()
        {
            return Err(StridedError::StrideLengthMismatch);
        }
        checked_total_len(operand_dims)?;
        checked_total_len(index_dims)?;
        checked_total_len(update_dims)?;
        if dest_dims != operand_dims {
            return Err(StridedError::ShapeMismatch(
                dest_dims.to_vec(),
                operand_dims.to_vec(),
            ));
        }
        if !crate::fused::is_injective_layout(dest_dims, dest_strides) {
            return Err(StridedError::NonInjectiveOutputLayout);
        }

        let operand_rank = operand_dims.len();
        validate_unique_axes(&spec.inserted_window_dims, operand_rank)?;
        validate_unique_axes(&spec.scatter_dims_to_operand_dims, operand_rank)?;
        if spec.index_vector_dim > index_dims.len() {
            return Err(StridedError::InvalidAxis {
                axis: spec.index_vector_dim,
                rank: index_dims.len() + 1,
            });
        }
        let index_vector_size = if spec.index_vector_dim == index_dims.len() {
            1
        } else {
            index_dims[spec.index_vector_dim]
        };
        if index_vector_size != spec.scatter_dims_to_operand_dims.len() {
            return Err(StridedError::RankMismatch(
                index_vector_size,
                spec.scatter_dims_to_operand_dims.len(),
            ));
        }

        let batch_shape = index_batch_shape(index_dims, spec.index_vector_dim);
        let window_dims = operand_window_dims(operand_rank, &spec.inserted_window_dims);
        if spec.update_window_dims.len() != window_dims.len() {
            return Err(StridedError::RankMismatch(
                spec.update_window_dims.len(),
                window_dims.len(),
            ));
        }

        let update_rank = update_dims.len();
        let expected_batch_rank = update_rank
            .checked_sub(spec.update_window_dims.len())
            .ok_or(StridedError::RankMismatch(
                spec.update_window_dims.len(),
                update_rank,
            ))?;
        if expected_batch_rank != batch_shape.len() {
            return Err(StridedError::RankMismatch(
                expected_batch_rank,
                batch_shape.len(),
            ));
        }
        validate_unique_axes(&spec.update_window_dims, update_rank)?;

        let mut is_update_window_dim: AxisVec<bool> = (0..update_rank).map(|_| false).collect();
        for &axis in &spec.update_window_dims {
            is_update_window_dim[axis] = true;
        }

        let mut batch_axis = 0usize;
        for axis in 0..update_rank {
            if !is_update_window_dim[axis] {
                if update_dims[axis] != batch_shape[batch_axis] {
                    return Err(StridedError::ShapeMismatch(
                        update_dims.to_vec(),
                        expected_scatter_update_shape(&batch_shape, &spec, update_dims).to_vec(),
                    ));
                }
                batch_axis += 1;
            }
        }

        let mut window_shape: AxisVec<usize> = (0..operand_rank).map(|_| 1).collect();
        let mut window_shape_updates: AxisVec<usize> =
            AxisVec::with_capacity(spec.update_window_dims.len());
        for (pos, &update_axis) in spec.update_window_dims.iter().enumerate() {
            let dim = update_dims[update_axis];
            window_shape_updates.push(dim);
            window_shape[window_dims[pos]] = dim;
        }
        validate_window_sizes(operand_dims, &window_shape)?;

        let batch_elems = checked_total_len(&batch_shape)?;
        let window_elems = checked_total_len(&window_shape_updates)?;
        let copy_plan = CopyPlan::compile(operand_dims, dest_strides, operand_strides)?;

        Ok(Self {
            operand_dims: operand_dims.into(),
            operand_strides: operand_strides.into(),
            index_dims: index_dims.into(),
            index_strides: index_strides.into(),
            update_dims: update_dims.into(),
            update_strides: update_strides.into(),
            dest_dims: dest_dims.into(),
            dest_strides: dest_strides.into(),
            spec,
            batch_shape,
            window_dims,
            window_shape,
            window_shape_updates,
            is_update_window_dim,
            batch_elems,
            window_elems,
            copy_plan,
        })
    }

    /// Execute the prepared additive scatter traversal.
    pub fn execute<T, I>(
        &self,
        dest: &mut RawStridedMut<'_, T>,
        operand: &RawStridedRef<'_, T>,
        scatter_indices: &RawStridedRef<'_, I>,
        updates: &RawStridedRef<'_, T>,
    ) -> Result<()>
    where
        T: Copy + Add<Output = T> + MaybeSendSync,
        I: GatherIndex,
    {
        self.check_call(dest, operand, scatter_indices, updates)?;
        self.copy_plan.execute(dest, operand)?;
        self.execute_updates(dest, scatter_indices, updates, |a, b| a + b)
    }

    /// Execute additive scatter into a destination whose reachable slots may
    /// be uninitialized. The operand copy completes before any RMW access.
    pub(crate) fn execute_uninit<'a, T, I>(
        &self,
        dest: &'a mut RawStridedMut<'a, MaybeUninit<T>>,
        operand: &RawStridedRef<'_, T>,
        scatter_indices: &RawStridedRef<'_, I>,
        updates: &RawStridedRef<'_, T>,
        combine: fn(T, T) -> T,
    ) -> Result<()>
    where
        T: Copy + Add<Output = T> + MaybeSendSync,
        I: GatherIndex,
    {
        self.check_call(dest, operand, scatter_indices, updates)?;
        self.copy_plan
            .execute_uninit_then(dest, operand, |mut receipt| {
                self.execute_updates(&mut receipt, scatter_indices, updates, combine)
            })?
    }

    fn execute_updates<T, I, W>(
        &self,
        dest: &mut W,
        scatter_indices: &RawStridedRef<'_, I>,
        updates: &RawStridedRef<'_, T>,
        combine: fn(T, T) -> T,
    ) -> Result<()>
    where
        T: Copy + MaybeSendSync,
        I: GatherIndex,
        W: ReadModifyWrite<T>,
    {
        if self.batch_elems == 0 || self.window_elems == 0 {
            return Ok(());
        }

        // Overlapping additive updates are order-sensitive, so this remains a
        // deterministic serial replay until a combine-aware parallel plan exists.
        let mut batch_idx_storage = CoordScratch::new(self.batch_shape.len());
        let mut window_idx_storage = CoordScratch::new(self.window_shape_updates.len());
        let mut update_idx_storage = CoordScratch::new(self.update_dims.len());
        let mut operand_base_storage = CoordScratch::new(self.operand_dims.len());
        let mut operand_idx_storage = CoordScratch::new(self.operand_dims.len());
        let batch_idx = batch_idx_storage.as_mut_slice();
        let window_idx = window_idx_storage.as_mut_slice();
        let update_idx = update_idx_storage.as_mut_slice();
        let operand_base = operand_base_storage.as_mut_slice();
        let operand_idx = operand_idx_storage.as_mut_slice();

        let index_offset_base = scatter_indices.offset();
        let index_strides = scatter_indices.strides();
        let index_data = scatter_indices.data();
        let update_offset_base = updates.offset();
        let update_strides = updates.strides();
        let update_data = updates.data();
        let dest_offset_base = dest.offset();

        for _ in 0..self.batch_elems {
            operand_base.fill(0);
            for (component, &operand_dim) in
                self.spec.scatter_dims_to_operand_dims.iter().enumerate()
            {
                let start = index_component(
                    scatter_indices.dims(),
                    index_strides,
                    index_offset_base,
                    index_data,
                    self.spec.index_vector_dim,
                    batch_idx,
                    component,
                )?;
                operand_base[operand_dim] = clamp_window_start(
                    start,
                    self.operand_dims[operand_dim],
                    self.window_shape[operand_dim],
                );
            }

            window_idx.fill(0);
            for _ in 0..self.window_elems {
                let mut batch_axis = 0usize;
                let mut window_axis = 0usize;
                for axis in 0..self.update_dims.len() {
                    if self.is_update_window_dim[axis] {
                        update_idx[axis] = window_idx[window_axis];
                        window_axis += 1;
                    } else {
                        update_idx[axis] = batch_idx[batch_axis];
                        batch_axis += 1;
                    }
                }

                operand_idx.copy_from_slice(operand_base);
                for (window_axis, &operand_axis) in self.window_dims.iter().enumerate() {
                    operand_idx[operand_axis] += window_idx[window_axis];
                }

                let update_offset =
                    checked_strided_offset(update_offset_base, update_strides, update_idx)?;
                let dest_offset =
                    checked_strided_offset(dest_offset_base, dest.strides(), operand_idx)?;
                let value = unsafe { *update_data.as_ptr().offset(update_offset) };
                // SAFETY: copy completion and serial scatter traversal prove
                // this initialized logical slot is in-bounds.
                unsafe { dest.add_at(dest_offset, value, combine) };
                advance_col_major_index(window_idx, &self.window_shape_updates);
            }
            advance_col_major_index(batch_idx, &self.batch_shape);
        }
        Ok(())
    }

    fn check_call<T, I, W>(
        &self,
        dest: &W,
        operand: &RawStridedRef<'_, T>,
        scatter_indices: &RawStridedRef<'_, I>,
        updates: &RawStridedRef<'_, T>,
    ) -> Result<()>
    where
        W: OverwriteWriter<T>,
    {
        if dest.dims() != &self.dest_dims[..]
            || dest.strides() != &self.dest_strides[..]
            || operand.dims() != &self.operand_dims[..]
            || operand.strides() != &self.operand_strides[..]
            || scatter_indices.dims() != &self.index_dims[..]
            || scatter_indices.strides() != &self.index_strides[..]
            || updates.dims() != &self.update_dims[..]
            || updates.strides() != &self.update_strides[..]
        {
            return Err(StridedError::PlanLayoutMismatch);
        }
        Ok(())
    }
}

fn validate_unique_axes(axes: &[usize], rank: usize) -> Result<()> {
    let mut seen = vec![false; rank];
    for &axis in axes {
        if axis >= rank {
            return Err(StridedError::InvalidAxis { axis, rank });
        }
        if seen[axis] {
            return Err(StridedError::InvalidAxis { axis, rank });
        }
        seen[axis] = true;
    }
    Ok(())
}

fn validate_start_vector(start_dims: &[usize], operand_rank: usize) -> Result<()> {
    if start_dims.len() != 1 {
        return Err(StridedError::RankMismatch(start_dims.len(), 1));
    }
    if start_dims[0] != operand_rank {
        return Err(StridedError::RankMismatch(start_dims[0], operand_rank));
    }
    Ok(())
}

fn validate_window_sizes(operand_dims: &[usize], window_sizes: &[usize]) -> Result<()> {
    if operand_dims.len() != window_sizes.len() {
        return Err(StridedError::RankMismatch(
            window_sizes.len(),
            operand_dims.len(),
        ));
    }
    for (axis, (&window, &dim)) in window_sizes.iter().zip(operand_dims.iter()).enumerate() {
        if window > dim {
            return Err(StridedError::InvalidAxis {
                axis,
                rank: operand_dims.len(),
            });
        }
    }
    Ok(())
}

fn read_clamped_starts<I>(
    starts: &RawStridedRef<'_, I>,
    operand_dims: &[usize],
    window_sizes: &[usize],
    out: &mut [usize],
) -> Result<()>
where
    I: GatherIndex,
{
    debug_assert_eq!(operand_dims.len(), window_sizes.len());
    debug_assert_eq!(operand_dims.len(), out.len());
    for axis in 0..operand_dims.len() {
        let offset = checked_offset_add(starts.offset(), starts.strides()[0], axis)?;
        let start = unsafe { *starts.data().as_ptr().offset(offset) }.to_i64();
        out[axis] = clamp_window_start(start, operand_dims[axis], window_sizes[axis]);
    }
    Ok(())
}

fn index_component<I>(
    index_dims: &[usize],
    index_strides: &[isize],
    index_offset_base: isize,
    index_data: &[I],
    index_vector_dim: usize,
    batch_idx: &[usize],
    component: usize,
) -> Result<i64>
where
    I: GatherIndex,
{
    let mut offset = index_offset_base;
    let mut batch_axis = 0usize;
    for axis in 0..index_dims.len() {
        let coord = if axis == index_vector_dim {
            component
        } else {
            let coord = batch_idx[batch_axis];
            batch_axis += 1;
            coord
        };
        offset = checked_offset_add(offset, index_strides[axis], coord)?;
    }
    Ok(unsafe { *index_data.as_ptr().offset(offset) }.to_i64())
}

#[inline]
fn clamp_window_start(start: i64, dim_size: usize, window_size: usize) -> usize {
    let max_start = dim_size.saturating_sub(window_size) as i64;
    start.clamp(0, max_start) as usize
}

fn expected_scatter_update_shape(
    batch_shape: &[usize],
    spec: &ScatterSpec,
    update_dims: &[usize],
) -> AxisVec<usize> {
    let mut expected: AxisVec<usize> = AxisVec::with_capacity(update_dims.len());
    let mut batch_axis = 0usize;
    for axis in 0..update_dims.len() {
        if spec.update_window_dims.contains(&axis) {
            expected.push(update_dims[axis]);
        } else {
            expected.push(batch_shape[batch_axis]);
            batch_axis += 1;
        }
    }
    expected
}

fn operand_window_dims(rank: usize, collapsed_slice_dims: &[usize]) -> AxisVec<usize> {
    (0..rank)
        .filter(|axis| !collapsed_slice_dims.contains(axis))
        .collect()
}

fn index_batch_shape(index_dims: &[usize], index_vector_dim: usize) -> AxisVec<usize> {
    if index_vector_dim == index_dims.len() {
        return index_dims.into();
    }
    index_dims
        .iter()
        .enumerate()
        .filter_map(|(axis, &dim)| (axis != index_vector_dim).then_some(dim))
        .collect()
}

fn checked_total_len(dims: &[usize]) -> Result<usize> {
    if dims.is_empty() {
        return Ok(1);
    }
    dims.iter()
        .try_fold(1usize, |acc, &dim| acc.checked_mul(dim))
        .ok_or(StridedError::OffsetOverflow)
}

fn checked_strided_offset(base: isize, strides: &[isize], index: &[usize]) -> Result<isize> {
    let mut offset = base;
    for (&stride, &coord) in strides.iter().zip(index.iter()) {
        offset = checked_offset_add(offset, stride, coord)?;
    }
    Ok(offset)
}

fn checked_offset_add(base: isize, stride: isize, coord: usize) -> Result<isize> {
    let coord = isize::try_from(coord).map_err(|_| StridedError::OffsetOverflow)?;
    let scaled = stride
        .checked_mul(coord)
        .ok_or(StridedError::OffsetOverflow)?;
    base.checked_add(scaled).ok_or(StridedError::OffsetOverflow)
}

fn advance_col_major_index(index: &mut [usize], shape: &[usize]) {
    for axis in 0..index.len() {
        index[axis] += 1;
        if index[axis] < shape[axis] {
            return;
        }
        index[axis] = 0;
    }
}

#[cfg(feature = "parallel")]
fn fill_col_major_index(mut linear: usize, shape: &[usize], out: &mut [usize]) {
    for (axis, coord) in out.iter_mut().enumerate() {
        let dim = shape[axis];
        *coord = linear % dim;
        linear /= dim;
    }
}

struct CoordScratch {
    inline: [usize; RAW_FUSED_RANK_LIMIT],
    heap: Option<Vec<usize>>,
    len: usize,
}

impl CoordScratch {
    fn new(len: usize) -> Self {
        if len <= RAW_FUSED_RANK_LIMIT {
            Self {
                inline: [0; RAW_FUSED_RANK_LIMIT],
                heap: None,
                len,
            }
        } else {
            Self {
                inline: [0; RAW_FUSED_RANK_LIMIT],
                heap: Some(vec![0; len]),
                len,
            }
        }
    }

    fn as_mut_slice(&mut self) -> &mut [usize] {
        match &mut self.heap {
            Some(heap) => heap,
            None => &mut self.inline[..self.len],
        }
    }
}

#[cfg(test)]
mod tests {
    use super::{DynamicSlicePlan, DynamicUpdateSlicePlan};

    #[test]
    fn dynamic_slice_fast_path_is_limited_to_rank_one_contiguous_layouts() {
        let contiguous =
            DynamicSlicePlan::compile(&[16], &[1], &[1], &[1], &[8], &[1], &[8]).unwrap();
        assert!(contiguous.uses_rank_one_contiguous_path());

        let higher_rank =
            DynamicSlicePlan::compile(&[4, 4], &[1, 4], &[2], &[1], &[2, 2], &[1, 2], &[2, 2])
                .unwrap();
        assert!(!higher_rank.uses_rank_one_contiguous_path());

        let strided = DynamicSlicePlan::compile(&[16], &[2], &[1], &[1], &[8], &[2], &[8]).unwrap();
        assert!(!strided.uses_rank_one_contiguous_path());
    }

    #[test]
    fn dynamic_update_fast_path_is_limited_to_rank_one_contiguous_layouts() {
        let contiguous =
            DynamicUpdateSlicePlan::compile(&[16], &[1], &[1], &[1], &[8], &[1], &[16], &[1])
                .unwrap();
        assert!(contiguous.uses_rank_one_contiguous_path());

        let higher_rank = DynamicUpdateSlicePlan::compile(
            &[4, 4],
            &[1, 4],
            &[2],
            &[1],
            &[2, 2],
            &[1, 2],
            &[4, 4],
            &[1, 4],
        )
        .unwrap();
        assert!(!higher_rank.uses_rank_one_contiguous_path());

        let strided =
            DynamicUpdateSlicePlan::compile(&[16], &[2], &[1], &[1], &[8], &[2], &[16], &[2])
                .unwrap();
        assert!(!strided.uses_rank_one_contiguous_path());
    }
}