cubecl_tile/layout/projection/
dims_builder.rs

1//! Writing a [`Projection`] one buffer dim at a time.
2//!
3//! ```text
4//! Projection::dims()
5//!     .dim(B)                                        // dim 0 = b
6//!     .dim(stencil(&[(OH, stride), (RH, dilation)]).pad(pad)) // dim 1 = oh*stride+rh*dilation-pad
7//!     .dim(C)                                        // dim 2 = c
8//!     .build()
9//! ```
10
11use cubecl::zspace::SmallVec;
12
13#[cfg(test)]
14use crate::Composition;
15use crate::{Axis, PhysicalAxisMap, Projection, Space};
16
17/// A [`Projection`] under construction.
18#[derive(Clone, Debug)]
19pub struct DimsBuilder {
20    physical: SmallVec<[PhysicalAxisMap; Space::MAX_RANK]>,
21    spanning: SmallVec<[Axis; Space::MAX_RANK]>,
22}
23
24impl DimsBuilder {
25    /// No dim stated yet.
26    pub(crate) fn new() -> Self {
27        DimsBuilder {
28            physical: SmallVec::new(),
29            spanning: SmallVec::new(),
30        }
31    }
32
33    /// The next buffer dim, coarsest first: an [`Axis`] or a [`split`].
34    pub fn dim(mut self, dim: impl Into<PhysicalAxisMap>) -> Self {
35        self.physical.push(dim.into());
36        self
37    }
38
39    /// An axis this operand is defined over but addressed by no dim.
40    pub fn spanning(mut self, axis: Axis) -> Self {
41        self.spanning.push(axis);
42        self
43    }
44
45    /// The projection, its logical axes derived from the dims, the innermost dim's last.
46    /// Panics on no dims, or on an axis both [`spanning`](Self::spanning) and addressed.
47    pub fn build(self) -> Projection {
48        assert!(
49            !self.physical.is_empty(),
50            "Projection::dims: an operand has at least one dim"
51        );
52        let (innermost, outer) = self.physical.split_last().expect("checked non-empty");
53        let mut axes: SmallVec<[Axis; Space::MAX_RANK]> = SmallVec::new();
54        let mention = |axis: Axis, axes: &mut SmallVec<[Axis; Space::MAX_RANK]>| {
55            if !axes.contains(&axis) {
56                axes.push(axis);
57            }
58        };
59        for map in outer {
60            for term in map.terms() {
61                mention(term.axis, &mut axes);
62            }
63        }
64        for &axis in &self.spanning {
65            assert!(
66                !self.physical.iter().any(|map| map.addresses(axis)),
67                "Projection::dims: {axis:?} is stated as spanning (addressed by no dim) but a dim \
68                 addresses it"
69            );
70            mention(axis, &mut axes);
71        }
72        // Last, so the line runs along it; a storage-tiled axis keeps its earlier position.
73        for term in innermost.terms() {
74            mention(term.axis, &mut axes);
75        }
76        Projection::new(&axes, &self.physical)
77    }
78}
79
80/// A dim that *is* one coordinate: `.dim(K)`.
81impl From<Axis> for PhysicalAxisMap {
82    fn from(axis: Axis) -> Self {
83        PhysicalAxisMap::of(axis)
84    }
85}
86
87/// One dim cut into blocks, stated in extents, coarsest first:
88/// `split(&[(KB, blocks), (KI, block)])` is `kb·block + ki`. Panics on no terms or extent 0.
89pub fn split(extents: &[(Axis, usize)]) -> PhysicalAxisMap {
90    assert!(
91        !extents.is_empty(),
92        "split: a dim is cut by at least one axis"
93    );
94    let mut terms: SmallVec<[(Axis, usize); Space::MAX_RANK]> = SmallVec::new();
95    let mut coefficient = 1;
96    for &(axis, extent) in extents.iter().rev() {
97        assert!(extent > 0, "split: {axis:?} has extent 0");
98        terms.push((axis, coefficient));
99        coefficient *= extent;
100    }
101    terms.reverse();
102    PhysicalAxisMap::disjoint(&terms)
103}
104
105/// One dim several axes slide over, in coefficients: `oh·stride + rh·dilation`.
106#[cfg(test)]
107pub(crate) fn stencil(coefficients: &[(Axis, usize)]) -> StencilDim {
108    StencilDim {
109        coefficients: SmallVec::from_slice(coefficients),
110        pad: 0,
111    }
112}
113
114/// A [`stencil`] before its padding is stated.
115#[cfg(test)]
116#[derive(Clone, Debug)]
117pub(crate) struct StencilDim {
118    coefficients: SmallVec<[(Axis, usize); Space::MAX_RANK]>,
119    pad: usize,
120}
121
122#[cfg(test)]
123impl StencilDim {
124    /// How many cells before the buffer's first the window's origin sits.
125    pub(crate) fn pad(mut self, pad: usize) -> Self {
126        self.pad = pad;
127        self
128    }
129}
130
131#[cfg(test)]
132impl From<StencilDim> for PhysicalAxisMap {
133    fn from(window: StencilDim) -> Self {
134        let map = PhysicalAxisMap::affine(&window.coefficients).shifted(-(window.pad as isize));
135        debug_assert!(
136            window.coefficients.len() == 1 || map.composition() == Composition::Overlapping,
137            "a window over several axes takes the overlapping reading"
138        );
139        map
140    }
141}
142
143#[cfg(test)]
144mod tests {
145    use super::*;
146    use crate::Offset;
147
148    const B: Axis = Axis(0);
149    const H: Axis = Axis(1);
150    const G: Axis = Axis(2);
151    const T: Axis = Axis(3);
152    const D: Axis = Axis(4);
153    const K: Axis = Axis(5);
154    const N: Axis = Axis(6);
155    const KB: Axis = Axis(7);
156    const KI: Axis = Axis(8);
157    const OH: Axis = Axis(9);
158    const RH: Axis = Axis(10);
159    const C: Axis = Axis(11);
160    const M: Axis = Axis(12);
161
162    /// One dim's index for a coordinate: `Σ axis·coefficient + offset`.
163    fn index(p: &Projection, dim: usize, coordinate: &[(Axis, usize)]) -> isize {
164        let sum: usize = coordinate
165            .iter()
166            .map(|&(axis, at)| at * p.scale(dim, axis))
167            .sum();
168        let offset = match p.offset(dim) {
169            Offset::Static(o) => o,
170            Offset::Dynamic => unreachable!("nothing here builds a dynamic offset"),
171        };
172        sum as isize + offset
173    }
174
175    #[test]
176    fn one_axis_per_dim_is_direct() {
177        let p = Projection::dims().dim(K).dim(N).build();
178        assert_eq!(p, Projection::direct(&[K, N]));
179    }
180
181    /// `kb·32 + ki` as 4 blocks of 32: `k = 100` lands at block 3, position 4.
182    #[test]
183    fn split_derives_the_coefficients_from_extents() {
184        let p = Projection::dims()
185            .dim(N)
186            .dim(split(&[(KB, 4), (KI, 32)]))
187            .build();
188
189        assert_eq!(p.scale(1, KB), 32);
190        assert_eq!(p.scale(1, KI), 1);
191        assert_eq!(index(&p, 1, &[(KB, 3), (KI, 4)]), 100);
192        assert_eq!(p.composition(), Composition::Disjoint);
193        assert_eq!(p.logical_axes(), &[N, KB, KI]);
194    }
195
196    #[test]
197    fn split_composes_more_than_two_axes() {
198        let p = Projection::dims()
199            .dim(split(&[(B, 2), (H, 3), (D, 5)]))
200            .build();
201        assert_eq!(p.scale(0, B), 15);
202        assert_eq!(p.scale(0, H), 5);
203        assert_eq!(p.scale(0, D), 1);
204        assert_eq!(index(&p, 0, &[(B, 1), (H, 2), (D, 4)]), 29);
205    }
206
207    /// Split-merge partials, group-major: `g·splits + t`.
208    #[test]
209    fn a_partition_in_the_middle_keeps_the_innermost_last() {
210        let (group, splits) = (4, 8);
211        let p = Projection::dims()
212            .dim(B)
213            .dim(H)
214            .dim(split(&[(G, group), (T, splits)]))
215            .dim(D)
216            .build();
217
218        assert_eq!(index(&p, 2, &[(G, 2), (T, 5)]), 21);
219        assert_eq!(index(&p, 2, &[(G, 2), (T, 6)]), 22);
220        assert_eq!(p.logical_axes(), &[B, H, G, T, D]);
221    }
222
223    /// `oh·2 + rh·1 − 1`: adjacent output steps overlap, and `oh = 0` reaches into the pad.
224    #[test]
225    fn stencil_is_an_overlapping_affine_map_with_padding() {
226        let p = Projection::dims()
227            .dim(B)
228            .dim(stencil(&[(OH, 2), (RH, 1)]).pad(1))
229            .dim(C)
230            .build();
231
232        assert_eq!(index(&p, 1, &[(OH, 3), (RH, 2)]), 7);
233        assert_eq!(index(&p, 1, &[(OH, 4), (RH, 0)]), 7);
234        assert_eq!(index(&p, 1, &[(OH, 0), (RH, 0)]), -1);
235        assert_eq!(p.composition(), Composition::Overlapping);
236        assert_eq!(p.logical_axes(), &[B, OH, RH, C]);
237    }
238
239    /// A grouped-query cache: `g` is addressed by no dim and stays before the innermost axis.
240    #[test]
241    fn a_spanning_axis_sits_before_the_innermost() {
242        let p = Projection::dims()
243            .dim(B)
244            .dim(H)
245            .dim(T)
246            .dim(D)
247            .spanning(G)
248            .build();
249
250        assert!(!p.addresses(G));
251        assert_eq!(p.logical_axes(), &[B, H, T, G, D]);
252        p.validate(4).unwrap();
253    }
254
255    #[test]
256    fn a_repeated_axis_is_storage_tiled() {
257        let p = Projection::dims()
258            .dim(B)
259            .dim(M)
260            .dim(K)
261            .dim(M)
262            .dim(K)
263            .dim(C)
264            .build();
265
266        assert_eq!(p.logical_axes(), &[B, M, K, C]);
267        assert_eq!(p.tiling().fragments(4).as_slice(), &[1, 2, 2, 1]);
268    }
269
270    #[test]
271    #[should_panic(expected = "stated as spanning")]
272    fn spanning_an_addressed_axis_is_refused() {
273        let _ = Projection::dims().dim(B).dim(D).spanning(D).build();
274    }
275}
cubecl_tile/layout/projection/dims_builder.rs

cubecl_tile/layout/projection/
dims_builder.rs
