cubecl_std/throughput/
base.rs

1use cubecl_core::ir::ElemType;
2use cubecl_runtime::{
3    client::ComputeClient,
4    runtime::Runtime,
5    server::CubeDim,
6    throughput::{ThroughputKey, ThroughputMode, ThroughputValue},
7    tune::{Bounds, Thresholds, Work, calculate_bounds},
8};
9
10use crate::throughput::{compute_cmma, compute_direct, launch_overhead, memory_direct};
11
12/// Measure peak throughput on `device` for each of the given `keys`.
13pub fn device_throughput<R: Runtime>(
14    device: &R::Device,
15    keys: &[ThroughputKey],
16) -> alloc::vec::Vec<ThroughputValue> {
17    let client = R::client(device);
18    keys.iter()
19        .map(|key| measure_peak_throughput::<R>(&client, *key))
20        .collect()
21}
22
23/// Computes the peak throughput for a given runtime and key.
24///
25/// Native only, panics on WASM
26pub fn measure_peak_throughput<R: Runtime>(
27    client: &ComputeClient<R>,
28    key: ThroughputKey,
29) -> ThroughputValue {
30    let launch_config = launch_config(client, key.dtype());
31
32    let kernel_config = match key.mode {
33        ThroughputMode::ComputeDirect { .. } => {
34            compute_direct::build_kernel(client, key, launch_config)
35        }
36        ThroughputMode::ComputeCmma {
37            config: cmma_config,
38            ..
39        } => {
40            if client.properties().features.matmul.cmma.is_empty() {
41                return ThroughputValue::ZERO;
42            }
43            compute_cmma::build_kernel(client, key, cmma_config, launch_config)
44        }
45        ThroughputMode::Memory => memory_direct::build_kernel(client, key, launch_config),
46        ThroughputMode::Launch => launch_overhead::build_kernel(client, key, launch_config),
47    };
48
49    let value = client.measure_throughput(key, kernel_config);
50
51    client.memory_cleanup();
52
53    value
54}
55
56/// Calculates roofline autotune bounds for a given [`Work`] amount and compute throughput key.
57///
58/// Measures compute and memory peak throughputs along with launch overhead for the runtime client.
59pub fn roofline_bounds<R: Runtime>(
60    client: &ComputeClient<R>,
61    compute_key: ThroughputKey,
62    work: Work,
63    thresholds: Thresholds,
64) -> Bounds {
65    let memory_key = ThroughputKey {
66        mode: ThroughputMode::Memory,
67    };
68    let launch_key = ThroughputKey {
69        mode: ThroughputMode::Launch,
70    };
71
72    Bounds {
73        bounds: calculate_bounds(
74            work,
75            thresholds,
76            &measure_peak_throughput(client, compute_key),
77            &measure_peak_throughput(client, memory_key),
78            &memory_key,
79        ),
80        launch_overhead: measure_peak_throughput(client, launch_key).duration_per_op(),
81    }
82}
83
84/// Hardware execution parameters for launching a compute kernel.
85#[derive(Clone, Copy)]
86pub struct LaunchConfig {
87    /// The number of threads per cube.
88    pub cube_dim: usize,
89    /// The total number of cubes to dispatch.
90    pub cube_count: usize,
91    /// The vectorization factor (e.g., 4 for `vec4` operations).
92    pub vector_size: usize,
93    /// The number of threads in a hardware execution plane.
94    pub plane_size: usize,
95}
96
97fn launch_config<R: Runtime>(client: &ComputeClient<R>, dtype: ElemType) -> LaunchConfig {
98    let hardware = &client.properties().hardware;
99
100    let plane_size = hardware.plane_size_max.max(1);
101    let requested = (hardware.max_units_per_cube / plane_size * plane_size)
102        .max(plane_size)
103        .min(hardware.max_cube_dim.0);
104
105    let cube_dim = CubeDim::new(client, requested as usize).num_elems();
106
107    let sms = hardware.num_streaming_multiprocessors.unwrap_or(64);
108    let cube_count = (sms * 32).min(hardware.max_cube_count.0);
109
110    let vector_size = client
111        .io_optimized_vector_sizes(dtype.size())
112        .next()
113        .unwrap_or(1);
114
115    LaunchConfig {
116        cube_dim: cube_dim as usize,
117        cube_count: cube_count as usize,
118        vector_size,
119        plane_size: plane_size as usize,
120    }
121}
cubecl_std/throughput/base.rs

cubecl_std/throughput/
base.rs
