cubecl-cpp 0.11.0-pre.4

CPP transpiler for CubeCL
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109

use core::cell::Ref;

use cubecl_core::{
    cmma::{MatrixLayout, MatrixShape, MatrixType},
    ir::{
        dialect::matrix::{CastOp, FillOp, LoadOp, MultiplyAccumulateOp, StoreOp},
        interfaces::TypedExt,
        prelude::*,
    },
};

use crate::{
    metal::{metal_op, ty::metal_ty},
    shared::{
        CompilationOptions, CppValue, DeclareMatrixOp,
        ty::{TypeExtCPP, TypeToCPP},
        wmma_api_base::{self, as_scalar_ptr},
    },
    target::Metal,
};

pub fn compile_cmma_includes_metal() -> String {
    "#include <metal_simdgroup_matrix>\n".into()
}

metal_ty!(MatrixType, |ty, ctx| {
    let MatrixShape { m, n, k } = ty.shape;
    let ty = ty.elem_ty.to_cpp(ctx);
    // currently as of Metal 3.2 only fragments of 8x8x8 are supported
    if m != 8 || n != 8 || k != 8 {
        panic!("{m}x{n}x{k} fragments not supported. Only 8x8x8 fragments are supported.");
    }
    format!("simdgroup_{ty}8x8")
});

metal_op!(DeclareMatrixOp, |op, ctx| {
    wmma_api_base::compile_matrix_declaration(
        ctx,
        op.get_result(ctx),
        op.value_ty(ctx).get_type(ctx),
    )
});

metal_op!(FillOp, |op, ctx| {
    let mat = op.matrix(ctx).name(ctx);
    let value = op.value(ctx).name(ctx);
    let ty = op.value(ctx).get_type(ctx).to_cpp(ctx);

    format!("*{mat} = make_filled_simdgroup_matrix<{ty}, 8, 8>({value});\n",)
});

metal_op!(LoadOp, |op, ctx| {
    let frag = op.matrix(ctx).name(ctx);
    let ptr = as_scalar_ptr(ctx, op.source(ctx));
    let stride = op.stride(ctx).name(ctx);
    let mat_ty = matrix_ty(ctx, op.matrix(ctx));
    let transpose = match mat_ty.layout {
        MatrixLayout::RowMajor | MatrixLayout::Undefined => false,
        MatrixLayout::ColMajor => true,
    };
    format!("simdgroup_load(*{frag}, {ptr}, {stride}, 0, {transpose});\n")
});

metal_op!(StoreOp, |op, ctx| {
    let mat = op.matrix(ctx).name(ctx);
    let destination = as_scalar_ptr(ctx, op.destination(ctx));
    let stride = op.stride(ctx).name(ctx);
    format!(
        "
simdgroup_store(*{mat}, {destination}, {stride});
simdgroup_barrier(mem_flags::mem_threadgroup);"
    )
});

metal_op!(MultiplyAccumulateOp, |op, ctx| {
    let a = op.mat_a(ctx).name(ctx);
    let b = op.mat_b(ctx).name(ctx);
    let c = op.mat_c(ctx).name(ctx);
    let d = op.mat_d(ctx).name(ctx);
    format!("simdgroup_multiply_accumulate(*{d}, *{a}, *{b}, *{c});\n")
});

metal_op!(CastOp, |op, ctx| {
    let input = op.input(ctx).name(ctx);
    let output = op.output(ctx).name(ctx);
    let output_ty = matrix_ty(ctx, op.output(ctx));
    let ty = output_ty.elem_ty.to_cpp(ctx);
    let threads_per_simdgroup = ctx.aux_ty::<CompilationOptions>().warp_size;
    let elements_held_by_each_thread =
        elements_held_by_each_thread(&output_ty.shape, threads_per_simdgroup);
    format!(
        "
simdgroup_barrier(mem_flags::mem_none);
for(int e=0; e<{elements_held_by_each_thread}; e++) {{
    {output}->thread_elements()[e] = {ty}({input}->thread_elements()[e]);
}}"
    )
});

fn elements_held_by_each_thread(shape: &MatrixShape, threads_per_simdgroup: usize) -> usize {
    (shape.m * shape.n) / threads_per_simdgroup
}

fn matrix_ty(ctx: &Context, ty: impl Typed) -> Ref<'_, MatrixType> {
    let ty = ty.unwrap_ptr(ctx).deref(ctx);
    Ref::map(ty, |ty| {
        ty.downcast_ref::<MatrixType>().expect("Should be matrix")
    })
}