subetha_pointers/
kstep_pointer.rs

1//! `KStepPointer<T>` - pointer with first-class stride encoding.
2//!
3//! Storage: `(base: *const T, k_step: u8)`. The stride between
4//! consecutive elements is `sizeof(T) << k_step`. So:
5//!
6//! | K_step | Stride                     | Use case                       |
7//! |--------|----------------------------|--------------------------------|
8//! | 0      | sizeof(T) (tight pack)     | Default, contiguous Vec/array  |
9//! | 1      | 2 * sizeof(T)              | Every other element            |
10//! | 2      | 4 * sizeof(T)              | Sub-quarter access pattern     |
11//! | 3      | 8 * sizeof(T)              | SIMD lane stride               |
12//! | 6      | 64 * sizeof(T)             | Cache-line stride              |
13//! | 12     | 4096 * sizeof(T)           | Page-aligned stride            |
14//!
15//! K_step is the pointer-side analog of quartz's `K_inner` axis -
16//! it controls the granularity of iteration. The advantage over a
17//! runtime `stride: usize` is that K_step is a const-encoded shift
18//! amount; the compiler can fold `<< k_step` into address generation
19//! and SIMD ops know the stride at codegen time.
20//!
21//! # Architectural rationale
22//!
23//! BLAS GEMM iterates over matrix rows AND columns with potentially
24//! different strides. NumPy's strided arrays do the same in higher
25//! dimensions. Today these are all encoded as runtime `stride: usize`
26//! fields - the compiler has to emit IMUL for each step. With KStep
27//! the stride is `1 << k_step` so the codegen is SHL (one cycle),
28//! and the compiler can hoist the shift amount as an immediate.
29
30use std::marker::PhantomData;
31
32/// Strided pointer to `T`. Stride = `sizeof(T) << k_step` bytes.
33#[derive(Debug)]
34pub struct KStepPointer<T> {
35    base: *const T,
36    k_step: u8,
37    _phantom: PhantomData<*const T>,
38}
39
40unsafe impl<T: Send> Send for KStepPointer<T> {}
41unsafe impl<T: Sync> Sync for KStepPointer<T> {}
42
43impl<T> KStepPointer<T> {
44    /// Direction signature of `KStepPointer<T>`. Engages the
45    /// `K_stride` axis (stride encoded as log2 shift count, no
46    /// runtime `stride: usize` field).
47    pub const SIGNATURE: subetha_core::AxisMask = subetha_core::AxisMask::from_axes(
48        &[subetha_core::Axis::Stride],
49    );
50
51    /// New strided pointer. `k_step` is the log2 of the stride
52    /// multiplier; the actual byte stride is `sizeof(T) << k_step`.
53    ///
54    /// # Safety
55    ///
56    /// `base` must point to a valid `T` and the strided sequence
57    /// `base + i * stride` must remain valid for all indices the
58    /// caller will use.
59    pub const unsafe fn new(base: *const T, k_step: u8) -> Self {
60        Self { base, k_step, _phantom: PhantomData }
61    }
62
63    /// Tight-packed strided pointer (K_step = 0).
64    ///
65    /// # Safety
66    ///
67    /// Same as [`Self::new`]: `base` must point to a contiguous run of `T`
68    /// for every index the caller will use.
69    pub const unsafe fn tight(base: *const T) -> Self {
70        unsafe { Self::new(base, 0) }
71    }
72
73    /// Cache-line strided pointer (K_step chosen so stride >= 64
74    /// bytes). Walks every 64/sizeof(T) elements.
75    ///
76    /// # Safety
77    ///
78    /// Same as [`Self::new`]: `base` must remain valid for every index
79    /// the caller will use under the chosen `k_step` stride.
80    pub const unsafe fn cache_line(base: *const T) -> Self {
81        // ceil(log2(64 / sizeof(T))) capped at 6.
82        let t_size = std::mem::size_of::<T>();
83        let k = if t_size >= 64 { 0 }
84            else if t_size >= 32 { 1 }
85            else if t_size >= 16 { 2 }
86            else if t_size >= 8 { 3 }
87            else if t_size >= 4 { 4 }
88            else if t_size >= 2 { 5 }
89            else { 6 };
90        unsafe { Self::new(base, k) }
91    }
92
93    #[inline]
94    pub const fn base(&self) -> *const T { self.base }
95
96    #[inline]
97    pub const fn k_step(&self) -> u8 { self.k_step }
98
99    /// Byte stride between consecutive elements.
100    #[inline]
101    pub const fn stride(&self) -> usize {
102        std::mem::size_of::<T>() << self.k_step
103    }
104
105    /// Pointer to the i-th element under the current stride.
106    ///
107    /// # Safety
108    ///
109    /// The caller must ensure `base + i * stride` is in bounds.
110    #[inline]
111    pub unsafe fn at(&self, i: usize) -> *const T {
112        let offset_bytes = i * self.stride();
113        unsafe { (self.base as *const u8).add(offset_bytes) as *const T }
114    }
115
116    /// Borrow the i-th element.
117    ///
118    /// # Safety
119    ///
120    /// `i` must be in bounds; the strided sequence at this index
121    /// must point to a valid `T`.
122    #[inline]
123    pub unsafe fn get(&self, i: usize) -> &T {
124        unsafe { &*self.at(i) }
125    }
126
127    /// Iterator over `count` strided elements.
128    ///
129    /// # Safety
130    ///
131    /// All `count` strided positions must be in bounds.
132    pub unsafe fn iter(&self, count: usize) -> StridedIter<'_, T> {
133        StridedIter { ptr: *self, i: 0, count, _life: PhantomData }
134    }
135}
136
137impl<T> Clone for KStepPointer<T> {
138    fn clone(&self) -> Self { *self }
139}
140impl<T> Copy for KStepPointer<T> {}
141
142/// Iterator yielding strided elements.
143pub struct StridedIter<'a, T> {
144    ptr: KStepPointer<T>,
145    i: usize,
146    count: usize,
147    _life: PhantomData<&'a T>,
148}
149
150impl<'a, T> Iterator for StridedIter<'a, T> {
151    type Item = &'a T;
152    fn next(&mut self) -> Option<&'a T> {
153        if self.i >= self.count { return None; }
154        let r = unsafe { &*self.ptr.at(self.i) };
155        self.i += 1;
156        Some(r)
157    }
158    fn size_hint(&self) -> (usize, Option<usize>) {
159        let n = self.count - self.i;
160        (n, Some(n))
161    }
162}
163
164impl<'a, T> ExactSizeIterator for StridedIter<'a, T> {}
165
166#[cfg(test)]
167mod tests {
168    use super::*;
169
170    #[test]
171    fn tight_stride_walks_contiguous_array() {
172        let data: Vec<u64> = (0..10u64).collect();
173        let p = unsafe { KStepPointer::tight(data.as_ptr()) };
174        assert_eq!(p.k_step(), 0);
175        assert_eq!(p.stride(), 8);
176        for i in 0..10 {
177            assert_eq!(unsafe { *p.get(i) }, i as u64);
178        }
179    }
180
181    #[test]
182    fn k_step_1_skips_every_other() {
183        let data: Vec<u64> = (0..10u64).collect();
184        let p = unsafe { KStepPointer::new(data.as_ptr(), 1) };
185        assert_eq!(p.stride(), 16);
186        // i=0 -> data[0]=0; i=1 -> data[2]=2; i=2 -> data[4]=4; ...
187        for i in 0..5 {
188            assert_eq!(unsafe { *p.get(i) }, (i * 2) as u64);
189        }
190    }
191
192    #[test]
193    fn k_step_2_skips_by_four() {
194        let data: Vec<u64> = (0..16u64).collect();
195        let p = unsafe { KStepPointer::new(data.as_ptr(), 2) };
196        assert_eq!(p.stride(), 32);
197        for i in 0..4 {
198            assert_eq!(unsafe { *p.get(i) }, (i * 4) as u64);
199        }
200    }
201
202    #[test]
203    fn cache_line_stride_picks_correct_k() {
204        // For u64 (8 bytes) cache_line() should pick k=3 -> stride 64.
205        let data: Vec<u64> = (0..64u64).collect();
206        let p = unsafe { KStepPointer::cache_line(data.as_ptr()) };
207        assert_eq!(p.k_step(), 3, "u64 cache_line should be k=3");
208        assert_eq!(p.stride(), 64);
209        // i=0 -> data[0]; i=1 -> data[8]; i=2 -> data[16]; ...
210        for i in 0..8 {
211            assert_eq!(unsafe { *p.get(i) }, (i * 8) as u64);
212        }
213    }
214
215    #[test]
216    fn cache_line_stride_for_u8() {
217        // For u8 (1 byte) cache_line should pick k=6 -> stride 64.
218        let data: Vec<u8> = (0u8..64).collect();
219        let p = unsafe { KStepPointer::<u8>::cache_line(data.as_ptr()) };
220        assert_eq!(p.k_step(), 6);
221        assert_eq!(p.stride(), 64);
222    }
223
224    #[test]
225    fn strided_iter_yields_correct_elements() {
226        let data: Vec<u64> = (0..20u64).collect();
227        let p = unsafe { KStepPointer::new(data.as_ptr(), 2) };
228        let collected: Vec<u64> = unsafe { p.iter(5) }.copied().collect();
229        assert_eq!(collected, vec![0, 4, 8, 12, 16]);
230    }
231
232    #[test]
233    fn matrix_row_stride_workflow() {
234        // 4x4 matrix in row-major, walking col 0 with k_step=2 (stride 4*8=32).
235        let matrix: Vec<u64> = (0..16u64).collect();
236        let col0 = unsafe { KStepPointer::new(matrix.as_ptr(), 2) };
237        let col0_vals: Vec<u64> = unsafe { col0.iter(4) }.copied().collect();
238        assert_eq!(col0_vals, vec![0, 4, 8, 12]);
239    }
240
241    #[test]
242    fn layout_is_16_bytes() {
243        // *const T = 8, u8 + PhantomData -> pad to 16 with align(8).
244        assert_eq!(std::mem::size_of::<KStepPointer<u64>>(), 16);
245    }
246}
subetha_pointers/kstep_pointer.rs

subetha_pointers/
kstep_pointer.rs
