cubecl_cpp/metal/
mma.rs

1use core::cell::Ref;
2
3use cubecl_core::{
4    cmma::{MatrixLayout, MatrixShape, MatrixType},
5    ir::{
6        dialect::matrix::{CastOp, FillOp, LoadOp, MultiplyAccumulateOp, StoreOp},
7        interfaces::TypedExt,
8        prelude::*,
9    },
10};
11
12use crate::{
13    metal::{metal_op, ty::metal_ty},
14    shared::{
15        CompilationOptions, CppValue, DeclareMatrixOp,
16        ty::{TypeExtCPP, TypeToCPP},
17        wmma_api_base::{self, as_scalar_ptr},
18    },
19    target::Metal,
20};
21
22pub fn compile_cmma_includes_metal() -> String {
23    "#include <metal_simdgroup_matrix>\n".into()
24}
25
26metal_ty!(MatrixType, |ty, ctx| {
27    let MatrixShape { m, n, k } = ty.shape;
28    let ty = ty.elem_ty.to_cpp(ctx);
29    // currently as of Metal 3.2 only fragments of 8x8x8 are supported
30    if m != 8 || n != 8 || k != 8 {
31        panic!("{m}x{n}x{k} fragments not supported. Only 8x8x8 fragments are supported.");
32    }
33    format!("simdgroup_{ty}8x8")
34});
35
36metal_op!(DeclareMatrixOp, |op, ctx| {
37    wmma_api_base::compile_matrix_declaration(
38        ctx,
39        op.get_result(ctx),
40        op.value_ty(ctx).get_type(ctx),
41    )
42});
43
44metal_op!(FillOp, |op, ctx| {
45    let mat = op.matrix(ctx).name(ctx);
46    let value = op.value(ctx).name(ctx);
47    let ty = op.value(ctx).get_type(ctx).to_cpp(ctx);
48
49    format!("*{mat} = make_filled_simdgroup_matrix<{ty}, 8, 8>({value});\n",)
50});
51
52metal_op!(LoadOp, |op, ctx| {
53    let frag = op.matrix(ctx).name(ctx);
54    let ptr = as_scalar_ptr(ctx, op.source(ctx));
55    let stride = op.stride(ctx).name(ctx);
56    let mat_ty = matrix_ty(ctx, op.matrix(ctx));
57    let transpose = match mat_ty.layout {
58        MatrixLayout::RowMajor | MatrixLayout::Undefined => false,
59        MatrixLayout::ColMajor => true,
60    };
61    format!("simdgroup_load(*{frag}, {ptr}, {stride}, 0, {transpose});\n")
62});
63
64metal_op!(StoreOp, |op, ctx| {
65    let mat = op.matrix(ctx).name(ctx);
66    let destination = as_scalar_ptr(ctx, op.destination(ctx));
67    let stride = op.stride(ctx).name(ctx);
68    format!(
69        "
70simdgroup_store(*{mat}, {destination}, {stride});
71simdgroup_barrier(mem_flags::mem_threadgroup);"
72    )
73});
74
75metal_op!(MultiplyAccumulateOp, |op, ctx| {
76    let a = op.mat_a(ctx).name(ctx);
77    let b = op.mat_b(ctx).name(ctx);
78    let c = op.mat_c(ctx).name(ctx);
79    let d = op.mat_d(ctx).name(ctx);
80    format!("simdgroup_multiply_accumulate(*{d}, *{a}, *{b}, *{c});\n")
81});
82
83metal_op!(CastOp, |op, ctx| {
84    let input = op.input(ctx).name(ctx);
85    let output = op.output(ctx).name(ctx);
86    let output_ty = matrix_ty(ctx, op.output(ctx));
87    let ty = output_ty.elem_ty.to_cpp(ctx);
88    let threads_per_simdgroup = ctx.aux_ty::<CompilationOptions>().warp_size;
89    let elements_held_by_each_thread =
90        elements_held_by_each_thread(&output_ty.shape, threads_per_simdgroup);
91    format!(
92        "
93simdgroup_barrier(mem_flags::mem_none);
94for(int e=0; e<{elements_held_by_each_thread}; e++) {{
95    {output}->thread_elements()[e] = {ty}({input}->thread_elements()[e]);
96}}"
97    )
98});
99
100fn elements_held_by_each_thread(shape: &MatrixShape, threads_per_simdgroup: usize) -> usize {
101    (shape.m * shape.n) / threads_per_simdgroup
102}
103
104fn matrix_ty(ctx: &Context, ty: impl Typed) -> Ref<'_, MatrixType> {
105    let ty = ty.unwrap_ptr(ctx).deref(ctx);
106    Ref::map(ty, |ty| {
107        ty.downcast_ref::<MatrixType>().expect("Should be matrix")
108    })
109}
cubecl_cpp/metal/mma.rs

cubecl_cpp/metal/
mma.rs
