cubecl_tile/space/walk/
portion.rs

1//! This instance's share of a level distributed as one flat index.
2
3use cubecl::prelude::*;
4
5use crate::{ComputeScope, Coverage, CubeAxis, Level, Region, Walk, WalkExpand};
6
7/// The regions this instance touches and how much of the first and last is its own.
8/// Counted in the level below's joint steps.
9///
10/// The steps are laid end to end and cut into one run an instance, every run `ceil(steps /
11/// instances)` long but the last that holds any, which holds what is left: the runs holding steps
12/// come first, so no run between two that hold steps is empty, and the longest run, which is what
13/// the level takes, is as short as an even split's.
14#[derive(CubeType)]
15pub struct Portion {
16    /// The regions the portion touches, in order.
17    regions: Walk,
18    /// The first of them, in the level's flat index.
19    first: usize,
20    /// The portion's bounds on the joint step index.
21    start: usize,
22    end: usize,
23    /// Steps of the level below per region.
24    stride: usize,
25    /// Steps every run holds but the last that holds any.
26    run_length: usize,
27    /// This instance's run, counted from the first.
28    run: usize,
29}
30
31#[cube]
32impl Portion {
33    /// How many regions the portion touches.
34    pub fn touched(&self) -> usize {
35        self.regions.total()
36    }
37
38    /// The `i`-th region the portion touches.
39    pub fn region(&self, i: usize) -> Region {
40        self.regions.region(i)
41    }
42
43    /// The `i`-th region the portion touches, in the level's flat index: one for every region of
44    /// the level, whichever instance touches it.
45    pub fn index(&self, i: usize) -> usize {
46        self.first + i
47    }
48
49    /// Start and length of this portion's steps within the `i`-th region's walk below.
50    pub fn steps(&self, i: usize) -> (usize, usize) {
51        let base = (self.first + i) * self.stride;
52        let from = select(base < self.start, self.start - base, 0);
53        let to = select(self.end < base + self.stride, self.end - base, self.stride);
54        (from, to - from)
55    }
56
57    /// The first run holding steps of the `i`-th region the portion touches, and how many runs
58    /// hold them. This run is one of them.
59    pub fn holders(&self, i: usize) -> (usize, usize) {
60        let base = (self.first + i) * self.stride;
61        let first_run = self.run_holding(base);
62        let last_run = self.run_holding(base + self.stride - 1);
63        (first_run, last_run + 1 - first_run)
64    }
65
66    /// This instance's run, counted from the first.
67    pub fn run(&self) -> usize {
68        self.run
69    }
70
71    /// Whether `run`, one of the [`holders`](Portion::holders) of the `i`-th region the portion
72    /// touches, starts there rather than in a region before it.
73    pub fn starts_in(&self, i: usize, run: usize) -> bool {
74        run * self.run_length >= (self.first + i) * self.stride
75    }
76
77    /// The run holding `step`.
78    fn run_holding(&self, step: usize) -> usize {
79        step / self.run_length
80    }
81}
82
83#[cube]
84impl Walk {
85    /// This instance's portion of the level's flat index, counted in steps of `below`.
86    pub fn portion(self, #[comptime] below: Level) -> Portion {
87        let instances = comptime!(
88            self.level
89                .shared_by()
90                .expect("Walk::portion: this level distributes no grid as one; say `shared_by`")
91        );
92        let stride = self.region(0).over(&below).total();
93        let steps = self.total() * stride;
94        let run = match comptime!(self.level.coverage()) {
95            Coverage::Distribute(ComputeScope::Cube) => CubeAxis::position(CubeAxis::X),
96            Coverage::Distribute(ComputeScope::Plane) => {
97                ComputeScope::position(ComputeScope::Plane)
98            }
99            Coverage::Distribute(ComputeScope::PlaneGroup { planes }) => {
100                ComputeScope::position(comptime!(ComputeScope::PlaneGroup { planes }))
101            }
102            Coverage::Distribute(ComputeScope::Unit) | Coverage::Walk => {
103                panic!(
104                    "Walk::portion: only cubes, plane groups or planes share a grid as one index"
105                )
106            }
107        };
108        let run_length = steps.div_ceil(instances);
109        let start = (run * run_length).min(steps);
110        let end = (start + run_length).min(steps);
111        let first = start / stride;
112        // `end` is exclusive; an empty portion touches nothing.
113        let touched = select(start < end, (end - 1) / stride + 1 - first, 0);
114        Portion {
115            regions: self.range(first, touched),
116            first,
117            start,
118            end,
119            stride,
120            run_length,
121            run,
122        }
123    }
124}
cubecl_tile/space/walk/portion.rs

cubecl_tile/space/walk/
portion.rs
