Skip to main content

cubecl_runtime/throughput/
cmma.rs

1use cubecl_ir::ElemType;
2
3/// Configuration for a matrix multiplication (CMMA) operation.
4#[derive(Eq, PartialEq, Clone, Hash, Debug, Copy)]
5#[cfg_attr(std_io, derive(serde::Serialize, serde::Deserialize))]
6pub struct ComputeCmmaConfig {
7    /// The data type used to store the running sum.
8    pub accumulator_type: AccumulatorType,
9    /// The spatial dimensions of the operation.
10    pub cmma_dims: CmmaDims,
11}
12
13/// The element type of the accumulator.
14pub type AccumulatorType = ElemType;
15
16/// The M, N, and K dimensions of a matrix multiplication.
17#[derive(Eq, PartialEq, Clone, Hash, Debug, Copy)]
18#[cfg_attr(std_io, derive(serde::Serialize, serde::Deserialize))]
19pub struct CmmaDims {
20    /// Rows in the output matrix.
21    pub m: usize,
22    /// Columns in the output matrix.
23    pub n: usize,
24    /// The shared inner dimension of the input matrices.
25    pub k: usize,
26}
27
28impl CmmaDims {
29    /// Returns the total iteration volume (M * N * K).
30    pub fn num_elems(&self) -> usize {
31        self.m * self.n * self.k
32    }
33}