cubecl_runtime/throughput/
cmma.rs

1use crate::{client::ComputeClient, runtime::Runtime};
2use cubecl_ir::{ElemType, StorageType};
3
4/// Configuration for a matrix multiplication (CMMA) operation.
5#[derive(Eq, PartialEq, Clone, Hash, Debug, Copy)]
6#[cfg_attr(std_io, derive(serde::Serialize, serde::Deserialize))]
7pub struct ComputeCmmaConfig {
8    /// The data type used to store the running sum.
9    pub accumulator_type: AccumulatorType,
10    /// The spatial dimensions of the operation.
11    pub cmma_dims: CmmaDims,
12}
13
14/// The element type of the accumulator.
15pub type AccumulatorType = ElemType;
16
17/// The M, N, and K dimensions of a matrix multiplication.
18#[derive(Eq, PartialEq, Clone, Hash, Debug, Copy)]
19#[cfg_attr(std_io, derive(serde::Serialize, serde::Deserialize))]
20pub struct CmmaDims {
21    /// Rows in the output matrix.
22    pub m: usize,
23    /// Columns in the output matrix.
24    pub n: usize,
25    /// The shared inner dimension of the input matrices.
26    pub k: usize,
27}
28
29impl CmmaDims {
30    /// Returns the total iteration volume (M * N * K).
31    pub fn num_elems(&self) -> usize {
32        self.m * self.n * self.k
33    }
34}
35
36/// Resolves the largest supported CMMA or MMA tile size `(m, n, k)`.
37pub fn select_cmma_tile<R: Runtime>(
38    client: &ComputeClient<R>,
39    lhs: StorageType,
40    rhs: StorageType,
41    acc: StorageType,
42    (m, n, k): (usize, usize, usize),
43) -> Option<(u32, u32, u32)> {
44    let props = client.properties();
45
46    props
47        .features
48        .matmul
49        .cmma
50        .iter()
51        // Combine both CMMA and MMA supported hardware features.
52        .chain(props.features.matmul.mma.iter())
53        // Filter for instructions matching the exact input/accumulator data types.
54        .filter(|it| it.a_type == lhs && it.b_type == rhs && it.cd_type == acc)
55        // Ensure the hardware tile size actually fits within our problem dimensions.
56        .filter(|it| m >= it.m as usize && n >= it.n as usize && k >= it.k as usize)
57        // Select the tile with the largest volume to maximize throughput.
58        .max_by_key(|it| it.m as u64 * it.n as u64 * it.k as u64)
59        .map(|it| (it.m, it.n, it.k))
60}
cubecl_runtime/throughput/cmma.rs

cubecl_runtime/throughput/
cmma.rs
