cubek_std/tile/ops/
partition.rs

1//! `Tile::partition` and the [`Partitioner`] strategies. A partitioner takes
2//! a tile at one [`TileScope`] and yields a per-primitive view at a lower
3//! scope.
4
5use std::marker::PhantomData;
6
7use cubecl::{prelude::*, std::tensor::layout::Coords2d};
8
9use crate::tile::{
10    PartitionTile, Plane, Tile, TileExpand, TileKind, TileKindExpand, TileScope, Unit,
11};
12
13/// Maps the current compute primitive to `(row, col)` in a partition grid.
14#[cube]
15pub trait Partitioner: 'static + Send + Sync {
16    type OutputScope: TileScope;
17
18    fn coordinates(
19        compute_index: u32,
20        #[comptime] plane_dim: u32,
21        #[comptime] num_partitions_col: u32,
22    ) -> Coords2d;
23}
24
25/// Per-unit views of a higher-scope tile.
26#[derive(Clone, Copy)]
27pub struct UnitPartitioner;
28
29#[cube]
30impl Partitioner for UnitPartitioner {
31    type OutputScope = Unit;
32
33    fn coordinates(
34        compute_index: u32,
35        #[comptime] plane_dim: u32,
36        #[comptime] num_partitions_col: u32,
37    ) -> Coords2d {
38        let absolute_index = UNIT_POS_X + plane_dim * compute_index;
39
40        (
41            absolute_index / num_partitions_col,
42            absolute_index % num_partitions_col,
43        )
44    }
45}
46
47/// Per-plane views of a higher-scope tile.
48#[derive(Clone, Copy)]
49pub struct PlanePartitioner;
50
51#[cube]
52impl Partitioner for PlanePartitioner {
53    type OutputScope = Plane;
54
55    fn coordinates(
56        compute_index: u32,
57        #[comptime] _plane_dim: u32,
58        #[comptime] num_partitions_col: u32,
59    ) -> Coords2d {
60        (
61            compute_index / num_partitions_col,
62            compute_index % num_partitions_col,
63        )
64    }
65}
66
67#[cube]
68impl<N: Numeric, Sc: TileScope> Tile<N, Sc> {
69    /// View of `self` at the partitioner's output scope. Source must be a
70    /// `TileKind::Stage`.
71    pub fn partition<P: Partitioner>(
72        &self,
73        compute_index: u32,
74        #[comptime] plane_dim: u32,
75        #[comptime] num_partitions_col: u32,
76    ) -> Tile<N, P::OutputScope> {
77        let (p_row, p_col) = P::coordinates(compute_index, plane_dim, num_partitions_col);
78        match &self.kind {
79            TileKind::Stage(stage) => {
80                let m_tiles = comptime!(stage.config.tiles_per_partition_along_row);
81                let n_tiles = comptime!(stage.config.tiles_per_partition_along_col);
82
83                let mut tiles = Sequence::new();
84
85                #[unroll]
86                for m in 0..m_tiles {
87                    #[unroll]
88                    for n in 0..n_tiles {
89                        let global = (p_row * m_tiles + m, p_col * n_tiles + n);
90                        let shared = stage.get_tile(global);
91                        tiles.push(Tile::<N, P::OutputScope>::new_SharedTile(shared));
92                    }
93                }
94
95                Tile::new_Partition(PartitionTile::<N, P::OutputScope> {
96                    tiles,
97                    rows: m_tiles,
98                    cols: n_tiles,
99                    _phantom: PhantomData,
100                })
101            }
102            TileKind::SharedTile(_)
103            | TileKind::Cmma(_)
104            | TileKind::Mma(_)
105            | TileKind::Register(_)
106            | TileKind::PlaneVec(_)
107            | TileKind::Interleaved(_)
108            | TileKind::Unit(_)
109            | TileKind::WhiteboxFragment(_)
110            | TileKind::RowWise(_)
111            | TileKind::Bounce(_)
112            | TileKind::Partition(_)
113            | TileKind::Pipelined(_)
114            | TileKind::None => {
115                panic!("Tile::partition: source variant cannot be partitioned")
116            }
117        }
118    }
119}
cubek_std/tile/ops/partition.rs

cubek_std/tile/ops/
partition.rs
