cubecl_cpp/cuda/mma/
mod.rs

1use cubecl_core::{
2    cmma::MatrixType,
3    ir::{
4        ContextExt,
5        dialect::matrix::{CastOp, FillOp, LoadOp, MultiplyAccumulateOp, StoreOp},
6        prelude::*,
7    },
8};
9
10use crate::{
11    cuda::{arch::CudaArchitecture, cuda_op, ty::cuda_ty},
12    shared::{DeclareMatrixOp, SupportedMmaCombinations, wmma_api_base},
13};
14
15pub mod cuda_compiler;
16pub mod manual;
17pub mod ptx_wmma_compiler;
18
19use cuda_compiler::*;
20use ptx_wmma_compiler::*;
21
22const WMMA_NAMESPACE: &str = "nvcuda::wmma";
23const WMMA_MINIMUM_VERSION: u32 = 70;
24
25// Maybe this should be a separate dialect but lets keep it simple for now
26
27#[derive(Clone, Copy, Debug)]
28pub enum CudaCmmaCompiler {
29    Cpp,
30    Ptx,
31}
32
33impl CudaCmmaCompiler {
34    pub fn supported_cmma_combinations(&self, arch: &CudaArchitecture) -> SupportedMmaCombinations {
35        match self {
36            CudaCmmaCompiler::Cpp => supported_cmma_combinations_wmma(arch),
37            CudaCmmaCompiler::Ptx => supported_cmma_combinations_ptx(arch),
38        }
39    }
40
41    pub fn imports(&self) -> String {
42        "#include <mma.h>\n".into()
43    }
44}
45
46impl CudaCmmaExt for Context {}
47pub trait CudaCmmaExt: ContextExt {
48    fn cuda_cmma(&self) -> CudaCmmaCompiler {
49        *self.aux_ty::<CudaCmmaCompiler>()
50    }
51    fn set_cuda_cmma(&mut self, value: CudaCmmaCompiler) {
52        self.set_aux_ty(value);
53    }
54}
55
56cuda_ty!(MatrixType, |ty, ctx| match ctx.cuda_cmma() {
57    CudaCmmaCompiler::Cpp => wmma_api_base::compile_matrix(ctx, ty, WMMA_NAMESPACE),
58    CudaCmmaCompiler::Ptx => compile_matrix_ptx(ctx, ty),
59});
60
61cuda_op!(DeclareMatrixOp, |op, ctx| match ctx.cuda_cmma() {
62    CudaCmmaCompiler::Cpp => wmma_api_base::compile_matrix_declaration(
63        ctx,
64        op.get_result(ctx),
65        op.value_ty(ctx).get_type(ctx),
66    ),
67    CudaCmmaCompiler::Ptx =>
68        compile_matrix_declaration_ptx(ctx, op.get_result(ctx), op.value_ty(ctx).get_type(ctx)),
69});
70
71cuda_op!(FillOp, |op, ctx| match ctx.cuda_cmma() {
72    CudaCmmaCompiler::Cpp => wmma_api_base::fill(ctx, op, WMMA_NAMESPACE),
73    CudaCmmaCompiler::Ptx => fill_ptx(ctx, op),
74});
75
76cuda_op!(LoadOp, |op, ctx| match ctx.cuda_cmma() {
77    CudaCmmaCompiler::Cpp => wmma_api_base::load(ctx, op, WMMA_NAMESPACE),
78    CudaCmmaCompiler::Ptx => load_ptx(ctx, op),
79});
80
81cuda_op!(StoreOp, |op, ctx| match ctx.cuda_cmma() {
82    CudaCmmaCompiler::Cpp => wmma_api_base::store(ctx, op, WMMA_NAMESPACE),
83    CudaCmmaCompiler::Ptx => store_ptx(ctx, op),
84});
85
86cuda_op!(MultiplyAccumulateOp, |op, ctx| match ctx.cuda_cmma() {
87    CudaCmmaCompiler::Cpp => wmma_api_base::execute(ctx, op, WMMA_NAMESPACE),
88    CudaCmmaCompiler::Ptx => execute_ptx(ctx, op),
89});
90
91cuda_op!(CastOp, |op, ctx| match ctx.cuda_cmma() {
92    CudaCmmaCompiler::Cpp => wmma_api_base::cast(ctx, op),
93    CudaCmmaCompiler::Ptx => cast_ptx(ctx, op),
94});
cubecl_cpp/cuda/mma/mod.rs

cubecl_cpp/cuda/mma/
mod.rs
