t4a_cubecl_std/throughput/runners/
compute_cmma.rs

1use cubecl::prelude::*;
2use cubecl_core as cubecl;
3use cubecl_runtime::throughput::{
4    CmmaDims, ComputeCmmaConfig, KernelConfig, ThroughputError, ThroughputKey,
5};
6
7use crate::throughput::{LaunchConfig, verify::verify};
8
9pub fn build_kernel(
10    client: &Client,
11    key: ThroughputKey,
12    cmma_config: ComputeCmmaConfig,
13    config: LaunchConfig,
14) -> Result<KernelConfig, ThroughputError> {
15    let client = client.clone();
16    let dtype = key.dtype();
17
18    let ops_per_cmma = 2 * cmma_config.cmma_dims.num_elems();
19    let out_bytes =
20        cmma_config.cmma_dims.m * cmma_config.cmma_dims.n * cmma_config.accumulator_type.size();
21
22    let out = client.empty(out_bytes);
23
24    let (verifier, written) = (client.clone(), out.clone());
25    let sample = Box::new(move |iterations: usize| {
26        let start = cubecl_common::profile::Instant::now();
27        unsafe {
28            compute_cmma_throughput::launch_unchecked(
29                &client,
30                CubeCount::Static(config.cube_count as u32, 1, 1),
31                config.cube_dim,
32                config.vector_size,
33                BufferArg::from_raw_parts(out.clone(), 1),
34                iterations,
35                cmma_config.cmma_dims,
36                dtype,
37                cmma_config.accumulator_type,
38            )
39        };
40        // A failure is not this sync's to report: `verify` asked the output.
41        let _ = cubecl_core::future::block_on(client.sync());
42        start.elapsed()
43    });
44    verify(&verifier, &sample, &written)?;
45
46    let planes_per_cube = config.cube_dim.num_elems() as usize / config.plane_size;
47    let ops_count = config.cube_count * planes_per_cube * ops_per_cmma;
48
49    Ok(KernelConfig {
50        sample,
51        ops_count,
52        min_iterations: 1,
53    })
54}
55
56#[cube(launch_unchecked)]
57pub fn compute_cmma_throughput<I: Numeric, ACC: Numeric, N: Size>(
58    output: &mut [Vector<ACC, N>],
59    n_iter: usize,
60    #[comptime] cmm_dims: CmmaDims,
61    #[define(I)] _dtype: ElemType,
62    #[define(ACC)] _acc: ElemType,
63) {
64    let CmmaDims { m, n, k } = cmm_dims;
65
66    let a = cmma::Matrix::<I>::from_value(
67        cmma::MatrixIdent::A,
68        m,
69        n,
70        k,
71        cmma::MatrixLayout::RowMajor,
72        I::cast_from(1),
73    );
74
75    let b = cmma::Matrix::<I>::from_value(
76        cmma::MatrixIdent::B,
77        m,
78        n,
79        k,
80        cmma::MatrixLayout::ColMajor,
81        I::cast_from(1),
82    );
83
84    let acc = cmma::Matrix::<ACC>::from_value(
85        cmma::MatrixIdent::Accumulator,
86        m,
87        n,
88        k,
89        cmma::MatrixLayout::Undefined,
90        ACC::cast_from(0.0),
91    );
92
93    for _ in 0..n_iter {
94        cmma::execute(&a, &b, &acc, &acc);
95    }
96
97    if ABSOLUTE_POS == 0 {
98        cmma::store(output, &acc, n as u32, cmma::MatrixLayout::RowMajor);
99    }
100}
t4a_cubecl_std/throughput/runners/compute_cmma.rs

t4a_cubecl_std/throughput/runners/
compute_cmma.rs
