cubecl_cpp/cuda/ptx/
mma.rs

1use crate::{
2    Dialect,
3    shared::{Component, Elem, FP4Kind, FP6Kind, FP8Kind, Value},
4};
5
6#[allow(clippy::too_many_arguments)]
7pub fn mma_template<D: Dialect>(
8    a_elem: Elem<D>,
9    b_elem: Elem<D>,
10    cd_elem: Elem<D>,
11    k: u32,
12    n_a_registers: usize,
13    n_b_registers: usize,
14    n_c_registers: usize,
15    n_d_registers: usize,
16) -> String {
17    let a_ty = mma_ty(a_elem);
18    let b_ty = mma_ty(b_elem);
19    let cd_ty = mma_ty(cd_elem);
20
21    let ab_arg_ty = match a_elem {
22        Elem::F32 => &format!("{}", Elem::<D>::F32),
23        _ => &format!("{}", Elem::<D>::U32),
24    };
25    let cd_arg_ty = match cd_elem {
26        Elem::F32 => &format!("{}", Elem::<D>::F32),
27        _ => &format!("{}", Elem::<D>::U32),
28    };
29
30    let args_a = (0..n_a_registers).map(|i| format!("{ab_arg_ty} const &reg_a_{i}"));
31    let args_b = (0..n_b_registers).map(|i| format!("{ab_arg_ty} const &reg_b_{i}"));
32    let args_c = (0..n_c_registers).map(|i| format!("{cd_arg_ty} const &reg_c_{i}"));
33    let args_d = (0..n_d_registers).map(|i| format!("{cd_arg_ty} &reg_d_{i}"));
34    let args = args_a
35        .chain(args_b)
36        .chain(args_c)
37        .chain(args_d)
38        .collect::<Vec<_>>()
39        .join(", ");
40
41    let kind = if is_fp6_fp4(a_elem) || is_fp6_fp4(b_elem) {
42        ".kind::f8f6f4"
43    } else {
44        ""
45    };
46
47    let mut idx = 0usize;
48
49    let placeholders_d = comma_separated((0..n_d_registers).map(|_| placeholder(&mut idx)));
50    let placeholders_d = format!("{{{placeholders_d}}}");
51
52    let placeholders_a = comma_separated((0..n_a_registers).map(|_| placeholder(&mut idx)));
53    let placeholders_a = format!("{{{placeholders_a}}}");
54
55    let placeholders_b = comma_separated((0..n_b_registers).map(|_| placeholder(&mut idx)));
56    let placeholders_b = format!("{{{placeholders_b}}}");
57
58    let placeholders_c = comma_separated((0..n_c_registers).map(|_| placeholder(&mut idx)));
59    let placeholders_c = format!("{{{placeholders_c}}}");
60
61    let params_out =
62        comma_separated((0..n_d_registers).map(|i| as_reg(&format!("reg_d_{i}"), cd_elem, true)));
63    let params_a = (0..n_a_registers).map(|i| as_reg(&format!("reg_a_{i}"), a_elem, false));
64    let params_b = (0..n_b_registers).map(|i| as_reg(&format!("reg_b_{i}"), b_elem, false));
65    let params_c = (0..n_c_registers).map(|i| as_reg(&format!("reg_c_{i}"), cd_elem, false));
66    let params_in = comma_separated(params_a.chain(params_b).chain(params_c));
67
68    format!(
69        r#"
70inline __device__ void
71__mma_m16n8k{k}_{a_elem}_{b_elem}_{cd_elem}({args}) {{
72  asm volatile("mma.sync.aligned.m16n8k{k}.row.col{kind}.{cd_ty}.{a_ty}.{b_ty}.{cd_ty}"
73               " {placeholders_d}, {placeholders_a}, {placeholders_b}, {placeholders_c};"
74               : {params_out}
75               : {params_in});
76    }}
77    "#
78    )
79}
80
81fn is_fp6_fp4<D: Dialect>(elem: Elem<D>) -> bool {
82    matches!(elem, Elem::<D>::FP4(_) | Elem::<D>::FP6(_))
83}
84
85#[allow(clippy::too_many_arguments)]
86pub fn mma_scaled_template<D: Dialect>(
87    a_elem: Elem<D>,
88    b_elem: Elem<D>,
89    cd_elem: Elem<D>,
90    k: u32,
91    n_a_registers: usize,
92    n_b_registers: usize,
93    n_c_registers: usize,
94    n_d_registers: usize,
95    scales_elem: Elem<D>,
96    scales_factor: u32,
97) -> String {
98    let a_ty = mma_ty(a_elem);
99    let b_ty = mma_ty(b_elem);
100    let cd_ty = mma_ty(cd_elem);
101    // Needs custom mapping because of the ignored sign bit
102    let s_ty = match scales_elem {
103        Elem::FP8(FP8Kind::UE8M0) => "ue8m0",
104        Elem::FP8(FP8Kind::E4M3) => "ue4m3",
105        _ => panic!("Unsupported scales type"),
106    };
107
108    let kind = match scales_factor {
109        1 => "mxf8f6f4",
110        2 | 4 => "mxf4nvf4",
111        _ => panic!("Unsupported scales factor"),
112    };
113
114    let ab_arg_ty = match a_elem {
115        Elem::F32 => &format!("{}", Elem::<D>::F32),
116        _ => &format!("{}", Elem::<D>::U32),
117    };
118    let cd_arg_ty = match cd_elem {
119        Elem::F32 => &format!("{}", Elem::<D>::F32),
120        _ => &format!("{}", Elem::<D>::U32),
121    };
122
123    // Note: Scaled MMA actually requires float registers for C/D, unlike normal MMA
124    let args_a = (0..n_a_registers).map(|i| format!("{ab_arg_ty} const &reg_a_{i}"));
125    let args_b = (0..n_b_registers).map(|i| format!("{ab_arg_ty} const &reg_b_{i}"));
126    let args_c = (0..n_c_registers).map(|i| format!("{cd_arg_ty} const &reg_c_{i}"));
127    let args_d = (0..n_d_registers).map(|i| format!("{cd_arg_ty} &reg_d_{i}"));
128    let args = args_a
129        .chain(args_b)
130        .chain(args_c)
131        .chain(args_d)
132        .collect::<Vec<_>>()
133        .join(", ");
134
135    let mut idx = 0usize;
136
137    let placeholders_d = comma_separated((0..n_d_registers).map(|_| placeholder(&mut idx)));
138    let placeholders_d = format!("{{{placeholders_d}}}");
139
140    let placeholders_a = comma_separated((0..n_a_registers).map(|_| placeholder(&mut idx)));
141    let placeholders_a = format!("{{{placeholders_a}}}");
142
143    let placeholders_b = comma_separated((0..n_b_registers).map(|_| placeholder(&mut idx)));
144    let placeholders_b = format!("{{{placeholders_b}}}");
145
146    let placeholders_c = comma_separated((0..n_c_registers).map(|_| placeholder(&mut idx)));
147    let placeholders_c = format!("{{{placeholders_c}}}");
148
149    let placeholder_scales_a = format!(
150        "{{{}}}, {{{}, {}}}",
151        placeholder(&mut idx),
152        placeholder(&mut idx),
153        placeholder(&mut idx)
154    );
155    let placeholder_scales_b = format!(
156        "{{{}}}, {{{}, {}}}",
157        placeholder(&mut idx),
158        placeholder(&mut idx),
159        placeholder(&mut idx)
160    );
161
162    let params_out =
163        comma_separated((0..n_d_registers).map(|i| as_reg(&format!("reg_d_{i}"), cd_elem, true)));
164    let params_a = (0..n_a_registers).map(|i| as_reg(&format!("reg_a_{i}"), a_elem, false));
165    let params_b = (0..n_b_registers).map(|i| as_reg(&format!("reg_b_{i}"), b_elem, false));
166    let params_c = (0..n_c_registers).map(|i| as_reg(&format!("reg_c_{i}"), cd_elem, false));
167    let params_in = comma_separated(params_a.chain(params_b).chain(params_c));
168
169    format!(
170        r#"
171inline __device__ void
172__mma_scaled_{scales_factor}x_m16n8k{k}_{a_elem}_{b_elem}_{cd_elem}({args}, uint32 const &scales_a, uint32 const &scales_b) {{
173    static constexpr uint16 tidA = 0;
174    static constexpr uint16 bidA = 0;
175    static constexpr uint16 tidB = 0;
176    static constexpr uint16 bidB = 0;
177
178    asm volatile("mma.sync.aligned.kind::{kind}.block_scale.scale_vec::{scales_factor}X.m16n8k{k}.row.col.{cd_ty}.{a_ty}.{b_ty}.{cd_ty}.{s_ty} "
179               "{placeholders_d}, {placeholders_a}, {placeholders_b}, {placeholders_c}, {placeholder_scales_a}, {placeholder_scales_b};"
180               : {params_out}
181               : {params_in}, "r"(scales_a), "h"(bidA), "h"(tidA), "r"(scales_b), "h"(bidB), "h"(tidB));
182    }}
183    "#
184    )
185}
186
187pub(crate) fn comma_separated(it: impl IntoIterator<Item = String>) -> String {
188    it.into_iter().collect::<Vec<_>>().join(", ")
189}
190
191fn placeholder(idx: &mut usize) -> String {
192    let placeholder = format!("%{idx}");
193    *idx += 1;
194    placeholder
195}
196
197fn as_reg<D: Dialect>(ident: &str, ty: Elem<D>, output: bool) -> String {
198    let ty = match ty {
199        Elem::F32 => "f",
200        Elem::F64 => "d",
201        Elem::U64 => "l",
202        _ => "r",
203    };
204    if output {
205        format!(r#""={ty}"({ident})"#)
206    } else {
207        format!(r#""{ty}"({ident})"#)
208    }
209}
210
211fn mma_ty<D: Dialect>(elem: Elem<D>) -> &'static str {
212    match elem {
213        Elem::TF32 => "tf32",
214        Elem::F32 => "f32",
215        Elem::F64 => "f64",
216        Elem::F16 => "f16",
217        Elem::BF16 => "bf16",
218        Elem::FP4(FP4Kind::E2M1) => "e2m1",
219        // For packed MMA this will always exist as fp4x2, since 4-bit values can't exist
220        Elem::FP4x2(FP4Kind::E2M1) => "e2m1",
221        Elem::FP6(FP6Kind::E2M3) => "e2m3",
222        Elem::FP6(FP6Kind::E3M2) => "e3m2",
223        Elem::FP8(FP8Kind::E4M3) => "e4m3",
224        Elem::FP8(FP8Kind::E5M2) => "e5m2",
225        Elem::FP8(FP8Kind::UE8M0) => "ue8m0",
226        Elem::I8 => "s8",
227        Elem::I16 => "s16",
228        Elem::I32 => "s32",
229        Elem::I64 => "s64",
230        Elem::U8 => "u8",
231        Elem::U16 => "u16",
232        Elem::U32 => "u32",
233        Elem::U64 => "u64",
234        Elem::Bool => "b1",
235        other => panic!("{other} not supported for MMA"),
236    }
237}
238
239pub fn ldmatrix_call<D: Dialect>(
240    output: &Value<D>,
241    ptr: &Value<D>,
242    factor: &u32,
243    transpose: &bool,
244) -> String {
245    let elem = output.elem();
246    let width = 16 / output.elem().size();
247    let is_transposed = if *transpose { "_trans" } else { "" };
248    let regs = comma_separated(
249        (0..*factor as usize).map(|i| format!("reinterpret_cast<uint32&>({})", output.index(i))),
250    );
251    let ptr = ptr.fmt_ptr();
252
253    format!("__ldmatrix_m{width}n8_{elem}_{factor}x{is_transposed}({regs}, {ptr});\n")
254}
255
256pub fn ldmatrix_template<D: Dialect>(elem: Elem<D>, factor: u32, transpose: bool) -> String {
257    let width = 16 / elem.size();
258    let arg_ty = Elem::<D>::U32;
259
260    let args_regs = (0..factor).map(|i| format!("{arg_ty} &reg_{i}"));
261    let arg_addr = ["void const *row_addr".to_string()];
262    let args = args_regs.chain(arg_addr).collect::<Vec<_>>().join(", ");
263
264    let mut idx = 0usize;
265
266    let placeholders_regs = comma_separated((0..factor).map(|_| placeholder(&mut idx)));
267    let placeholders_regs = format!("{{{placeholders_regs}}}");
268
269    let placeholder_addr = format!("[{}]", placeholder(&mut idx));
270
271    let params_regs = comma_separated((0..factor).map(|i| format!(r#""=r"(reg_{i})"#)));
272    let param_addr = r#""r"(addr)"#;
273
274    let is_transposed = if transpose { "_trans" } else { "" };
275    let transposed_arg = if transpose { ".trans" } else { "" };
276    let num = format!("x{factor}");
277
278    let ty = match elem.size() {
279        2 => "b16",
280        1 => "b8",
281        _ => unreachable!(),
282    };
283
284    format!(
285        r#"
286inline __device__ void
287__ldmatrix_m{width}n8_{elem}_{factor}x{is_transposed}({args}) {{
288  uint32 addr = static_cast<uint32>(__cvta_generic_to_shared(row_addr));
289  asm volatile("ldmatrix.sync.aligned.m8n{width}.{num}{transposed_arg}.shared::cta.{ty}"
290               " {placeholders_regs}, {placeholder_addr};"
291               : {params_regs}
292               : {param_addr});
293    }}
294    "#
295    )
296}
297
298pub fn stmatrix_call<D: Dialect>(
299    registers: &Value<D>,
300    ptr: &Value<D>,
301    factor: &u32,
302    transpose: &bool,
303) -> String {
304    let elem = registers.elem();
305    let width = 16 / registers.elem().size();
306    let is_transposed = if *transpose { "_trans" } else { "" };
307    let regs = comma_separated(
308        (0..*factor as usize)
309            .map(|i| format!("reinterpret_cast<const uint32&>({})", registers.index(i))),
310    );
311    let ptr = ptr.fmt_ptr();
312
313    format!("__stmatrix_m{width}n8_{elem}_{factor}x{is_transposed}({regs}, {ptr});\n")
314}
315
316pub fn stmatrix_template<D: Dialect>(elem: Elem<D>, factor: u32, transpose: bool) -> String {
317    let width = 16 / elem.size();
318    let arg_ty = Elem::<D>::U32;
319
320    let args_regs = (0..factor).map(|i| format!("{arg_ty} const &reg_{i}"));
321    let arg_addr = ["void *row_addr".to_string()];
322    let args = args_regs.chain(arg_addr).collect::<Vec<_>>().join(", ");
323
324    let mut idx = 0usize;
325
326    let placeholder_addr = format!("[{}]", placeholder(&mut idx));
327
328    let placeholders_regs = comma_separated((0..factor).map(|_| placeholder(&mut idx)));
329    let placeholders_regs = format!("{{{placeholders_regs}}}");
330
331    let params_regs = comma_separated((0..factor).map(|i| format!(r#""r"(reg_{i})"#)));
332    let param_addr = r#""r"(addr)"#;
333
334    let is_transposed = if transpose { "_trans" } else { "" };
335    let transposed_arg = if transpose { ".trans" } else { "" };
336    let num = format!("x{factor}");
337
338    let ty = match elem.size() {
339        2 => "b16",
340        1 => "b8",
341        _ => unreachable!(),
342    };
343
344    // Note: smem technically an input
345    format!(
346        r#"
347inline __device__ void
348__stmatrix_m{width}n8_{elem}_{factor}x{is_transposed}({args}) {{
349  uint32 addr = static_cast<uint32>(__cvta_generic_to_shared(row_addr));
350  asm volatile("stmatrix.sync.aligned.m8n{width}.{num}{transposed_arg}.shared::cta.{ty}"
351               " {placeholder_addr}, {placeholders_regs};"
352               :: {param_addr}, {params_regs});
353    }}
354    "#
355    )
356}
cubecl_cpp/cuda/ptx/mma.rs

cubecl_cpp/cuda/ptx/
mma.rs
