ruda_test_utils/test_tensor/
strides.rs

1use ruda_kernel::dsl as kernel_dsl;
2use ruda_kernel::dsl::zspace::Shape;
3use ruda_kernel::dsl::zspace::Strides;
4
5#[derive(Debug, PartialEq, Eq, Default)]
6pub enum StrideSpec {
7    #[default]
8    RowMajor,
9    ColMajor,
10    Custom(Vec<usize>),
11}
12
13/// Number of elements in the physical buffer required to cover every logical
14/// index in `shape` under `strides`, assuming element 0 is at offset 0.
15///
16/// Exceeds `shape.iter().product()` for jumpy strides (e.g. a slice stepping
17/// over padding) and is less than it for broadcast strides (a stride of 0
18/// makes every index in that dim share the same physical offset).
19pub fn physical_extent(shape: &Shape, strides: &Strides) -> usize {
20    let mut max_offset = 0usize;
21    for (s, d) in strides.iter().zip(shape.iter()) {
22        if *d > 0 && *s > 0 {
23            max_offset += (d - 1) * s;
24        }
25    }
26    max_offset + 1
27}
28
29#[cfg(test)]
30mod tests {
31    use super::*;
32
33    #[test]
34    fn physical_extent_contiguous_row_major() {
35        // Row-major 2x3 → strides (3, 1) → 6 elements covered.
36        let shape = Shape::from(vec![2, 3]);
37        let strides = Strides::new(&[3, 1]);
38        assert_eq!(physical_extent(&shape, &strides), 6);
39    }
40
41    #[test]
42    fn physical_extent_jumpy_strides_exceed_logical() {
43        // 256x256 logical view of a wider 256x512 buffer (stride 512 on dim 0).
44        // Last reachable offset is 255*512 + 255*1 = 130815 → +1 = 130816.
45        let shape = Shape::from(vec![256, 256]);
46        let strides = Strides::new(&[512, 1]);
47        assert_eq!(physical_extent(&shape, &strides), 130816);
48        // And it strictly exceeds the logical element count.
49        assert!(physical_extent(&shape, &strides) > 256 * 256);
50    }
51
52    #[test]
53    fn physical_extent_broadcast_strides_undercount_logical() {
54        // Broadcast dim: stride 0 means every index along that dim shares the
55        // same physical offset. A 4x3 tensor broadcasting dim 0 only needs 3
56        // elements of physical storage, not 12.
57        let shape = Shape::from(vec![4, 3]);
58        let strides = Strides::new(&[0, 1]);
59        assert_eq!(physical_extent(&shape, &strides), 3);
60        // ...and it's less than the logical element count.
61        assert!(physical_extent(&shape, &strides) < 4 * 3);
62    }
63}
64
65impl StrideSpec {
66    pub fn compute_strides(&self, shape: &Shape) -> Strides {
67        let n = shape.len();
68        match self {
69            StrideSpec::RowMajor => {
70                assert!(n >= 2, "RowMajor requires at least 2 dimensions");
71                let mut strides = vec![0; n];
72                strides[n - 1] = 1;
73                for i in (0..n - 1).rev() {
74                    strides[i] = strides[i + 1] * shape[i + 1];
75                }
76                Strides::new(&strides)
77            }
78            StrideSpec::ColMajor => {
79                assert!(n >= 2, "ColMajor requires at least 2 dimensions");
80                let mut strides = vec![0; n];
81                strides[n - 2] = 1;
82                strides[n - 1] = shape[n - 2];
83                for i in (0..n - 2).rev() {
84                    strides[i] = strides[i + 1] * shape[i + 1];
85                }
86                Strides::new(&strides)
87            }
88            StrideSpec::Custom(strides) => {
89                assert!(
90                    strides.len() == n,
91                    "Custom strides must have the same rank as the shape"
92                );
93                strides.clone().into()
94            }
95        }
96    }
97}
ruda_test_utils/test_tensor/strides.rs

ruda_test_utils/test_tensor/
strides.rs
