cubecl_cpp/cuda/mma/
ptx_wmma_compiler.rs

1use super::WMMA_MINIMUM_VERSION;
2use crate::{
3    Dialect,
4    cuda::{
5        CudaDialect,
6        arch::CudaArchitecture,
7        ptx::{comma_separated, ldmatrix_call, stmatrix_call},
8    },
9    shared::{
10        Architecture, Component, DialectWmmaCompiler, Elem, Flags, FmtLeft, FragmentIdent,
11        FragmentLayout, FragmentType, Item, ManualMma, SupportedMmaCombinations,
12        SupportedScaledMmaCombinations, Value, WmmaInstruction,
13    },
14};
15use cubecl_core::ir::{
16    self as gpu, ConstantValue, MatrixIdent, MatrixType,
17    features::{MmaConfig, ScaledMmaConfig},
18};
19use itertools::Itertools;
20use std::fmt::Display;
21
22#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
23pub struct PtxWmmaCompiler {}
24
25impl DialectWmmaCompiler<CudaDialect<Self>> for PtxWmmaCompiler {
26    fn compile_wmma_includes(
27        f: &mut std::fmt::Formatter<'_>,
28        flags: &Flags<CudaDialect<Self>>,
29    ) -> std::fmt::Result {
30        // We need mma header for conversion
31        if flags.elem_tf32 {
32            f.write_str("#include <mma.h>\n")?;
33        }
34        Ok(())
35    }
36
37    fn compile_wmma_fragment_declaration(
38        f: &mut std::fmt::Formatter<'_>,
39        val: &Value<CudaDialect<Self>>,
40        ty: &Item<CudaDialect<Self>>,
41    ) -> std::fmt::Result {
42        let frag = match ty {
43            Item::Fragment(frag) => frag,
44            _ => panic!("load instruction expects a WmmaFragment"),
45        };
46        let reg_count = get_fragment_register_total_count(frag);
47        let ty = match frag.elem {
48            Elem::U8 | Elem::I8 | Elem::F16 | Elem::BF16 | Elem::TF32 => "unsigned int",
49            Elem::F32 => "float",
50            Elem::F64 => "double",
51            _ => panic!("unsupported type"),
52        };
53        writeln!(f, "{ty} {val}_store[{reg_count}];")
54    }
55
56    fn compile_wmma_instruction(
57        f: &mut std::fmt::Formatter<'_>,
58        instruction: &WmmaInstruction<CudaDialect<Self>>,
59    ) -> std::fmt::Result {
60        match instruction {
61            WmmaInstruction::Fill { frag, value } => {
62                let frag_ty = match frag.item() {
63                    Item::Fragment(frag) => frag,
64                    _ => panic!("load instruction expects a WmmaFragment"),
65                };
66                let reg_count = get_fragment_register_total_count(&frag_ty);
67                write!(
68                    f,
69                    "// fill
70for (uint i = 0; i < uint({reg_count}); ++i) {{
71  {frag}[i] = {value};
72}}
73 "
74                )
75            }
76            WmmaInstruction::Load {
77                frag,
78                ptr,
79                stride,
80                layout,
81            } => {
82                let frag_ty = match frag.item() {
83                    Item::Fragment(frag) => frag,
84                    _ => panic!("load instruction expects a WmmaFragment"),
85                };
86                // Important note: the current frontend has been designed around
87                // CUDA wmma which is not optimal in the case of PTX wmma and mma
88                // We choose here to use the layout defined in the fragment first,
89                // if it is unknown and we look into the layout passed to the instruction.
90                let layout = if frag_ty.layout.is_some() {
91                    get_fragment_layout_qualifier(frag)
92                } else if let Some(layout) = layout {
93                    get_qualifier_from_layout(layout)
94                } else {
95                    panic!("unknown matrix layout for wmma load instruction");
96                };
97                // instruction qualifiers
98                let ty = get_type_qualifier(ptr);
99                let matrix = match frag_ty.ident {
100                    FragmentIdent::A => "a",
101                    FragmentIdent::B => "b",
102                    FragmentIdent::Accumulator => "c",
103                    FragmentIdent::_Dialect(_) => unreachable!(),
104                };
105                let value_ptr_ty = ptr.item().value_ptr();
106                let opcode = match frag_ty.elem {
107                    Elem::U8 | Elem::I8 | Elem::F16 | Elem::BF16 | Elem::F32 | Elem::TF32 => {
108                        format!(
109                            "wmma.load.{matrix}.sync.aligned.{layout}.m{}n{}k{}.{ty}",
110                            frag_ty.m, frag_ty.n, frag_ty.k,
111                        )
112                    }
113                    other => panic!("{other} fragment type not supported"),
114                };
115                // constraints
116                let mut reg_count = 0;
117                let (regs_decl, out_constraints) =
118                    get_value_regs_decl_constraints(frag, true, &mut reg_count);
119                let buffer_reg = format_reg_and_inc(&mut reg_count);
120                let (stride_reg, stride_constraint) =
121                    get_value_regs_decl_constraints(stride, false, &mut reg_count);
122                let tmp_ptr = Value::tmp(value_ptr_ty);
123                let tmp_ptr_left = tmp_ptr.fmt_left();
124                write!(
125                    f,
126                    r#"// load
127{tmp_ptr_left} = ({value_ptr_ty}){ptr};
128asm volatile(
129    "{opcode} "
130    "{{{regs_decl}}}, [{buffer_reg}], {stride_reg};\n"
131    : {out_constraints}
132    : "l"({tmp_ptr}){stride_constraint}
133);
134"#
135                )
136            }
137            WmmaInstruction::LdMatrix {
138                output,
139                ptr,
140                factor,
141                transpose,
142            } => f.write_str(&ldmatrix_call(output, ptr, factor, transpose)),
143            WmmaInstruction::StMatrix {
144                registers,
145                ptr,
146                factor,
147                transpose,
148            } => f.write_str(&stmatrix_call(registers, ptr, factor, transpose)),
149            WmmaInstruction::Execute {
150                frag_a,
151                frag_b,
152                frag_c,
153                frag_d,
154                ..
155            } => {
156                let frag_a_ty = match frag_a.item() {
157                    Item::Fragment(frag) => frag,
158                    _ => panic!("value should be WmmaFragment"),
159                };
160                let layout_a = get_fragment_layout_qualifier(frag_a);
161                let layout_b = get_fragment_layout_qualifier(frag_b);
162                let type_c = get_type_qualifier(frag_c);
163                let type_d = get_type_qualifier(frag_d);
164                let opcode = match frag_a.elem() {
165                    Elem::U8 | Elem::I8 | Elem::F16 | Elem::F32 => format!(
166                        "wmma.mma.sync.aligned.m{}n{}k{}.{layout_a}.{layout_b}.{type_d}.{type_c}",
167                        frag_a_ty.m, frag_a_ty.n, frag_a_ty.k,
168                    ),
169                    Elem::BF16 => format!(
170                        "wmma.mma.sync.aligned.{layout_a}.{layout_b}.m{}n{}k{}.f32.bf16.bf16.f32",
171                        frag_a_ty.m, frag_a_ty.n, frag_a_ty.k,
172                    ),
173                    Elem::TF32 => format!(
174                        "wmma.mma.sync.aligned.{layout_a}.{layout_b}.m{}n{}k{}.f32.tf32.tf32.f32",
175                        frag_a_ty.m, frag_a_ty.n, frag_a_ty.k,
176                    ),
177                    other => panic!("{other} fragment type not supported"),
178                };
179                let mut reg_count = 0;
180                // order matters, declare the registers in the same order as the intrinsic
181                let (regs_decl_d, out_constraints_d) =
182                    get_value_regs_decl_constraints(frag_d, true, &mut reg_count);
183                let (regs_decl_a, in_constraints_a) =
184                    get_value_regs_decl_constraints(frag_a, false, &mut reg_count);
185                let (regs_decl_b, in_constraints_b) =
186                    get_value_regs_decl_constraints(frag_b, false, &mut reg_count);
187                let (regs_decl_c, in_constraints_c) =
188                    get_value_regs_decl_constraints(frag_c, false, &mut reg_count);
189                write!(
190                    f,
191                    r#"// execute
192asm volatile(
193    "{opcode} "
194    "{{{regs_decl_d}}}, "
195    "{{{regs_decl_a}}}, "
196    "{{{regs_decl_b}}}, "
197    "{{{regs_decl_c}}};\n"
198    : {out_constraints_d}
199    : {in_constraints_a}, {in_constraints_b}, {in_constraints_c}
200);
201"#
202                )
203            }
204            WmmaInstruction::ExecuteManual {
205                shape,
206                frag_a,
207                frag_b,
208                frag_c,
209                frag_d,
210            } => {
211                Self::compile_manual_mma(f, ManualMma::new(*shape, frag_a, frag_b, frag_c, frag_d))
212            }
213            WmmaInstruction::ExecuteScaled {
214                shape,
215                frag_a,
216                frag_b,
217                frag_c,
218                frag_d,
219
220                scales_a,
221                scales_b,
222                scales_factor,
223            } => Self::compile_scaled_mma(
224                f,
225                ManualMma::new(*shape, frag_a, frag_b, frag_c, frag_d),
226                *scales_a,
227                *scales_b,
228                *scales_factor,
229            ),
230            WmmaInstruction::Store {
231                frag,
232                stride,
233                destination,
234                layout,
235            } => {
236                let frag_acc = match frag.item() {
237                    Item::Fragment(frag) => frag,
238                    _ => panic!("value should be WmmaFragment"),
239                };
240                // instruction qualifiers
241                let layout = match layout {
242                    FragmentLayout::ColMajor => "col",
243                    FragmentLayout::RowMajor => "row",
244                    FragmentLayout::_Dialect(..) => unreachable!(),
245                };
246                let opcode = match frag.elem() {
247                    Elem::F16 | Elem::BF16 => format!(
248                        // hack because wmma.store does not support bf16
249                        // f16 should still work correctly for bf16 as long
250                        // as the input registers are in correct format
251                        "wmma.store.d.sync.aligned.{layout}.m{}n{}k{}.f16",
252                        frag_acc.m, frag_acc.n, frag_acc.k,
253                    ),
254                    Elem::TF32 | Elem::F32 => format!(
255                        // same hack for tf32
256                        "wmma.store.d.sync.aligned.{layout}.m{}n{}k{}.f32",
257                        frag_acc.m, frag_acc.n, frag_acc.k,
258                    ),
259                    Elem::I32 => format!(
260                        // same hack for tf32
261                        "wmma.store.d.sync.aligned.{layout}.m{}n{}k{}.s32",
262                        frag_acc.m, frag_acc.n, frag_acc.k,
263                    ),
264                    other => panic!("{other} fragment type not supported"),
265                };
266                // constraints
267                let mut reg_count = 0;
268                let buffer_reg = format_reg_and_inc(&mut reg_count);
269                // offset and stride can be passed as local const or as const scalar
270                // we need to handle both cases correctly in the asm.
271                let (stride_reg, stride_constraint) =
272                    get_value_regs_decl_constraints(stride, false, &mut reg_count);
273                // we start at 2 because of the buffer address calculation
274                let (regs_decl, in_constraints) =
275                    get_value_regs_decl_constraints(frag, false, &mut reg_count);
276                write!(
277                    f,
278                    r#"// store
279asm volatile(
280    "{opcode} "
281    "[{buffer_reg}], {{{regs_decl}}}, {stride_reg};\n"
282    :
283    : "l"({destination}),
284      {in_constraints}{stride_constraint}
285);
286"#
287                )
288            }
289            WmmaInstruction::Cast { input, output } => {
290                let frag = match input.item() {
291                    Item::Fragment(frag) => frag,
292                    _ => panic!("value should be WmmaFragment"),
293                };
294                let reg_count = get_fragment_register_total_count(&frag);
295                match output.elem() {
296                    Elem::F16 => {
297                        write!(
298                            f,
299                            "// cast
300for (int i = 0; i < {reg_count}; ++i) {{
301    __half h_lo = __float2half_rn({input}[2*i + 0]);
302    __half h_hi = __float2half_rn({input}[2*i + 1]);
303    __half2 h2 = __halves2half2(h_lo, h_hi);
304    {output}[i] = *reinterpret_cast<unsigned int*>(&h2);
305}}
306"
307                        )
308                    }
309                    Elem::BF16 => {
310                        write!(
311                            f,
312                            "// cast
313for (int i = 0; i < {reg_count}; ++i) {{
314    __nv_bfloat16 b_lo = __float2bfloat16({input}[2*i + 0]);
315    __nv_bfloat16 b_hi = __float2bfloat16({input}[2*i + 1]);
316    __nv_bfloat162 bf2 = __halves2bfloat162(b_lo, b_hi);
317    {output}[i] = *reinterpret_cast<unsigned int*>(&bf2);
318}}
319"
320                        )
321                    }
322                    other => panic!("casting fragment to {other} not supported"),
323                }
324            }
325        }
326    }
327
328    fn compile_manual_mma(
329        f: &mut std::fmt::Formatter<'_>,
330        mma: ManualMma<CudaDialect<Self>>,
331    ) -> std::fmt::Result {
332        compile_manual_mma(f, mma)
333    }
334
335    fn compile_scaled_mma(
336        f: &mut std::fmt::Formatter<'_>,
337        mma: ManualMma<CudaDialect<Self>>,
338        scales_a: Value<CudaDialect<Self>>,
339        scales_b: Value<CudaDialect<Self>>,
340        scales_factor: u32,
341    ) -> std::fmt::Result {
342        compile_scaled_mma(f, mma, scales_a, scales_b, scales_factor)
343    }
344
345    fn supported_wmma_combinations(arch: &CudaArchitecture) -> SupportedMmaCombinations {
346        let mut result: SupportedMmaCombinations = vec![];
347        if arch.get_version() >= WMMA_MINIMUM_VERSION {
348            // Types fully supported.
349            let types = vec![
350                (
351                    gpu::ElemType::Float(gpu::FloatKind::F16), // m
352                    gpu::ElemType::Float(gpu::FloatKind::F16), // n
353                    gpu::ElemType::Float(gpu::FloatKind::F16), // k
354                ),
355                (
356                    gpu::ElemType::Float(gpu::FloatKind::F16),
357                    gpu::ElemType::Float(gpu::FloatKind::F16),
358                    gpu::ElemType::Float(gpu::FloatKind::F32),
359                ),
360                (
361                    gpu::ElemType::Float(gpu::FloatKind::BF16),
362                    gpu::ElemType::Float(gpu::FloatKind::BF16),
363                    gpu::ElemType::Float(gpu::FloatKind::F32),
364                ),
365            ];
366            let combinations: SupportedMmaCombinations = types
367                .into_iter()
368                .map(|(a, b, cd)| MmaConfig {
369                    a_type: a.into(),
370                    b_type: b.into(),
371                    cd_type: cd.into(),
372                    m: 16,
373                    n: 16,
374                    k: 16,
375                })
376                .collect();
377            result.extend(combinations);
378            if arch.get_version() >= 72 {
379                result.extend([
380                    MmaConfig {
381                        a_type: gpu::ElemType::UInt(gpu::UIntKind::U8).into(),
382                        b_type: gpu::ElemType::UInt(gpu::UIntKind::U8).into(),
383                        cd_type: gpu::ElemType::Int(gpu::IntKind::I32).into(),
384                        m: 16,
385                        n: 16,
386                        k: 16,
387                    },
388                    MmaConfig {
389                        a_type: gpu::ElemType::Int(gpu::IntKind::I8).into(),
390                        b_type: gpu::ElemType::Int(gpu::IntKind::I8).into(),
391                        cd_type: gpu::ElemType::Int(gpu::IntKind::I32).into(),
392                        m: 16,
393                        n: 16,
394                        k: 16,
395                    },
396                ]);
397            }
398            if arch.get_version() >= 80 {
399                result.push(MmaConfig {
400                    a_type: gpu::ElemType::Float(gpu::FloatKind::TF32).into(),
401                    b_type: gpu::ElemType::Float(gpu::FloatKind::TF32).into(),
402                    cd_type: gpu::ElemType::Float(gpu::FloatKind::F32).into(),
403                    m: 16,
404                    n: 16,
405                    k: 8,
406                });
407            }
408        }
409        result
410    }
411
412    fn supported_mma_combinations(arch: &CudaArchitecture) -> SupportedMmaCombinations {
413        supported_mma_combinations(arch)
414    }
415
416    fn supported_scaled_mma_combinations(
417        arch: &CudaArchitecture,
418    ) -> SupportedScaledMmaCombinations {
419        supported_scaled_mma_combinations(arch)
420    }
421}
422
423fn get_fragment_register_total_count(frag: &FragmentType<CudaDialect<PtxWmmaCompiler>>) -> u32 {
424    let FragmentType {
425        ident,
426        m,
427        n,
428        k,
429        elem,
430        ..
431    } = frag;
432    let elements = match ident {
433        FragmentIdent::A => m * k,
434        FragmentIdent::B => k * n,
435        FragmentIdent::Accumulator => m * n,
436        _ => unreachable!(),
437    };
438    let bits_per_elem = elem.size_bits() as u32;
439    // TODO: retrieve the warp size from the compiler CompilationOptions
440    let lanes_per_reg = 32 / bits_per_elem;
441    // choose threads-per-frag:
442    // - accumulators always use 32 lanes
443    // - A/B use 16 lanes _except_ TF32 (k=8) which also uses 32 lanes
444    let threads_per_frag = match ident {
445        FragmentIdent::Accumulator => 32,
446        FragmentIdent::A | FragmentIdent::B => {
447            if frag.elem == Elem::TF32 {
448                32
449            } else {
450                16
451            }
452        }
453        _ => unreachable!(),
454    };
455
456    elements / (lanes_per_reg * threads_per_frag)
457}
458
459fn get_type_qualifier(val: &Value<CudaDialect<PtxWmmaCompiler>>) -> String {
460    match val.elem() {
461        Elem::U8 => "u8",
462        Elem::I8 => "s8",
463        Elem::F16 => "f16",
464        Elem::BF16 => "bf16",
465        Elem::F32 => "f32",
466        Elem::TF32 => "tf32",
467        Elem::I32 => "s32",
468        Elem::F64 => "f64",
469        _ => panic!("unsupported WMMA fragment type"),
470    }
471    .to_string()
472}
473
474fn get_fragment_layout_qualifier(val: &Value<CudaDialect<PtxWmmaCompiler>>) -> String {
475    let frag = match val.item() {
476        Item::Fragment(frag) => frag,
477        _ => panic!("value should be WmmaFragment"),
478    };
479    match frag.layout {
480        Some(layout) => get_qualifier_from_layout(&layout),
481        None => "".to_string(),
482    }
483}
484
485fn get_qualifier_from_layout(layout: &FragmentLayout<CudaDialect<PtxWmmaCompiler>>) -> String {
486    match layout {
487        FragmentLayout::ColMajor => "col",
488        FragmentLayout::RowMajor => "row",
489        FragmentLayout::_Dialect(..) => unreachable!(),
490    }
491    .to_string()
492}
493
494fn get_value_regs_decl_constraints(
495    val: &Value<CudaDialect<PtxWmmaCompiler>>,
496    output: bool,
497    reg_count: &mut u8,
498) -> (String, String) {
499    match val {
500        _ if let Item::Fragment(frag) = val.item() => {
501            let reg_total_count = get_fragment_register_total_count(&frag);
502            let reg_decl = (0..reg_total_count)
503                .map(|_| format_reg_and_inc(reg_count))
504                .collect::<Vec<_>>()
505                .join(",");
506            let frag_elem = frag.elem;
507            let modifier = format!(
508                "{}{}",
509                if output { "=" } else { "" },
510                match frag_elem {
511                    Elem::F32 => "f",
512                    Elem::F64 => "d",
513                    _ => "r",
514                },
515            );
516            let constraints = (0..reg_total_count)
517                .map(|i| format!("\"{modifier}\"({val}[{i}])"))
518                .collect::<Vec<_>>()
519                .join(", ");
520            (reg_decl, constraints)
521        }
522        Value::Constant(number, ..) => match number {
523            ConstantValue::UInt(val, ..) => (val.to_string(), "".to_string()),
524            _ => panic!("value should be an unsigned integer"),
525        },
526        _ => (format_reg_and_inc(reg_count), format!(r#", "r"({val})"#)),
527    }
528}
529
530fn format_reg_and_inc(count: &mut u8) -> String {
531    let res = format!("%{count}");
532    *count += 1;
533    res
534}
535
536fn as_ty_idx<D: Dialect>(val: &Value<D>, idx: impl Display, ty: impl Display) -> String {
537    format!("reinterpret_cast<{ty}*>({})[{idx}]", val.fmt_ptr())
538}
539
540fn as_const_ty_idx<D: Dialect>(val: &Value<D>, idx: impl Display, ty: impl Display) -> String {
541    format!("reinterpret_cast<const {ty}*>({})[{idx}]", val.fmt_ptr())
542}
543
544pub(super) fn compile_manual_mma<D: Dialect>(
545    f: &mut core::fmt::Formatter<'_>,
546    mma: ManualMma<D>,
547) -> std::fmt::Result {
548    let ManualMma {
549        shape,
550        frag_a,
551        frag_b,
552        frag_c,
553        frag_d,
554    } = mma;
555
556    let a_elem = frag_a.elem().unpacked();
557    let b_elem = frag_b.elem().unpacked();
558    let cd_elem = frag_c.elem().unpacked();
559
560    let ab_ty = match a_elem {
561        Elem::F32 => &format!("{}", Elem::<D>::F32),
562        _ => &format!("{}", Elem::<D>::U32),
563    };
564    let cd_ty = match cd_elem {
565        Elem::F32 => &format!("{}", Elem::<D>::F32),
566        _ => &format!("{}", Elem::<D>::U32),
567    };
568
569    let a_elems = shape.num_elems(FragmentIdent::<D>::A) / 32;
570    let b_elems = shape.num_elems(FragmentIdent::<D>::B) / 32;
571    let cd_elems = shape.num_elems(FragmentIdent::<D>::Accumulator) / 32;
572
573    let a_regs = a_elems as usize / (32 / frag_a.elem().unpacked().size_bits());
574    let b_regs = b_elems as usize / (32 / frag_b.elem().unpacked().size_bits());
575    let cd_regs = cd_elems as usize / (32 / frag_c.elem().unpacked().size_bits());
576
577    let frag_a = (0..a_regs).map(|i| as_const_ty_idx(frag_a, i, ab_ty));
578    let frag_b = (0..b_regs).map(|i| as_const_ty_idx(frag_b, i, ab_ty));
579    let frag_c = (0..cd_regs).map(|i| as_const_ty_idx(frag_c, i, cd_ty));
580    let frag_d = (0..cd_regs).map(|i| as_ty_idx(frag_d, i, cd_ty));
581
582    let args = comma_separated(frag_a.chain(frag_b).chain(frag_c).chain(frag_d));
583    write!(
584        f,
585        "__mma_m16n8k{}_{}_{}_{}({args});",
586        shape.k, a_elem, b_elem, cd_elem
587    )
588}
589
590pub(super) fn compile_scaled_mma<D: Dialect>(
591    f: &mut core::fmt::Formatter<'_>,
592    mma: ManualMma<D>,
593    scales_a: Value<D>,
594    scales_b: Value<D>,
595    scales_factor: u32,
596) -> std::fmt::Result {
597    let ManualMma {
598        shape,
599        frag_a,
600        frag_b,
601        frag_c,
602        frag_d,
603    } = mma;
604
605    let a_elem = frag_a.elem().unpacked();
606    let b_elem = frag_b.elem().unpacked();
607    let cd_elem = frag_c.elem().unpacked();
608
609    let ab_ty = &format!("{}", Elem::<D>::U32);
610    let cd_ty = &format!("{}", Elem::<D>::F32);
611
612    let a_elems = shape.num_elems(FragmentIdent::<D>::A) / 32;
613    let b_elems = shape.num_elems(FragmentIdent::<D>::B) / 32;
614    let cd_elems = shape.num_elems(FragmentIdent::<D>::Accumulator) / 32;
615
616    let a_regs = a_elems as usize / (32 / frag_a.elem().unpacked().size_bits());
617    let b_regs = b_elems as usize / (32 / frag_b.elem().unpacked().size_bits());
618    let cd_regs = cd_elems as usize / (32 / frag_c.elem().unpacked().size_bits());
619
620    let frag_a = (0..a_regs).map(|i| as_const_ty_idx(frag_a, i, ab_ty));
621    let frag_b = (0..b_regs).map(|i| as_const_ty_idx(frag_b, i, ab_ty));
622    let frag_c = (0..cd_regs).map(|i| as_const_ty_idx(frag_c, i, cd_ty));
623    let frag_d = (0..cd_regs).map(|i| as_ty_idx(frag_d, i, cd_ty));
624
625    let scales_a = scales_a.ensure_lvalue(f)?;
626    let scales_b = scales_b.ensure_lvalue(f)?;
627
628    let fragments = comma_separated(frag_a.chain(frag_b).chain(frag_c).chain(frag_d));
629    write!(
630        f,
631        "__mma_scaled_{scales_factor}x_m16n8k{}_{}_{}_{}({fragments}, reinterpret_cast<const uint32&>({scales_a}), reinterpret_cast<const uint32&>({scales_b}));",
632        shape.k, a_elem, b_elem, cd_elem
633    )
634}
635
636pub(super) fn supported_mma_combinations(arch: &CudaArchitecture) -> SupportedMmaCombinations {
637    let mut result: SupportedMmaCombinations = vec![];
638    // Higher than WMMA because we only support the newest shapes. Other shapes would make things
639    // very complicated.
640    // Also only use f32 accumulators for now
641    if arch.get_version() >= 80 {
642        result.extend([
643            MmaConfig {
644                a_type: gpu::ElemType::Float(gpu::FloatKind::F16).into(), // a
645                b_type: gpu::ElemType::Float(gpu::FloatKind::F16).into(), // b
646                cd_type: gpu::ElemType::Float(gpu::FloatKind::F32).into(), // cd
647                m: 16,
648                n: 8,
649                k: 16,
650            },
651            MmaConfig {
652                a_type: gpu::ElemType::Float(gpu::FloatKind::BF16).into(),
653                b_type: gpu::ElemType::Float(gpu::FloatKind::BF16).into(),
654                cd_type: gpu::ElemType::Float(gpu::FloatKind::F32).into(),
655                m: 16,
656                n: 8,
657                k: 16,
658            },
659            MmaConfig {
660                a_type: gpu::ElemType::Float(gpu::FloatKind::TF32).into(),
661                b_type: gpu::ElemType::Float(gpu::FloatKind::TF32).into(),
662                cd_type: gpu::ElemType::Float(gpu::FloatKind::F32).into(),
663                m: 16,
664                n: 8,
665                k: 8,
666            },
667            MmaConfig {
668                a_type: gpu::ElemType::Int(gpu::IntKind::I8).into(),
669                b_type: gpu::ElemType::Int(gpu::IntKind::I8).into(),
670                cd_type: gpu::ElemType::Int(gpu::IntKind::I32).into(),
671                m: 16,
672                n: 8,
673                k: 32,
674            },
675            MmaConfig {
676                a_type: gpu::ElemType::UInt(gpu::UIntKind::U8).into(),
677                b_type: gpu::ElemType::UInt(gpu::UIntKind::U8).into(),
678                cd_type: gpu::ElemType::Int(gpu::IntKind::I32).into(),
679                m: 16,
680                n: 8,
681                k: 32,
682            },
683            MmaConfig {
684                a_type: gpu::ElemType::Int(gpu::IntKind::I8).into(),
685                b_type: gpu::ElemType::UInt(gpu::UIntKind::U8).into(),
686                cd_type: gpu::ElemType::Int(gpu::IntKind::I32).into(),
687                m: 16,
688                n: 8,
689                k: 32,
690            },
691            MmaConfig {
692                a_type: gpu::ElemType::UInt(gpu::UIntKind::U8).into(),
693                b_type: gpu::ElemType::Int(gpu::IntKind::I8).into(),
694                cd_type: gpu::ElemType::Int(gpu::IntKind::I32).into(),
695                m: 16,
696                n: 8,
697                k: 32,
698            },
699            // TODO: u4/i4/b1, there's no types for them yet
700        ]);
701    }
702    if arch.get_version() >= 89 {
703        let f8f6f4_types = [
704            gpu::FloatKind::E4M3,
705            gpu::FloatKind::E5M2,
706            gpu::FloatKind::E3M2,
707            gpu::FloatKind::E2M3,
708            gpu::FloatKind::E2M1,
709        ];
710        let combinations = f8f6f4_types.iter().cartesian_product(f8f6f4_types.iter());
711        result.extend(combinations.map(|(t1, t2)| MmaConfig {
712            a_type: gpu::ElemType::Float(*t1).into(),
713            b_type: gpu::ElemType::Float(*t2).into(),
714            cd_type: gpu::ElemType::Float(gpu::FloatKind::F32).into(),
715            m: 16,
716            n: 8,
717            k: 32,
718        }));
719    }
720    // Warning: this likely does not follow the same layout pattern as those after 80
721    if arch.get_version() >= 70 && arch.get_version() < 80 {
722        result.push(MmaConfig {
723            a_type: gpu::ElemType::Float(gpu::FloatKind::F16).into(),
724            b_type: gpu::ElemType::Float(gpu::FloatKind::F16).into(),
725            cd_type: gpu::ElemType::Float(gpu::FloatKind::F32).into(),
726            m: 16,
727            n: 8,
728            k: 8,
729        });
730    }
731    result
732}
733
734pub(super) fn supported_scaled_mma_combinations(
735    arch: &CudaArchitecture,
736) -> SupportedScaledMmaCombinations {
737    let mut result: SupportedScaledMmaCombinations = vec![];
738    // sm_120f
739    if arch.get_version() >= 120 && arch.get_version() < 130 {
740        let f8f6f4_types = [
741            gpu::FloatKind::E4M3,
742            gpu::FloatKind::E5M2,
743            gpu::FloatKind::E3M2,
744            gpu::FloatKind::E2M3,
745            gpu::FloatKind::E2M1,
746        ];
747        let combinations = f8f6f4_types
748            .iter()
749            .flat_map(|t1| f8f6f4_types.iter().map(move |t2| (t1, t2)));
750
751        result.extend(combinations.map(|(t1, t2)| ScaledMmaConfig {
752            a_type: gpu::ElemType::Float(*t1).into(),
753            b_type: gpu::ElemType::Float(*t2).into(),
754            cd_type: gpu::ElemType::Float(gpu::FloatKind::F32).into(),
755            scales_type: gpu::ElemType::Float(gpu::FloatKind::UE8M0).into(),
756            m: 16,
757            n: 8,
758            k: 32,
759            scales_factor: 1,
760        }));
761
762        result.extend([
763            ScaledMmaConfig {
764                a_type: gpu::StorageType::Packed(gpu::ElemType::Float(gpu::FloatKind::E2M1), 2),
765                b_type: gpu::StorageType::Packed(gpu::ElemType::Float(gpu::FloatKind::E2M1), 2),
766                cd_type: gpu::ElemType::Float(gpu::FloatKind::F32).into(),
767                scales_type: gpu::ElemType::Float(gpu::FloatKind::UE8M0).into(),
768                m: 16,
769                n: 8,
770                k: 64,
771                scales_factor: 2,
772            },
773            // Sign of scales is ignored
774            ScaledMmaConfig {
775                a_type: gpu::StorageType::Packed(gpu::ElemType::Float(gpu::FloatKind::E2M1), 2),
776                b_type: gpu::StorageType::Packed(gpu::ElemType::Float(gpu::FloatKind::E2M1), 2),
777                cd_type: gpu::ElemType::Float(gpu::FloatKind::F32).into(),
778                scales_type: gpu::ElemType::Float(gpu::FloatKind::E4M3).into(),
779                m: 16,
780                n: 8,
781                k: 64,
782                scales_factor: 4,
783            },
784        ]);
785    }
786    result
787}
788
789pub fn contiguous_elements_cuda(ident: MatrixIdent, matrix: MatrixType) -> usize {
790    match ident {
791        MatrixIdent::A | MatrixIdent::B => 32 / matrix.storage.size_bits(),
792        MatrixIdent::Accumulator => 2,
793    }
794}
cubecl_cpp/cuda/mma/ptx_wmma_compiler.rs

cubecl_cpp/cuda/mma/
ptx_wmma_compiler.rs
