cubecl_ir/
cmma.rs

1use alloc::{format, string::String};
2use derive_more::Display;
3use derive_new::new;
4
5use super::Value;
6use crate::{Closure, OperationReflect};
7use crate::{StorageType, TypeHash};
8use core::fmt::Display;
9
10#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
11#[derive(Debug, Clone, Copy, TypeHash, PartialEq, Eq, Hash, PartialOrd, Ord, Display)]
12#[allow(missing_docs)]
13pub enum MatrixIdent {
14    #[display("IdentA")]
15    A,
16    #[display("IdentB")]
17    B,
18    #[display("IdentAcc")]
19    Accumulator,
20}
21
22#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
23#[derive(Debug, Clone, Copy, TypeHash, PartialEq, Eq, Hash, PartialOrd, Ord, Display)]
24#[display(rename_all = "snake_case")]
25#[allow(missing_docs)]
26pub enum MatrixLayout {
27    ColMajor,
28    RowMajor,
29    Undefined,
30}
31
32#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
33#[derive(Debug, Clone, Copy, TypeHash, PartialEq, Eq, Hash, PartialOrd, Ord, Display)]
34#[display(rename_all = "snake_case")]
35#[allow(missing_docs)]
36pub enum MatrixScope {
37    Plane,
38    Cube,
39}
40
41#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
42#[derive(new, Debug, Clone, Copy, TypeHash, PartialEq, Eq, Hash, PartialOrd, Ord)]
43#[allow(missing_docs)]
44pub struct MatrixType {
45    pub ident: MatrixIdent,
46    pub m: usize,
47    pub n: usize,
48    pub k: usize,
49    pub storage: StorageType,
50    pub layout: MatrixLayout,
51    pub scope: MatrixScope,
52}
53
54impl MatrixType {
55    /// Number of elements in terms of the physical storage type, accounting for packed elements
56    pub fn num_elems(&self) -> usize {
57        let elems = match self.ident {
58            MatrixIdent::A => self.m * self.k,
59            MatrixIdent::B => self.k * self.n,
60            MatrixIdent::Accumulator => self.m * self.n,
61        };
62        elems / self.storage.packing_factor()
63    }
64}
65
66#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
67#[derive(Debug, Clone, Copy, TypeHash, PartialEq, Eq, Hash, PartialOrd, Ord)]
68pub enum ClampMode {
69    Undefined,
70    Constant(u32),
71    ClampToEdge,
72    Repeat,
73    RepeatMirrored,
74}
75
76/// Cooperative Matrix-Multiply and Accumulate Instruction.
77#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
78#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationReflect)]
79#[operation(opcode_name = CmmaOpCode)]
80#[allow(missing_docs)]
81pub enum CoopMma {
82    /// Fill the matrix with the value.
83    Fill { value: Value },
84    /// Load the value into the matrix given the stride.
85    Load {
86        #[args(allow_ptr, ptr_read)]
87        ptr: Value,
88        stride: Value,
89        #[args(skip)]
90        layout: Option<MatrixLayout>,
91    },
92    /// Load the value into the matrix given the tensor layout.
93    LoadTensor {
94        #[args(allow_ptr, ptr_read)]
95        buffer: Value,
96        layout: Value,
97        #[args(skip)]
98        view: Option<Value>,
99    },
100    /// Executes D=A*B+C;
101    ///
102    /// For implementing a matmul, `D=C` : `C+=A*B`
103    Execute {
104        mat_a: Value,
105        mat_b: Value,
106        mat_c: Value,
107    },
108    /// Store the matrix in an output variable following the stride and the layout.
109    Store {
110        mat: Value,
111        stride: Value,
112        #[args(allow_ptr, ptr_write)]
113        destination: Value,
114        #[args(skip)]
115        layout: MatrixLayout,
116    },
117    /// Store the matrix in an output variable following the tensor layout.
118    StoreTensor {
119        mat: Value,
120        layout: Value,
121        #[args(skip)]
122        view: Option<Value>,
123    },
124    /// Cast a fragment to another type.
125    Cast { input: Value },
126
127    /// Row index of nth element in the lane
128    RowIndex {
129        lane_id: Value,
130        i: Value,
131        #[args(skip)]
132        matrix: MatrixType,
133    },
134    /// Column index of nth element in the lane
135    ColIndex {
136        lane_id: Value,
137        i: Value,
138        #[args(skip)]
139        matrix: MatrixType,
140    },
141    /// Execute a CUDA `ldmatrix` instruction
142    LoadMatrix {
143        #[args(allow_ptr, ptr_read)]
144        ptr: Value,
145        factor: usize,
146        transpose: bool,
147    },
148    /// Execute a CUDA `stmatrix` instruction
149    StoreMatrix {
150        registers: Value,
151        factor: usize,
152        transpose: bool,
153        #[args(allow_ptr, ptr_write)]
154        destination: Value,
155    },
156    /// Manual execute.
157    ExecuteManual {
158        #[args(skip)]
159        matrix: MatrixType,
160        registers_a: Value,
161        registers_b: Value,
162        registers_c: Value,
163    },
164    /// Scaled manual execute.
165    ExecuteScaled {
166        #[args(skip)]
167        matrix: MatrixType,
168        registers_a: Value,
169        registers_b: Value,
170        registers_c: Value,
171        scales_a: Value,
172        scales_b: Value,
173        scales_factor: usize,
174    },
175    ExecuteElementwise {
176        matrix: Value,
177        #[args(skip)]
178        op: Closure,
179    },
180}
181
182impl Display for CoopMma {
183    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
184        match self {
185            CoopMma::Fill { value } => write!(f, "fill({value})"),
186            CoopMma::Load {
187                ptr,
188                stride,
189                layout,
190            } => {
191                let layout = layout
192                    .map(|it| format!(", layout: {it:?}"))
193                    .unwrap_or(String::new());
194                write!(f, "matrix_load({ptr}, stride: {stride}{layout})")
195            }
196            CoopMma::LoadTensor {
197                buffer,
198                layout,
199                view,
200            } => {
201                let view = view.map(|it| format!(", view: {it}")).unwrap_or_default();
202                write!(f, "matrix_load_tensor({buffer}, layout: {layout}{view})")
203            }
204            CoopMma::Execute {
205                mat_a,
206                mat_b,
207                mat_c,
208            } => write!(f, "execute_cmma({mat_a}, {mat_b}, {mat_c})"),
209            CoopMma::ExecuteManual {
210                matrix,
211                registers_a,
212                registers_b,
213                registers_c,
214            } => {
215                write!(
216                    f,
217                    "execute_manual_mma(
218                    matrix: {matrix:?},
219                    frag_a: {registers_a},
220                    frag_b: {registers_b},
221                    frag_c: {registers_c},
222                )"
223                )
224            }
225            CoopMma::ExecuteScaled {
226                matrix,
227                registers_a,
228                registers_b,
229                registers_c,
230                scales_a,
231                scales_b,
232                scales_factor,
233            } => {
234                write!(
235                    f,
236                    "execute_scaled_mma_{scales_factor}x(
237                    matrix: {matrix:?},
238                    frag_a: {registers_a},
239                    frag_b: {registers_b},
240                    frag_c: {registers_c},
241                    scales_a: {scales_a},
242                    scales_b: {scales_b}
243                )"
244                )
245            }
246            CoopMma::ExecuteElementwise { matrix, op } => {
247                write!(f, "execute_elementwise({matrix}, {op})")
248            }
249            CoopMma::Store {
250                mat,
251                stride,
252                destination,
253                layout,
254            } => write!(
255                f,
256                "matrix_store({mat}, stride: {stride}, layout: {layout:?}, dest: {destination})"
257            ),
258            CoopMma::StoreTensor { mat, layout, .. } => {
259                write!(f, "matrix_store_tensor({mat}, layout: {layout})")
260            }
261            CoopMma::Cast { input } => {
262                write!(f, "matrix_cast(input: {input})")
263            }
264            CoopMma::RowIndex { lane_id, i, matrix } => {
265                write!(f, "row_idx(lane_id: {lane_id}, i: {i}, matrix: {matrix:?})",)
266            }
267            CoopMma::ColIndex { lane_id, i, matrix } => {
268                write!(f, "col_idx(lane_id: {lane_id}, i: {i}, matrix: {matrix:?})",)
269            }
270            CoopMma::LoadMatrix {
271                ptr,
272                factor,
273                transpose,
274            } => {
275                write!(f, "ldmatrix_{factor}x({ptr}, transpose: {transpose})")
276            }
277            CoopMma::StoreMatrix {
278                registers,
279                factor,
280                transpose,
281                destination,
282            } => {
283                write!(
284                    f,
285                    "stmatrix_{factor}x({registers}, dest: {destination}, transpose: {transpose})"
286                )
287            }
288        }
289    }
290}
cubecl_ir/cmma.rs

cubecl_ir/
cmma.rs
