Skip to main content

cubek_std/tile/
base.rs

1use cubecl::prelude::*;
2
3use crate::tile::{
4    BounceTile, CmmaTile, InterleavedTile, MmaTile, PlaneVecTile, RegisterTile, ScopeMarker,
5    SharedTile, TileScope, UnitTile, WhiteboxFragment,
6};
7
8#[derive(CubeType)]
9pub enum Tile<N: Numeric, Sc: TileScope, IO: SliceVisibility> {
10    SharedMemory(SharedTile<N, IO>),
11    Cmma(CmmaTile<N>),
12    Mma(MmaTile<N>),
13    Register(RegisterTile<N>),
14    PlaneVec(PlaneVecTile<N>),
15    Interleaved(InterleavedTile<N>),
16    /// Each unit holds a full row-major copy of the tile in registers.
17    /// Only valid when `Sc = Unit`.
18    Unit(UnitTile<N>),
19    /// The tile is fragmented across plane units, with the layout exposed.
20    /// Only valid when `Sc = Plane`.
21    WhiteboxFragment(WhiteboxFragment<N>),
22    /// Bundles a cmma fragment, an smem scratch slice, and a `WhiteboxFragment` view.
23    /// From the caller's perspective it is a single tile; the smem round-trip
24    /// is internal to ops dispatch. Only valid when `Sc = Plane`.
25    Bounce(BounceTile<N>),
26    Broadcasted(Value<N>),
27    None,
28    _Phantom(ScopeMarker<Sc>),
29}
30
31/// Wrapper over val to make enum work
32#[derive(CubeType)]
33pub struct Value<E: Numeric> {
34    pub val: E,
35}