cubecl_runtime/tune/
bounds_generator.rs

1use core::time::Duration;
2
3use alloc::vec::Vec;
4
5use crate::tune::TuneInputs;
6
7/// A set of [`AutotuneBound`]s for a given key and reference inputs, with a launch overhead.
8pub struct Bounds<B: TimeBound> {
9    /// The bounds for autotuning.
10    pub bounds: Vec<B>,
11    /// The launch overhead for autotuning.
12    pub launch_overhead: Duration,
13}
14
15/// Produces a set of [`AutotuneBound`]s for a given key and reference inputs.
16#[diagnostic::on_unimplemented(
17    message = "`{Self}` is not a valid bounds generator",
18    label = "invalid bounds generator"
19)]
20pub trait BoundsGenerator<K, I: TuneInputs, B: TimeBound>: Send + Sync + 'static {
21    /// Generate a set of bounds for a given key and reference inputs.
22    fn generate<'a>(&self, key: &K, inputs: &I::At<'a>) -> Bounds<B>;
23}
24
25/// `Fn(&K, &A) -> Bounds<B>` acts as a [`BoundsGenerator`] when `A` is an owned type. For
26/// multi-input kernels, `A` is a tuple that the closure destructures internally.
27impl<K, Func, A, B: TimeBound> BoundsGenerator<K, A, B> for Func
28where
29    A: Clone + Send + Sync + 'static,
30    K: 'static,
31    Func: Send + Sync + 'static + Fn(&K, &A) -> Bounds<B>,
32{
33    #[inline]
34    fn generate<'a>(&self, key: &K, inputs: &<A as TuneInputs>::At<'a>) -> Bounds<B> {
35        (self)(key, inputs)
36    }
37}
38
39/// A calculator that determines the time limit for autotune bounds.
40pub trait TimeBound {
41    /// Returns the time limit for autotune bounds.
42    fn time_limit(&self) -> Option<Duration>;
43}
44
45/// A bound for autotuning a throughput kernel, specifying the key, threshold, and number of operations.
46pub struct AutotuneBound {
47    /// Peak throughput of the reference kernel, in ops (or bytes) per second.
48    pub throughput: f64,
49    /// The threshold for this bound, over which the kernel will be considered accurate.
50    pub threshold: f32,
51    /// The number of operations the kernel will run.
52    pub ops_count: usize,
53}
54
55impl TimeBound for AutotuneBound {
56    fn time_limit(&self) -> Option<Duration> {
57        if self.throughput.is_normal() && self.threshold.is_normal() {
58            Some(Duration::from_secs_f64(
59                (self.ops_count as f64 / self.throughput) / self.threshold as f64,
60            ))
61        } else {
62            None
63        }
64    }
65}
66
67impl<B: TimeBound> TimeBound for Vec<B> {
68    fn time_limit(&self) -> Option<Duration> {
69        self.iter().filter_map(|b| b.time_limit()).max()
70    }
71}
72
73impl<B: TimeBound> TimeBound for Bounds<B> {
74    fn time_limit(&self) -> Option<Duration> {
75        self.bounds
76            .time_limit()
77            .map(|limit| limit + self.launch_overhead)
78    }
79}
80
81#[cfg(test)]
82mod tests {
83    use super::*;
84    use alloc::vec;
85
86    fn bound(ops_count: usize, throughput: f64, threshold: f32) -> AutotuneBound {
87        AutotuneBound {
88            throughput,
89            threshold,
90            ops_count,
91        }
92    }
93
94    #[test]
95    fn time_limit_is_ops_over_throughput_scaled_by_threshold() {
96        // (8 ops / 4 ops/s) / 0.5 = 4s. Powers of two keep the f64 math exact.
97        let limit = bound(8, 4.0, 0.5).time_limit();
98        assert_eq!(limit, Some(Duration::from_secs(4)));
99    }
100
101    #[test]
102    fn time_limit_is_none_when_inputs_are_not_normal() {
103        // A zero/NaN/inf throughput or a zero threshold would divide by zero or blow up,
104        // so the bound disables the short-circuit instead of producing a garbage limit.
105        assert_eq!(bound(8, 0.0, 0.5).time_limit(), None);
106        assert_eq!(bound(8, f64::NAN, 0.5).time_limit(), None);
107        assert_eq!(bound(8, f64::INFINITY, 0.5).time_limit(), None);
108        assert_eq!(bound(8, 4.0, 0.0).time_limit(), None);
109    }
110
111    #[test]
112    fn vec_time_limit_takes_the_roofline_max_not_min() {
113        // Two simultaneous resource bounds (e.g. compute vs memory): the achievable floor
114        // is the *slower* one, so the reduction must be `max`. `min` would pick the
115        // unreachable 1s and the short-circuit would never fire.
116        let compute = bound(8, 4.0, 1.0); // 2s
117        let memory = bound(8, 8.0, 1.0); // 1s
118        let limit = vec![compute, memory].time_limit();
119        assert_eq!(limit, Some(Duration::from_secs(2)));
120    }
121
122    #[test]
123    fn vec_time_limit_skips_non_normal_bounds_and_is_none_when_empty() {
124        // A non-normal bound is filtered out rather than poisoning the reduction.
125        let limit = vec![bound(8, 0.0, 1.0), bound(8, 4.0, 1.0)].time_limit();
126        assert_eq!(limit, Some(Duration::from_secs(2)));
127
128        assert_eq!(Vec::<AutotuneBound>::new().time_limit(), None);
129    }
130
131    #[test]
132    fn bounds_time_limit_adds_launch_overhead() {
133        let bounds = Bounds {
134            bounds: vec![bound(8, 4.0, 1.0)], // 2s
135            launch_overhead: Duration::from_millis(500),
136        };
137        assert_eq!(bounds.time_limit(), Some(Duration::from_millis(2500)));
138    }
139
140    #[test]
141    fn bounds_time_limit_is_none_without_usable_bounds() {
142        // No usable bound means no limit at all — the launch overhead is not a limit on
143        // its own, so the short-circuit stays disabled.
144        let bounds = Bounds::<AutotuneBound> {
145            bounds: vec![],
146            launch_overhead: Duration::from_millis(500),
147        };
148        assert_eq!(bounds.time_limit(), None);
149    }
150}
cubecl_runtime/tune/bounds_generator.rs

cubecl_runtime/tune/
bounds_generator.rs
