cubecl_std/throughput/runners/
compute_cmma.rs

1use cubecl::prelude::*;
2use cubecl_core as cubecl;
3use cubecl_runtime::throughput::{CmmaDims, ComputeCmmaConfig, KernelConfig, ThroughputKey};
4
5use crate::throughput::LaunchConfig;
6
7pub fn build_kernel<R: Runtime>(
8    client: &ComputeClient<R>,
9    key: ThroughputKey,
10    cmma_config: ComputeCmmaConfig,
11    config: LaunchConfig,
12) -> KernelConfig {
13    let client = client.clone();
14    let dtype = key.dtype;
15
16    let ops_per_cmma = 2 * cmma_config.cmma_dims.num_elems();
17    let out_bytes =
18        cmma_config.cmma_dims.m * cmma_config.cmma_dims.n * cmma_config.accumulator_type.size();
19
20    let kernel = Box::new(move |iterations: usize| unsafe {
21        let out = client.empty(out_bytes);
22
23        compute_cmma_throughput::launch_unchecked(
24            &client,
25            CubeCount::Static(config.cube_count as u32, 1, 1),
26            CubeDim::new(&client, config.cube_dim),
27            config.vector_size,
28            BufferArg::from_raw_parts(out, 1),
29            iterations,
30            cmma_config.cmma_dims,
31            dtype.into(),
32            cmma_config.accumulator_type.into(),
33        )
34    });
35
36    let planes_per_cube = config.cube_dim / config.plane_size;
37    let ops_count = config.cube_count * planes_per_cube * ops_per_cmma;
38
39    KernelConfig { kernel, ops_count }
40}
41
42#[cube(launch_unchecked)]
43pub fn compute_cmma_throughput<I: Numeric, ACC: Numeric, N: Size>(
44    output: &mut [Vector<ACC, N>],
45    n_iter: usize,
46    #[comptime] cmm_dims: CmmaDims,
47    #[define(I)] _dtype: StorageType,
48    #[define(ACC)] _acc: StorageType,
49) {
50    let CmmaDims { m, n, k } = cmm_dims;
51
52    let a = cmma::Matrix::<I>::from_value(
53        cmma::MatrixIdent::A,
54        m,
55        n,
56        k,
57        cmma::MatrixLayout::RowMajor,
58        I::cast_from(1),
59    );
60
61    let b = cmma::Matrix::<I>::from_value(
62        cmma::MatrixIdent::B,
63        m,
64        n,
65        k,
66        cmma::MatrixLayout::ColMajor,
67        I::cast_from(1),
68    );
69
70    let acc = cmma::Matrix::<ACC>::from_value(
71        cmma::MatrixIdent::Accumulator,
72        m,
73        n,
74        k,
75        cmma::MatrixLayout::Undefined,
76        ACC::cast_from(0.0),
77    );
78
79    for _ in 0..n_iter {
80        cmma::execute(&a, &b, &acc, &acc);
81    }
82
83    if ABSOLUTE_POS == 0 {
84        cmma::store(output, &acc, n as u32, cmma::MatrixLayout::RowMajor);
85    }
86}
cubecl_std/throughput/runners/compute_cmma.rs

cubecl_std/throughput/runners/
compute_cmma.rs
