cubecl_cpp/shared/
mma.rs

1use crate::shared::Item;
2
3use super::{Component, Dialect, Elem, Value};
4use cubecl_core::ir::{
5    DeviceProperties,
6    features::{MmaConfig, ScaledMmaConfig},
7};
8use std::{
9    fmt::{Debug, Display, Formatter},
10    marker::PhantomData,
11};
12
13pub type SupportedMmaCombinations = Vec<MmaConfig>;
14pub type SupportedScaledMmaCombinations = Vec<ScaledMmaConfig>;
15
16pub trait Architecture {
17    fn warp_size(&self) -> u32;
18    fn is_wmma_capable(&self) -> bool;
19    fn is_mfma_capable(&self) -> bool;
20    fn get_version(&self) -> u32 {
21        0
22    }
23}
24
25pub fn register_wmma_features(
26    supported_combinations: SupportedMmaCombinations,
27    properties: &mut DeviceProperties,
28) {
29    for config in supported_combinations {
30        properties.features.matmul.cmma.insert(config);
31    }
32}
33
34pub fn register_mma_features(
35    supported_combinations: SupportedMmaCombinations,
36    properties: &mut DeviceProperties,
37) {
38    for config in supported_combinations {
39        properties.features.matmul.mma.insert(config);
40    }
41}
42
43pub fn register_scaled_mma_features(
44    supported_combinations: SupportedScaledMmaCombinations,
45    properties: &mut DeviceProperties,
46) {
47    for config in supported_combinations {
48        properties.features.matmul.scaled_mma.insert(config);
49    }
50}
51
52#[derive(Debug, Clone, PartialEq, Eq, Copy, Hash)]
53pub enum FragmentIdent<D: Dialect> {
54    A,
55    B,
56    Accumulator,
57    _Dialect(PhantomData<D>),
58}
59
60#[derive(Debug, Clone, PartialEq, Eq, Copy, Hash)]
61pub enum FragmentLayout<D: Dialect> {
62    ColMajor,
63    RowMajor,
64    _Dialect(PhantomData<D>),
65}
66
67#[derive(Debug, Clone, PartialEq, Eq, Copy, Hash)]
68pub struct FragmentType<D: Dialect> {
69    pub ident: FragmentIdent<D>,
70    pub m: u32,
71    pub n: u32,
72    pub k: u32,
73    pub elem: Elem<D>,
74    pub layout: Option<FragmentLayout<D>>,
75}
76
77#[derive(new, Debug, Clone, PartialEq, Eq, Copy)]
78pub struct MmaShape<D: Dialect> {
79    pub m: u32,
80    pub n: u32,
81    pub k: u32,
82    _d: PhantomData<D>,
83}
84
85impl<D: Dialect> MmaShape<D> {
86    pub fn num_elems(&self, ident: FragmentIdent<D>) -> u32 {
87        match ident {
88            FragmentIdent::A => self.m * self.k,
89            FragmentIdent::B => self.k * self.n,
90            FragmentIdent::Accumulator => self.m * self.n,
91            _ => unimplemented!(),
92        }
93    }
94}
95
96/// Warp Matrix-Multiply and Accumulate Instruction.
97#[derive(Debug, Clone, PartialEq)]
98pub enum WmmaInstruction<D: Dialect> {
99    /// Fill the fragment with the value.
100    Fill { frag: Value<D>, value: Value<D> },
101    /// Load the value into the fragment given the stride.
102    Load {
103        frag: Value<D>,
104        ptr: Value<D>,
105        stride: Value<D>,
106        layout: Option<FragmentLayout<D>>,
107    },
108    /// Executes D=A*B+C;
109    ///
110    /// For implementing a matmul, `D=C` : `C+=A*B`
111    Execute {
112        frag_a: Value<D>,
113        frag_b: Value<D>,
114        frag_c: Value<D>,
115        frag_d: Value<D>,
116        warp_size: u32,
117    },
118    /// Executes D=A*B+C using manually managed registers;
119    ///
120    /// For implementing a matmul, `D=C` : `C+=A*B`
121    /// Takes a sequence of registers for the inputs, and returns an array of registers for the
122    /// output. PTX requires output registers to be non-overlapping, so we use array to ensure that
123    /// and handle potentially destructuring it internally.
124    ExecuteManual {
125        shape: MmaShape<D>,
126        frag_a: Value<D>,
127        frag_b: Value<D>,
128        frag_c: Value<D>,
129        frag_d: Value<D>,
130    },
131    /// Executes D=A*B+C using manually managed registers;
132    ///
133    /// For implementing a matmul, `D=C` : `C+=A*B`
134    /// Takes a sequence of registers for the inputs, and returns an array of registers for the
135    /// output. PTX requires output registers to be non-overlapping, so we use array to ensure that
136    /// and handle potentially destructuring it internally.
137    ExecuteScaled {
138        shape: MmaShape<D>,
139        frag_a: Value<D>,
140        frag_b: Value<D>,
141        frag_c: Value<D>,
142        frag_d: Value<D>,
143
144        scales_a: Value<D>,
145        scales_b: Value<D>,
146        scales_factor: u32,
147    },
148    /// Store the fragment in an output variable following the stride and the layout.
149    Store {
150        frag: Value<D>,
151        stride: Value<D>,
152        destination: Value<D>,
153        layout: FragmentLayout<D>,
154    },
155    /// Load a part of a fragment into registers, either 1, 2, or 4 at once.
156    LdMatrix {
157        output: Value<D>,
158        ptr: Value<D>,
159        factor: u32,
160        transpose: bool,
161    },
162    /// Store a part of a fragment into smem, either 1, 2, or 4 at once.
163    StMatrix {
164        registers: Value<D>,
165        ptr: Value<D>,
166        factor: u32,
167        transpose: bool,
168    },
169    /// Cast
170    Cast { input: Value<D>, output: Value<D> },
171}
172
173impl<D: Dialect> Display for FragmentLayout<D> {
174    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
175        D::compile_wmma_fragment_layout(f, self)
176    }
177}
178
179impl<D: Dialect> Display for FragmentIdent<D> {
180    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
181        D::compile_wwma_fragment_ident(f, self)
182    }
183}
184
185impl<D: Dialect> Display for FragmentType<D> {
186    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
187        D::compile_wmma_fragment(f, self)
188    }
189}
190
191impl<D: Dialect> Display for WmmaInstruction<D> {
192    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
193        D::compile_wmma_instruction(f, self)
194    }
195}
196
197pub mod wmma_api_base {
198    use crate::{
199        cuda::ptx::{ldmatrix_call, stmatrix_call},
200        shared::ManualMma,
201    };
202
203    use super::*;
204
205    pub fn compile_fragment_declaration<D: Dialect>(
206        f: &mut std::fmt::Formatter<'_>,
207        val: &Value<D>,
208        ty: &Item<D>,
209    ) -> std::fmt::Result {
210        match ty {
211            Item::Fragment(frag) => writeln!(f, "{frag} {val}_store;"),
212            _ => panic!("value must be a fragment"),
213        }
214    }
215
216    pub fn compile_fragment_ident<D: Dialect>(
217        f: &mut std::fmt::Formatter<'_>,
218        namespace: &str,
219        ident: &FragmentIdent<D>,
220    ) -> std::fmt::Result {
221        match ident {
222            FragmentIdent::A => write!(f, "{namespace}::matrix_a"),
223            FragmentIdent::B => write!(f, "{namespace}::matrix_b"),
224            FragmentIdent::Accumulator => write!(f, "{namespace}::accumulator"),
225            FragmentIdent::_Dialect(_) => Ok(()),
226        }
227    }
228
229    pub fn compile_fragment_layout<D: Dialect>(
230        f: &mut std::fmt::Formatter<'_>,
231        namespace: &str,
232        layout: &FragmentLayout<D>,
233    ) -> std::fmt::Result {
234        match layout {
235            FragmentLayout::ColMajor => f.write_str(format!("{namespace}::col_major").as_str()),
236            FragmentLayout::RowMajor => f.write_str(format!("{namespace}::row_major").as_str()),
237            FragmentLayout::_Dialect(_) => Ok(()),
238        }
239    }
240
241    pub fn compile_fragment<D: Dialect>(
242        f: &mut std::fmt::Formatter<'_>,
243        namespace: &str,
244        fragment: &FragmentType<D>,
245    ) -> std::fmt::Result {
246        let elem = match fragment.elem {
247            Elem::TF32 => format!("{namespace}::precision::tf32"),
248            Elem::BF16 => {
249                if fragment.ident == FragmentIdent::Accumulator {
250                    format!("{}", Elem::<D>::F16) // Normally not supported except for cast.
251                } else {
252                    format!("{}", fragment.elem)
253                }
254            }
255            elem => format!("{elem}"),
256        };
257        match fragment.layout {
258            Some(layout) => write!(
259                f,
260                "{namespace}::fragment<{}, {}, {}, {}, {}, {}>",
261                fragment.ident, fragment.m, fragment.n, fragment.k, elem, layout
262            ),
263            None => write!(
264                f,
265                "{namespace}::fragment<{}, {}, {}, {}, {}>",
266                fragment.ident, fragment.m, fragment.n, fragment.k, elem,
267            ),
268        }
269    }
270
271    pub fn compile_instruction<D: Dialect>(
272        f: &mut std::fmt::Formatter<'_>,
273        namespace: &str,
274        instruction: &WmmaInstruction<D>,
275    ) -> std::fmt::Result {
276        match instruction {
277            WmmaInstruction::Fill { frag, value } => {
278                let frag = frag.fmt_ref();
279                writeln!(f, "{namespace}::fill_fragment({frag}, {value});")
280            }
281            WmmaInstruction::Load {
282                frag,
283                ptr,
284                stride,
285                layout: None,
286            } => {
287                let item = *ptr.item().value_ty();
288                let frag = frag.fmt_ref();
289                if item.vectorization() > 1 {
290                    let elem = item.elem();
291                    let qualifier = ptr.const_qualifier();
292                    writeln!(
293                        f,
294                        "{namespace}::load_matrix_sync({frag}, reinterpret_cast<{elem}{qualifier}*>({ptr}), {stride});"
295                    )
296                } else {
297                    writeln!(f, "{namespace}::load_matrix_sync({frag}, {ptr}, {stride});")
298                }
299            }
300            WmmaInstruction::Load {
301                frag,
302                ptr,
303                stride,
304                layout: Some(layout),
305            } => {
306                let frag = frag.fmt_ref();
307                let layout = match layout {
308                    FragmentLayout::ColMajor => format!("{namespace}::mem_col_major"),
309                    FragmentLayout::RowMajor => format!("{namespace}::mem_row_major"),
310                    FragmentLayout::_Dialect(_) => "".to_string(),
311                };
312                let item = *ptr.item().value_ty();
313                if item.vectorization() > 1 {
314                    let elem = item.elem();
315                    writeln!(
316                        f,
317                        "{namespace}::load_matrix_sync({frag}, reinterpret_cast<{elem} *>({ptr}), {stride}, {layout});"
318                    )
319                } else {
320                    writeln!(
321                        f,
322                        "{namespace}::load_matrix_sync({frag}, {ptr}, {stride}, {layout});"
323                    )
324                }
325            }
326            WmmaInstruction::LdMatrix {
327                output,
328                ptr,
329                factor,
330                transpose,
331            } => f.write_str(&ldmatrix_call(output, ptr, factor, transpose)),
332            WmmaInstruction::StMatrix {
333                registers,
334                ptr,
335                factor,
336                transpose,
337            } => f.write_str(&stmatrix_call(registers, ptr, factor, transpose)),
338            WmmaInstruction::Execute {
339                frag_a,
340                frag_b,
341                frag_c,
342                frag_d,
343                ..
344            } => {
345                let frag_a = frag_a.fmt_ref();
346                let frag_b = frag_b.fmt_ref();
347                let frag_c = frag_c.fmt_ref();
348                let frag_d = frag_d.fmt_ref();
349                writeln!(
350                    f,
351                    "{namespace}::mma_sync({frag_d}, {frag_a}, {frag_b}, {frag_c});"
352                )
353            }
354            WmmaInstruction::Store {
355                frag,
356                stride,
357                destination,
358                layout,
359            } => {
360                let frag = frag.fmt_ref();
361                let layout = match layout {
362                    FragmentLayout::ColMajor => format!("{namespace}::mem_col_major"),
363                    FragmentLayout::RowMajor => format!("{namespace}::mem_row_major"),
364                    FragmentLayout::_Dialect(_) => "".to_string(),
365                };
366
367                let item = *destination.item().value_ty();
368                let mut reinterpret_cast = item.vectorization() > 1;
369                let elem = match item.elem() {
370                    Elem::BF16 => {
371                        reinterpret_cast = true;
372                        Elem::F16
373                    }
374                    _ => *item.elem(),
375                };
376                if reinterpret_cast {
377                    writeln!(
378                        f,
379                        "{namespace}::store_matrix_sync(reinterpret_cast<{elem} *>({destination}), {frag}, {stride}, {layout});"
380                    )
381                } else {
382                    writeln!(
383                        f,
384                        "{namespace}::store_matrix_sync({destination}, {frag}, {stride}, {layout});"
385                    )
386                }
387            }
388            WmmaInstruction::Cast { input, output } => {
389                let input = input.ensure_lvalue(f)?;
390                let output = output.ensure_lvalue(f)?;
391                let ty = match *output.item().value_ty() {
392                    Item::Fragment(frag) => frag.elem,
393                    _ => panic!("Should be a fragment"),
394                };
395                match ty {
396                    Elem::BF16 => {
397                        let elem = Elem::<D>::F16;
398                        write!(
399                            f,
400                            "// cast
401for(int t=0; t<{input}.num_elements; t++) {{
402  {ty} elem = {ty}({input}.x[t]);
403  {output}.x[t] = *reinterpret_cast<{elem} *>(&elem);
404}}
405"
406                        )
407                    }
408                    _ => {
409                        write!(
410                            f,
411                            "// cast
412for(int t=0; t<{input}.num_elements; t++) {{ {output}.x[t] = {ty}({input}.x[t]); }}
413"
414                        )
415                    }
416                }
417            }
418            WmmaInstruction::ExecuteManual {
419                shape,
420                frag_a,
421                frag_b,
422                frag_c,
423                frag_d,
424            } => D::compile_manual_mma(f, ManualMma::new(*shape, frag_a, frag_b, frag_c, frag_d)),
425            WmmaInstruction::ExecuteScaled {
426                shape,
427                frag_a,
428                frag_b,
429                frag_c,
430                frag_d,
431                scales_a,
432                scales_b,
433                scales_factor,
434            } => D::compile_scaled_mma(
435                f,
436                ManualMma::new(*shape, frag_a, frag_b, frag_c, frag_d),
437                *scales_a,
438                *scales_b,
439                *scales_factor,
440            ),
441        }
442    }
443}
444
445pub fn frag_as_ptr<D: Dialect>(f: &mut Formatter<'_>, ptr: &Value<D>) -> Value<D> {
446    let item = ptr.item();
447    if item.vectorization() > 1 {
448        let item_value = item.as_scalar();
449        ptr.reinterpret_ptr(f, item_value)
450    } else {
451        *ptr
452    }
453}
454
455pub fn frag_ident_str<D: Dialect>(frag: &FragmentIdent<D>) -> &str {
456    match frag {
457        FragmentIdent::A => "a",
458        FragmentIdent::B => "b",
459        FragmentIdent::Accumulator => "c",
460        FragmentIdent::_Dialect(_) => "d",
461    }
462}
463
464pub fn frag_layout_str<D: Dialect>(frag: &Option<FragmentLayout<D>>) -> &str {
465    match frag {
466        Some(layout) => match layout {
467            FragmentLayout::ColMajor => "col",
468            FragmentLayout::RowMajor => "row",
469            FragmentLayout::_Dialect(_) => "",
470        },
471        None => "",
472    }
473}
474
475pub fn value_to_frag<D: Dialect>(frag: &Value<D>) -> FragmentType<D> {
476    match frag.item().unwrap_ptr() {
477        Item::Fragment(frag) => frag,
478        _ => panic!(),
479    }
480}
cubecl_cpp/shared/mma.rs

cubecl_cpp/shared/
mma.rs
