cubecl_cpp/hip/mma/
wmma_intrinsics_compiler.rs

1use std::fmt::Formatter;
2
3use crate::{
4    Dialect,
5    hip::{HipDialect, arch::AMDArchitecture},
6    shared::{
7        Architecture, Component, DialectWmmaCompiler, Elem, Flags, FmtLeft, FragmentIdent,
8        FragmentLayout, FragmentType, Item, ManualMma, MmaShape, SupportedMmaCombinations, Value,
9        WmmaInstruction, frag_as_ptr, frag_ident_str, frag_layout_str, value_to_frag,
10        wmma_api_base,
11    },
12};
13use cubecl_core::ir::{self as gpu, MatrixIdent, MatrixType, features::MmaConfig};
14
15#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
16pub struct WmmaIntrinsicCompiler {}
17
18#[derive(new, Debug, Clone, PartialEq)]
19pub struct WmmaFill<D: Dialect> {
20    frag: FragmentType<D>,
21}
22
23#[derive(new, Debug, Clone, PartialEq)]
24pub struct WmmaLoad<D: Dialect> {
25    frag: FragmentType<D>,
26    layout: Option<FragmentLayout<D>>,
27}
28
29#[derive(new, Debug, Clone, PartialEq)]
30pub struct WmmaStore<D: Dialect> {
31    frag: FragmentType<D>,
32    layout: FragmentLayout<D>,
33}
34
35#[derive(new, Debug, Clone, PartialEq)]
36pub struct WmmaExecute<D: Dialect> {
37    frag_a: FragmentType<D>,
38    frag_b: FragmentType<D>,
39    frag_c: FragmentType<D>,
40    frag_d: FragmentType<D>,
41}
42
43#[derive(new, Debug, Clone, PartialEq)]
44pub struct WmmaCast<D: Dialect> {
45    frag_input: FragmentType<D>,
46    frag_output: FragmentType<D>,
47}
48
49impl<D: Dialect> WmmaFill<D> {
50    pub fn fn_name(&self) -> String {
51        let layout = frag_layout_str(&self.frag.layout);
52        let ident = frag_ident_str(&self.frag.ident);
53        let (m, n, k) = (self.frag.m, self.frag.n, self.frag.k);
54        let elem = self.frag.elem;
55
56        format!("wmma_fill_{elem}_{ident}_{m}x{n}x{k}_{layout}",)
57    }
58
59    pub fn format_extension(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
60        let elem = self.frag.elem;
61        let frag = self.frag;
62        let name = self.fn_name();
63
64        write!(
65            f,
66            "
67// Fill the fragment.
68__device__ void {name}({frag}& frag, {elem} value) {{
69    #pragma unroll
70    for (uint i = 0; i < 8; ++i) {{
71      frag[i] = value;
72    }}
73}}
74        "
75        )
76    }
77}
78
79impl<D: Dialect> WmmaLoad<D> {
80    pub fn fn_name(&self) -> String {
81        let layout_frag = frag_layout_str(&self.frag.layout);
82        let layout = frag_layout_str(&self.layout);
83        let ident = frag_ident_str(&self.frag.ident);
84        let elem = self.frag.elem;
85        let (m, n, k) = (self.frag.m, self.frag.n, self.frag.k);
86
87        format!("wmma_load_{elem}_{ident}_{m}x{n}x{k}_{layout_frag}_{layout}",)
88    }
89
90    /// Matrix A must be in column major layout (so fragments correspond to a row)
91    /// Matrices B, C and D must be in row major layout (so fragments correspond to a column)
92    ///
93    /// Each lane is a thread so each column get 8 VGPRs used to store fragments
94    /// Here is the layout for C and D matrices and how they map to registers
95    ///
96    /// Lane index   0      1      2      3      ...     13     14     15     ...     17     18     ...     30     31
97    /// --------------------------------------------------------------------------------------------------------------
98    /// VGPR0      | 1,1  | 1,2  | 1,3  | 1,4  | ...  | 1,13 | 1,14 | 1,15 | ...  | 2,1  | 2,2  | ...  | 2,15 | 2,16 |
99    /// --------------------------------------------------------------------------------------------------------------
100    /// VGPR1      | 3,1  | 3,2  | 3,3  | 3,4  | ...  | 3,13 | 3,14 | 3,15 | ...  | 4,1  | 4,2  | ...  | 4,15 | 4,16 |
101    /// --------------------------------------------------------------------------------------------------------------
102    /// VGPR2      | 5,1  | 5,2  | 5,3  | 5,4  | ...  | 5,13 | 5,14 | 5,15 | ...  | 6,1  | 6,2  | ...  | 6,15 | 6,16 |
103    /// --------------------------------------------------------------------------------------------------------------
104    /// VGPR3      | 7,1  | 7,2  | 7,3  | 7,4  | ...  | 7,13 | 7,14 | 7,15 | ...  | 8,1  | 8,2  | ...  | 8,15 | 8,16 |
105    /// --------------------------------------------------------------------------------------------------------------
106    /// VGPR4      | 9,1  | 9,2  | 9,3  | 9,4  | ...  | 9,13 | 9,14 | 9,15 | ...  | 10,1 | 10,2 | ...  | 10,15| 10,16|
107    /// --------------------------------------------------------------------------------------------------------------
108    /// VGPR5      | 11,1 | 11,2 | 11,3 | 11,4 | ...  | 11,13| 11,14| 11,15| ...  | 12,1 | 12,2 | ...  | 12,15| 12,16|
109    /// --------------------------------------------------------------------------------------------------------------
110    /// VGPR6      | 13,1 | 13,2 | 13,3 | 13,4 | ...  | 13,13| 13,14| 13,15| ...  | 14,1 | 14,2 | ...  | 14,15| 14,16|
111    /// --------------------------------------------------------------------------------------------------------------
112    /// VGPR7      | 15,1 | 15,2 | 15,3 | 15,4 | ...  | 15,13| 15,14| 15,15| ...  | 16,1 | 16,2 | ...  | 16,15| 16,16|
113    /// --------------------------------------------------------------------------------------------------------------
114    pub fn format_extension(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
115        let elem = self.frag.elem;
116        let frag = self.frag;
117        let name = self.fn_name();
118
119        let (index_body, length, step) = match frag.ident {
120            FragmentIdent::A | FragmentIdent::B => {
121                let length = 16;
122                let step = 1;
123                // fragment a and b are always in half precision and they don't require special attention
124                // to how they are stored in memory as matrix A and B are also in half precision
125                let index = if (frag.ident == FragmentIdent::A
126                    && frag.layout.unwrap() == FragmentLayout::ColMajor)
127                    || (frag.ident == FragmentIdent::B
128                        && frag.layout.unwrap() == FragmentLayout::RowMajor)
129                {
130                    "i * stride + wmmaLane".to_string()
131                } else {
132                    "i + wmmaLane * stride".to_string()
133                };
134                (index, length, step)
135            }
136            FragmentIdent::Accumulator => {
137                let length = 8;
138                let step = get_output_accumulator_index_step(&elem, &frag);
139                let index = match self.layout {
140                    Some(FragmentLayout::ColMajor) => {
141                        "(i * uint(2) + threadIdx.x / uint(16)) + wmmaLane * stride".to_string()
142                    }
143                    Some(FragmentLayout::RowMajor) => {
144                        "(i * uint(2) + threadIdx.x / uint(16)) * stride + wmmaLane".to_string()
145                    }
146                    _ => panic!(
147                        "cannot load data to an accumulator without knowing the layout of the data"
148                    ),
149                };
150                (index, length, step)
151            }
152            other => panic!("unknown matrix identifier {other}"),
153        };
154
155        write!(
156            f,
157            "
158// Load the fragment.
159__device__ void {name}({frag}& frag, const {elem}* value_ptr, const uint stride) {{
160    {WMMA_LANE_DEF}
161
162    #pragma unroll
163    for (uint i = 0; i < {length}; ++i) {{
164      const uint index = {index_body};
165      frag[i * {step}] = value_ptr[index];
166    }}
167}}
168        "
169        )
170    }
171}
172
173impl<D: Dialect> WmmaStore<D> {
174    pub fn fn_name(&self) -> String {
175        let layout_frag = frag_layout_str(&self.frag.layout);
176        let layout_option = Some(self.layout);
177        let layout = frag_layout_str(&layout_option);
178        let ident = frag_ident_str(&self.frag.ident);
179        let (m, n, k) = (self.frag.m, self.frag.n, self.frag.k);
180        let elem = self.frag.elem;
181
182        format!("wmma_store_{elem}_{ident}_{m}x{n}x{k}_{layout_frag}_{layout}",)
183    }
184
185    pub fn format_extension(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
186        let elem = self.frag.elem;
187        let frag = self.frag;
188        let name = self.fn_name();
189        // frag holds a result column where threads 0-15 of the wavefront have the even rows and threads 16-31 the odd rows
190        // moreover, since we use OPSEL to false in the Execute instruction in f16 output format, the output elements are
191        // stored in even indexes (0, 2, 4, ...) (low 16-bits of the VGPR) in frag
192        let frag_idx = match elem {
193            Elem::F16 | Elem::BF16 => "elemIdx * 2",
194            Elem::F32 => "elemIdx",
195            other => {
196                panic!("C fragment format cannot be {other}. Only f16, bf16 and f32 are supported.")
197            }
198        };
199        // FragmentLayout here represents the desired layout of the matrix C
200        let output_idx = match self.layout {
201            FragmentLayout::ColMajor => "wmmaLane * stride + rowIdx".to_string(),
202            FragmentLayout::RowMajor => "wmmaLane + rowIdx * stride".to_string(),
203            FragmentLayout::_Dialect(_) => String::new(),
204        };
205
206        write!(
207            f,
208            "
209// Store the fragment.
210__device__ void {name}(const {frag}& frag, {elem}* output_ptr, uint stride) {{
211    {WMMA_LANE_DEF}
212
213    #pragma unroll
214    for (uint elemIdx = 0; elemIdx < uint(8); ++elemIdx) {{
215      const uint rowIdx = elemIdx * uint(2) + threadIdx.x / uint(16);
216      output_ptr[{output_idx}] = frag[{frag_idx}];
217    }}
218}}
219        "
220        )
221    }
222}
223
224impl<D: Dialect> WmmaExecute<D> {
225    pub fn from_manual(shape: MmaShape<D>, ab_elem: Elem<D>, cd_elem: Elem<D>) -> Self {
226        let frag_a = FragmentType {
227            ident: FragmentIdent::A,
228            m: shape.m,
229            n: shape.n,
230            k: shape.k,
231            elem: ab_elem,
232            layout: Some(FragmentLayout::ColMajor),
233        };
234        let frag_b = FragmentType {
235            ident: FragmentIdent::B,
236            layout: Some(FragmentLayout::RowMajor),
237            ..frag_a
238        };
239        let frag_cd = FragmentType {
240            ident: FragmentIdent::Accumulator,
241            elem: cd_elem,
242            ..frag_b
243        };
244        WmmaExecute::new(frag_a, frag_b, frag_cd, frag_cd)
245    }
246
247    pub fn fn_name(&self) -> String {
248        format!(
249            "wmma_execute_16x16x16_{}_{}",
250            self.frag_a.elem, self.frag_c.elem
251        )
252    }
253
254    pub fn format_extension(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
255        let name = self.fn_name();
256        let ab_format = match self.frag_a.elem {
257            Elem::F32 => "f32",
258            Elem::BF16 => "bf16",
259            Elem::F16 => "f16",
260            _ => panic!(),
261        };
262        let (cd_format, opsel) = match self.frag_c.elem {
263            Elem::F32 => ("f32", ""),
264            Elem::BF16 => ("bf16", ", false"),
265            Elem::F16 => ("f16", ", false"),
266            _ => panic!(),
267        };
268        let warp_size = 32;
269        write!(
270            f,
271            "
272// Execute wmma.
273__device__ void {name}(const {}& frag_a, const {}& frag_b, const {}& frag_c, {}& frag_d) {{
274    frag_d = __builtin_amdgcn_wmma_{cd_format}_16x16x16_{ab_format}_w{warp_size}(frag_a, frag_b, frag_c{opsel});
275}}
276        ", self.frag_a, self.frag_b, self.frag_c, self.frag_d
277        )
278    }
279}
280
281impl<D: Dialect> WmmaCast<D> {
282    pub fn fn_name(&self) -> String {
283        let layout = frag_layout_str(&self.frag_input.layout);
284        let ident = frag_ident_str(&self.frag_input.ident);
285        let (m, n, k) = (self.frag_input.m, self.frag_input.n, self.frag_input.k);
286        let elem = self.frag_input.elem;
287        let elem_out = self.frag_output.elem;
288
289        format!("wmma_cast_{elem}_to_{elem_out}_{ident}_{m}x{n}x{k}_{layout}",)
290    }
291
292    pub fn format_extension(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
293        let input = self.frag_input;
294        let output = self.frag_output;
295        let name = self.fn_name();
296        let step = match output.ident {
297            FragmentIdent::Accumulator => {
298                get_output_accumulator_index_step(&self.frag_input.elem, &output)
299            }
300            _ => 1,
301        };
302
303        write!(
304            f,
305            "
306// Cast the fragment.
307__device__ void {name}(const {input}& input, {output}& output) {{
308    #pragma unroll
309    for (uint elemIdx = 0; elemIdx < uint(8); ++elemIdx) {{
310      output[elemIdx * {step}] = input[elemIdx];
311    }}
312}}
313        "
314        )
315    }
316}
317
318impl DialectWmmaCompiler<HipDialect<Self>> for WmmaIntrinsicCompiler {
319    fn compile_wmma_type_definitions(
320        f: &mut std::fmt::Formatter<'_>,
321        flags: &Flags<HipDialect<Self>>,
322    ) -> std::fmt::Result {
323        if flags.elem_bf16 {
324            f.write_str("typedef __bf16 bhalf8_t __attribute__((ext_vector_type(8)));\n")?;
325            f.write_str("typedef __bf16 bhalf16_t __attribute__((ext_vector_type(16)));\n")?;
326        }
327        if flags.elem_f16 {
328            f.write_str("typedef _Float16 half8_t __attribute__((ext_vector_type(8)));\n")?;
329            f.write_str("typedef _Float16 half16_t __attribute__((ext_vector_type(16)));\n")?;
330        }
331        f.write_str("typedef float float8_t __attribute__((ext_vector_type(8)));\n")
332    }
333
334    fn compile_wmma_fragment_declaration(
335        f: &mut std::fmt::Formatter<'_>,
336        val: &crate::shared::Value<HipDialect<Self>>,
337        ty: &crate::shared::Item<HipDialect<Self>>,
338    ) -> std::fmt::Result {
339        wmma_api_base::compile_fragment_declaration(f, val, ty)
340    }
341
342    fn compile_wmma_fragment(
343        f: &mut std::fmt::Formatter<'_>,
344        fragment: &FragmentType<HipDialect<Self>>,
345    ) -> std::fmt::Result {
346        match fragment.ident {
347            FragmentIdent::A | FragmentIdent::B => match fragment.elem {
348                Elem::F16 => write!(f, "half16_t"),
349                Elem::BF16 => write!(f, "bhalf16_t"),
350                other => panic!(
351                    "unsupported type {other} for fragment ident {:?}",
352                    fragment.ident
353                ),
354            },
355            FragmentIdent::Accumulator => match fragment.elem {
356                Elem::F16 => write!(f, "half16_t"),
357                Elem::BF16 => write!(f, "bhalf16_t"),
358                Elem::F32 => write!(f, "float8_t"),
359                other => panic!(
360                    "unsupported type {other} for fragment ident {:?}",
361                    fragment.ident
362                ),
363            },
364            FragmentIdent::_Dialect(_) => Ok(()),
365        }
366    }
367
368    fn compile_wmma_instruction(
369        f: &mut std::fmt::Formatter<'_>,
370        instruction: &WmmaInstruction<HipDialect<Self>>,
371    ) -> std::fmt::Result {
372        match instruction {
373            WmmaInstruction::Fill { frag, value } => {
374                let extension = WmmaFill::new(match frag.item().unwrap_ptr() {
375                    Item::Fragment(frag) => frag,
376                    _ => panic!(),
377                });
378                let name = extension.fn_name();
379                let frag = frag.fmt_ref();
380                writeln!(f, "{name}({frag}, {value});")
381            }
382            WmmaInstruction::Load {
383                frag,
384                ptr,
385                layout,
386                stride,
387            } => {
388                let extension = WmmaLoad::new(value_to_frag(frag), *layout);
389                let name = extension.fn_name();
390                let value_ptr = frag_as_ptr(f, ptr);
391                let frag = frag.fmt_ref();
392                writeln!(f, "{name}({frag}, {value_ptr}, {stride});")
393            }
394            WmmaInstruction::LdMatrix { .. } | WmmaInstruction::StMatrix { .. } => {
395                f.write_str("#error LdMatrix & StMatrix are not supported on HIP\n")
396            }
397            WmmaInstruction::Execute {
398                frag_a,
399                frag_b,
400                frag_c,
401                frag_d,
402                warp_size,
403            } => {
404                if *warp_size != 32 {
405                    f.write_str(
406                        "#error Only warp size of 32 supported for Wmma::Execute on HIP\n",
407                    )?;
408                }
409
410                let extension = WmmaExecute::new(
411                    value_to_frag(frag_a),
412                    value_to_frag(frag_b),
413                    value_to_frag(frag_c),
414                    value_to_frag(frag_d),
415                );
416                let name = extension.fn_name();
417                let frag_d = frag_d.fmt_ref();
418                writeln!(f, "{name}({frag_a}, {frag_b}, {frag_c}, {frag_d});")
419            }
420            WmmaInstruction::ExecuteManual {
421                shape,
422                frag_a,
423                frag_b,
424                frag_c,
425                frag_d,
426            } => {
427                Self::compile_manual_mma(f, ManualMma::new(*shape, frag_a, frag_b, frag_c, frag_d))
428            }
429            WmmaInstruction::ExecuteScaled {
430                shape,
431                frag_a,
432                frag_b,
433                frag_c,
434                frag_d,
435                scales_a,
436                scales_b,
437                scales_factor,
438            } => Self::compile_scaled_mma(
439                f,
440                ManualMma::new(*shape, frag_a, frag_b, frag_c, frag_d),
441                *scales_a,
442                *scales_b,
443                *scales_factor,
444            ),
445            WmmaInstruction::Store {
446                frag,
447                layout,
448                destination,
449                stride,
450            } => {
451                let extension = WmmaStore::new(value_to_frag(frag), *layout);
452                let name = extension.fn_name();
453                let output_ptr = frag_as_ptr(f, destination);
454                let frag = frag.fmt_ref();
455                writeln!(f, "{name}({frag}, {output_ptr}, {stride});")
456            }
457            WmmaInstruction::Cast { input, output } => {
458                let extension = WmmaCast::new(value_to_frag(input), value_to_frag(output));
459                let name = extension.fn_name();
460                let input = input.fmt_ref();
461                let output = output.fmt_ref();
462                writeln!(f, "{name}({input}, {output});")
463            }
464        }
465    }
466
467    fn compile_manual_mma(
468        f: &mut std::fmt::Formatter<'_>,
469        mma: ManualMma<HipDialect<Self>>,
470    ) -> std::fmt::Result {
471        compile_manual_mma(f, mma.shape, mma.frag_a, mma.frag_b, mma.frag_c, mma.frag_d)
472    }
473
474    fn compile_scaled_mma(
475        f: &mut std::fmt::Formatter<'_>,
476        _mma: ManualMma<HipDialect<Self>>,
477        _scales_a: Value<HipDialect<Self>>,
478        _scales_b: Value<HipDialect<Self>>,
479        _scales_factor: u32,
480    ) -> std::fmt::Result {
481        f.write_str("#error scaled mma not supported in HIP\n")
482    }
483
484    fn supported_wmma_combinations(arch: &AMDArchitecture) -> SupportedMmaCombinations {
485        // Reference: https://gpuopen.com/learn/wmma_on_rdna3/
486        let mut result: SupportedMmaCombinations = vec![];
487        if arch.is_wmma_capable() {
488            // Types fully supported.
489            let types = vec![
490                (
491                    gpu::ElemType::Float(gpu::FloatKind::F16), // m
492                    gpu::ElemType::Float(gpu::FloatKind::F16), // n
493                    gpu::ElemType::Float(gpu::FloatKind::F16), // k
494                ),
495                (
496                    gpu::ElemType::Float(gpu::FloatKind::F16),
497                    gpu::ElemType::Float(gpu::FloatKind::F16),
498                    gpu::ElemType::Float(gpu::FloatKind::F32),
499                ),
500                (
501                    gpu::ElemType::Float(gpu::FloatKind::BF16),
502                    gpu::ElemType::Float(gpu::FloatKind::BF16),
503                    gpu::ElemType::Float(gpu::FloatKind::F32),
504                ),
505            ];
506            let combinations: SupportedMmaCombinations = types
507                .into_iter()
508                .map(|(a, b, c)| MmaConfig {
509                    a_type: a.into(),
510                    b_type: b.into(),
511                    cd_type: c.into(),
512                    m: 16,
513                    n: 16,
514                    k: 16,
515                })
516                .collect();
517            result.extend(combinations);
518        }
519        result
520    }
521
522    fn supported_mma_combinations(arch: &AMDArchitecture) -> SupportedMmaCombinations {
523        supported_mma_combinations(arch)
524    }
525}
526
527fn get_output_accumulator_index_step<D: Dialect>(
528    input_elem: &Elem<D>,
529    output: &FragmentType<D>,
530) -> u32 {
531    // Each VGPR is 32 bit wide and there is 8 VGPR per lane, an accumulator can then be either:
532    // - a vector of 8 float
533    // - a vector of 16 half
534    // Depending on the precision used for the input, the whole 32 bits per register will be used or
535    // just only 16 bits. In such a case we always use the lower 16 bits (opsel set to false) which means
536    // that we only assign values to even indexes of the accumulator (0, 2, 4, ...)
537
538    assert_eq!(output.ident, FragmentIdent::<D>::Accumulator);
539
540    match input_elem {
541        Elem::F16 | Elem::BF16 | Elem::F32 => {
542            match output.elem {
543                // loading into accumulator of 16 half precision
544                Elem::F16 | Elem::BF16 => 2,
545                // loading into accumulator of 8 full precision
546                Elem::F32 => 1,
547                other => panic!("unsupported format {other} for {output}"),
548            }
549        }
550        other => panic!("unsupported format {other} for {input_elem}"),
551    }
552}
553
554pub(super) fn compile_manual_mma<D: Dialect>(
555    f: &mut std::fmt::Formatter<'_>,
556    shape: MmaShape<D>,
557    frag_a: &Value<D>,
558    frag_b: &Value<D>,
559    frag_c: &Value<D>,
560    frag_d: &Value<D>,
561) -> std::fmt::Result {
562    let extension = WmmaExecute::from_manual(shape, frag_a.elem(), frag_c.elem());
563
564    let cd_elems = shape.num_elems(FragmentIdent::<D>::Accumulator) / 32;
565
566    let frag_cd_step = 4usize.div_ceil(frag_c.elem().size());
567    let frag_d_tmp = Value::tmp_declared(Item::Scalar(Elem::<D>::I32)).fmt_left();
568
569    // Need to reconstruct the fragments from an array of vectors to a single vector type.
570    // This requires double indexing over both the array index and the vector index.
571    // Will generate something like
572    // `float8_t {arr[0].i_0, arr[0].i_1, arr[1].i_0, ...}`
573    let frag = |val: &Value<D>, len: usize| {
574        let frag: Vec<_> = if let Item::Vector(_, vec) = *val.item().value_ty() {
575            (0..len)
576                .map(|i| format!("{}.i_{}", val.index(i / vec), i % vec))
577                .collect()
578        } else {
579            (0..len).map(|i| format!("{}", val.index(i))).collect()
580        };
581        frag.join(", ")
582    };
583
584    let frag_a = frag(frag_a, 16);
585    let frag_b = frag(frag_b, 16);
586    // C matrix needs to be padded for f16, because it only uses the low bytes. The simplest way is
587    // to just replicate the same f16 in both halves of the register.
588    let frag_c = {
589        let frag: Vec<_> = if let Item::Vector(_, vec) = frag_c.item() {
590            (0..cd_elems as usize)
591                .flat_map(|i| {
592                    (0..frag_cd_step).map(move |_| format!("{frag_c}[{}].i_{}", i / vec, i % vec))
593                })
594                .collect()
595        } else {
596            (0..cd_elems as usize)
597                .flat_map(|i| (0..frag_cd_step).map(move |_| format!("{frag_c}[{}]", i)))
598                .collect()
599        };
600        frag.join(", ")
601    };
602
603    // Should optimize out
604    let name = extension.fn_name();
605
606    // Item is irrelevant
607    writeln!(f, "{} {frag_d_tmp} = {{}};", extension.frag_d)?;
608
609    writeln!(
610        f,
611        "{name}({}{{{frag_a}}}, {}{{{frag_b}}}, {}{{{frag_c}}}, {frag_d_tmp});",
612        extension.frag_a, extension.frag_b, extension.frag_c
613    )?;
614
615    for i in 0..cd_elems as usize {
616        if let Item::Vector(_, vec) = frag_d.item() {
617            writeln!(
618                f,
619                "{}.i_{} = {frag_d_tmp}[{i} * {frag_cd_step}];",
620                frag_d.index(i / vec),
621                i % vec
622            )?;
623        } else {
624            writeln!(
625                f,
626                "{} = {frag_d_tmp}[{i} * {frag_cd_step}];",
627                frag_d.index(i)
628            )?;
629        }
630    }
631
632    Ok(())
633}
634
635pub(super) fn supported_mma_combinations(arch: &AMDArchitecture) -> SupportedMmaCombinations {
636    // Correctness is wrong.
637    const ENABLED: bool = true;
638
639    if !ENABLED {
640        return Vec::new();
641    }
642
643    // Reference: https://gpuopen.com/learn/wmma_on_rdna3/
644    // Feel free to add more if additional intrinsics are supported for execute
645    let mut result: SupportedMmaCombinations = vec![];
646    if arch.is_wmma_capable() {
647        // Types fully supported.
648        let types = vec![
649            (
650                gpu::ElemType::Float(gpu::FloatKind::F16),
651                gpu::ElemType::Float(gpu::FloatKind::F32),
652            ),
653            (
654                gpu::ElemType::Float(gpu::FloatKind::BF16),
655                gpu::ElemType::Float(gpu::FloatKind::F32),
656            ),
657        ];
658        let combinations = types.into_iter().map(|(ab_elem, cd_elem)| MmaConfig {
659            a_type: ab_elem.into(),
660            b_type: ab_elem.into(),
661            cd_type: cd_elem.into(),
662            m: 16,
663            n: 16,
664            k: 16,
665        });
666        result.extend(combinations);
667    }
668    result
669}
670
671pub fn contiguous_elements_rdna3(ident: MatrixIdent, matrix: MatrixType) -> usize {
672    // Don't exceed swizzle atom and load width
673    let max_vector_size = 16 / matrix.storage.size();
674    match ident {
675        MatrixIdent::A | MatrixIdent::B => 16.min(max_vector_size),
676        MatrixIdent::Accumulator => 1,
677    }
678}
679
680// threads 0-15 and threads 16-31 of the wavefront hold the same fragments respectively
681// in other words fragments are duplicated
682// so lanes 0,16 / 1,17 / ... / 15, 31 are the same
683static WMMA_LANE_DEF: &str = "uint wmmaLane = uint(threadIdx.x % 16);";
cubecl_cpp/hip/mma/wmma_intrinsics_compiler.rs

cubecl_cpp/hip/mma/
wmma_intrinsics_compiler.rs
