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 sample = Box::new(move |iterations: usize| {
21        let start = cubecl_common::profile::Instant::now();
22        unsafe {
23            let out = client.empty(out_bytes);
24
25            compute_cmma_throughput::launch_unchecked(
26                &client,
27                CubeCount::Static(config.cube_count as u32, 1, 1),
28                CubeDim::new(&client, config.cube_dim),
29                config.vector_size,
30                BufferArg::from_raw_parts(out, 1),
31                iterations,
32                cmma_config.cmma_dims,
33                dtype.into(),
34                cmma_config.accumulator_type.into(),
35            )
36        };
37        let _ = cubecl_core::future::block_on(client.sync());
38        start.elapsed()
39    });
40
41    let planes_per_cube = config.cube_dim / config.plane_size;
42    let ops_count = config.cube_count * planes_per_cube * ops_per_cmma;
43
44    KernelConfig { sample, ops_count }
45}
46
47#[cube(launch_unchecked)]
48pub fn compute_cmma_throughput<I: Numeric, ACC: Numeric, N: Size>(
49    output: &mut [Vector<ACC, N>],
50    n_iter: usize,
51    #[comptime] cmm_dims: CmmaDims,
52    #[define(I)] _dtype: StorageType,
53    #[define(ACC)] _acc: StorageType,
54) {
55    let CmmaDims { m, n, k } = cmm_dims;
56
57    let a = cmma::Matrix::<I>::from_value(
58        cmma::MatrixIdent::A,
59        m,
60        n,
61        k,
62        cmma::MatrixLayout::RowMajor,
63        I::cast_from(1),
64    );
65
66    let b = cmma::Matrix::<I>::from_value(
67        cmma::MatrixIdent::B,
68        m,
69        n,
70        k,
71        cmma::MatrixLayout::ColMajor,
72        I::cast_from(1),
73    );
74
75    let acc = cmma::Matrix::<ACC>::from_value(
76        cmma::MatrixIdent::Accumulator,
77        m,
78        n,
79        k,
80        cmma::MatrixLayout::Undefined,
81        ACC::cast_from(0.0),
82    );
83
84    for _ in 0..n_iter {
85        cmma::execute(&a, &b, &acc, &acc);
86    }
87
88    if ABSOLUTE_POS == 0 {
89        cmma::store(output, &acc, n as u32, cmma::MatrixLayout::RowMajor);
90    }
91}
cubecl_std/throughput/runners/compute_cmma.rs

cubecl_std/throughput/runners/
compute_cmma.rs
