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 crate::{Result, StridedView, StridedViewMut};
11
12/// Borrowed raw strided input layout.
13///
14/// Use [`RawStridedRef::new`] for checked construction, or
15/// [`RawStridedRef::new_unchecked`] when a higher-level compiled plan has
16/// already validated bounds.
17#[derive(Clone, Copy, Debug)]
18pub struct RawStridedRef<'a, T> {
19    data: &'a [T],
20    dims: &'a [usize],
21    strides: &'a [isize],
22    offset: isize,
23}
24
25impl<'a, T> RawStridedRef<'a, T> {
26    /// Create a raw strided input after validating reachable offsets.
27    pub fn new(
28        data: &'a [T],
29        dims: &'a [usize],
30        strides: &'a [isize],
31        offset: isize,
32    ) -> Result<Self> {
33        validate_bounds(data.len(), dims, strides, offset)?;
34        Ok(Self {
35            data,
36            dims,
37            strides,
38            offset,
39        })
40    }
41
42    /// Create a raw strided input without bounds checking.
43    ///
44    /// # Safety
45    /// The caller must ensure every index reachable by `dims`/`strides` from
46    /// `offset` lies inside `data`.
47    pub unsafe fn new_unchecked(
48        data: &'a [T],
49        dims: &'a [usize],
50        strides: &'a [isize],
51        offset: isize,
52    ) -> Self {
53        Self {
54            data,
55            dims,
56            strides,
57            offset,
58        }
59    }
60
61    #[inline]
62    pub fn data(&self) -> &'a [T] {
63        self.data
64    }
65
66    #[inline]
67    pub fn dims(&self) -> &'a [usize] {
68        self.dims
69    }
70
71    #[inline]
72    pub fn strides(&self) -> &'a [isize] {
73        self.strides
74    }
75
76    #[inline]
77    pub fn offset(&self) -> isize {
78        self.offset
79    }
80
81    #[inline]
82    pub fn ptr(&self) -> *const T {
83        unsafe { self.data.as_ptr().offset(self.offset) }
84    }
85
86    /// Convert to an immutable owning-metadata view.
87    ///
88    /// This is for compatibility paths. Hot prepared paths should use the raw
89    /// accessors directly and avoid this conversion.
90    #[inline]
91    pub fn as_view(&self) -> StridedView<'a, T, Identity> {
92        unsafe { StridedView::new_unchecked(self.data, self.dims, self.strides, self.offset) }
93    }
94}
95
96/// Borrowed raw strided output layout.
97///
98/// This is the mutable counterpart to [`RawStridedRef`]. It avoids allocating
99/// owned shape/stride metadata in prepared replay paths.
100#[derive(Debug)]
101pub struct RawStridedMut<'a, T> {
102    data: &'a mut [T],
103    dims: &'a [usize],
104    strides: &'a [isize],
105    offset: isize,
106}
107
108impl<'a, T> RawStridedMut<'a, T> {
109    /// Create a raw strided output after validating reachable offsets.
110    pub fn new(
111        data: &'a mut [T],
112        dims: &'a [usize],
113        strides: &'a [isize],
114        offset: isize,
115    ) -> Result<Self> {
116        validate_bounds(data.len(), dims, strides, offset)?;
117        Ok(Self {
118            data,
119            dims,
120            strides,
121            offset,
122        })
123    }
124
125    /// Create a raw strided output without bounds checking.
126    ///
127    /// # Safety
128    /// The caller must ensure every index reachable by `dims`/`strides` from
129    /// `offset` lies inside `data`, and no aliases violate mutable access.
130    pub unsafe fn new_unchecked(
131        data: &'a mut [T],
132        dims: &'a [usize],
133        strides: &'a [isize],
134        offset: isize,
135    ) -> Self {
136        Self {
137            data,
138            dims,
139            strides,
140            offset,
141        }
142    }
143
144    #[inline]
145    pub fn data(&self) -> &[T] {
146        self.data
147    }
148
149    #[inline]
150    pub fn data_mut(&mut self) -> &mut [T] {
151        self.data
152    }
153
154    #[inline]
155    pub fn dims(&self) -> &'a [usize] {
156        self.dims
157    }
158
159    #[inline]
160    pub fn strides(&self) -> &'a [isize] {
161        self.strides
162    }
163
164    #[inline]
165    pub fn offset(&self) -> isize {
166        self.offset
167    }
168
169    #[inline]
170    pub fn ptr(&self) -> *const T {
171        unsafe { self.data.as_ptr().offset(self.offset) }
172    }
173
174    #[inline]
175    pub fn as_mut_ptr(&mut self) -> *mut T {
176        unsafe { self.data.as_mut_ptr().offset(self.offset) }
177    }
178
179    /// Convert to an immutable owning-metadata view.
180    ///
181    /// This is for compatibility paths. Hot prepared paths should use the raw
182    /// accessors directly and avoid this conversion.
183    #[inline]
184    pub fn as_view(&self) -> StridedView<'_, T, Identity> {
185        unsafe { StridedView::new_unchecked(self.data, self.dims, self.strides, self.offset) }
186    }
187
188    /// Convert to a mutable owning-metadata view.
189    ///
190    /// This is for compatibility paths. Hot prepared paths should use the raw
191    /// accessors directly and avoid this conversion.
192    #[inline]
193    pub fn as_view_mut(&mut self) -> StridedViewMut<'_, T> {
194        unsafe { StridedViewMut::new_unchecked(self.data, self.dims, self.strides, self.offset) }
195    }
196}
197
198#[cfg(test)]
199mod tests {
200    use super::*;
201
202    #[test]
203    fn raw_ref_rejects_out_of_bounds_layout() {
204        let data = [0.0f64; 4];
205        let err = RawStridedRef::new(&data, &[2, 3], &[3, 1], 0).unwrap_err();
206        assert!(matches!(err, crate::StridedError::OffsetOverflow));
207    }
208
209    #[test]
210    fn raw_mut_can_reborrow_as_view() {
211        let mut data = [1, 2, 3, 4];
212        let mut raw = RawStridedMut::new(&mut data, &[2, 2], &[2, 1], 0).unwrap();
213        {
214            let view = raw.as_view();
215            assert_eq!(view.dims(), &[2, 2]);
216        }
217        let view_mut = raw.as_view_mut();
218        assert_eq!(view_mut.get(&[1, 1]), 4);
219    }
220}
strided_view/raw.rs

strided_view/
raw.rs
