Skip to main content

strided_view/
raw.rs

1//! Borrowed raw strided layout types.
2//!
3//! These are the prepared-replay counterparts to [`crate::StridedView`] and
4//! [`crate::StridedViewMut`]. They borrow shape/stride metadata instead of
5//! owning it, so compiled kernels can reuse already-validated layout
6//! descriptors without rebuilding dynamic-rank view wrappers.
7
8use crate::element_op::Identity;
9use crate::view::validate_bounds;
10use core::marker::PhantomData;
11use core::mem::MaybeUninit;
12use core::ptr::NonNull;
13use num_complex::{Complex32, Complex64};
14
15mod private {
16    pub trait Sealed {}
17}
18
19/// A scalar type that can safely witness storage for an erased kernel view.
20pub trait KernelStorageElement: private::Sealed + Copy + 'static {
21    const DTYPE: KernelDType;
22}
23
24macro_rules! kernel_storage_element {
25    ($($ty:ty => $dtype:ident),* $(,)?) => {$ (
26        impl private::Sealed for $ty {}
27        impl KernelStorageElement for $ty {
28            const DTYPE: KernelDType = KernelDType::$dtype;
29        }
30    )* };
31}
32
33kernel_storage_element! {
34    f32 => F32,
35    f64 => F64,
36    i32 => I32,
37    i64 => I64,
38    bool => Bool,
39    Complex32 => C32,
40    Complex64 => C64,
41}
42
43use crate::{Result, StridedError, StridedView, StridedViewMut};
44
45/// Dtypes supported by dtype-erased kernel entry points.
46///
47/// The enum is intentionally limited to the scalar set currently used by the
48/// tensor runtime callers. Later FFI layers should map their ABI dtype tags to
49/// this enum before dispatching into prepared kernels.
50#[non_exhaustive]
51#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
52#[repr(u8)]
53pub enum KernelDType {
54    F32 = 1,
55    F64 = 2,
56    I32 = 3,
57    I64 = 4,
58    Bool = 5,
59    C32 = 6,
60    C64 = 7,
61}
62
63impl KernelDType {
64    #[inline]
65    pub const fn label(self) -> &'static str {
66        match self {
67            Self::F32 => "f32",
68            Self::F64 => "f64",
69            Self::I32 => "i32",
70            Self::I64 => "i64",
71            Self::Bool => "bool",
72            Self::C32 => "c32",
73            Self::C64 => "c64",
74        }
75    }
76
77    #[inline]
78    pub const fn size_of(self) -> usize {
79        match self {
80            Self::F32 => core::mem::size_of::<f32>(),
81            Self::F64 => core::mem::size_of::<f64>(),
82            Self::I32 => core::mem::size_of::<i32>(),
83            Self::I64 => core::mem::size_of::<i64>(),
84            Self::Bool => core::mem::size_of::<bool>(),
85            Self::C32 => core::mem::size_of::<Complex32>(),
86            Self::C64 => core::mem::size_of::<Complex64>(),
87        }
88    }
89
90    #[inline]
91    pub const fn alignment(self) -> usize {
92        match self {
93            Self::F32 => core::mem::align_of::<f32>(),
94            Self::F64 => core::mem::align_of::<f64>(),
95            Self::I32 => core::mem::align_of::<i32>(),
96            Self::I64 => core::mem::align_of::<i64>(),
97            Self::Bool => core::mem::align_of::<bool>(),
98            Self::C32 => core::mem::align_of::<Complex32>(),
99            Self::C64 => core::mem::align_of::<Complex64>(),
100        }
101    }
102
103    #[inline]
104    pub const fn requires_valid_byte_values(self) -> bool {
105        matches!(self, Self::Bool)
106    }
107}
108
109fn validate_erased_buffer(dtype: KernelDType, data: &[u8]) -> Result<usize> {
110    let element_size = dtype.size_of();
111    if data.len() % element_size != 0 {
112        return Err(StridedError::ByteLengthMismatch {
113            dtype: dtype.label(),
114            byte_len: data.len(),
115            element_size,
116        });
117    }
118
119    let element_count = data.len() / element_size;
120    if element_count == 0 {
121        return Ok(0);
122    }
123
124    let alignment = dtype.alignment();
125    if data.as_ptr() as usize % alignment != 0 {
126        return Err(StridedError::DataAlignmentMismatch {
127            dtype: dtype.label(),
128            alignment,
129        });
130    }
131    if dtype.requires_valid_byte_values() {
132        if let Some(&value) = data.iter().find(|&&value| value > 1) {
133            return Err(StridedError::InvalidBoolByte { value });
134        }
135    }
136    Ok(element_count)
137}
138
139fn validate_erased_buffer_layout(
140    dtype: KernelDType,
141    data: NonNull<u8>,
142    byte_len: usize,
143) -> Result<usize> {
144    let element_size = dtype.size_of();
145    if byte_len % element_size != 0 {
146        return Err(StridedError::ByteLengthMismatch {
147            dtype: dtype.label(),
148            byte_len,
149            element_size,
150        });
151    }
152    if byte_len != 0 && data.as_ptr() as usize % dtype.alignment() != 0 {
153        return Err(StridedError::DataAlignmentMismatch {
154            dtype: dtype.label(),
155            alignment: dtype.alignment(),
156        });
157    }
158    Ok(byte_len / element_size)
159}
160
161/// Pointer-backed dtype-erased input used by one-shot write entry points.
162///
163/// Unlike [`ErasedRawStridedRef`], this descriptor does not create a shared
164/// Rust reference at construction time. That lets the entry point reject an
165/// input/output overlap before forming references to either side.
166#[derive(Clone, Copy, Debug)]
167pub struct ErasedRawStridedPtr<'a> {
168    dtype: KernelDType,
169    data: NonNull<u8>,
170    byte_len: usize,
171    dims: &'a [usize],
172    strides: &'a [isize],
173    offset: isize,
174    marker: PhantomData<&'a [MaybeUninit<u8>]>,
175}
176
177impl<'a> ErasedRawStridedPtr<'a> {
178    /// Create a pointer-backed descriptor from erased storage.
179    ///
180    /// # Safety
181    ///
182    /// `data` must point into an allocation whose alignment is suitable for
183    /// `dtype`, independently of the observed address. The allocation must
184    /// provide `byte_len` initialized, readable bytes with valid provenance
185    /// for `'a`, and remain alive for `'a`. For `bool`, the bytes may contain
186    /// invalid values temporarily. They must not be read or used to form a
187    /// typed reference until a caller has rejected overlap and
188    /// [`Self::try_as_ref_after_no_overlap`] has validated the complete extent.
189    /// The allocation may overlap a destination.
190    /// The original owner may perform synchronized sequential mutation through
191    /// the same-provenance raw pointer before conversion. No concurrent
192    /// mutation or conflicting access is permitted during overlap checking,
193    /// conversion, or consumer access.
194    pub unsafe fn from_raw_parts(
195        dtype: KernelDType,
196        data: NonNull<u8>,
197        byte_len: usize,
198        dims: &'a [usize],
199        strides: &'a [isize],
200        offset: isize,
201    ) -> Result<Self> {
202        let element_count = validate_erased_buffer_layout(dtype, data, byte_len)?;
203        validate_bounds(element_count, dims, strides, offset)?;
204        Ok(Self {
205            dtype,
206            data,
207            byte_len,
208            dims,
209            strides,
210            offset,
211            marker: PhantomData,
212        })
213    }
214
215    /// Borrow a safe erased input as a pointer descriptor.
216    pub fn from_ref(input: &ErasedRawStridedRef<'a>) -> Self {
217        Self {
218            dtype: input.dtype,
219            data: NonNull::new(input.data.as_ptr().cast_mut()).unwrap_or_else(NonNull::dangling),
220            byte_len: input.data.len(),
221            dims: input.dims,
222            strides: input.strides,
223            offset: input.offset,
224            marker: PhantomData,
225        }
226    }
227
228    /// Convert this pointer to an initialized descriptor after overlap checks.
229    ///
230    /// # Safety
231    ///
232    /// The caller must have proved that no mutable allocation overlaps this
233    /// pointer's complete byte extent.
234    /// The allocation contract from [`Self::from_raw_parts`] must still hold.
235    /// No mutation or other conflicting access may occur during this conversion
236    /// or for the duration of the returned descriptor's consumer access.
237    /// Bool bytes are validated here, immediately before typed access.
238    pub unsafe fn try_as_ref_after_no_overlap(&self) -> Result<ErasedRawStridedRef<'_>> {
239        let bytes = core::slice::from_raw_parts(self.data.as_ptr(), self.byte_len);
240        validate_erased_buffer(self.dtype, bytes)?;
241        ErasedRawStridedRef::from_raw_parts(
242            self.dtype,
243            self.data,
244            self.byte_len,
245            self.dims,
246            self.strides,
247            self.offset,
248        )
249    }
250
251    #[inline]
252    pub fn dtype(&self) -> KernelDType {
253        self.dtype
254    }
255
256    /// Check overlap with initialized mutable storage without reading it.
257    pub fn overlaps_mut(&self, dest: &ErasedRawStridedMut<'_>) -> Result<bool> {
258        ranges_overlap(
259            self.data.as_ptr() as usize,
260            self.byte_len,
261            dest.data.as_ptr() as usize,
262            dest.data.len(),
263        )
264    }
265
266    /// Check overlap with uninitialized mutable storage without reading it.
267    pub fn overlaps_uninit_mut(&self, dest: &ErasedRawStridedUninitMut<'_>) -> Result<bool> {
268        ranges_overlap(
269            self.data.as_ptr() as usize,
270            self.byte_len,
271            dest.data.as_ptr() as usize,
272            dest.data.len(),
273        )
274    }
275
276    #[inline]
277    pub fn dims(&self) -> &'a [usize] {
278        self.dims
279    }
280
281    #[inline]
282    pub fn strides(&self) -> &'a [isize] {
283        self.strides
284    }
285
286    #[inline]
287    pub fn offset(&self) -> isize {
288        self.offset
289    }
290}
291
292/// Borrowed dtype-erased raw strided input layout.
293///
294/// `dims`, `strides`, and `offset` are expressed in dtype elements, not bytes.
295#[derive(Clone, Copy, Debug)]
296pub struct ErasedRawStridedRef<'a> {
297    dtype: KernelDType,
298    data: &'a [u8],
299    dims: &'a [usize],
300    strides: &'a [isize],
301    offset: isize,
302}
303
304impl<'a> ErasedRawStridedRef<'a> {
305    /// Create an erased input descriptor from typed, initialized storage.
306    pub fn from_slice<T: KernelStorageElement>(
307        data: &'a [T],
308        dims: &'a [usize],
309        strides: &'a [isize],
310        offset: isize,
311    ) -> Result<Self> {
312        let dtype = T::DTYPE;
313        let byte_len = data
314            .len()
315            .checked_mul(core::mem::size_of::<T>())
316            .ok_or(StridedError::OffsetOverflow)?;
317        let bytes = unsafe { core::slice::from_raw_parts(data.as_ptr().cast::<u8>(), byte_len) };
318        let element_count = validate_erased_buffer(dtype, bytes)?;
319        validate_bounds(element_count, dims, strides, offset)?;
320        Ok(Self {
321            dtype,
322            data: bytes,
323            dims,
324            strides,
325            offset,
326        })
327    }
328
329    /// Construct from erased storage.
330    ///
331    /// # Safety
332    /// `data` must be the start of an allocation aligned for `dtype`; the
333    /// allocation must contain `byte_len` initialized, readable bytes for
334    /// `'a`, and every byte extent must represent valid values for `dtype`.
335    /// The allocation and metadata must outlive `'a`; no mutable alias may
336    /// exist while this descriptor is used. Alignment is an allocation
337    /// property and must not be inferred from the observed address alone.
338    pub unsafe fn from_raw_parts(
339        dtype: KernelDType,
340        data: NonNull<u8>,
341        byte_len: usize,
342        dims: &'a [usize],
343        strides: &'a [isize],
344        offset: isize,
345    ) -> Result<Self> {
346        let count = validate_erased_buffer_layout(dtype, data, byte_len)?;
347        let bytes = core::slice::from_raw_parts(data.as_ptr(), byte_len);
348        validate_bounds(count, dims, strides, offset)?;
349        Ok(Self {
350            dtype,
351            data: bytes,
352            dims,
353            strides,
354            offset,
355        })
356    }
357
358    /// Borrow the initialized storage as its concrete scalar type.
359    pub fn data_as<T: KernelStorageElement>(&self) -> Result<&[T]> {
360        if self.dtype != T::DTYPE {
361            return Err(StridedError::DTypeMismatch {
362                expected: T::DTYPE.label(),
363                actual: self.dtype.label(),
364            });
365        }
366        Ok(unsafe {
367            core::slice::from_raw_parts(
368                self.data.as_ptr().cast::<T>(),
369                self.data.len() / core::mem::size_of::<T>(),
370            )
371        })
372    }
373
374    #[inline]
375    pub fn dtype(&self) -> KernelDType {
376        self.dtype
377    }
378
379    #[inline]
380    pub fn dims(&self) -> &'a [usize] {
381        self.dims
382    }
383
384    #[inline]
385    pub fn strides(&self) -> &'a [isize] {
386        self.strides
387    }
388
389    #[inline]
390    pub fn offset(&self) -> isize {
391        self.offset
392    }
393}
394
395/// Borrowed dtype-erased raw strided output layout.
396///
397/// `dims`, `strides`, and `offset` are expressed in dtype elements, not bytes.
398#[derive(Debug)]
399pub struct ErasedRawStridedMut<'a> {
400    dtype: KernelDType,
401    data: &'a mut [u8],
402    dims: &'a [usize],
403    strides: &'a [isize],
404    offset: isize,
405}
406
407impl<'a> ErasedRawStridedMut<'a> {
408    /// Create an erased output descriptor from typed, initialized storage.
409    pub fn from_slice_mut<T: KernelStorageElement>(
410        data: &'a mut [T],
411        dims: &'a [usize],
412        strides: &'a [isize],
413        offset: isize,
414    ) -> Result<Self> {
415        let dtype = T::DTYPE;
416        let byte_len = data
417            .len()
418            .checked_mul(core::mem::size_of::<T>())
419            .ok_or(StridedError::OffsetOverflow)?;
420        let data =
421            unsafe { core::slice::from_raw_parts_mut(data.as_mut_ptr().cast::<u8>(), byte_len) };
422        let element_count = validate_erased_buffer(dtype, data)?;
423        validate_bounds(element_count, dims, strides, offset)?;
424        Ok(Self {
425            dtype,
426            data,
427            dims,
428            strides,
429            offset,
430        })
431    }
432
433    /// Construct from erased storage.
434    ///
435    /// # Safety
436    /// `data..data + byte_len` must be an aligned, writable allocation valid
437    /// for `'a`, with valid initialized values for `dtype` throughout the
438    /// complete byte extent; its provenance, extent, lifetime, and exclusive
439    /// aliasing must be upheld by the caller.
440    /// The alignment requirement is on the allocation, not merely the observed
441    /// address. The caller must retain exclusive access: no mutable alias or
442    /// concurrent mutation may exist while this descriptor is used.
443    pub unsafe fn from_raw_parts(
444        dtype: KernelDType,
445        data: NonNull<u8>,
446        byte_len: usize,
447        dims: &'a [usize],
448        strides: &'a [isize],
449        offset: isize,
450    ) -> Result<Self> {
451        let count = validate_erased_buffer_layout(dtype, data, byte_len)?;
452        let data = core::slice::from_raw_parts_mut(data.as_ptr(), byte_len);
453        validate_erased_buffer(dtype, data)?;
454        validate_bounds(count, dims, strides, offset)?;
455        Ok(Self {
456            dtype,
457            data,
458            dims,
459            strides,
460            offset,
461        })
462    }
463
464    /// Borrow initialized storage as its concrete scalar type.
465    pub fn data_as<T: KernelStorageElement>(&self) -> Result<&[T]> {
466        if self.dtype != T::DTYPE {
467            return Err(StridedError::DTypeMismatch {
468                expected: T::DTYPE.label(),
469                actual: self.dtype.label(),
470            });
471        }
472        Ok(unsafe {
473            core::slice::from_raw_parts(
474                self.data.as_ptr().cast::<T>(),
475                self.data.len() / core::mem::size_of::<T>(),
476            )
477        })
478    }
479
480    /// Borrow initialized storage mutably as its concrete scalar type.
481    pub fn data_as_mut<T: KernelStorageElement>(&mut self) -> Result<&mut [T]> {
482        if self.dtype != T::DTYPE {
483            return Err(StridedError::DTypeMismatch {
484                expected: T::DTYPE.label(),
485                actual: self.dtype.label(),
486            });
487        }
488        Ok(unsafe {
489            core::slice::from_raw_parts_mut(
490                self.data.as_mut_ptr().cast::<T>(),
491                self.data.len() / core::mem::size_of::<T>(),
492            )
493        })
494    }
495
496    #[inline]
497    pub fn dtype(&self) -> KernelDType {
498        self.dtype
499    }
500
501    #[inline]
502    pub fn dims(&self) -> &'a [usize] {
503        self.dims
504    }
505
506    #[inline]
507    pub fn strides(&self) -> &'a [isize] {
508        self.strides
509    }
510
511    #[inline]
512    pub fn offset(&self) -> isize {
513        self.offset
514    }
515}
516
517/// Borrowed dtype-erased raw strided output whose reachable elements may be
518/// uninitialized.
519///
520/// This descriptor is only accepted by operations that prove and perform a
521/// full overwrite of every reachable logical destination element. The backing
522/// allocation may contain non-reachable holes, which remain uninitialized.
523#[derive(Debug)]
524pub struct ErasedRawStridedUninitMut<'a> {
525    dtype: KernelDType,
526    data: &'a mut [MaybeUninit<u8>],
527    dims: &'a [usize],
528    strides: &'a [isize],
529    offset: isize,
530}
531
532impl<'a> ErasedRawStridedUninitMut<'a> {
533    /// Create an uninitialized erased output descriptor from typed storage.
534    pub fn from_uninit_slice<T: KernelStorageElement>(
535        data: &'a mut [MaybeUninit<T>],
536        dims: &'a [usize],
537        strides: &'a [isize],
538        offset: isize,
539    ) -> Result<Self> {
540        let dtype = T::DTYPE;
541        let byte_len = data
542            .len()
543            .checked_mul(core::mem::size_of::<T>())
544            .ok_or(StridedError::OffsetOverflow)?;
545        let data = unsafe {
546            core::slice::from_raw_parts_mut(data.as_mut_ptr().cast::<MaybeUninit<u8>>(), byte_len)
547        };
548        let data_ptr =
549            NonNull::new(data.as_mut_ptr().cast::<u8>()).unwrap_or_else(NonNull::dangling);
550        let element_count = validate_erased_buffer_layout(dtype, data_ptr, byte_len)?;
551        validate_bounds(element_count, dims, strides, offset)?;
552        Ok(Self {
553            dtype,
554            data,
555            dims,
556            strides,
557            offset,
558        })
559    }
560
561    /// Construct from erased uninitialized storage.
562    ///
563    /// # Safety
564    /// `data..data + byte_len` must be an aligned, writable allocation valid
565    /// for `'a`, with provenance and extent sufficient for all reachable
566    /// elements. The caller must ensure every reachable element is completely
567    /// overwritten before it is read or exposed as initialized storage. The
568    /// allocation's alignment is an allocation property independent of the
569    /// observed address, and the descriptor must have exclusive access with no
570    /// mutable alias or concurrent mutation during use.
571    pub unsafe fn from_raw_parts(
572        dtype: KernelDType,
573        data: NonNull<u8>,
574        byte_len: usize,
575        dims: &'a [usize],
576        strides: &'a [isize],
577        offset: isize,
578    ) -> Result<Self> {
579        let count = validate_erased_buffer_layout(dtype, data, byte_len)?;
580        let data =
581            core::slice::from_raw_parts_mut(data.as_ptr().cast::<MaybeUninit<u8>>(), byte_len);
582        validate_bounds(count, dims, strides, offset)?;
583        Ok(Self {
584            dtype,
585            data,
586            dims,
587            strides,
588            offset,
589        })
590    }
591
592    /// Borrow the uninitialized storage as its concrete scalar type.
593    pub fn data_as_uninit_mut<T: KernelStorageElement>(&mut self) -> Result<&mut [MaybeUninit<T>]> {
594        if self.dtype != T::DTYPE {
595            return Err(StridedError::DTypeMismatch {
596                expected: T::DTYPE.label(),
597                actual: self.dtype.label(),
598            });
599        }
600        Ok(unsafe {
601            core::slice::from_raw_parts_mut(
602                self.data.as_mut_ptr().cast::<MaybeUninit<T>>(),
603                self.data.len() / core::mem::size_of::<T>(),
604            )
605        })
606    }
607
608    #[inline]
609    pub fn dtype(&self) -> KernelDType {
610        self.dtype
611    }
612
613    #[inline]
614    pub fn dims(&self) -> &'a [usize] {
615        self.dims
616    }
617
618    #[inline]
619    pub fn strides(&self) -> &'a [isize] {
620        self.strides
621    }
622
623    #[inline]
624    pub fn offset(&self) -> isize {
625        self.offset
626    }
627}
628
629fn ranges_overlap(a_start: usize, a_len: usize, b_start: usize, b_len: usize) -> Result<bool> {
630    let a_end = a_start
631        .checked_add(a_len)
632        .ok_or(StridedError::OffsetOverflow)?;
633    let b_end = b_start
634        .checked_add(b_len)
635        .ok_or(StridedError::OffsetOverflow)?;
636    Ok(a_start < b_end && b_start < a_end)
637}
638
639/// Borrowed raw strided input layout.
640///
641/// Use [`RawStridedRef::new`] for checked construction, or
642/// [`RawStridedRef::new_unchecked`] when a higher-level compiled plan has
643/// already validated bounds.
644#[derive(Clone, Copy, Debug)]
645pub struct RawStridedRef<'a, T> {
646    data: &'a [T],
647    dims: &'a [usize],
648    strides: &'a [isize],
649    offset: isize,
650}
651
652impl<'a, T> RawStridedRef<'a, T> {
653    /// Create a raw strided input after validating reachable offsets.
654    pub fn new(
655        data: &'a [T],
656        dims: &'a [usize],
657        strides: &'a [isize],
658        offset: isize,
659    ) -> Result<Self> {
660        validate_bounds(data.len(), dims, strides, offset)?;
661        Ok(Self {
662            data,
663            dims,
664            strides,
665            offset,
666        })
667    }
668
669    /// Create a raw strided input without bounds checking.
670    ///
671    /// # Safety
672    /// The caller must ensure every index reachable by `dims`/`strides` from
673    /// `offset` lies inside `data`.
674    pub unsafe fn new_unchecked(
675        data: &'a [T],
676        dims: &'a [usize],
677        strides: &'a [isize],
678        offset: isize,
679    ) -> Self {
680        Self {
681            data,
682            dims,
683            strides,
684            offset,
685        }
686    }
687
688    #[inline]
689    pub fn data(&self) -> &'a [T] {
690        self.data
691    }
692
693    #[inline]
694    pub fn dims(&self) -> &'a [usize] {
695        self.dims
696    }
697
698    #[inline]
699    pub fn strides(&self) -> &'a [isize] {
700        self.strides
701    }
702
703    #[inline]
704    pub fn offset(&self) -> isize {
705        self.offset
706    }
707
708    #[inline]
709    pub fn ptr(&self) -> *const T {
710        if self.dims.iter().any(|&dim| dim == 0) {
711            NonNull::<T>::dangling().as_ptr()
712        } else {
713            unsafe { self.data.as_ptr().offset(self.offset) }
714        }
715    }
716
717    /// Convert to an immutable owning-metadata view.
718    ///
719    /// This is for compatibility paths. Hot prepared paths should use the raw
720    /// accessors directly and avoid this conversion.
721    #[inline]
722    pub fn as_view(&self) -> StridedView<'a, T, Identity> {
723        unsafe { StridedView::new_unchecked(self.data, self.dims, self.strides, self.offset) }
724    }
725}
726
727/// Borrowed raw strided output layout.
728///
729/// This is the mutable counterpart to [`RawStridedRef`]. It avoids allocating
730/// owned shape/stride metadata in prepared replay paths.
731#[derive(Debug)]
732pub struct RawStridedMut<'a, T> {
733    data: &'a mut [T],
734    dims: &'a [usize],
735    strides: &'a [isize],
736    offset: isize,
737}
738
739impl<'a, T> RawStridedMut<'a, T> {
740    /// Create a raw strided output after validating reachable offsets.
741    pub fn new(
742        data: &'a mut [T],
743        dims: &'a [usize],
744        strides: &'a [isize],
745        offset: isize,
746    ) -> Result<Self> {
747        validate_bounds(data.len(), dims, strides, offset)?;
748        Ok(Self {
749            data,
750            dims,
751            strides,
752            offset,
753        })
754    }
755
756    /// Create a raw strided output without bounds checking.
757    ///
758    /// # Safety
759    /// The caller must ensure every index reachable by `dims`/`strides` from
760    /// `offset` lies inside `data`, and no aliases violate mutable access.
761    pub unsafe fn new_unchecked(
762        data: &'a mut [T],
763        dims: &'a [usize],
764        strides: &'a [isize],
765        offset: isize,
766    ) -> Self {
767        Self {
768            data,
769            dims,
770            strides,
771            offset,
772        }
773    }
774
775    #[inline]
776    pub fn data(&self) -> &[T] {
777        self.data
778    }
779
780    #[inline]
781    pub fn data_mut(&mut self) -> &mut [T] {
782        self.data
783    }
784
785    #[inline]
786    pub fn dims(&self) -> &'a [usize] {
787        self.dims
788    }
789
790    #[inline]
791    pub fn strides(&self) -> &'a [isize] {
792        self.strides
793    }
794
795    #[inline]
796    pub fn offset(&self) -> isize {
797        self.offset
798    }
799
800    #[inline]
801    pub fn ptr(&self) -> *const T {
802        if self.dims.iter().any(|&dim| dim == 0) {
803            NonNull::<T>::dangling().as_ptr()
804        } else {
805            unsafe { self.data.as_ptr().offset(self.offset) }
806        }
807    }
808
809    #[inline]
810    pub fn as_mut_ptr(&mut self) -> *mut T {
811        if self.dims.iter().any(|&dim| dim == 0) {
812            NonNull::<T>::dangling().as_ptr()
813        } else {
814            unsafe { self.data.as_mut_ptr().offset(self.offset) }
815        }
816    }
817
818    /// Convert to an immutable owning-metadata view.
819    ///
820    /// This is for compatibility paths. Hot prepared paths should use the raw
821    /// accessors directly and avoid this conversion.
822    #[inline]
823    pub fn as_view(&self) -> StridedView<'_, T, Identity> {
824        unsafe { StridedView::new_unchecked(self.data, self.dims, self.strides, self.offset) }
825    }
826
827    /// Convert to a mutable owning-metadata view.
828    ///
829    /// This is for compatibility paths. Hot prepared paths should use the raw
830    /// accessors directly and avoid this conversion.
831    #[inline]
832    pub fn as_view_mut(&mut self) -> StridedViewMut<'_, T> {
833        unsafe { StridedViewMut::new_unchecked(self.data, self.dims, self.strides, self.offset) }
834    }
835}
836
837#[cfg(test)]
838mod tests {
839    use super::*;
840
841    #[test]
842    fn typed_erased_storage_covers_all_kernel_dtypes() {
843        macro_rules! check {
844            ($ty:ty, $value:expr) => {{
845                let mut initialized = [$value, $value];
846                let dims = [2usize];
847                let strides = [1isize];
848                let reference =
849                    ErasedRawStridedRef::from_slice(&initialized, &dims, &strides, 0).unwrap();
850                assert_eq!(reference.data_as::<$ty>().unwrap().len(), 2);
851                let mut mutable =
852                    ErasedRawStridedMut::from_slice_mut(&mut initialized, &dims, &strides, 0)
853                        .unwrap();
854                assert_eq!(mutable.data_as::<$ty>().unwrap().len(), 2);
855                assert_eq!(mutable.data_as_mut::<$ty>().unwrap().len(), 2);
856                let mut uninitialized = [MaybeUninit::<$ty>::uninit(); 2];
857                let mut destination = ErasedRawStridedUninitMut::from_uninit_slice(
858                    &mut uninitialized,
859                    &dims,
860                    &strides,
861                    0,
862                )
863                .unwrap();
864                assert_eq!(destination.data_as_uninit_mut::<$ty>().unwrap().len(), 2);
865            }};
866        }
867
868        check!(f32, 1.0f32);
869        check!(f64, 1.0f64);
870        check!(i32, 1i32);
871        check!(i64, 1i64);
872        check!(bool, true);
873        check!(Complex32, Complex32::new(1.0, 0.0));
874        check!(Complex64, Complex64::new(1.0, 0.0));
875    }
876
877    #[test]
878    fn rejects_usize_max_dimension_before_isize_truncation() {
879        let data = [false; 1];
880        let dims = [usize::MAX];
881        let strides = [-1isize];
882        assert!(matches!(
883            RawStridedRef::new(&data, &dims, &strides, 0),
884            Err(crate::StridedError::OffsetOverflow)
885        ));
886        let mut data = [false; 1];
887        assert!(matches!(
888            RawStridedMut::new(&mut data, &dims, &strides, 0),
889            Err(crate::StridedError::OffsetOverflow)
890        ));
891        assert!(matches!(
892            ErasedRawStridedRef::from_slice(&data, &dims, &strides, 0),
893            Err(crate::StridedError::OffsetOverflow)
894        ));
895        let mut data = [false; 1];
896        assert!(matches!(
897            ErasedRawStridedMut::from_slice_mut(&mut data, &dims, &strides, 0),
898            Err(crate::StridedError::OffsetOverflow)
899        ));
900        let ptr = NonNull::new(data.as_mut_ptr()).unwrap();
901        assert!(matches!(
902            unsafe {
903                ErasedRawStridedPtr::from_raw_parts(
904                    KernelDType::Bool,
905                    ptr.cast(),
906                    data.len(),
907                    &dims,
908                    &strides,
909                    0,
910                )
911            },
912            Err(crate::StridedError::OffsetOverflow)
913        ));
914        let mut data = vec![MaybeUninit::<bool>::uninit(); 1];
915        assert!(matches!(
916            ErasedRawStridedUninitMut::from_uninit_slice(&mut data, &dims, &strides, 0),
917            Err(crate::StridedError::OffsetOverflow)
918        ));
919    }
920
921    #[test]
922    fn empty_raw_views_use_dangling_base_pointers_for_extreme_offsets() {
923        let dims = [0usize];
924        let strides = [1isize];
925        let data: [f64; 0] = [];
926        let raw = RawStridedRef::new(&data, &dims, &strides, isize::MAX).unwrap();
927        assert_eq!(raw.ptr(), NonNull::<f64>::dangling().as_ptr());
928
929        let mut data: [f64; 0] = [];
930        let mut raw = RawStridedMut::new(&mut data, &dims, &strides, isize::MAX).unwrap();
931        assert_eq!(raw.ptr(), NonNull::<f64>::dangling().as_ptr());
932        assert_eq!(raw.as_mut_ptr(), NonNull::<f64>::dangling().as_ptr());
933    }
934
935    #[test]
936    fn raw_ref_rejects_out_of_bounds_layout() {
937        let data = [0.0f64; 4];
938        let err = RawStridedRef::new(&data, &[2, 3], &[3, 1], 0).unwrap_err();
939        assert!(matches!(err, crate::StridedError::OffsetOverflow));
940    }
941
942    #[test]
943    fn raw_mut_can_reborrow_as_view() {
944        let mut data = [1, 2, 3, 4];
945        let mut raw = RawStridedMut::new(&mut data, &[2, 2], &[2, 1], 0).unwrap();
946        {
947            let view = raw.as_view();
948            assert_eq!(view.dims(), &[2, 2]);
949        }
950        let view_mut = raw.as_view_mut();
951        assert_eq!(view_mut.get(&[1, 1]), 4);
952    }
953}