cubecl_tile/layout/compaction/
base.rs

1//! The smallest physical box a gathered operand's sub-tile can be staged in, and its mapping.
2
3use cubecl::zspace::SmallVec;
4
5use crate::algebra::gcd;
6use crate::{Axis, PhysicalAxisMap, Projection, Scale, Space};
7
8/// The compacted stage of a [`Projection`]: per physical axis, its cell count and source step,
9/// plus the projection addressing it.
10#[derive(Clone, PartialEq, Eq, Debug)]
11pub struct Compaction {
12    steps: SmallVec<[usize; Space::MAX_RANK]>,
13    extents: SmallVec<[usize; Space::MAX_RANK]>,
14    projection: Projection,
15}
16
17impl Compaction {
18    /// Compact `projection`'s window over the sub-tile whose extents `extent_of` gives.
19    pub fn new(
20        projection: &Projection,
21        vector_size: usize,
22        extent_of: impl Fn(Axis) -> usize,
23    ) -> Compaction {
24        let rank = projection.physical_rank();
25        let mut steps = SmallVec::new();
26        let mut extents = SmallVec::new();
27        let mut physical = Vec::with_capacity(rank);
28
29        for pa in 0..rank {
30            let axis_map = projection.physical_axis(pa);
31            // At offset 0 on the stage, a static divisor may reduce even under a Dynamic offset.
32            let scaled: Vec<(Axis, Scale)> =
33                axis_map.terms().iter().map(|t| (t.axis, t.scale)).collect();
34            let stage_map = PhysicalAxisMap::scaled(&scaled).over(axis_map.divisor());
35            let terms = stage_map.terms();
36
37            if stage_map.is_rational() {
38                // A rational axis has no single step: dense, sized at its widest bounds.
39                let d = stage_map.divisor().bound();
40                let field: usize = terms
41                    .iter()
42                    .filter(|t| extent_of(t.axis) > 1)
43                    .map(|t| (extent_of(t.axis) - 1) * t.scale.bound())
44                    .sum();
45                extents.push(1 + field.div_ceil(d));
46                steps.push(1);
47                physical.push(stage_map);
48            } else {
49                // Only moving terms constrain the step; a moving Dynamic coefficient forces step 1.
50                let step = terms
51                    .iter()
52                    .filter(|t| extent_of(t.axis) > 1)
53                    .try_fold(0, |g, t| match t.scale {
54                        Scale::Static(s) => Some(gcd(g, s)),
55                        Scale::Dynamic { .. } => None,
56                    })
57                    .unwrap_or(1);
58                let step = step.max(1);
59                // Non-moving static terms pin to 1; Dynamic ones keep their coefficient-array slot.
60                let scaled: Vec<(Axis, Scale)> = terms
61                    .iter()
62                    .map(|t| {
63                        let scale = match t.scale {
64                            Scale::Static(s) if extent_of(t.axis) > 1 => Scale::Static(s / step),
65                            Scale::Static(_) => Scale::Static(1),
66                            dynamic => dynamic,
67                        };
68                        (t.axis, scale)
69                    })
70                    .collect();
71                extents.push(
72                    1 + scaled
73                        .iter()
74                        .map(|&(axis, scale)| (extent_of(axis) - 1) * scale.bound())
75                        .sum::<usize>(),
76                );
77                steps.push(step);
78                physical.push(PhysicalAxisMap::scaled(&scaled));
79            }
80        }
81
82        let projection = Projection::new(projection.logical_axes(), &physical);
83        // Compacted at expansion as well as on the host, where no caller takes a refusal.
84        if let Err(refusal) = projection.validate(vector_size) {
85            panic!("{refusal}");
86        }
87
88        Compaction {
89            steps,
90            extents,
91            projection,
92        }
93    }
94
95    /// How the stage's own logical axes address its cells.
96    pub fn projection(&self) -> &Projection {
97        &self.projection
98    }
99
100    /// The step per physical axis: stage coordinate times step is the source coordinate.
101    pub fn steps(&self) -> &[usize] {
102        &self.steps
103    }
104
105    /// The cell count per physical axis, in elements.
106    pub fn extents(&self) -> &[usize] {
107        &self.extents
108    }
109
110    /// Whether the window has no holes, so a fill reads the source box straight through.
111    pub(crate) fn is_dense(&self) -> bool {
112        self.steps.iter().all(|&g| g == 1)
113    }
114
115    /// [`extents`](Compaction::extents) with the innermost in `vector_size`-wide lines.
116    pub(crate) fn line_extents(&self, vector_size: usize) -> Vec<usize> {
117        let last = self.extents.len() - 1;
118        assert!(
119            self.steps[last] == 1,
120            "Compaction: the innermost physical axis is addressed in {vector_size}-wide lines, so \
121             it must be an ungathered axis (extent {}, step {})",
122            self.extents[last],
123            self.steps[last]
124        );
125        let mut lines: Vec<usize> = self.extents.to_vec();
126        lines[last] = lines[last].div_ceil(vector_size);
127        lines
128    }
129
130    /// How many cells the stage allocates, in lines.
131    pub fn cells(&self, vector_size: usize) -> usize {
132        self.line_extents(vector_size).iter().product()
133    }
134}
135
136#[cfg(test)]
137mod tests {
138    use super::*;
139    use crate::{Divisor, Offset};
140
141    const OH: Axis = Axis(0);
142    const RH: Axis = Axis(1);
143    const CI: Axis = Axis(2);
144
145    /// `Ih <- Oh*stride + Rh*dilation`, `Ci` passed through.
146    fn conv(stride: usize, dilation: usize) -> Projection {
147        Projection::new(
148            &[OH, RH, CI],
149            &[
150                PhysicalAxisMap::affine(&[(OH, stride), (RH, dilation)]),
151                PhysicalAxisMap::of(CI),
152            ],
153        )
154    }
155
156    fn extents(oh: usize, rh: usize, ci: usize) -> impl Fn(Axis) -> usize {
157        move |a| match a {
158            OH => oh,
159            RH => rh,
160            _ => ci,
161        }
162    }
163
164    fn extents2(oh: usize, rh: usize) -> impl Fn(Axis) -> usize {
165        move |a| match a {
166            OH => oh,
167            RH => rh,
168            _ => unreachable!("axis not spanned by this projection"),
169        }
170    }
171
172    #[test]
173    fn direct_compacts_to_itself() {
174        let p = Projection::direct(&[OH, CI]);
175        let c = Compaction::new(&p, 4, extents(8, 1, 16));
176        assert!(c.is_dense());
177        assert_eq!(c.steps(), &[1, 1]);
178        assert_eq!(c.extents(), &[8, 16]);
179        assert_eq!(c.projection(), &p);
180    }
181
182    #[test]
183    fn a_unit_stride_window_is_the_receptive_field() {
184        let c = Compaction::new(&conv(1, 1), 4, extents(8, 3, 16));
185        assert!(c.is_dense());
186        assert_eq!(c.extents(), &[10, 16]);
187        assert_eq!(c.cells(4), 10 * 4);
188    }
189
190    #[test]
191    fn a_strided_window_with_adjacent_taps_stays_dense() {
192        let c = Compaction::new(&conv(2, 1), 4, extents(8, 3, 16));
193        assert!(c.is_dense());
194        // 1 + 7*2 + 2*1
195        assert_eq!(c.extents(), &[17, 16]);
196    }
197
198    #[test]
199    fn a_single_tap_at_stride_two_halves_the_stage() {
200        for dilation in [1, 3] {
201            let c = Compaction::new(&conv(2, dilation), 4, extents(8, 1, 16));
202            assert!(!c.is_dense());
203            assert_eq!(c.steps(), &[2, 1]);
204            assert_eq!(c.extents(), &[8, 16]);
205            assert_eq!(c.projection().scale(0, OH), 1);
206            assert_eq!(c.projection().scale(0, RH), 1);
207        }
208    }
209
210    #[test]
211    fn a_shared_factor_quotients_the_window() {
212        let c = Compaction::new(&conv(2, 2), 4, extents(8, 3, 16));
213        assert_eq!(c.steps(), &[2, 1]);
214        // Bounding box 1 + 7*2 + 2*2 = 19, on the even lattice: 1 + 7 + 2 = 10.
215        assert_eq!(c.extents(), &[10, 16]);
216        assert_eq!(c.projection().scale(0, OH), 1);
217        assert_eq!(c.projection().scale(0, RH), 1);
218    }
219
220    /// `{0, 1}` at stride 3 reaches `{0, 1, 3, 4, 6, 7}`: gcd 1, holes kept as padding.
221    #[test]
222    fn unreachable_offsets_inside_the_box_stay_as_padding() {
223        let c = Compaction::new(&conv(3, 1), 4, extents(3, 2, 16));
224        assert!(c.is_dense());
225        assert_eq!(c.extents(), &[8, 16]);
226    }
227
228    #[test]
229    fn a_ragged_innermost_extent_is_padded() {
230        let lines = Compaction::new(&conv(1, 1), 4, extents(8, 3, 6)).line_extents(4);
231        assert_eq!(lines, &[10, 2]);
232    }
233
234    #[test]
235    fn a_scalar_gather_may_compact_its_only_axis() {
236        let p = Projection::new(&[OH, RH], &[PhysicalAxisMap::affine(&[(OH, 2), (RH, 1)])]);
237        let c = Compaction::new(&p, 1, extents2(8, 3));
238        assert!(c.is_dense());
239        assert_eq!(c.extents(), &[17]);
240    }
241
242    #[test]
243    fn padded_projection_compacts_identically() {
244        let p = Projection::new(
245            &[OH, RH, CI],
246            &[
247                PhysicalAxisMap::affine(&[(OH, 2), (RH, 1)]).shifted(-2),
248                PhysicalAxisMap::of(CI),
249            ],
250        );
251        let c = Compaction::new(&p, 4, extents(8, 3, 16));
252        assert!(c.is_dense());
253        assert_eq!(c.steps(), &[1, 1]);
254        assert_eq!(c.extents(), &[17, 16]);
255        assert_eq!(c.projection().offset(0), Offset::Static(0));
256    }
257
258    #[test]
259    fn dynamic_offset_projection_compacts_identically() {
260        let p = Projection::new(
261            &[OH, RH, CI],
262            &[
263                PhysicalAxisMap::affine(&[(OH, 2), (RH, 1)]).shifted(Offset::Dynamic),
264                PhysicalAxisMap::of(CI),
265            ],
266        );
267        let c = Compaction::new(&p, 4, extents(8, 3, 16));
268        assert!(c.is_dense());
269        assert_eq!(c.steps(), &[1, 1]);
270        assert_eq!(c.extents(), &[17, 16]);
271        assert_eq!(c.projection().offset(0), Offset::Static(0));
272    }
273
274    #[test]
275    fn a_dynamic_coefficient_compacts_dense_against_its_max() {
276        let p = Projection::new(
277            &[OH, RH, CI],
278            &[
279                PhysicalAxisMap::scaled(&[(OH, Scale::Dynamic { max: 2 }), (RH, Scale::Static(1))]),
280                PhysicalAxisMap::of(CI),
281            ],
282        );
283        let c = Compaction::new(&p, 4, extents(8, 3, 16));
284        assert!(c.is_dense());
285        // 1 + 7*2 + 2*1, the box `conv(2, 1)` fills.
286        assert_eq!(c.extents(), &[17, 16]);
287        assert_eq!(
288            Compaction::new(&conv(2, 1), 4, extents(8, 3, 16)).extents(),
289            c.extents()
290        );
291        assert!(c.projection().physical_axis(0).has_dynamic_scale());
292        assert_eq!(c.projection().dynamic_scale_index(0, 0), Some(0));
293    }
294
295    #[test]
296    fn a_moving_dynamic_coefficient_gives_up_the_lattice() {
297        let p = Projection::new(
298            &[OH, RH, CI],
299            &[
300                PhysicalAxisMap::scaled(&[(OH, Scale::Static(2)), (RH, Scale::Dynamic { max: 2 })]),
301                PhysicalAxisMap::of(CI),
302            ],
303        );
304        let c = Compaction::new(&p, 4, extents(8, 3, 16));
305        assert!(c.is_dense());
306        assert_eq!(c.steps(), &[1, 1]);
307        // 1 + 7*2 + 2*2, the bounding box `conv(2, 2)` quotients but this one cannot.
308        assert_eq!(c.extents(), &[19, 16]);
309    }
310
311    #[test]
312    fn a_non_moving_dynamic_coefficient_keeps_its_array_slot() {
313        let p = Projection::new(
314            &[OH, RH, CI],
315            &[
316                PhysicalAxisMap::scaled(&[(OH, Scale::Static(2)), (RH, Scale::Dynamic { max: 4 })]),
317                PhysicalAxisMap::of(CI),
318            ],
319        );
320        let c = Compaction::new(&p, 4, extents(8, 1, 16));
321        assert_eq!(c.steps(), &[2, 1]);
322        assert_eq!(c.extents(), &[8, 16]);
323        assert_eq!(p.dynamic_scale_index(0, 1), Some(0));
324        assert_eq!(c.projection().dynamic_scale_index(0, 1), Some(0));
325    }
326
327    #[test]
328    fn rational_axis_compacts_dense_with_conservative_extent() {
329        let p = Projection::new(
330            &[OH, RH, CI],
331            &[
332                PhysicalAxisMap::affine(&[(OH, 2), (RH, 3)]).over(3),
333                PhysicalAxisMap::of(CI),
334            ],
335        );
336        // field = 7*2 + 2*3 = 20, extent = 1 + ⌈20 / 3⌉ = 8
337        let c = Compaction::new(&p, 4, extents(8, 3, 16));
338        assert!(c.is_dense());
339        assert_eq!(c.steps(), &[1, 1]);
340        assert_eq!(c.extents(), &[8, 16]);
341        assert!(c.projection().physical_axis(0).is_rational());
342        assert_eq!(c.projection().divisor(0), Divisor::Static(3));
343    }
344
345    #[test]
346    fn a_dynamic_divisor_compacts_against_its_min() {
347        let p = Projection::new(
348            &[OH, RH, CI],
349            &[
350                PhysicalAxisMap::affine(&[(OH, 2), (RH, 3)]).over(Divisor::Dynamic { min: 3 }),
351                PhysicalAxisMap::of(CI),
352            ],
353        );
354        let c = Compaction::new(&p, 4, extents(8, 3, 16));
355        assert!(c.is_dense());
356        assert_eq!(c.extents(), &[8, 16]);
357        assert_eq!(c.projection().divisor(0), Divisor::Dynamic { min: 3 });
358        assert_eq!(c.projection().dynamic_divisor_index(0), Some(0));
359    }
360
361    #[test]
362    fn a_dynamic_divisor_box_holds_every_divisor_above_its_min() {
363        let bounded = Projection::new(
364            &[OH, RH, CI],
365            &[
366                PhysicalAxisMap::affine(&[(OH, 2), (RH, 3)]).over(Divisor::Dynamic { min: 2 }),
367                PhysicalAxisMap::of(CI),
368            ],
369        );
370        let box_extent = Compaction::new(&bounded, 4, extents(8, 3, 16)).extents()[0];
371        for d in 2..24 {
372            let exact = Projection::new(
373                &[OH, RH, CI],
374                &[
375                    PhysicalAxisMap::affine(&[(OH, 2), (RH, 3)]).over(d),
376                    PhysicalAxisMap::of(CI),
377                ],
378            );
379            if exact.physical_axis(0).is_rational() {
380                assert!(Compaction::new(&exact, 4, extents(8, 3, 16)).extents()[0] <= box_extent);
381            }
382        }
383    }
384
385    #[test]
386    fn dynamic_offset_with_cancelling_divisor_compacts_as_integers() {
387        let p_dynamic_offset = Projection::new(
388            &[OH, RH, CI],
389            &[
390                PhysicalAxisMap::scaled(&[(OH, Scale::Static(16)), (RH, Scale::Static(8))])
391                    .shifted(Offset::Dynamic)
392                    .over(4),
393                PhysicalAxisMap::of(CI),
394            ],
395        );
396        let p_integer = Projection::new(
397            &[OH, RH, CI],
398            &[
399                PhysicalAxisMap::affine(&[(OH, 4), (RH, 2)]),
400                PhysicalAxisMap::of(CI),
401            ],
402        );
403        let c_dynamic = Compaction::new(&p_dynamic_offset, 4, extents(8, 3, 16));
404        let c_integer = Compaction::new(&p_integer, 4, extents(8, 3, 16));
405
406        assert_eq!(c_dynamic.steps(), c_integer.steps());
407        assert_eq!(c_dynamic.extents(), c_integer.extents());
408        assert_eq!(c_dynamic.projection(), c_integer.projection());
409        assert_eq!(c_dynamic.steps(), &[2, 1]);
410    }
411}
cubecl_tile/layout/compaction/base.rs

cubecl_tile/layout/compaction/
base.rs
