strided_kernel/erased/
uninit.rs

1//! Dtype-erased elementwise entry points that fully overwrite uninitialized
2//! output storage.
3//!
4//! Every entry point validates dtype, shape, destination injectivity, and
5//! input/output overlap before the first write. `Ok(())` means every reachable
6//! destination element was initialized. A panic during replay may leave the
7//! destination partially initialized, which remains safe to drop because the
8//! storage is `MaybeUninit`.
9
10use super::{erased_raw_ref, map_output_dtype, raw_any, OneShotScalar};
11use crate::*;
12use core::mem::MaybeUninit;
13use core::num::Wrapping;
14use num_complex::{Complex32, Complex64};
15use strided_basic::execution::{
16    check_dtype, ensure_same_shape, map_raw_into_validated,
17    validate_destination_layout_without_alloc, validate_uninit_no_overlap,
18    zip_map2_raw_into_validated, ValidatedDestinationLayout,
19};
20
21/// Apply one runtime-selected unary operation into uninitialized storage.
22///
23/// The dtype contract matches [`erased_map_into`](crate::erased_map_into):
24/// complex absolute value writes `c32 -> f32` and `c64 -> f64`, `bool`
25/// supports only [`ErasedMapOp::Conj`], and signed integers use wrapping
26/// negate/abs.
27///
28/// # Examples
29///
30/// ```
31/// use core::mem::MaybeUninit;
32/// use strided_kernel::{
33///     erased_map_into_uninit, ErasedMapOp, ErasedRawStridedPtr, ErasedRawStridedRef,
34///     ErasedRawStridedUninitMut, Exec, KernelDType,
35/// };
36///
37/// let input = [-1.5_f64, 2.0];
38/// let input = ErasedRawStridedRef::from_slice(&input, &[2], &[1], 0).unwrap();
39/// let mut out = [MaybeUninit::<f64>::uninit(); 2];
40/// let mut dest = ErasedRawStridedUninitMut::from_uninit_slice(&mut out, &[2], &[1], 0).unwrap();
41/// erased_map_into_uninit(
42///     KernelDType::F64,
43///     ErasedMapOp::Abs,
44///     &mut dest,
45///     &ErasedRawStridedPtr::from_ref(&input),
46///     Exec::serial(),
47/// )
48/// .unwrap();
49/// // SAFETY: a successful call initializes every reachable element.
50/// assert_eq!(unsafe { [out[0].assume_init(), out[1].assume_init()] }, [1.5, 2.0]);
51/// ```
52///
53/// # Errors
54///
55/// Returns a typed [`StridedError`] for dtype, shape, output-layout, overlap,
56/// or unsupported dtype/op contracts. Validation completes before any write.
57pub fn erased_map_into_uninit(
58    input_dtype: KernelDType,
59    op: ErasedMapOp,
60    dest: &mut ErasedRawStridedUninitMut<'_>,
61    input: &ErasedRawStridedPtr<'_>,
62    exec: Exec<'_>,
63) -> Result<()> {
64    check_dtype(input_dtype, input.dtype())?;
65    check_dtype(map_output_dtype(input_dtype, op)?, dest.dtype())?;
66    validate_uninit_no_overlap(dest, input, 0)?;
67    // SAFETY: input/output overlap was rejected before forming references.
68    let input = unsafe { input.try_as_ref_after_no_overlap() }?;
69
70    match (input_dtype, op) {
71        (KernelDType::C32, ErasedMapOp::Abs) => {
72            uninit_map_with::<f32, Complex32>(dest, &input, |value| value.norm(), exec)
73        }
74        (KernelDType::C64, ErasedMapOp::Abs) => {
75            uninit_map_with::<f64, Complex64>(dest, &input, |value| value.norm(), exec)
76        }
77        (KernelDType::F32, _) => uninit_map::<f32>(op, dest, &input, exec),
78        (KernelDType::F64, _) => uninit_map::<f64>(op, dest, &input, exec),
79        (KernelDType::I32, _) => uninit_map::<i32>(op, dest, &input, exec),
80        (KernelDType::I64, _) => uninit_map::<i64>(op, dest, &input, exec),
81        (KernelDType::Bool, _) => uninit_map::<bool>(op, dest, &input, exec),
82        (KernelDType::C32, _) => uninit_map::<Complex32>(op, dest, &input, exec),
83        (KernelDType::C64, _) => uninit_map::<Complex64>(op, dest, &input, exec),
84        _ => Err(StridedError::UnsupportedDType {
85            dtype: input_dtype.label(),
86        }),
87    }
88}
89
90/// Apply one runtime-selected binary operation into uninitialized storage.
91///
92/// The dtype contract matches [`erased_zip_into`](crate::erased_zip_into).
93/// Signed-integer divide and remainder scan the divisor before any write and
94/// return [`StridedError::IntegerDivisionByZero`] when it contains zero; all
95/// other integer arithmetic wraps. Real maximum/minimum propagate NaN.
96///
97/// # Examples
98///
99/// ```
100/// use core::mem::MaybeUninit;
101/// use strided_kernel::{
102///     erased_zip_into_uninit, ErasedRawStridedPtr, ErasedRawStridedRef,
103///     ErasedRawStridedUninitMut, ErasedZipOp, Exec, KernelDType,
104/// };
105///
106/// let lhs = [i32::MAX, 7];
107/// let rhs = [1_i32, 2];
108/// let lhs = ErasedRawStridedRef::from_slice(&lhs, &[2], &[1], 0).unwrap();
109/// let rhs = ErasedRawStridedRef::from_slice(&rhs, &[2], &[1], 0).unwrap();
110/// let mut out = [MaybeUninit::<i32>::uninit(); 2];
111/// let mut dest = ErasedRawStridedUninitMut::from_uninit_slice(&mut out, &[2], &[1], 0).unwrap();
112/// erased_zip_into_uninit(
113///     KernelDType::I32,
114///     ErasedZipOp::Add,
115///     &mut dest,
116///     &ErasedRawStridedPtr::from_ref(&lhs),
117///     &ErasedRawStridedPtr::from_ref(&rhs),
118///     Exec::serial(),
119/// )
120/// .unwrap();
121/// // SAFETY: a successful call initializes every reachable element.
122/// assert_eq!(unsafe { [out[0].assume_init(), out[1].assume_init()] }, [i32::MIN, 9]);
123/// ```
124///
125/// # Errors
126///
127/// Returns a typed [`StridedError`] for dtype, shape, output-layout, overlap,
128/// unsupported dtype/op contracts, or an integer zero divisor. Validation
129/// completes before any write.
130pub fn erased_zip_into_uninit(
131    dtype: KernelDType,
132    op: ErasedZipOp,
133    dest: &mut ErasedRawStridedUninitMut<'_>,
134    lhs: &ErasedRawStridedPtr<'_>,
135    rhs: &ErasedRawStridedPtr<'_>,
136    exec: Exec<'_>,
137) -> Result<()> {
138    check_dtype(dtype, dest.dtype())?;
139    check_dtype(dtype, lhs.dtype())?;
140    check_dtype(dtype, rhs.dtype())?;
141    validate_uninit_no_overlap(dest, lhs, 0)?;
142    validate_uninit_no_overlap(dest, rhs, 1)?;
143    // SAFETY: input/output overlap was rejected before forming references.
144    let lhs = unsafe { lhs.try_as_ref_after_no_overlap() }?;
145    // SAFETY: input/output overlap was rejected before forming references.
146    let rhs = unsafe { rhs.try_as_ref_after_no_overlap() }?;
147
148    match dtype {
149        KernelDType::F32 => uninit_zip::<f32>(op, dest, &lhs, &rhs, exec),
150        KernelDType::F64 => uninit_zip::<f64>(op, dest, &lhs, &rhs, exec),
151        KernelDType::I32 => uninit_zip::<i32>(op, dest, &lhs, &rhs, exec),
152        KernelDType::I64 => uninit_zip::<i64>(op, dest, &lhs, &rhs, exec),
153        KernelDType::Bool => uninit_zip::<bool>(op, dest, &lhs, &rhs, exec),
154        KernelDType::C32 => uninit_zip::<Complex32>(op, dest, &lhs, &rhs, exec),
155        KernelDType::C64 => uninit_zip::<Complex64>(op, dest, &lhs, &rhs, exec),
156        _ => Err(StridedError::UnsupportedDType {
157            dtype: dtype.label(),
158        }),
159    }
160}
161
162/// Compare two same-dtype operands elementwise into uninitialized `bool`
163/// storage.
164///
165/// Real, signed-integer, and `bool` operands support every [`CompareOp`] with
166/// `PartialOrd` semantics, so any comparison involving NaN is `false`.
167/// Complex operands support only [`CompareOp::Eq`]; ordered comparisons return
168/// [`StridedError::UnsupportedOp`]. The destination dtype must be `bool`.
169///
170/// # Examples
171///
172/// ```
173/// use core::mem::MaybeUninit;
174/// use strided_kernel::{
175///     erased_compare_into_uninit, CompareOp, ErasedRawStridedPtr, ErasedRawStridedRef,
176///     ErasedRawStridedUninitMut, Exec, KernelDType,
177/// };
178///
179/// let lhs = [1.0_f64, f64::NAN];
180/// let rhs = [2.0_f64, 0.0];
181/// let lhs = ErasedRawStridedRef::from_slice(&lhs, &[2], &[1], 0).unwrap();
182/// let rhs = ErasedRawStridedRef::from_slice(&rhs, &[2], &[1], 0).unwrap();
183/// let mut out = [MaybeUninit::<bool>::uninit(); 2];
184/// let mut dest = ErasedRawStridedUninitMut::from_uninit_slice(&mut out, &[2], &[1], 0).unwrap();
185/// erased_compare_into_uninit(
186///     KernelDType::F64,
187///     CompareOp::Lt,
188///     &mut dest,
189///     &ErasedRawStridedPtr::from_ref(&lhs),
190///     &ErasedRawStridedPtr::from_ref(&rhs),
191///     Exec::serial(),
192/// )
193/// .unwrap();
194/// // SAFETY: a successful call initializes every reachable element.
195/// assert_eq!(unsafe { [out[0].assume_init(), out[1].assume_init()] }, [true, false]);
196/// ```
197///
198/// # Errors
199///
200/// Returns a typed [`StridedError`] for dtype, shape, output-layout, overlap,
201/// or unsupported dtype/op contracts. Validation completes before any write.
202pub fn erased_compare_into_uninit(
203    dtype: KernelDType,
204    op: CompareOp,
205    dest: &mut ErasedRawStridedUninitMut<'_>,
206    lhs: &ErasedRawStridedPtr<'_>,
207    rhs: &ErasedRawStridedPtr<'_>,
208    exec: Exec<'_>,
209) -> Result<()> {
210    check_dtype(KernelDType::Bool, dest.dtype())?;
211    check_dtype(dtype, lhs.dtype())?;
212    check_dtype(dtype, rhs.dtype())?;
213    validate_uninit_no_overlap(dest, lhs, 0)?;
214    validate_uninit_no_overlap(dest, rhs, 1)?;
215    // SAFETY: input/output overlap was rejected before forming references.
216    let lhs = unsafe { lhs.try_as_ref_after_no_overlap() }?;
217    // SAFETY: input/output overlap was rejected before forming references.
218    let rhs = unsafe { rhs.try_as_ref_after_no_overlap() }?;
219
220    match dtype {
221        KernelDType::F32 => uninit_compare_ordered::<f32>(op, dest, &lhs, &rhs, exec),
222        KernelDType::F64 => uninit_compare_ordered::<f64>(op, dest, &lhs, &rhs, exec),
223        KernelDType::I32 => uninit_compare_ordered::<i32>(op, dest, &lhs, &rhs, exec),
224        KernelDType::I64 => uninit_compare_ordered::<i64>(op, dest, &lhs, &rhs, exec),
225        KernelDType::Bool => uninit_compare_ordered::<bool>(op, dest, &lhs, &rhs, exec),
226        KernelDType::C32 => uninit_compare_eq::<Complex32>(op, dest, &lhs, &rhs, exec),
227        KernelDType::C64 => uninit_compare_eq::<Complex64>(op, dest, &lhs, &rhs, exec),
228        _ => Err(StridedError::UnsupportedDType {
229            dtype: dtype.label(),
230        }),
231    }
232}
233
234/// Select elementwise between two same-dtype operands by a `bool` predicate,
235/// writing uninitialized storage: `dest[i] = if pred[i] { on_true[i] } else
236/// { on_false[i] }`.
237///
238/// Every dtype is supported. The predicate dtype must be `bool`.
239///
240/// # Examples
241///
242/// ```
243/// use core::mem::MaybeUninit;
244/// use strided_kernel::{
245///     erased_select_into_uninit, ErasedRawStridedPtr, ErasedRawStridedRef,
246///     ErasedRawStridedUninitMut, Exec, KernelDType,
247/// };
248///
249/// let pred = [true, false];
250/// let on_true = [1_i64, 2];
251/// let on_false = [10_i64, 20];
252/// let pred = ErasedRawStridedRef::from_slice(&pred, &[2], &[1], 0).unwrap();
253/// let on_true = ErasedRawStridedRef::from_slice(&on_true, &[2], &[1], 0).unwrap();
254/// let on_false = ErasedRawStridedRef::from_slice(&on_false, &[2], &[1], 0).unwrap();
255/// let mut out = [MaybeUninit::<i64>::uninit(); 2];
256/// let mut dest = ErasedRawStridedUninitMut::from_uninit_slice(&mut out, &[2], &[1], 0).unwrap();
257/// erased_select_into_uninit(
258///     KernelDType::I64,
259///     &mut dest,
260///     &ErasedRawStridedPtr::from_ref(&pred),
261///     &ErasedRawStridedPtr::from_ref(&on_true),
262///     &ErasedRawStridedPtr::from_ref(&on_false),
263///     Exec::serial(),
264/// )
265/// .unwrap();
266/// // SAFETY: a successful call initializes every reachable element.
267/// assert_eq!(unsafe { [out[0].assume_init(), out[1].assume_init()] }, [1, 20]);
268/// ```
269///
270/// # Errors
271///
272/// Returns a typed [`StridedError`] for dtype, shape, output-layout, or
273/// overlap contracts. Validation completes before any write.
274pub fn erased_select_into_uninit(
275    dtype: KernelDType,
276    dest: &mut ErasedRawStridedUninitMut<'_>,
277    pred: &ErasedRawStridedPtr<'_>,
278    on_true: &ErasedRawStridedPtr<'_>,
279    on_false: &ErasedRawStridedPtr<'_>,
280    exec: Exec<'_>,
281) -> Result<()> {
282    check_dtype(dtype, dest.dtype())?;
283    check_dtype(KernelDType::Bool, pred.dtype())?;
284    check_dtype(dtype, on_true.dtype())?;
285    check_dtype(dtype, on_false.dtype())?;
286    validate_uninit_no_overlap(dest, pred, 0)?;
287    validate_uninit_no_overlap(dest, on_true, 1)?;
288    validate_uninit_no_overlap(dest, on_false, 2)?;
289    // SAFETY: input/output overlap was rejected before forming references.
290    let pred = unsafe { pred.try_as_ref_after_no_overlap() }?;
291    // SAFETY: input/output overlap was rejected before forming references.
292    let on_true = unsafe { on_true.try_as_ref_after_no_overlap() }?;
293    // SAFETY: input/output overlap was rejected before forming references.
294    let on_false = unsafe { on_false.try_as_ref_after_no_overlap() }?;
295
296    match dtype {
297        KernelDType::F32 => uninit_select::<f32>(dest, &pred, &on_true, &on_false, exec),
298        KernelDType::F64 => uninit_select::<f64>(dest, &pred, &on_true, &on_false, exec),
299        KernelDType::I32 => uninit_select::<i32>(dest, &pred, &on_true, &on_false, exec),
300        KernelDType::I64 => uninit_select::<i64>(dest, &pred, &on_true, &on_false, exec),
301        KernelDType::Bool => uninit_select::<bool>(dest, &pred, &on_true, &on_false, exec),
302        KernelDType::C32 => uninit_select::<Complex32>(dest, &pred, &on_true, &on_false, exec),
303        KernelDType::C64 => uninit_select::<Complex64>(dest, &pred, &on_true, &on_false, exec),
304        _ => Err(StridedError::UnsupportedDType {
305            dtype: dtype.label(),
306        }),
307    }
308}
309
310/// Clamp elementwise into uninitialized storage:
311/// `dest[i] = minimum(hi[i], maximum(lo[i], x[i]))`.
312///
313/// `maximum` and `minimum` are [`ErasedZipOp::Maximum`] and
314/// [`ErasedZipOp::Minimum`], so a NaN in any real operand yields NaN and
315/// `lo > hi` yields `hi`. Supported dtypes are `f32`, `f64`, `i32`, and `i64`.
316///
317/// # Examples
318///
319/// ```
320/// use core::mem::MaybeUninit;
321/// use strided_kernel::{
322///     erased_clamp_into_uninit, ErasedRawStridedPtr, ErasedRawStridedRef,
323///     ErasedRawStridedUninitMut, Exec, KernelDType,
324/// };
325///
326/// let x = [-3.0_f32, 0.5, 9.0];
327/// let lo = [0.0_f32; 1];
328/// let hi = [1.0_f32; 1];
329/// let x = ErasedRawStridedRef::from_slice(&x, &[3], &[1], 0).unwrap();
330/// // Stride-0 descriptors broadcast one bound over the output.
331/// let lo = ErasedRawStridedRef::from_slice(&lo, &[3], &[0], 0).unwrap();
332/// let hi = ErasedRawStridedRef::from_slice(&hi, &[3], &[0], 0).unwrap();
333/// let mut out = [MaybeUninit::<f32>::uninit(); 3];
334/// let mut dest = ErasedRawStridedUninitMut::from_uninit_slice(&mut out, &[3], &[1], 0).unwrap();
335/// erased_clamp_into_uninit(
336///     KernelDType::F32,
337///     &mut dest,
338///     &ErasedRawStridedPtr::from_ref(&x),
339///     &ErasedRawStridedPtr::from_ref(&lo),
340///     &ErasedRawStridedPtr::from_ref(&hi),
341///     Exec::serial(),
342/// )
343/// .unwrap();
344/// // SAFETY: a successful call initializes every reachable element.
345/// let out = unsafe { out.map(|value| value.assume_init()) };
346/// assert_eq!(out, [0.0, 0.5, 1.0]);
347/// ```
348///
349/// # Errors
350///
351/// Returns a typed [`StridedError`] for dtype, shape, output-layout, or
352/// overlap contracts. Validation completes before any write.
353pub fn erased_clamp_into_uninit(
354    dtype: KernelDType,
355    dest: &mut ErasedRawStridedUninitMut<'_>,
356    x: &ErasedRawStridedPtr<'_>,
357    lo: &ErasedRawStridedPtr<'_>,
358    hi: &ErasedRawStridedPtr<'_>,
359    exec: Exec<'_>,
360) -> Result<()> {
361    check_dtype(dtype, dest.dtype())?;
362    check_dtype(dtype, x.dtype())?;
363    check_dtype(dtype, lo.dtype())?;
364    check_dtype(dtype, hi.dtype())?;
365    validate_uninit_no_overlap(dest, x, 0)?;
366    validate_uninit_no_overlap(dest, lo, 1)?;
367    validate_uninit_no_overlap(dest, hi, 2)?;
368    // SAFETY: input/output overlap was rejected before forming references.
369    let x = unsafe { x.try_as_ref_after_no_overlap() }?;
370    // SAFETY: input/output overlap was rejected before forming references.
371    let lo = unsafe { lo.try_as_ref_after_no_overlap() }?;
372    // SAFETY: input/output overlap was rejected before forming references.
373    let hi = unsafe { hi.try_as_ref_after_no_overlap() }?;
374
375    match dtype {
376        KernelDType::F32 => uninit_clamp::<f32>(dest, &x, &lo, &hi, exec),
377        KernelDType::F64 => uninit_clamp::<f64>(dest, &x, &lo, &hi, exec),
378        KernelDType::I32 => uninit_clamp::<i32>(dest, &x, &lo, &hi, exec),
379        KernelDType::I64 => uninit_clamp::<i64>(dest, &x, &lo, &hi, exec),
380        _ => Err(StridedError::UnsupportedDType {
381            dtype: dtype.label(),
382        }),
383    }
384}
385
386/// Broadcast two same-dtype operands onto the destination axes and multiply
387/// them into uninitialized storage.
388///
389/// `lhs_axes[k]` names the destination axis that `lhs` axis `k` maps to; each
390/// mapped extent must equal the destination extent or be one, and unmapped
391/// destination axes broadcast. Supported dtypes are `f32`, `f64`, `c32`,
392/// `c64`, and the signed integers, whose products wrap on overflow.
393///
394/// # Examples
395///
396/// ```
397/// use core::mem::MaybeUninit;
398/// use strided_kernel::{
399///     erased_broadcast_mul_into_uninit, ErasedRawStridedPtr, ErasedRawStridedRef,
400///     ErasedRawStridedUninitMut, Exec, KernelDType,
401/// };
402///
403/// // Outer product: out[i, j] = a[i] * b[j], column-major 2 x 3 output.
404/// let a = [1.0_f64, 2.0];
405/// let b = [10.0_f64, 20.0, 30.0];
406/// let a = ErasedRawStridedRef::from_slice(&a, &[2], &[1], 0).unwrap();
407/// let b = ErasedRawStridedRef::from_slice(&b, &[3], &[1], 0).unwrap();
408/// let mut out = [MaybeUninit::<f64>::uninit(); 6];
409/// let mut dest =
410///     ErasedRawStridedUninitMut::from_uninit_slice(&mut out, &[2, 3], &[1, 2], 0).unwrap();
411/// erased_broadcast_mul_into_uninit(
412///     KernelDType::F64,
413///     &mut dest,
414///     &ErasedRawStridedPtr::from_ref(&a),
415///     &[0],
416///     &ErasedRawStridedPtr::from_ref(&b),
417///     &[1],
418///     Exec::serial(),
419/// )
420/// .unwrap();
421/// // SAFETY: a successful call initializes every reachable element.
422/// let out = unsafe { out.map(|value| value.assume_init()) };
423/// assert_eq!(out, [10.0, 20.0, 20.0, 40.0, 30.0, 60.0]);
424/// ```
425///
426/// # Errors
427///
428/// Returns a typed [`StridedError`] for dtype, axis mapping, shape,
429/// output-layout, or overlap contracts. Validation completes before any write.
430pub fn erased_broadcast_mul_into_uninit(
431    dtype: KernelDType,
432    dest: &mut ErasedRawStridedUninitMut<'_>,
433    lhs: &ErasedRawStridedPtr<'_>,
434    lhs_axes: &[usize],
435    rhs: &ErasedRawStridedPtr<'_>,
436    rhs_axes: &[usize],
437    exec: Exec<'_>,
438) -> Result<()> {
439    check_dtype(dtype, dest.dtype())?;
440    check_dtype(dtype, lhs.dtype())?;
441    check_dtype(dtype, rhs.dtype())?;
442    validate_uninit_no_overlap(dest, lhs, 0)?;
443    validate_uninit_no_overlap(dest, rhs, 1)?;
444    // SAFETY: input/output overlap was rejected before forming references.
445    let lhs = unsafe { lhs.try_as_ref_after_no_overlap() }?;
446    // SAFETY: input/output overlap was rejected before forming references.
447    let rhs = unsafe { rhs.try_as_ref_after_no_overlap() }?;
448
449    match dtype {
450        KernelDType::F32 => uninit_broadcast_mul::<f32>(dest, &lhs, lhs_axes, &rhs, rhs_axes, exec),
451        KernelDType::F64 => uninit_broadcast_mul::<f64>(dest, &lhs, lhs_axes, &rhs, rhs_axes, exec),
452        KernelDType::C32 => {
453            uninit_broadcast_mul::<Complex32>(dest, &lhs, lhs_axes, &rhs, rhs_axes, exec)
454        }
455        KernelDType::C64 => {
456            uninit_broadcast_mul::<Complex64>(dest, &lhs, lhs_axes, &rhs, rhs_axes, exec)
457        }
458        KernelDType::I32 => {
459            uninit_broadcast_mul_wrapping::<i32>(dest, &lhs, lhs_axes, &rhs, rhs_axes, exec)
460        }
461        KernelDType::I64 => {
462            uninit_broadcast_mul_wrapping::<i64>(dest, &lhs, lhs_axes, &rhs, rhs_axes, exec)
463        }
464        _ => Err(StridedError::UnsupportedDType {
465            dtype: dtype.label(),
466        }),
467    }
468}
469
470/// Validate the destination layout and the input shapes.
471///
472/// Returns `None` for an empty destination, which needs no writes.
473fn prepare_destination(
474    dest: &ErasedRawStridedUninitMut<'_>,
475    inputs: &[&[usize]],
476) -> Result<Option<ValidatedDestinationLayout>> {
477    let validated = validate_destination_layout_without_alloc(dest.dims(), dest.strides())?;
478    for input_dims in inputs {
479        ensure_same_shape(dest.dims(), input_dims)?;
480    }
481    if dest.dims().contains(&0) {
482        Ok(None)
483    } else {
484        Ok(Some(validated))
485    }
486}
487
488fn uninit_raw_mut<'a, T: KernelStorageElement>(
489    dest: &'a mut ErasedRawStridedUninitMut<'_>,
490) -> Result<RawStridedMut<'a, MaybeUninit<T>>> {
491    let dims = dest.dims();
492    let strides = dest.strides();
493    let offset = dest.offset();
494    let data = dest.data_as_uninit_mut::<T>()?;
495    // SAFETY: the uninit descriptor validated every reachable offset against
496    // its storage at construction.
497    Ok(unsafe { RawStridedMut::new_unchecked(data, dims, strides, offset) })
498}
499
500fn uninit_view_mut<'a, T: KernelStorageElement>(
501    dest: &'a mut ErasedRawStridedUninitMut<'_>,
502) -> Result<StridedViewMut<'a, MaybeUninit<T>>> {
503    let dims = dest.dims();
504    let strides = dest.strides();
505    let offset = dest.offset();
506    let data = dest.data_as_uninit_mut::<T>()?;
507    // SAFETY: the uninit descriptor validated every reachable offset against
508    // its storage at construction.
509    Ok(unsafe { StridedViewMut::new_unchecked(data, dims, strides, offset) })
510}
511
512fn erased_view<'a, T: KernelStorageElement>(
513    src: &'a ErasedRawStridedRef<'a>,
514) -> Result<StridedView<'a, T>> {
515    let data = src.data_as::<T>()?;
516    // SAFETY: the descriptor validated every reachable offset at construction.
517    Ok(unsafe { StridedView::new_unchecked(data, src.dims(), src.strides(), src.offset()) })
518}
519
520fn uninit_map<T: OneShotScalar>(
521    op: ErasedMapOp,
522    dest: &mut ErasedRawStridedUninitMut<'_>,
523    input: &ErasedRawStridedRef<'_>,
524    exec: Exec<'_>,
525) -> Result<()> {
526    if !T::supports_map(op) {
527        return Err(StridedError::UnsupportedOp {
528            op: op.label(),
529            dtype: T::one_shot_dtype_label(),
530        });
531    }
532    // Select the operation once, outside the element loop, so each replay
533    // closure is monomorphic and the inner loop can vectorize.
534    match op {
535        ErasedMapOp::Negate => uninit_map_with::<T, T>(
536            dest,
537            input,
538            |value| T::map(ErasedMapOp::Negate, value),
539            exec,
540        ),
541        ErasedMapOp::Conj => {
542            uninit_map_with::<T, T>(dest, input, |value| T::map(ErasedMapOp::Conj, value), exec)
543        }
544        ErasedMapOp::Abs => {
545            uninit_map_with::<T, T>(dest, input, |value| T::map(ErasedMapOp::Abs, value), exec)
546        }
547        ErasedMapOp::Sign => {
548            uninit_map_with::<T, T>(dest, input, |value| T::map(ErasedMapOp::Sign, value), exec)
549        }
550    }
551}
552
553fn uninit_map_with<D, A>(
554    dest: &mut ErasedRawStridedUninitMut<'_>,
555    input: &ErasedRawStridedRef<'_>,
556    map: impl Fn(A) -> D + Sync,
557    exec: Exec<'_>,
558) -> Result<()>
559where
560    D: Copy + Send + Sync + KernelStorageElement,
561    A: Copy + Send + Sync + KernelStorageElement,
562{
563    let Some(validated) = prepare_destination(dest, &[input.dims()])? else {
564        return Ok(());
565    };
566    let input = erased_raw_ref::<A>(input)?;
567    let mut dest = uninit_raw_mut::<D>(dest)?;
568    // SAFETY: matching shapes and this destination layout were validated above,
569    // and the caller rejected input/output overlap.
570    unsafe {
571        map_raw_into_validated::<MaybeUninit<D>, A, Identity>(
572            &mut dest,
573            &input,
574            |value| MaybeUninit::new(map(value)),
575            validated,
576            exec,
577        )
578    }
579}
580
581fn uninit_zip<T: OneShotScalar>(
582    op: ErasedZipOp,
583    dest: &mut ErasedRawStridedUninitMut<'_>,
584    lhs: &ErasedRawStridedRef<'_>,
585    rhs: &ErasedRawStridedRef<'_>,
586    exec: Exec<'_>,
587) -> Result<()> {
588    if !T::supports_zip(op) {
589        return Err(StridedError::UnsupportedOp {
590            op: op.label(),
591            dtype: T::one_shot_dtype_label(),
592        });
593    }
594    let Some(validated) = prepare_destination(dest, &[lhs.dims(), rhs.dims()])? else {
595        return Ok(());
596    };
597    let lhs = erased_raw_ref::<T>(lhs)?;
598    let rhs = erased_raw_ref::<T>(rhs)?;
599    if matches!(op, ErasedZipOp::Divide | ErasedZipOp::Remainder)
600        && T::INTEGER
601        && raw_any(&rhs, T::is_zero)?
602    {
603        return Err(StridedError::IntegerDivisionByZero { op: op.label() });
604    }
605    let mut dest = uninit_raw_mut::<T>(dest)?;
606    // Select the operation once, outside the element loop, so each replay
607    // closure is monomorphic and the inner loop can vectorize.
608    macro_rules! replay {
609        ($op:ident) => {
610            // SAFETY: matching shapes and this destination layout were
611            // validated above, and the caller rejected input/output overlap.
612            unsafe {
613                zip_map2_raw_into_validated::<MaybeUninit<T>, T, T, Identity, Identity>(
614                    &mut dest,
615                    &lhs,
616                    &rhs,
617                    |lhs, rhs| MaybeUninit::new(T::zip(ErasedZipOp::$op, lhs, rhs)),
618                    validated,
619                    exec,
620                )
621            }
622        };
623    }
624    match op {
625        ErasedZipOp::Add => replay!(Add),
626        ErasedZipOp::Subtract => replay!(Subtract),
627        ErasedZipOp::Multiply => replay!(Multiply),
628        ErasedZipOp::Divide => replay!(Divide),
629        ErasedZipOp::Remainder => replay!(Remainder),
630        ErasedZipOp::Maximum => replay!(Maximum),
631        ErasedZipOp::Minimum => replay!(Minimum),
632    }
633}
634
635fn uninit_compare_with<T>(
636    dest: &mut ErasedRawStridedUninitMut<'_>,
637    lhs: &ErasedRawStridedRef<'_>,
638    rhs: &ErasedRawStridedRef<'_>,
639    compare: impl Fn(T, T) -> bool + Sync,
640    exec: Exec<'_>,
641) -> Result<()>
642where
643    T: Copy + Send + Sync + KernelStorageElement,
644{
645    let Some(validated) = prepare_destination(dest, &[lhs.dims(), rhs.dims()])? else {
646        return Ok(());
647    };
648    let lhs = erased_raw_ref::<T>(lhs)?;
649    let rhs = erased_raw_ref::<T>(rhs)?;
650    let mut dest = uninit_raw_mut::<bool>(dest)?;
651    // SAFETY: matching shapes and this destination layout were validated above,
652    // and the caller rejected input/output overlap.
653    unsafe {
654        zip_map2_raw_into_validated::<MaybeUninit<bool>, T, T, Identity, Identity>(
655            &mut dest,
656            &lhs,
657            &rhs,
658            |lhs, rhs| MaybeUninit::new(compare(lhs, rhs)),
659            validated,
660            exec,
661        )
662    }
663}
664
665fn uninit_compare_ordered<T>(
666    op: CompareOp,
667    dest: &mut ErasedRawStridedUninitMut<'_>,
668    lhs: &ErasedRawStridedRef<'_>,
669    rhs: &ErasedRawStridedRef<'_>,
670    exec: Exec<'_>,
671) -> Result<()>
672where
673    T: Copy + Send + Sync + KernelStorageElement + PartialOrd,
674{
675    // The comparison is selected once, outside the element loop.
676    match op {
677        CompareOp::Eq => uninit_compare_with::<T>(dest, lhs, rhs, |a, b| a == b, exec),
678        CompareOp::Lt => uninit_compare_with::<T>(dest, lhs, rhs, |a, b| a < b, exec),
679        CompareOp::Le => uninit_compare_with::<T>(dest, lhs, rhs, |a, b| a <= b, exec),
680        CompareOp::Gt => uninit_compare_with::<T>(dest, lhs, rhs, |a, b| a > b, exec),
681        CompareOp::Ge => uninit_compare_with::<T>(dest, lhs, rhs, |a, b| a >= b, exec),
682        _ => Err(StridedError::UnsupportedOp {
683            op: compare_label(op),
684            dtype: T::DTYPE.label(),
685        }),
686    }
687}
688
689fn uninit_compare_eq<T>(
690    op: CompareOp,
691    dest: &mut ErasedRawStridedUninitMut<'_>,
692    lhs: &ErasedRawStridedRef<'_>,
693    rhs: &ErasedRawStridedRef<'_>,
694    exec: Exec<'_>,
695) -> Result<()>
696where
697    T: Copy + Send + Sync + KernelStorageElement + PartialEq,
698{
699    match op {
700        CompareOp::Eq => uninit_compare_with::<T>(dest, lhs, rhs, |a, b| a == b, exec),
701        _ => Err(StridedError::UnsupportedOp {
702            op: compare_label(op),
703            dtype: T::DTYPE.label(),
704        }),
705    }
706}
707
708fn compare_label(op: CompareOp) -> &'static str {
709    match op {
710        CompareOp::Eq => "eq",
711        CompareOp::Lt => "lt",
712        CompareOp::Le => "le",
713        CompareOp::Gt => "gt",
714        CompareOp::Ge => "ge",
715        _ => "compare",
716    }
717}
718
719fn uninit_select<T>(
720    dest: &mut ErasedRawStridedUninitMut<'_>,
721    pred: &ErasedRawStridedRef<'_>,
722    on_true: &ErasedRawStridedRef<'_>,
723    on_false: &ErasedRawStridedRef<'_>,
724    exec: Exec<'_>,
725) -> Result<()>
726where
727    T: Copy + Send + Sync + KernelStorageElement,
728{
729    let dims: [&[usize]; 3] = [pred.dims(), on_true.dims(), on_false.dims()];
730    if prepare_destination(dest, &dims)?.is_none() {
731        return Ok(());
732    }
733    let pred = erased_view::<bool>(pred)?;
734    let on_true = erased_view::<T>(on_true)?;
735    let on_false = erased_view::<T>(on_false)?;
736    let mut dest = uninit_view_mut::<T>(dest)?;
737    zip_map3_into(
738        &mut dest,
739        &pred,
740        &on_true,
741        &on_false,
742        |pred, a, b| MaybeUninit::new(if pred { a } else { b }),
743        exec,
744    )
745}
746
747fn uninit_clamp<T: OneShotScalar>(
748    dest: &mut ErasedRawStridedUninitMut<'_>,
749    x: &ErasedRawStridedRef<'_>,
750    lo: &ErasedRawStridedRef<'_>,
751    hi: &ErasedRawStridedRef<'_>,
752    exec: Exec<'_>,
753) -> Result<()> {
754    let dims: [&[usize]; 3] = [x.dims(), lo.dims(), hi.dims()];
755    if prepare_destination(dest, &dims)?.is_none() {
756        return Ok(());
757    }
758    let x = erased_view::<T>(x)?;
759    let lo = erased_view::<T>(lo)?;
760    let hi = erased_view::<T>(hi)?;
761    let mut dest = uninit_view_mut::<T>(dest)?;
762    zip_map3_into(
763        &mut dest,
764        &x,
765        &lo,
766        &hi,
767        |x, lo, hi| MaybeUninit::new(T::clamp(x, lo, hi)),
768        exec,
769    )
770}
771
772fn uninit_broadcast_mul<T>(
773    dest: &mut ErasedRawStridedUninitMut<'_>,
774    lhs: &ErasedRawStridedRef<'_>,
775    lhs_axes: &[usize],
776    rhs: &ErasedRawStridedRef<'_>,
777    rhs_axes: &[usize],
778    exec: Exec<'_>,
779) -> Result<()>
780where
781    T: Copy + Send + Sync + KernelStorageElement + core::ops::Mul<Output = T>,
782{
783    let lhs = erased_view::<T>(lhs)?;
784    let rhs = erased_view::<T>(rhs)?;
785    let mut dest = uninit_view_mut::<T>(dest)?;
786    broadcast_mul_into_uninit(&mut dest, &lhs, lhs_axes, &rhs, rhs_axes, exec)
787}
788
789fn uninit_broadcast_mul_wrapping<T>(
790    dest: &mut ErasedRawStridedUninitMut<'_>,
791    lhs: &ErasedRawStridedRef<'_>,
792    lhs_axes: &[usize],
793    rhs: &ErasedRawStridedRef<'_>,
794    rhs_axes: &[usize],
795    exec: Exec<'_>,
796) -> Result<()>
797where
798    T: Copy + Send + Sync + KernelStorageElement,
799    Wrapping<T>: core::ops::Mul<Output = Wrapping<T>> + Send + Sync,
800{
801    let lhs_data = wrapping_slice(lhs.data_as::<T>()?);
802    let rhs_data = wrapping_slice(rhs.data_as::<T>()?);
803    // SAFETY: the descriptors validated every reachable offset at construction,
804    // and the reinterpreted slices keep their element count.
805    let lhs = unsafe {
806        StridedView::<Wrapping<T>, Identity>::new_unchecked(
807            lhs_data,
808            lhs.dims(),
809            lhs.strides(),
810            lhs.offset(),
811        )
812    };
813    // SAFETY: as above.
814    let rhs = unsafe {
815        StridedView::<Wrapping<T>, Identity>::new_unchecked(
816            rhs_data,
817            rhs.dims(),
818            rhs.strides(),
819            rhs.offset(),
820        )
821    };
822    let dims = dest.dims();
823    let strides = dest.strides();
824    let offset = dest.offset();
825    let data = dest.data_as_uninit_mut::<T>()?;
826    // SAFETY: `Wrapping<T>` is `repr(transparent)` over `T`, so
827    // `MaybeUninit<Wrapping<T>>` has the layout of `MaybeUninit<T>`, and the
828    // slice keeps its element count and exclusive borrow.
829    let data = unsafe {
830        core::slice::from_raw_parts_mut(
831            data.as_mut_ptr().cast::<MaybeUninit<Wrapping<T>>>(),
832            data.len(),
833        )
834    };
835    // SAFETY: the uninit descriptor validated every reachable offset.
836    let mut dest = unsafe { StridedViewMut::new_unchecked(data, dims, strides, offset) };
837    broadcast_mul_into_uninit(&mut dest, &lhs, lhs_axes, &rhs, rhs_axes, exec)
838}
839
840fn wrapping_slice<T>(data: &[T]) -> &[Wrapping<T>] {
841    // SAFETY: `Wrapping<T>` is `repr(transparent)` over `T`; the slice keeps
842    // its element count and shared borrow.
843    unsafe { core::slice::from_raw_parts(data.as_ptr().cast::<Wrapping<T>>(), data.len()) }
844}
strided_kernel/erased/uninit.rs

strided_kernel/erased/
uninit.rs
