cubecl_cpp/shared/
mma.rs

1use cubecl_core::{
2    cmma::{MatrixIdent, MatrixLayout, MatrixShape},
3    ir::{
4        DeviceProperties,
5        features::{MmaConfig, ScaledMmaConfig},
6    },
7};
8
9pub type SupportedMmaCombinations = Vec<MmaConfig>;
10pub type SupportedScaledMmaCombinations = Vec<ScaledMmaConfig>;
11
12pub trait Architecture {
13    fn warp_size(&self) -> u32;
14    fn is_wmma_capable(&self) -> bool;
15    fn is_mfma_capable(&self) -> bool;
16    fn get_version(&self) -> u32 {
17        0
18    }
19}
20
21pub fn register_wmma_features(
22    supported_combinations: SupportedMmaCombinations,
23    properties: &mut DeviceProperties,
24) {
25    for config in supported_combinations {
26        properties.features.matmul.cmma.insert(config);
27    }
28}
29
30pub fn register_mma_features(
31    supported_combinations: SupportedMmaCombinations,
32    properties: &mut DeviceProperties,
33) {
34    for config in supported_combinations {
35        properties.features.matmul.mma.insert(config);
36    }
37}
38
39pub fn register_scaled_mma_features(
40    supported_combinations: SupportedScaledMmaCombinations,
41    properties: &mut DeviceProperties,
42) {
43    for config in supported_combinations {
44        properties.features.matmul.scaled_mma.insert(config);
45    }
46}
47
48pub mod wmma_api_base {
49    use cubecl_core::{
50        cmma::{MatrixIdent, MatrixLayout, MatrixType},
51        ir::{
52            dialect::matrix::{CastOp, FillOp, LoadOp, MultiplyAccumulateOp, StoreOp},
53            interfaces::{TypeExt, TypedExt},
54            types::{PointerType, scalar::TFloat32Type},
55        },
56    };
57    use pliron::{
58        context::Context,
59        r#type::{TypeHandle, Typed},
60        value::Value,
61    };
62
63    use crate::shared::{CppValue, ty::TypeExtCPP};
64
65    use super::*;
66
67    pub fn compile_matrix_declaration(ctx: &Context, val: Value, value_ty: TypeHandle) -> String {
68        format!(
69            "{} {id}_store; {} {id} = &{id}_store;",
70            value_ty.to_cpp(ctx),
71            val.get_type(ctx).to_cpp(ctx),
72            id = val.name(ctx),
73        )
74    }
75
76    pub fn compile_matrix(ctx: &Context, ty: &MatrixType, ns: &str) -> String {
77        let elem = match ty.elem_ty.deref(ctx).is::<TFloat32Type>() {
78            true => format!("{ns}::precision::tf32"),
79            false => ty.elem_ty.to_cpp(ctx),
80        };
81        let ident = match ty.ident {
82            MatrixIdent::A => format!("{ns}::matrix_a"),
83            MatrixIdent::B => format!("{ns}::matrix_b"),
84            MatrixIdent::Accumulator => format!("{ns}::accumulator"),
85        };
86        let MatrixShape { m, n, k } = ty.shape;
87        // The layout-free fragment specialization exists for the accumulator and for nothing else.
88        let layout = match (ty.ident, ty.layout) {
89            (MatrixIdent::Accumulator, _) => {
90                return format!("{ns}::fragment<{ident}, {m}, {n}, {k}, {elem}>");
91            }
92            (_, MatrixLayout::ColMajor) => format!("{ns}::col_major"),
93            (_, MatrixLayout::RowMajor) => format!("{ns}::row_major"),
94            (_, MatrixLayout::Undefined) => {
95                panic!("An A or B fragment names the layout of the data it is loaded from.")
96            }
97        };
98        format!("{ns}::fragment<{ident}, {m}, {n}, {k}, {elem}, {layout}>")
99    }
100
101    pub fn fill(ctx: &Context, op: &FillOp, namespace: &str) -> String {
102        let mat = op.matrix(ctx).name(ctx);
103        let value = op.value(ctx).name(ctx);
104        format!("{namespace}::fill_fragment(*{mat}, {value});")
105    }
106
107    pub fn load(ctx: &Context, op: &LoadOp, namespace: &str) -> String {
108        let mat = op.matrix(ctx).name(ctx);
109        let stride = op.stride(ctx).name(ctx);
110        let ptr = as_scalar_ptr(ctx, op.source(ctx));
111        let mat_ty = matrix_ty(ctx, op.matrix(ctx));
112        // CUDA is annoying and doesn't allow layout on A/B even though the PTX equivalent takes one
113        let layout = match mat_ty.ident {
114            MatrixIdent::A | MatrixIdent::B => String::new(),
115            MatrixIdent::Accumulator => match op.layout(ctx).0 {
116                MatrixLayout::RowMajor => format!(", {namespace}::mem_row_major"),
117                MatrixLayout::ColMajor => format!(", {namespace}::mem_col_major"),
118                _ => String::new(),
119            },
120        };
121        format!("{namespace}::load_matrix_sync(*{mat}, {ptr}, {stride}{layout});")
122    }
123
124    pub fn store(ctx: &Context, op: &StoreOp, namespace: &str) -> String {
125        let mat = op.matrix(ctx).name(ctx);
126        let stride = op.stride(ctx).name(ctx);
127        let destination = as_scalar_ptr(ctx, op.destination(ctx));
128        let mat_ty = matrix_ty(ctx, op.matrix(ctx));
129        // CUDA is annoying and doesn't allow layout on A/B even though the PTX equivalent takes one
130        let layout = match mat_ty.ident {
131            MatrixIdent::A | MatrixIdent::B => String::new(),
132            MatrixIdent::Accumulator => match op.layout(ctx).0 {
133                MatrixLayout::RowMajor => format!(", {namespace}::mem_row_major"),
134                MatrixLayout::ColMajor => format!(", {namespace}::mem_col_major"),
135                _ => String::new(),
136            },
137        };
138
139        format!("{namespace}::store_matrix_sync({destination}, *{mat}, {stride}{layout});")
140    }
141
142    pub fn execute(ctx: &Context, op: &MultiplyAccumulateOp, namespace: &str) -> String {
143        let mat_a = op.mat_a(ctx).name(ctx);
144        let mat_b = op.mat_b(ctx).name(ctx);
145        let mat_c = op.mat_c(ctx).name(ctx);
146        let mat_d = op.mat_d(ctx).name(ctx);
147
148        format!("{namespace}::mma_sync(*{mat_d}, *{mat_a}, *{mat_b}, *{mat_c});")
149    }
150
151    pub fn cast(ctx: &Context, op: &CastOp) -> String {
152        let input = op.input(ctx).name(ctx);
153        let output = op.output(ctx).name(ctx);
154        let mat_ty = matrix_ty(ctx, op.output(ctx));
155        let out_elem = mat_ty.elem_ty.to_cpp(ctx);
156        format!(
157            "for(int t=0; t<{input}->num_elements; t++) {{ {output}->x[t] = {out_elem}({input}->x[t]); }}"
158        )
159    }
160
161    pub fn as_scalar_ptr(ctx: &Context, value: Value) -> String {
162        let PointerType {
163            inner,
164            address_space,
165        } = value.get_type(ctx).as_ptr(ctx);
166        let new_ty = PointerType::get(ctx, inner.scalar_ty(ctx), address_space).to_handle();
167        format!(
168            "reinterpret_cast<{}>({})",
169            new_ty.to_cpp(ctx),
170            value.name(ctx)
171        )
172    }
173
174    pub fn matrix_ty(ctx: &Context, value: impl Typed) -> MatrixType {
175        let ty = value.unwrap_ptr(ctx).deref(ctx);
176        *ty.downcast_ref::<MatrixType>().unwrap()
177    }
178}
179
180pub fn frag_ident_str(frag: &MatrixIdent) -> &str {
181    match frag {
182        MatrixIdent::A => "a",
183        MatrixIdent::B => "b",
184        MatrixIdent::Accumulator => "c",
185    }
186}
187
188pub fn frag_layout_str(frag: &MatrixLayout) -> &str {
189    match frag {
190        MatrixLayout::ColMajor => "col",
191        MatrixLayout::RowMajor => "row",
192        MatrixLayout::Undefined => "",
193    }
194}
cubecl_cpp/shared/mma.rs

cubecl_cpp/shared/
mma.rs
