cubecl_std/throughput/
base.rs

1use cubecl_core::{
2    CubeCount,
3    frontend::BufferArg,
4    ir::{ElemType, IntKind},
5};
6use cubecl_runtime::{
7    client::ComputeClient,
8    runtime::Runtime,
9    server::CubeDim,
10    throughput::{ThroughputKey, ThroughputMode, ThroughputValue},
11};
12
13use crate::throughput::{compute_cmma, compute_direct, launch_overhead, memory_direct};
14
15/// Computes the peak throughput for a given runtime and key.
16///
17/// Native only, panics on WASM
18pub fn measure_peak_throughput<R: Runtime>(
19    client: &ComputeClient<R>,
20    key: ThroughputKey,
21) -> ThroughputValue {
22    let launch_config = launch_config(client, key.dtype);
23
24    let kernel_config = match key.mode {
25        ThroughputMode::ComputeDirect => compute_direct::build_kernel(client, key, launch_config),
26        ThroughputMode::ComputeCmma(cmma_config) => {
27            if client.properties().features.matmul.cmma.is_empty() {
28                return ThroughputValue::ZERO;
29            }
30            compute_cmma::build_kernel(client, key, cmma_config, launch_config)
31        }
32        ThroughputMode::Memory => memory_direct::build_kernel(client, key, launch_config),
33    };
34
35    client.measure_throughput(key, kernel_config)
36}
37
38/// Hardware execution parameters for launching a compute kernel.
39#[derive(Clone, Copy)]
40pub struct LaunchConfig {
41    /// The number of threads per cube.
42    pub cube_dim: usize,
43    /// The total number of cubes to dispatch.
44    pub cube_count: usize,
45    /// The vectorization factor (e.g., 4 for `vec4` operations).
46    pub vector_size: usize,
47    /// The number of threads in a hardware execution plane.
48    pub plane_size: usize,
49}
50
51fn launch_config<R: Runtime>(client: &ComputeClient<R>, dtype: ElemType) -> LaunchConfig {
52    let hardware = &client.properties().hardware;
53
54    let plane_size = hardware.plane_size_max.max(1);
55    let requested = (hardware.max_units_per_cube / plane_size * plane_size)
56        .max(plane_size)
57        .min(hardware.max_cube_dim.0);
58
59    let cube_dim = CubeDim::new(client, requested as usize).num_elems();
60
61    let sms = hardware.num_streaming_multiprocessors.unwrap_or(64);
62    let cube_count = (sms * 32).min(hardware.max_cube_count.0);
63
64    let vector_size = client
65        .io_optimized_vector_sizes(dtype.size())
66        .next()
67        .unwrap_or(1);
68
69    LaunchConfig {
70        cube_dim: cube_dim as usize,
71        cube_count: cube_count as usize,
72        vector_size,
73        plane_size: plane_size as usize,
74    }
75}
76
77/// Measures the fixed cost of a single kernel launch.
78///
79/// Native only, panics on WASM
80pub fn measure_launch_overhead<R: Runtime>(client: &ComputeClient<R>) -> core::time::Duration {
81    client.measure_launch_overhead(|| {
82        let input = client.empty(size_of::<i32>());
83        let output = client.empty(size_of::<i32>());
84
85        let (_, duration) = client
86            .profile(
87                || unsafe {
88                    launch_overhead::launch_overhead::launch_unchecked::<R>(
89                        client,
90                        CubeCount::new_single(),
91                        CubeDim::new_single(),
92                        1,
93                        BufferArg::from_raw_parts(input.clone(), 1),
94                        BufferArg::from_raw_parts(output.clone(), 1),
95                        ElemType::Int(IntKind::I32).into(),
96                    );
97                },
98                "launch_overhead",
99            )
100            .expect("should succeed launch_overhead");
101
102        cubecl_core::future::block_on(duration.into_future()).duration()
103    })
104}
cubecl_std/throughput/base.rs

cubecl_std/throughput/
base.rs
