cubecl_tile/layout/projection/
axis_map.rs

1//! One physical axis of a buffer as an affine combination of logical axes, and its coefficients.
2
3use cubecl::zspace::SmallVec;
4
5use crate::space::extent::Extent;
6use crate::{Axis, Space, algebra::gcd};
7
8/// How far one unit of a logical axis moves along one physical axis.
9/// A `Dynamic` coefficient declares `max`, the largest value the launch may pass.
10#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
11pub enum Scale {
12    Static(usize),
13    Dynamic { max: usize },
14}
15
16impl Scale {
17    /// The comptime coefficient; panics on `Dynamic`.
18    pub fn get(self) -> usize {
19        match self {
20            Scale::Static(n) => n,
21            Scale::Dynamic { .. } => {
22                panic!(
23                    "Scale::get: this coefficient is Dynamic; its value is only known at runtime"
24                )
25            }
26        }
27    }
28
29    /// The largest value this coefficient can take: itself, or the declared `max`.
30    pub fn bound(self) -> usize {
31        match self {
32            Scale::Static(n) => n,
33            Scale::Dynamic { max } => max,
34        }
35    }
36
37    pub fn is_dynamic(self) -> bool {
38        matches!(self, Scale::Dynamic { .. })
39    }
40}
41
42/// The constant term of one physical axis's affine combination.
43#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
44pub enum Offset {
45    Static(isize),
46    Dynamic,
47}
48
49impl Offset {
50    /// The comptime offset; panics on `Dynamic`.
51    pub fn get(self) -> isize {
52        match self {
53            Offset::Static(n) => n,
54            Offset::Dynamic => {
55                panic!("Offset::get: this offset is Dynamic; its value is only known at runtime")
56            }
57        }
58    }
59
60    pub fn is_dynamic(self) -> bool {
61        matches!(self, Offset::Dynamic)
62    }
63}
64
65impl From<isize> for Offset {
66    fn from(n: isize) -> Self {
67        Offset::Static(n)
68    }
69}
70
71/// What one physical axis's affine combination is divided by, floored.
72/// A `Dynamic` divisor declares `min`, the smallest value the launch may pass.
73#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
74pub enum Divisor {
75    Static(usize),
76    Dynamic { min: usize },
77}
78
79impl Divisor {
80    pub fn is_dynamic(self) -> bool {
81        matches!(self, Divisor::Dynamic { .. })
82    }
83
84    /// The smallest value this divisor can take: itself, or the declared `min`.
85    pub fn bound(self) -> usize {
86        match self {
87            Divisor::Static(d) => d,
88            Divisor::Dynamic { min } => min,
89        }
90    }
91
92    /// Whether this divides by `1`; a `Dynamic` divisor never is.
93    pub(crate) fn is_unit(self) -> bool {
94        self == Divisor::Static(1)
95    }
96}
97
98impl From<usize> for Divisor {
99    fn from(n: usize) -> Self {
100        Divisor::Static(n)
101    }
102}
103
104/// One logical axis's contribution to one physical axis: `coordinate * scale`.
105#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
106pub(crate) struct AxisTerm {
107    pub axis: Axis,
108    pub scale: Scale,
109}
110
111/// Whether a physical axis's terms can land on the same cell, as stated by the caller.
112#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
113pub(crate) enum Composition {
114    /// No two logical positions share a cell: the terms partition the axis.
115    Disjoint,
116    /// Two positions may land on the same cell, as with a stencil or a resample.
117    Overlapping,
118}
119
120/// One buffer dim as `(Σ coordinate(axis) * scale + offset) / divisor`, floored.
121#[derive(Clone, PartialEq, Eq, Hash, Debug)]
122pub struct PhysicalAxisMap {
123    terms: SmallVec<[AxisTerm; Space::MAX_RANK]>,
124    offset: Offset,
125    divisor: Divisor,
126    composition: Composition,
127}
128
129/// What a physical axis is addressed by.
130#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
131pub(crate) enum Addressed {
132    /// The leading logical axis of this map's terms, naming its coordinate group.
133    By(Axis),
134    /// Nothing: every position resolves to the same element.
135    Broadcast,
136}
137
138impl PhysicalAxisMap {
139    /// A map addressing no logical axis: every position resolves to the same element.
140    pub fn broadcast() -> Self {
141        PhysicalAxisMap::scaled(&[])
142    }
143
144    pub fn of(axis: Axis) -> Self {
145        PhysicalAxisMap {
146            terms: SmallVec::from_slice(&[AxisTerm {
147                axis,
148                scale: Scale::Static(1),
149            }]),
150            offset: Offset::Static(0),
151            divisor: Divisor::Static(1),
152            composition: Composition::Disjoint,
153        }
154    }
155
156    /// This physical axis partitioned by several logical ones, coarsest first:
157    /// `disjoint(&[(KB, 32), (KI, 1)])` is `k = kb * 32 + ki`.
158    pub fn disjoint(terms: &[(Axis, usize)]) -> Self {
159        let mut map = Self::affine(terms);
160        map.composition = Composition::Disjoint;
161        map
162    }
163
164    /// An affine combination, e.g. `affine(&[(Oh, stride), (Rh, dilation)])`.
165    pub fn affine(terms: &[(Axis, usize)]) -> Self {
166        let terms: SmallVec<[(Axis, Scale); Space::MAX_RANK]> = terms
167            .iter()
168            .map(|&(axis, scale)| (axis, Scale::Static(scale)))
169            .collect();
170        Self::scaled(&terms)
171    }
172
173    /// [`affine`](Self::affine) over explicit [`Scale`]s, for runtime coefficients.
174    pub fn scaled(terms: &[(Axis, Scale)]) -> Self {
175        for &(axis, scale) in terms {
176            assert!(
177                scale.bound() > 0,
178                "PhysicalAxisMap: {axis:?}'s coefficient is bounded at 0, so it addresses no cell; \
179                 drop the term instead"
180            );
181        }
182        // Only the identity is known to partition; anything else must be claimed `disjoint`.
183        let composition = match terms {
184            [(_, Scale::Static(1))] => Composition::Disjoint,
185            _ => Composition::Overlapping,
186        };
187        PhysicalAxisMap {
188            terms: terms
189                .iter()
190                .map(|&(axis, scale)| AxisTerm { axis, scale })
191                .collect(),
192            offset: Offset::Static(0),
193            divisor: Divisor::Static(1),
194            composition,
195        }
196    }
197
198    /// The same combination plus a signed constant or dynamic `offset`, added before any
199    /// [`over`](Self::over).
200    pub fn shifted(mut self, offset: impl Into<Offset>) -> Self {
201        let offset = offset.into();
202        assert!(
203            self.divisor.is_unit() && self.offset == Offset::Static(0),
204            "PhysicalAxisMap::shifted: shift once, before `over`"
205        );
206        if offset != Offset::Static(0) {
207            self.composition = Composition::Overlapping;
208        }
209        self.offset = offset;
210        self
211    }
212
213    /// The same combination divided by `divisor`, floored; a divisor every coefficient cancels
214    /// is reduced away.
215    pub fn over(mut self, divisor: impl Into<Divisor>) -> Self {
216        let divisor = divisor.into();
217        assert!(
218            divisor.bound() > 0,
219            "PhysicalAxisMap::over: a divisor of 0 does not map anywhere"
220        );
221        assert!(
222            self.divisor.is_unit(),
223            "PhysicalAxisMap::over: this map already divides by {:?}; state the whole divisor once",
224            self.divisor
225        );
226        self.divisor = divisor;
227        self.composition = Composition::Overlapping;
228        self.reduced()
229    }
230
231    /// The same mapping without a divisor, when every static coefficient is divisible by it.
232    fn reduced(mut self) -> Self {
233        let (Divisor::Static(d), Offset::Static(o)) = (self.divisor, self.offset) else {
234            return self;
235        };
236        let common = self.terms.iter().try_fold(0, |g, t| match t.scale {
237            Scale::Static(s) => Some(gcd(g, s)),
238            Scale::Dynamic { .. } => None,
239        });
240        // `gcd(0, s)` seeds from the first coefficient; an all-`0` combination reports `0`.
241        if !matches!(common, Some(g) if g % d == 0) {
242            return self;
243        }
244        for term in self.terms.iter_mut() {
245            term.scale = Scale::Static(term.scale.get() / d);
246        }
247        self.offset = Offset::Static(o.div_euclid(d as isize));
248        self.divisor = Divisor::Static(1);
249        self
250    }
251
252    /// What this physical axis is addressed by.
253    pub(crate) fn addressed(&self) -> Addressed {
254        match self.terms().first() {
255            Some(term) => Addressed::By(term.axis),
256            None => Addressed::Broadcast,
257        }
258    }
259
260    pub(crate) fn terms(&self) -> &[AxisTerm] {
261        &self.terms
262    }
263
264    /// How this axis's terms sit on it ([`Composition`]).
265    pub(crate) fn composition(&self) -> Composition {
266        self.composition
267    }
268
269    /// The `(axis, coefficient)` radices a disjoint map claims, coarsest first; empty otherwise.
270    pub(crate) fn claimed_radices(&self) -> SmallVec<[(Axis, usize); Space::MAX_RANK]> {
271        match self.composition {
272            Composition::Overlapping => SmallVec::new(),
273            Composition::Disjoint => self
274                .terms
275                .iter()
276                .map(|t| (t.axis, t.scale.bound()))
277                .collect(),
278        }
279    }
280
281    /// The offset of this physical axis.
282    pub fn offset(&self) -> Offset {
283        self.offset
284    }
285
286    /// The divisor, [`Static(1)`](Divisor::Static) unless [`over`](Self::over) made it rational.
287    pub(crate) fn divisor(&self) -> Divisor {
288        self.divisor
289    }
290
291    /// Term `term`'s static additive step, `None` if dynamic; panics for one kept inside a floor.
292    pub(crate) fn additive_step(&self, term: usize) -> Option<usize> {
293        if self.is_rational() {
294            return Some(self.static_offset_step(term).unwrap_or_else(|| {
295                panic!(
296                    "ProjectionInKernel::advance: {:?} stays inside this axis's floor, so stepping \
297                     it is not an addition and the map has to be summed at every position",
298                    self.terms()[term].axis
299                )
300            }));
301        }
302        match self.terms()[term].scale {
303            Scale::Static(s) => Some(s),
304            Scale::Dynamic { .. } => None,
305        }
306    }
307
308    /// Term `term`'s static step outside the rational floor, where `scale / divisor` is exact.
309    pub(crate) fn static_offset_step(&self, term: usize) -> Option<usize> {
310        match (self.divisor, self.terms[term].scale) {
311            (Divisor::Static(d), Scale::Static(s)) if d > 1 && s.is_multiple_of(d) => Some(s / d),
312            _ => None,
313        }
314    }
315
316    /// Whether this axis divides by anything but `1`.
317    pub(crate) fn is_rational(&self) -> bool {
318        !self.divisor.is_unit()
319    }
320
321    /// The physical cell the logical origin lands on, `⌊offset / divisor⌋`; `None` if dynamic.
322    pub(crate) fn origin(&self) -> Option<isize> {
323        match (self.offset, self.divisor) {
324            (Offset::Static(o), Divisor::Static(d)) => Some(o.div_euclid(d as isize)),
325            _ => None,
326        }
327    }
328
329    /// The division's starting phase, `offset - divisor * origin`; `None` if dynamic.
330    pub(crate) fn residue(&self) -> Option<usize> {
331        match (self.offset, self.divisor) {
332            (Offset::Static(o), Divisor::Static(d)) => Some(o.rem_euclid(d as isize) as usize),
333            _ => None,
334        }
335    }
336
337    /// How many of this axis's coefficients are [`Dynamic`](Scale::Dynamic).
338    pub(crate) fn dynamic_scale_count(&self) -> usize {
339        self.terms.iter().filter(|t| t.scale.is_dynamic()).count()
340    }
341
342    /// Whether any of this axis's coefficients are [`Dynamic`](Scale::Dynamic).
343    pub(crate) fn has_dynamic_scale(&self) -> bool {
344        self.terms.iter().any(|t| t.scale.is_dynamic())
345    }
346
347    /// Whether `axis` addresses this physical axis at all.
348    pub(crate) fn addresses(&self, axis: Axis) -> bool {
349        self.terms.iter().any(|t| t.axis == axis)
350    }
351
352    /// `axis`'s coefficient, `0` when it does not address this axis. Panics if it is dynamic.
353    pub fn scale(&self, axis: Axis) -> usize {
354        self.terms
355            .iter()
356            .find(|t| t.axis == axis)
357            .map_or(0, |t| t.scale.get())
358    }
359
360    /// The one [`Axis`] this map is the identity of, `None` otherwise.
361    pub(crate) fn identity_axis(&self) -> Option<Axis> {
362        match self.terms.as_slice() {
363            [
364                AxisTerm {
365                    axis,
366                    scale: Scale::Static(1),
367                },
368            ] if self.offset == Offset::Static(0) && self.divisor.is_unit() => Some(*axis),
369            _ => None,
370        }
371    }
372
373    /// The axes this buffer dim is made of, when it is one index over several axes laid end to
374    /// end, like `K = KB·16 + KI` over `KI` and `KB`. Listed finest first: the finest at scale
375    /// `1`, each next one at the product of the extents before it, no offset, no division. Such a
376    /// dim stays in the buffer whenever each of its axes stays in its extent. A single axis at
377    /// scale `1` is the simplest case ([`identity_axis`](Self::identity_axis)). `None` for any
378    /// other map, or when an axis other than the last has no comptime extent in `space`.
379    pub(crate) fn flattened_axes(
380        &self,
381        space: &Space,
382    ) -> Option<SmallVec<[Axis; Space::MAX_RANK]>> {
383        if self.terms.is_empty() || self.offset != Offset::Static(0) || !self.divisor.is_unit() {
384            return None;
385        }
386        let mut terms = self
387            .terms
388            .iter()
389            .map(|term| match term.scale {
390                Scale::Static(scale) => Some((term.axis, scale)),
391                Scale::Dynamic { .. } => None,
392            })
393            .collect::<Option<SmallVec<[(Axis, usize); Space::MAX_RANK]>>>()?;
394        terms.sort_by_key(|&(_, scale)| scale);
395        let mut next_scale = Some(1);
396        for &(axis, scale) in &terms {
397            if next_scale != Some(scale) || !space.contains(axis) {
398                return None;
399            }
400            next_scale = match space.extent_raw(axis) {
401                Extent::Static(extent) => Some(scale * extent),
402                Extent::Dynamic => None,
403            };
404        }
405        Some(terms.iter().map(|&(axis, _)| axis).collect())
406    }
407
408    /// Whether this physical axis is exactly `axis` at coefficient `1`, offset `0`, no division.
409    pub(crate) fn is_identity(&self, axis: Axis) -> bool {
410        self.identity_axis() == Some(axis)
411    }
412}
413
414#[cfg(test)]
415mod tests {
416    use super::*;
417
418    const A: Axis = Axis(0);
419    const B: Axis = Axis(1);
420
421    #[test]
422    fn a_plain_affine_term_steps_by_its_coefficient() {
423        let map = PhysicalAxisMap::affine(&[(A, 2), (B, 3)]);
424        assert_eq!(map.additive_step(0), Some(2));
425        assert_eq!(map.additive_step(1), Some(3));
426    }
427
428    #[test]
429    fn a_dynamic_coefficient_has_no_static_step() {
430        let map = PhysicalAxisMap::scaled(&[(A, Scale::Static(2)), (B, Scale::Dynamic { max: 4 })]);
431        assert_eq!(map.additive_step(1), None);
432    }
433
434    /// `⌊(x + m·4)/2⌋` moves by `2` per `m`.
435    #[test]
436    fn a_divisible_term_steps_by_what_the_floor_factors_out() {
437        let map = PhysicalAxisMap::affine(&[(A, 3), (B, 4)]).over(2);
438        assert_eq!(map.additive_step(1), Some(2));
439    }
440
441    #[test]
442    #[should_panic(expected = "stays inside this axis's floor")]
443    fn an_indivisible_term_under_a_floor_cannot_be_stepped() {
444        let map = PhysicalAxisMap::affine(&[(A, 3), (B, 4)]).over(2);
445        map.additive_step(0);
446    }
447
448    #[test]
449    fn identity_is_exactly_one_axis_at_coefficient_one() {
450        let id = PhysicalAxisMap::of(A);
451        assert!(id.is_identity(A));
452        assert!(!id.is_identity(B));
453        assert_eq!(id.scale(A), 1);
454        assert_eq!(id.scale(B), 0);
455        assert_eq!(id.offset(), Offset::Static(0));
456
457        let affine = PhysicalAxisMap::affine(&[(A, 2), (B, 3)]);
458        assert!(!affine.is_identity(A));
459        assert_eq!(affine.scale(A), 2);
460        assert_eq!(affine.scale(B), 3);
461        assert_eq!(affine.offset(), Offset::Static(0));
462        assert!(!PhysicalAxisMap::affine(&[(A, 2)]).is_identity(A));
463        assert!(PhysicalAxisMap::affine(&[(A, 1)]).is_identity(A));
464
465        let with_offset = PhysicalAxisMap::affine(&[(A, 1)]).shifted(-2);
466        assert!(!with_offset.is_identity(A));
467        assert_eq!(with_offset.scale(A), 1);
468        assert_eq!(with_offset.offset(), Offset::Static(-2));
469
470        let with_dynamic_offset = PhysicalAxisMap::affine(&[(A, 1)]).shifted(Offset::Dynamic);
471        assert!(!with_dynamic_offset.is_identity(A));
472        assert_eq!(with_dynamic_offset.offset(), Offset::Dynamic);
473        assert!(with_dynamic_offset.offset().is_dynamic());
474        assert!(!with_dynamic_offset.has_dynamic_scale());
475        assert_eq!(with_dynamic_offset.dynamic_scale_count(), 0);
476    }
477
478    /// The axes come back finest first, whatever order the terms are written in. A gap or
479    /// overlap between them, an offset, a division, a dynamic scale, or a finer axis with no
480    /// comptime extent gives `None`.
481    #[test]
482    fn flattened_axes_come_finest_first() {
483        const C: Axis = Axis(2);
484        let space = Space::new(&[(A, 4), (B, 16), (C, 3)]);
485        let axes = |map: &PhysicalAxisMap| map.flattened_axes(&space).map(|axes| axes.to_vec());
486
487        assert_eq!(axes(&PhysicalAxisMap::of(A)), Some(vec![A]));
488        assert_eq!(
489            axes(&PhysicalAxisMap::affine(&[(A, 16), (B, 1)])),
490            Some(vec![B, A])
491        );
492        assert_eq!(
493            axes(&PhysicalAxisMap::affine(&[(C, 64), (A, 16), (B, 1)])),
494            Some(vec![B, A, C])
495        );
496
497        assert_eq!(axes(&PhysicalAxisMap::affine(&[(A, 32), (B, 1)])), None);
498        assert_eq!(axes(&PhysicalAxisMap::affine(&[(A, 8), (B, 1)])), None);
499        assert_eq!(axes(&PhysicalAxisMap::affine(&[(A, 16), (B, 2)])), None);
500        assert_eq!(axes(&PhysicalAxisMap::of(A).shifted(1)), None);
501        assert_eq!(axes(&PhysicalAxisMap::of(A).over(2)), None);
502        assert_eq!(
503            axes(&PhysicalAxisMap::scaled(&[(A, Scale::Dynamic { max: 4 })])),
504            None
505        );
506        assert_eq!(axes(&PhysicalAxisMap::of(Axis(7))), None);
507
508        let finer_dynamic = Space::new(&[(A, 4), (B, 16)]).with_dynamic(&[B]);
509        let blocked = PhysicalAxisMap::affine(&[(A, 16), (B, 1)]);
510        assert_eq!(blocked.flattened_axes(&finer_dynamic), None);
511        let coarser_dynamic = Space::new(&[(A, 4), (B, 16)]).with_dynamic(&[A]);
512        assert_eq!(
513            blocked
514                .flattened_axes(&coarser_dynamic)
515                .map(|axes| axes.to_vec()),
516            Some(vec![B, A])
517        );
518    }
519
520    #[test]
521    fn rational_axis_map_properties() {
522        let map = PhysicalAxisMap::affine(&[(A, 100)]).shifted(-50).over(133);
523        assert!(map.is_rational());
524        assert_eq!(map.divisor(), Divisor::Static(133));
525        assert_eq!(map.offset(), Offset::Static(-50));
526        assert_eq!(map.origin(), Some((-50isize).div_euclid(133)));
527        assert_eq!(map.residue(), Some((-50isize).rem_euclid(133) as usize));
528
529        let dynamic_div = PhysicalAxisMap::affine(&[(A, 100)]).over(Divisor::Dynamic { min: 3 });
530        assert!(dynamic_div.is_rational());
531        assert!(dynamic_div.divisor().is_dynamic());
532        assert_eq!(dynamic_div.origin(), None);
533        assert_eq!(dynamic_div.residue(), None);
534    }
535
536    /// A rational map's terms whose coefficient the divisor divides step exactly.
537    #[test]
538    fn a_rational_map_factors_exact_static_terms_into_offsets() {
539        let map = PhysicalAxisMap::affine(&[(A, 5), (B, 6)])
540            .shifted(-2)
541            .over(6);
542
543        assert!(map.is_rational());
544        assert_eq!(map.static_offset_step(0), None);
545        assert_eq!(map.static_offset_step(1), Some(1));
546        for output in 0..8isize {
547            for tap in 0..3isize {
548                assert_eq!(
549                    (output * 5 + tap * 6 - 2).div_euclid(6),
550                    (output * 5 - 2).div_euclid(6) + tap
551                );
552            }
553        }
554
555        let strided_tap = PhysicalAxisMap::affine(&[(A, 5), (B, 12)]).over(6);
556        assert_eq!(strided_tap.static_offset_step(1), Some(2));
557
558        let fractional_tap = PhysicalAxisMap::affine(&[(A, 5), (B, 2)]).over(3);
559        assert_eq!(fractional_tap.static_offset_step(1), None);
560
561        let dynamic_divisor =
562            PhysicalAxisMap::affine(&[(A, 5), (B, 6)]).over(Divisor::Dynamic { min: 3 });
563        assert_eq!(dynamic_divisor.static_offset_step(1), None);
564    }
565
566    #[test]
567    #[should_panic(expected = "divisor of 0 does not map anywhere")]
568    fn divisor_zero_panics() {
569        PhysicalAxisMap::affine(&[(A, 1)]).over(0);
570    }
571
572    #[test]
573    #[should_panic(expected = "state the whole divisor once")]
574    fn a_second_over_is_refused() {
575        PhysicalAxisMap::affine(&[(A, 4), (B, 6)]).over(4).over(3);
576    }
577
578    #[test]
579    fn a_divisor_the_coefficients_cancel_reduces_away() {
580        let map = PhysicalAxisMap::affine(&[(A, 8), (B, 4)]).over(4);
581        assert!(!map.is_rational());
582        assert_eq!(map.divisor(), Divisor::Static(1));
583        assert_eq!(map.scale(A), 2);
584        assert_eq!(map.scale(B), 1);
585        assert_eq!(map.offset(), Offset::Static(0));
586
587        assert!(PhysicalAxisMap::affine(&[(A, 4)]).over(4).is_identity(A));
588    }
589
590    /// `⌊(8a - 3)/4⌋` is `2a - 1`.
591    #[test]
592    fn reducing_floors_the_offset() {
593        let map = PhysicalAxisMap::affine(&[(A, 8)]).shifted(-3).over(4);
594        assert!(!map.is_rational());
595        assert_eq!(map.scale(A), 2);
596        assert_eq!(map.offset(), Offset::Static(-1));
597        assert_eq!(map.origin(), Some(-1));
598        assert_eq!(map.residue(), Some(0));
599    }
600
601    #[test]
602    fn a_divisor_the_coefficients_do_not_cancel_stays() {
603        assert!(
604            PhysicalAxisMap::affine(&[(A, 4), (B, 6)])
605                .over(6)
606                .is_rational()
607        );
608        assert!(
609            PhysicalAxisMap::scaled(&[(A, Scale::Dynamic { max: 2 }), (B, Scale::Static(4))])
610                .over(4)
611                .is_rational()
612        );
613        assert!(
614            PhysicalAxisMap::scaled(&[(A, Scale::Static(4))])
615                .shifted(Offset::Dynamic)
616                .over(4)
617                .is_rational()
618        );
619    }
620
621    #[test]
622    fn the_same_coefficients_partition_or_overlap() {
623        assert_eq!(
624            PhysicalAxisMap::disjoint(&[(A, 2), (B, 1)]).composition(),
625            Composition::Disjoint
626        );
627        assert_eq!(
628            PhysicalAxisMap::affine(&[(A, 2), (B, 1)]).composition(),
629            Composition::Overlapping
630        );
631    }
632
633    #[test]
634    fn the_identity_partitions_whichever_door_minted_it() {
635        assert_eq!(PhysicalAxisMap::of(A).composition(), Composition::Disjoint);
636        assert_eq!(
637            PhysicalAxisMap::affine(&[(A, 1)]).composition(),
638            Composition::Disjoint
639        );
640        assert_eq!(
641            PhysicalAxisMap::scaled(&[(A, Scale::Static(1))]).composition(),
642            Composition::Disjoint
643        );
644        assert_eq!(
645            PhysicalAxisMap::affine(&[(A, 2)]).composition(),
646            Composition::Overlapping
647        );
648    }
649
650    #[test]
651    fn a_coarsening_is_never_a_partition() {
652        assert_eq!(
653            PhysicalAxisMap::of(A).over(4).composition(),
654            Composition::Overlapping
655        );
656    }
657
658    #[test]
659    fn a_partition_claims_its_radices() {
660        let map = PhysicalAxisMap::disjoint(&[(A, 32), (B, 1)]);
661        assert_eq!(&map.claimed_radices()[..], &[(A, 32), (B, 1)]);
662        assert!(
663            PhysicalAxisMap::affine(&[(A, 32), (B, 1)])
664                .claimed_radices()
665                .is_empty()
666        );
667    }
668}
cubecl_tile/layout/projection/axis_map.rs

cubecl_tile/layout/projection/
axis_map.rs
