cubecl_tile/space/partition/
levels.rs

1//! [`Levels`]: stating a partitioning's levels from the leaf up, each as a count of the one below.
2//!
3//! ```ignore
4//! Levels::leaf(&[(M, 16), (N, 8), (K, 16)])   // the instruction
5//!     .walk(&[(M, 2), (N, 4)])                // fragments a plane holds
6//!     .walk(&[(K, 2)])                        // instruction steps in one stage
7//!     .planes(&[(M, 2), (N, 4)])              // planes a cube holds
8//!     .walk_every(&[K])                       // every stage `k` holds
9//!     .cubes(&[M, N]).batches(&[B])           // every box `m x n` holds
10//!     .build()
11//! ```
12//!
13//! An `every` level ([`Count::All`]) closes its axis; only [`across`](Levels::across) may name it
14//! again, on the cube level.
15
16use super::{CubeOrder, LaunchSpan, PlaneOrder, UnitCube};
17use crate::{Axis, ComputeScope, Count, Coverage, Level, Spread};
18
19/// One level as the builder holds it before it is a [`Level`].
20#[derive(Clone, Debug)]
21struct StatedLevel {
22    coverage: Coverage,
23    /// `(axis, tile size, count)`; an every-level's count is unused.
24    tiles: Vec<(Axis, usize, usize)>,
25    takes: Takes,
26    /// A cube level distributing one of its axes across several cubes.
27    across: Option<(Axis, usize)>,
28    shared_by: Option<usize>,
29    fillers: usize,
30    batches: Vec<Axis>,
31    interleaved: Vec<Axis>,
32    reversed: Vec<Axis>,
33    /// Closed axes this cube level names again, to distribute across cubes.
34    reopened: Vec<Axis>,
35    order: CubeOrder,
36    /// Who across the whole launch takes this level's tiles, where it spans the launch.
37    launch_span: Option<LaunchSpan>,
38}
39
40/// Which [`Count`] a stated level's entries carry.
41#[derive(Clone, Copy, PartialEq, Eq, Debug)]
42enum Takes {
43    Stated,
44    Every,
45    DistributedToUnits,
46}
47
48/// A partitioning stated from the leaf up.
49#[derive(Clone, Debug)]
50pub struct Levels {
51    /// What one tile of each axis covers so far, leaf up.
52    sizes: Vec<(Axis, usize)>,
53    /// Axes an `every` level closed.
54    closed: Vec<Axis>,
55    /// Innermost first; [`build`](Self::build) reverses.
56    stated: Vec<StatedLevel>,
57}
58
59impl Levels {
60    /// The tile one worker holds at the bottom, e.g. an instruction's shape.
61    pub fn leaf(tile: &[(Axis, usize)]) -> Self {
62        Levels {
63            sizes: tile.to_vec(),
64            closed: Vec::new(),
65            stated: Vec::new(),
66        }
67    }
68
69    /// Every worker steps through this many of the level below.
70    pub fn walk(self, counts: &[(Axis, usize)]) -> Self {
71        self.state(Coverage::Walk, counts)
72    }
73
74    /// Every worker steps through every tile of the level below along `axes`.
75    pub fn walk_every(self, axes: &[Axis]) -> Self {
76        self.every(Coverage::Walk, axes)
77    }
78
79    /// This many of the thing below, one per unit of the plane.
80    pub fn units(self, counts: &[(Axis, usize)]) -> Self {
81        self.state(Coverage::Distribute(ComputeScope::Unit), counts)
82    }
83
84    /// This many of the level below, taken in turns by however many units the launch runs.
85    pub fn units_distributed(mut self, axis: Axis, count: usize) -> Self {
86        let tiles = vec![(axis, self.size(axis), count)];
87        let mut stated = StatedLevel::new(
88            Coverage::Distribute(ComputeScope::Unit),
89            tiles,
90            Takes::DistributedToUnits,
91        );
92        stated.interleaved.push(axis);
93        self.stated.push(stated);
94        self.grow(axis, count);
95        self
96    }
97
98    /// This many of the level below, one per plane of the cube.
99    pub fn planes(self, counts: &[(Axis, usize)]) -> Self {
100        self.state(Coverage::Distribute(ComputeScope::Plane), counts)
101    }
102
103    /// This many of the level below, one per group of `planes` consecutive planes of the cube: an
104    /// instruction the group's planes issue together, a warpgroup MMA's four, takes each tile.
105    pub fn plane_groups(self, planes: usize, counts: &[(Axis, usize)]) -> Self {
106        assert!(
107            planes > 1,
108            "Levels::plane_groups: a group of {planes} plane is a plane; state `planes`"
109        );
110        self.state(
111            Coverage::Distribute(ComputeScope::PlaneGroup { planes }),
112            counts,
113        )
114    }
115
116    /// Every tile of the level below along `axis`, dealt to `planes` planes of the cube in runs:
117    /// the planes need not divide the tiles, and a plane past the last whole run takes the rest.
118    /// What splits a walk between the planes where its steps have no factor the planes fit.
119    pub fn planes_across(self, axis: Axis, planes: usize) -> Self {
120        let mut levels = self.every(Coverage::Distribute(ComputeScope::Plane), &[axis]);
121        levels.last("planes_across").across = Some((axis, planes));
122        levels
123    }
124
125    /// One cube for every box these axes hold; the entries ride the grid's `X`, `Y`, `Z` in order.
126    pub fn cubes(self, axes: &[Axis]) -> Self {
127        assert!(
128            axes.len() <= 3,
129            "Levels::cubes: {} axes, but a launch grid has three dimensions",
130            axes.len()
131        );
132        self.every(Coverage::Distribute(ComputeScope::Cube), axes)
133    }
134
135    /// One plane of the launch for every tile these axes hold, the planes taking them in `order`:
136    /// the outermost level of a partitioning with no cube level. A plane routine's planes share
137    /// nothing, so which cube a plane sits in is the launch's grouping, not a level of the
138    /// partitioning; the order sets the cubes' size ([`PlaneOrder::planes_per_cube`]).
139    pub fn launch_planes(mut self, axes: &[Axis], order: PlaneOrder) -> Self {
140        self = self.every(Coverage::Distribute(ComputeScope::Plane), axes);
141        let stated = self.last("launch_planes");
142        assert!(
143            stated.reopened.is_empty(),
144            "Levels::launch_planes: {:?} was taken whole by a level below",
145            stated.reopened
146        );
147        stated.launch_span = Some(LaunchSpan::Planes(order));
148        self
149    }
150
151    /// One unit of the launch for every tile these axes hold, in row order, the units grouped in
152    /// cubes of `cube`: the outermost level of a partitioning with no cube or plane level. A unit
153    /// routine's units share nothing, so which cube or plane a unit sits in is the launch's
154    /// grouping, not a level of the partitioning, and `for unit in &partitioning` walks them.
155    pub fn launch_units(mut self, axes: &[Axis], cube: UnitCube) -> Self {
156        self = self.every(Coverage::Distribute(ComputeScope::Unit), axes);
157        let stated = self.last("launch_units");
158        assert!(
159            stated.reopened.is_empty(),
160            "Levels::launch_units: {:?} was taken whole by a level below",
161            stated.reopened
162        );
163        stated.launch_span = Some(LaunchSpan::Units(cube));
164        self
165    }
166
167    /// Distribute `axis` of the cube level just stated across `cubes` cubes, each a run (split-K).
168    pub fn across(mut self, axis: Axis, cubes: usize) -> Self {
169        let stated = self.last("across");
170        assert!(
171            stated.coverage == Coverage::Distribute(ComputeScope::Cube),
172            "Levels::across: only cubes take an axis in runs; the level just stated is {:?}",
173            stated.coverage
174        );
175        assert!(
176            stated.tiles.iter().any(|&(a, _, _)| a == axis),
177            "Levels::across: {axis:?} is not an axis of the cube level just stated; name it there"
178        );
179        stated.across = Some((axis, cubes));
180        self
181    }
182
183    /// Distribute this level's tiles to `cubes` cubes as one flat index (stream-K).
184    pub fn shared_by(mut self, cubes: usize) -> Self {
185        self.last("shared_by").shared_by = Some(cubes);
186        self
187    }
188
189    /// The order the cube level just stated distributes its boxes to the grid.
190    pub fn ordered(mut self, order: CubeOrder) -> Self {
191        let stated = self.last("ordered");
192        assert!(
193            stated.coverage == Coverage::Distribute(ComputeScope::Cube),
194            "Levels::ordered: only cubes are distributed to a grid; the level just stated is {:?}",
195            stated.coverage
196        );
197        stated.order = order.canonicalize();
198        self
199    }
200
201    /// This walk's stages are filled by `n` extra planes that take no tile.
202    pub fn filled_by(mut self, n: usize) -> Self {
203        self.last("filled_by").fillers = n;
204        self
205    }
206
207    /// Batch axes this cube level hands out one of.
208    pub fn batches(mut self, axes: &[Axis]) -> Self {
209        self.last("batches").batches.extend_from_slice(axes);
210        self
211    }
212
213    /// The workers of the level just stated take `axis`'s tiles in turns rather than in runs.
214    pub fn interleaved(mut self, axis: Axis) -> Self {
215        self.last("interleaved").interleaved.push(axis);
216        self
217    }
218
219    /// The level just stated, a cube level or a plane level spanning the launch, hands out
220    /// `axis`'s tiles from the last to the first: the first workers launched take the last tiles.
221    /// Where the later tiles hold more work, the heaviest start first and the light ones fill in
222    /// behind them.
223    pub fn reversed(mut self, axis: Axis) -> Self {
224        self.last("reversed").reversed.push(axis);
225        self
226    }
227
228    /// The levels, outermost first.
229    pub fn build(self) -> Vec<Level> {
230        for stated in &self.stated {
231            for &axis in &stated.reopened {
232                assert!(
233                    matches!(stated.across, Some((a, _)) if a == axis),
234                    "Levels::cubes: {axis:?} was taken whole by a level below; a cube level names \
235                     it again only to distribute it across cubes (`.across({axis:?}, n)`)"
236                );
237            }
238        }
239        let groups = self
240            .stated
241            .iter()
242            .filter(|stated| {
243                matches!(
244                    stated.coverage,
245                    Coverage::Distribute(ComputeScope::PlaneGroup { .. })
246                )
247            })
248            .count();
249        let planes = self
250            .stated
251            .iter()
252            .any(|stated| stated.coverage == Coverage::Distribute(ComputeScope::Plane));
253        assert!(
254            groups == 0 || (groups == 1 && !planes),
255            "Levels: a cube's planes are dealt to plane groups by one level, beside no level of \
256             single planes; these levels state {groups} plane group level(s){}",
257            if planes { " and a plane level" } else { "" }
258        );
259        self.stated.iter().rev().map(StatedLevel::level).collect()
260    }
261
262    /// The single level this tiling states. Panics if it states several.
263    pub fn level(self) -> Level {
264        assert!(
265            self.stated.len() == 1,
266            "Levels::level: {} levels are stated; a walk over a region takes exactly one",
267            self.stated.len()
268        );
269        self.build().remove(0)
270    }
271
272    /// State a level of `counts` of the level below, then grow each axis by its count.
273    fn state(mut self, coverage: Coverage, counts: &[(Axis, usize)]) -> Self {
274        let tiles = counts
275            .iter()
276            .map(|&(axis, count)| (axis, self.size(axis), count));
277        self.stated
278            .push(StatedLevel::new(coverage, tiles.collect(), Takes::Stated));
279        for &(axis, count) in counts {
280            self.grow(axis, count);
281        }
282        self
283    }
284
285    /// State a level covering every tile these axes hold, and close them.
286    fn every(mut self, coverage: Coverage, axes: &[Axis]) -> Self {
287        let reopened: Vec<Axis> = axes
288            .iter()
289            .copied()
290            .filter(|axis| self.closed.contains(axis))
291            .collect();
292        if let (Coverage::Walk, Some(axis)) = (coverage, reopened.first()) {
293            panic!(
294                "Levels::walk_every: {axis:?} was taken whole by a level below, so a walk above \
295                 it has nothing to step through"
296            );
297        }
298        let tiles = axes.iter().map(|&axis| (axis, self.size(axis), 1));
299        let mut stated = StatedLevel::new(coverage, tiles.collect(), Takes::Every);
300        stated.reopened = reopened;
301        self.stated.push(stated);
302        self.closed.extend_from_slice(axes);
303        self
304    }
305
306    /// What one tile of `axis` covers so far; `1` if unbuilt.
307    fn size(&self, axis: Axis) -> usize {
308        self.sizes
309            .iter()
310            .find(|&&(a, _)| a == axis)
311            .map_or(1, |&(_, size)| size)
312    }
313
314    fn grow(&mut self, axis: Axis, count: usize) {
315        assert!(
316            !self.closed.contains(&axis),
317            "Levels: {axis:?} was taken whole by a level below, so nothing above it may state a \
318             count of its tiles"
319        );
320        match self.sizes.iter_mut().find(|(a, _)| *a == axis) {
321            Some((_, size)) => *size *= count,
322            None => self.sizes.push((axis, count)),
323        }
324    }
325
326    fn last(&mut self, verb: &str) -> &mut StatedLevel {
327        self.stated
328            .last_mut()
329            .unwrap_or_else(|| panic!("Levels::{verb}: no level has been stated for it to modify"))
330    }
331}
332
333impl StatedLevel {
334    fn new(coverage: Coverage, tiles: Vec<(Axis, usize, usize)>, takes: Takes) -> Self {
335        StatedLevel {
336            coverage,
337            tiles,
338            takes,
339            across: None,
340            shared_by: None,
341            fillers: 0,
342            batches: Vec::new(),
343            interleaved: Vec::new(),
344            reversed: Vec::new(),
345            reopened: Vec::new(),
346            order: CubeOrder::RowMajor,
347            launch_span: None,
348        }
349    }
350
351    /// How the level's workers take `axis`'s tiles.
352    fn spread(&self, axis: Axis) -> Spread {
353        match self.interleaved.contains(&axis) {
354            true => Spread::Interleaved,
355            false => Spread::Contiguous,
356        }
357    }
358
359    /// The [`Level`] this states.
360    fn level(&self) -> Level {
361        let tile = || -> Vec<_> {
362            self.tiles
363                .iter()
364                .map(|&(axis, tile, _)| (axis, tile))
365                .collect()
366        };
367        if let Some(span) = &self.launch_span {
368            let level = match span {
369                LaunchSpan::Planes(order) => Level::spanning_launch(&tile(), order.clone()),
370                LaunchSpan::Units(cube) => Level::units_spanning_launch(&tile(), *cube),
371            };
372            return level.reversed(&self.reversed);
373        }
374        let cuts: Vec<(Axis, usize, Count, Spread)> = self
375            .tiles
376            .iter()
377            .map(|&(axis, tile, count)| {
378                let count = match (self.takes, self.across) {
379                    (Takes::Every, Some((across, cubes))) if across == axis => {
380                        Count::AllAcross(cubes)
381                    }
382                    (Takes::Every, _) => Count::All,
383                    (Takes::Stated, _) => Count::Stated(count),
384                    (Takes::DistributedToUnits, _) => Count::Distributed(count),
385                };
386                (axis, tile, count, self.spread(axis))
387            })
388            .collect();
389        let level = Level::new(self.coverage, &cuts);
390        let level = match self.batches.is_empty() {
391            true => level,
392            false => level.batching(&self.batches),
393        };
394        let level = match self.fillers {
395            0 => level,
396            n => level.filling(n),
397        };
398        let level = match self.order.swizzles() {
399            false => level,
400            true => level.distributed_in(self.order),
401        };
402        let level = match self.shared_by {
403            None => level,
404            Some(cubes) => level.sharing(cubes),
405        };
406        match self.reversed.is_empty() {
407            true => level,
408            false => level.reversed(&self.reversed),
409        }
410    }
411}
412
413#[cfg(test)]
414mod tests {
415    use super::*;
416    use crate::{Partitioning, Space};
417
418    const M: Axis = Axis(0);
419    const N: Axis = Axis(1);
420    const K: Axis = Axis(2);
421    const B: Axis = Axis(3);
422
423    /// A level's tiles stay inside the box it is handed exactly where its tile divides that box,
424    /// which is what lets a window cut by it keep the page of the window it was cut from: pages of
425    /// 16 keys over 20 overhang, and steps of 4 inside a page's box of 16 do not.
426    #[test]
427    fn a_level_overhangs_only_a_box_its_tile_does_not_divide() {
428        let levels = Levels::leaf(&[(K, 4)])
429            .walk(&[(K, 4)])
430            .walk_every(&[K])
431            .build();
432        let keys = Space::new(&[(K, 20)]);
433        assert!(levels[0].overhangs(&keys, K));
434        let page = levels[0].child(&keys);
435        assert_eq!(page.extent(K), 16);
436        assert!(!levels[1].overhangs(&page, K));
437    }
438
439    /// Every level's tile is the product of what was stated below it.
440    #[test]
441    fn every_level_is_a_count_of_the_level_before() {
442        let levels = Levels::leaf(&[(M, 16), (N, 8), (K, 16)])
443            .walk(&[(M, 2), (N, 4)])
444            .walk(&[(K, 2)])
445            .planes(&[(M, 2), (N, 4)])
446            .walk_every(&[K])
447            .cubes(&[M, N])
448            .batches(&[B])
449            .build();
450        assert_eq!(levels.len(), 5);
451        let tiles = |i: usize, axis| levels[i].tile(axis);
452        let counts = |i: usize, axis| levels[i].count(axis);
453
454        assert_eq!((tiles(0, M), tiles(0, N)), (Some(64), Some(128)));
455        assert_eq!(
456            (counts(0, M), counts(0, N)),
457            (Some(Count::All), Some(Count::All))
458        );
459        assert_eq!((tiles(0, B), counts(0, B)), (Some(1), Some(Count::All)));
460        assert_eq!(tiles(0, K), None);
461
462        assert_eq!((tiles(1, K), counts(1, K)), (Some(32), Some(Count::All)));
463        assert_eq!(
464            (tiles(2, M), counts(2, M)),
465            (Some(32), Some(Count::Stated(2)))
466        );
467        assert_eq!(
468            (tiles(2, N), counts(2, N)),
469            (Some(32), Some(Count::Stated(4)))
470        );
471        assert_eq!(
472            (tiles(3, K), counts(3, K)),
473            (Some(16), Some(Count::Stated(2)))
474        );
475        assert_eq!(
476            (tiles(4, M), counts(4, M)),
477            (Some(16), Some(Count::Stated(2)))
478        );
479        assert_eq!(
480            (tiles(4, N), counts(4, N)),
481            (Some(8), Some(Count::Stated(4)))
482        );
483    }
484
485    /// The planes take a walk's steps in runs at a count that need not divide them: the plane
486    /// level states the count of planes, not a tile of planes times a step.
487    #[test]
488    fn the_planes_take_every_step_across_in_runs() {
489        let levels = Levels::leaf(&[(N, 8), (K, 16)])
490            .planes_across(K, 4)
491            .cubes(&[N])
492            .build();
493        assert_eq!(
494            levels[1].coverage(),
495            Coverage::Distribute(ComputeScope::Plane)
496        );
497        assert_eq!(levels[1].count(K), Some(Count::AllAcross(4)));
498        assert_eq!(levels[1].tile(K), Some(16));
499    }
500
501    /// Plane groups are dealt by one level, beside no level of single planes, which the cube's
502    /// count of planes would otherwise take twice.
503    #[test]
504    #[should_panic(expected = "a cube's planes are dealt to plane groups by one level")]
505    fn plane_groups_take_no_plane_level_beside_them() {
506        Levels::leaf(&[(M, 64), (N, 128), (K, 64)])
507            .planes(&[(N, 2)])
508            .plane_groups(4, &[(M, 2)])
509            .walk_every(&[K])
510            .build();
511    }
512
513    /// A plane group's level counts its groups, and the cube holds each group's planes.
514    #[test]
515    fn a_plane_group_holds_its_planes() {
516        let levels = Levels::leaf(&[(M, 64), (N, 128), (K, 64)])
517            .plane_groups(4, &[(M, 2)])
518            .walk_every(&[K])
519            .filled_by(1)
520            .cubes(&[M, N])
521            .build();
522        let groups = Coverage::Distribute(ComputeScope::PlaneGroup { planes: 4 });
523        assert!(levels.iter().any(|level| level.coverage() == groups));
524        let partitioning = Partitioning::new(Space::new(&[(M, 1024), (N, 1024), (K, 512)]), levels);
525        assert_eq!(partitioning.instances(groups), 2);
526        assert_eq!(partitioning.planes_per_cube(), 9);
527    }
528
529    /// A closed axis returns to the cube level to be split across cubes.
530    #[test]
531    fn a_closed_axis_returns_to_the_cube_level_to_be_distributed_across() {
532        let levels = Levels::leaf(&[(M, 16), (N, 8), (K, 16)])
533            .walk_every(&[K])
534            .cubes(&[M, N, K])
535            .across(K, 4)
536            .build();
537        assert_eq!(levels[0].count(K), Some(Count::AllAcross(4)));
538        assert_eq!(levels[0].tile(K), Some(16));
539        assert_eq!(levels[0].cube_axis(K), Some(crate::CubeAxis::Z));
540        assert_eq!(levels[1].count(K), Some(Count::All));
541    }
542
543    /// A count builds its tile, so nothing can fail to divide.
544    #[test]
545    fn a_count_builds_its_tile_and_nothing_refuses() {
546        let levels = Levels::leaf(&[(M, 5)])
547            .planes(&[(M, 3)])
548            .walk(&[(M, 7)])
549            .cubes(&[M])
550            .build();
551        assert_eq!(levels[0].tile(M), Some(105));
552        assert_eq!(levels[1].tile(M), Some(15));
553        assert_eq!(levels[1].count(M), Some(Count::Stated(7)));
554        assert_eq!(levels[2].count(M), Some(Count::Stated(3)));
555        // Only the top every-level can overhang.
556        let space = Space::new(&[(M, 100)]);
557        assert!(levels[0].overhangs(&space, M));
558        let cube = levels[0].child(&space);
559        assert!(!levels[1].overhangs(&cube, M));
560        assert!(!levels[2].overhangs(&levels[1].child(&cube), M));
561    }
562
563    /// Walking every tile is not walking one.
564    #[test]
565    fn every_is_not_one() {
566        let every = Levels::leaf(&[(K, 16)]).walk_every(&[K]).level();
567        let one = Levels::leaf(&[(K, 16)]).walk(&[(K, 1)]).level();
568        assert_ne!(every, one);
569        assert_eq!(every.count(K), Some(Count::All));
570        assert_eq!(one.count(K), Some(Count::Stated(1)));
571    }
572
573    /// A count distributed to the units builds its tile like any other.
574    #[test]
575    fn a_count_distributed_to_the_units_builds_its_tile() {
576        let levels = Levels::leaf(&[(M, 4), (N, 4)])
577            .units_distributed(N, 64)
578            .planes(&[(M, 2)])
579            .cubes(&[M, N])
580            .build();
581        assert_eq!(levels[0].tile(N), Some(256));
582        assert_eq!(levels[2].count(N), Some(Count::Distributed(64)));
583        assert_eq!(
584            levels[2].coverage(),
585            Coverage::Distribute(ComputeScope::Unit)
586        );
587        assert_eq!(levels[2].spread(N), Some(Spread::Interleaved));
588        let partitioning = Partitioning::new(Space::new(&[(M, 100), (N, 1000)]), levels);
589        assert_eq!(partitioning.units(), 1);
590        assert_eq!(partitioning.planes_per_cube(), 2);
591    }
592
593    #[test]
594    #[should_panic(expected = "distributes no other axis to the units")]
595    fn a_count_distributed_to_the_units_is_their_only_axis() {
596        let _ = Level::new(
597            Coverage::Distribute(ComputeScope::Unit),
598            &[
599                (M, 4, Count::Distributed(8), Spread::Interleaved),
600                (N, 4, Count::Stated(4), Spread::Contiguous),
601            ],
602        );
603    }
604
605    /// A region walk is the chain it stands for.
606    #[test]
607    fn a_region_walk_is_stated_once() {
608        assert_eq!(
609            Level::every(&[(M, 8), (K, 16)]),
610            Levels::leaf(&[(M, 8), (K, 16)]).walk_every(&[M, K]).level()
611        );
612    }
613
614    #[test]
615    #[should_panic(
616        expected = "taken whole by a level below, so nothing above it may state a count"
617    )]
618    fn nothing_above_all_of_it_may_count_it() {
619        let _ = Levels::leaf(&[(K, 16)]).walk_every(&[K]).walk(&[(K, 2)]);
620    }
621
622    #[test]
623    #[should_panic(expected = "a walk above it has nothing to step through")]
624    fn all_of_it_is_walked_once() {
625        let _ = Levels::leaf(&[(K, 16)]).walk_every(&[K]).walk_every(&[K]);
626    }
627
628    #[test]
629    #[should_panic(expected = "names it again only to distribute it across cubes")]
630    fn a_closed_axis_on_the_cube_level_must_be_distributed_across() {
631        let _ = Levels::leaf(&[(M, 16), (K, 16)])
632            .walk_every(&[K])
633            .cubes(&[M, K])
634            .build();
635    }
636
637    #[test]
638    #[should_panic(expected = "only cubes take an axis in runs")]
639    fn only_cubes_take_an_axis_in_runs() {
640        let _ = Levels::leaf(&[(K, 16)]).planes(&[(K, 4)]).across(K, 2);
641    }
642
643    /// A reversed axis reaches the cube level that states it, and no level below.
644    #[test]
645    fn a_reversed_axis_reaches_the_cube_level() {
646        let levels = Levels::leaf(&[(M, 16), (N, 8), (K, 16)])
647            .walk_every(&[K])
648            .planes(&[(M, 2)])
649            .cubes(&[M, N])
650            .reversed(M)
651            .build();
652        assert!(levels[0].is_reversed(M));
653        assert!(!levels[0].is_reversed(N));
654        assert!(levels[1..].iter().all(|level| !level.is_reversed(M)));
655    }
656
657    /// The planes of one cube are not dealt in launch order, so there is nothing to reverse.
658    #[test]
659    #[should_panic(expected = "only a cube level or a level spanning the launch")]
660    fn only_a_launch_order_is_reversed() {
661        let _ = Levels::leaf(&[(M, 16)])
662            .planes(&[(M, 2)])
663            .reversed(M)
664            .build();
665    }
666
667    #[test]
668    #[should_panic(expected = "a walk over a region takes exactly one")]
669    fn a_walk_over_a_region_takes_one_level() {
670        let _ = Levels::leaf(&[(K, 16)])
671            .walk(&[(K, 2)])
672            .walk_every(&[K])
673            .level();
674    }
675}
cubecl_tile/space/partition/levels.rs

cubecl_tile/space/partition/
levels.rs
