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

cubecl_std/throughput/runners/
compute_cmma.rs
