cubecl_std/throughput/
base.rs

1use cubecl_core::ir::ElemType;
2use cubecl_runtime::{
3    client::ComputeClient,
4    runtime::Runtime,
5    server::CubeDim,
6    throughput::{
7        DEFAULT_BUFFER_BYTES, MemoryAccess, MemoryCurve, MemoryPoint, ThroughputKey,
8        ThroughputMode, ThroughputValue, working_set_sweep,
9    },
10    tune::{Bounds, Thresholds, Work, calculate_bounds},
11};
12
13use crate::throughput::{
14    compute_cmma, compute_direct, launch_overhead, memory_direct, memory_read, memory_write,
15};
16
17/// Independent cube positions each CPU worker interleaves, so a compute
18/// pass pipelines past instruction latency instead of serializing on one
19/// dependency chain. A depth, not a machine guess: a handful hides any
20/// core's fma latency, excess is free because the iteration budget is
21/// time-calibrated, and nothing about the launch scales with it. Memory
22/// probes with blocked addressing pin their own count back to one (see
23/// [`MemoryProbe::new`](crate::throughput::memory_probe::MemoryProbe::new)).
24const CPU_CHAIN_DEPTH: usize = 64;
25
26/// Measure peak throughput on `device` for each of the given `keys`.
27pub fn device_throughput<R: Runtime>(
28    device: &R::Device,
29    keys: &[ThroughputKey],
30) -> alloc::vec::Vec<ThroughputValue> {
31    let client = R::client(device);
32    keys.iter()
33        .map(|key| measure_peak_throughput::<R>(&client, *key))
34        .collect()
35}
36
37/// Measure the memory ceiling across a range of working sets, from a few
38/// kilobytes up to as much as the device will allocate.
39///
40/// One point per size in [`working_set_sweep`], each measured and cached
41/// exactly like the single-size probe — [`measure_peak_throughput`] with a
42/// [`ThroughputMode::MemoryWorkingSet`] key — so a curve costs one probe per
43/// size on the first run and nothing afterwards.
44///
45/// Native only, panics on WASM
46pub fn measure_memory_curve<R: Runtime>(
47    client: &ComputeClient<R>,
48    access: MemoryAccess,
49) -> MemoryCurve {
50    let points = working_set_sweep(working_set_cap(client, access))
51        .into_iter()
52        .map(|bytes| {
53            let key = ThroughputKey {
54                mode: ThroughputMode::MemoryWorkingSet { access, bytes },
55            };
56
57            MemoryPoint {
58                bytes,
59                value: measure_peak_throughput::<R>(client, key),
60            }
61        });
62
63    MemoryCurve::new(access, points)
64}
65
66/// The largest working set `access` can be probed at: as much as one buffer can
67/// hold, times the buffers the access touches.
68fn working_set_cap<R: Runtime>(client: &ComputeClient<R>, access: MemoryAccess) -> u64 {
69    let max_alloc = client.properties().memory.max_page_size;
70
71    DEFAULT_BUFFER_BYTES.min(max_alloc) * access.buffers()
72}
73
74/// Computes the peak throughput for a given runtime and key.
75///
76/// Native only, panics on WASM
77pub fn measure_peak_throughput<R: Runtime>(
78    client: &ComputeClient<R>,
79    key: ThroughputKey,
80) -> ThroughputValue {
81    // A throughput probe is a measurement: inside a dry run its launches must
82    // still execute, or they would be timed anyway and cache a garbage peak in
83    // the device-level throughput store. The guard is read where the launch is
84    // issued, which for these is this thread.
85    let _measurement = cubecl_runtime::dry_run::RealRun::new();
86
87    let launch_config = launch_config(client, key.dtype());
88
89    let kernel_config = match key.mode {
90        ThroughputMode::ComputeDirect { .. } => {
91            compute_direct::build_kernel(client, key, launch_config)
92        }
93        ThroughputMode::ComputeCmma {
94            config: cmma_config,
95            ..
96        } => {
97            if client.properties().features.matmul.cmma.is_empty() {
98                return ThroughputValue::ZERO;
99            }
100            compute_cmma::build_kernel(client, key, cmma_config, launch_config)
101        }
102        ThroughputMode::Memory
103        | ThroughputMode::MemoryRead
104        | ThroughputMode::MemoryWrite
105        | ThroughputMode::MemoryWorkingSet { .. } => {
106            // The memory modes differ only in access and working set, and
107            // `memory_probe` is the one place that mapping lives.
108            let (access, working_set) = key
109                .mode
110                .memory_probe()
111                .expect("A memory mode describes a probe");
112            let working_set = working_set.min(usize::MAX as u64) as usize;
113
114            match access {
115                MemoryAccess::Copy => {
116                    memory_direct::build_kernel(client, key, launch_config, working_set)
117                }
118                MemoryAccess::Read => {
119                    memory_read::build_kernel(client, key, launch_config, working_set)
120                }
121                MemoryAccess::Write => {
122                    memory_write::build_kernel(client, key, launch_config, working_set)
123                }
124            }
125        }
126        ThroughputMode::Launch => launch_overhead::build_kernel(client, key, launch_config),
127    };
128
129    let value = client.measure_throughput(key, kernel_config);
130
131    client.memory_cleanup();
132
133    value
134}
135
136/// Calculates roofline autotune bounds for a given [`Work`] amount and compute throughput key.
137///
138/// Measures compute and memory peak throughputs along with launch overhead for the runtime client.
139pub fn roofline_bounds<R: Runtime>(
140    client: &ComputeClient<R>,
141    compute_key: ThroughputKey,
142    work: Work,
143    thresholds: Thresholds,
144) -> Bounds {
145    let memory_key = ThroughputKey {
146        mode: ThroughputMode::Memory,
147    };
148    let launch_key = ThroughputKey {
149        mode: ThroughputMode::Launch,
150    };
151
152    Bounds {
153        bounds: calculate_bounds(
154            work,
155            thresholds,
156            &measure_peak_throughput(client, compute_key),
157            &measure_peak_throughput(client, memory_key),
158            &memory_key,
159        ),
160        launch_overhead: measure_peak_throughput(client, launch_key).duration_per_op(),
161    }
162}
163
164/// Hardware execution parameters for launching a compute kernel.
165#[derive(Clone, Copy)]
166pub struct LaunchConfig {
167    /// The number of threads per cube.
168    pub cube_dim: usize,
169    /// The total number of cubes to dispatch.
170    pub cube_count: usize,
171    /// The vectorization factor (e.g., 4 for `vec4` operations).
172    pub vector_size: usize,
173    /// The number of threads in a hardware execution plane.
174    pub plane_size: usize,
175}
176
177fn launch_config<R: Runtime>(client: &ComputeClient<R>, dtype: ElemType) -> LaunchConfig {
178    let hardware = &client.properties().hardware;
179
180    let plane_size = hardware.plane_size_max.max(1);
181    let vector_size = client
182        .io_optimized_vector_sizes(dtype.size())
183        .next()
184        .unwrap_or(1);
185
186    // A CPU has no SMs, so `sms * 32` cubes is the wrong grid to size from:
187    // a cube's units are its real dispatched workers here, while its cube
188    // count is only a loop inside each of them. `num_cpu_cores` units, one
189    // per core, is the real worker count.
190    if let Some(cores) = hardware.num_cpu_cores {
191        return LaunchConfig {
192            cube_dim: cores as usize,
193            cube_count: CPU_CHAIN_DEPTH,
194            vector_size,
195            plane_size: plane_size as usize,
196        };
197    }
198
199    let requested = (hardware.max_units_per_cube / plane_size * plane_size)
200        .max(plane_size)
201        .min(hardware.max_cube_dim.0);
202
203    let cube_dim = CubeDim::new(client, requested as usize).num_elems();
204
205    let sms = hardware.num_streaming_multiprocessors.unwrap_or(64);
206    let cube_count = (sms * 32).min(hardware.max_cube_count.0);
207
208    LaunchConfig {
209        cube_dim: cube_dim as usize,
210        cube_count: cube_count as usize,
211        vector_size,
212        plane_size: plane_size as usize,
213    }
214}
cubecl_std/throughput/base.rs

cubecl_std/throughput/
base.rs
