cubecl_runtime/throughput/
benchmarker.rs

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

cubecl_runtime/throughput/
benchmarker.rs
