strided_kernel/
erased.rs

1use crate::*;
2use core::ops::Add;
3use num_complex::{Complex32, Complex64};
4use strided_basic::execution::check_static_indexing_dtype;
5use strided_basic::execution::{check_dtype, validate_uninit_no_overlap};
6
7mod uninit;
8pub use uninit::{
9    erased_broadcast_mul_into_uninit, erased_clamp_into_uninit, erased_compare_into_uninit,
10    erased_map_into_uninit, erased_select_into_uninit, erased_zip_into_uninit,
11};
12/// Runtime unary operation for [`erased_map_into`].
13#[non_exhaustive]
14#[derive(Clone, Copy, Debug, Eq, PartialEq)]
15pub enum ErasedMapOp {
16    Negate,
17    Conj,
18    Abs,
19    Sign,
20}
21
22impl ErasedMapOp {
23    const fn label(self) -> &'static str {
24        match self {
25            Self::Negate => "negate",
26            Self::Conj => "conj",
27            Self::Abs => "abs",
28            Self::Sign => "sign",
29        }
30    }
31}
32
33/// Runtime binary operation for [`erased_zip_into`].
34#[non_exhaustive]
35#[derive(Clone, Copy, Debug, Eq, PartialEq)]
36pub enum ErasedZipOp {
37    Add,
38    Subtract,
39    Multiply,
40    Divide,
41    Remainder,
42    Maximum,
43    Minimum,
44}
45
46impl ErasedZipOp {
47    const fn label(self) -> &'static str {
48        match self {
49            Self::Add => "add",
50            Self::Subtract => "subtract",
51            Self::Multiply => "multiply",
52            Self::Divide => "divide",
53            Self::Remainder => "remainder",
54            Self::Maximum => "maximum",
55            Self::Minimum => "minimum",
56        }
57    }
58}
59
60/// Apply one runtime-selected unary operation without compiling a plan.
61///
62/// The destination must not overlap the input. Real and complex dtypes support
63/// every [`ErasedMapOp`], signed integers use wrapping negate/abs semantics,
64/// and `bool` supports only [`ErasedMapOp::Conj`]. Complex absolute value has
65/// the real output contract `c32 -> f32` and `c64 -> f64`; all other supported
66/// unary operations preserve dtype.
67///
68/// # Errors
69///
70/// Returns a typed [`StridedError`] for dtype, shape, output-layout, overlap,
71/// or unsupported dtype/op contracts. Validation completes before any write.
72pub fn erased_map_into(
73    input_dtype: KernelDType,
74    op: ErasedMapOp,
75    dest: &mut ErasedRawStridedMut<'_>,
76    input: &ErasedRawStridedPtr<'_>,
77    exec: Exec<'_>,
78) -> Result<()> {
79    check_dtype(input_dtype, input.dtype())?;
80    check_dtype(map_output_dtype(input_dtype, op)?, dest.dtype())?;
81    validate_no_overlap(dest, input, 0)?;
82    // SAFETY: input/output overlap was rejected before forming references.
83    let input = unsafe { input.try_as_ref_after_no_overlap() }?;
84
85    match (input_dtype, op) {
86        (KernelDType::C32, ErasedMapOp::Abs) => {
87            execute_one_shot_map_with::<f32, Complex32>(dest, &input, |value| value.norm(), exec)
88        }
89        (KernelDType::C64, ErasedMapOp::Abs) => {
90            execute_one_shot_map_with::<f64, Complex64>(dest, &input, |value| value.norm(), exec)
91        }
92        (KernelDType::F32, _) => execute_one_shot_map::<f32>(op, dest, &input, exec),
93        (KernelDType::F64, _) => execute_one_shot_map::<f64>(op, dest, &input, exec),
94        (KernelDType::I32, _) => execute_one_shot_map::<i32>(op, dest, &input, exec),
95        (KernelDType::I64, _) => execute_one_shot_map::<i64>(op, dest, &input, exec),
96        (KernelDType::Bool, _) => execute_one_shot_map::<bool>(op, dest, &input, exec),
97        (KernelDType::C32, _) => execute_one_shot_map::<Complex32>(op, dest, &input, exec),
98        (KernelDType::C64, _) => execute_one_shot_map::<Complex64>(op, dest, &input, exec),
99        _ => Err(StridedError::UnsupportedDType {
100            dtype: input_dtype.label(),
101        }),
102    }
103}
104
105/// Apply one runtime-selected binary operation without compiling a plan.
106///
107/// The destination must not overlap either input. Real dtypes support every
108/// [`ErasedZipOp`]. Signed integers support every operation with wrapping
109/// arithmetic and a pre-write zero-divisor check. Complex dtypes support
110/// add/subtract/multiply/divide. `bool` has no binary one-shot operations.
111///
112/// # Errors
113///
114/// Returns a typed [`StridedError`] for dtype, shape, output-layout, overlap,
115/// or unsupported dtype/op contracts. Validation completes before any write.
116pub fn erased_zip_into(
117    dtype: KernelDType,
118    op: ErasedZipOp,
119    dest: &mut ErasedRawStridedMut<'_>,
120    lhs: &ErasedRawStridedPtr<'_>,
121    rhs: &ErasedRawStridedPtr<'_>,
122    exec: Exec<'_>,
123) -> Result<()> {
124    check_dtype(dtype, dest.dtype())?;
125    check_dtype(dtype, lhs.dtype())?;
126    check_dtype(dtype, rhs.dtype())?;
127    validate_no_overlap(dest, lhs, 0)?;
128    validate_no_overlap(dest, rhs, 1)?;
129    // SAFETY: the owning erased entry rejected all input/output overlap before conversion.
130    let lhs = unsafe { lhs.try_as_ref_after_no_overlap() }?;
131    // SAFETY: the owning erased entry rejected all input/output overlap before conversion.
132    let rhs = unsafe { rhs.try_as_ref_after_no_overlap() }?;
133
134    match dtype {
135        KernelDType::F32 => execute_one_shot_zip::<f32>(op, dest, &lhs, &rhs, exec),
136        KernelDType::F64 => execute_one_shot_zip::<f64>(op, dest, &lhs, &rhs, exec),
137        KernelDType::I32 => execute_one_shot_zip::<i32>(op, dest, &lhs, &rhs, exec),
138        KernelDType::I64 => execute_one_shot_zip::<i64>(op, dest, &lhs, &rhs, exec),
139        KernelDType::Bool => execute_one_shot_zip::<bool>(op, dest, &lhs, &rhs, exec),
140        KernelDType::C32 => execute_one_shot_zip::<Complex32>(op, dest, &lhs, &rhs, exec),
141        KernelDType::C64 => execute_one_shot_zip::<Complex64>(op, dest, &lhs, &rhs, exec),
142        _ => Err(StridedError::UnsupportedDType {
143            dtype: dtype.label(),
144        }),
145    }
146}
147
148/// Dtype-erased static-slice wrapper.
149#[derive(Clone, Debug)]
150pub struct ErasedSlicePlan {
151    dtype: KernelDType,
152    plan: SlicePlan,
153}
154
155/// Dtype-erased reverse wrapper.
156#[derive(Clone, Debug)]
157pub struct ErasedReversePlan {
158    dtype: KernelDType,
159    plan: ReversePlan,
160}
161
162/// Dtype-erased pad wrapper.
163#[derive(Clone, Debug)]
164pub struct ErasedPadPlan {
165    dtype: KernelDType,
166    plan: PadPlan,
167}
168
169impl ErasedSlicePlan {
170    /// Validate and store a static slice plan for one dtype and fixed layout set.
171    #[allow(clippy::too_many_arguments)]
172    pub fn compile(
173        dtype: KernelDType,
174        operand_dims: &[usize],
175        operand_strides: &[isize],
176        dest_dims: &[usize],
177        dest_strides: &[isize],
178        starts: &[usize],
179        limits: &[usize],
180        slice_strides: &[usize],
181    ) -> Result<Self> {
182        check_static_indexing_dtype(dtype)?;
183        Ok(Self {
184            dtype,
185            plan: SlicePlan::compile(
186                operand_dims,
187                operand_strides,
188                dest_dims,
189                dest_strides,
190                starts,
191                limits,
192                slice_strides,
193            )?,
194        })
195    }
196
197    #[inline]
198    pub fn dtype(&self) -> KernelDType {
199        self.dtype
200    }
201
202    #[inline]
203    pub fn plan(&self) -> &SlicePlan {
204        &self.plan
205    }
206
207    /// Execute a static slice into an erased output descriptor.
208    pub fn execute(
209        &self,
210        dest: &mut ErasedRawStridedMut<'_>,
211        operand: &ErasedRawStridedRef<'_>,
212        exec: Exec<'_>,
213    ) -> Result<()> {
214        check_dtype(self.dtype, dest.dtype())?;
215        check_dtype(self.dtype, operand.dtype())?;
216
217        match self.dtype {
218            KernelDType::F32 => execute_slice::<f32>(&self.plan, dest, operand, exec),
219            KernelDType::F64 => execute_slice::<f64>(&self.plan, dest, operand, exec),
220            KernelDType::I32 => execute_slice::<i32>(&self.plan, dest, operand, exec),
221            KernelDType::I64 => execute_slice::<i64>(&self.plan, dest, operand, exec),
222            KernelDType::Bool => execute_slice::<bool>(&self.plan, dest, operand, exec),
223            KernelDType::C32 => execute_slice::<Complex32>(&self.plan, dest, operand, exec),
224            KernelDType::C64 => execute_slice::<Complex64>(&self.plan, dest, operand, exec),
225            _ => Err(StridedError::UnsupportedDType {
226                dtype: self.dtype.label(),
227            }),
228        }
229    }
230
231    /// Execute a static slice as a full overwrite of uninitialized output storage.
232    /// On success, every reachable destination slot is fully overwritten;
233    /// unreachable holes are neither read nor initialized. Validation errors
234    /// are returned before any destination write. A panic during execution
235    /// may leave a partially initialized `MaybeUninit` destination, which is
236    /// still safely droppable; no readable value is promised for unwritten
237    /// reachable slots.
238    pub fn execute_uninit(
239        &self,
240        dest: &mut ErasedRawStridedUninitMut<'_>,
241        operand: &ErasedRawStridedPtr<'_>,
242        exec: Exec<'_>,
243    ) -> Result<()> {
244        check_dtype(self.dtype, dest.dtype())?;
245        check_dtype(self.dtype, operand.dtype())?;
246        validate_uninit_no_overlap(dest, operand, 0)?;
247        // SAFETY: the owning erased entry rejected all input/output overlap before conversion.
248        let operand = unsafe { operand.try_as_ref_after_no_overlap() }?;
249
250        match self.dtype {
251            KernelDType::F32 => execute_slice_uninit::<f32>(&self.plan, dest, &operand, exec),
252            KernelDType::F64 => execute_slice_uninit::<f64>(&self.plan, dest, &operand, exec),
253            KernelDType::I32 => execute_slice_uninit::<i32>(&self.plan, dest, &operand, exec),
254            KernelDType::I64 => execute_slice_uninit::<i64>(&self.plan, dest, &operand, exec),
255            KernelDType::Bool => execute_slice_uninit::<bool>(&self.plan, dest, &operand, exec),
256            KernelDType::C32 => execute_slice_uninit::<Complex32>(&self.plan, dest, &operand, exec),
257            KernelDType::C64 => execute_slice_uninit::<Complex64>(&self.plan, dest, &operand, exec),
258            _ => Err(StridedError::UnsupportedDType {
259                dtype: self.dtype.label(),
260            }),
261        }
262    }
263}
264
265impl ErasedReversePlan {
266    /// Validate and store a reverse plan for one dtype and fixed layout set.
267    pub fn compile(
268        dtype: KernelDType,
269        operand_dims: &[usize],
270        operand_strides: &[isize],
271        dest_strides: &[isize],
272        axes: &[usize],
273    ) -> Result<Self> {
274        check_static_indexing_dtype(dtype)?;
275        Ok(Self {
276            dtype,
277            plan: ReversePlan::compile(operand_dims, operand_strides, dest_strides, axes)?,
278        })
279    }
280
281    #[inline]
282    pub fn dtype(&self) -> KernelDType {
283        self.dtype
284    }
285
286    #[inline]
287    pub fn plan(&self) -> &ReversePlan {
288        &self.plan
289    }
290
291    /// Execute a reverse into an erased output descriptor.
292    pub fn execute(
293        &self,
294        dest: &mut ErasedRawStridedMut<'_>,
295        operand: &ErasedRawStridedRef<'_>,
296        exec: Exec<'_>,
297    ) -> Result<()> {
298        check_dtype(self.dtype, dest.dtype())?;
299        check_dtype(self.dtype, operand.dtype())?;
300
301        match self.dtype {
302            KernelDType::F32 => execute_reverse::<f32>(&self.plan, dest, operand, exec),
303            KernelDType::F64 => execute_reverse::<f64>(&self.plan, dest, operand, exec),
304            KernelDType::I32 => execute_reverse::<i32>(&self.plan, dest, operand, exec),
305            KernelDType::I64 => execute_reverse::<i64>(&self.plan, dest, operand, exec),
306            KernelDType::Bool => execute_reverse::<bool>(&self.plan, dest, operand, exec),
307            KernelDType::C32 => execute_reverse::<Complex32>(&self.plan, dest, operand, exec),
308            KernelDType::C64 => execute_reverse::<Complex64>(&self.plan, dest, operand, exec),
309            _ => Err(StridedError::UnsupportedDType {
310                dtype: self.dtype.label(),
311            }),
312        }
313    }
314
315    /// Execute reverse as a full overwrite of uninitialized output storage.
316    /// On success, every reachable destination slot is fully overwritten;
317    /// unreachable holes are neither read nor initialized. Validation errors
318    /// are returned before any destination write. A panic during execution
319    /// may leave a partially initialized `MaybeUninit` destination, which is
320    /// still safely droppable; no readable value is promised for unwritten
321    /// reachable slots.
322    pub fn execute_uninit(
323        &self,
324        dest: &mut ErasedRawStridedUninitMut<'_>,
325        operand: &ErasedRawStridedPtr<'_>,
326        exec: Exec<'_>,
327    ) -> Result<()> {
328        check_dtype(self.dtype, dest.dtype())?;
329        check_dtype(self.dtype, operand.dtype())?;
330        validate_uninit_no_overlap(dest, operand, 0)?;
331        // SAFETY: the owning erased entry rejected all input/output overlap before conversion.
332        let operand = unsafe { operand.try_as_ref_after_no_overlap() }?;
333
334        match self.dtype {
335            KernelDType::F32 => execute_reverse_uninit::<f32>(&self.plan, dest, &operand, exec),
336            KernelDType::F64 => execute_reverse_uninit::<f64>(&self.plan, dest, &operand, exec),
337            KernelDType::I32 => execute_reverse_uninit::<i32>(&self.plan, dest, &operand, exec),
338            KernelDType::I64 => execute_reverse_uninit::<i64>(&self.plan, dest, &operand, exec),
339            KernelDType::Bool => execute_reverse_uninit::<bool>(&self.plan, dest, &operand, exec),
340            KernelDType::C32 => {
341                execute_reverse_uninit::<Complex32>(&self.plan, dest, &operand, exec)
342            }
343            KernelDType::C64 => {
344                execute_reverse_uninit::<Complex64>(&self.plan, dest, &operand, exec)
345            }
346            _ => Err(StridedError::UnsupportedDType {
347                dtype: self.dtype.label(),
348            }),
349        }
350    }
351}
352
353impl ErasedPadPlan {
354    /// Validate and store a pad plan for one dtype and fixed layout set.
355    #[allow(clippy::too_many_arguments)]
356    pub fn compile(
357        dtype: KernelDType,
358        operand_dims: &[usize],
359        operand_strides: &[isize],
360        dest_dims: &[usize],
361        dest_strides: &[isize],
362        edge_padding_low: &[i64],
363        edge_padding_high: &[i64],
364        interior_padding: &[i64],
365    ) -> Result<Self> {
366        check_static_indexing_dtype(dtype)?;
367        Ok(Self {
368            dtype,
369            plan: PadPlan::compile(
370                operand_dims,
371                operand_strides,
372                dest_dims,
373                dest_strides,
374                edge_padding_low,
375                edge_padding_high,
376                interior_padding,
377            )?,
378        })
379    }
380
381    #[inline]
382    pub fn dtype(&self) -> KernelDType {
383        self.dtype
384    }
385
386    #[inline]
387    pub fn plan(&self) -> &PadPlan {
388        &self.plan
389    }
390
391    /// Execute pad into an erased output descriptor using one dtype scalar as fill.
392    pub fn execute(
393        &self,
394        dest: &mut ErasedRawStridedMut<'_>,
395        operand: &ErasedRawStridedRef<'_>,
396        fill: &[u8],
397        exec: Exec<'_>,
398    ) -> Result<()> {
399        check_dtype(self.dtype, dest.dtype())?;
400        check_dtype(self.dtype, operand.dtype())?;
401        validate_scalar_bytes(self.dtype, fill)?;
402
403        match self.dtype {
404            KernelDType::F32 => execute_pad::<f32>(&self.plan, dest, operand, fill, exec),
405            KernelDType::F64 => execute_pad::<f64>(&self.plan, dest, operand, fill, exec),
406            KernelDType::I32 => execute_pad::<i32>(&self.plan, dest, operand, fill, exec),
407            KernelDType::I64 => execute_pad::<i64>(&self.plan, dest, operand, fill, exec),
408            KernelDType::Bool => execute_pad::<bool>(&self.plan, dest, operand, fill, exec),
409            KernelDType::C32 => execute_pad::<Complex32>(&self.plan, dest, operand, fill, exec),
410            KernelDType::C64 => execute_pad::<Complex64>(&self.plan, dest, operand, fill, exec),
411            _ => Err(StridedError::UnsupportedDType {
412                dtype: self.dtype.label(),
413            }),
414        }
415    }
416
417    /// Execute pad as a full overwrite of uninitialized output storage.
418    /// On success, every reachable destination slot is fully overwritten;
419    /// unreachable holes are neither read nor initialized. Validation errors
420    /// are returned before any destination write. A panic during execution
421    /// may leave a partially initialized `MaybeUninit` destination, which is
422    /// still safely droppable; no readable value is promised for unwritten
423    /// reachable slots.
424    pub fn execute_uninit(
425        &self,
426        dest: &mut ErasedRawStridedUninitMut<'_>,
427        operand: &ErasedRawStridedPtr<'_>,
428        fill: &[u8],
429        exec: Exec<'_>,
430    ) -> Result<()> {
431        check_dtype(self.dtype, dest.dtype())?;
432        check_dtype(self.dtype, operand.dtype())?;
433        validate_scalar_bytes(self.dtype, fill)?;
434        validate_uninit_no_overlap(dest, operand, 0)?;
435        // SAFETY: the owning erased entry rejected all input/output overlap before conversion.
436        let operand = unsafe { operand.try_as_ref_after_no_overlap() }?;
437
438        match self.dtype {
439            KernelDType::F32 => execute_pad_uninit::<f32>(&self.plan, dest, &operand, fill, exec),
440            KernelDType::F64 => execute_pad_uninit::<f64>(&self.plan, dest, &operand, fill, exec),
441            KernelDType::I32 => execute_pad_uninit::<i32>(&self.plan, dest, &operand, fill, exec),
442            KernelDType::I64 => execute_pad_uninit::<i64>(&self.plan, dest, &operand, fill, exec),
443            KernelDType::Bool => execute_pad_uninit::<bool>(&self.plan, dest, &operand, fill, exec),
444            KernelDType::C32 => {
445                execute_pad_uninit::<Complex32>(&self.plan, dest, &operand, fill, exec)
446            }
447            KernelDType::C64 => {
448                execute_pad_uninit::<Complex64>(&self.plan, dest, &operand, fill, exec)
449            }
450            _ => Err(StridedError::UnsupportedDType {
451                dtype: self.dtype.label(),
452            }),
453        }
454    }
455}
456
457/// Dtype-erased gather wrapper.
458///
459/// This is the erased replay boundary for indexed reads. Value buffers use the
460/// configured value dtype, while the index descriptor must use `i32` or `i64`.
461#[derive(Clone, Debug)]
462pub struct ErasedGatherPlan {
463    dtype: KernelDType,
464    index_dtype: KernelDType,
465    plan: GatherPlan,
466}
467
468/// Dtype-erased fixed-window dynamic-slice wrapper.
469#[derive(Clone, Debug)]
470pub struct ErasedDynamicSlicePlan {
471    dtype: KernelDType,
472    index_dtype: KernelDType,
473    plan: DynamicSlicePlan,
474}
475
476/// Dtype-erased dynamic-update-slice wrapper.
477#[derive(Clone, Debug)]
478pub struct ErasedDynamicUpdateSlicePlan {
479    dtype: KernelDType,
480    index_dtype: KernelDType,
481    plan: DynamicUpdateSlicePlan,
482}
483
484/// Dtype-erased additive scatter wrapper.
485#[derive(Clone, Debug)]
486pub struct ErasedScatterPlan {
487    dtype: KernelDType,
488    index_dtype: KernelDType,
489    plan: ScatterPlan,
490}
491
492impl ErasedGatherPlan {
493    /// Validate and store a gather plan for one value dtype, index dtype, and layout set.
494    #[allow(clippy::too_many_arguments)]
495    pub fn compile(
496        dtype: KernelDType,
497        index_dtype: KernelDType,
498        operand_dims: &[usize],
499        operand_strides: &[isize],
500        index_dims: &[usize],
501        index_strides: &[isize],
502        dest_dims: &[usize],
503        dest_strides: &[isize],
504        spec: GatherSpec,
505    ) -> Result<Self> {
506        check_index_dtype(index_dtype)?;
507        check_gather_value_dtype(dtype)?;
508        Ok(Self {
509            dtype,
510            index_dtype,
511            plan: GatherPlan::compile(
512                operand_dims,
513                operand_strides,
514                index_dims,
515                index_strides,
516                dest_dims,
517                dest_strides,
518                spec,
519            )?,
520        })
521    }
522
523    #[inline]
524    pub fn dtype(&self) -> KernelDType {
525        self.dtype
526    }
527
528    #[inline]
529    pub fn index_dtype(&self) -> KernelDType {
530        self.index_dtype
531    }
532
533    #[inline]
534    pub fn plan(&self) -> &GatherPlan {
535        &self.plan
536    }
537
538    /// Execute an indexed read into an erased output descriptor.
539    pub fn execute(
540        &self,
541        dest: &mut ErasedRawStridedMut<'_>,
542        operand: &ErasedRawStridedRef<'_>,
543        start_indices: &ErasedRawStridedRef<'_>,
544        exec: Exec<'_>,
545    ) -> Result<()> {
546        check_dtype(self.dtype, dest.dtype())?;
547        check_dtype(self.dtype, operand.dtype())?;
548        check_dtype(self.index_dtype, start_indices.dtype())?;
549
550        match self.dtype {
551            KernelDType::F32 => dispatch_gather_index::<f32>(
552                &self.plan,
553                self.index_dtype,
554                dest,
555                &operand,
556                &start_indices,
557                exec,
558            ),
559            KernelDType::F64 => dispatch_gather_index::<f64>(
560                &self.plan,
561                self.index_dtype,
562                dest,
563                operand,
564                start_indices,
565                exec,
566            ),
567            KernelDType::I32 => dispatch_gather_index::<i32>(
568                &self.plan,
569                self.index_dtype,
570                dest,
571                operand,
572                start_indices,
573                exec,
574            ),
575            KernelDType::I64 => dispatch_gather_index::<i64>(
576                &self.plan,
577                self.index_dtype,
578                dest,
579                operand,
580                start_indices,
581                exec,
582            ),
583            KernelDType::Bool => dispatch_gather_index::<bool>(
584                &self.plan,
585                self.index_dtype,
586                dest,
587                operand,
588                start_indices,
589                exec,
590            ),
591            KernelDType::C32 => dispatch_gather_index::<Complex32>(
592                &self.plan,
593                self.index_dtype,
594                dest,
595                operand,
596                start_indices,
597                exec,
598            ),
599            KernelDType::C64 => dispatch_gather_index::<Complex64>(
600                &self.plan,
601                self.index_dtype,
602                dest,
603                operand,
604                start_indices,
605                exec,
606            ),
607            _ => Err(StridedError::UnsupportedDType {
608                dtype: self.dtype.label(),
609            }),
610        }
611    }
612
613    /// Execute gather into a destination whose reachable slots may be
614    /// uninitialized. All validation precedes the first destination write.
615    /// On success, every reachable destination slot is fully overwritten;
616    /// unreachable holes are neither read nor initialized. Validation errors
617    /// are returned before any destination write. A panic during execution
618    /// may leave a partially initialized `MaybeUninit` destination, which is
619    /// still safely droppable; no readable value is promised for unwritten
620    /// reachable slots.
621    pub fn execute_uninit(
622        &self,
623        dest: &mut ErasedRawStridedUninitMut<'_>,
624        operand: &ErasedRawStridedPtr<'_>,
625        start_indices: &ErasedRawStridedPtr<'_>,
626        exec: Exec<'_>,
627    ) -> Result<()> {
628        check_dtype(self.dtype, dest.dtype())?;
629        check_dtype(self.dtype, operand.dtype())?;
630        check_dtype(self.index_dtype, start_indices.dtype())?;
631        validate_uninit_no_overlap(dest, operand, 0)?;
632        validate_uninit_no_overlap(dest, start_indices, 1)?;
633        // SAFETY: the owning erased entry rejected all input/output overlap before conversion.
634        let operand = &unsafe { operand.try_as_ref_after_no_overlap() }?;
635        // SAFETY: the owning erased entry rejected all input/output overlap before conversion.
636        let start_indices = &unsafe { start_indices.try_as_ref_after_no_overlap() }?;
637        match self.dtype {
638            KernelDType::F32 => execute_gather_uninit_dispatch::<f32>(
639                &self.plan,
640                self.index_dtype,
641                dest,
642                operand,
643                start_indices,
644                exec,
645            ),
646            KernelDType::F64 => execute_gather_uninit_dispatch::<f64>(
647                &self.plan,
648                self.index_dtype,
649                dest,
650                operand,
651                start_indices,
652                exec,
653            ),
654            KernelDType::I32 => execute_gather_uninit_dispatch::<i32>(
655                &self.plan,
656                self.index_dtype,
657                dest,
658                operand,
659                start_indices,
660                exec,
661            ),
662            KernelDType::I64 => execute_gather_uninit_dispatch::<i64>(
663                &self.plan,
664                self.index_dtype,
665                dest,
666                operand,
667                start_indices,
668                exec,
669            ),
670            KernelDType::Bool => execute_gather_uninit_dispatch::<bool>(
671                &self.plan,
672                self.index_dtype,
673                dest,
674                operand,
675                start_indices,
676                exec,
677            ),
678            KernelDType::C32 => execute_gather_uninit_dispatch::<Complex32>(
679                &self.plan,
680                self.index_dtype,
681                dest,
682                operand,
683                start_indices,
684                exec,
685            ),
686            KernelDType::C64 => execute_gather_uninit_dispatch::<Complex64>(
687                &self.plan,
688                self.index_dtype,
689                dest,
690                operand,
691                start_indices,
692                exec,
693            ),
694            _ => Err(StridedError::UnsupportedDType {
695                dtype: self.dtype.label(),
696            }),
697        }
698    }
699}
700
701impl ErasedDynamicSlicePlan {
702    /// Validate and store a dynamic-slice plan for one value dtype, index dtype, and layout set.
703    #[allow(clippy::too_many_arguments)]
704    pub fn compile(
705        dtype: KernelDType,
706        index_dtype: KernelDType,
707        operand_dims: &[usize],
708        operand_strides: &[isize],
709        start_dims: &[usize],
710        start_strides: &[isize],
711        dest_dims: &[usize],
712        dest_strides: &[isize],
713        slice_sizes: &[usize],
714    ) -> Result<Self> {
715        check_index_dtype(index_dtype)?;
716        check_gather_value_dtype(dtype)?;
717        Ok(Self {
718            dtype,
719            index_dtype,
720            plan: DynamicSlicePlan::compile(
721                operand_dims,
722                operand_strides,
723                start_dims,
724                start_strides,
725                dest_dims,
726                dest_strides,
727                slice_sizes,
728            )?,
729        })
730    }
731
732    #[inline]
733    pub fn dtype(&self) -> KernelDType {
734        self.dtype
735    }
736
737    #[inline]
738    pub fn index_dtype(&self) -> KernelDType {
739        self.index_dtype
740    }
741
742    #[inline]
743    pub fn plan(&self) -> &DynamicSlicePlan {
744        &self.plan
745    }
746
747    /// Execute a fixed-window dynamic slice into an erased output descriptor.
748    pub fn execute(
749        &self,
750        dest: &mut ErasedRawStridedMut<'_>,
751        operand: &ErasedRawStridedRef<'_>,
752        starts: &ErasedRawStridedRef<'_>,
753        exec: Exec<'_>,
754    ) -> Result<()> {
755        check_dtype(self.dtype, dest.dtype())?;
756        check_dtype(self.dtype, operand.dtype())?;
757        check_dtype(self.index_dtype, starts.dtype())?;
758
759        match self.dtype {
760            KernelDType::F32 => dispatch_dynamic_slice_index::<f32>(
761                &self.plan,
762                self.index_dtype,
763                dest,
764                &operand,
765                &starts,
766                exec,
767            ),
768            KernelDType::F64 => dispatch_dynamic_slice_index::<f64>(
769                &self.plan,
770                self.index_dtype,
771                dest,
772                operand,
773                starts,
774                exec,
775            ),
776            KernelDType::I32 => dispatch_dynamic_slice_index::<i32>(
777                &self.plan,
778                self.index_dtype,
779                dest,
780                operand,
781                starts,
782                exec,
783            ),
784            KernelDType::I64 => dispatch_dynamic_slice_index::<i64>(
785                &self.plan,
786                self.index_dtype,
787                dest,
788                operand,
789                starts,
790                exec,
791            ),
792            KernelDType::Bool => dispatch_dynamic_slice_index::<bool>(
793                &self.plan,
794                self.index_dtype,
795                dest,
796                operand,
797                starts,
798                exec,
799            ),
800            KernelDType::C32 => dispatch_dynamic_slice_index::<Complex32>(
801                &self.plan,
802                self.index_dtype,
803                dest,
804                operand,
805                starts,
806                exec,
807            ),
808            KernelDType::C64 => dispatch_dynamic_slice_index::<Complex64>(
809                &self.plan,
810                self.index_dtype,
811                dest,
812                operand,
813                starts,
814                exec,
815            ),
816            _ => Err(StridedError::UnsupportedDType {
817                dtype: self.dtype.label(),
818            }),
819        }
820    }
821
822    /// Execute dynamic slice into a destination whose reachable slots may be
823    /// uninitialized.
824    /// On success, every reachable destination slot is fully overwritten;
825    /// unreachable holes are neither read nor initialized. Validation errors
826    /// are returned before any destination write. A panic during execution may
827    /// leave reachable slots partially initialized, but the `MaybeUninit`
828    /// destination remains safely droppable.
829    pub fn execute_uninit(
830        &self,
831        dest: &mut ErasedRawStridedUninitMut<'_>,
832        operand: &ErasedRawStridedPtr<'_>,
833        starts: &ErasedRawStridedPtr<'_>,
834        exec: Exec<'_>,
835    ) -> Result<()> {
836        check_dtype(self.dtype, dest.dtype())?;
837        check_dtype(self.dtype, operand.dtype())?;
838        check_dtype(self.index_dtype, starts.dtype())?;
839        validate_uninit_no_overlap(dest, operand, 0)?;
840        validate_uninit_no_overlap(dest, starts, 1)?;
841        // SAFETY: the owning erased entry rejected all input/output overlap before conversion.
842        let operand = &unsafe { operand.try_as_ref_after_no_overlap() }?;
843        // SAFETY: the owning erased entry rejected all input/output overlap before conversion.
844        let starts = &unsafe { starts.try_as_ref_after_no_overlap() }?;
845        match self.dtype {
846            KernelDType::F32 => execute_dynamic_slice_uninit_dispatch::<f32>(
847                &self.plan,
848                self.index_dtype,
849                dest,
850                operand,
851                starts,
852                exec,
853            ),
854            KernelDType::F64 => execute_dynamic_slice_uninit_dispatch::<f64>(
855                &self.plan,
856                self.index_dtype,
857                dest,
858                operand,
859                starts,
860                exec,
861            ),
862            KernelDType::I32 => execute_dynamic_slice_uninit_dispatch::<i32>(
863                &self.plan,
864                self.index_dtype,
865                dest,
866                operand,
867                starts,
868                exec,
869            ),
870            KernelDType::I64 => execute_dynamic_slice_uninit_dispatch::<i64>(
871                &self.plan,
872                self.index_dtype,
873                dest,
874                operand,
875                starts,
876                exec,
877            ),
878            KernelDType::Bool => execute_dynamic_slice_uninit_dispatch::<bool>(
879                &self.plan,
880                self.index_dtype,
881                dest,
882                operand,
883                starts,
884                exec,
885            ),
886            KernelDType::C32 => execute_dynamic_slice_uninit_dispatch::<Complex32>(
887                &self.plan,
888                self.index_dtype,
889                dest,
890                operand,
891                starts,
892                exec,
893            ),
894            KernelDType::C64 => execute_dynamic_slice_uninit_dispatch::<Complex64>(
895                &self.plan,
896                self.index_dtype,
897                dest,
898                operand,
899                starts,
900                exec,
901            ),
902            _ => Err(StridedError::UnsupportedDType {
903                dtype: self.dtype.label(),
904            }),
905        }
906    }
907}
908
909impl ErasedDynamicUpdateSlicePlan {
910    /// Validate and store a dynamic-update-slice plan for one value dtype, index dtype, and layout set.
911    #[allow(clippy::too_many_arguments)]
912    pub fn compile(
913        dtype: KernelDType,
914        index_dtype: KernelDType,
915        operand_dims: &[usize],
916        operand_strides: &[isize],
917        start_dims: &[usize],
918        start_strides: &[isize],
919        update_dims: &[usize],
920        update_strides: &[isize],
921        dest_dims: &[usize],
922        dest_strides: &[isize],
923    ) -> Result<Self> {
924        check_index_dtype(index_dtype)?;
925        check_gather_value_dtype(dtype)?;
926        Ok(Self {
927            dtype,
928            index_dtype,
929            plan: DynamicUpdateSlicePlan::compile(
930                operand_dims,
931                operand_strides,
932                start_dims,
933                start_strides,
934                update_dims,
935                update_strides,
936                dest_dims,
937                dest_strides,
938            )?,
939        })
940    }
941
942    #[inline]
943    pub fn dtype(&self) -> KernelDType {
944        self.dtype
945    }
946
947    #[inline]
948    pub fn index_dtype(&self) -> KernelDType {
949        self.index_dtype
950    }
951
952    #[inline]
953    pub fn plan(&self) -> &DynamicUpdateSlicePlan {
954        &self.plan
955    }
956
957    /// Execute a dynamic update slice into an erased output descriptor.
958    pub fn execute(
959        &self,
960        dest: &mut ErasedRawStridedMut<'_>,
961        operand: &ErasedRawStridedRef<'_>,
962        update: &ErasedRawStridedRef<'_>,
963        starts: &ErasedRawStridedRef<'_>,
964        exec: Exec<'_>,
965    ) -> Result<()> {
966        check_dtype(self.dtype, dest.dtype())?;
967        check_dtype(self.dtype, operand.dtype())?;
968        check_dtype(self.dtype, update.dtype())?;
969        check_dtype(self.index_dtype, starts.dtype())?;
970
971        match self.dtype {
972            KernelDType::F32 => dispatch_dynamic_update_slice_index::<f32>(
973                &self.plan,
974                self.index_dtype,
975                dest,
976                &operand,
977                &update,
978                &starts,
979                exec,
980            ),
981            KernelDType::F64 => dispatch_dynamic_update_slice_index::<f64>(
982                &self.plan,
983                self.index_dtype,
984                dest,
985                operand,
986                update,
987                starts,
988                exec,
989            ),
990            KernelDType::I32 => dispatch_dynamic_update_slice_index::<i32>(
991                &self.plan,
992                self.index_dtype,
993                dest,
994                operand,
995                update,
996                starts,
997                exec,
998            ),
999            KernelDType::I64 => dispatch_dynamic_update_slice_index::<i64>(
1000                &self.plan,
1001                self.index_dtype,
1002                dest,
1003                operand,
1004                update,
1005                starts,
1006                exec,
1007            ),
1008            KernelDType::Bool => dispatch_dynamic_update_slice_index::<bool>(
1009                &self.plan,
1010                self.index_dtype,
1011                dest,
1012                operand,
1013                update,
1014                starts,
1015                exec,
1016            ),
1017            KernelDType::C32 => dispatch_dynamic_update_slice_index::<Complex32>(
1018                &self.plan,
1019                self.index_dtype,
1020                dest,
1021                operand,
1022                update,
1023                starts,
1024                exec,
1025            ),
1026            KernelDType::C64 => dispatch_dynamic_update_slice_index::<Complex64>(
1027                &self.plan,
1028                self.index_dtype,
1029                dest,
1030                operand,
1031                update,
1032                starts,
1033                exec,
1034            ),
1035            _ => Err(StridedError::UnsupportedDType {
1036                dtype: self.dtype.label(),
1037            }),
1038        }
1039    }
1040
1041    /// On success, the copy phase initializes every reachable destination
1042    /// slot before the read-modify-write phase. Unreachable holes are neither
1043    /// read nor initialized. Validation errors before the copy leave the
1044    /// destination untouched; an error or panic after the copy may leave a
1045    /// mixture of old and new reachable values, all initialized and safely
1046    /// droppable.
1047    pub fn execute_uninit(
1048        &self,
1049        dest: &mut ErasedRawStridedUninitMut<'_>,
1050        operand: &ErasedRawStridedPtr<'_>,
1051        update: &ErasedRawStridedPtr<'_>,
1052        starts: &ErasedRawStridedPtr<'_>,
1053        exec: Exec<'_>,
1054    ) -> Result<()> {
1055        check_dtype(self.dtype, dest.dtype())?;
1056        check_dtype(self.dtype, operand.dtype())?;
1057        check_dtype(self.dtype, update.dtype())?;
1058        check_dtype(self.index_dtype, starts.dtype())?;
1059        validate_uninit_no_overlap(dest, operand, 0)?;
1060        validate_uninit_no_overlap(dest, update, 1)?;
1061        validate_uninit_no_overlap(dest, starts, 2)?;
1062        // SAFETY: the owning erased entry rejected all input/output overlap before conversion.
1063        let operand = &unsafe { operand.try_as_ref_after_no_overlap() }?;
1064        // SAFETY: the owning erased entry rejected all input/output overlap before conversion.
1065        let update = &unsafe { update.try_as_ref_after_no_overlap() }?;
1066        // SAFETY: the owning erased entry rejected all input/output overlap before conversion.
1067        let starts = &unsafe { starts.try_as_ref_after_no_overlap() }?;
1068        match self.dtype {
1069            KernelDType::F32 => execute_dynamic_update_uninit_dispatch::<f32>(
1070                &self.plan,
1071                self.index_dtype,
1072                dest,
1073                operand,
1074                update,
1075                starts,
1076                exec,
1077            ),
1078            KernelDType::F64 => execute_dynamic_update_uninit_dispatch::<f64>(
1079                &self.plan,
1080                self.index_dtype,
1081                dest,
1082                operand,
1083                update,
1084                starts,
1085                exec,
1086            ),
1087            KernelDType::I32 => execute_dynamic_update_uninit_dispatch::<i32>(
1088                &self.plan,
1089                self.index_dtype,
1090                dest,
1091                operand,
1092                update,
1093                starts,
1094                exec,
1095            ),
1096            KernelDType::I64 => execute_dynamic_update_uninit_dispatch::<i64>(
1097                &self.plan,
1098                self.index_dtype,
1099                dest,
1100                operand,
1101                update,
1102                starts,
1103                exec,
1104            ),
1105            KernelDType::Bool => execute_dynamic_update_uninit_dispatch::<bool>(
1106                &self.plan,
1107                self.index_dtype,
1108                dest,
1109                operand,
1110                update,
1111                starts,
1112                exec,
1113            ),
1114            KernelDType::C32 => execute_dynamic_update_uninit_dispatch::<Complex32>(
1115                &self.plan,
1116                self.index_dtype,
1117                dest,
1118                operand,
1119                update,
1120                starts,
1121                exec,
1122            ),
1123            KernelDType::C64 => execute_dynamic_update_uninit_dispatch::<Complex64>(
1124                &self.plan,
1125                self.index_dtype,
1126                dest,
1127                operand,
1128                update,
1129                starts,
1130                exec,
1131            ),
1132            _ => Err(StridedError::UnsupportedDType {
1133                dtype: self.dtype.label(),
1134            }),
1135        }
1136    }
1137}
1138
1139impl ErasedScatterPlan {
1140    /// Validate and store an additive scatter plan for one value dtype, index dtype, and layout set.
1141    #[allow(clippy::too_many_arguments)]
1142    pub fn compile(
1143        dtype: KernelDType,
1144        index_dtype: KernelDType,
1145        operand_dims: &[usize],
1146        operand_strides: &[isize],
1147        index_dims: &[usize],
1148        index_strides: &[isize],
1149        update_dims: &[usize],
1150        update_strides: &[isize],
1151        dest_dims: &[usize],
1152        dest_strides: &[isize],
1153        spec: ScatterSpec,
1154    ) -> Result<Self> {
1155        check_index_dtype(index_dtype)?;
1156        check_scatter_value_dtype(dtype)?;
1157        Ok(Self {
1158            dtype,
1159            index_dtype,
1160            plan: ScatterPlan::compile(
1161                operand_dims,
1162                operand_strides,
1163                index_dims,
1164                index_strides,
1165                update_dims,
1166                update_strides,
1167                dest_dims,
1168                dest_strides,
1169                spec,
1170            )?,
1171        })
1172    }
1173
1174    #[inline]
1175    pub fn dtype(&self) -> KernelDType {
1176        self.dtype
1177    }
1178
1179    #[inline]
1180    pub fn index_dtype(&self) -> KernelDType {
1181        self.index_dtype
1182    }
1183
1184    #[inline]
1185    pub fn plan(&self) -> &ScatterPlan {
1186        &self.plan
1187    }
1188
1189    /// Execute additive scatter into an erased output descriptor.
1190    pub fn execute(
1191        &self,
1192        dest: &mut ErasedRawStridedMut<'_>,
1193        operand: &ErasedRawStridedRef<'_>,
1194        scatter_indices: &ErasedRawStridedRef<'_>,
1195        updates: &ErasedRawStridedRef<'_>,
1196        exec: Exec<'_>,
1197    ) -> Result<()> {
1198        check_dtype(self.dtype, dest.dtype())?;
1199        check_dtype(self.dtype, operand.dtype())?;
1200        check_dtype(self.dtype, updates.dtype())?;
1201        check_dtype(self.index_dtype, scatter_indices.dtype())?;
1202
1203        match self.dtype {
1204            KernelDType::F32 => dispatch_scatter_index::<f32>(
1205                &self.plan,
1206                self.index_dtype,
1207                dest,
1208                &operand,
1209                &scatter_indices,
1210                &updates,
1211                exec,
1212            ),
1213            KernelDType::F64 => dispatch_scatter_index::<f64>(
1214                &self.plan,
1215                self.index_dtype,
1216                dest,
1217                operand,
1218                scatter_indices,
1219                updates,
1220                exec,
1221            ),
1222            KernelDType::I32 => dispatch_scatter_index::<i32>(
1223                &self.plan,
1224                self.index_dtype,
1225                dest,
1226                operand,
1227                scatter_indices,
1228                updates,
1229                exec,
1230            ),
1231            KernelDType::I64 => dispatch_scatter_index::<i64>(
1232                &self.plan,
1233                self.index_dtype,
1234                dest,
1235                operand,
1236                scatter_indices,
1237                updates,
1238                exec,
1239            ),
1240            KernelDType::C32 => dispatch_scatter_index::<Complex32>(
1241                &self.plan,
1242                self.index_dtype,
1243                dest,
1244                operand,
1245                scatter_indices,
1246                updates,
1247                exec,
1248            ),
1249            KernelDType::C64 => dispatch_scatter_index::<Complex64>(
1250                &self.plan,
1251                self.index_dtype,
1252                dest,
1253                operand,
1254                scatter_indices,
1255                updates,
1256                exec,
1257            ),
1258            _ => Err(StridedError::UnsupportedDType {
1259                dtype: self.dtype.label(),
1260            }),
1261        }
1262    }
1263
1264    /// On success, the copy phase initializes every reachable destination
1265    /// slot before the read-modify-write phase. Unreachable holes are neither
1266    /// read nor initialized. Validation errors before the copy leave the
1267    /// destination untouched; an error or panic after the copy may leave a
1268    /// mixture of old and new reachable values, all initialized and safely
1269    /// droppable.
1270    pub fn execute_uninit(
1271        &self,
1272        dest: &mut ErasedRawStridedUninitMut<'_>,
1273        operand: &ErasedRawStridedPtr<'_>,
1274        scatter_indices: &ErasedRawStridedPtr<'_>,
1275        updates: &ErasedRawStridedPtr<'_>,
1276        exec: Exec<'_>,
1277    ) -> Result<()> {
1278        check_dtype(self.dtype, dest.dtype())?;
1279        check_dtype(self.dtype, operand.dtype())?;
1280        check_dtype(self.dtype, updates.dtype())?;
1281        check_dtype(self.index_dtype, scatter_indices.dtype())?;
1282        validate_uninit_no_overlap(dest, operand, 0)?;
1283        validate_uninit_no_overlap(dest, scatter_indices, 1)?;
1284        validate_uninit_no_overlap(dest, updates, 2)?;
1285        // SAFETY: the owning erased entry rejected all input/output overlap before conversion.
1286        let operand = &unsafe { operand.try_as_ref_after_no_overlap() }?;
1287        // SAFETY: the owning erased entry rejected all input/output overlap before conversion.
1288        let scatter_indices = &unsafe { scatter_indices.try_as_ref_after_no_overlap() }?;
1289        // SAFETY: the owning erased entry rejected all input/output overlap before conversion.
1290        let updates = &unsafe { updates.try_as_ref_after_no_overlap() }?;
1291        match self.dtype {
1292            KernelDType::F32 => execute_scatter_uninit_dispatch::<f32>(
1293                &self.plan,
1294                self.index_dtype,
1295                dest,
1296                operand,
1297                scatter_indices,
1298                updates,
1299                add_values::<f32>,
1300                exec,
1301            ),
1302            KernelDType::F64 => execute_scatter_uninit_dispatch::<f64>(
1303                &self.plan,
1304                self.index_dtype,
1305                dest,
1306                operand,
1307                scatter_indices,
1308                updates,
1309                add_values::<f64>,
1310                exec,
1311            ),
1312            KernelDType::I32 => execute_scatter_uninit_dispatch::<i32>(
1313                &self.plan,
1314                self.index_dtype,
1315                dest,
1316                operand,
1317                scatter_indices,
1318                updates,
1319                i32::wrapping_add,
1320                exec,
1321            ),
1322            KernelDType::I64 => execute_scatter_uninit_dispatch::<i64>(
1323                &self.plan,
1324                self.index_dtype,
1325                dest,
1326                operand,
1327                scatter_indices,
1328                updates,
1329                i64::wrapping_add,
1330                exec,
1331            ),
1332            KernelDType::C32 => execute_scatter_uninit_dispatch::<Complex32>(
1333                &self.plan,
1334                self.index_dtype,
1335                dest,
1336                operand,
1337                scatter_indices,
1338                updates,
1339                add_values::<Complex32>,
1340                exec,
1341            ),
1342            KernelDType::C64 => execute_scatter_uninit_dispatch::<Complex64>(
1343                &self.plan,
1344                self.index_dtype,
1345                dest,
1346                operand,
1347                scatter_indices,
1348                updates,
1349                add_values::<Complex64>,
1350                exec,
1351            ),
1352            _ => Err(StridedError::UnsupportedDType {
1353                dtype: self.dtype.label(),
1354            }),
1355        }
1356    }
1357}
1358
1359fn add_values<T: Add<Output = T>>(lhs: T, rhs: T) -> T {
1360    lhs + rhs
1361}
1362
1363fn execute_one_shot_map<T: OneShotScalar>(
1364    op: ErasedMapOp,
1365    dest: &mut ErasedRawStridedMut<'_>,
1366    input: &ErasedRawStridedRef<'_>,
1367    exec: Exec<'_>,
1368) -> Result<()> {
1369    if !T::supports_map(op) {
1370        return Err(StridedError::UnsupportedOp {
1371            op: op.label(),
1372            dtype: T::one_shot_dtype_label(),
1373        });
1374    }
1375    let validated = strided_basic::execution::validate_destination_layout_without_alloc(
1376        dest.dims(),
1377        dest.strides(),
1378    )?;
1379    strided_basic::execution::ensure_same_shape(dest.dims(), input.dims())?;
1380    if dest.dims().contains(&0) {
1381        return Ok(());
1382    }
1383
1384    let dest_dims = dest.dims();
1385    let dest_strides = dest.strides();
1386    let dest_offset = dest.offset();
1387    let dest_data = dest.data_as_mut::<T>()?;
1388    let mut dest =
1389        unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) };
1390    let input = erased_raw_ref::<T>(input)?;
1391
1392    // The op is matched once here so each arm replays a loop whose closure
1393    // names a constant op. After inlining, `T::map` folds to one arithmetic
1394    // expression and the contiguous and threaded leaf loops can vectorize.
1395    macro_rules! replay {
1396        ($op:ident) => {
1397            // SAFETY: matching shapes and this destination layout were validated before specialization/replay.
1398            unsafe {
1399                strided_basic::execution::map_raw_into_validated::<T, T, Identity>(
1400                    &mut dest,
1401                    &input,
1402                    |value| T::map(ErasedMapOp::$op, value),
1403                    validated, exec,
1404)
1405            }
1406        };
1407    }
1408    match op {
1409        ErasedMapOp::Negate => replay!(Negate),
1410        ErasedMapOp::Conj => replay!(Conj),
1411        ErasedMapOp::Abs => replay!(Abs),
1412        ErasedMapOp::Sign => replay!(Sign),
1413    }
1414}
1415
1416fn execute_one_shot_map_with<D, A>(
1417    dest: &mut ErasedRawStridedMut<'_>,
1418    input: &ErasedRawStridedRef<'_>,
1419    map: impl Fn(A) -> D + Sync,
1420    exec: Exec<'_>,
1421) -> Result<()>
1422where
1423    D: Copy + Send + Sync + KernelStorageElement,
1424    A: Copy + Send + Sync + KernelStorageElement,
1425{
1426    let validated = strided_basic::execution::validate_destination_layout_without_alloc(
1427        dest.dims(),
1428        dest.strides(),
1429    )?;
1430    strided_basic::execution::ensure_same_shape(dest.dims(), input.dims())?;
1431    if dest.dims().contains(&0) {
1432        return Ok(());
1433    }
1434    let dest_dims = dest.dims();
1435    let dest_strides = dest.strides();
1436    let dest_offset = dest.offset();
1437    let dest_data = dest.data_as_mut::<D>()?;
1438    let mut dest =
1439        unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) };
1440    let input = erased_raw_ref::<A>(input)?;
1441    // SAFETY: matching shapes and this destination layout were validated before specialization/replay.
1442    unsafe {
1443        strided_basic::execution::map_raw_into_validated::<D, A, Identity>(
1444            &mut dest, &input, map, validated, exec,
1445        )
1446    }
1447}
1448
1449fn execute_one_shot_zip<T: OneShotScalar>(
1450    op: ErasedZipOp,
1451    dest: &mut ErasedRawStridedMut<'_>,
1452    lhs: &ErasedRawStridedRef<'_>,
1453    rhs: &ErasedRawStridedRef<'_>,
1454    exec: Exec<'_>,
1455) -> Result<()> {
1456    if !T::supports_zip(op) {
1457        return Err(StridedError::UnsupportedOp {
1458            op: op.label(),
1459            dtype: T::one_shot_dtype_label(),
1460        });
1461    }
1462    let validated = strided_basic::execution::validate_destination_layout_without_alloc(
1463        dest.dims(),
1464        dest.strides(),
1465    )?;
1466    strided_basic::execution::ensure_same_shape(dest.dims(), lhs.dims())?;
1467    strided_basic::execution::ensure_same_shape(dest.dims(), rhs.dims())?;
1468    if dest.dims().contains(&0) {
1469        return Ok(());
1470    }
1471
1472    let lhs = erased_raw_ref::<T>(lhs)?;
1473    let rhs = erased_raw_ref::<T>(rhs)?;
1474    if matches!(op, ErasedZipOp::Divide | ErasedZipOp::Remainder)
1475        && T::INTEGER
1476        && raw_any(&rhs, T::is_zero)?
1477    {
1478        return Err(StridedError::IntegerDivisionByZero { op: op.label() });
1479    }
1480    let dest_dims = dest.dims();
1481    let dest_strides = dest.strides();
1482    let dest_offset = dest.offset();
1483    let dest_data = dest.data_as_mut::<T>()?;
1484    let mut dest =
1485        unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) };
1486    // The op is matched once here so each arm replays a loop whose closure
1487    // names a constant op; see `execute_one_shot_map`.
1488    macro_rules! replay {
1489        ($op:ident) => {
1490            // SAFETY: matching shapes and this destination layout were validated before specialization/replay.
1491            unsafe {
1492                strided_basic::execution::zip_map2_raw_into_validated::<
1493                    T,
1494                    T,
1495                    T,
1496                    Identity,
1497                    Identity,
1498                >(
1499                    &mut dest,
1500                    &lhs,
1501                    &rhs,
1502                    |lhs, rhs| T::zip(ErasedZipOp::$op, lhs, rhs),
1503                    validated, exec,
1504)
1505            }
1506        };
1507    }
1508    match op {
1509        ErasedZipOp::Add => replay!(Add),
1510        ErasedZipOp::Subtract => replay!(Subtract),
1511        ErasedZipOp::Multiply => replay!(Multiply),
1512        ErasedZipOp::Divide => replay!(Divide),
1513        ErasedZipOp::Remainder => replay!(Remainder),
1514        ErasedZipOp::Maximum => replay!(Maximum),
1515        ErasedZipOp::Minimum => replay!(Minimum),
1516    }
1517}
1518
1519fn validate_no_overlap(
1520    dest: &ErasedRawStridedMut<'_>,
1521    input: &ErasedRawStridedPtr<'_>,
1522    input_index: usize,
1523) -> Result<()> {
1524    if input.overlaps_mut(dest)? {
1525        Err(StridedError::OverlappingInputOutput { input: input_index })
1526    } else {
1527        Ok(())
1528    }
1529}
1530
1531trait OneShotScalar: Copy + Send + Sync + KernelStorageElement + 'static {
1532    const INTEGER: bool = false;
1533    fn is_zero(_value: Self) -> bool {
1534        false
1535    }
1536    fn one_shot_dtype_label() -> &'static str;
1537    fn supports_map(op: ErasedMapOp) -> bool;
1538    fn supports_zip(op: ErasedZipOp) -> bool;
1539    fn map(op: ErasedMapOp, value: Self) -> Self;
1540    fn zip(op: ErasedZipOp, lhs: Self, rhs: Self) -> Self;
1541    /// `minimum(hi, maximum(lo, x))`, the composition clamp promises.
1542    #[inline(always)]
1543    fn clamp(x: Self, lo: Self, hi: Self) -> Self {
1544        let raised = Self::zip(ErasedZipOp::Maximum, lo, x);
1545        Self::zip(ErasedZipOp::Minimum, hi, raised)
1546    }
1547}
1548
1549macro_rules! impl_real_one_shot_scalar {
1550    ($ty:ty, $label:literal) => {
1551        impl OneShotScalar for $ty {
1552            fn one_shot_dtype_label() -> &'static str {
1553                $label
1554            }
1555
1556            fn supports_map(_op: ErasedMapOp) -> bool {
1557                true
1558            }
1559
1560            fn supports_zip(_op: ErasedZipOp) -> bool {
1561                true
1562            }
1563
1564            #[inline(always)]
1565            fn map(op: ErasedMapOp, value: Self) -> Self {
1566                match op {
1567                    ErasedMapOp::Negate => -value,
1568                    ErasedMapOp::Conj => value,
1569                    ErasedMapOp::Abs => value.abs(),
1570                    ErasedMapOp::Sign => {
1571                        if value == 0.0 {
1572                            0.0
1573                        } else {
1574                            value.signum()
1575                        }
1576                    }
1577                }
1578            }
1579
1580            #[inline(always)]
1581            fn zip(op: ErasedZipOp, lhs: Self, rhs: Self) -> Self {
1582                match op {
1583                    ErasedZipOp::Add => lhs + rhs,
1584                    ErasedZipOp::Subtract => lhs - rhs,
1585                    ErasedZipOp::Multiply => lhs * rhs,
1586                    ErasedZipOp::Divide => lhs / rhs,
1587                    ErasedZipOp::Remainder => lhs % rhs,
1588                    // Shared NaN-propagating selects, so the fused path
1589                    // cannot drift from the ordinary binary kernel.
1590                    ErasedZipOp::Maximum => strided_basic::nan_propagating_maximum(lhs, rhs),
1591                    ErasedZipOp::Minimum => strided_basic::nan_propagating_minimum(lhs, rhs),
1592                }
1593            }
1594
1595            // Same result as the default composition: any NaN operand gives
1596            // the canonical NaN and ties return the bound. Testing NaN once
1597            // over the three operands, instead of once per step, keeps the
1598            // loop at memory bandwidth.
1599            #[inline(always)]
1600            fn clamp(x: Self, lo: Self, hi: Self) -> Self {
1601                let raised = if lo >= x { lo } else { x };
1602                let lowered = if hi <= raised { hi } else { raised };
1603                if x.is_nan() | lo.is_nan() | hi.is_nan() {
1604                    <$ty>::NAN
1605                } else {
1606                    lowered
1607                }
1608            }
1609        }
1610    };
1611}
1612
1613macro_rules! impl_integer_one_shot_scalar {
1614    ($ty:ty, $label:literal) => {
1615        impl OneShotScalar for $ty {
1616            const INTEGER: bool = true;
1617
1618            fn is_zero(value: Self) -> bool {
1619                value == 0
1620            }
1621            fn one_shot_dtype_label() -> &'static str {
1622                $label
1623            }
1624
1625            fn supports_map(_op: ErasedMapOp) -> bool {
1626                true
1627            }
1628
1629            fn supports_zip(_op: ErasedZipOp) -> bool {
1630                true
1631            }
1632
1633            #[inline(always)]
1634            fn map(op: ErasedMapOp, value: Self) -> Self {
1635                match op {
1636                    ErasedMapOp::Negate => value.wrapping_neg(),
1637                    ErasedMapOp::Conj => value,
1638                    ErasedMapOp::Abs => value.wrapping_abs(),
1639                    ErasedMapOp::Sign => value.signum(),
1640                }
1641            }
1642
1643            #[inline(always)]
1644            fn zip(op: ErasedZipOp, lhs: Self, rhs: Self) -> Self {
1645                match op {
1646                    ErasedZipOp::Add => lhs.wrapping_add(rhs),
1647                    ErasedZipOp::Subtract => lhs.wrapping_sub(rhs),
1648                    ErasedZipOp::Multiply => lhs.wrapping_mul(rhs),
1649                    ErasedZipOp::Maximum => lhs.max(rhs),
1650                    ErasedZipOp::Minimum => lhs.min(rhs),
1651                    ErasedZipOp::Divide => lhs.wrapping_div(rhs),
1652                    ErasedZipOp::Remainder => lhs.wrapping_rem(rhs),
1653                }
1654            }
1655        }
1656    };
1657}
1658
1659macro_rules! impl_complex_one_shot_scalar {
1660    ($ty:ty, $label:literal, $div:path) => {
1661        impl OneShotScalar for $ty {
1662            fn one_shot_dtype_label() -> &'static str {
1663                $label
1664            }
1665
1666            fn supports_map(_op: ErasedMapOp) -> bool {
1667                true
1668            }
1669
1670            fn supports_zip(op: ErasedZipOp) -> bool {
1671                !matches!(
1672                    op,
1673                    ErasedZipOp::Remainder | ErasedZipOp::Maximum | ErasedZipOp::Minimum
1674                )
1675            }
1676
1677            #[inline(always)]
1678            fn map(op: ErasedMapOp, value: Self) -> Self {
1679                match op {
1680                    ErasedMapOp::Negate => -value,
1681                    ErasedMapOp::Conj => value.conj(),
1682                    ErasedMapOp::Abs => Self::new(value.norm(), 0.0),
1683                    ErasedMapOp::Sign => {
1684                        if value.re == 0.0 && value.im == 0.0 {
1685                            Self::new(0.0, 0.0)
1686                        } else {
1687                            // Divide the components by the real modulus: complex
1688                            // division squares the divisor's components, which
1689                            // underflows for tiny magnitudes such as 1e-200 in
1690                            // f64 and yields NaN instead of the unit phase.
1691                            let norm = value.norm();
1692                            Self::new(value.re / norm, value.im / norm)
1693                        }
1694                    }
1695                }
1696            }
1697
1698            #[inline(always)]
1699            fn zip(op: ErasedZipOp, lhs: Self, rhs: Self) -> Self {
1700                match op {
1701                    ErasedZipOp::Add => lhs + rhs,
1702                    ErasedZipOp::Subtract => lhs - rhs,
1703                    ErasedZipOp::Multiply => lhs * rhs,
1704                    ErasedZipOp::Divide => $div(lhs, rhs),
1705                    ErasedZipOp::Remainder | ErasedZipOp::Maximum | ErasedZipOp::Minimum => {
1706                        unreachable!("unsupported complex one-shot op")
1707                    }
1708                }
1709            }
1710        }
1711    };
1712}
1713
1714impl_real_one_shot_scalar!(f32, "f32");
1715
1716impl_real_one_shot_scalar!(f64, "f64");
1717
1718impl_integer_one_shot_scalar!(i32, "i32");
1719
1720impl_integer_one_shot_scalar!(i64, "i64");
1721
1722impl_complex_one_shot_scalar!(Complex32, "c32", strided_basic::robust_complex_divide_f32);
1723
1724impl_complex_one_shot_scalar!(Complex64, "c64", strided_basic::robust_complex_divide_f64);
1725
1726impl OneShotScalar for bool {
1727    fn one_shot_dtype_label() -> &'static str {
1728        "bool"
1729    }
1730
1731    fn supports_map(op: ErasedMapOp) -> bool {
1732        matches!(op, ErasedMapOp::Conj)
1733    }
1734
1735    fn supports_zip(_op: ErasedZipOp) -> bool {
1736        false
1737    }
1738
1739    fn map(op: ErasedMapOp, value: Self) -> Self {
1740        match op {
1741            ErasedMapOp::Conj => value,
1742            _ => unreachable!("unsupported bool one-shot op"),
1743        }
1744    }
1745
1746    fn zip(_op: ErasedZipOp, _lhs: Self, _rhs: Self) -> Self {
1747        unreachable!("unsupported bool one-shot op")
1748    }
1749}
1750
1751fn erased_raw_ref<'a, T: KernelStorageElement>(
1752    src: &'a ErasedRawStridedRef<'a>,
1753) -> Result<RawStridedRef<'a, T>> {
1754    let data = src.data_as::<T>()?;
1755    Ok(unsafe { RawStridedRef::new_unchecked(data, src.dims(), src.strides(), src.offset()) })
1756}
1757
1758fn map_output_dtype(dtype: KernelDType, op: ErasedMapOp) -> Result<KernelDType> {
1759    match (dtype, op) {
1760        (KernelDType::C32, ErasedMapOp::Abs) => Ok(KernelDType::F32),
1761        (KernelDType::C64, ErasedMapOp::Abs) => Ok(KernelDType::F64),
1762        (KernelDType::Bool, ErasedMapOp::Conj) => Ok(KernelDType::Bool),
1763        (KernelDType::Bool, _) => Err(StridedError::UnsupportedOp {
1764            op: op.label(),
1765            dtype: dtype.label(),
1766        }),
1767        _ => Ok(dtype),
1768    }
1769}
1770
1771fn raw_any<T: Copy>(
1772    src: &RawStridedRef<'_, T>,
1773    predicate: impl Fn(T) -> bool + Copy,
1774) -> Result<bool> {
1775    let total = src
1776        .dims()
1777        .iter()
1778        .try_fold(1usize, |total, &dim| total.checked_mul(dim))
1779        .ok_or(StridedError::OffsetOverflow)?;
1780    if total == 0 {
1781        return Ok(false);
1782    }
1783
1784    let rank = src.dims().len();
1785    if rank <= RAW_FUSED_RANK_LIMIT {
1786        let mut coordinates = [0usize; RAW_FUSED_RANK_LIMIT];
1787        let mut resets = [0isize; RAW_FUSED_RANK_LIMIT];
1788        raw_any_odometer(
1789            src,
1790            predicate,
1791            &mut coordinates[..rank],
1792            &mut resets[..rank],
1793        )
1794    } else {
1795        // Ranks above the fused limit keep the same incremental-offset
1796        // odometer; only the cursor storage moves to the heap.
1797        let mut coordinates = vec![0usize; rank];
1798        let mut resets = vec![0isize; rank];
1799        raw_any_odometer(src, predicate, &mut coordinates, &mut resets)
1800    }
1801}
1802
1803fn raw_any_odometer<T: Copy>(
1804    src: &RawStridedRef<'_, T>,
1805    predicate: impl Fn(T) -> bool + Copy,
1806    coordinates: &mut [usize],
1807    resets: &mut [isize],
1808) -> Result<bool> {
1809    let dims = src.dims();
1810    let strides = src.strides();
1811    // INVARIANT: the caller rejected empty extents, so every `dim - 1` is valid.
1812    for axis in 0..dims.len() {
1813        let last = isize::try_from(dims[axis] - 1).map_err(|_| StridedError::OffsetOverflow)?;
1814        resets[axis] = strides[axis]
1815            .checked_mul(last)
1816            .and_then(isize::checked_neg)
1817            .ok_or(StridedError::OffsetOverflow)?;
1818    }
1819
1820    let mut offset = src.offset();
1821    loop {
1822        // SAFETY: RawStridedRef construction validated every reachable offset.
1823        if predicate(unsafe { *src.data().as_ptr().offset(offset) }) {
1824            return Ok(true);
1825        }
1826
1827        let mut axis = 0;
1828        while axis < dims.len() && coordinates[axis] == dims[axis] - 1 {
1829            axis += 1;
1830        }
1831        if axis == dims.len() {
1832            return Ok(false);
1833        }
1834        for reset_axis in 0..axis {
1835            coordinates[reset_axis] = 0;
1836            offset = offset
1837                .checked_add(resets[reset_axis])
1838                .ok_or(StridedError::OffsetOverflow)?;
1839        }
1840        coordinates[axis] = coordinates[axis]
1841            .checked_add(1)
1842            .ok_or(StridedError::OffsetOverflow)?;
1843        offset = offset
1844            .checked_add(strides[axis])
1845            .ok_or(StridedError::OffsetOverflow)?;
1846    }
1847}
1848
1849fn check_index_dtype(dtype: KernelDType) -> Result<()> {
1850    match dtype {
1851        KernelDType::I32 | KernelDType::I64 => Ok(()),
1852        _ => Err(StridedError::UnsupportedDType {
1853            dtype: dtype.label(),
1854        }),
1855    }
1856}
1857
1858fn check_gather_value_dtype(dtype: KernelDType) -> Result<()> {
1859    match dtype {
1860        KernelDType::F32
1861        | KernelDType::F64
1862        | KernelDType::I32
1863        | KernelDType::I64
1864        | KernelDType::Bool
1865        | KernelDType::C32
1866        | KernelDType::C64 => Ok(()),
1867        _ => Err(StridedError::UnsupportedDType {
1868            dtype: dtype.label(),
1869        }),
1870    }
1871}
1872
1873fn check_scatter_value_dtype(dtype: KernelDType) -> Result<()> {
1874    match dtype {
1875        KernelDType::F32
1876        | KernelDType::F64
1877        | KernelDType::I32
1878        | KernelDType::I64
1879        | KernelDType::C32
1880        | KernelDType::C64 => Ok(()),
1881        _ => Err(StridedError::UnsupportedDType {
1882            dtype: dtype.label(),
1883        }),
1884    }
1885}
1886
1887fn validate_scalar_bytes(dtype: KernelDType, bytes: &[u8]) -> Result<()> {
1888    let element_size = dtype.size_of();
1889    if bytes.len() != element_size {
1890        return Err(StridedError::ByteLengthMismatch {
1891            dtype: dtype.label(),
1892            byte_len: bytes.len(),
1893            element_size,
1894        });
1895    }
1896    if dtype.requires_valid_byte_values() {
1897        if let Some(&value) = bytes.iter().find(|&&value| value > 1) {
1898            return Err(StridedError::InvalidBoolByte { value });
1899        }
1900    }
1901    Ok(())
1902}
1903
1904fn execute_slice<T>(
1905    plan: &SlicePlan,
1906    dest: &mut ErasedRawStridedMut<'_>,
1907    operand: &ErasedRawStridedRef<'_>,
1908    exec: Exec<'_>,
1909) -> Result<()>
1910where
1911    T: Copy + Send + Sync + KernelStorageElement,
1912{
1913    let operand_data = operand.data_as::<T>()?;
1914    let dest_dims = dest.dims();
1915    let dest_strides = dest.strides();
1916    let dest_offset = dest.offset();
1917    let dest_data = dest.data_as_mut::<T>()?;
1918    let operand_ref = unsafe {
1919        RawStridedRef::new_unchecked(
1920            operand_data,
1921            operand.dims(),
1922            operand.strides(),
1923            operand.offset(),
1924        )
1925    };
1926    let mut dest_ref =
1927        unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) };
1928    plan.execute(&mut dest_ref, &operand_ref, exec)
1929}
1930
1931fn execute_gather_uninit_dispatch<T>(
1932    plan: &GatherPlan,
1933    index_dtype: KernelDType,
1934    dest: &mut ErasedRawStridedUninitMut<'_>,
1935    operand: &ErasedRawStridedRef<'_>,
1936    start_indices: &ErasedRawStridedRef<'_>,
1937    exec: Exec<'_>,
1938) -> Result<()>
1939where
1940    T: Copy + Send + Sync + KernelStorageElement,
1941{
1942    match index_dtype {
1943        KernelDType::I32 => {
1944            execute_gather_uninit::<T, i32>(plan, index_dtype, dest, operand, start_indices, exec)
1945        }
1946        KernelDType::I64 => {
1947            execute_gather_uninit::<T, i64>(plan, index_dtype, dest, operand, start_indices, exec)
1948        }
1949        _ => Err(StridedError::UnsupportedDType {
1950            dtype: index_dtype.label(),
1951        }),
1952    }
1953}
1954
1955fn execute_gather_uninit<T, I>(
1956    plan: &GatherPlan,
1957    _index_dtype: KernelDType,
1958    dest: &mut ErasedRawStridedUninitMut<'_>,
1959    operand: &ErasedRawStridedRef<'_>,
1960    start_indices: &ErasedRawStridedRef<'_>,
1961    exec: Exec<'_>,
1962) -> Result<()>
1963where
1964    T: Copy + Send + Sync + KernelStorageElement,
1965    I: GatherIndex + KernelStorageElement,
1966{
1967    let operand_data = operand.data_as::<T>()?;
1968    let index_data = start_indices.data_as::<I>()?;
1969    let dest_dims = dest.dims();
1970    let dest_strides = dest.strides();
1971    let dest_offset = dest.offset();
1972    let dest_data = dest.data_as_uninit_mut::<T>()?;
1973    let operand_ref = unsafe {
1974        RawStridedRef::new_unchecked(
1975            operand_data,
1976            operand.dims(),
1977            operand.strides(),
1978            operand.offset(),
1979        )
1980    };
1981    let index_ref = unsafe {
1982        RawStridedRef::new_unchecked(
1983            index_data,
1984            start_indices.dims(),
1985            start_indices.strides(),
1986            start_indices.offset(),
1987        )
1988    };
1989    let mut dest_ref =
1990        unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) };
1991    // SAFETY: the erased entry rejected overlap before constructing these descriptors.
1992    unsafe {
1993        strided_basic::execution::gather_into_uninit(
1994            plan,
1995            &mut dest_ref,
1996            &operand_ref,
1997            &index_ref,
1998            exec,
1999        )
2000    }
2001}
2002
2003fn execute_dynamic_slice_uninit_dispatch<T>(
2004    plan: &DynamicSlicePlan,
2005    index_dtype: KernelDType,
2006    dest: &mut ErasedRawStridedUninitMut<'_>,
2007    operand: &ErasedRawStridedRef<'_>,
2008    starts: &ErasedRawStridedRef<'_>,
2009    exec: Exec<'_>,
2010) -> Result<()>
2011where
2012    T: Copy + Send + Sync + KernelStorageElement,
2013{
2014    match index_dtype {
2015        KernelDType::I32 => {
2016            execute_dynamic_slice_uninit::<T, i32>(plan, dest, operand, starts, exec)
2017        }
2018        KernelDType::I64 => {
2019            execute_dynamic_slice_uninit::<T, i64>(plan, dest, operand, starts, exec)
2020        }
2021        _ => Err(StridedError::UnsupportedDType {
2022            dtype: index_dtype.label(),
2023        }),
2024    }
2025}
2026
2027fn execute_dynamic_slice_uninit<T, I>(
2028    plan: &DynamicSlicePlan,
2029    dest: &mut ErasedRawStridedUninitMut<'_>,
2030    operand: &ErasedRawStridedRef<'_>,
2031    starts: &ErasedRawStridedRef<'_>,
2032    exec: Exec<'_>,
2033) -> Result<()>
2034where
2035    T: Copy + Send + Sync + KernelStorageElement,
2036    I: GatherIndex + KernelStorageElement,
2037{
2038    let operand_data = operand.data_as::<T>()?;
2039    let starts_data = starts.data_as::<I>()?;
2040    let dest_dims = dest.dims();
2041    let dest_strides = dest.strides();
2042    let dest_offset = dest.offset();
2043    let dest_data = dest.data_as_uninit_mut::<T>()?;
2044    let operand_ref = unsafe {
2045        RawStridedRef::new_unchecked(
2046            operand_data,
2047            operand.dims(),
2048            operand.strides(),
2049            operand.offset(),
2050        )
2051    };
2052    let starts_ref = unsafe {
2053        RawStridedRef::new_unchecked(
2054            starts_data,
2055            starts.dims(),
2056            starts.strides(),
2057            starts.offset(),
2058        )
2059    };
2060    let mut dest_ref =
2061        unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) };
2062    // SAFETY: the erased entry rejected overlap before constructing these descriptors.
2063    unsafe {
2064        strided_basic::execution::dynamic_slice_into_uninit(
2065            plan,
2066            &mut dest_ref,
2067            &operand_ref,
2068            &starts_ref,
2069            exec,
2070        )
2071    }
2072}
2073
2074fn execute_dynamic_update_uninit_dispatch<T>(
2075    plan: &DynamicUpdateSlicePlan,
2076    index_dtype: KernelDType,
2077    dest: &mut ErasedRawStridedUninitMut<'_>,
2078    operand: &ErasedRawStridedRef<'_>,
2079    update: &ErasedRawStridedRef<'_>,
2080    starts: &ErasedRawStridedRef<'_>,
2081    exec: Exec<'_>,
2082) -> Result<()>
2083where
2084    T: Copy + Send + Sync + KernelStorageElement,
2085{
2086    match index_dtype {
2087        KernelDType::I32 => {
2088            execute_dynamic_update_uninit::<T, i32>(plan, dest, operand, update, starts, exec)
2089        }
2090        KernelDType::I64 => {
2091            execute_dynamic_update_uninit::<T, i64>(plan, dest, operand, update, starts, exec)
2092        }
2093        _ => Err(StridedError::UnsupportedDType {
2094            dtype: index_dtype.label(),
2095        }),
2096    }
2097}
2098
2099fn execute_dynamic_update_uninit<T, I>(
2100    plan: &DynamicUpdateSlicePlan,
2101    dest: &mut ErasedRawStridedUninitMut<'_>,
2102    operand: &ErasedRawStridedRef<'_>,
2103    update: &ErasedRawStridedRef<'_>,
2104    starts: &ErasedRawStridedRef<'_>,
2105    exec: Exec<'_>,
2106) -> Result<()>
2107where
2108    T: Copy + Send + Sync + KernelStorageElement,
2109    I: GatherIndex + KernelStorageElement,
2110{
2111    let operand_data = operand.data_as::<T>()?;
2112    let update_data = update.data_as::<T>()?;
2113    let starts_data = starts.data_as::<I>()?;
2114    let dest_dims = dest.dims();
2115    let dest_strides = dest.strides();
2116    let dest_offset = dest.offset();
2117    let dest_data = dest.data_as_uninit_mut::<T>()?;
2118    let operand_ref = unsafe {
2119        RawStridedRef::new_unchecked(
2120            operand_data,
2121            operand.dims(),
2122            operand.strides(),
2123            operand.offset(),
2124        )
2125    };
2126    let update_ref = unsafe {
2127        RawStridedRef::new_unchecked(
2128            update_data,
2129            update.dims(),
2130            update.strides(),
2131            update.offset(),
2132        )
2133    };
2134    let starts_ref = unsafe {
2135        RawStridedRef::new_unchecked(
2136            starts_data,
2137            starts.dims(),
2138            starts.strides(),
2139            starts.offset(),
2140        )
2141    };
2142    let mut dest_ref =
2143        unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) };
2144    // SAFETY: the erased entry rejected overlap before constructing these descriptors.
2145    unsafe {
2146        strided_basic::execution::dynamic_update_into_uninit(
2147            plan,
2148            &mut dest_ref,
2149            &operand_ref,
2150            &update_ref,
2151            &starts_ref,
2152            exec,
2153        )
2154    }
2155}
2156
2157fn execute_scatter_uninit_dispatch<T>(
2158    plan: &ScatterPlan,
2159    index_dtype: KernelDType,
2160    dest: &mut ErasedRawStridedUninitMut<'_>,
2161    operand: &ErasedRawStridedRef<'_>,
2162    scatter_indices: &ErasedRawStridedRef<'_>,
2163    updates: &ErasedRawStridedRef<'_>,
2164    combine: fn(T, T) -> T,
2165    exec: Exec<'_>,
2166) -> Result<()>
2167where
2168    T: Copy + Add<Output = T> + Send + Sync + KernelStorageElement,
2169{
2170    match index_dtype {
2171        KernelDType::I32 => execute_scatter_uninit::<T, i32>(
2172            plan,
2173            dest,
2174            operand,
2175            scatter_indices,
2176            updates,
2177            combine,
2178            exec,
2179        ),
2180        KernelDType::I64 => execute_scatter_uninit::<T, i64>(
2181            plan,
2182            dest,
2183            operand,
2184            scatter_indices,
2185            updates,
2186            combine,
2187            exec,
2188        ),
2189        _ => Err(StridedError::UnsupportedDType {
2190            dtype: index_dtype.label(),
2191        }),
2192    }
2193}
2194
2195fn execute_scatter_uninit<T, I>(
2196    plan: &ScatterPlan,
2197    dest: &mut ErasedRawStridedUninitMut<'_>,
2198    operand: &ErasedRawStridedRef<'_>,
2199    scatter_indices: &ErasedRawStridedRef<'_>,
2200    updates: &ErasedRawStridedRef<'_>,
2201    combine: fn(T, T) -> T,
2202    exec: Exec<'_>,
2203) -> Result<()>
2204where
2205    T: Copy + Add<Output = T> + Send + Sync + KernelStorageElement,
2206    I: GatherIndex + KernelStorageElement,
2207{
2208    let indices = scatter_indices;
2209    let operand_data = operand.data_as::<T>()?;
2210    let index_data = indices.data_as::<I>()?;
2211    let update_data = updates.data_as::<T>()?;
2212    let dest_dims = dest.dims();
2213    let dest_strides = dest.strides();
2214    let dest_offset = dest.offset();
2215    let dest_data = dest.data_as_uninit_mut::<T>()?;
2216    let operand_ref = unsafe {
2217        RawStridedRef::new_unchecked(
2218            operand_data,
2219            operand.dims(),
2220            operand.strides(),
2221            operand.offset(),
2222        )
2223    };
2224    let index_ref = unsafe {
2225        RawStridedRef::new_unchecked(
2226            index_data,
2227            indices.dims(),
2228            indices.strides(),
2229            indices.offset(),
2230        )
2231    };
2232    let update_ref = unsafe {
2233        RawStridedRef::new_unchecked(
2234            update_data,
2235            updates.dims(),
2236            updates.strides(),
2237            updates.offset(),
2238        )
2239    };
2240    let mut dest_ref =
2241        unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) };
2242    // SAFETY: the erased entry rejected overlap before constructing these descriptors.
2243    unsafe {
2244        strided_basic::execution::scatter_into_uninit(
2245            plan,
2246            &mut dest_ref,
2247            &operand_ref,
2248            &index_ref,
2249            &update_ref,
2250            combine,
2251            exec,
2252        )
2253    }
2254}
2255
2256fn execute_slice_uninit<T>(
2257    plan: &SlicePlan,
2258    dest: &mut ErasedRawStridedUninitMut<'_>,
2259    operand: &ErasedRawStridedRef<'_>,
2260    exec: Exec<'_>,
2261) -> Result<()>
2262where
2263    T: Copy + Send + Sync + KernelStorageElement,
2264{
2265    let operand_data = operand.data_as::<T>()?;
2266    let dest_dims = dest.dims();
2267    let dest_strides = dest.strides();
2268    let dest_offset = dest.offset();
2269    let dest_data = dest.data_as_uninit_mut::<T>()?;
2270    let operand_ref = unsafe {
2271        RawStridedRef::new_unchecked(
2272            operand_data,
2273            operand.dims(),
2274            operand.strides(),
2275            operand.offset(),
2276        )
2277    };
2278    let mut dest_ref =
2279        unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) };
2280    plan.execute_uninit(&mut dest_ref, &operand_ref, exec)
2281}
2282
2283fn execute_reverse<T>(
2284    plan: &ReversePlan,
2285    dest: &mut ErasedRawStridedMut<'_>,
2286    operand: &ErasedRawStridedRef<'_>,
2287    exec: Exec<'_>,
2288) -> Result<()>
2289where
2290    T: Copy + Send + Sync + KernelStorageElement,
2291{
2292    let operand_data = operand.data_as::<T>()?;
2293    let dest_dims = dest.dims();
2294    let dest_strides = dest.strides();
2295    let dest_offset = dest.offset();
2296    let dest_data = dest.data_as_mut::<T>()?;
2297    let operand_ref = unsafe {
2298        RawStridedRef::new_unchecked(
2299            operand_data,
2300            operand.dims(),
2301            operand.strides(),
2302            operand.offset(),
2303        )
2304    };
2305    let mut dest_ref =
2306        unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) };
2307    plan.execute(&mut dest_ref, &operand_ref, exec)
2308}
2309
2310fn execute_reverse_uninit<T>(
2311    plan: &ReversePlan,
2312    dest: &mut ErasedRawStridedUninitMut<'_>,
2313    operand: &ErasedRawStridedRef<'_>,
2314    exec: Exec<'_>,
2315) -> Result<()>
2316where
2317    T: Copy + Send + Sync + KernelStorageElement,
2318{
2319    let operand_data = operand.data_as::<T>()?;
2320    let dest_dims = dest.dims();
2321    let dest_strides = dest.strides();
2322    let dest_offset = dest.offset();
2323    let dest_data = dest.data_as_uninit_mut::<T>()?;
2324    let operand_ref = unsafe {
2325        RawStridedRef::new_unchecked(
2326            operand_data,
2327            operand.dims(),
2328            operand.strides(),
2329            operand.offset(),
2330        )
2331    };
2332    let mut dest_ref =
2333        unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) };
2334    plan.execute_uninit(&mut dest_ref, &operand_ref, exec)
2335}
2336
2337fn execute_pad<T>(
2338    plan: &PadPlan,
2339    dest: &mut ErasedRawStridedMut<'_>,
2340    operand: &ErasedRawStridedRef<'_>,
2341    fill: &[u8],
2342    exec: Exec<'_>,
2343) -> Result<()>
2344where
2345    T: Copy + Send + Sync + KernelStorageElement,
2346{
2347    let fill = read_unaligned_scalar::<T>(fill);
2348    let operand_data = operand.data_as::<T>()?;
2349    let dest_dims = dest.dims();
2350    let dest_strides = dest.strides();
2351    let dest_offset = dest.offset();
2352    let dest_data = dest.data_as_mut::<T>()?;
2353    let operand_ref = unsafe {
2354        RawStridedRef::new_unchecked(
2355            operand_data,
2356            operand.dims(),
2357            operand.strides(),
2358            operand.offset(),
2359        )
2360    };
2361    let mut dest_ref =
2362        unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) };
2363    plan.execute(&mut dest_ref, &operand_ref, fill, exec)
2364}
2365
2366fn execute_pad_uninit<T>(
2367    plan: &PadPlan,
2368    dest: &mut ErasedRawStridedUninitMut<'_>,
2369    operand: &ErasedRawStridedRef<'_>,
2370    fill: &[u8],
2371    exec: Exec<'_>,
2372) -> Result<()>
2373where
2374    T: Copy + Send + Sync + KernelStorageElement,
2375{
2376    let fill = read_unaligned_scalar::<T>(fill);
2377    let operand_data = operand.data_as::<T>()?;
2378    let dest_dims = dest.dims();
2379    let dest_strides = dest.strides();
2380    let dest_offset = dest.offset();
2381    let dest_data = dest.data_as_uninit_mut::<T>()?;
2382    let operand_ref = unsafe {
2383        RawStridedRef::new_unchecked(
2384            operand_data,
2385            operand.dims(),
2386            operand.strides(),
2387            operand.offset(),
2388        )
2389    };
2390    let mut dest_ref =
2391        unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) };
2392    plan.execute_uninit(&mut dest_ref, &operand_ref, fill, exec)
2393}
2394
2395fn dispatch_gather_index<T>(
2396    plan: &GatherPlan,
2397    index_dtype: KernelDType,
2398    dest: &mut ErasedRawStridedMut<'_>,
2399    operand: &ErasedRawStridedRef<'_>,
2400    start_indices: &ErasedRawStridedRef<'_>,
2401    exec: Exec<'_>,
2402) -> Result<()>
2403where
2404    T: Copy + Send + Sync + KernelStorageElement,
2405{
2406    match index_dtype {
2407        KernelDType::I32 => execute_gather::<T, i32>(plan, dest, operand, start_indices, exec),
2408        KernelDType::I64 => execute_gather::<T, i64>(plan, dest, operand, start_indices, exec),
2409        _ => Err(StridedError::UnsupportedDType {
2410            dtype: index_dtype.label(),
2411        }),
2412    }
2413}
2414
2415fn execute_gather<T, I>(
2416    plan: &GatherPlan,
2417    dest: &mut ErasedRawStridedMut<'_>,
2418    operand: &ErasedRawStridedRef<'_>,
2419    start_indices: &ErasedRawStridedRef<'_>,
2420    exec: Exec<'_>,
2421) -> Result<()>
2422where
2423    T: Copy + Send + Sync + KernelStorageElement,
2424    I: GatherIndex + KernelStorageElement,
2425{
2426    let operand_data = operand.data_as::<T>()?;
2427    let index_data = start_indices.data_as::<I>()?;
2428    let dest_dims = dest.dims();
2429    let dest_strides = dest.strides();
2430    let dest_offset = dest.offset();
2431    let dest_data = dest.data_as_mut::<T>()?;
2432    let operand_ref = unsafe {
2433        RawStridedRef::new_unchecked(
2434            operand_data,
2435            operand.dims(),
2436            operand.strides(),
2437            operand.offset(),
2438        )
2439    };
2440    let index_ref = unsafe {
2441        RawStridedRef::new_unchecked(
2442            index_data,
2443            start_indices.dims(),
2444            start_indices.strides(),
2445            start_indices.offset(),
2446        )
2447    };
2448    let mut dest_ref =
2449        unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) };
2450    plan.execute(&mut dest_ref, &operand_ref, &index_ref, exec)
2451}
2452
2453fn dispatch_dynamic_slice_index<T>(
2454    plan: &DynamicSlicePlan,
2455    index_dtype: KernelDType,
2456    dest: &mut ErasedRawStridedMut<'_>,
2457    operand: &ErasedRawStridedRef<'_>,
2458    starts: &ErasedRawStridedRef<'_>,
2459    exec: Exec<'_>,
2460) -> Result<()>
2461where
2462    T: Copy + Send + Sync + KernelStorageElement,
2463{
2464    match index_dtype {
2465        KernelDType::I32 => execute_dynamic_slice::<T, i32>(plan, dest, operand, starts, exec),
2466        KernelDType::I64 => execute_dynamic_slice::<T, i64>(plan, dest, operand, starts, exec),
2467        _ => Err(StridedError::UnsupportedDType {
2468            dtype: index_dtype.label(),
2469        }),
2470    }
2471}
2472
2473fn execute_dynamic_slice<T, I>(
2474    plan: &DynamicSlicePlan,
2475    dest: &mut ErasedRawStridedMut<'_>,
2476    operand: &ErasedRawStridedRef<'_>,
2477    starts: &ErasedRawStridedRef<'_>,
2478    exec: Exec<'_>,
2479) -> Result<()>
2480where
2481    T: Copy + Send + Sync + KernelStorageElement,
2482    I: GatherIndex + KernelStorageElement,
2483{
2484    let operand_data = operand.data_as::<T>()?;
2485    let start_data = starts.data_as::<I>()?;
2486    let dest_dims = dest.dims();
2487    let dest_strides = dest.strides();
2488    let dest_offset = dest.offset();
2489    let dest_data = dest.data_as_mut::<T>()?;
2490    let operand_ref = unsafe {
2491        RawStridedRef::new_unchecked(
2492            operand_data,
2493            operand.dims(),
2494            operand.strides(),
2495            operand.offset(),
2496        )
2497    };
2498    let start_ref = unsafe {
2499        RawStridedRef::new_unchecked(start_data, starts.dims(), starts.strides(), starts.offset())
2500    };
2501    let mut dest_ref =
2502        unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) };
2503    plan.execute(&mut dest_ref, &operand_ref, &start_ref, exec)
2504}
2505
2506fn dispatch_dynamic_update_slice_index<T>(
2507    plan: &DynamicUpdateSlicePlan,
2508    index_dtype: KernelDType,
2509    dest: &mut ErasedRawStridedMut<'_>,
2510    operand: &ErasedRawStridedRef<'_>,
2511    update: &ErasedRawStridedRef<'_>,
2512    starts: &ErasedRawStridedRef<'_>,
2513    exec: Exec<'_>,
2514) -> Result<()>
2515where
2516    T: Copy + Send + Sync + KernelStorageElement,
2517{
2518    match index_dtype {
2519        KernelDType::I32 => {
2520            execute_dynamic_update_slice::<T, i32>(plan, dest, operand, update, starts, exec)
2521        }
2522        KernelDType::I64 => {
2523            execute_dynamic_update_slice::<T, i64>(plan, dest, operand, update, starts, exec)
2524        }
2525        _ => Err(StridedError::UnsupportedDType {
2526            dtype: index_dtype.label(),
2527        }),
2528    }
2529}
2530
2531fn execute_dynamic_update_slice<T, I>(
2532    plan: &DynamicUpdateSlicePlan,
2533    dest: &mut ErasedRawStridedMut<'_>,
2534    operand: &ErasedRawStridedRef<'_>,
2535    update: &ErasedRawStridedRef<'_>,
2536    starts: &ErasedRawStridedRef<'_>,
2537    exec: Exec<'_>,
2538) -> Result<()>
2539where
2540    T: Copy + Send + Sync + KernelStorageElement,
2541    I: GatherIndex + KernelStorageElement,
2542{
2543    let operand_data = operand.data_as::<T>()?;
2544    let update_data = update.data_as::<T>()?;
2545    let start_data = starts.data_as::<I>()?;
2546    let dest_dims = dest.dims();
2547    let dest_strides = dest.strides();
2548    let dest_offset = dest.offset();
2549    let dest_data = dest.data_as_mut::<T>()?;
2550    let operand_ref = unsafe {
2551        RawStridedRef::new_unchecked(
2552            operand_data,
2553            operand.dims(),
2554            operand.strides(),
2555            operand.offset(),
2556        )
2557    };
2558    let update_ref = unsafe {
2559        RawStridedRef::new_unchecked(
2560            update_data,
2561            update.dims(),
2562            update.strides(),
2563            update.offset(),
2564        )
2565    };
2566    let start_ref = unsafe {
2567        RawStridedRef::new_unchecked(start_data, starts.dims(), starts.strides(), starts.offset())
2568    };
2569    let mut dest_ref =
2570        unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) };
2571    plan.execute(&mut dest_ref, &operand_ref, &update_ref, &start_ref, exec)
2572}
2573
2574fn dispatch_scatter_index<T>(
2575    plan: &ScatterPlan,
2576    index_dtype: KernelDType,
2577    dest: &mut ErasedRawStridedMut<'_>,
2578    operand: &ErasedRawStridedRef<'_>,
2579    scatter_indices: &ErasedRawStridedRef<'_>,
2580    updates: &ErasedRawStridedRef<'_>,
2581    exec: Exec<'_>,
2582) -> Result<()>
2583where
2584    T: Copy + Add<Output = T> + Send + Sync + KernelStorageElement,
2585{
2586    match index_dtype {
2587        KernelDType::I32 => {
2588            execute_scatter::<T, i32>(plan, dest, operand, scatter_indices, updates, exec)
2589        }
2590        KernelDType::I64 => {
2591            execute_scatter::<T, i64>(plan, dest, operand, scatter_indices, updates, exec)
2592        }
2593        _ => Err(StridedError::UnsupportedDType {
2594            dtype: index_dtype.label(),
2595        }),
2596    }
2597}
2598
2599fn execute_scatter<T, I>(
2600    plan: &ScatterPlan,
2601    dest: &mut ErasedRawStridedMut<'_>,
2602    operand: &ErasedRawStridedRef<'_>,
2603    scatter_indices: &ErasedRawStridedRef<'_>,
2604    updates: &ErasedRawStridedRef<'_>,
2605    exec: Exec<'_>,
2606) -> Result<()>
2607where
2608    T: Copy + Add<Output = T> + Send + Sync + KernelStorageElement,
2609    I: GatherIndex + KernelStorageElement,
2610{
2611    let operand_data = operand.data_as::<T>()?;
2612    let index_data = scatter_indices.data_as::<I>()?;
2613    let update_data = updates.data_as::<T>()?;
2614    let dest_dims = dest.dims();
2615    let dest_strides = dest.strides();
2616    let dest_offset = dest.offset();
2617    let dest_data = dest.data_as_mut::<T>()?;
2618    let operand_ref = unsafe {
2619        RawStridedRef::new_unchecked(
2620            operand_data,
2621            operand.dims(),
2622            operand.strides(),
2623            operand.offset(),
2624        )
2625    };
2626    let index_ref = unsafe {
2627        RawStridedRef::new_unchecked(
2628            index_data,
2629            scatter_indices.dims(),
2630            scatter_indices.strides(),
2631            scatter_indices.offset(),
2632        )
2633    };
2634    let update_ref = unsafe {
2635        RawStridedRef::new_unchecked(
2636            update_data,
2637            updates.dims(),
2638            updates.strides(),
2639            updates.offset(),
2640        )
2641    };
2642    let mut dest_ref =
2643        unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) };
2644    plan.execute(&mut dest_ref, &operand_ref, &index_ref, &update_ref, exec)
2645}
2646
2647fn read_unaligned_scalar<T>(bytes: &[u8]) -> T
2648where
2649    T: Copy,
2650{
2651    unsafe { core::ptr::read_unaligned(bytes.as_ptr().cast::<T>()) }
2652}
strided_kernel/erased.rs

strided_kernel/
erased.rs
