cubek_std/tile/
base.rs

1#![allow(non_snake_case)]
2
3use std::marker::PhantomData;
4
5use cubecl::prelude::*;
6
7use crate::tile::{
8    BounceTile, CmmaTile, InterleavedTile, MmaTile, PartitionTile, PipelinedTile, PlaneVecTile,
9    RegisterTile, RowWise, ScopeMarker, SharedTile, StageTile, TileScope, UnitTile,
10    WhiteboxFragment, variants::stage::partition::partition_get_at_mut,
11};
12
13/// Public tile type. Wraps a [`TileKind`] payload; the inner enum is
14/// crate-private and callers construct via `Tile::new_*`.
15#[derive(CubeType)]
16pub struct Tile<N: Numeric, Sc: TileScope> {
17    pub(crate) kind: TileKind<N, Sc>,
18    pub(crate) _scope: ScopeMarker<Sc>,
19}
20
21/// Storage variants of a tile.
22#[derive(CubeType)]
23#[allow(dead_code)]
24pub(crate) enum TileKind<N: Numeric, Sc: TileScope> {
25    /// Whole-stage view, used for partition-level dispatch.
26    Stage(StageTile<N>),
27    /// Sequence of per-tile accumulators.
28    Partition(PartitionTile<N, Sc>),
29    /// Stage slot exposed as a tile (no distribution, no compute).
30    SharedTile(SharedTile<N>),
31
32    /// CMMA fragment.
33    Cmma(CmmaTile<N>),
34    /// MMA fragment; operand role (Lhs/Rhs/Acc) carried inside.
35    Mma(MmaTile<N>),
36    /// Register-resident tile for the software register matmul.
37    Register(RegisterTile<N>),
38    /// Plane-vector matmul tile.
39    PlaneVec(PlaneVecTile<N>),
40    /// Plane-interleaved-on-k matmul tile.
41    Interleaved(InterleavedTile<N>),
42    /// Per-unit register array. `Sc = Unit`.
43    Unit(UnitTile<N>),
44    /// Plane-exposed fragment with a visible layout. `Sc = Plane`.
45    WhiteboxFragment(WhiteboxFragment<N>),
46    /// Per-row vector tile (softmax max/sum state). `Sc = Plane`.
47    RowWise(RowWise<N>),
48
49    /// Rhs fragments for the partition matmul (1 = single-buffered, 2 = double).
50    Pipelined(PipelinedTile<N, Sc>),
51    /// CMMA fragment + smem scratch + whitebox view. `Sc = Plane`.
52    Bounce(BounceTile<N>),
53
54    /// Sentinel for zero-init via `copy_from`.
55    None,
56}
57
58#[cube]
59impl<N: Numeric, Sc: TileScope> Tile<N, Sc> {
60    pub(crate) fn from_kind(kind: TileKind<N, Sc>) -> Tile<N, Sc> {
61        Tile::<N, Sc> {
62            kind,
63            _scope: ScopeMarker::<Sc> {
64                _phantom: PhantomData,
65            },
66        }
67    }
68
69    pub fn new_SharedTile(t: SharedTile<N>) -> Tile<N, Sc> {
70        Self::from_kind(TileKind::new_SharedTile(t))
71    }
72
73    pub fn new_Stage(t: StageTile<N>) -> Tile<N, Sc> {
74        Self::from_kind(TileKind::new_Stage(t))
75    }
76
77    pub fn new_Partition(t: PartitionTile<N, Sc>) -> Tile<N, Sc> {
78        Self::from_kind(TileKind::new_Partition(t))
79    }
80
81    pub fn new_Pipelined(t: PipelinedTile<N, Sc>) -> Tile<N, Sc> {
82        Self::from_kind(TileKind::new_Pipelined(t))
83    }
84
85    pub fn new_None() -> Tile<N, Sc> {
86        Self::from_kind(TileKind::new_None())
87    }
88
89    pub fn new_RowWise(t: RowWise<N>) -> Tile<N, Sc> {
90        Self::from_kind(TileKind::new_RowWise(t))
91    }
92
93    /// Mutable reference to the `(m, n)` element of a `Partition` tile.
94    pub fn partition_tile_at_mut(
95        &mut self,
96        #[comptime] m: usize,
97        #[comptime] n: usize,
98        #[comptime] n_cols: usize,
99    ) -> &mut Tile<N, Sc> {
100        match &mut self.kind {
101            TileKind::Partition(p) => partition_get_at_mut::<N, Sc>(p, m, n, n_cols),
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::Stage(_)
113            | TileKind::Pipelined(_)
114            | TileKind::None => {
115                panic!("Tile::partition_tile_at_mut: self.kind is not Partition")
116            }
117        }
118    }
119}
cubek_std/tile/base.rs

cubek_std/tile/
base.rs
