cubecl_tile/space/partition/
count.rs

1//! [`Count`]: how many tiles a level takes along one of its axes.
2
3use crate::Extent;
4
5/// How many tiles a level takes along one of its axes.
6#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
7pub enum Count {
8    /// This many, stated: a walk's steps, or a scope's workers with one tile each.
9    Stated(usize),
10    /// Every tile the level above hands down; the one count whose tile can overhang.
11    All,
12    /// Every tile, distributed across this many workers in runs.
13    AllAcross(usize),
14    /// This many tiles, taken in turns by however many units the launch runs.
15    Distributed(usize),
16}
17
18/// What a walk counts along one axis of a level: a constant, or the extent handed down.
19#[derive(Clone, Copy, PartialEq, Eq, Debug)]
20pub(crate) enum GridCount {
21    Const(usize),
22    Extent(usize),
23}
24
25impl Count {
26    /// The workers the tiles are distributed to, where that is a stated number.
27    pub(crate) fn stated(self) -> Option<usize> {
28        match self {
29            Count::Stated(n) | Count::AllAcross(n) => Some(n),
30            Count::All | Count::Distributed(_) => None,
31        }
32    }
33
34    /// Tiles this count takes along an axis of `extent`; a partial last tile counts.
35    pub(crate) fn tiles(self, extent: usize, tile: usize) -> usize {
36        match self {
37            Count::Stated(n) | Count::Distributed(n) => n,
38            Count::All | Count::AllAcross(_) => extent.div_ceil(tile),
39        }
40    }
41
42    /// [`tiles`](Count::tiles) where the host can prove it.
43    pub(crate) fn tiles_const(self, extent: Extent, tile: usize) -> Option<usize> {
44        match (self, extent) {
45            (Count::Stated(n) | Count::Distributed(n), _) => Some(n),
46            (_, Extent::Static(extent)) => Some(self.tiles(extent, tile)),
47            (_, Extent::Dynamic) => None,
48        }
49    }
50
51    /// What a walk counts along an axis stepped in `tile`.
52    pub(crate) fn grid(self, tile: usize) -> GridCount {
53        match self {
54            Count::Stated(n) | Count::Distributed(n) => GridCount::Const(n),
55            Count::All | Count::AllAcross(_) => GridCount::Extent(tile),
56        }
57    }
58}
59
60/// A count a units level states along an axis, as the units it takes.
61impl From<Count> for usize {
62    fn from(count: Count) -> usize {
63        count
64            .stated()
65            .expect("a units level's count is stated; every tile is the launch's")
66    }
67}
cubecl_tile/space/partition/count.rs

cubecl_tile/space/partition/
count.rs
