strided_kernel/
erased.rs

1//! Dtype-erased prepared kernel entry points.
2//!
3//! These wrappers keep dtype-specific monomorphization inside `strided-kernel`
4//! so downstream runtime crates can replay prepared kernels through stable,
5//! non-generic entry points.
6//!
7//! C ABI symbols are intentionally out of scope here. A future ABI layer must
8//! pass an explicit execution context, preserve the non-overlap contract for
9//! descriptors used by one replay call, and validate ABI dtype tags before
10//! constructing these Rust descriptors.
11//!
12use core::{mem::MaybeUninit, ops::Add};
13
14use num_complex::{Complex32, Complex64};
15use num_traits::{One, Zero};
16
17use crate::{
18    fused_elementwise_into, ConcatenatePlan, CopyPlan, DynamicSlicePlan, DynamicUpdateSlicePlan,
19    ErasedRawStridedMut, ErasedRawStridedPtr, ErasedRawStridedRef, ErasedRawStridedUninitMut,
20    ExecContext, FusedPlan, FusedScalar, GatherIndex, GatherPlan, GatherSpec, Identity,
21    KernelDType, KernelStorageElement, PadPlan, RawStridedMut, RawStridedRef, Result, ReversePlan,
22    ScatterPlan, ScatterSpec, SlicePlan, StridedError, StridedView, StridedViewMut,
23    RAW_FUSED_RANK_LIMIT,
24};
25
26const ERASED_FUSED_INPUT_LIMIT: usize = 4;
27const SERIAL_REDUCE_LANES: usize = 8;
28
29trait ReduceWriter<T> {
30    fn offset(&self) -> isize;
31    /// # Safety
32    /// The pointer may only be used within the validated destination extent.
33    unsafe fn ptr(&mut self) -> *mut T;
34    fn extent(&self) -> usize;
35    /// # Safety
36    /// The offset must be an in-bounds logical reduction destination offset.
37    unsafe fn write_at(&mut self, offset: isize, value: T) {
38        debug_assert!(offset >= 0 && (offset as usize) < self.extent());
39        // SAFETY: reduction layout validation proves the logical offset.
40        unsafe { self.ptr().offset(offset).write(value) }
41    }
42}
43
44struct RawReduceWriter<'a, T> {
45    ptr: *mut T,
46    extent: usize,
47    offset: isize,
48    _marker: core::marker::PhantomData<&'a mut [MaybeUninit<T>]>,
49}
50
51impl<'a, T> ReduceWriter<T> for RawReduceWriter<'a, T> {
52    fn offset(&self) -> isize {
53        self.offset
54    }
55    unsafe fn ptr(&mut self) -> *mut T {
56        self.ptr
57    }
58    fn extent(&self) -> usize {
59        self.extent
60    }
61}
62
63/// Runtime unary operation for [`erased_map_into`].
64#[non_exhaustive]
65#[derive(Clone, Copy, Debug, Eq, PartialEq)]
66pub enum ErasedMapOp {
67    Negate,
68    Conj,
69    Abs,
70    Sign,
71}
72
73impl ErasedMapOp {
74    const fn label(self) -> &'static str {
75        match self {
76            Self::Negate => "negate",
77            Self::Conj => "conj",
78            Self::Abs => "abs",
79            Self::Sign => "sign",
80        }
81    }
82}
83
84/// Runtime binary operation for [`erased_zip_into`].
85#[non_exhaustive]
86#[derive(Clone, Copy, Debug, Eq, PartialEq)]
87pub enum ErasedZipOp {
88    Add,
89    Subtract,
90    Multiply,
91    Divide,
92    Remainder,
93    Maximum,
94    Minimum,
95}
96
97impl ErasedZipOp {
98    const fn label(self) -> &'static str {
99        match self {
100            Self::Add => "add",
101            Self::Subtract => "subtract",
102            Self::Multiply => "multiply",
103            Self::Divide => "divide",
104            Self::Remainder => "remainder",
105            Self::Maximum => "maximum",
106            Self::Minimum => "minimum",
107        }
108    }
109}
110
111/// Apply one runtime-selected unary operation without compiling a plan.
112///
113/// The destination must not overlap the input. Real and complex dtypes support
114/// every [`ErasedMapOp`], signed integers use wrapping negate/abs semantics,
115/// and `bool` supports only [`ErasedMapOp::Conj`]. Complex absolute value has
116/// the real output contract `c32 -> f32` and `c64 -> f64`; all other supported
117/// unary operations preserve dtype.
118///
119/// # Errors
120///
121/// Returns a typed [`StridedError`] for dtype, shape, output-layout, overlap,
122/// or unsupported dtype/op contracts. Validation completes before any write.
123pub fn erased_map_into(
124    input_dtype: KernelDType,
125    op: ErasedMapOp,
126    ctx: &ExecContext,
127    dest: &mut ErasedRawStridedMut<'_>,
128    input: &ErasedRawStridedPtr<'_>,
129) -> Result<()> {
130    check_dtype(input_dtype, input.dtype())?;
131    check_dtype(map_output_dtype(input_dtype, op)?, dest.dtype())?;
132    validate_no_overlap(dest, input, 0)?;
133    let input = validated_input_ref(input)?;
134
135    let result = ctx.run(|| match (input_dtype, op) {
136        (KernelDType::C32, ErasedMapOp::Abs) => {
137            execute_one_shot_map_with::<f32, Complex32>(dest, &input, |value| value.norm())
138        }
139        (KernelDType::C64, ErasedMapOp::Abs) => {
140            execute_one_shot_map_with::<f64, Complex64>(dest, &input, |value| value.norm())
141        }
142        (KernelDType::F32, _) => execute_one_shot_map::<f32>(op, dest, &input),
143        (KernelDType::F64, _) => execute_one_shot_map::<f64>(op, dest, &input),
144        (KernelDType::I32, _) => execute_one_shot_map::<i32>(op, dest, &input),
145        (KernelDType::I64, _) => execute_one_shot_map::<i64>(op, dest, &input),
146        (KernelDType::Bool, _) => execute_one_shot_map::<bool>(op, dest, &input),
147        (KernelDType::C32, _) => execute_one_shot_map::<Complex32>(op, dest, &input),
148        (KernelDType::C64, _) => execute_one_shot_map::<Complex64>(op, dest, &input),
149        _ => Err(StridedError::UnsupportedDType {
150            dtype: input_dtype.label(),
151        }),
152    });
153    result
154}
155
156/// Apply one runtime-selected binary operation without compiling a plan.
157///
158/// The destination must not overlap either input. Real dtypes support every
159/// [`ErasedZipOp`]. Signed integers support every operation with wrapping
160/// arithmetic and a pre-write zero-divisor check. Complex dtypes support
161/// add/subtract/multiply/divide. `bool` has no binary one-shot operations.
162///
163/// # Errors
164///
165/// Returns a typed [`StridedError`] for dtype, shape, output-layout, overlap,
166/// or unsupported dtype/op contracts. Validation completes before any write.
167pub fn erased_zip_into(
168    dtype: KernelDType,
169    op: ErasedZipOp,
170    ctx: &ExecContext,
171    dest: &mut ErasedRawStridedMut<'_>,
172    lhs: &ErasedRawStridedPtr<'_>,
173    rhs: &ErasedRawStridedPtr<'_>,
174) -> Result<()> {
175    check_dtype(dtype, dest.dtype())?;
176    check_dtype(dtype, lhs.dtype())?;
177    check_dtype(dtype, rhs.dtype())?;
178    validate_no_overlap(dest, lhs, 0)?;
179    validate_no_overlap(dest, rhs, 1)?;
180    let lhs = validated_input_ref(lhs)?;
181    let rhs = validated_input_ref(rhs)?;
182
183    let result = ctx.run(|| match dtype {
184        KernelDType::F32 => execute_one_shot_zip::<f32>(op, dest, &lhs, &rhs),
185        KernelDType::F64 => execute_one_shot_zip::<f64>(op, dest, &lhs, &rhs),
186        KernelDType::I32 => execute_one_shot_zip::<i32>(op, dest, &lhs, &rhs),
187        KernelDType::I64 => execute_one_shot_zip::<i64>(op, dest, &lhs, &rhs),
188        KernelDType::Bool => execute_one_shot_zip::<bool>(op, dest, &lhs, &rhs),
189        KernelDType::C32 => execute_one_shot_zip::<Complex32>(op, dest, &lhs, &rhs),
190        KernelDType::C64 => execute_one_shot_zip::<Complex64>(op, dest, &lhs, &rhs),
191        _ => Err(StridedError::UnsupportedDType {
192            dtype: dtype.label(),
193        }),
194    });
195    result
196}
197
198/// Dtype-erased wrapper around [`CopyPlan`].
199#[derive(Clone, Debug)]
200pub struct ErasedCopyPlan {
201    dtype: KernelDType,
202    plan: CopyPlan,
203}
204
205/// Dtype-erased static-slice wrapper.
206#[derive(Clone, Debug)]
207pub struct ErasedSlicePlan {
208    dtype: KernelDType,
209    plan: SlicePlan,
210}
211
212/// Dtype-erased reverse wrapper.
213#[derive(Clone, Debug)]
214pub struct ErasedReversePlan {
215    dtype: KernelDType,
216    plan: ReversePlan,
217}
218
219/// Dtype-erased pad wrapper.
220#[derive(Clone, Debug)]
221pub struct ErasedPadPlan {
222    dtype: KernelDType,
223    plan: PadPlan,
224}
225
226/// Dtype-erased concatenate wrapper.
227#[derive(Clone, Debug)]
228pub struct ErasedConcatenatePlan {
229    dtype: KernelDType,
230    plan: ConcatenatePlan,
231}
232
233impl ErasedCopyPlan {
234    /// Compile a copy plan for one dtype and layout pair.
235    pub fn compile(
236        dtype: KernelDType,
237        dims: &[usize],
238        dst_strides: &[isize],
239        src_strides: &[isize],
240    ) -> Result<Self> {
241        Ok(Self {
242            dtype,
243            plan: CopyPlan::compile(dims, dst_strides, src_strides)?,
244        })
245    }
246
247    #[inline]
248    pub fn dtype(&self) -> KernelDType {
249        self.dtype
250    }
251
252    /// `dest = src` through a non-generic dtype-erased replay boundary.
253    pub fn execute(
254        &self,
255        ctx: &ExecContext,
256        dest: &mut ErasedRawStridedMut<'_>,
257        src: &ErasedRawStridedRef<'_>,
258    ) -> Result<()> {
259        self.check_dtype(dest.dtype())?;
260        self.check_dtype(src.dtype())?;
261
262        let result = ctx.run(|| match self.dtype {
263            KernelDType::F32 => execute_copy::<f32>(&self.plan, dest, src),
264            KernelDType::F64 => execute_copy::<f64>(&self.plan, dest, src),
265            KernelDType::I32 => execute_copy::<i32>(&self.plan, dest, src),
266            KernelDType::I64 => execute_copy::<i64>(&self.plan, dest, src),
267            KernelDType::Bool => execute_copy::<bool>(&self.plan, dest, src),
268            KernelDType::C32 => execute_copy::<Complex32>(&self.plan, dest, src),
269            KernelDType::C64 => execute_copy::<Complex64>(&self.plan, dest, src),
270            _ => Err(StridedError::UnsupportedDType {
271                dtype: self.dtype.label(),
272            }),
273        });
274        result
275    }
276
277    fn check_dtype(&self, actual: KernelDType) -> Result<()> {
278        if actual != self.dtype {
279            return Err(StridedError::DTypeMismatch {
280                expected: self.dtype.label(),
281                actual: actual.label(),
282            });
283        }
284        Ok(())
285    }
286}
287
288impl ErasedSlicePlan {
289    /// Validate and store a static slice plan for one dtype and fixed layout set.
290    #[allow(clippy::too_many_arguments)]
291    pub fn compile(
292        dtype: KernelDType,
293        operand_dims: &[usize],
294        operand_strides: &[isize],
295        dest_dims: &[usize],
296        dest_strides: &[isize],
297        starts: &[usize],
298        limits: &[usize],
299        slice_strides: &[usize],
300    ) -> Result<Self> {
301        check_static_indexing_dtype(dtype)?;
302        Ok(Self {
303            dtype,
304            plan: SlicePlan::compile(
305                operand_dims,
306                operand_strides,
307                dest_dims,
308                dest_strides,
309                starts,
310                limits,
311                slice_strides,
312            )?,
313        })
314    }
315
316    #[inline]
317    pub fn dtype(&self) -> KernelDType {
318        self.dtype
319    }
320
321    #[inline]
322    pub fn plan(&self) -> &SlicePlan {
323        &self.plan
324    }
325
326    /// Execute a static slice into an erased output descriptor.
327    pub fn execute(
328        &self,
329        ctx: &ExecContext,
330        dest: &mut ErasedRawStridedMut<'_>,
331        operand: &ErasedRawStridedRef<'_>,
332    ) -> Result<()> {
333        check_dtype(self.dtype, dest.dtype())?;
334        check_dtype(self.dtype, operand.dtype())?;
335
336        let result = ctx.run(|| match self.dtype {
337            KernelDType::F32 => execute_slice::<f32>(&self.plan, dest, operand),
338            KernelDType::F64 => execute_slice::<f64>(&self.plan, dest, operand),
339            KernelDType::I32 => execute_slice::<i32>(&self.plan, dest, operand),
340            KernelDType::I64 => execute_slice::<i64>(&self.plan, dest, operand),
341            KernelDType::Bool => execute_slice::<bool>(&self.plan, dest, operand),
342            KernelDType::C32 => execute_slice::<Complex32>(&self.plan, dest, operand),
343            KernelDType::C64 => execute_slice::<Complex64>(&self.plan, dest, operand),
344            _ => Err(StridedError::UnsupportedDType {
345                dtype: self.dtype.label(),
346            }),
347        });
348        result
349    }
350
351    /// Execute a static slice as a full overwrite of uninitialized output storage.
352    /// On success, every reachable destination slot is fully overwritten;
353    /// unreachable holes are neither read nor initialized. Validation errors
354    /// are returned before any destination write. A panic during execution
355    /// may leave a partially initialized `MaybeUninit` destination, which is
356    /// still safely droppable; no readable value is promised for unwritten
357    /// reachable slots.
358    pub fn execute_uninit(
359        &self,
360        ctx: &ExecContext,
361        dest: &mut ErasedRawStridedUninitMut<'_>,
362        operand: &ErasedRawStridedPtr<'_>,
363    ) -> Result<()> {
364        check_dtype(self.dtype, dest.dtype())?;
365        check_dtype(self.dtype, operand.dtype())?;
366        validate_uninit_no_overlap(dest, operand, 0)?;
367        let operand = validated_input_ref(operand)?;
368
369        ctx.run(|| match self.dtype {
370            KernelDType::F32 => execute_slice_uninit::<f32>(&self.plan, dest, &operand),
371            KernelDType::F64 => execute_slice_uninit::<f64>(&self.plan, dest, &operand),
372            KernelDType::I32 => execute_slice_uninit::<i32>(&self.plan, dest, &operand),
373            KernelDType::I64 => execute_slice_uninit::<i64>(&self.plan, dest, &operand),
374            KernelDType::Bool => execute_slice_uninit::<bool>(&self.plan, dest, &operand),
375            KernelDType::C32 => execute_slice_uninit::<Complex32>(&self.plan, dest, &operand),
376            KernelDType::C64 => execute_slice_uninit::<Complex64>(&self.plan, dest, &operand),
377            _ => Err(StridedError::UnsupportedDType {
378                dtype: self.dtype.label(),
379            }),
380        })
381    }
382}
383
384impl ErasedReversePlan {
385    /// Validate and store a reverse plan for one dtype and fixed layout set.
386    pub fn compile(
387        dtype: KernelDType,
388        operand_dims: &[usize],
389        operand_strides: &[isize],
390        dest_strides: &[isize],
391        axes: &[usize],
392    ) -> Result<Self> {
393        check_static_indexing_dtype(dtype)?;
394        Ok(Self {
395            dtype,
396            plan: ReversePlan::compile(operand_dims, operand_strides, dest_strides, axes)?,
397        })
398    }
399
400    #[inline]
401    pub fn dtype(&self) -> KernelDType {
402        self.dtype
403    }
404
405    #[inline]
406    pub fn plan(&self) -> &ReversePlan {
407        &self.plan
408    }
409
410    /// Execute a reverse into an erased output descriptor.
411    pub fn execute(
412        &self,
413        ctx: &ExecContext,
414        dest: &mut ErasedRawStridedMut<'_>,
415        operand: &ErasedRawStridedRef<'_>,
416    ) -> Result<()> {
417        check_dtype(self.dtype, dest.dtype())?;
418        check_dtype(self.dtype, operand.dtype())?;
419
420        let result = ctx.run(|| match self.dtype {
421            KernelDType::F32 => execute_reverse::<f32>(&self.plan, dest, operand),
422            KernelDType::F64 => execute_reverse::<f64>(&self.plan, dest, operand),
423            KernelDType::I32 => execute_reverse::<i32>(&self.plan, dest, operand),
424            KernelDType::I64 => execute_reverse::<i64>(&self.plan, dest, operand),
425            KernelDType::Bool => execute_reverse::<bool>(&self.plan, dest, operand),
426            KernelDType::C32 => execute_reverse::<Complex32>(&self.plan, dest, operand),
427            KernelDType::C64 => execute_reverse::<Complex64>(&self.plan, dest, operand),
428            _ => Err(StridedError::UnsupportedDType {
429                dtype: self.dtype.label(),
430            }),
431        });
432        result
433    }
434
435    /// Execute reverse as a full overwrite of uninitialized output storage.
436    /// On success, every reachable destination slot is fully overwritten;
437    /// unreachable holes are neither read nor initialized. Validation errors
438    /// are returned before any destination write. A panic during execution
439    /// may leave a partially initialized `MaybeUninit` destination, which is
440    /// still safely droppable; no readable value is promised for unwritten
441    /// reachable slots.
442    pub fn execute_uninit(
443        &self,
444        ctx: &ExecContext,
445        dest: &mut ErasedRawStridedUninitMut<'_>,
446        operand: &ErasedRawStridedPtr<'_>,
447    ) -> Result<()> {
448        check_dtype(self.dtype, dest.dtype())?;
449        check_dtype(self.dtype, operand.dtype())?;
450        validate_uninit_no_overlap(dest, operand, 0)?;
451        let operand = validated_input_ref(operand)?;
452
453        ctx.run(|| match self.dtype {
454            KernelDType::F32 => execute_reverse_uninit::<f32>(&self.plan, dest, &operand),
455            KernelDType::F64 => execute_reverse_uninit::<f64>(&self.plan, dest, &operand),
456            KernelDType::I32 => execute_reverse_uninit::<i32>(&self.plan, dest, &operand),
457            KernelDType::I64 => execute_reverse_uninit::<i64>(&self.plan, dest, &operand),
458            KernelDType::Bool => execute_reverse_uninit::<bool>(&self.plan, dest, &operand),
459            KernelDType::C32 => execute_reverse_uninit::<Complex32>(&self.plan, dest, &operand),
460            KernelDType::C64 => execute_reverse_uninit::<Complex64>(&self.plan, dest, &operand),
461            _ => Err(StridedError::UnsupportedDType {
462                dtype: self.dtype.label(),
463            }),
464        })
465    }
466}
467
468impl ErasedPadPlan {
469    /// Validate and store a pad plan for one dtype and fixed layout set.
470    #[allow(clippy::too_many_arguments)]
471    pub fn compile(
472        dtype: KernelDType,
473        operand_dims: &[usize],
474        operand_strides: &[isize],
475        dest_dims: &[usize],
476        dest_strides: &[isize],
477        edge_padding_low: &[i64],
478        edge_padding_high: &[i64],
479        interior_padding: &[i64],
480    ) -> Result<Self> {
481        check_static_indexing_dtype(dtype)?;
482        Ok(Self {
483            dtype,
484            plan: PadPlan::compile(
485                operand_dims,
486                operand_strides,
487                dest_dims,
488                dest_strides,
489                edge_padding_low,
490                edge_padding_high,
491                interior_padding,
492            )?,
493        })
494    }
495
496    #[inline]
497    pub fn dtype(&self) -> KernelDType {
498        self.dtype
499    }
500
501    #[inline]
502    pub fn plan(&self) -> &PadPlan {
503        &self.plan
504    }
505
506    /// Execute pad into an erased output descriptor using one dtype scalar as fill.
507    pub fn execute(
508        &self,
509        ctx: &ExecContext,
510        dest: &mut ErasedRawStridedMut<'_>,
511        operand: &ErasedRawStridedRef<'_>,
512        fill: &[u8],
513    ) -> Result<()> {
514        check_dtype(self.dtype, dest.dtype())?;
515        check_dtype(self.dtype, operand.dtype())?;
516        validate_scalar_bytes(self.dtype, fill)?;
517
518        let result = ctx.run(|| match self.dtype {
519            KernelDType::F32 => execute_pad::<f32>(&self.plan, dest, operand, fill),
520            KernelDType::F64 => execute_pad::<f64>(&self.plan, dest, operand, fill),
521            KernelDType::I32 => execute_pad::<i32>(&self.plan, dest, operand, fill),
522            KernelDType::I64 => execute_pad::<i64>(&self.plan, dest, operand, fill),
523            KernelDType::Bool => execute_pad::<bool>(&self.plan, dest, operand, fill),
524            KernelDType::C32 => execute_pad::<Complex32>(&self.plan, dest, operand, fill),
525            KernelDType::C64 => execute_pad::<Complex64>(&self.plan, dest, operand, fill),
526            _ => Err(StridedError::UnsupportedDType {
527                dtype: self.dtype.label(),
528            }),
529        });
530        result
531    }
532
533    /// Execute pad as a full overwrite of uninitialized output storage.
534    /// On success, every reachable destination slot is fully overwritten;
535    /// unreachable holes are neither read nor initialized. Validation errors
536    /// are returned before any destination write. A panic during execution
537    /// may leave a partially initialized `MaybeUninit` destination, which is
538    /// still safely droppable; no readable value is promised for unwritten
539    /// reachable slots.
540    pub fn execute_uninit(
541        &self,
542        ctx: &ExecContext,
543        dest: &mut ErasedRawStridedUninitMut<'_>,
544        operand: &ErasedRawStridedPtr<'_>,
545        fill: &[u8],
546    ) -> Result<()> {
547        check_dtype(self.dtype, dest.dtype())?;
548        check_dtype(self.dtype, operand.dtype())?;
549        validate_scalar_bytes(self.dtype, fill)?;
550        validate_uninit_no_overlap(dest, operand, 0)?;
551        let operand = validated_input_ref(operand)?;
552
553        ctx.run(|| match self.dtype {
554            KernelDType::F32 => execute_pad_uninit::<f32>(&self.plan, dest, &operand, fill),
555            KernelDType::F64 => execute_pad_uninit::<f64>(&self.plan, dest, &operand, fill),
556            KernelDType::I32 => execute_pad_uninit::<i32>(&self.plan, dest, &operand, fill),
557            KernelDType::I64 => execute_pad_uninit::<i64>(&self.plan, dest, &operand, fill),
558            KernelDType::Bool => execute_pad_uninit::<bool>(&self.plan, dest, &operand, fill),
559            KernelDType::C32 => execute_pad_uninit::<Complex32>(&self.plan, dest, &operand, fill),
560            KernelDType::C64 => execute_pad_uninit::<Complex64>(&self.plan, dest, &operand, fill),
561            _ => Err(StridedError::UnsupportedDType {
562                dtype: self.dtype.label(),
563            }),
564        })
565    }
566}
567
568impl ErasedConcatenatePlan {
569    /// Validate and store a concatenate plan for one dtype and fixed layout set.
570    pub fn compile(
571        dtype: KernelDType,
572        input_dims: &[&[usize]],
573        input_strides: &[&[isize]],
574        dest_dims: &[usize],
575        dest_strides: &[isize],
576        axis: usize,
577    ) -> Result<Self> {
578        check_static_indexing_dtype(dtype)?;
579        Ok(Self {
580            dtype,
581            plan: ConcatenatePlan::compile(
582                input_dims,
583                input_strides,
584                dest_dims,
585                dest_strides,
586                axis,
587            )?,
588        })
589    }
590
591    #[inline]
592    pub fn dtype(&self) -> KernelDType {
593        self.dtype
594    }
595
596    #[inline]
597    pub fn plan(&self) -> &ConcatenatePlan {
598        &self.plan
599    }
600
601    /// Execute concatenate into an erased output descriptor.
602    pub fn execute(
603        &self,
604        ctx: &ExecContext,
605        dest: &mut ErasedRawStridedMut<'_>,
606        inputs: &[ErasedRawStridedRef<'_>],
607    ) -> Result<()> {
608        check_dtype(self.dtype, dest.dtype())?;
609        for input in inputs {
610            check_dtype(self.dtype, input.dtype())?;
611        }
612
613        let result = ctx.run(|| match self.dtype {
614            KernelDType::F32 => execute_concatenate::<f32>(&self.plan, dest, inputs),
615            KernelDType::F64 => execute_concatenate::<f64>(&self.plan, dest, inputs),
616            KernelDType::I32 => execute_concatenate::<i32>(&self.plan, dest, inputs),
617            KernelDType::I64 => execute_concatenate::<i64>(&self.plan, dest, inputs),
618            KernelDType::Bool => execute_concatenate::<bool>(&self.plan, dest, inputs),
619            KernelDType::C32 => execute_concatenate::<Complex32>(&self.plan, dest, inputs),
620            KernelDType::C64 => execute_concatenate::<Complex64>(&self.plan, dest, inputs),
621            _ => Err(StridedError::UnsupportedDType {
622                dtype: self.dtype.label(),
623            }),
624        });
625        result
626    }
627
628    /// Execute concatenate as a full overwrite of uninitialized output storage.
629    /// On success, every reachable destination slot is fully overwritten;
630    /// unreachable holes are neither read nor initialized. Validation errors
631    /// are returned before any destination write. A panic during execution
632    /// may leave a partially initialized `MaybeUninit` destination, which is
633    /// still safely droppable; no readable value is promised for unwritten
634    /// reachable slots.
635    pub fn execute_uninit(
636        &self,
637        ctx: &ExecContext,
638        dest: &mut ErasedRawStridedUninitMut<'_>,
639        inputs: &[ErasedRawStridedPtr<'_>],
640    ) -> Result<()> {
641        check_dtype(self.dtype, dest.dtype())?;
642        if inputs.len() != self.plan.input_count() {
643            return Err(StridedError::RankMismatch(
644                inputs.len(),
645                self.plan.input_count(),
646            ));
647        }
648        for input in inputs {
649            check_dtype(self.dtype, input.dtype())?;
650        }
651        for (position, input) in inputs.iter().enumerate() {
652            validate_uninit_no_overlap(dest, input, position)?;
653        }
654        for input in inputs {
655            validated_input_ref(input)?;
656        }
657
658        ctx.run(|| match self.dtype {
659            KernelDType::F32 => execute_concatenate_uninit::<f32>(&self.plan, dest, inputs),
660            KernelDType::F64 => execute_concatenate_uninit::<f64>(&self.plan, dest, inputs),
661            KernelDType::I32 => execute_concatenate_uninit::<i32>(&self.plan, dest, inputs),
662            KernelDType::I64 => execute_concatenate_uninit::<i64>(&self.plan, dest, inputs),
663            KernelDType::Bool => execute_concatenate_uninit::<bool>(&self.plan, dest, inputs),
664            KernelDType::C32 => execute_concatenate_uninit::<Complex32>(&self.plan, dest, inputs),
665            KernelDType::C64 => execute_concatenate_uninit::<Complex64>(&self.plan, dest, inputs),
666            _ => Err(StridedError::UnsupportedDType {
667                dtype: self.dtype.label(),
668            }),
669        })
670    }
671}
672
673/// Dtype-erased single-output wrapper around [`FusedPlan`].
674///
675/// This is the erased replay boundary for unary map and zip-map elementwise
676/// families. It supports the same runtime op-code vocabulary as [`FusedPlan`],
677/// but only for the scalar dtypes currently implementing [`FusedScalar`].
678#[derive(Clone, Debug)]
679pub struct ErasedFusedPlan {
680    dtype: KernelDType,
681    plan: FusedPlan,
682}
683
684/// Runtime reduction operation for dtype-erased full reductions.
685#[non_exhaustive]
686#[derive(Clone, Copy, Debug, Eq, PartialEq)]
687pub enum ReduceOp {
688    Sum,
689    Product,
690    /// Sum of same-dtype rounded squares.
691    ///
692    /// Each input is first multiplied by itself without FMA contraction, then
693    /// accumulated under the same association policy as [`Self::Sum`].
694    SumSquares,
695}
696
697/// Dtype-erased reduction wrapper.
698///
699/// This is the erased replay boundary for full-tensor scalar reductions and
700/// axis reductions with a fixed output layout. It supports only operations with
701/// an unambiguous identity value in the selected dtype.
702#[derive(Clone, Debug)]
703pub struct ErasedReducePlan {
704    dtype: KernelDType,
705    op: ReduceOp,
706    layout: ReduceLayout,
707}
708
709#[derive(Clone, Debug)]
710enum ReduceLayout {
711    Full {
712        dims: Vec<usize>,
713        src_strides: Vec<isize>,
714    },
715    Axes {
716        src_dims: Vec<usize>,
717        src_strides: Vec<isize>,
718        dest_dims: Vec<usize>,
719        dest_strides: Vec<isize>,
720        axes: Vec<usize>,
721        kept_axes: Vec<usize>,
722        reduce_dims: Vec<usize>,
723        dest_total: usize,
724        reduce_total: usize,
725    },
726}
727
728impl ReduceLayout {
729    fn src_dims(&self) -> &[usize] {
730        match self {
731            Self::Full { dims, .. } => dims,
732            Self::Axes { src_dims, .. } => src_dims,
733        }
734    }
735
736    fn src_strides(&self) -> &[isize] {
737        match self {
738            Self::Full { src_strides, .. } | Self::Axes { src_strides, .. } => src_strides,
739        }
740    }
741
742    fn check_src_layout(&self, src: &ErasedRawStridedRef<'_>) -> Result<()> {
743        if src.dims() != self.src_dims() || src.strides() != self.src_strides() {
744            return Err(StridedError::PlanLayoutMismatch);
745        }
746        Ok(())
747    }
748}
749
750#[derive(Clone, Copy, Debug)]
751struct AxesLayout<'a> {
752    src_dims: &'a [usize],
753    src_strides: &'a [isize],
754    dest_dims: &'a [usize],
755    dest_strides: &'a [isize],
756    axes: &'a [usize],
757    kept_axes: &'a [usize],
758    reduce_dims: &'a [usize],
759    dest_total: usize,
760    reduce_total: usize,
761}
762
763/// Dtype-erased gather wrapper.
764///
765/// This is the erased replay boundary for indexed reads. Value buffers use the
766/// configured value dtype, while the index descriptor must use `i32` or `i64`.
767#[derive(Clone, Debug)]
768pub struct ErasedGatherPlan {
769    dtype: KernelDType,
770    index_dtype: KernelDType,
771    plan: GatherPlan,
772}
773
774/// Dtype-erased fixed-window dynamic-slice wrapper.
775#[derive(Clone, Debug)]
776pub struct ErasedDynamicSlicePlan {
777    dtype: KernelDType,
778    index_dtype: KernelDType,
779    plan: DynamicSlicePlan,
780}
781
782/// Dtype-erased dynamic-update-slice wrapper.
783#[derive(Clone, Debug)]
784pub struct ErasedDynamicUpdateSlicePlan {
785    dtype: KernelDType,
786    index_dtype: KernelDType,
787    plan: DynamicUpdateSlicePlan,
788}
789
790/// Dtype-erased additive scatter wrapper.
791#[derive(Clone, Debug)]
792pub struct ErasedScatterPlan {
793    dtype: KernelDType,
794    index_dtype: KernelDType,
795    plan: ScatterPlan,
796}
797
798impl ErasedFusedPlan {
799    /// Validate and store a single-output fused elementwise plan for one dtype.
800    pub fn compile(dtype: KernelDType, plan: FusedPlan) -> Result<Self> {
801        check_fused_dtype(dtype)?;
802        if plan.input_count == 0 || plan.input_count > ERASED_FUSED_INPUT_LIMIT {
803            return Err(StridedError::UnsupportedArity {
804                arity: plan.input_count,
805                max: ERASED_FUSED_INPUT_LIMIT,
806            });
807        }
808        if plan.outputs.len() != 1 {
809            return Err(StridedError::RankMismatch(plan.outputs.len(), 1));
810        }
811        validate_fused_plan_for_dtype(dtype, &plan)?;
812        Ok(Self { dtype, plan })
813    }
814
815    #[inline]
816    pub fn dtype(&self) -> KernelDType {
817        self.dtype
818    }
819
820    #[inline]
821    pub fn plan(&self) -> &FusedPlan {
822        &self.plan
823    }
824
825    /// Execute a single-output fused elementwise plan through erased descriptors.
826    pub fn execute(
827        &self,
828        ctx: &ExecContext,
829        dest: &mut ErasedRawStridedMut<'_>,
830        inputs: &[ErasedRawStridedRef<'_>],
831    ) -> Result<()> {
832        if inputs.len() != self.plan.input_count {
833            return Err(StridedError::RankMismatch(
834                inputs.len(),
835                self.plan.input_count,
836            ));
837        }
838        check_dtype(self.dtype, dest.dtype())?;
839        for input in inputs {
840            check_dtype(self.dtype, input.dtype())?;
841        }
842
843        let result = match self.dtype {
844            KernelDType::F32 => execute_fused::<f32>(&self.plan, ctx, dest, inputs),
845            KernelDType::F64 => execute_fused::<f64>(&self.plan, ctx, dest, inputs),
846            KernelDType::I32 => execute_fused::<i32>(&self.plan, ctx, dest, inputs),
847            KernelDType::I64 => execute_fused::<i64>(&self.plan, ctx, dest, inputs),
848            KernelDType::Bool => execute_fused::<bool>(&self.plan, ctx, dest, inputs),
849            KernelDType::C32 => execute_fused::<Complex32>(&self.plan, ctx, dest, inputs),
850            KernelDType::C64 => execute_fused::<Complex64>(&self.plan, ctx, dest, inputs),
851            _ => Err(StridedError::UnsupportedDType {
852                dtype: self.dtype.label(),
853            }),
854        };
855        result
856    }
857
858    /// Execute a single-output fused plan into fully overwritten uninitialized storage.
859    ///
860    /// Dtype, shape, destination injectivity, bounds, and input/output overlap
861    /// are validated before any shared typed input descriptor is formed or any
862    /// destination byte is written. On `Ok(())`, every logical destination
863    /// element is initialized. An error leaves the destination untouched; a
864    /// panic during execution may leave partial initialization, but the backing
865    /// `MaybeUninit` storage remains safe to drop.
866    ///
867    /// # Errors
868    ///
869    /// Returns a typed dtype, input-count, shape, bounds, destination
870    /// injectivity, unsupported-operation, or input/output-overlap error. All
871    /// error-producing validation completes before execution starts.
872    /// On success, every reachable destination slot is fully overwritten;
873    /// unreachable holes are neither read nor initialized. Validation errors
874    /// are returned before any destination write. A panic during execution
875    /// may leave a partially initialized `MaybeUninit` destination, which is
876    /// still safely droppable; no readable value is promised for unwritten
877    /// reachable slots.
878    pub fn execute_uninit(
879        &self,
880        ctx: &ExecContext,
881        dest: &mut ErasedRawStridedUninitMut<'_>,
882        inputs: &[ErasedRawStridedPtr<'_>],
883    ) -> Result<()> {
884        if inputs.len() != self.plan.input_count {
885            return Err(StridedError::RankMismatch(
886                inputs.len(),
887                self.plan.input_count,
888            ));
889        }
890        check_dtype(self.dtype, dest.dtype())?;
891        for input in inputs {
892            check_dtype(self.dtype, input.dtype())?;
893        }
894        for (index, input) in inputs.iter().enumerate() {
895            validate_uninit_no_overlap(dest, input, index)?;
896            if input.dims() != dest.dims() {
897                return Err(StridedError::ShapeMismatch(
898                    input.dims().to_vec(),
899                    dest.dims().to_vec(),
900                ));
901            }
902        }
903        let validated = crate::map_view::validate_destination_layout_without_alloc(
904            dest.dims(),
905            dest.strides(),
906        )?;
907
908        let run = |dest: &mut ErasedRawStridedUninitMut<'_>| match self.dtype {
909            KernelDType::F32 => execute_fused_uninit_ptrs::<f32>(
910                &self.plan,
911                dest,
912                inputs,
913                ctx.is_serial(),
914                validated,
915            ),
916            KernelDType::F64 => execute_fused_uninit_ptrs::<f64>(
917                &self.plan,
918                dest,
919                inputs,
920                ctx.is_serial(),
921                validated,
922            ),
923            KernelDType::I32 => execute_fused_uninit_ptrs::<i32>(
924                &self.plan,
925                dest,
926                inputs,
927                ctx.is_serial(),
928                validated,
929            ),
930            KernelDType::I64 => execute_fused_uninit_ptrs::<i64>(
931                &self.plan,
932                dest,
933                inputs,
934                ctx.is_serial(),
935                validated,
936            ),
937            KernelDType::Bool => execute_fused_uninit_ptrs::<bool>(
938                &self.plan,
939                dest,
940                inputs,
941                ctx.is_serial(),
942                validated,
943            ),
944            KernelDType::C32 => execute_fused_uninit_ptrs::<Complex32>(
945                &self.plan,
946                dest,
947                inputs,
948                ctx.is_serial(),
949                validated,
950            ),
951            KernelDType::C64 => execute_fused_uninit_ptrs::<Complex64>(
952                &self.plan,
953                dest,
954                inputs,
955                ctx.is_serial(),
956                validated,
957            ),
958            _ => Err(StridedError::UnsupportedDType {
959                dtype: self.dtype.label(),
960            }),
961        };
962        if ctx.is_serial() {
963            run(dest)
964        } else {
965            ctx.run(|| run(dest))
966        }
967    }
968}
969
970impl ErasedReducePlan {
971    /// Validate and store a full-reduction plan for one dtype and source layout.
972    pub fn compile(
973        dtype: KernelDType,
974        op: ReduceOp,
975        dims: &[usize],
976        src_strides: &[isize],
977    ) -> Result<Self> {
978        check_reduce_op_dtype(dtype, op)?;
979        if dims.len() != src_strides.len() {
980            return Err(StridedError::StrideLengthMismatch);
981        }
982        checked_total_len(dims)?;
983        Ok(Self {
984            dtype,
985            op,
986            layout: ReduceLayout::Full {
987                dims: dims.to_vec(),
988                src_strides: src_strides.to_vec(),
989            },
990        })
991    }
992
993    /// Validate and store an axis-reduction plan for one dtype and fixed source/output layouts.
994    ///
995    /// `axes` names the source axes reduced away. Output dimensions must be the
996    /// remaining source dimensions in source-axis order. When all axes are
997    /// reduced, any output layout with exactly one reachable element is accepted.
998    #[allow(clippy::too_many_arguments)]
999    pub fn compile_axes(
1000        dtype: KernelDType,
1001        op: ReduceOp,
1002        src_dims: &[usize],
1003        src_strides: &[isize],
1004        dest_dims: &[usize],
1005        dest_strides: &[isize],
1006        axes: &[usize],
1007    ) -> Result<Self> {
1008        check_reduce_op_dtype(dtype, op)?;
1009        if src_dims.len() != src_strides.len() || dest_dims.len() != dest_strides.len() {
1010            return Err(StridedError::StrideLengthMismatch);
1011        }
1012        checked_total_len(src_dims)?;
1013        let dest_total = checked_total_len(dest_dims)?;
1014        if !crate::fused::is_injective_layout(dest_dims, dest_strides) {
1015            return Err(StridedError::NonInjectiveOutputLayout);
1016        }
1017        validate_unique_axes(axes, src_dims.len())?;
1018
1019        let kept_axes: Vec<usize> = (0..src_dims.len())
1020            .filter(|axis| !axes.contains(axis))
1021            .collect();
1022        let expected_dest_dims: Vec<usize> = kept_axes.iter().map(|&axis| src_dims[axis]).collect();
1023        if expected_dest_dims.is_empty() {
1024            if dest_total != 1 {
1025                return Err(StridedError::ShapeMismatch(
1026                    dest_dims.to_vec(),
1027                    expected_dest_dims,
1028                ));
1029            }
1030        } else if dest_dims != expected_dest_dims.as_slice() {
1031            return Err(StridedError::ShapeMismatch(
1032                dest_dims.to_vec(),
1033                expected_dest_dims,
1034            ));
1035        }
1036
1037        let reduce_dims = axes.iter().map(|&axis| src_dims[axis]).collect::<Vec<_>>();
1038        let reduce_total = checked_total_len(&reduce_dims)?;
1039        Ok(Self {
1040            dtype,
1041            op,
1042            layout: ReduceLayout::Axes {
1043                src_dims: src_dims.to_vec(),
1044                src_strides: src_strides.to_vec(),
1045                dest_dims: dest_dims.to_vec(),
1046                dest_strides: dest_strides.to_vec(),
1047                axes: axes.to_vec(),
1048                kept_axes,
1049                reduce_dims,
1050                dest_total,
1051                reduce_total,
1052            },
1053        })
1054    }
1055
1056    #[inline]
1057    pub fn dtype(&self) -> KernelDType {
1058        self.dtype
1059    }
1060
1061    #[inline]
1062    pub fn op(&self) -> ReduceOp {
1063        self.op
1064    }
1065
1066    /// Execute the reduction into an erased output descriptor.
1067    pub fn execute(
1068        &self,
1069        ctx: &ExecContext,
1070        dest: &mut ErasedRawStridedMut<'_>,
1071        src: &ErasedRawStridedRef<'_>,
1072    ) -> Result<()> {
1073        check_dtype(self.dtype, dest.dtype())?;
1074        check_dtype(self.dtype, src.dtype())?;
1075        self.layout.check_src_layout(src)?;
1076        match &self.layout {
1077            ReduceLayout::Full { .. } => {
1078                let dest_len = checked_total_len(dest.dims())?;
1079                if dest_len != 1 {
1080                    return Err(StridedError::RankMismatch(dest_len, 1));
1081                }
1082            }
1083            ReduceLayout::Axes {
1084                dest_dims,
1085                dest_strides,
1086                ..
1087            } => {
1088                if dest.dims() != dest_dims.as_slice() || dest.strides() != dest_strides.as_slice()
1089                {
1090                    return Err(StridedError::PlanLayoutMismatch);
1091                }
1092            }
1093        }
1094
1095        let result = match self.dtype {
1096            KernelDType::F32 => {
1097                let mut writer = reduce_writer::<f32>(dest)?;
1098                dispatch_reduce::<f32, _>(self.op, &self.layout, ctx, &mut writer, src)
1099            }
1100            KernelDType::F64 => {
1101                let mut writer = reduce_writer::<f64>(dest)?;
1102                dispatch_reduce::<f64, _>(self.op, &self.layout, ctx, &mut writer, src)
1103            }
1104            KernelDType::I32 => {
1105                let mut writer = reduce_writer::<i32>(dest)?;
1106                dispatch_reduce::<i32, _>(self.op, &self.layout, ctx, &mut writer, src)
1107            }
1108            KernelDType::I64 => {
1109                let mut writer = reduce_writer::<i64>(dest)?;
1110                dispatch_reduce::<i64, _>(self.op, &self.layout, ctx, &mut writer, src)
1111            }
1112            KernelDType::C32 => {
1113                let mut writer = reduce_writer::<Complex32>(dest)?;
1114                dispatch_reduce::<Complex32, _>(self.op, &self.layout, ctx, &mut writer, src)
1115            }
1116            KernelDType::C64 => {
1117                let mut writer = reduce_writer::<Complex64>(dest)?;
1118                dispatch_reduce::<Complex64, _>(self.op, &self.layout, ctx, &mut writer, src)
1119            }
1120            _ => Err(StridedError::UnsupportedDType {
1121                dtype: self.dtype.label(),
1122            }),
1123        };
1124        result
1125    }
1126
1127    /// On success, every reachable destination slot is fully overwritten;
1128    /// unreachable holes are neither read nor initialized. Validation errors
1129    /// are returned before any destination write. A panic during execution
1130    /// may leave a partially initialized `MaybeUninit` destination, which is
1131    /// still safely droppable; no readable value is promised for unwritten
1132    /// reachable slots.
1133    pub fn execute_uninit(
1134        &self,
1135        ctx: &ExecContext,
1136        dest: &mut ErasedRawStridedUninitMut<'_>,
1137        src: &ErasedRawStridedPtr<'_>,
1138    ) -> Result<()> {
1139        check_dtype(self.dtype, dest.dtype())?;
1140        check_dtype(self.dtype, src.dtype())?;
1141        validate_uninit_no_overlap(dest, src, 0)?;
1142        let src = validated_input_ref(src)?;
1143        self.layout.check_src_layout(&src)?;
1144        match &self.layout {
1145            ReduceLayout::Full { .. } => {
1146                let total = checked_total_len(dest.dims())?;
1147                if total != 1 {
1148                    return Err(StridedError::RankMismatch(total, 1));
1149                }
1150            }
1151            ReduceLayout::Axes {
1152                dest_dims,
1153                dest_strides,
1154                ..
1155            } => {
1156                if dest.dims() != dest_dims.as_slice() || dest.strides() != dest_strides.as_slice()
1157                {
1158                    return Err(StridedError::PlanLayoutMismatch);
1159                }
1160            }
1161        }
1162        macro_rules! run {
1163            ($ty:ty) => {{
1164                let mut writer = reduce_uninit_writer::<$ty>(dest)?;
1165                dispatch_reduce::<$ty, _>(self.op, &self.layout, ctx, &mut writer, &src)
1166            }};
1167        }
1168        match self.dtype {
1169            KernelDType::F32 => run!(f32),
1170            KernelDType::F64 => run!(f64),
1171            KernelDType::I32 => run!(i32),
1172            KernelDType::I64 => run!(i64),
1173            KernelDType::C32 => run!(Complex32),
1174            KernelDType::C64 => run!(Complex64),
1175            _ => Err(StridedError::UnsupportedDType {
1176                dtype: self.dtype.label(),
1177            }),
1178        }
1179    }
1180}
1181
1182impl ErasedGatherPlan {
1183    /// Validate and store a gather plan for one value dtype, index dtype, and layout set.
1184    #[allow(clippy::too_many_arguments)]
1185    pub fn compile(
1186        dtype: KernelDType,
1187        index_dtype: KernelDType,
1188        operand_dims: &[usize],
1189        operand_strides: &[isize],
1190        index_dims: &[usize],
1191        index_strides: &[isize],
1192        dest_dims: &[usize],
1193        dest_strides: &[isize],
1194        spec: GatherSpec,
1195    ) -> Result<Self> {
1196        check_index_dtype(index_dtype)?;
1197        check_gather_value_dtype(dtype)?;
1198        Ok(Self {
1199            dtype,
1200            index_dtype,
1201            plan: GatherPlan::compile(
1202                operand_dims,
1203                operand_strides,
1204                index_dims,
1205                index_strides,
1206                dest_dims,
1207                dest_strides,
1208                spec,
1209            )?,
1210        })
1211    }
1212
1213    #[inline]
1214    pub fn dtype(&self) -> KernelDType {
1215        self.dtype
1216    }
1217
1218    #[inline]
1219    pub fn index_dtype(&self) -> KernelDType {
1220        self.index_dtype
1221    }
1222
1223    #[inline]
1224    pub fn plan(&self) -> &GatherPlan {
1225        &self.plan
1226    }
1227
1228    /// Execute an indexed read into an erased output descriptor.
1229    pub fn execute(
1230        &self,
1231        ctx: &ExecContext,
1232        dest: &mut ErasedRawStridedMut<'_>,
1233        operand: &ErasedRawStridedRef<'_>,
1234        start_indices: &ErasedRawStridedRef<'_>,
1235    ) -> Result<()> {
1236        check_dtype(self.dtype, dest.dtype())?;
1237        check_dtype(self.dtype, operand.dtype())?;
1238        check_dtype(self.index_dtype, start_indices.dtype())?;
1239
1240        let result = ctx.run(|| match self.dtype {
1241            KernelDType::F32 => dispatch_gather_index::<f32>(
1242                &self.plan,
1243                self.index_dtype,
1244                dest,
1245                &operand,
1246                &start_indices,
1247            ),
1248            KernelDType::F64 => dispatch_gather_index::<f64>(
1249                &self.plan,
1250                self.index_dtype,
1251                dest,
1252                operand,
1253                start_indices,
1254            ),
1255            KernelDType::I32 => dispatch_gather_index::<i32>(
1256                &self.plan,
1257                self.index_dtype,
1258                dest,
1259                operand,
1260                start_indices,
1261            ),
1262            KernelDType::I64 => dispatch_gather_index::<i64>(
1263                &self.plan,
1264                self.index_dtype,
1265                dest,
1266                operand,
1267                start_indices,
1268            ),
1269            KernelDType::Bool => dispatch_gather_index::<bool>(
1270                &self.plan,
1271                self.index_dtype,
1272                dest,
1273                operand,
1274                start_indices,
1275            ),
1276            KernelDType::C32 => dispatch_gather_index::<Complex32>(
1277                &self.plan,
1278                self.index_dtype,
1279                dest,
1280                operand,
1281                start_indices,
1282            ),
1283            KernelDType::C64 => dispatch_gather_index::<Complex64>(
1284                &self.plan,
1285                self.index_dtype,
1286                dest,
1287                operand,
1288                start_indices,
1289            ),
1290            _ => Err(StridedError::UnsupportedDType {
1291                dtype: self.dtype.label(),
1292            }),
1293        });
1294        result
1295    }
1296
1297    /// Execute gather into a destination whose reachable slots may be
1298    /// uninitialized. All validation precedes the first destination write.
1299    /// On success, every reachable destination slot is fully overwritten;
1300    /// unreachable holes are neither read nor initialized. Validation errors
1301    /// are returned before any destination write. A panic during execution
1302    /// may leave a partially initialized `MaybeUninit` destination, which is
1303    /// still safely droppable; no readable value is promised for unwritten
1304    /// reachable slots.
1305    pub fn execute_uninit(
1306        &self,
1307        ctx: &ExecContext,
1308        dest: &mut ErasedRawStridedUninitMut<'_>,
1309        operand: &ErasedRawStridedPtr<'_>,
1310        start_indices: &ErasedRawStridedPtr<'_>,
1311    ) -> Result<()> {
1312        check_dtype(self.dtype, dest.dtype())?;
1313        check_dtype(self.dtype, operand.dtype())?;
1314        check_dtype(self.index_dtype, start_indices.dtype())?;
1315        validate_uninit_no_overlap(dest, operand, 0)?;
1316        validate_uninit_no_overlap(dest, start_indices, 1)?;
1317        let operand = &validated_input_ref(operand)?;
1318        let start_indices = &validated_input_ref(start_indices)?;
1319        let run = |dest: &mut ErasedRawStridedUninitMut<'_>| match self.dtype {
1320            KernelDType::F32 => execute_gather_uninit_dispatch::<f32>(
1321                &self.plan,
1322                self.index_dtype,
1323                dest,
1324                operand,
1325                start_indices,
1326            ),
1327            KernelDType::F64 => execute_gather_uninit_dispatch::<f64>(
1328                &self.plan,
1329                self.index_dtype,
1330                dest,
1331                operand,
1332                start_indices,
1333            ),
1334            KernelDType::I32 => execute_gather_uninit_dispatch::<i32>(
1335                &self.plan,
1336                self.index_dtype,
1337                dest,
1338                operand,
1339                start_indices,
1340            ),
1341            KernelDType::I64 => execute_gather_uninit_dispatch::<i64>(
1342                &self.plan,
1343                self.index_dtype,
1344                dest,
1345                operand,
1346                start_indices,
1347            ),
1348            KernelDType::Bool => execute_gather_uninit_dispatch::<bool>(
1349                &self.plan,
1350                self.index_dtype,
1351                dest,
1352                operand,
1353                start_indices,
1354            ),
1355            KernelDType::C32 => execute_gather_uninit_dispatch::<Complex32>(
1356                &self.plan,
1357                self.index_dtype,
1358                dest,
1359                operand,
1360                start_indices,
1361            ),
1362            KernelDType::C64 => execute_gather_uninit_dispatch::<Complex64>(
1363                &self.plan,
1364                self.index_dtype,
1365                dest,
1366                operand,
1367                start_indices,
1368            ),
1369            _ => Err(StridedError::UnsupportedDType {
1370                dtype: self.dtype.label(),
1371            }),
1372        };
1373        if ctx.is_serial() {
1374            run(dest)
1375        } else {
1376            ctx.run(|| run(dest))
1377        }
1378    }
1379}
1380
1381impl ErasedDynamicSlicePlan {
1382    /// Validate and store a dynamic-slice plan for one value dtype, index dtype, and layout set.
1383    #[allow(clippy::too_many_arguments)]
1384    pub fn compile(
1385        dtype: KernelDType,
1386        index_dtype: KernelDType,
1387        operand_dims: &[usize],
1388        operand_strides: &[isize],
1389        start_dims: &[usize],
1390        start_strides: &[isize],
1391        dest_dims: &[usize],
1392        dest_strides: &[isize],
1393        slice_sizes: &[usize],
1394    ) -> Result<Self> {
1395        check_index_dtype(index_dtype)?;
1396        check_gather_value_dtype(dtype)?;
1397        Ok(Self {
1398            dtype,
1399            index_dtype,
1400            plan: DynamicSlicePlan::compile(
1401                operand_dims,
1402                operand_strides,
1403                start_dims,
1404                start_strides,
1405                dest_dims,
1406                dest_strides,
1407                slice_sizes,
1408            )?,
1409        })
1410    }
1411
1412    #[inline]
1413    pub fn dtype(&self) -> KernelDType {
1414        self.dtype
1415    }
1416
1417    #[inline]
1418    pub fn index_dtype(&self) -> KernelDType {
1419        self.index_dtype
1420    }
1421
1422    #[inline]
1423    pub fn plan(&self) -> &DynamicSlicePlan {
1424        &self.plan
1425    }
1426
1427    /// Execute a fixed-window dynamic slice into an erased output descriptor.
1428    pub fn execute(
1429        &self,
1430        ctx: &ExecContext,
1431        dest: &mut ErasedRawStridedMut<'_>,
1432        operand: &ErasedRawStridedRef<'_>,
1433        starts: &ErasedRawStridedRef<'_>,
1434    ) -> Result<()> {
1435        check_dtype(self.dtype, dest.dtype())?;
1436        check_dtype(self.dtype, operand.dtype())?;
1437        check_dtype(self.index_dtype, starts.dtype())?;
1438
1439        let result = ctx.run(|| match self.dtype {
1440            KernelDType::F32 => dispatch_dynamic_slice_index::<f32>(
1441                &self.plan,
1442                self.index_dtype,
1443                dest,
1444                &operand,
1445                &starts,
1446            ),
1447            KernelDType::F64 => dispatch_dynamic_slice_index::<f64>(
1448                &self.plan,
1449                self.index_dtype,
1450                dest,
1451                operand,
1452                starts,
1453            ),
1454            KernelDType::I32 => dispatch_dynamic_slice_index::<i32>(
1455                &self.plan,
1456                self.index_dtype,
1457                dest,
1458                operand,
1459                starts,
1460            ),
1461            KernelDType::I64 => dispatch_dynamic_slice_index::<i64>(
1462                &self.plan,
1463                self.index_dtype,
1464                dest,
1465                operand,
1466                starts,
1467            ),
1468            KernelDType::Bool => dispatch_dynamic_slice_index::<bool>(
1469                &self.plan,
1470                self.index_dtype,
1471                dest,
1472                operand,
1473                starts,
1474            ),
1475            KernelDType::C32 => dispatch_dynamic_slice_index::<Complex32>(
1476                &self.plan,
1477                self.index_dtype,
1478                dest,
1479                operand,
1480                starts,
1481            ),
1482            KernelDType::C64 => dispatch_dynamic_slice_index::<Complex64>(
1483                &self.plan,
1484                self.index_dtype,
1485                dest,
1486                operand,
1487                starts,
1488            ),
1489            _ => Err(StridedError::UnsupportedDType {
1490                dtype: self.dtype.label(),
1491            }),
1492        });
1493        result
1494    }
1495
1496    /// Execute dynamic slice into a destination whose reachable slots may be
1497    /// uninitialized.
1498    /// On success, every reachable destination slot is fully overwritten;
1499    /// unreachable holes are neither read nor initialized. Validation errors
1500    /// are returned before any destination write. A panic during execution may
1501    /// leave reachable slots partially initialized, but the `MaybeUninit`
1502    /// destination remains safely droppable.
1503    pub fn execute_uninit(
1504        &self,
1505        ctx: &ExecContext,
1506        dest: &mut ErasedRawStridedUninitMut<'_>,
1507        operand: &ErasedRawStridedPtr<'_>,
1508        starts: &ErasedRawStridedPtr<'_>,
1509    ) -> Result<()> {
1510        check_dtype(self.dtype, dest.dtype())?;
1511        check_dtype(self.dtype, operand.dtype())?;
1512        check_dtype(self.index_dtype, starts.dtype())?;
1513        validate_uninit_no_overlap(dest, operand, 0)?;
1514        validate_uninit_no_overlap(dest, starts, 1)?;
1515        let operand = &validated_input_ref(operand)?;
1516        let starts = &validated_input_ref(starts)?;
1517        let run = |dest: &mut ErasedRawStridedUninitMut<'_>| match self.dtype {
1518            KernelDType::F32 => execute_dynamic_slice_uninit_dispatch::<f32>(
1519                &self.plan,
1520                self.index_dtype,
1521                dest,
1522                operand,
1523                starts,
1524            ),
1525            KernelDType::F64 => execute_dynamic_slice_uninit_dispatch::<f64>(
1526                &self.plan,
1527                self.index_dtype,
1528                dest,
1529                operand,
1530                starts,
1531            ),
1532            KernelDType::I32 => execute_dynamic_slice_uninit_dispatch::<i32>(
1533                &self.plan,
1534                self.index_dtype,
1535                dest,
1536                operand,
1537                starts,
1538            ),
1539            KernelDType::I64 => execute_dynamic_slice_uninit_dispatch::<i64>(
1540                &self.plan,
1541                self.index_dtype,
1542                dest,
1543                operand,
1544                starts,
1545            ),
1546            KernelDType::Bool => execute_dynamic_slice_uninit_dispatch::<bool>(
1547                &self.plan,
1548                self.index_dtype,
1549                dest,
1550                operand,
1551                starts,
1552            ),
1553            KernelDType::C32 => execute_dynamic_slice_uninit_dispatch::<Complex32>(
1554                &self.plan,
1555                self.index_dtype,
1556                dest,
1557                operand,
1558                starts,
1559            ),
1560            KernelDType::C64 => execute_dynamic_slice_uninit_dispatch::<Complex64>(
1561                &self.plan,
1562                self.index_dtype,
1563                dest,
1564                operand,
1565                starts,
1566            ),
1567            _ => Err(StridedError::UnsupportedDType {
1568                dtype: self.dtype.label(),
1569            }),
1570        };
1571        if ctx.is_serial() {
1572            run(dest)
1573        } else {
1574            ctx.run(|| run(dest))
1575        }
1576    }
1577}
1578
1579impl ErasedDynamicUpdateSlicePlan {
1580    /// Validate and store a dynamic-update-slice plan for one value dtype, index dtype, and layout set.
1581    #[allow(clippy::too_many_arguments)]
1582    pub fn compile(
1583        dtype: KernelDType,
1584        index_dtype: KernelDType,
1585        operand_dims: &[usize],
1586        operand_strides: &[isize],
1587        start_dims: &[usize],
1588        start_strides: &[isize],
1589        update_dims: &[usize],
1590        update_strides: &[isize],
1591        dest_dims: &[usize],
1592        dest_strides: &[isize],
1593    ) -> Result<Self> {
1594        check_index_dtype(index_dtype)?;
1595        check_gather_value_dtype(dtype)?;
1596        Ok(Self {
1597            dtype,
1598            index_dtype,
1599            plan: DynamicUpdateSlicePlan::compile(
1600                operand_dims,
1601                operand_strides,
1602                start_dims,
1603                start_strides,
1604                update_dims,
1605                update_strides,
1606                dest_dims,
1607                dest_strides,
1608            )?,
1609        })
1610    }
1611
1612    #[inline]
1613    pub fn dtype(&self) -> KernelDType {
1614        self.dtype
1615    }
1616
1617    #[inline]
1618    pub fn index_dtype(&self) -> KernelDType {
1619        self.index_dtype
1620    }
1621
1622    #[inline]
1623    pub fn plan(&self) -> &DynamicUpdateSlicePlan {
1624        &self.plan
1625    }
1626
1627    /// Execute a dynamic update slice into an erased output descriptor.
1628    pub fn execute(
1629        &self,
1630        ctx: &ExecContext,
1631        dest: &mut ErasedRawStridedMut<'_>,
1632        operand: &ErasedRawStridedRef<'_>,
1633        update: &ErasedRawStridedRef<'_>,
1634        starts: &ErasedRawStridedRef<'_>,
1635    ) -> Result<()> {
1636        check_dtype(self.dtype, dest.dtype())?;
1637        check_dtype(self.dtype, operand.dtype())?;
1638        check_dtype(self.dtype, update.dtype())?;
1639        check_dtype(self.index_dtype, starts.dtype())?;
1640
1641        let result = ctx.run(|| match self.dtype {
1642            KernelDType::F32 => dispatch_dynamic_update_slice_index::<f32>(
1643                &self.plan,
1644                self.index_dtype,
1645                dest,
1646                &operand,
1647                &update,
1648                &starts,
1649            ),
1650            KernelDType::F64 => dispatch_dynamic_update_slice_index::<f64>(
1651                &self.plan,
1652                self.index_dtype,
1653                dest,
1654                operand,
1655                update,
1656                starts,
1657            ),
1658            KernelDType::I32 => dispatch_dynamic_update_slice_index::<i32>(
1659                &self.plan,
1660                self.index_dtype,
1661                dest,
1662                operand,
1663                update,
1664                starts,
1665            ),
1666            KernelDType::I64 => dispatch_dynamic_update_slice_index::<i64>(
1667                &self.plan,
1668                self.index_dtype,
1669                dest,
1670                operand,
1671                update,
1672                starts,
1673            ),
1674            KernelDType::Bool => dispatch_dynamic_update_slice_index::<bool>(
1675                &self.plan,
1676                self.index_dtype,
1677                dest,
1678                operand,
1679                update,
1680                starts,
1681            ),
1682            KernelDType::C32 => dispatch_dynamic_update_slice_index::<Complex32>(
1683                &self.plan,
1684                self.index_dtype,
1685                dest,
1686                operand,
1687                update,
1688                starts,
1689            ),
1690            KernelDType::C64 => dispatch_dynamic_update_slice_index::<Complex64>(
1691                &self.plan,
1692                self.index_dtype,
1693                dest,
1694                operand,
1695                update,
1696                starts,
1697            ),
1698            _ => Err(StridedError::UnsupportedDType {
1699                dtype: self.dtype.label(),
1700            }),
1701        });
1702        result
1703    }
1704
1705    /// On success, the copy phase initializes every reachable destination
1706    /// slot before the read-modify-write phase. Unreachable holes are neither
1707    /// read nor initialized. Validation errors before the copy leave the
1708    /// destination untouched; an error or panic after the copy may leave a
1709    /// mixture of old and new reachable values, all initialized and safely
1710    /// droppable.
1711    pub fn execute_uninit(
1712        &self,
1713        ctx: &ExecContext,
1714        dest: &mut ErasedRawStridedUninitMut<'_>,
1715        operand: &ErasedRawStridedPtr<'_>,
1716        update: &ErasedRawStridedPtr<'_>,
1717        starts: &ErasedRawStridedPtr<'_>,
1718    ) -> Result<()> {
1719        check_dtype(self.dtype, dest.dtype())?;
1720        check_dtype(self.dtype, operand.dtype())?;
1721        check_dtype(self.dtype, update.dtype())?;
1722        check_dtype(self.index_dtype, starts.dtype())?;
1723        validate_uninit_no_overlap(dest, operand, 0)?;
1724        validate_uninit_no_overlap(dest, update, 1)?;
1725        validate_uninit_no_overlap(dest, starts, 2)?;
1726        let operand = &validated_input_ref(operand)?;
1727        let update = &validated_input_ref(update)?;
1728        let starts = &validated_input_ref(starts)?;
1729        let run = |dest: &mut ErasedRawStridedUninitMut<'_>| match self.dtype {
1730            KernelDType::F32 => execute_dynamic_update_uninit_dispatch::<f32>(
1731                &self.plan,
1732                self.index_dtype,
1733                dest,
1734                operand,
1735                update,
1736                starts,
1737            ),
1738            KernelDType::F64 => execute_dynamic_update_uninit_dispatch::<f64>(
1739                &self.plan,
1740                self.index_dtype,
1741                dest,
1742                operand,
1743                update,
1744                starts,
1745            ),
1746            KernelDType::I32 => execute_dynamic_update_uninit_dispatch::<i32>(
1747                &self.plan,
1748                self.index_dtype,
1749                dest,
1750                operand,
1751                update,
1752                starts,
1753            ),
1754            KernelDType::I64 => execute_dynamic_update_uninit_dispatch::<i64>(
1755                &self.plan,
1756                self.index_dtype,
1757                dest,
1758                operand,
1759                update,
1760                starts,
1761            ),
1762            KernelDType::Bool => execute_dynamic_update_uninit_dispatch::<bool>(
1763                &self.plan,
1764                self.index_dtype,
1765                dest,
1766                operand,
1767                update,
1768                starts,
1769            ),
1770            KernelDType::C32 => execute_dynamic_update_uninit_dispatch::<Complex32>(
1771                &self.plan,
1772                self.index_dtype,
1773                dest,
1774                operand,
1775                update,
1776                starts,
1777            ),
1778            KernelDType::C64 => execute_dynamic_update_uninit_dispatch::<Complex64>(
1779                &self.plan,
1780                self.index_dtype,
1781                dest,
1782                operand,
1783                update,
1784                starts,
1785            ),
1786            _ => Err(StridedError::UnsupportedDType {
1787                dtype: self.dtype.label(),
1788            }),
1789        };
1790        if ctx.is_serial() {
1791            run(dest)
1792        } else {
1793            ctx.run(|| run(dest))
1794        }
1795    }
1796}
1797
1798impl ErasedScatterPlan {
1799    /// Validate and store an additive scatter plan for one value dtype, index dtype, and layout set.
1800    #[allow(clippy::too_many_arguments)]
1801    pub fn compile(
1802        dtype: KernelDType,
1803        index_dtype: KernelDType,
1804        operand_dims: &[usize],
1805        operand_strides: &[isize],
1806        index_dims: &[usize],
1807        index_strides: &[isize],
1808        update_dims: &[usize],
1809        update_strides: &[isize],
1810        dest_dims: &[usize],
1811        dest_strides: &[isize],
1812        spec: ScatterSpec,
1813    ) -> Result<Self> {
1814        check_index_dtype(index_dtype)?;
1815        check_scatter_value_dtype(dtype)?;
1816        Ok(Self {
1817            dtype,
1818            index_dtype,
1819            plan: ScatterPlan::compile(
1820                operand_dims,
1821                operand_strides,
1822                index_dims,
1823                index_strides,
1824                update_dims,
1825                update_strides,
1826                dest_dims,
1827                dest_strides,
1828                spec,
1829            )?,
1830        })
1831    }
1832
1833    #[inline]
1834    pub fn dtype(&self) -> KernelDType {
1835        self.dtype
1836    }
1837
1838    #[inline]
1839    pub fn index_dtype(&self) -> KernelDType {
1840        self.index_dtype
1841    }
1842
1843    #[inline]
1844    pub fn plan(&self) -> &ScatterPlan {
1845        &self.plan
1846    }
1847
1848    /// Execute additive scatter into an erased output descriptor.
1849    pub fn execute(
1850        &self,
1851        ctx: &ExecContext,
1852        dest: &mut ErasedRawStridedMut<'_>,
1853        operand: &ErasedRawStridedRef<'_>,
1854        scatter_indices: &ErasedRawStridedRef<'_>,
1855        updates: &ErasedRawStridedRef<'_>,
1856    ) -> Result<()> {
1857        check_dtype(self.dtype, dest.dtype())?;
1858        check_dtype(self.dtype, operand.dtype())?;
1859        check_dtype(self.dtype, updates.dtype())?;
1860        check_dtype(self.index_dtype, scatter_indices.dtype())?;
1861
1862        let result = ctx.run(|| match self.dtype {
1863            KernelDType::F32 => dispatch_scatter_index::<f32>(
1864                &self.plan,
1865                self.index_dtype,
1866                dest,
1867                &operand,
1868                &scatter_indices,
1869                &updates,
1870            ),
1871            KernelDType::F64 => dispatch_scatter_index::<f64>(
1872                &self.plan,
1873                self.index_dtype,
1874                dest,
1875                operand,
1876                scatter_indices,
1877                updates,
1878            ),
1879            KernelDType::I32 => dispatch_scatter_index::<i32>(
1880                &self.plan,
1881                self.index_dtype,
1882                dest,
1883                operand,
1884                scatter_indices,
1885                updates,
1886            ),
1887            KernelDType::I64 => dispatch_scatter_index::<i64>(
1888                &self.plan,
1889                self.index_dtype,
1890                dest,
1891                operand,
1892                scatter_indices,
1893                updates,
1894            ),
1895            KernelDType::C32 => dispatch_scatter_index::<Complex32>(
1896                &self.plan,
1897                self.index_dtype,
1898                dest,
1899                operand,
1900                scatter_indices,
1901                updates,
1902            ),
1903            KernelDType::C64 => dispatch_scatter_index::<Complex64>(
1904                &self.plan,
1905                self.index_dtype,
1906                dest,
1907                operand,
1908                scatter_indices,
1909                updates,
1910            ),
1911            _ => Err(StridedError::UnsupportedDType {
1912                dtype: self.dtype.label(),
1913            }),
1914        });
1915        result
1916    }
1917
1918    /// On success, the copy phase initializes every reachable destination
1919    /// slot before the read-modify-write phase. Unreachable holes are neither
1920    /// read nor initialized. Validation errors before the copy leave the
1921    /// destination untouched; an error or panic after the copy may leave a
1922    /// mixture of old and new reachable values, all initialized and safely
1923    /// droppable.
1924    pub fn execute_uninit(
1925        &self,
1926        ctx: &ExecContext,
1927        dest: &mut ErasedRawStridedUninitMut<'_>,
1928        operand: &ErasedRawStridedPtr<'_>,
1929        scatter_indices: &ErasedRawStridedPtr<'_>,
1930        updates: &ErasedRawStridedPtr<'_>,
1931    ) -> Result<()> {
1932        check_dtype(self.dtype, dest.dtype())?;
1933        check_dtype(self.dtype, operand.dtype())?;
1934        check_dtype(self.dtype, updates.dtype())?;
1935        check_dtype(self.index_dtype, scatter_indices.dtype())?;
1936        validate_uninit_no_overlap(dest, operand, 0)?;
1937        validate_uninit_no_overlap(dest, scatter_indices, 1)?;
1938        validate_uninit_no_overlap(dest, updates, 2)?;
1939        let operand = &validated_input_ref(operand)?;
1940        let scatter_indices = &validated_input_ref(scatter_indices)?;
1941        let updates = &validated_input_ref(updates)?;
1942        let run = |dest: &mut ErasedRawStridedUninitMut<'_>| match self.dtype {
1943            KernelDType::F32 => execute_scatter_uninit_dispatch::<f32>(
1944                &self.plan,
1945                self.index_dtype,
1946                dest,
1947                operand,
1948                scatter_indices,
1949                updates,
1950                add_values::<f32>,
1951            ),
1952            KernelDType::F64 => execute_scatter_uninit_dispatch::<f64>(
1953                &self.plan,
1954                self.index_dtype,
1955                dest,
1956                operand,
1957                scatter_indices,
1958                updates,
1959                add_values::<f64>,
1960            ),
1961            KernelDType::I32 => execute_scatter_uninit_dispatch::<i32>(
1962                &self.plan,
1963                self.index_dtype,
1964                dest,
1965                operand,
1966                scatter_indices,
1967                updates,
1968                i32::wrapping_add,
1969            ),
1970            KernelDType::I64 => execute_scatter_uninit_dispatch::<i64>(
1971                &self.plan,
1972                self.index_dtype,
1973                dest,
1974                operand,
1975                scatter_indices,
1976                updates,
1977                i64::wrapping_add,
1978            ),
1979            KernelDType::C32 => execute_scatter_uninit_dispatch::<Complex32>(
1980                &self.plan,
1981                self.index_dtype,
1982                dest,
1983                operand,
1984                scatter_indices,
1985                updates,
1986                add_values::<Complex32>,
1987            ),
1988            KernelDType::C64 => execute_scatter_uninit_dispatch::<Complex64>(
1989                &self.plan,
1990                self.index_dtype,
1991                dest,
1992                operand,
1993                scatter_indices,
1994                updates,
1995                add_values::<Complex64>,
1996            ),
1997            _ => Err(StridedError::UnsupportedDType {
1998                dtype: self.dtype.label(),
1999            }),
2000        };
2001        if ctx.is_serial() {
2002            run(dest)
2003        } else {
2004            ctx.run(|| run(dest))
2005        }
2006    }
2007}
2008
2009fn add_values<T: Add<Output = T>>(lhs: T, rhs: T) -> T {
2010    lhs + rhs
2011}
2012
2013fn execute_one_shot_map<T: OneShotScalar>(
2014    op: ErasedMapOp,
2015    dest: &mut ErasedRawStridedMut<'_>,
2016    input: &ErasedRawStridedRef<'_>,
2017) -> Result<()> {
2018    if !T::supports_map(op) {
2019        return Err(StridedError::UnsupportedOp {
2020            op: op.label(),
2021            dtype: T::one_shot_dtype_label(),
2022        });
2023    }
2024    let validated =
2025        crate::map_view::validate_destination_layout_without_alloc(dest.dims(), dest.strides())?;
2026    crate::kernel::ensure_same_shape(dest.dims(), input.dims())?;
2027    if dest.dims().contains(&0) {
2028        return Ok(());
2029    }
2030
2031    let dest_dims = dest.dims();
2032    let dest_strides = dest.strides();
2033    let dest_offset = dest.offset();
2034    let dest_data = dest.data_as_mut::<T>()?;
2035    let mut dest =
2036        unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) };
2037    let input = erased_raw_ref::<T>(input)?;
2038
2039    crate::map_view::map_raw_into_validated::<T, T, Identity>(
2040        &mut dest,
2041        &input,
2042        |value| T::map(op, value),
2043        validated,
2044    )
2045}
2046
2047fn execute_one_shot_map_with<D, A>(
2048    dest: &mut ErasedRawStridedMut<'_>,
2049    input: &ErasedRawStridedRef<'_>,
2050    map: impl Fn(A) -> D + crate::MaybeSync,
2051) -> Result<()>
2052where
2053    D: Copy + crate::MaybeSendSync + KernelStorageElement,
2054    A: Copy + crate::MaybeSendSync + KernelStorageElement,
2055{
2056    let validated =
2057        crate::map_view::validate_destination_layout_without_alloc(dest.dims(), dest.strides())?;
2058    crate::kernel::ensure_same_shape(dest.dims(), input.dims())?;
2059    if dest.dims().contains(&0) {
2060        return Ok(());
2061    }
2062    let dest_dims = dest.dims();
2063    let dest_strides = dest.strides();
2064    let dest_offset = dest.offset();
2065    let dest_data = dest.data_as_mut::<D>()?;
2066    let mut dest =
2067        unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) };
2068    let input = erased_raw_ref::<A>(input)?;
2069    crate::map_view::map_raw_into_validated::<D, A, Identity>(&mut dest, &input, map, validated)
2070}
2071
2072fn execute_one_shot_zip<T: OneShotScalar>(
2073    op: ErasedZipOp,
2074    dest: &mut ErasedRawStridedMut<'_>,
2075    lhs: &ErasedRawStridedRef<'_>,
2076    rhs: &ErasedRawStridedRef<'_>,
2077) -> Result<()> {
2078    if !T::supports_zip(op) {
2079        return Err(StridedError::UnsupportedOp {
2080            op: op.label(),
2081            dtype: T::one_shot_dtype_label(),
2082        });
2083    }
2084    let validated =
2085        crate::map_view::validate_destination_layout_without_alloc(dest.dims(), dest.strides())?;
2086    crate::kernel::ensure_same_shape(dest.dims(), lhs.dims())?;
2087    crate::kernel::ensure_same_shape(dest.dims(), rhs.dims())?;
2088    if dest.dims().contains(&0) {
2089        return Ok(());
2090    }
2091
2092    let lhs = erased_raw_ref::<T>(lhs)?;
2093    let rhs = erased_raw_ref::<T>(rhs)?;
2094    if matches!(op, ErasedZipOp::Divide | ErasedZipOp::Remainder)
2095        && T::INTEGER
2096        && raw_any(&rhs, T::is_zero)?
2097    {
2098        return Err(StridedError::IntegerDivisionByZero { op: op.label() });
2099    }
2100    let dest_dims = dest.dims();
2101    let dest_strides = dest.strides();
2102    let dest_offset = dest.offset();
2103    let dest_data = dest.data_as_mut::<T>()?;
2104    let mut dest =
2105        unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) };
2106    crate::map_view::zip_map2_raw_into_validated::<T, T, T, Identity, Identity>(
2107        &mut dest,
2108        &lhs,
2109        &rhs,
2110        |lhs, rhs| T::zip(op, lhs, rhs),
2111        validated,
2112    )
2113}
2114
2115fn validate_no_overlap(
2116    dest: &ErasedRawStridedMut<'_>,
2117    input: &ErasedRawStridedPtr<'_>,
2118    input_index: usize,
2119) -> Result<()> {
2120    if input.overlaps_mut(dest)? {
2121        Err(StridedError::OverlappingInputOutput { input: input_index })
2122    } else {
2123        Ok(())
2124    }
2125}
2126
2127fn validate_uninit_no_overlap(
2128    dest: &ErasedRawStridedUninitMut<'_>,
2129    input: &ErasedRawStridedPtr<'_>,
2130    input_index: usize,
2131) -> Result<()> {
2132    if input.overlaps_uninit_mut(dest)? {
2133        Err(StridedError::OverlappingInputOutput { input: input_index })
2134    } else {
2135        Ok(())
2136    }
2137}
2138
2139trait OneShotScalar: Copy + crate::MaybeSendSync + KernelStorageElement + 'static {
2140    const INTEGER: bool = false;
2141    fn is_zero(_value: Self) -> bool {
2142        false
2143    }
2144    fn one_shot_dtype_label() -> &'static str;
2145    fn supports_map(op: ErasedMapOp) -> bool;
2146    fn supports_zip(op: ErasedZipOp) -> bool;
2147    fn map(op: ErasedMapOp, value: Self) -> Self;
2148    fn zip(op: ErasedZipOp, lhs: Self, rhs: Self) -> Self;
2149}
2150
2151macro_rules! impl_real_one_shot_scalar {
2152    ($ty:ty, $label:literal) => {
2153        impl OneShotScalar for $ty {
2154            fn one_shot_dtype_label() -> &'static str {
2155                $label
2156            }
2157
2158            fn supports_map(_op: ErasedMapOp) -> bool {
2159                true
2160            }
2161
2162            fn supports_zip(_op: ErasedZipOp) -> bool {
2163                true
2164            }
2165
2166            #[inline(always)]
2167            fn map(op: ErasedMapOp, value: Self) -> Self {
2168                match op {
2169                    ErasedMapOp::Negate => -value,
2170                    ErasedMapOp::Conj => value,
2171                    ErasedMapOp::Abs => value.abs(),
2172                    ErasedMapOp::Sign => {
2173                        if value == 0.0 {
2174                            0.0
2175                        } else {
2176                            value.signum()
2177                        }
2178                    }
2179                }
2180            }
2181
2182            #[inline(always)]
2183            fn zip(op: ErasedZipOp, lhs: Self, rhs: Self) -> Self {
2184                match op {
2185                    ErasedZipOp::Add => lhs + rhs,
2186                    ErasedZipOp::Subtract => lhs - rhs,
2187                    ErasedZipOp::Multiply => lhs * rhs,
2188                    ErasedZipOp::Divide => lhs / rhs,
2189                    ErasedZipOp::Remainder => lhs % rhs,
2190                    ErasedZipOp::Maximum => {
2191                        if lhs.is_nan() || rhs.is_nan() {
2192                            <$ty>::NAN
2193                        } else if lhs >= rhs {
2194                            lhs
2195                        } else {
2196                            rhs
2197                        }
2198                    }
2199                    ErasedZipOp::Minimum => {
2200                        if lhs.is_nan() || rhs.is_nan() {
2201                            <$ty>::NAN
2202                        } else if lhs <= rhs {
2203                            lhs
2204                        } else {
2205                            rhs
2206                        }
2207                    }
2208                }
2209            }
2210        }
2211    };
2212}
2213
2214macro_rules! impl_integer_one_shot_scalar {
2215    ($ty:ty, $label:literal) => {
2216        impl OneShotScalar for $ty {
2217            const INTEGER: bool = true;
2218
2219            fn is_zero(value: Self) -> bool {
2220                value == 0
2221            }
2222            fn one_shot_dtype_label() -> &'static str {
2223                $label
2224            }
2225
2226            fn supports_map(_op: ErasedMapOp) -> bool {
2227                true
2228            }
2229
2230            fn supports_zip(_op: ErasedZipOp) -> bool {
2231                true
2232            }
2233
2234            #[inline(always)]
2235            fn map(op: ErasedMapOp, value: Self) -> Self {
2236                match op {
2237                    ErasedMapOp::Negate => value.wrapping_neg(),
2238                    ErasedMapOp::Conj => value,
2239                    ErasedMapOp::Abs => value.wrapping_abs(),
2240                    ErasedMapOp::Sign => value.signum(),
2241                }
2242            }
2243
2244            #[inline(always)]
2245            fn zip(op: ErasedZipOp, lhs: Self, rhs: Self) -> Self {
2246                match op {
2247                    ErasedZipOp::Add => lhs.wrapping_add(rhs),
2248                    ErasedZipOp::Subtract => lhs.wrapping_sub(rhs),
2249                    ErasedZipOp::Multiply => lhs.wrapping_mul(rhs),
2250                    ErasedZipOp::Maximum => lhs.max(rhs),
2251                    ErasedZipOp::Minimum => lhs.min(rhs),
2252                    ErasedZipOp::Divide => lhs.wrapping_div(rhs),
2253                    ErasedZipOp::Remainder => lhs.wrapping_rem(rhs),
2254                }
2255            }
2256        }
2257    };
2258}
2259
2260macro_rules! impl_complex_one_shot_scalar {
2261    ($ty:ty, $label:literal) => {
2262        impl OneShotScalar for $ty {
2263            fn one_shot_dtype_label() -> &'static str {
2264                $label
2265            }
2266
2267            fn supports_map(_op: ErasedMapOp) -> bool {
2268                true
2269            }
2270
2271            fn supports_zip(op: ErasedZipOp) -> bool {
2272                !matches!(
2273                    op,
2274                    ErasedZipOp::Remainder | ErasedZipOp::Maximum | ErasedZipOp::Minimum
2275                )
2276            }
2277
2278            #[inline(always)]
2279            fn map(op: ErasedMapOp, value: Self) -> Self {
2280                match op {
2281                    ErasedMapOp::Negate => -value,
2282                    ErasedMapOp::Conj => value.conj(),
2283                    ErasedMapOp::Abs => Self::new(value.norm(), 0.0),
2284                    ErasedMapOp::Sign => {
2285                        let norm = value.norm();
2286                        if norm == 0.0 {
2287                            Self::new(0.0, 0.0)
2288                        } else {
2289                            value / Self::new(norm, 0.0)
2290                        }
2291                    }
2292                }
2293            }
2294
2295            #[inline(always)]
2296            fn zip(op: ErasedZipOp, lhs: Self, rhs: Self) -> Self {
2297                match op {
2298                    ErasedZipOp::Add => lhs + rhs,
2299                    ErasedZipOp::Subtract => lhs - rhs,
2300                    ErasedZipOp::Multiply => lhs * rhs,
2301                    ErasedZipOp::Divide => lhs / rhs,
2302                    ErasedZipOp::Remainder | ErasedZipOp::Maximum | ErasedZipOp::Minimum => {
2303                        unreachable!("unsupported complex one-shot op")
2304                    }
2305                }
2306            }
2307        }
2308    };
2309}
2310
2311impl_real_one_shot_scalar!(f32, "f32");
2312impl_real_one_shot_scalar!(f64, "f64");
2313impl_integer_one_shot_scalar!(i32, "i32");
2314impl_integer_one_shot_scalar!(i64, "i64");
2315impl_complex_one_shot_scalar!(Complex32, "c32");
2316impl_complex_one_shot_scalar!(Complex64, "c64");
2317
2318impl OneShotScalar for bool {
2319    fn one_shot_dtype_label() -> &'static str {
2320        "bool"
2321    }
2322
2323    fn supports_map(op: ErasedMapOp) -> bool {
2324        matches!(op, ErasedMapOp::Conj)
2325    }
2326
2327    fn supports_zip(_op: ErasedZipOp) -> bool {
2328        false
2329    }
2330
2331    fn map(op: ErasedMapOp, value: Self) -> Self {
2332        match op {
2333            ErasedMapOp::Conj => value,
2334            _ => unreachable!("unsupported bool one-shot op"),
2335        }
2336    }
2337
2338    fn zip(_op: ErasedZipOp, _lhs: Self, _rhs: Self) -> Self {
2339        unreachable!("unsupported bool one-shot op")
2340    }
2341}
2342
2343fn erased_raw_ref<'a, T: KernelStorageElement>(
2344    src: &'a ErasedRawStridedRef<'a>,
2345) -> Result<RawStridedRef<'a, T>> {
2346    let data = src.data_as::<T>()?;
2347    Ok(unsafe { RawStridedRef::new_unchecked(data, src.dims(), src.strides(), src.offset()) })
2348}
2349
2350fn validated_input_ref<'a>(input: &'a ErasedRawStridedPtr<'a>) -> Result<ErasedRawStridedRef<'a>> {
2351    // SAFETY: callers reject overlap before this conversion.
2352    unsafe { input.try_as_ref_after_no_overlap() }
2353}
2354
2355fn map_output_dtype(dtype: KernelDType, op: ErasedMapOp) -> Result<KernelDType> {
2356    match (dtype, op) {
2357        (KernelDType::C32, ErasedMapOp::Abs) => Ok(KernelDType::F32),
2358        (KernelDType::C64, ErasedMapOp::Abs) => Ok(KernelDType::F64),
2359        (KernelDType::Bool, ErasedMapOp::Conj) => Ok(KernelDType::Bool),
2360        (KernelDType::Bool, _) => Err(StridedError::UnsupportedOp {
2361            op: op.label(),
2362            dtype: dtype.label(),
2363        }),
2364        _ => Ok(dtype),
2365    }
2366}
2367
2368fn raw_any<T: Copy>(
2369    src: &RawStridedRef<'_, T>,
2370    predicate: impl Fn(T) -> bool + Copy,
2371) -> Result<bool> {
2372    let total = src
2373        .dims()
2374        .iter()
2375        .try_fold(1usize, |total, &dim| total.checked_mul(dim))
2376        .ok_or(StridedError::OffsetOverflow)?;
2377    for linear in 0..total {
2378        let mut remainder = linear;
2379        let mut offset = src.offset();
2380        for (&dim, &stride) in src.dims().iter().zip(src.strides()) {
2381            let index = remainder % dim;
2382            remainder /= dim;
2383            offset = offset
2384                .checked_add(
2385                    stride
2386                        .checked_mul(index as isize)
2387                        .ok_or(StridedError::OffsetOverflow)?,
2388                )
2389                .ok_or(StridedError::OffsetOverflow)?;
2390        }
2391        // SAFETY: RawStridedRef construction validated every reachable offset.
2392        if predicate(unsafe { *src.data().as_ptr().offset(offset) }) {
2393            return Ok(true);
2394        }
2395    }
2396    Ok(false)
2397}
2398
2399fn check_dtype(expected: KernelDType, actual: KernelDType) -> Result<()> {
2400    if actual != expected {
2401        return Err(StridedError::DTypeMismatch {
2402            expected: expected.label(),
2403            actual: actual.label(),
2404        });
2405    }
2406    Ok(())
2407}
2408
2409fn reduce_writer<'a, T>(dest: &'a mut ErasedRawStridedMut<'_>) -> Result<RawReduceWriter<'a, T>>
2410where
2411    T: KernelStorageElement,
2412{
2413    let offset = dest.offset();
2414    let data = dest.data_as_mut::<T>()?;
2415    let ptr = data.as_mut_ptr();
2416    let extent = data.len();
2417    Ok(RawReduceWriter {
2418        ptr,
2419        extent,
2420        offset,
2421        _marker: core::marker::PhantomData,
2422    })
2423}
2424
2425fn reduce_uninit_writer<'a, T>(
2426    dest: &'a mut ErasedRawStridedUninitMut<'_>,
2427) -> Result<RawReduceWriter<'a, T>>
2428where
2429    T: KernelStorageElement,
2430{
2431    let offset = dest.offset();
2432    let data = dest.data_as_uninit_mut::<T>()?;
2433    let ptr = data.as_mut_ptr().cast::<T>();
2434    let extent = data.len();
2435    Ok(RawReduceWriter {
2436        ptr,
2437        extent,
2438        offset,
2439        _marker: core::marker::PhantomData,
2440    })
2441}
2442
2443fn check_fused_dtype(dtype: KernelDType) -> Result<()> {
2444    match dtype {
2445        KernelDType::F32
2446        | KernelDType::F64
2447        | KernelDType::I32
2448        | KernelDType::I64
2449        | KernelDType::Bool
2450        | KernelDType::C32
2451        | KernelDType::C64 => Ok(()),
2452        _ => Err(StridedError::UnsupportedDType {
2453            dtype: dtype.label(),
2454        }),
2455    }
2456}
2457
2458fn validate_fused_plan_for_dtype(dtype: KernelDType, plan: &FusedPlan) -> Result<()> {
2459    match dtype {
2460        KernelDType::F32 => {
2461            crate::fused::validate_plan_for_scalar::<f32>(plan, plan.input_count, 1)
2462        }
2463        KernelDType::F64 => {
2464            crate::fused::validate_plan_for_scalar::<f64>(plan, plan.input_count, 1)
2465        }
2466        KernelDType::I32 => {
2467            crate::fused::validate_plan_for_scalar::<i32>(plan, plan.input_count, 1)
2468        }
2469        KernelDType::I64 => {
2470            crate::fused::validate_plan_for_scalar::<i64>(plan, plan.input_count, 1)
2471        }
2472        KernelDType::Bool => {
2473            crate::fused::validate_plan_for_scalar::<bool>(plan, plan.input_count, 1)
2474        }
2475        KernelDType::C32 => {
2476            crate::fused::validate_plan_for_scalar::<Complex32>(plan, plan.input_count, 1)
2477        }
2478        KernelDType::C64 => {
2479            crate::fused::validate_plan_for_scalar::<Complex64>(plan, plan.input_count, 1)
2480        }
2481        _ => Err(StridedError::UnsupportedDType {
2482            dtype: dtype.label(),
2483        }),
2484    }
2485}
2486
2487fn check_reduce_dtype(dtype: KernelDType) -> Result<()> {
2488    match dtype {
2489        KernelDType::F32
2490        | KernelDType::F64
2491        | KernelDType::I32
2492        | KernelDType::I64
2493        | KernelDType::C32
2494        | KernelDType::C64 => Ok(()),
2495        _ => Err(StridedError::UnsupportedDType {
2496            dtype: dtype.label(),
2497        }),
2498    }
2499}
2500
2501fn check_reduce_op_dtype(dtype: KernelDType, op: ReduceOp) -> Result<()> {
2502    if op == ReduceOp::SumSquares && !matches!(dtype, KernelDType::F32 | KernelDType::F64) {
2503        return Err(StridedError::UnsupportedDType {
2504            dtype: dtype.label(),
2505        });
2506    }
2507    check_reduce_dtype(dtype)
2508}
2509
2510fn check_index_dtype(dtype: KernelDType) -> Result<()> {
2511    match dtype {
2512        KernelDType::I32 | KernelDType::I64 => Ok(()),
2513        _ => Err(StridedError::UnsupportedDType {
2514            dtype: dtype.label(),
2515        }),
2516    }
2517}
2518
2519fn check_gather_value_dtype(dtype: KernelDType) -> Result<()> {
2520    match dtype {
2521        KernelDType::F32
2522        | KernelDType::F64
2523        | KernelDType::I32
2524        | KernelDType::I64
2525        | KernelDType::Bool
2526        | KernelDType::C32
2527        | KernelDType::C64 => Ok(()),
2528        _ => Err(StridedError::UnsupportedDType {
2529            dtype: dtype.label(),
2530        }),
2531    }
2532}
2533
2534fn check_scatter_value_dtype(dtype: KernelDType) -> Result<()> {
2535    match dtype {
2536        KernelDType::F32
2537        | KernelDType::F64
2538        | KernelDType::I32
2539        | KernelDType::I64
2540        | KernelDType::C32
2541        | KernelDType::C64 => Ok(()),
2542        _ => Err(StridedError::UnsupportedDType {
2543            dtype: dtype.label(),
2544        }),
2545    }
2546}
2547
2548fn check_static_indexing_dtype(dtype: KernelDType) -> Result<()> {
2549    match dtype {
2550        KernelDType::F32
2551        | KernelDType::F64
2552        | KernelDType::I32
2553        | KernelDType::I64
2554        | KernelDType::Bool
2555        | KernelDType::C32
2556        | KernelDType::C64 => Ok(()),
2557        _ => Err(StridedError::UnsupportedDType {
2558            dtype: dtype.label(),
2559        }),
2560    }
2561}
2562
2563fn validate_scalar_bytes(dtype: KernelDType, bytes: &[u8]) -> Result<()> {
2564    let element_size = dtype.size_of();
2565    if bytes.len() != element_size {
2566        return Err(StridedError::ByteLengthMismatch {
2567            dtype: dtype.label(),
2568            byte_len: bytes.len(),
2569            element_size,
2570        });
2571    }
2572    if dtype.requires_valid_byte_values() {
2573        if let Some(&value) = bytes.iter().find(|&&value| value > 1) {
2574            return Err(StridedError::InvalidBoolByte { value });
2575        }
2576    }
2577    Ok(())
2578}
2579
2580fn checked_total_len(dims: &[usize]) -> Result<usize> {
2581    if dims.is_empty() {
2582        return Ok(1);
2583    }
2584    dims.iter()
2585        .try_fold(1usize, |acc, &dim| acc.checked_mul(dim))
2586        .ok_or(StridedError::OffsetOverflow)
2587}
2588
2589fn execute_copy<T>(
2590    plan: &CopyPlan,
2591    dest: &mut ErasedRawStridedMut<'_>,
2592    src: &ErasedRawStridedRef<'_>,
2593) -> Result<()>
2594where
2595    T: Copy + crate::MaybeSendSync + KernelStorageElement,
2596{
2597    let source_data = src.data_as::<T>()?;
2598    let dest_dims = dest.dims();
2599    let dest_strides = dest.strides();
2600    let dest_offset = dest.offset();
2601    let dest_data = dest.data_as_mut::<T>()?;
2602    let source = unsafe {
2603        RawStridedRef::new_unchecked(source_data, src.dims(), src.strides(), src.offset())
2604    };
2605    let mut dest =
2606        unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) };
2607    plan.execute(&mut dest, &source)
2608}
2609
2610fn execute_slice<T>(
2611    plan: &SlicePlan,
2612    dest: &mut ErasedRawStridedMut<'_>,
2613    operand: &ErasedRawStridedRef<'_>,
2614) -> Result<()>
2615where
2616    T: Copy + crate::MaybeSendSync + KernelStorageElement,
2617{
2618    let operand_data = operand.data_as::<T>()?;
2619    let dest_dims = dest.dims();
2620    let dest_strides = dest.strides();
2621    let dest_offset = dest.offset();
2622    let dest_data = dest.data_as_mut::<T>()?;
2623    let operand_ref = unsafe {
2624        RawStridedRef::new_unchecked(
2625            operand_data,
2626            operand.dims(),
2627            operand.strides(),
2628            operand.offset(),
2629        )
2630    };
2631    let mut dest_ref =
2632        unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) };
2633    plan.execute(&mut dest_ref, &operand_ref)
2634}
2635
2636fn execute_gather_uninit_dispatch<T>(
2637    plan: &GatherPlan,
2638    index_dtype: KernelDType,
2639    dest: &mut ErasedRawStridedUninitMut<'_>,
2640    operand: &ErasedRawStridedRef<'_>,
2641    start_indices: &ErasedRawStridedRef<'_>,
2642) -> Result<()>
2643where
2644    T: Copy + crate::MaybeSendSync + KernelStorageElement,
2645{
2646    match index_dtype {
2647        KernelDType::I32 => {
2648            execute_gather_uninit::<T, i32>(plan, index_dtype, dest, operand, start_indices)
2649        }
2650        KernelDType::I64 => {
2651            execute_gather_uninit::<T, i64>(plan, index_dtype, dest, operand, start_indices)
2652        }
2653        _ => Err(StridedError::UnsupportedDType {
2654            dtype: index_dtype.label(),
2655        }),
2656    }
2657}
2658
2659fn execute_gather_uninit<T, I>(
2660    plan: &GatherPlan,
2661    _index_dtype: KernelDType,
2662    dest: &mut ErasedRawStridedUninitMut<'_>,
2663    operand: &ErasedRawStridedRef<'_>,
2664    start_indices: &ErasedRawStridedRef<'_>,
2665) -> Result<()>
2666where
2667    T: Copy + crate::MaybeSendSync + KernelStorageElement,
2668    I: GatherIndex + KernelStorageElement,
2669{
2670    let operand_data = operand.data_as::<T>()?;
2671    let index_data = start_indices.data_as::<I>()?;
2672    let dest_dims = dest.dims();
2673    let dest_strides = dest.strides();
2674    let dest_offset = dest.offset();
2675    let dest_data = dest.data_as_uninit_mut::<T>()?;
2676    let operand_ref = unsafe {
2677        RawStridedRef::new_unchecked(
2678            operand_data,
2679            operand.dims(),
2680            operand.strides(),
2681            operand.offset(),
2682        )
2683    };
2684    let index_ref = unsafe {
2685        RawStridedRef::new_unchecked(
2686            index_data,
2687            start_indices.dims(),
2688            start_indices.strides(),
2689            start_indices.offset(),
2690        )
2691    };
2692    let mut dest_ref =
2693        unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) };
2694    plan.execute_uninit(&mut dest_ref, &operand_ref, &index_ref)
2695}
2696
2697fn execute_dynamic_slice_uninit_dispatch<T>(
2698    plan: &DynamicSlicePlan,
2699    index_dtype: KernelDType,
2700    dest: &mut ErasedRawStridedUninitMut<'_>,
2701    operand: &ErasedRawStridedRef<'_>,
2702    starts: &ErasedRawStridedRef<'_>,
2703) -> Result<()>
2704where
2705    T: Copy + crate::MaybeSendSync + KernelStorageElement,
2706{
2707    match index_dtype {
2708        KernelDType::I32 => execute_dynamic_slice_uninit::<T, i32>(plan, dest, operand, starts),
2709        KernelDType::I64 => execute_dynamic_slice_uninit::<T, i64>(plan, dest, operand, starts),
2710        _ => Err(StridedError::UnsupportedDType {
2711            dtype: index_dtype.label(),
2712        }),
2713    }
2714}
2715
2716fn execute_dynamic_slice_uninit<T, I>(
2717    plan: &DynamicSlicePlan,
2718    dest: &mut ErasedRawStridedUninitMut<'_>,
2719    operand: &ErasedRawStridedRef<'_>,
2720    starts: &ErasedRawStridedRef<'_>,
2721) -> Result<()>
2722where
2723    T: Copy + crate::MaybeSendSync + KernelStorageElement,
2724    I: GatherIndex + KernelStorageElement,
2725{
2726    let operand_data = operand.data_as::<T>()?;
2727    let starts_data = starts.data_as::<I>()?;
2728    let dest_dims = dest.dims();
2729    let dest_strides = dest.strides();
2730    let dest_offset = dest.offset();
2731    let dest_data = dest.data_as_uninit_mut::<T>()?;
2732    let operand_ref = unsafe {
2733        RawStridedRef::new_unchecked(
2734            operand_data,
2735            operand.dims(),
2736            operand.strides(),
2737            operand.offset(),
2738        )
2739    };
2740    let starts_ref = unsafe {
2741        RawStridedRef::new_unchecked(
2742            starts_data,
2743            starts.dims(),
2744            starts.strides(),
2745            starts.offset(),
2746        )
2747    };
2748    let mut dest_ref =
2749        unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) };
2750    plan.execute_uninit(&mut dest_ref, &operand_ref, &starts_ref)
2751}
2752
2753fn execute_dynamic_update_uninit_dispatch<T>(
2754    plan: &DynamicUpdateSlicePlan,
2755    index_dtype: KernelDType,
2756    dest: &mut ErasedRawStridedUninitMut<'_>,
2757    operand: &ErasedRawStridedRef<'_>,
2758    update: &ErasedRawStridedRef<'_>,
2759    starts: &ErasedRawStridedRef<'_>,
2760) -> Result<()>
2761where
2762    T: Copy + crate::MaybeSendSync + KernelStorageElement,
2763{
2764    match index_dtype {
2765        KernelDType::I32 => {
2766            execute_dynamic_update_uninit::<T, i32>(plan, dest, operand, update, starts)
2767        }
2768        KernelDType::I64 => {
2769            execute_dynamic_update_uninit::<T, i64>(plan, dest, operand, update, starts)
2770        }
2771        _ => Err(StridedError::UnsupportedDType {
2772            dtype: index_dtype.label(),
2773        }),
2774    }
2775}
2776
2777fn execute_dynamic_update_uninit<T, I>(
2778    plan: &DynamicUpdateSlicePlan,
2779    dest: &mut ErasedRawStridedUninitMut<'_>,
2780    operand: &ErasedRawStridedRef<'_>,
2781    update: &ErasedRawStridedRef<'_>,
2782    starts: &ErasedRawStridedRef<'_>,
2783) -> Result<()>
2784where
2785    T: Copy + crate::MaybeSendSync + KernelStorageElement,
2786    I: GatherIndex + KernelStorageElement,
2787{
2788    let operand_data = operand.data_as::<T>()?;
2789    let update_data = update.data_as::<T>()?;
2790    let starts_data = starts.data_as::<I>()?;
2791    let dest_dims = dest.dims();
2792    let dest_strides = dest.strides();
2793    let dest_offset = dest.offset();
2794    let dest_data = dest.data_as_uninit_mut::<T>()?;
2795    let operand_ref = unsafe {
2796        RawStridedRef::new_unchecked(
2797            operand_data,
2798            operand.dims(),
2799            operand.strides(),
2800            operand.offset(),
2801        )
2802    };
2803    let update_ref = unsafe {
2804        RawStridedRef::new_unchecked(
2805            update_data,
2806            update.dims(),
2807            update.strides(),
2808            update.offset(),
2809        )
2810    };
2811    let starts_ref = unsafe {
2812        RawStridedRef::new_unchecked(
2813            starts_data,
2814            starts.dims(),
2815            starts.strides(),
2816            starts.offset(),
2817        )
2818    };
2819    let mut dest_ref =
2820        unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) };
2821    plan.execute_uninit(&mut dest_ref, &operand_ref, &update_ref, &starts_ref)
2822}
2823
2824fn execute_scatter_uninit_dispatch<T>(
2825    plan: &ScatterPlan,
2826    index_dtype: KernelDType,
2827    dest: &mut ErasedRawStridedUninitMut<'_>,
2828    operand: &ErasedRawStridedRef<'_>,
2829    scatter_indices: &ErasedRawStridedRef<'_>,
2830    updates: &ErasedRawStridedRef<'_>,
2831    combine: fn(T, T) -> T,
2832) -> Result<()>
2833where
2834    T: Copy + Add<Output = T> + crate::MaybeSendSync + KernelStorageElement,
2835{
2836    match index_dtype {
2837        KernelDType::I32 => {
2838            execute_scatter_uninit::<T, i32>(plan, dest, operand, scatter_indices, updates, combine)
2839        }
2840        KernelDType::I64 => {
2841            execute_scatter_uninit::<T, i64>(plan, dest, operand, scatter_indices, updates, combine)
2842        }
2843        _ => Err(StridedError::UnsupportedDType {
2844            dtype: index_dtype.label(),
2845        }),
2846    }
2847}
2848
2849fn execute_scatter_uninit<T, I>(
2850    plan: &ScatterPlan,
2851    dest: &mut ErasedRawStridedUninitMut<'_>,
2852    operand: &ErasedRawStridedRef<'_>,
2853    scatter_indices: &ErasedRawStridedRef<'_>,
2854    updates: &ErasedRawStridedRef<'_>,
2855    combine: fn(T, T) -> T,
2856) -> Result<()>
2857where
2858    T: Copy + Add<Output = T> + crate::MaybeSendSync + KernelStorageElement,
2859    I: GatherIndex + KernelStorageElement,
2860{
2861    let indices = scatter_indices;
2862    let operand_data = operand.data_as::<T>()?;
2863    let index_data = indices.data_as::<I>()?;
2864    let update_data = updates.data_as::<T>()?;
2865    let dest_dims = dest.dims();
2866    let dest_strides = dest.strides();
2867    let dest_offset = dest.offset();
2868    let dest_data = dest.data_as_uninit_mut::<T>()?;
2869    let operand_ref = unsafe {
2870        RawStridedRef::new_unchecked(
2871            operand_data,
2872            operand.dims(),
2873            operand.strides(),
2874            operand.offset(),
2875        )
2876    };
2877    let index_ref = unsafe {
2878        RawStridedRef::new_unchecked(
2879            index_data,
2880            indices.dims(),
2881            indices.strides(),
2882            indices.offset(),
2883        )
2884    };
2885    let update_ref = unsafe {
2886        RawStridedRef::new_unchecked(
2887            update_data,
2888            updates.dims(),
2889            updates.strides(),
2890            updates.offset(),
2891        )
2892    };
2893    let mut dest_ref =
2894        unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) };
2895    plan.execute_uninit(
2896        &mut dest_ref,
2897        &operand_ref,
2898        &index_ref,
2899        &update_ref,
2900        combine,
2901    )
2902}
2903
2904fn execute_slice_uninit<T>(
2905    plan: &SlicePlan,
2906    dest: &mut ErasedRawStridedUninitMut<'_>,
2907    operand: &ErasedRawStridedRef<'_>,
2908) -> Result<()>
2909where
2910    T: Copy + crate::MaybeSendSync + KernelStorageElement,
2911{
2912    let operand_data = operand.data_as::<T>()?;
2913    let dest_dims = dest.dims();
2914    let dest_strides = dest.strides();
2915    let dest_offset = dest.offset();
2916    let dest_data = dest.data_as_uninit_mut::<T>()?;
2917    let operand_ref = unsafe {
2918        RawStridedRef::new_unchecked(
2919            operand_data,
2920            operand.dims(),
2921            operand.strides(),
2922            operand.offset(),
2923        )
2924    };
2925    let mut dest_ref =
2926        unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) };
2927    plan.execute_uninit(&mut dest_ref, &operand_ref)
2928}
2929
2930fn execute_reverse<T>(
2931    plan: &ReversePlan,
2932    dest: &mut ErasedRawStridedMut<'_>,
2933    operand: &ErasedRawStridedRef<'_>,
2934) -> Result<()>
2935where
2936    T: Copy + crate::MaybeSendSync + KernelStorageElement,
2937{
2938    let operand_data = operand.data_as::<T>()?;
2939    let dest_dims = dest.dims();
2940    let dest_strides = dest.strides();
2941    let dest_offset = dest.offset();
2942    let dest_data = dest.data_as_mut::<T>()?;
2943    let operand_ref = unsafe {
2944        RawStridedRef::new_unchecked(
2945            operand_data,
2946            operand.dims(),
2947            operand.strides(),
2948            operand.offset(),
2949        )
2950    };
2951    let mut dest_ref =
2952        unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) };
2953    plan.execute(&mut dest_ref, &operand_ref)
2954}
2955
2956fn execute_reverse_uninit<T>(
2957    plan: &ReversePlan,
2958    dest: &mut ErasedRawStridedUninitMut<'_>,
2959    operand: &ErasedRawStridedRef<'_>,
2960) -> Result<()>
2961where
2962    T: Copy + crate::MaybeSendSync + KernelStorageElement,
2963{
2964    let operand_data = operand.data_as::<T>()?;
2965    let dest_dims = dest.dims();
2966    let dest_strides = dest.strides();
2967    let dest_offset = dest.offset();
2968    let dest_data = dest.data_as_uninit_mut::<T>()?;
2969    let operand_ref = unsafe {
2970        RawStridedRef::new_unchecked(
2971            operand_data,
2972            operand.dims(),
2973            operand.strides(),
2974            operand.offset(),
2975        )
2976    };
2977    let mut dest_ref =
2978        unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) };
2979    plan.execute_uninit(&mut dest_ref, &operand_ref)
2980}
2981
2982fn execute_pad<T>(
2983    plan: &PadPlan,
2984    dest: &mut ErasedRawStridedMut<'_>,
2985    operand: &ErasedRawStridedRef<'_>,
2986    fill: &[u8],
2987) -> Result<()>
2988where
2989    T: Copy + crate::MaybeSendSync + KernelStorageElement,
2990{
2991    let fill = read_unaligned_scalar::<T>(fill);
2992    let operand_data = operand.data_as::<T>()?;
2993    let dest_dims = dest.dims();
2994    let dest_strides = dest.strides();
2995    let dest_offset = dest.offset();
2996    let dest_data = dest.data_as_mut::<T>()?;
2997    let operand_ref = unsafe {
2998        RawStridedRef::new_unchecked(
2999            operand_data,
3000            operand.dims(),
3001            operand.strides(),
3002            operand.offset(),
3003        )
3004    };
3005    let mut dest_ref =
3006        unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) };
3007    plan.execute(&mut dest_ref, &operand_ref, fill)
3008}
3009
3010fn execute_pad_uninit<T>(
3011    plan: &PadPlan,
3012    dest: &mut ErasedRawStridedUninitMut<'_>,
3013    operand: &ErasedRawStridedRef<'_>,
3014    fill: &[u8],
3015) -> Result<()>
3016where
3017    T: Copy + crate::MaybeSendSync + KernelStorageElement,
3018{
3019    let fill = read_unaligned_scalar::<T>(fill);
3020    let operand_data = operand.data_as::<T>()?;
3021    let dest_dims = dest.dims();
3022    let dest_strides = dest.strides();
3023    let dest_offset = dest.offset();
3024    let dest_data = dest.data_as_uninit_mut::<T>()?;
3025    let operand_ref = unsafe {
3026        RawStridedRef::new_unchecked(
3027            operand_data,
3028            operand.dims(),
3029            operand.strides(),
3030            operand.offset(),
3031        )
3032    };
3033    let mut dest_ref =
3034        unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) };
3035    plan.execute_uninit(&mut dest_ref, &operand_ref, fill)
3036}
3037
3038fn execute_concatenate<T>(
3039    plan: &ConcatenatePlan,
3040    dest: &mut ErasedRawStridedMut<'_>,
3041    inputs: &[ErasedRawStridedRef<'_>],
3042) -> Result<()>
3043where
3044    T: Copy + crate::MaybeSendSync + KernelStorageElement,
3045{
3046    if inputs.len() != plan.input_count() {
3047        return Err(StridedError::RankMismatch(inputs.len(), plan.input_count()));
3048    }
3049    let dest_dims = dest.dims();
3050    let dest_strides = dest.strides();
3051    let dest_offset = dest.offset();
3052    let dest_data = dest.data_as_mut::<T>()?;
3053    let mut dest_ref =
3054        unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) };
3055    plan.check_dest_layout(&dest_ref)?;
3056
3057    for (position, input) in inputs.iter().enumerate() {
3058        let input_data = input.data_as::<T>()?;
3059        let input_ref = unsafe {
3060            RawStridedRef::new_unchecked(input_data, input.dims(), input.strides(), input.offset())
3061        };
3062        plan.check_input_layout(position, &input_ref)?;
3063        plan.segment_offset(position, dest_offset)?;
3064    }
3065    for (position, input) in inputs.iter().enumerate() {
3066        let input_data = input.data_as::<T>()?;
3067        let input_ref = unsafe {
3068            RawStridedRef::new_unchecked(input_data, input.dims(), input.strides(), input.offset())
3069        };
3070        plan.execute_segment(position, &mut dest_ref, &input_ref)?;
3071    }
3072    Ok(())
3073}
3074
3075fn execute_concatenate_uninit<T>(
3076    plan: &ConcatenatePlan,
3077    dest: &mut ErasedRawStridedUninitMut<'_>,
3078    inputs: &[ErasedRawStridedPtr<'_>],
3079) -> Result<()>
3080where
3081    T: Copy + crate::MaybeSendSync + KernelStorageElement,
3082{
3083    let dest_dims = dest.dims();
3084    let dest_strides = dest.strides();
3085    let dest_offset = dest.offset();
3086    let dest_data = dest.data_as_uninit_mut::<T>()?;
3087    let mut dest_ref =
3088        unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) };
3089    plan.check_dest_layout(&dest_ref)?;
3090
3091    for (position, input) in inputs.iter().enumerate() {
3092        let input = validated_input_ref(input)?;
3093        let input_data = input.data_as::<T>()?;
3094        let input_ref = unsafe {
3095            RawStridedRef::new_unchecked(input_data, input.dims(), input.strides(), input.offset())
3096        };
3097        plan.check_input_layout(position, &input_ref)?;
3098        plan.segment_offset(position, dest_offset)?;
3099    }
3100    for (position, input) in inputs.iter().enumerate() {
3101        let input = validated_input_ref(input)?;
3102        let input_data = input.data_as::<T>()?;
3103        let input_ref = unsafe {
3104            RawStridedRef::new_unchecked(input_data, input.dims(), input.strides(), input.offset())
3105        };
3106        plan.execute_segment_uninit(position, &mut dest_ref, &input_ref)?;
3107    }
3108    Ok(())
3109}
3110
3111fn execute_reduce<T, W>(
3112    op: ReduceOp,
3113    ctx: &ExecContext,
3114    dest: &mut W,
3115    src: &ErasedRawStridedRef<'_>,
3116) -> Result<()>
3117where
3118    T: ErasedReduceScalar,
3119    W: ReduceWriter<T>,
3120{
3121    let use_serial = ctx.is_serial()
3122        || ctx
3123            .max_threads_limit()
3124            .is_some_and(|max_threads| max_threads.get() == 1);
3125    let value = if use_serial {
3126        if let Some(value) = reduce_contiguous_serial(op, src) {
3127            value
3128        } else {
3129            let source = erased_view::<T>(src)?;
3130            crate::reduce_view::reduce_serial(
3131                &source,
3132                |value| reduce_map_value(op, value),
3133                |a, b| reduce_values(op, a, b),
3134                reduce_identity(op),
3135            )?
3136        }
3137    } else {
3138        let source = erased_view::<T>(src)?;
3139        ctx.run(|| {
3140            crate::reduce(
3141                &source,
3142                |value| reduce_map_value(op, value),
3143                |a, b| reduce_values(op, a, b),
3144                reduce_identity(op),
3145            )
3146        })?
3147    };
3148
3149    // SAFETY: validated rank-zero destination layout proves the offset.
3150    unsafe { dest.write_at(dest.offset(), value) };
3151    Ok(())
3152}
3153
3154fn reduce_contiguous_serial<T>(op: ReduceOp, src: &ErasedRawStridedRef<'_>) -> Option<T>
3155where
3156    T: ErasedReduceScalar,
3157{
3158    crate::kernel::same_contiguous_layout(src.dims(), &[src.strides()])?;
3159    let len = checked_total_len(src.dims()).ok()?;
3160    if len == 0 {
3161        return Some(reduce_identity(op));
3162    }
3163
3164    let source_data = src.data_as::<T>().ok()?;
3165    let start = usize::try_from(src.offset()).ok()?;
3166    let end = start.checked_add(len)?;
3167    let values = source_data.get(start..end)?;
3168    Some(match op {
3169        ReduceOp::Sum => T::try_simd_sum(values)
3170            .unwrap_or_else(|| reduce_contiguous_lanes(values, T::zero(), T::reduce_sum)),
3171        ReduceOp::Product => T::try_simd_product(values)
3172            .unwrap_or_else(|| reduce_contiguous_lanes(values, T::one(), T::reduce_product)),
3173        ReduceOp::SumSquares => T::try_simd_sum_squares(values).unwrap_or_else(|| {
3174            reduce_contiguous_mapped_lanes(
3175                values,
3176                T::zero(),
3177                |value| T::reduce_product(value, value),
3178                T::reduce_sum,
3179            )
3180        }),
3181    })
3182}
3183
3184#[inline]
3185fn reduce_contiguous_lanes<T>(values: &[T], identity: T, combine: impl Fn(T, T) -> T) -> T
3186where
3187    T: Copy,
3188{
3189    reduce_contiguous_mapped_lanes(values, identity, |value| value, combine)
3190}
3191
3192#[inline]
3193fn reduce_contiguous_mapped_lanes<T>(
3194    values: &[T],
3195    identity: T,
3196    map: impl Fn(T) -> T,
3197    combine: impl Fn(T, T) -> T,
3198) -> T
3199where
3200    T: Copy,
3201{
3202    let mut lanes = [identity; SERIAL_REDUCE_LANES];
3203    let mut chunks = values.chunks_exact(SERIAL_REDUCE_LANES);
3204    for chunk in chunks.by_ref() {
3205        for lane in 0..SERIAL_REDUCE_LANES {
3206            lanes[lane] = combine(lanes[lane], map(chunk[lane]));
3207        }
3208    }
3209    for (lane, &value) in chunks.remainder().iter().enumerate() {
3210        lanes[lane] = combine(lanes[lane], map(value));
3211    }
3212    lanes.into_iter().fold(identity, combine)
3213}
3214
3215fn dispatch_reduce<T, W>(
3216    op: ReduceOp,
3217    layout: &ReduceLayout,
3218    ctx: &ExecContext,
3219    dest: &mut W,
3220    src: &ErasedRawStridedRef<'_>,
3221) -> Result<()>
3222where
3223    T: ErasedReduceScalar,
3224    W: ReduceWriter<T>,
3225{
3226    match layout {
3227        ReduceLayout::Full { .. } => execute_reduce::<T, W>(op, ctx, dest, src),
3228        ReduceLayout::Axes {
3229            src_dims,
3230            src_strides,
3231            dest_dims,
3232            dest_strides,
3233            axes,
3234            kept_axes,
3235            reduce_dims,
3236            dest_total,
3237            reduce_total,
3238        } => execute_reduce_axes::<T, W>(
3239            op,
3240            ctx,
3241            dest,
3242            src,
3243            AxesLayout {
3244                src_dims,
3245                src_strides,
3246                dest_dims,
3247                dest_strides,
3248                axes,
3249                kept_axes,
3250                reduce_dims,
3251                dest_total: *dest_total,
3252                reduce_total: *reduce_total,
3253            },
3254        ),
3255    }
3256}
3257
3258fn execute_reduce_axes<T, W>(
3259    op: ReduceOp,
3260    ctx: &ExecContext,
3261    dest: &mut W,
3262    src: &ErasedRawStridedRef<'_>,
3263    layout: AxesLayout<'_>,
3264) -> Result<()>
3265where
3266    T: ErasedReduceScalar,
3267    W: ReduceWriter<T>,
3268{
3269    if layout.kept_axes.is_empty()
3270        && layout.axes.len() == layout.src_dims.len()
3271        && layout.dest_total == 1
3272    {
3273        return execute_reduce::<T, W>(op, ctx, dest, src);
3274    }
3275
3276    if layout.dest_total == 0 {
3277        return Ok(());
3278    }
3279
3280    if ctx.is_serial() {
3281        execute_reduce_axes_serial::<T, W>(op, dest, src, layout)
3282    } else {
3283        ctx.run(|| execute_reduce_axes_policy::<T, W>(op, dest, src, layout))
3284    }
3285}
3286
3287fn execute_reduce_axes_policy<T, W>(
3288    op: ReduceOp,
3289    dest: &mut W,
3290    src: &ErasedRawStridedRef<'_>,
3291    layout: AxesLayout<'_>,
3292) -> Result<()>
3293where
3294    T: ErasedReduceScalar,
3295    W: ReduceWriter<T>,
3296{
3297    let source_data = src.data_as::<T>()?;
3298    let dest_offset_base = dest.offset();
3299    #[cfg(feature = "parallel")]
3300    {
3301        let nthreads = crate::threading::parallel_threads_for_len(layout.dest_total);
3302        if nthreads > 1 {
3303            return execute_reduce_axes_parallel(
3304                op,
3305                dest_offset_base,
3306                dest,
3307                src.offset(),
3308                source_data,
3309                layout,
3310                nthreads,
3311            );
3312        }
3313    }
3314
3315    execute_reduce_axes_serial_data(
3316        op,
3317        dest_offset_base,
3318        dest,
3319        src.offset(),
3320        source_data,
3321        layout,
3322    )
3323}
3324
3325fn execute_reduce_axes_serial<T, W>(
3326    op: ReduceOp,
3327    dest: &mut W,
3328    src: &ErasedRawStridedRef<'_>,
3329    layout: AxesLayout<'_>,
3330) -> Result<()>
3331where
3332    T: ErasedReduceScalar,
3333    W: ReduceWriter<T>,
3334{
3335    let source_data = src.data_as::<T>()?;
3336    let dest_offset_base = dest.offset();
3337    execute_reduce_axes_serial_data(
3338        op,
3339        dest_offset_base,
3340        dest,
3341        src.offset(),
3342        source_data,
3343        layout,
3344    )
3345}
3346
3347fn execute_reduce_axes_serial_data<T, W>(
3348    op: ReduceOp,
3349    dest_offset_base: isize,
3350    dest: &mut W,
3351    source_offset_base: isize,
3352    source_data: &[T],
3353    layout: AxesLayout<'_>,
3354) -> Result<()>
3355where
3356    T: ErasedReduceScalar,
3357    W: ReduceWriter<T>,
3358{
3359    let mut out_idx_storage = CoordScratch::new(layout.dest_dims.len());
3360    let mut reduce_idx_storage = CoordScratch::new(layout.reduce_dims.len());
3361    let mut src_idx_storage = CoordScratch::new(layout.src_dims.len());
3362    let out_idx = out_idx_storage.as_mut_slice();
3363    let reduce_idx = reduce_idx_storage.as_mut_slice();
3364    let src_idx = src_idx_storage.as_mut_slice();
3365
3366    for _ in 0..layout.dest_total {
3367        src_idx.fill(0);
3368        for (dest_axis, &src_axis) in layout.kept_axes.iter().enumerate() {
3369            src_idx[src_axis] = out_idx[dest_axis];
3370        }
3371
3372        let mut acc = reduce_identity(op);
3373        reduce_idx.fill(0);
3374        for _ in 0..layout.reduce_total {
3375            for (reduce_axis, &src_axis) in layout.axes.iter().enumerate() {
3376                src_idx[src_axis] = reduce_idx[reduce_axis];
3377            }
3378            let source_offset =
3379                checked_strided_offset(source_offset_base, layout.src_strides, src_idx)?;
3380            let value = unsafe { *source_data.as_ptr().offset(source_offset) };
3381            acc = reduce_values(op, acc, reduce_map_value(op, value));
3382            advance_col_major_index(reduce_idx, layout.reduce_dims);
3383        }
3384
3385        let dest_offset = checked_strided_offset(dest_offset_base, layout.dest_strides, out_idx)?;
3386        // SAFETY: reduction layout and extent validation prove the offset.
3387        unsafe { dest.write_at(dest_offset, acc) };
3388        advance_col_major_index(out_idx, layout.dest_dims);
3389    }
3390    Ok(())
3391}
3392
3393#[cfg(feature = "parallel")]
3394fn execute_reduce_axes_parallel<T, W>(
3395    op: ReduceOp,
3396    dest_offset_base: isize,
3397    dest: &mut W,
3398    source_offset_base: isize,
3399    source_data: &[T],
3400    layout: AxesLayout<'_>,
3401    nthreads: usize,
3402) -> Result<()>
3403where
3404    T: ErasedReduceScalar,
3405    W: ReduceWriter<T>,
3406{
3407    // SAFETY: the validated reduction writer owns the destination allocation.
3408    let dest_ptr = crate::threading::SendPtr(unsafe { dest.ptr() });
3409    let source_ptr = crate::threading::SendPtr(source_data.as_ptr() as *mut T);
3410    crate::threading::parallel_map_reduce(
3411        0..layout.dest_total,
3412        nthreads,
3413        &|range| {
3414            let mut out_idx_storage = CoordScratch::new(layout.dest_dims.len());
3415            let mut reduce_idx_storage = CoordScratch::new(layout.reduce_dims.len());
3416            let mut src_idx_storage = CoordScratch::new(layout.src_dims.len());
3417            let out_idx = out_idx_storage.as_mut_slice();
3418            let reduce_idx = reduce_idx_storage.as_mut_slice();
3419            let src_idx = src_idx_storage.as_mut_slice();
3420            fill_col_major_index(range.start, layout.dest_dims, out_idx);
3421            let dest_ptr = dest_ptr.as_ptr();
3422            let source_ptr = source_ptr.as_const();
3423
3424            for _ in range {
3425                src_idx.fill(0);
3426                for (dest_axis, &src_axis) in layout.kept_axes.iter().enumerate() {
3427                    src_idx[src_axis] = out_idx[dest_axis];
3428                }
3429
3430                let mut acc = reduce_identity(op);
3431                reduce_idx.fill(0);
3432                for _ in 0..layout.reduce_total {
3433                    for (reduce_axis, &src_axis) in layout.axes.iter().enumerate() {
3434                        src_idx[src_axis] = reduce_idx[reduce_axis];
3435                    }
3436                    let source_offset =
3437                        checked_strided_offset(source_offset_base, layout.src_strides, src_idx)?;
3438                    let value = unsafe { *source_ptr.offset(source_offset) };
3439                    acc = reduce_values(op, acc, reduce_map_value(op, value));
3440                    advance_col_major_index(reduce_idx, layout.reduce_dims);
3441                }
3442
3443                let dest_offset =
3444                    checked_strided_offset(dest_offset_base, layout.dest_strides, out_idx)?;
3445                unsafe {
3446                    // SAFETY: axis reduction writes exactly one scalar per
3447                    // logical output position, and compile rejected
3448                    // non-injective destination layouts.
3449                    dest_ptr.offset(dest_offset).write(acc);
3450                }
3451                advance_col_major_index(out_idx, layout.dest_dims);
3452            }
3453            Ok(())
3454        },
3455        &|left, right| left.and(right),
3456    )
3457}
3458
3459#[inline]
3460fn reduce_identity<T>(op: ReduceOp) -> T
3461where
3462    T: One + Zero,
3463{
3464    match op {
3465        ReduceOp::Sum => T::zero(),
3466        ReduceOp::Product => T::one(),
3467        ReduceOp::SumSquares => T::zero(),
3468    }
3469}
3470
3471#[inline]
3472fn reduce_values<T>(op: ReduceOp, a: T, b: T) -> T
3473where
3474    T: ErasedReduceScalar,
3475{
3476    match op {
3477        ReduceOp::Sum => T::reduce_sum(a, b),
3478        ReduceOp::Product => T::reduce_product(a, b),
3479        ReduceOp::SumSquares => T::reduce_sum(a, b),
3480    }
3481}
3482
3483#[inline]
3484fn reduce_map_value<T>(op: ReduceOp, value: T) -> T
3485where
3486    T: ErasedReduceScalar,
3487{
3488    match op {
3489        ReduceOp::Sum | ReduceOp::Product => value,
3490        ReduceOp::SumSquares => T::reduce_product(value, value),
3491    }
3492}
3493
3494trait ErasedReduceScalar:
3495    KernelStorageElement
3496    + Copy
3497    + One
3498    + Zero
3499    + crate::MaybeSendSync
3500    + crate::simd::MaybeSimdOps
3501    + crate::simd::MaybeSimdProduct
3502    + crate::simd::MaybeSimdSumSquares
3503{
3504    fn reduce_sum(lhs: Self, rhs: Self) -> Self;
3505    fn reduce_product(lhs: Self, rhs: Self) -> Self;
3506}
3507
3508macro_rules! impl_default_erased_reduce_scalar {
3509    ($($ty:ty),* $(,)?) => {
3510        $(
3511            impl ErasedReduceScalar for $ty {
3512                #[inline(always)]
3513                fn reduce_sum(lhs: Self, rhs: Self) -> Self {
3514                    lhs + rhs
3515                }
3516
3517                #[inline(always)]
3518                fn reduce_product(lhs: Self, rhs: Self) -> Self {
3519                    lhs * rhs
3520                }
3521            }
3522        )*
3523    };
3524}
3525
3526macro_rules! impl_wrapping_erased_reduce_scalar {
3527    ($($ty:ty),* $(,)?) => {
3528        $(
3529            impl ErasedReduceScalar for $ty {
3530                #[inline(always)]
3531                fn reduce_sum(lhs: Self, rhs: Self) -> Self {
3532                    lhs.wrapping_add(rhs)
3533                }
3534
3535                #[inline(always)]
3536                fn reduce_product(lhs: Self, rhs: Self) -> Self {
3537                    lhs.wrapping_mul(rhs)
3538                }
3539            }
3540        )*
3541    };
3542}
3543
3544impl_default_erased_reduce_scalar!(f32, f64, Complex32, Complex64);
3545impl_wrapping_erased_reduce_scalar!(i32, i64);
3546
3547fn validate_unique_axes(axes: &[usize], rank: usize) -> Result<()> {
3548    let mut seen = vec![false; rank];
3549    for &axis in axes {
3550        if axis >= rank {
3551            return Err(StridedError::InvalidAxis { axis, rank });
3552        }
3553        if seen[axis] {
3554            return Err(StridedError::InvalidAxis { axis, rank });
3555        }
3556        seen[axis] = true;
3557    }
3558    Ok(())
3559}
3560
3561fn checked_strided_offset(base: isize, strides: &[isize], index: &[usize]) -> Result<isize> {
3562    let mut offset = base;
3563    for (&stride, &coord) in strides.iter().zip(index.iter()) {
3564        offset = checked_offset_add(offset, stride, coord)?;
3565    }
3566    Ok(offset)
3567}
3568
3569fn checked_offset_add(base: isize, stride: isize, coord: usize) -> Result<isize> {
3570    let coord = isize::try_from(coord).map_err(|_| StridedError::OffsetOverflow)?;
3571    let scaled = stride
3572        .checked_mul(coord)
3573        .ok_or(StridedError::OffsetOverflow)?;
3574    base.checked_add(scaled).ok_or(StridedError::OffsetOverflow)
3575}
3576
3577fn advance_col_major_index(index: &mut [usize], shape: &[usize]) {
3578    for axis in 0..index.len() {
3579        index[axis] += 1;
3580        if index[axis] < shape[axis] {
3581            return;
3582        }
3583        index[axis] = 0;
3584    }
3585}
3586
3587#[cfg(feature = "parallel")]
3588fn fill_col_major_index(mut linear: usize, shape: &[usize], out: &mut [usize]) {
3589    for (axis, coord) in out.iter_mut().enumerate() {
3590        let dim = shape[axis];
3591        *coord = linear % dim;
3592        linear /= dim;
3593    }
3594}
3595
3596struct CoordScratch {
3597    inline: [usize; RAW_FUSED_RANK_LIMIT],
3598    heap: Option<Vec<usize>>,
3599    len: usize,
3600}
3601
3602impl CoordScratch {
3603    fn new(len: usize) -> Self {
3604        if len <= RAW_FUSED_RANK_LIMIT {
3605            Self {
3606                inline: [0; RAW_FUSED_RANK_LIMIT],
3607                heap: None,
3608                len,
3609            }
3610        } else {
3611            Self {
3612                inline: [0; RAW_FUSED_RANK_LIMIT],
3613                heap: Some(vec![0; len]),
3614                len,
3615            }
3616        }
3617    }
3618
3619    fn as_mut_slice(&mut self) -> &mut [usize] {
3620        match &mut self.heap {
3621            Some(heap) => heap,
3622            None => &mut self.inline[..self.len],
3623        }
3624    }
3625}
3626
3627fn execute_fused<T>(
3628    plan: &FusedPlan,
3629    ctx: &ExecContext,
3630    dest: &mut ErasedRawStridedMut<'_>,
3631    inputs: &[ErasedRawStridedRef<'_>],
3632) -> Result<()>
3633where
3634    T: FusedScalar + KernelStorageElement,
3635{
3636    let dest_dims = dest.dims();
3637    let dest_strides = dest.strides();
3638    let dest_offset = dest.offset();
3639    let dest_data = dest.data_as_mut::<T>()?;
3640    let dest_view =
3641        unsafe { StridedViewMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) };
3642
3643    match inputs {
3644        [a] => {
3645            let input_views = [erased_view::<T>(a)?];
3646            let mut dests = [dest_view];
3647            execute_fused_views(ctx, &mut dests, &input_views, plan)
3648        }
3649        [a, b] => {
3650            let input_views = [erased_view::<T>(a)?, erased_view::<T>(b)?];
3651            let mut dests = [dest_view];
3652            execute_fused_views(ctx, &mut dests, &input_views, plan)
3653        }
3654        [a, b, c] => {
3655            let input_views = [
3656                erased_view::<T>(a)?,
3657                erased_view::<T>(b)?,
3658                erased_view::<T>(c)?,
3659            ];
3660            let mut dests = [dest_view];
3661            execute_fused_views(ctx, &mut dests, &input_views, plan)
3662        }
3663        [a, b, c, d] => {
3664            let input_views = [
3665                erased_view::<T>(a)?,
3666                erased_view::<T>(b)?,
3667                erased_view::<T>(c)?,
3668                erased_view::<T>(d)?,
3669            ];
3670            let mut dests = [dest_view];
3671            execute_fused_views(ctx, &mut dests, &input_views, plan)
3672        }
3673        _ => Err(StridedError::UnsupportedArity {
3674            arity: inputs.len(),
3675            max: ERASED_FUSED_INPUT_LIMIT,
3676        }),
3677    }
3678}
3679
3680fn execute_fused_uninit<T>(
3681    plan: &FusedPlan,
3682    dest: &mut ErasedRawStridedUninitMut<'_>,
3683    inputs: &[ErasedRawStridedRef<'_>],
3684    serial: bool,
3685    validated: crate::map_view::ValidatedDestinationLayout,
3686) -> Result<()>
3687where
3688    T: FusedScalar + KernelStorageElement,
3689{
3690    let dims = dest.dims();
3691    let strides = dest.strides();
3692    let offset = dest.offset();
3693    let dest_data = dest.data_as_uninit_mut::<T>()?;
3694    let mut dest_view = unsafe { StridedViewMut::new_unchecked(dest_data, dims, strides, offset) };
3695    match inputs {
3696        [a] => {
3697            let input_views = [erased_view::<T>(a)?];
3698            crate::fused::fused_elementwise_into_uninit(
3699                &mut dest_view,
3700                &input_views,
3701                plan,
3702                serial,
3703                validated,
3704            )
3705        }
3706        [a, b] => {
3707            let input_views = [erased_view::<T>(a)?, erased_view::<T>(b)?];
3708            crate::fused::fused_elementwise_into_uninit(
3709                &mut dest_view,
3710                &input_views,
3711                plan,
3712                serial,
3713                validated,
3714            )
3715        }
3716        [a, b, c] => {
3717            let input_views = [
3718                erased_view::<T>(a)?,
3719                erased_view::<T>(b)?,
3720                erased_view::<T>(c)?,
3721            ];
3722            crate::fused::fused_elementwise_into_uninit(
3723                &mut dest_view,
3724                &input_views,
3725                plan,
3726                serial,
3727                validated,
3728            )
3729        }
3730        [a, b, c, d] => {
3731            let input_views = [
3732                erased_view::<T>(a)?,
3733                erased_view::<T>(b)?,
3734                erased_view::<T>(c)?,
3735                erased_view::<T>(d)?,
3736            ];
3737            crate::fused::fused_elementwise_into_uninit(
3738                &mut dest_view,
3739                &input_views,
3740                plan,
3741                serial,
3742                validated,
3743            )
3744        }
3745        _ => Err(StridedError::UnsupportedArity {
3746            arity: inputs.len(),
3747            max: ERASED_FUSED_INPUT_LIMIT,
3748        }),
3749    }
3750}
3751
3752fn execute_fused_uninit_ptrs<T>(
3753    plan: &FusedPlan,
3754    dest: &mut ErasedRawStridedUninitMut<'_>,
3755    inputs: &[ErasedRawStridedPtr<'_>],
3756    serial: bool,
3757    validated: crate::map_view::ValidatedDestinationLayout,
3758) -> Result<()>
3759where
3760    T: FusedScalar + KernelStorageElement,
3761{
3762    match inputs {
3763        [a] => {
3764            let refs = [validated_input_ref(a)?];
3765            execute_fused_uninit::<T>(plan, dest, &refs, serial, validated)
3766        }
3767        [a, b] => {
3768            let refs = [validated_input_ref(a)?, validated_input_ref(b)?];
3769            execute_fused_uninit::<T>(plan, dest, &refs, serial, validated)
3770        }
3771        [a, b, c] => {
3772            let refs = [
3773                validated_input_ref(a)?,
3774                validated_input_ref(b)?,
3775                validated_input_ref(c)?,
3776            ];
3777            execute_fused_uninit::<T>(plan, dest, &refs, serial, validated)
3778        }
3779        [a, b, c, d] => {
3780            let refs = [
3781                validated_input_ref(a)?,
3782                validated_input_ref(b)?,
3783                validated_input_ref(c)?,
3784                validated_input_ref(d)?,
3785            ];
3786            execute_fused_uninit::<T>(plan, dest, &refs, serial, validated)
3787        }
3788        _ => Err(StridedError::UnsupportedArity {
3789            arity: inputs.len(),
3790            max: ERASED_FUSED_INPUT_LIMIT,
3791        }),
3792    }
3793}
3794
3795fn dispatch_gather_index<T>(
3796    plan: &GatherPlan,
3797    index_dtype: KernelDType,
3798    dest: &mut ErasedRawStridedMut<'_>,
3799    operand: &ErasedRawStridedRef<'_>,
3800    start_indices: &ErasedRawStridedRef<'_>,
3801) -> Result<()>
3802where
3803    T: Copy + crate::MaybeSendSync + KernelStorageElement,
3804{
3805    match index_dtype {
3806        KernelDType::I32 => execute_gather::<T, i32>(plan, dest, operand, start_indices),
3807        KernelDType::I64 => execute_gather::<T, i64>(plan, dest, operand, start_indices),
3808        _ => Err(StridedError::UnsupportedDType {
3809            dtype: index_dtype.label(),
3810        }),
3811    }
3812}
3813
3814fn execute_gather<T, I>(
3815    plan: &GatherPlan,
3816    dest: &mut ErasedRawStridedMut<'_>,
3817    operand: &ErasedRawStridedRef<'_>,
3818    start_indices: &ErasedRawStridedRef<'_>,
3819) -> Result<()>
3820where
3821    T: Copy + crate::MaybeSendSync + KernelStorageElement,
3822    I: GatherIndex + KernelStorageElement,
3823{
3824    let operand_data = operand.data_as::<T>()?;
3825    let index_data = start_indices.data_as::<I>()?;
3826    let dest_dims = dest.dims();
3827    let dest_strides = dest.strides();
3828    let dest_offset = dest.offset();
3829    let dest_data = dest.data_as_mut::<T>()?;
3830    let operand_ref = unsafe {
3831        RawStridedRef::new_unchecked(
3832            operand_data,
3833            operand.dims(),
3834            operand.strides(),
3835            operand.offset(),
3836        )
3837    };
3838    let index_ref = unsafe {
3839        RawStridedRef::new_unchecked(
3840            index_data,
3841            start_indices.dims(),
3842            start_indices.strides(),
3843            start_indices.offset(),
3844        )
3845    };
3846    let mut dest_ref =
3847        unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) };
3848    plan.execute(&mut dest_ref, &operand_ref, &index_ref)
3849}
3850
3851fn dispatch_dynamic_slice_index<T>(
3852    plan: &DynamicSlicePlan,
3853    index_dtype: KernelDType,
3854    dest: &mut ErasedRawStridedMut<'_>,
3855    operand: &ErasedRawStridedRef<'_>,
3856    starts: &ErasedRawStridedRef<'_>,
3857) -> Result<()>
3858where
3859    T: Copy + crate::MaybeSendSync + KernelStorageElement,
3860{
3861    match index_dtype {
3862        KernelDType::I32 => execute_dynamic_slice::<T, i32>(plan, dest, operand, starts),
3863        KernelDType::I64 => execute_dynamic_slice::<T, i64>(plan, dest, operand, starts),
3864        _ => Err(StridedError::UnsupportedDType {
3865            dtype: index_dtype.label(),
3866        }),
3867    }
3868}
3869
3870fn execute_dynamic_slice<T, I>(
3871    plan: &DynamicSlicePlan,
3872    dest: &mut ErasedRawStridedMut<'_>,
3873    operand: &ErasedRawStridedRef<'_>,
3874    starts: &ErasedRawStridedRef<'_>,
3875) -> Result<()>
3876where
3877    T: Copy + crate::MaybeSendSync + KernelStorageElement,
3878    I: GatherIndex + KernelStorageElement,
3879{
3880    let operand_data = operand.data_as::<T>()?;
3881    let start_data = starts.data_as::<I>()?;
3882    let dest_dims = dest.dims();
3883    let dest_strides = dest.strides();
3884    let dest_offset = dest.offset();
3885    let dest_data = dest.data_as_mut::<T>()?;
3886    let operand_ref = unsafe {
3887        RawStridedRef::new_unchecked(
3888            operand_data,
3889            operand.dims(),
3890            operand.strides(),
3891            operand.offset(),
3892        )
3893    };
3894    let start_ref = unsafe {
3895        RawStridedRef::new_unchecked(start_data, starts.dims(), starts.strides(), starts.offset())
3896    };
3897    let mut dest_ref =
3898        unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) };
3899    plan.execute(&mut dest_ref, &operand_ref, &start_ref)
3900}
3901
3902fn dispatch_dynamic_update_slice_index<T>(
3903    plan: &DynamicUpdateSlicePlan,
3904    index_dtype: KernelDType,
3905    dest: &mut ErasedRawStridedMut<'_>,
3906    operand: &ErasedRawStridedRef<'_>,
3907    update: &ErasedRawStridedRef<'_>,
3908    starts: &ErasedRawStridedRef<'_>,
3909) -> Result<()>
3910where
3911    T: Copy + crate::MaybeSendSync + KernelStorageElement,
3912{
3913    match index_dtype {
3914        KernelDType::I32 => {
3915            execute_dynamic_update_slice::<T, i32>(plan, dest, operand, update, starts)
3916        }
3917        KernelDType::I64 => {
3918            execute_dynamic_update_slice::<T, i64>(plan, dest, operand, update, starts)
3919        }
3920        _ => Err(StridedError::UnsupportedDType {
3921            dtype: index_dtype.label(),
3922        }),
3923    }
3924}
3925
3926fn execute_dynamic_update_slice<T, I>(
3927    plan: &DynamicUpdateSlicePlan,
3928    dest: &mut ErasedRawStridedMut<'_>,
3929    operand: &ErasedRawStridedRef<'_>,
3930    update: &ErasedRawStridedRef<'_>,
3931    starts: &ErasedRawStridedRef<'_>,
3932) -> Result<()>
3933where
3934    T: Copy + crate::MaybeSendSync + KernelStorageElement,
3935    I: GatherIndex + KernelStorageElement,
3936{
3937    let operand_data = operand.data_as::<T>()?;
3938    let update_data = update.data_as::<T>()?;
3939    let start_data = starts.data_as::<I>()?;
3940    let dest_dims = dest.dims();
3941    let dest_strides = dest.strides();
3942    let dest_offset = dest.offset();
3943    let dest_data = dest.data_as_mut::<T>()?;
3944    let operand_ref = unsafe {
3945        RawStridedRef::new_unchecked(
3946            operand_data,
3947            operand.dims(),
3948            operand.strides(),
3949            operand.offset(),
3950        )
3951    };
3952    let update_ref = unsafe {
3953        RawStridedRef::new_unchecked(
3954            update_data,
3955            update.dims(),
3956            update.strides(),
3957            update.offset(),
3958        )
3959    };
3960    let start_ref = unsafe {
3961        RawStridedRef::new_unchecked(start_data, starts.dims(), starts.strides(), starts.offset())
3962    };
3963    let mut dest_ref =
3964        unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) };
3965    plan.execute(&mut dest_ref, &operand_ref, &update_ref, &start_ref)
3966}
3967
3968fn dispatch_scatter_index<T>(
3969    plan: &ScatterPlan,
3970    index_dtype: KernelDType,
3971    dest: &mut ErasedRawStridedMut<'_>,
3972    operand: &ErasedRawStridedRef<'_>,
3973    scatter_indices: &ErasedRawStridedRef<'_>,
3974    updates: &ErasedRawStridedRef<'_>,
3975) -> Result<()>
3976where
3977    T: Copy + Add<Output = T> + crate::MaybeSendSync + KernelStorageElement,
3978{
3979    match index_dtype {
3980        KernelDType::I32 => {
3981            execute_scatter::<T, i32>(plan, dest, operand, scatter_indices, updates)
3982        }
3983        KernelDType::I64 => {
3984            execute_scatter::<T, i64>(plan, dest, operand, scatter_indices, updates)
3985        }
3986        _ => Err(StridedError::UnsupportedDType {
3987            dtype: index_dtype.label(),
3988        }),
3989    }
3990}
3991
3992fn execute_scatter<T, I>(
3993    plan: &ScatterPlan,
3994    dest: &mut ErasedRawStridedMut<'_>,
3995    operand: &ErasedRawStridedRef<'_>,
3996    scatter_indices: &ErasedRawStridedRef<'_>,
3997    updates: &ErasedRawStridedRef<'_>,
3998) -> Result<()>
3999where
4000    T: Copy + Add<Output = T> + crate::MaybeSendSync + KernelStorageElement,
4001    I: GatherIndex + KernelStorageElement,
4002{
4003    let operand_data = operand.data_as::<T>()?;
4004    let index_data = scatter_indices.data_as::<I>()?;
4005    let update_data = updates.data_as::<T>()?;
4006    let dest_dims = dest.dims();
4007    let dest_strides = dest.strides();
4008    let dest_offset = dest.offset();
4009    let dest_data = dest.data_as_mut::<T>()?;
4010    let operand_ref = unsafe {
4011        RawStridedRef::new_unchecked(
4012            operand_data,
4013            operand.dims(),
4014            operand.strides(),
4015            operand.offset(),
4016        )
4017    };
4018    let index_ref = unsafe {
4019        RawStridedRef::new_unchecked(
4020            index_data,
4021            scatter_indices.dims(),
4022            scatter_indices.strides(),
4023            scatter_indices.offset(),
4024        )
4025    };
4026    let update_ref = unsafe {
4027        RawStridedRef::new_unchecked(
4028            update_data,
4029            updates.dims(),
4030            updates.strides(),
4031            updates.offset(),
4032        )
4033    };
4034    let mut dest_ref =
4035        unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) };
4036    plan.execute(&mut dest_ref, &operand_ref, &index_ref, &update_ref)
4037}
4038
4039fn execute_fused_views<T>(
4040    ctx: &ExecContext,
4041    dests: &mut [StridedViewMut<'_, T>],
4042    inputs: &[StridedView<'_, T>],
4043    plan: &FusedPlan,
4044) -> Result<()>
4045where
4046    T: FusedScalar + KernelStorageElement,
4047{
4048    if ctx.is_serial() {
4049        crate::fused::fused_elementwise_into_serial(dests, inputs, plan)
4050    } else {
4051        ctx.run(|| fused_elementwise_into(dests, inputs, plan))
4052    }
4053}
4054
4055fn erased_view<'a, T: KernelStorageElement>(
4056    src: &'a ErasedRawStridedRef<'a>,
4057) -> Result<StridedView<'a, T>> {
4058    let data = src.data_as::<T>()?;
4059    Ok(unsafe { StridedView::new_unchecked(data, src.dims(), src.strides(), src.offset()) })
4060}
4061
4062fn read_unaligned_scalar<T>(bytes: &[u8]) -> T
4063where
4064    T: Copy,
4065{
4066    unsafe { core::ptr::read_unaligned(bytes.as_ptr().cast::<T>()) }
4067}
strided_kernel/erased.rs

strided_kernel/
erased.rs
