cubecl_std/throughput/
launch.rs

1use cubecl_core::ir::ElemType;
2use cubecl_runtime::{client::Client, server::CubeDim};
3
4/// Independent cube positions each CPU worker interleaves, so a compute
5/// pass pipelines past instruction latency instead of serializing on one
6/// dependency chain. A depth, not a machine guess: a handful hides any
7/// core's fma latency, excess is free because the iteration budget is
8/// time-calibrated, and nothing about the launch scales with it. Memory
9/// probes with blocked addressing pin their own count back to one (see
10/// [`MemoryProbe::new`](crate::throughput::memory_probe::MemoryProbe::new)).
11const CPU_CHAIN_DEPTH: usize = 64;
12
13/// Units a GPU probe asks for. A wider cube measures no faster, and makes the
14/// memory probes report several times the bus rate.
15const PROBE_UNITS_PER_CUBE: u32 = 256;
16
17/// Hardware execution parameters for launching a compute kernel.
18#[derive(Clone, Copy)]
19pub struct LaunchConfig {
20    /// The cube the probe launches, resolved once so `ops_count` cannot
21    /// describe a launch that did not happen.
22    pub cube_dim: CubeDim,
23    /// The total number of cubes to dispatch.
24    pub cube_count: usize,
25    /// The vectorization factor (e.g., 4 for `vec4` operations).
26    pub vector_size: usize,
27    /// The number of threads in a hardware execution plane.
28    pub plane_size: usize,
29}
30
31impl LaunchConfig {
32    /// The launch a probe of `dtype` is issued in on this device.
33    pub(super) fn for_device(client: &Client, dtype: ElemType) -> Self {
34        let hardware = &client.properties().hardware;
35
36        let plane_size = hardware.plane_size_max.max(1);
37        let vector_size = client
38            .io_optimized_vector_sizes(dtype.size())
39            .next()
40            .unwrap_or(1);
41
42        // A CPU has no SMs, so `sms * 32` cubes is the wrong grid to size from:
43        // a cube's units are its real dispatched workers here, while its cube
44        // count is only a loop inside each of them. `num_cpu_cores` units, one
45        // per core, is the real worker count.
46        let (units, cube_count) = match hardware.num_cpu_cores {
47            Some(cores) => (cores, CPU_CHAIN_DEPTH as u32),
48            None => {
49                let sms = hardware.num_streaming_multiprocessors.unwrap_or(64);
50                (
51                    PROBE_UNITS_PER_CUBE,
52                    (sms * 32).min(hardware.max_cube_count.0),
53                )
54            }
55        };
56
57        Self {
58            cube_dim: CubeDim::new(client, units as usize),
59            cube_count: cube_count as usize,
60            vector_size,
61            plane_size: plane_size as usize,
62        }
63    }
64
65    /// The same launch, dispatched across `units` workers.
66    pub(super) fn with_units(self, client: &Client, units: u32) -> Self {
67        Self {
68            cube_dim: CubeDim::new(client, units as usize),
69            ..self
70        }
71    }
72}
cubecl_std/throughput/launch.rs

cubecl_std/throughput/
launch.rs
