cubecl_runtime/tune/
bounds_generator.rs

1use core::time::Duration;
2
3use alloc::vec::Vec;
4
5use crate::throughput::{ThroughputKey, ThroughputValue};
6use crate::tune::TuneInputs;
7
8/// A set of [`AutotuneBound`]s for a given key and reference inputs, with a launch overhead.
9#[derive(Debug, Clone, PartialEq)]
10#[cfg_attr(autotune_persistence, derive(serde::Serialize, serde::Deserialize))]
11pub struct Bounds {
12    /// The bounds for autotuning.
13    pub bounds: Vec<AutotuneBound>,
14    /// The launch overhead for autotuning.
15    pub launch_overhead: Duration,
16}
17
18// Sound because [`AutotuneBound`] compares its floats bitwise, so equality stays reflexive.
19impl Eq for Bounds {}
20
21/// Produces a set of [`AutotuneBound`]s for a given key and reference inputs.
22#[diagnostic::on_unimplemented(
23    message = "`{Self}` is not a valid bounds generator",
24    label = "invalid bounds generator"
25)]
26pub trait BoundsGenerator<K, I: TuneInputs>: Send + Sync + 'static {
27    /// Generate a set of bounds for a given key and reference inputs.
28    fn generate<'a>(&self, key: &K, inputs: &I::At<'a>) -> Bounds;
29}
30
31/// `Fn(&K, &A) -> Bounds` acts as a [`BoundsGenerator`] when `A` is an owned type. For
32/// multi-input kernels, `A` is a tuple that the closure destructures internally.
33impl<K, Func, A> BoundsGenerator<K, A> for Func
34where
35    A: Clone + Send + Sync + 'static,
36    K: 'static,
37    Func: Send + Sync + 'static + Fn(&K, &A) -> Bounds,
38{
39    #[inline]
40    fn generate<'a>(&self, key: &K, inputs: &<A as TuneInputs>::At<'a>) -> Bounds {
41        (self)(key, inputs)
42    }
43}
44
45/// A calculator that determines the time limit for autotune bounds.
46pub trait TimeBound {
47    /// Returns the time limit for autotune bounds.
48    fn time_limit(&self) -> Option<Duration>;
49}
50
51/// A bound for autotuning a throughput kernel, specifying the key, threshold, and number of operations.
52#[derive(Debug, Clone)]
53#[cfg_attr(autotune_persistence, derive(serde::Serialize, serde::Deserialize))]
54pub struct AutotuneBound {
55    /// Peak throughput of the reference kernel, in ops (or bytes) per second.
56    pub throughput: f64,
57    /// The threshold for this bound, over which the kernel will be considered accurate.
58    pub threshold: f32,
59    /// The number of operations the kernel will run.
60    pub ops_count: usize,
61}
62
63/// Bitwise comparison of the measured throughputs, so that equality is reflexive even if a
64/// degenerate measurement ever produces a `NaN`, which is what makes the [`Eq`] below sound.
65impl PartialEq for AutotuneBound {
66    fn eq(&self, other: &Self) -> bool {
67        self.throughput.to_bits() == other.throughput.to_bits()
68            && self.threshold.to_bits() == other.threshold.to_bits()
69            && self.ops_count == other.ops_count
70    }
71}
72
73impl Eq for AutotuneBound {}
74
75/// Work required by a problem, specified in minimum compute operations and byte transfers.
76#[derive(Debug, Clone, Copy, PartialEq, Eq)]
77#[cfg_attr(std_io, derive(serde::Serialize, serde::Deserialize))]
78pub struct Work {
79    /// Compute operations required.
80    pub compute_ops: usize,
81    /// Memory bytes transferred (reads and writes).
82    pub bytes: usize,
83}
84
85/// Target fractions of modeled peak compute and memory roofline throughput.
86#[derive(Debug, Clone, Copy, PartialEq)]
87#[cfg_attr(std_io, derive(serde::Serialize, serde::Deserialize))]
88pub struct Thresholds {
89    /// Fraction of peak compute throughput expected.
90    pub compute: f32,
91    /// Fraction of peak memory bandwidth expected.
92    pub memory: f32,
93}
94
95impl Thresholds {
96    /// The same fraction for both bounds.
97    pub const fn uniform(fraction: f32) -> Self {
98        Self {
99            compute: fraction,
100            memory: fraction,
101        }
102    }
103}
104
105impl Default for Thresholds {
106    /// The roofline itself: a candidate is expected to reach 100% of the modeled peak, which
107    /// is the only fraction that needs no justification.
108    fn default() -> Self {
109        Self::uniform(1.0)
110    }
111}
112
113/// Standardizes the creation of compute and memory [`AutotuneBound`]s.
114pub fn calculate_bounds(
115    work: Work,
116    thresholds: Thresholds,
117    compute_throughput: &ThroughputValue,
118    memory_throughput: &ThroughputValue,
119    memory_key: &ThroughputKey,
120) -> Vec<AutotuneBound> {
121    alloc::vec![
122        AutotuneBound {
123            ops_count: work.compute_ops,
124            throughput: compute_throughput.ops_per_s(),
125            threshold: thresholds.compute,
126        },
127        AutotuneBound {
128            ops_count: work.bytes,
129            throughput: memory_throughput.bytes_per_s(memory_key),
130            threshold: thresholds.memory,
131        },
132    ]
133}
134
135impl TimeBound for AutotuneBound {
136    fn time_limit(&self) -> Option<Duration> {
137        if self.throughput.is_normal() && self.threshold.is_normal() {
138            Some(Duration::from_secs_f64(
139                (self.ops_count as f64 / self.throughput) / self.threshold as f64,
140            ))
141        } else {
142            None
143        }
144    }
145}
146
147impl<B: TimeBound> TimeBound for Vec<B> {
148    fn time_limit(&self) -> Option<Duration> {
149        self.iter().filter_map(|b| b.time_limit()).max()
150    }
151}
152
153impl TimeBound for Bounds {
154    fn time_limit(&self) -> Option<Duration> {
155        self.bounds
156            .time_limit()
157            .map(|limit| limit + self.launch_overhead)
158    }
159}
160
161#[cfg(test)]
162mod tests {
163    use crate::throughput::ThroughputMode;
164
165    use super::*;
166    use alloc::vec;
167
168    fn bound(ops_count: usize, throughput: f64, threshold: f32) -> AutotuneBound {
169        AutotuneBound {
170            throughput,
171            threshold,
172            ops_count,
173        }
174    }
175
176    #[test]
177    fn time_limit_is_ops_over_throughput_scaled_by_threshold() {
178        // (8 ops / 4 ops/s) / 0.5 = 4s. Powers of two keep the f64 math exact.
179        let limit = bound(8, 4.0, 0.5).time_limit();
180        assert_eq!(limit, Some(Duration::from_secs(4)));
181    }
182
183    #[test]
184    fn time_limit_is_none_when_inputs_are_not_normal() {
185        // A zero/NaN/inf throughput or a zero threshold would divide by zero or blow up,
186        // so the bound disables the short-circuit instead of producing a garbage limit.
187        assert_eq!(bound(8, 0.0, 0.5).time_limit(), None);
188        assert_eq!(bound(8, f64::NAN, 0.5).time_limit(), None);
189        assert_eq!(bound(8, f64::INFINITY, 0.5).time_limit(), None);
190        assert_eq!(bound(8, 4.0, 0.0).time_limit(), None);
191    }
192
193    #[test]
194    fn vec_time_limit_takes_the_roofline_max_not_min() {
195        // Two simultaneous resource bounds (e.g. compute vs memory): the achievable floor
196        // is the *slower* one, so the reduction must be `max`. `min` would pick the
197        // unreachable 1s and the short-circuit would never fire.
198        let compute = bound(8, 4.0, 1.0); // 2s
199        let memory = bound(8, 8.0, 1.0); // 1s
200        let limit = vec![compute, memory].time_limit();
201        assert_eq!(limit, Some(Duration::from_secs(2)));
202    }
203
204    #[test]
205    fn vec_time_limit_skips_non_normal_bounds_and_is_none_when_empty() {
206        // A non-normal bound is filtered out rather than poisoning the reduction.
207        let limit = vec![bound(8, 0.0, 1.0), bound(8, 4.0, 1.0)].time_limit();
208        assert_eq!(limit, Some(Duration::from_secs(2)));
209
210        assert_eq!(Vec::<AutotuneBound>::new().time_limit(), None);
211    }
212
213    #[test]
214    fn bounds_time_limit_adds_launch_overhead() {
215        let bounds = Bounds {
216            bounds: vec![bound(8, 4.0, 1.0)], // 2s
217            launch_overhead: Duration::from_millis(500),
218        };
219        assert_eq!(bounds.time_limit(), Some(Duration::from_millis(2500)));
220    }
221
222    #[test]
223    fn calculate_bounds_applies_a_threshold_per_resource() {
224        let work = Work {
225            compute_ops: 8,
226            bytes: 16,
227        };
228        let thresholds = Thresholds {
229            compute: 0.5,
230            memory: 1.0,
231        };
232        let key = ThroughputKey {
233            mode: ThroughputMode::Memory,
234        };
235
236        let bounds = calculate_bounds(
237            work,
238            thresholds,
239            &ThroughputValue::ZERO,
240            &ThroughputValue::ZERO,
241            &key,
242        );
243
244        assert_eq!(bounds[0].ops_count, 8);
245        assert_eq!(bounds[0].threshold, 0.5);
246        assert_eq!(bounds[1].ops_count, 16);
247        assert_eq!(bounds[1].threshold, 1.0);
248    }
249
250    #[test]
251    fn the_default_threshold_is_the_roofline_itself() {
252        assert_eq!(Thresholds::default(), Thresholds::uniform(1.0));
253        assert_eq!(Thresholds::default().compute, 1.0);
254    }
255
256    #[test]
257    fn bounds_time_limit_is_none_without_usable_bounds() {
258        // No usable bound means no limit at all — the launch overhead is not a limit on
259        // its own, so the short-circuit stays disabled.
260        let bounds = Bounds {
261            bounds: vec![],
262            launch_overhead: Duration::from_millis(500),
263        };
264        assert_eq!(bounds.time_limit(), None);
265    }
266}
cubecl_runtime/tune/bounds_generator.rs

cubecl_runtime/tune/
bounds_generator.rs
