cubek_std/tile/variants/instruction/mma/
writer.rs

1use cubecl::{
2    prelude::*,
3    {cmma::MmaDefinition, ir::MatrixIdent},
4};
5
6use crate::{
7    MatrixLayout, as_cmma_layout,
8    tile::{
9        StridedTile,
10        variants::{MmaIOConfig, StoreMethod},
11    },
12};
13
14/// Writer for storing the output registers.
15#[derive(CubeType)]
16pub struct MmaStageWriter {}
17
18#[cube]
19impl MmaStageWriter {
20    pub fn store_fragment<
21        E: Numeric,
22        N: Size,
23        V: Numeric,
24        NV: Size,
25        A: Numeric,
26        B: Numeric,
27        CD: Numeric,
28    >(
29        tile: &mut StridedTile<V, NV>,
30        fragment: &Array<Vector<E, N>>,
31        def: &MmaDefinition<A, B, CD>,
32        #[comptime] ident: MatrixIdent,
33        #[comptime] layout: MatrixLayout,
34        #[comptime] m: u32,
35        #[comptime] config: MmaIOConfig,
36    ) {
37        let vector_layout = def.vector_layout(ident);
38        let transposed = comptime![as_cmma_layout(layout) != vector_layout];
39
40        match config.store_method() {
41            StoreMethod::Manual => {
42                if transposed {
43                    store_manual_transposed(tile, fragment, def, ident, layout);
44                } else {
45                    store_manual_plain(tile, fragment, def, ident, layout);
46                }
47            }
48            StoreMethod::StoreMatrix => {
49                store_stmatrix::<E, N, V, NV, A, B, CD>(tile, fragment, def, transposed, ident, m)
50            }
51        }
52    }
53}
54
55#[cube]
56fn store_manual_transposed<
57    E: Numeric,
58    N: Size,
59    V: Numeric,
60    NV: Size,
61    A: Numeric,
62    B: Numeric,
63    CD: Numeric,
64>(
65    tile: &mut StridedTile<V, NV>,
66    fragment: &Array<Vector<E, N>>,
67    def: &MmaDefinition<A, B, CD>,
68    #[comptime] ident: MatrixIdent,
69    #[comptime] layout: MatrixLayout,
70) {
71    let num_vectors = def.vectors_per_lane(ident);
72    let vector_size = def.vector_size(ident);
73    let lane_id = UNIT_POS_PLANE;
74
75    let stride = tile.unvectorized_stride();
76    let mut tile = tile.with_vector_size::<Const<1>>();
77
78    let (stride_row, stride_col) = match layout {
79        MatrixLayout::RowMajor => (stride, 1),
80        MatrixLayout::ColMajor => (1, stride),
81    };
82
83    #[unroll]
84    for i in 0..num_vectors {
85        #[unroll]
86        for n in 0..vector_size {
87            let elem_idx = i * vector_size + n;
88            let (row, col) = def.position_of_nth(lane_id, elem_idx as u32, ident);
89            let offset = row * stride_row + col * stride_col;
90            let offset = tile.stage_offset(offset);
91
92            tile.container[offset as usize] = Vector::cast_from(fragment[i].extract(n));
93        }
94    }
95}
96
97#[cube]
98fn store_manual_plain<
99    E: Numeric,
100    N: Size,
101    V: Numeric,
102    NV: Size,
103    A: Numeric,
104    B: Numeric,
105    CD: Numeric,
106>(
107    tile: &mut StridedTile<V, NV>,
108    fragment: &Array<Vector<E, N>>,
109    def: &MmaDefinition<A, B, CD>,
110    #[comptime] ident: MatrixIdent,
111    #[comptime] layout: MatrixLayout,
112) {
113    let num_vectors = def.vectors_per_lane(ident);
114    let vector_size = def.vector_size(ident);
115    let lane_id = UNIT_POS_PLANE;
116    let stride = tile.unvectorized_stride();
117    // Supported on all targets that support manual MMA
118    let mut tile = tile.with_vector_size::<N>();
119
120    let (stride_row, stride_col) = match layout {
121        MatrixLayout::RowMajor => (stride, 1),
122        MatrixLayout::ColMajor => (1, stride),
123    };
124
125    #[unroll]
126    for i in 0..num_vectors {
127        let value = fragment[i];
128        let elem_idx = i * vector_size;
129        let (row, col) = def.position_of_nth(lane_id, elem_idx as u32, ident);
130        let offset = row * stride_row + col * stride_col;
131        let offset = tile.stage_offset(offset / vector_size as u32);
132
133        tile.container[offset as usize] = Vector::cast_from(value);
134    }
135}
136
137/// This is important to use on CUDA because CUDA's matrices are heavily permuted, being organized
138/// into 8x8 chunks with only 32 contiguous bits per thread. `stmatrix` uses warp shuffles to move
139/// the elements from the mma fragment positions for each thread to 8 consecutive elements in each
140/// thread (if executed with x4), then stores them in one transaction. This currently only supports
141/// f16, fp8 needs more handling and packed fp4 isn't supported at all. So these currently fall back
142/// to manual loading. tf32 isn't supported by the instruction at all.
143#[cube]
144fn store_stmatrix<
145    E: Numeric,
146    N: Size,
147    V: Numeric,
148    NV: Size,
149    A: Numeric,
150    B: Numeric,
151    CD: Numeric,
152>(
153    tile: &mut StridedTile<V, NV>,
154    fragment: &Array<Vector<E, N>>,
155    def: &MmaDefinition<A, B, CD>,
156    #[comptime] transposed: bool,
157    #[comptime] ident: MatrixIdent,
158    #[comptime] m: u32,
159) {
160    let stage_vector_size = tile.container.vector_size().comptime();
161    let stride = tile.unvectorized_stride();
162
163    let elem_size = E::type_size().comptime();
164    let num_regs = def.vectors_per_lane(ident);
165    let width = (16 / elem_size / stage_vector_size) as u32;
166
167    let start = stmatrix_offset::<V, A, B, CD>(stride, def, stage_vector_size, ident, m);
168    let start = tile.stage_offset(start);
169
170    let row_slice = &mut tile.container[start as usize..(start + width) as usize];
171
172    let stage_ty = V::as_type().comptime();
173    let frag_ty = E::as_type().comptime();
174    if stage_ty == frag_ty {
175        let row_slice = row_slice.downcast_mut();
176        def.store_matrix::<Vector<E, NV>, N>(row_slice, fragment, ident, num_regs, transposed);
177    } else {
178        let mut frag = Array::new(num_regs);
179        #[unroll]
180        for i in 0..num_regs {
181            frag[i] = Vector::cast_from(fragment[i]);
182        }
183        def.store_matrix::<_, N>(row_slice, &frag, ident, num_regs, transposed);
184    }
185}
186
187/// Very hardcoded, still haven't figured out the proper generic formula. So keep this separate from
188/// the read index for now, and ensure out is row-major.
189#[cube]
190pub(crate) fn stmatrix_offset<E: Numeric, A: Numeric, B: Numeric, CD: Numeric>(
191    stride: u32,
192    def: &MmaDefinition<A, B, CD>,
193    #[comptime] stage_vector_size: VectorSize,
194    #[comptime] ident: MatrixIdent,
195    #[comptime] m: u32,
196) -> u32 {
197    let (stride_row, stride_col) = (stride, 1);
198
199    let elem_size = E::type_size().comptime();
200    let num_regs = def.vectors_per_lane(ident);
201    let width = (16 / elem_size) as u32;
202    // Height is always 8, and lanes are divided into blocks of 8.
203    let height = 8;
204
205    //  Indices are wrapped for < 4 registers.
206    let lane = UNIT_POS_PLANE;
207    let sub_lane = lane % height;
208    let nth_matrix = lane / height % num_regs as u32;
209
210    let tiles_row = m / height;
211
212    // Tiles are arranged in column-major fashion
213    let row_offs = (nth_matrix % tiles_row) * 8;
214    let col_offs = (nth_matrix / tiles_row) * width;
215
216    let (row, col) = (row_offs + sub_lane, col_offs);
217
218    let start = row * stride_row + col * stride_col;
219    start / stage_vector_size as u32
220}
cubek_std/tile/variants/instruction/mma/writer.rs

cubek_std/tile/variants/instruction/mma/
writer.rs
