cubecl_runtime/throughput/
benchmarker.rs

1use crate::{
2    client::ComputeClient,
3    config::CubeClRuntimeConfig,
4    runtime::Runtime,
5    throughput::{ThroughputCache, ThroughputKey, ThroughputValue},
6};
7use alloc::boxed::Box;
8use alloc::sync::Arc;
9use cubecl_common::{
10    config::RuntimeConfig,
11    future::block_on,
12    profile::{Duration, Instant},
13};
14use spin::Mutex;
15
16type Cache = Arc<Mutex<ThroughputCache>>;
17
18/// Configuration and payload for a benchmarkable compute kernel.
19pub struct KernelConfig {
20    /// The executable kernel closure to be evaluated.
21    pub kernel: Box<dyn Fn(usize)>,
22    /// The number of operations processed in one iteration.
23    pub ops_count: usize,
24}
25
26/// A marker for measuring throughput of compute kernels.
27pub struct ThroughputBenchmarker {
28    cache: Cache,
29    cache_enabled: bool,
30}
31
32impl ThroughputBenchmarker {
33    /// Creates a new `ThroughputBenchmarker` with the given cache.
34    pub fn new(cache: Cache) -> Self {
35        let cache_enabled = !CubeClRuntimeConfig::get().throughput.disable_cache;
36        Self {
37            cache,
38            cache_enabled,
39        }
40    }
41
42    /// Measure the maximum compute throughput of the given kernel on the given client.
43    /// Warms up the kernel until it plateaus,
44    /// then measures the throughput over multiple iterations taking the minimum time per iteration (peak attained).
45    pub fn measure<R: Runtime>(
46        &mut self,
47        client: &ComputeClient<R>,
48        key: ThroughputKey,
49        kernel_config: KernelConfig,
50    ) -> ThroughputValue {
51        if self.cache_enabled
52            && let Some(cached_value) = self.cache.lock().get(&key)
53        {
54            return *cached_value;
55        }
56
57        let kernel = kernel_config.kernel;
58
59        let sample = |iterations: usize| {
60            let start = Instant::now();
61            kernel(iterations);
62            let _ = block_on(client.sync());
63            start.elapsed()
64        };
65
66        let iterations = self.warmup(sample);
67        let duration = self.sample_peak_duration(iterations, sample);
68
69        let value = ThroughputValue {
70            ops_count: kernel_config.ops_count,
71            duration,
72        };
73
74        if self.cache_enabled {
75            self.cache.lock().insert(key, value);
76        }
77
78        value
79    }
80
81    /// Warms up the device by running the kernel multiple times
82    /// and estimating the number of iterations needed to reach a stable duration.
83    fn warmup(&self, sample: impl Fn(usize) -> Duration) -> usize {
84        const MAX_WARMUP: usize = 50;
85        const PLATEAU_TOL: f64 = 0.03;
86        const PATIENCE: usize = 3;
87        const TARGET_DURATION_MS: f64 = 20.0;
88
89        let mut best = f64::INFINITY;
90        let mut stable = 0;
91        let mut iterations = 1;
92
93        for _ in 0..MAX_WARMUP {
94            let duration = sample(iterations).as_secs_f64() * 1000.0;
95            if duration < TARGET_DURATION_MS {
96                let extra_iters = if duration > 1e-6 {
97                    let duration_per_iter = duration / iterations as f64;
98                    ((TARGET_DURATION_MS - duration) / duration_per_iter).ceil() as usize
99                } else {
100                    iterations
101                };
102                iterations += extra_iters.max(1);
103                best = f64::INFINITY;
104                stable = 0;
105                continue;
106            }
107
108            let duration_per_iter = duration / iterations as f64;
109            if duration_per_iter < best * (1.0 - PLATEAU_TOL) {
110                best = duration_per_iter;
111                stable = 0;
112            } else {
113                best = best.min(duration_per_iter);
114                stable += 1;
115                if stable >= PATIENCE {
116                    break;
117                }
118            }
119        }
120
121        iterations
122    }
123
124    /// Sample the peak throughput of the kernel by running it multiple times
125    /// and measuring the duration of each iteration.
126    fn sample_peak_duration(
127        &self,
128        iterations: usize,
129        sample_once: impl Fn(usize) -> Duration,
130    ) -> Duration {
131        const MIN_SAMPLES: usize = 20;
132        const MAX_SAMPLES: usize = 200;
133        const REL_TOL: f64 = 0.01;
134        const PATIENCE: usize = 12;
135
136        let mut best = f64::INFINITY;
137        let mut stale = 0;
138
139        for i in 0..MAX_SAMPLES {
140            let s = sample_once(iterations).as_secs_f64();
141            if s < best * (1.0 - REL_TOL) {
142                best = s;
143                stale = 0;
144            } else {
145                best = best.min(s);
146                stale += 1;
147            }
148            if i > MIN_SAMPLES && stale >= PATIENCE {
149                break;
150            }
151        }
152
153        Duration::from_secs_f64(best / iterations as f64)
154    }
155}
cubecl_runtime/throughput/benchmarker.rs

cubecl_runtime/throughput/
benchmarker.rs
