cubecl_cpp/cuda/mma/
cuda_compiler.rs

1use super::{WMMA_MINIMUM_VERSION, WMMA_NAMESPACE};
2use crate::{
3    cuda::{
4        CudaDialect,
5        arch::CudaArchitecture,
6        mma::{
7            compile_manual_mma, compile_scaled_mma, supported_mma_combinations,
8            supported_scaled_mma_combinations,
9        },
10    },
11    shared::{
12        Architecture, DialectWmmaCompiler, Flags, FragmentIdent, FragmentLayout, FragmentType,
13        ManualMma, SupportedMmaCombinations, SupportedScaledMmaCombinations, Value,
14        WmmaInstruction, wmma_api_base,
15    },
16};
17use cubecl_core::ir::{self as gpu, features::MmaConfig};
18use itertools::Itertools;
19
20#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
21pub struct CudaWmmaCompiler {}
22
23impl DialectWmmaCompiler<CudaDialect<Self>> for CudaWmmaCompiler {
24    fn compile_wmma_includes(
25        f: &mut std::fmt::Formatter<'_>,
26        _flags: &Flags<CudaDialect<Self>>,
27    ) -> std::fmt::Result {
28        f.write_str("#include <mma.h>\n")
29    }
30
31    fn compile_wmma_fragment_declaration(
32        f: &mut std::fmt::Formatter<'_>,
33        val: &crate::shared::Value<CudaDialect<Self>>,
34        ty: &crate::shared::Item<CudaDialect<Self>>,
35    ) -> std::fmt::Result {
36        wmma_api_base::compile_fragment_declaration(f, val, ty)
37    }
38
39    fn compile_wwma_fragment_ident(
40        f: &mut std::fmt::Formatter<'_>,
41        ident: &FragmentIdent<CudaDialect<Self>>,
42    ) -> std::fmt::Result {
43        wmma_api_base::compile_fragment_ident(f, WMMA_NAMESPACE, ident)
44    }
45
46    fn compile_wmma_fragment_layout(
47        f: &mut std::fmt::Formatter<'_>,
48        layout: &FragmentLayout<CudaDialect<Self>>,
49    ) -> std::fmt::Result {
50        wmma_api_base::compile_fragment_layout(f, WMMA_NAMESPACE, layout)
51    }
52
53    fn compile_wmma_fragment(
54        f: &mut std::fmt::Formatter<'_>,
55        fragment: &FragmentType<CudaDialect<Self>>,
56    ) -> std::fmt::Result {
57        wmma_api_base::compile_fragment(f, WMMA_NAMESPACE, fragment)
58    }
59
60    fn compile_wmma_instruction(
61        f: &mut std::fmt::Formatter<'_>,
62        instruction: &WmmaInstruction<CudaDialect<Self>>,
63    ) -> std::fmt::Result {
64        wmma_api_base::compile_instruction(f, WMMA_NAMESPACE, instruction)
65    }
66
67    fn compile_manual_mma(
68        f: &mut std::fmt::Formatter<'_>,
69        mma: ManualMma<CudaDialect<Self>>,
70    ) -> std::fmt::Result {
71        compile_manual_mma(f, mma)
72    }
73
74    fn compile_scaled_mma(
75        f: &mut std::fmt::Formatter<'_>,
76        mma: ManualMma<CudaDialect<Self>>,
77        scales_a: Value<CudaDialect<Self>>,
78        scales_b: Value<CudaDialect<Self>>,
79        scales_factor: u32,
80    ) -> std::fmt::Result {
81        compile_scaled_mma(f, mma, scales_a, scales_b, scales_factor)
82    }
83
84    fn supported_wmma_combinations(arch: &CudaArchitecture) -> SupportedMmaCombinations {
85        let mut result: SupportedMmaCombinations = vec![];
86        if arch.get_version() >= WMMA_MINIMUM_VERSION {
87            let tdims = vec![(16, 16, 16), (32, 8, 16), (8, 32, 16)];
88            // Types fully supported.
89            let types = vec![
90                (
91                    gpu::ElemType::Float(gpu::FloatKind::F16), // m
92                    gpu::ElemType::Float(gpu::FloatKind::F16), // n
93                    gpu::ElemType::Float(gpu::FloatKind::F16), // k
94                ),
95                (
96                    gpu::ElemType::Float(gpu::FloatKind::F16),
97                    gpu::ElemType::Float(gpu::FloatKind::F16),
98                    gpu::ElemType::Float(gpu::FloatKind::F32),
99                ),
100                (
101                    gpu::ElemType::Float(gpu::FloatKind::BF16),
102                    gpu::ElemType::Float(gpu::FloatKind::BF16),
103                    gpu::ElemType::Float(gpu::FloatKind::F32),
104                ),
105                (
106                    gpu::ElemType::Int(gpu::IntKind::I8),
107                    gpu::ElemType::Int(gpu::IntKind::I8),
108                    gpu::ElemType::Int(gpu::IntKind::I32),
109                ),
110                (
111                    gpu::ElemType::UInt(gpu::UIntKind::U8),
112                    gpu::ElemType::UInt(gpu::UIntKind::U8),
113                    gpu::ElemType::Int(gpu::IntKind::I32),
114                ),
115            ];
116            let combinations: SupportedMmaCombinations = types
117                .into_iter()
118                .cartesian_product(tdims)
119                .map(|((a, b, c), (m, n, k))| MmaConfig {
120                    a_type: a.into(),
121                    b_type: b.into(),
122                    cd_type: c.into(),
123                    m,
124                    n,
125                    k,
126                })
127                .collect();
128            result.extend(combinations);
129            if arch.get_version() >= 80 {
130                result.push(MmaConfig {
131                    a_type: gpu::ElemType::Float(gpu::FloatKind::TF32).into(),
132                    b_type: gpu::ElemType::Float(gpu::FloatKind::TF32).into(),
133                    cd_type: gpu::ElemType::Float(gpu::FloatKind::F32).into(),
134                    m: 16,
135                    n: 16,
136                    k: 8,
137                });
138            }
139        }
140        result
141    }
142
143    fn supported_mma_combinations(arch: &CudaArchitecture) -> SupportedMmaCombinations {
144        supported_mma_combinations(arch)
145    }
146
147    fn supported_scaled_mma_combinations(
148        arch: &CudaArchitecture,
149    ) -> SupportedScaledMmaCombinations {
150        supported_scaled_mma_combinations(arch)
151    }
152}
cubecl_cpp/cuda/mma/cuda_compiler.rs

cubecl_cpp/cuda/mma/
cuda_compiler.rs
