cubek_std/tile/variants/instruction/cmma/
base.rs

1use cubecl;
2use cubecl::{
3    cmma::{self},
4    prelude::*,
5};
6
7use crate::{
8    MatrixLayout, StageIdent, TileSize, as_cmma_layout,
9    tile::{
10        SharedTile, Tile, TileKind, TileKindExpand, TileScope,
11        variants::instruction::cmma::{CmmaStageWriter, cmma_load_strided},
12    },
13};
14
15#[derive(CubeType)]
16pub struct CmmaTile<N: Numeric> {
17    pub matrix: cmma::Matrix<N>,
18    #[cube(comptime)]
19    pub matrix_layout: MatrixLayout,
20    #[cube(comptime)]
21    pub tile_size: TileSize,
22}
23
24#[cube]
25impl<E: Float> CmmaTile<E> {
26    pub fn fill_zero(&mut self) {
27        cubecl::cmma::fill(&mut self.matrix, E::from_int(0));
28    }
29}
30
31#[cube]
32impl<A: Numeric> CmmaTile<A> {
33    /// Executes `lhs · rhs`, accumulating into `self`.
34    pub fn mma<L: Numeric, R: Numeric>(&mut self, lhs: &CmmaTile<L>, rhs: &CmmaTile<R>) {
35        cmma_execute(&lhs.matrix, &rhs.matrix, &mut self.matrix);
36    }
37}
38
39#[cube]
40impl<N: Numeric> CmmaTile<N> {
41    /// Supported sources: `SharedTile` (load) and `None` (zero-init).
42    pub fn copy_from<SE: Numeric, SS: Size, Sc: TileScope>(
43        &mut self,
44        source: &Tile<SE, Sc>,
45        #[comptime] ident: StageIdent,
46    ) {
47        match &source.kind {
48            TileKind::SharedTile(shared) => {
49                cmma_load_from_shared::<SE, SS, N>(
50                    shared,
51                    &mut self.matrix,
52                    ident,
53                    self.matrix_layout,
54                );
55            }
56            TileKind::None => cmma_load_zeros::<N>(&mut self.matrix),
57            TileKind::Cmma(_)
58            | TileKind::Mma(_)
59            | TileKind::Register(_)
60            | TileKind::PlaneVec(_)
61            | TileKind::Interleaved(_)
62            | TileKind::Unit(_)
63            | TileKind::WhiteboxFragment(_)
64            | TileKind::RowWise(_)
65            | TileKind::Bounce(_)
66            | TileKind::Stage(_)
67            | TileKind::Partition(_)
68            | TileKind::Pipelined(_) => panic!("CmmaTile::copy_from: unsupported source variant"),
69        }
70    }
71
72    pub fn init_zero(&mut self) {
73        cmma_load_zeros::<N>(&mut self.matrix);
74    }
75}
76
77#[cube]
78pub fn cmma_allocate_lhs<L: Numeric, Sc: TileScope>(
79    #[comptime] layout: MatrixLayout,
80    #[comptime] tile_size: TileSize,
81) -> Tile<L, Sc> {
82    let fragment = unsafe {
83        cmma::Matrix::<L>::uninitialized(
84            cmma::MatrixIdent::A,
85            tile_size.m as usize,
86            tile_size.n as usize,
87            tile_size.k as usize,
88            as_cmma_layout(layout),
89        )
90    };
91    Tile::from_kind(TileKind::new_Cmma(CmmaTile::<L> {
92        matrix: fragment,
93        matrix_layout: layout,
94        tile_size,
95    }))
96}
97
98#[cube]
99pub fn cmma_allocate_rhs<R: Numeric, Sc: TileScope>(
100    #[comptime] layout: MatrixLayout,
101    #[comptime] tile_size: TileSize,
102) -> Tile<R, Sc> {
103    let fragment = unsafe {
104        cmma::Matrix::<R>::uninitialized(
105            cmma::MatrixIdent::B,
106            tile_size.m as usize,
107            tile_size.n as usize,
108            tile_size.k as usize,
109            as_cmma_layout(layout),
110        )
111    };
112    Tile::from_kind(TileKind::new_Cmma(CmmaTile::<R> {
113        matrix: fragment,
114        matrix_layout: layout,
115        tile_size,
116    }))
117}
118
119#[cube]
120pub fn cmma_allocate_acc<A: Numeric, Sc: TileScope>(
121    #[comptime] layout: MatrixLayout,
122    #[comptime] tile_size: TileSize,
123) -> Tile<A, Sc> {
124    let fragment = unsafe {
125        cmma::Matrix::<A>::uninitialized(
126            cmma::MatrixIdent::Accumulator,
127            tile_size.m as usize,
128            tile_size.n as usize,
129            tile_size.k as usize,
130            cmma::MatrixLayout::Undefined,
131        )
132    };
133    Tile::from_kind(TileKind::new_Cmma(CmmaTile::<A> {
134        matrix: fragment,
135        matrix_layout: layout,
136        tile_size,
137    }))
138}
139
140// ===========================================================================
141// Compute: matmul / load / write / zero-init
142// ===========================================================================
143
144#[cube]
145pub fn cmma_execute<L: Numeric, R: Numeric, A: Numeric>(
146    lhs: &cmma::Matrix<L>,
147    rhs: &cmma::Matrix<R>,
148    acc: &mut cmma::Matrix<A>,
149) {
150    cmma::execute(lhs, rhs, &*acc, &*acc);
151}
152
153#[cube]
154pub fn cmma_load_from_shared<E: Numeric, ES: Size, N: Numeric>(
155    shared: &SharedTile<E>,
156    matrix: &mut cmma::Matrix<N>,
157    #[comptime] ident: StageIdent,
158    #[comptime] matrix_layout: MatrixLayout,
159) {
160    let shared = shared.view::<ES>();
161    match ident {
162        StageIdent::Lhs | StageIdent::Rhs => {
163            cmma_load_strided(&shared, matrix, ComptimeOption::new_None());
164        }
165        StageIdent::Acc => {
166            cmma_load_strided(
167                &shared,
168                matrix,
169                ComptimeOption::new_Some(as_cmma_layout(matrix_layout)),
170            );
171        }
172        _ => panic!("Invalid ident for CMMA load"),
173    }
174}
175
176#[cube]
177pub fn cmma_load_zeros<N: Numeric>(matrix: &mut cmma::Matrix<N>) {
178    cmma::fill(matrix, N::from_int(0));
179}
180
181#[cube]
182pub fn cmma_write_to_shared<E: Numeric, ES: Size, A: Numeric>(
183    shared: &mut SharedTile<E>,
184    matrix: &cmma::Matrix<A>,
185) {
186    let mut shared = shared.view::<ES>();
187    let casted: cmma::Matrix<E> = cmma::cast(matrix);
188    CmmaStageWriter::store_fragment(&mut shared, &casted);
189}
cubek_std/tile/variants/instruction/cmma/base.rs

cubek_std/tile/variants/instruction/cmma/
base.rs
