cubek_std/tile/
scheduler.rs

1use cubecl::prelude::*;
2
3use crate::PartitionSize;
4
5/// Defines how partition indices are scheduled across axes.
6#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
7pub enum PartitionSchedulerScheme {
8    /// Rotates indices per axis to stagger partitions and reduce shared memory conflicts.
9    Offset,
10    /// Maps partitions in simple row-major order without rotations.
11    Naive,
12}
13
14/// Schedules global indices for M, N, and K axes in a partitioned matmul.
15/// Internally uses an `AxisScheduler` per axis.
16#[derive(CubeType)]
17pub struct PartitionScheduler {
18    pub m: AxisScheduler,
19    pub n: AxisScheduler,
20    pub k: AxisScheduler,
21}
22
23#[cube]
24impl PartitionScheduler {
25    /// Creates a `PartitionScheduler` for a partition at (partition_index_m, partition_index_n).
26    ///
27    /// - `Offset`: rotates indices per axis to spread memory access and avoid conflicts.
28    /// - `Naive`: uses simple row-major order for partitions.
29    pub fn new(
30        partition_index_m: u32,
31        partition_index_n: u32,
32        #[comptime] partition_size: PartitionSize,
33        #[comptime] partition_schedule_scheme: PartitionSchedulerScheme,
34    ) -> PartitionScheduler {
35        match partition_schedule_scheme {
36            PartitionSchedulerScheme::Offset => {
37                // M-axis rotation: ensures partitions in the same row start at different M tiles.
38                let m_offset = (partition_index_n / partition_size.k()) % partition_size.m();
39
40                // N-axis rotation: ensures partitions in the same column start at different N tiles.
41                let n_offset = (partition_index_m / partition_size.k()) % partition_size.n();
42
43                // K-axis rotation: simple offset; same diagonal can share K safely.
44                let k_offset = (partition_index_m + partition_index_n) % partition_size.k();
45
46                PartitionScheduler {
47                    m: AxisScheduler::new_Offset(OffsetAxisScheduler::new(
48                        m_offset,
49                        partition_index_m,
50                        partition_size.m(),
51                    )),
52                    n: AxisScheduler::new_Offset(OffsetAxisScheduler::new(
53                        n_offset,
54                        partition_index_n,
55                        partition_size.n(),
56                    )),
57                    k: AxisScheduler::new_Offset(OffsetAxisScheduler::new(
58                        k_offset,
59                        0u32,
60                        partition_size.k(),
61                    )),
62                }
63            }
64            PartitionSchedulerScheme::Naive => PartitionScheduler {
65                m: AxisScheduler::new_Naive(NaiveAxisScheduler::new(
66                    partition_index_m,
67                    partition_size.m(),
68                )),
69                n: AxisScheduler::new_Naive(NaiveAxisScheduler::new(
70                    partition_index_n,
71                    partition_size.n(),
72                )),
73                k: AxisScheduler::new_Naive(NaiveAxisScheduler::new(0u32, partition_size.k())),
74            },
75        }
76    }
77
78    /// Maps a local M index to a global index.
79    pub fn map_m(&self, i: u32) -> u32 {
80        self.m.map(i)
81    }
82
83    /// Maps a local N index to a global index.
84    pub fn map_n(&self, i: u32) -> u32 {
85        self.n.map(i)
86    }
87
88    /// Maps a local K index to a global index.
89    pub fn map_k(&self, i: u32) -> u32 {
90        self.k.map(i)
91    }
92}
93
94/// Axis-specific scheduler that delegates to either `OffsetAxisScheduler` or `NaiveAxisScheduler`.
95#[derive(CubeType)]
96#[allow(unused)]
97pub enum AxisScheduler {
98    Offset(OffsetAxisScheduler),
99    Naive(NaiveAxisScheduler),
100}
101
102/// Schedules indices for one axis with rotation and wrapping.
103///
104/// Combines:
105/// - `inner_offset`: rotation inside this partition.
106/// - `outer_offset`: global shift for skipping previous partitions.
107#[derive(CubeType)]
108pub struct OffsetAxisScheduler {
109    inner_offset: u32,
110    outer_offset: u32,
111    #[cube(comptime)]
112    len: u32,
113}
114
115/// Schedules indices for one axis in row-major order.
116/// Just adds a global shift based on the partition index.
117#[derive(CubeType)]
118pub struct NaiveAxisScheduler {
119    outer_offset: u32,
120}
121
122#[cube]
123impl AxisScheduler {
124    pub fn map(&self, i: u32) -> u32 {
125        match self {
126            AxisScheduler::Offset(offset_axis_scheduler) => offset_axis_scheduler.map(i),
127            AxisScheduler::Naive(naive_axis_scheduler) => naive_axis_scheduler.map(i),
128        }
129    }
130}
131
132#[cube]
133impl OffsetAxisScheduler {
134    pub fn new(
135        inner_offset: u32,
136        partition_index: u32,
137        #[comptime] len: u32,
138    ) -> OffsetAxisScheduler {
139        let outer_offset = partition_index * len;
140        OffsetAxisScheduler {
141            inner_offset,
142            outer_offset,
143            len,
144        }
145    }
146
147    pub fn map(&self, i: u32) -> u32 {
148        let relative = (i + self.inner_offset) % self.len;
149        relative + self.outer_offset
150    }
151}
152
153#[cube]
154impl NaiveAxisScheduler {
155    pub fn new(partition_index: u32, #[comptime] len: u32) -> NaiveAxisScheduler {
156        let outer_offset = partition_index * len;
157        NaiveAxisScheduler { outer_offset }
158    }
159
160    pub fn map(&self, i: u32) -> u32 {
161        i + self.outer_offset
162    }
163}
164
165/// Number of buffers held for the rhs side of a partition matmul. Single
166/// buffering keeps one fragment; double buffering keeps two and rotates
167/// to overlap loads with computes.
168#[derive(Default, Clone, Copy, PartialEq, Eq, Hash, Debug)]
169pub enum PartitionBuffering {
170    Single,
171    #[default]
172    Double,
173}
cubek_std/tile/scheduler.rs

cubek_std/tile/
scheduler.rs
