cubecl_runtime/tune/
bounds_generator.rs

1use core::time::Duration;
2
3use alloc::vec::Vec;
4
5use crate::config::autotune::AutotuneLevel;
6use crate::tune::TuneInputs;
7
8// A bound-builder constructs `AutotuneBound { resource: ResourceBound { .. }, .. }`
9// alongside this module's own types, so the neutral record is re-exported
10// here too, the same way `Work` is below.
11pub use crate::throughput::ResourceBound;
12
13/// A set of [`AutotuneBound`]s for a given key and reference inputs, with a launch overhead.
14#[derive(Debug, Clone, PartialEq)]
15#[cfg_attr(serializable, derive(serde::Serialize, serde::Deserialize))]
16pub struct Bounds {
17    /// The bounds for autotuning.
18    pub bounds: Vec<AutotuneBound>,
19    /// The launch overhead for autotuning.
20    pub launch_overhead: Duration,
21}
22
23// Sound because [`AutotuneBound`] compares its floats bitwise, so equality stays reflexive.
24impl Eq for Bounds {}
25
26/// Produces a set of [`AutotuneBound`]s for a given key and reference inputs.
27#[diagnostic::on_unimplemented(
28    message = "`{Self}` is not a valid bounds generator",
29    label = "invalid bounds generator"
30)]
31pub trait BoundsGenerator<K, I: TuneInputs>: Send + Sync + 'static {
32    /// Generate a set of bounds for a given key and reference inputs.
33    fn generate<'a>(&self, key: &K, inputs: &I::At<'a>) -> Bounds;
34}
35
36/// `Fn(&K, &A) -> Bounds` acts as a [`BoundsGenerator`] when `A` is an owned type. For
37/// multi-input kernels, `A` is a tuple that the closure destructures internally.
38impl<K, Func, A> BoundsGenerator<K, A> for Func
39where
40    A: Clone + Send + Sync + 'static,
41    K: 'static,
42    Func: Send + Sync + 'static + Fn(&K, &A) -> Bounds,
43{
44    #[inline]
45    fn generate<'a>(&self, key: &K, inputs: &<A as TuneInputs>::At<'a>) -> Bounds {
46        (self)(key, inputs)
47    }
48}
49
50/// A calculator that determines the time limit for autotune bounds.
51pub trait TimeBound {
52    /// Returns the time limit for autotune bounds.
53    fn time_limit(&self) -> Option<Duration>;
54}
55
56/// A bound for autotuning a throughput kernel: a [`ResourceBound`] plus the
57/// threshold over which the kernel is considered accurate.
58#[derive(Debug, Clone)]
59#[cfg_attr(serializable, derive(serde::Serialize, serde::Deserialize))]
60pub struct AutotuneBound {
61    /// How much work, against what peak throughput.
62    pub resource: ResourceBound,
63    /// The threshold for this bound, over which the kernel will be considered accurate.
64    pub threshold: f32,
65}
66
67/// Bitwise comparison of the measured throughputs, so that equality is reflexive even if a
68/// degenerate measurement ever produces a `NaN`, which is what makes the [`Eq`] below sound.
69impl PartialEq for AutotuneBound {
70    fn eq(&self, other: &Self) -> bool {
71        self.resource.peak_per_s.to_bits() == other.resource.peak_per_s.to_bits()
72            && self.threshold.to_bits() == other.threshold.to_bits()
73            && self.resource.amount == other.resource.amount
74    }
75}
76
77impl Eq for AutotuneBound {}
78
79// `cubecl-common` cannot depend on `cubecl-runtime`, so `Work` lives there and
80// autotune bounds and benchmark reporting share one definition.
81pub use cubecl_common::work::Work;
82
83/// Target fractions of modeled peak compute and memory roofline throughput.
84#[derive(Debug, Clone, Copy, PartialEq)]
85#[cfg_attr(serializable, derive(serde::Serialize, serde::Deserialize))]
86pub struct Thresholds {
87    /// Fraction of peak compute throughput expected.
88    pub compute: f32,
89    /// Fraction of peak memory bandwidth expected.
90    pub memory: f32,
91}
92
93impl Thresholds {
94    /// The same fraction for both bounds.
95    pub const fn uniform(fraction: f32) -> Self {
96        Self {
97            compute: fraction,
98            memory: fraction,
99        }
100    }
101
102    /// The threshold that sets no limit, so every candidate is measured.
103    /// [`AutotuneBound::time_limit`] declines it.
104    pub const UNBOUNDED: Self = Self::uniform(0.0);
105
106    /// The fraction of peak an [`AutotuneLevel`] settles for.
107    ///
108    /// A higher threshold is a tighter time limit, since the limit is the roofline time
109    /// divided by the threshold. The fractions are ad-hoc observations rather than a
110    /// systematic sweep, so nothing should depend on the exact values.
111    pub const fn for_level(level: &AutotuneLevel) -> Self {
112        match level {
113            AutotuneLevel::Minimal => Self::uniform(0.6),
114            AutotuneLevel::Balanced => Self::uniform(0.8),
115            AutotuneLevel::Extensive => Self::uniform(0.95),
116            AutotuneLevel::Full => Self::UNBOUNDED,
117        }
118    }
119}
120
121impl Default for Thresholds {
122    /// The roofline itself: a candidate is expected to reach 100% of the modeled peak, which
123    /// is the only fraction that needs no justification.
124    fn default() -> Self {
125        Self::uniform(1.0)
126    }
127}
128
129impl TimeBound for AutotuneBound {
130    fn time_limit(&self) -> Option<Duration> {
131        // The threshold divides the roofline time. A negative one panics `div_f64`, and
132        // zero or a subnormal divides the limit away, so neither is a limit to compute.
133        if self.threshold <= 0.0 || !self.threshold.is_normal() {
134            return None;
135        }
136        self.resource
137            .time_at_peak()
138            .map(|limit| limit.div_f64(self.threshold as f64))
139    }
140}
141
142impl<B: TimeBound> TimeBound for Vec<B> {
143    fn time_limit(&self) -> Option<Duration> {
144        self.iter().filter_map(|b| b.time_limit()).max()
145    }
146}
147
148impl TimeBound for Bounds {
149    fn time_limit(&self) -> Option<Duration> {
150        self.bounds
151            .time_limit()
152            .map(|limit| limit + self.launch_overhead)
153    }
154}
155
156#[cfg(test)]
157mod tests {
158    use super::*;
159    use alloc::vec;
160
161    fn bound(ops_count: usize, throughput: f64, threshold: f32) -> AutotuneBound {
162        AutotuneBound {
163            resource: ResourceBound {
164                amount: ops_count,
165                peak_per_s: throughput,
166            },
167            threshold,
168        }
169    }
170
171    #[test]
172    fn time_limit_is_ops_over_throughput_scaled_by_threshold() {
173        // (8 ops / 4 ops/s) / 0.5 = 4s. Powers of two keep the f64 math exact.
174        let limit = bound(8, 4.0, 0.5).time_limit();
175        assert_eq!(limit, Some(Duration::from_secs(4)));
176    }
177
178    #[test]
179    fn time_limit_is_none_when_inputs_are_not_normal() {
180        // A zero/NaN/inf throughput or a zero threshold would divide by zero or blow up,
181        // so the bound disables the short-circuit instead of producing a garbage limit.
182        assert_eq!(bound(8, 0.0, 0.5).time_limit(), None);
183        assert_eq!(bound(8, f64::NAN, 0.5).time_limit(), None);
184        assert_eq!(bound(8, f64::INFINITY, 0.5).time_limit(), None);
185        assert_eq!(bound(8, 4.0, 0.0).time_limit(), None);
186    }
187
188    #[test]
189    fn time_limit_declines_a_negative_threshold_rather_than_panicking() {
190        // A negative threshold is normal, so `is_normal` alone lets it reach
191        // `Duration::div_f64`, which panics on a negative divisor. The bound builders take
192        // the threshold straight from the caller, where it can be computed.
193        assert_eq!(bound(8, 4.0, -0.5).time_limit(), None);
194        assert_eq!(bound(8, 4.0, f32::NEG_INFINITY).time_limit(), None);
195    }
196
197    #[test]
198    fn vec_time_limit_takes_the_roofline_max_not_min() {
199        // Two simultaneous resource bounds (e.g. compute vs memory): the achievable floor
200        // is the *slower* one, so the reduction must be `max`. `min` would pick the
201        // unreachable 1s and the short-circuit would never fire.
202        let compute = bound(8, 4.0, 1.0); // 2s
203        let memory = bound(8, 8.0, 1.0); // 1s
204        let limit = vec![compute, memory].time_limit();
205        assert_eq!(limit, Some(Duration::from_secs(2)));
206    }
207
208    #[test]
209    fn vec_time_limit_skips_non_normal_bounds_and_is_none_when_empty() {
210        // A non-normal bound is filtered out rather than poisoning the reduction.
211        let limit = vec![bound(8, 0.0, 1.0), bound(8, 4.0, 1.0)].time_limit();
212        assert_eq!(limit, Some(Duration::from_secs(2)));
213
214        assert_eq!(Vec::<AutotuneBound>::new().time_limit(), None);
215    }
216
217    #[test]
218    fn bounds_time_limit_adds_launch_overhead() {
219        let bounds = Bounds {
220            bounds: vec![bound(8, 4.0, 1.0)], // 2s
221            launch_overhead: Duration::from_millis(500),
222        };
223        assert_eq!(bounds.time_limit(), Some(Duration::from_millis(2500)));
224    }
225
226    #[test]
227    fn a_level_tightens_the_limit_as_it_rises() {
228        let threshold = |level| Thresholds::for_level(&level).compute;
229        assert!(threshold(AutotuneLevel::Minimal) < threshold(AutotuneLevel::Balanced));
230        assert!(threshold(AutotuneLevel::Balanced) < threshold(AutotuneLevel::Extensive));
231        assert_eq!(
232            Thresholds::for_level(&AutotuneLevel::Full),
233            Thresholds::UNBOUNDED,
234            "full measures everything"
235        );
236    }
237
238    #[test]
239    fn an_unbounded_threshold_gives_no_time_limit() {
240        // What lets `Full` be handed to a bound like any other threshold.
241        let unbounded = bound(8, 4.0, Thresholds::UNBOUNDED.compute);
242        assert_eq!(unbounded.time_limit(), None);
243    }
244
245    #[test]
246    fn the_default_threshold_is_the_roofline_itself() {
247        assert_eq!(Thresholds::default(), Thresholds::uniform(1.0));
248        assert_eq!(Thresholds::default().compute, 1.0);
249    }
250
251    #[test]
252    fn bounds_time_limit_is_none_without_usable_bounds() {
253        // No usable bound means no limit at all — the launch overhead is not a limit on
254        // its own, so the short-circuit stays disabled.
255        let bounds = Bounds {
256            bounds: vec![],
257            launch_overhead: Duration::from_millis(500),
258        };
259        assert_eq!(bounds.time_limit(), None);
260    }
261}
cubecl_runtime/tune/bounds_generator.rs

cubecl_runtime/tune/
bounds_generator.rs
