cubecl_cpp/hip/mma/
mod.rs

1pub mod manual;
2pub mod rocwmma_compiler;
3pub mod wmma_intrinsics_compiler;
4
5use cubecl_core::{
6    cmma::MatrixType,
7    ir::{ContextExt, dialect::matrix::*},
8};
9use pliron::{
10    builtin::op_interfaces::OneResultInterface, context::Context, derive::type_interface_impl,
11    r#type::Typed,
12};
13
14use rocwmma_compiler::*;
15pub use wmma_intrinsics_compiler::*;
16
17use crate::{
18    hip::{
19        arch::{AMDArchitecture, AmdWmma},
20        hip_op,
21        ty::hip_ty,
22    },
23    shared::{DeclareMatrixOp, SupportedMmaCombinations, ty::TypeToCPP, wmma_api_base},
24    target::Hip,
25};
26
27const WMMA_NAMESPACE: &str = "rocwmma";
28
29#[derive(Clone, Copy, Debug)]
30pub enum HipCmmaCompiler {
31    RocWmma,
32    Intrinsics,
33}
34
35impl HipCmmaCompiler {
36    pub fn supported_cmma_combinations(&self, arch: &AMDArchitecture) -> SupportedMmaCombinations {
37        match self {
38            HipCmmaCompiler::RocWmma => supported_wmma_combinations_rocwmma(arch),
39            HipCmmaCompiler::Intrinsics => supported_wmma_combinations_intrinsic(arch),
40        }
41    }
42
43    pub fn imports(&self) -> String {
44        match self {
45            HipCmmaCompiler::RocWmma => compile_rocwmma_includes(),
46            HipCmmaCompiler::Intrinsics => String::new(),
47        }
48    }
49
50    // These used to be based on flags but extra type defs don't actually hurt anything so there's
51    // no point adding complexity
52    pub fn type_definitions(&self) -> String {
53        match self {
54            HipCmmaCompiler::RocWmma => String::new(),
55            HipCmmaCompiler::Intrinsics => r#"
56typedef __bf16 bhalf8_t __attribute__((ext_vector_type(8)));
57typedef __bf16 bhalf16_t __attribute__((ext_vector_type(16)));
58typedef _Float16 half8_t __attribute__((ext_vector_type(8)));
59typedef _Float16 half16_t __attribute__((ext_vector_type(16)));
60typedef float float8_t __attribute__((ext_vector_type(8)));
61        "#
62            .into(),
63        }
64    }
65}
66
67impl HipCmmaExt for Context {}
68pub trait HipCmmaExt: ContextExt {
69    fn hip_cmma(&self) -> HipCmmaCompiler {
70        *self.aux_ty::<HipCmmaCompiler>()
71    }
72    fn set_hip_cmma(&mut self, value: HipCmmaCompiler) {
73        self.set_aux_ty(value);
74    }
75}
76
77/// Which WMMA generation the intrinsic fragments and builtins should target.
78pub fn amd_wmma(ctx: &Context) -> AmdWmma {
79    ctx.aux_ty::<crate::shared::CompilationOptions>()
80        .amd_wmma
81        .expect("wmma should only be compiled for architectures that support it")
82}
83
84hip_ty!(MatrixType, |ty, ctx| match ctx.hip_cmma() {
85    HipCmmaCompiler::RocWmma => wmma_api_base::compile_matrix(ctx, ty, WMMA_NAMESPACE),
86    HipCmmaCompiler::Intrinsics => compile_fragment_intrinsic(ctx, ty),
87});
88
89hip_op!(DeclareMatrixOp, |op, ctx| {
90    wmma_api_base::compile_matrix_declaration(
91        ctx,
92        op.get_result(ctx),
93        op.value_ty(ctx).get_type(ctx),
94    )
95});
96
97hip_op!(FillOp, |op, ctx| match ctx.hip_cmma() {
98    HipCmmaCompiler::RocWmma => wmma_api_base::fill(ctx, op, WMMA_NAMESPACE),
99    HipCmmaCompiler::Intrinsics => compile_fill_intrinsic(ctx, op),
100});
101
102hip_op!(LoadOp, |op, ctx| match ctx.hip_cmma() {
103    HipCmmaCompiler::RocWmma => wmma_api_base::load(ctx, op, WMMA_NAMESPACE),
104    HipCmmaCompiler::Intrinsics => compile_load_intrinsic(ctx, op),
105});
106
107hip_op!(StoreOp, |op, ctx| match ctx.hip_cmma() {
108    HipCmmaCompiler::RocWmma => wmma_api_base::store(ctx, op, WMMA_NAMESPACE),
109    HipCmmaCompiler::Intrinsics => compile_store_intrinsic(ctx, op),
110});
111
112hip_op!(MultiplyAccumulateOp, |op, ctx| match ctx.hip_cmma() {
113    HipCmmaCompiler::RocWmma => wmma_api_base::execute(ctx, op, WMMA_NAMESPACE),
114    HipCmmaCompiler::Intrinsics => compile_execute_intrinsic(ctx, op),
115});
116
117hip_op!(CastOp, |op, ctx| match ctx.hip_cmma() {
118    HipCmmaCompiler::RocWmma => wmma_api_base::cast(ctx, op),
119    HipCmmaCompiler::Intrinsics => compile_cast_intrinsic(ctx, op),
120});
cubecl_cpp/hip/mma/mod.rs

cubecl_cpp/hip/mma/
mod.rs
