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 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(autotune_persistence, 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(autotune_persistence, 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(std_io, 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
103impl Default for Thresholds {
104    /// The roofline itself: a candidate is expected to reach 100% of the modeled peak, which
105    /// is the only fraction that needs no justification.
106    fn default() -> Self {
107        Self::uniform(1.0)
108    }
109}
110
111/// Standardizes the creation of compute and memory [`AutotuneBound`]s.
112pub fn calculate_bounds(
113    work: Work,
114    thresholds: Thresholds,
115    compute_throughput: &ThroughputValue,
116    memory_throughput: &ThroughputValue,
117    memory_key: &ThroughputKey,
118) -> Vec<AutotuneBound> {
119    alloc::vec![
120        AutotuneBound {
121            resource: ResourceBound {
122                amount: work.compute_ops,
123                peak_per_s: compute_throughput.ops_per_s(),
124            },
125            threshold: thresholds.compute,
126        },
127        AutotuneBound {
128            resource: ResourceBound {
129                amount: work.bytes,
130                peak_per_s: memory_throughput.bytes_per_s(memory_key),
131            },
132            threshold: thresholds.memory,
133        },
134    ]
135}
136
137impl TimeBound for AutotuneBound {
138    fn time_limit(&self) -> Option<Duration> {
139        if !self.threshold.is_normal() {
140            return None;
141        }
142        self.resource
143            .time_at_peak()
144            .map(|limit| limit.div_f64(self.threshold as f64))
145    }
146}
147
148impl<B: TimeBound> TimeBound for Vec<B> {
149    fn time_limit(&self) -> Option<Duration> {
150        self.iter().filter_map(|b| b.time_limit()).max()
151    }
152}
153
154impl TimeBound for Bounds {
155    fn time_limit(&self) -> Option<Duration> {
156        self.bounds
157            .time_limit()
158            .map(|limit| limit + self.launch_overhead)
159    }
160}
161
162#[cfg(test)]
163mod tests {
164    use crate::throughput::ThroughputMode;
165
166    use super::*;
167    use alloc::vec;
168
169    fn bound(ops_count: usize, throughput: f64, threshold: f32) -> AutotuneBound {
170        AutotuneBound {
171            resource: ResourceBound {
172                amount: ops_count,
173                peak_per_s: throughput,
174            },
175            threshold,
176        }
177    }
178
179    #[test]
180    fn time_limit_is_ops_over_throughput_scaled_by_threshold() {
181        // (8 ops / 4 ops/s) / 0.5 = 4s. Powers of two keep the f64 math exact.
182        let limit = bound(8, 4.0, 0.5).time_limit();
183        assert_eq!(limit, Some(Duration::from_secs(4)));
184    }
185
186    #[test]
187    fn time_limit_is_none_when_inputs_are_not_normal() {
188        // A zero/NaN/inf throughput or a zero threshold would divide by zero or blow up,
189        // so the bound disables the short-circuit instead of producing a garbage limit.
190        assert_eq!(bound(8, 0.0, 0.5).time_limit(), None);
191        assert_eq!(bound(8, f64::NAN, 0.5).time_limit(), None);
192        assert_eq!(bound(8, f64::INFINITY, 0.5).time_limit(), None);
193        assert_eq!(bound(8, 4.0, 0.0).time_limit(), None);
194    }
195
196    #[test]
197    fn vec_time_limit_takes_the_roofline_max_not_min() {
198        // Two simultaneous resource bounds (e.g. compute vs memory): the achievable floor
199        // is the *slower* one, so the reduction must be `max`. `min` would pick the
200        // unreachable 1s and the short-circuit would never fire.
201        let compute = bound(8, 4.0, 1.0); // 2s
202        let memory = bound(8, 8.0, 1.0); // 1s
203        let limit = vec![compute, memory].time_limit();
204        assert_eq!(limit, Some(Duration::from_secs(2)));
205    }
206
207    #[test]
208    fn vec_time_limit_skips_non_normal_bounds_and_is_none_when_empty() {
209        // A non-normal bound is filtered out rather than poisoning the reduction.
210        let limit = vec![bound(8, 0.0, 1.0), bound(8, 4.0, 1.0)].time_limit();
211        assert_eq!(limit, Some(Duration::from_secs(2)));
212
213        assert_eq!(Vec::<AutotuneBound>::new().time_limit(), None);
214    }
215
216    #[test]
217    fn bounds_time_limit_adds_launch_overhead() {
218        let bounds = Bounds {
219            bounds: vec![bound(8, 4.0, 1.0)], // 2s
220            launch_overhead: Duration::from_millis(500),
221        };
222        assert_eq!(bounds.time_limit(), Some(Duration::from_millis(2500)));
223    }
224
225    #[test]
226    fn calculate_bounds_applies_a_threshold_per_resource() {
227        let work = Work {
228            compute_ops: 8,
229            bytes: 16,
230        };
231        let thresholds = Thresholds {
232            compute: 0.5,
233            memory: 1.0,
234        };
235        let key = ThroughputKey {
236            mode: ThroughputMode::Memory,
237        };
238
239        let bounds = calculate_bounds(
240            work,
241            thresholds,
242            &ThroughputValue::ZERO,
243            &ThroughputValue::ZERO,
244            &key,
245        );
246
247        assert_eq!(bounds[0].resource.amount, 8);
248        assert_eq!(bounds[0].threshold, 0.5);
249        assert_eq!(bounds[1].resource.amount, 16);
250        assert_eq!(bounds[1].threshold, 1.0);
251    }
252
253    #[test]
254    fn the_default_threshold_is_the_roofline_itself() {
255        assert_eq!(Thresholds::default(), Thresholds::uniform(1.0));
256        assert_eq!(Thresholds::default().compute, 1.0);
257    }
258
259    #[test]
260    fn bounds_time_limit_is_none_without_usable_bounds() {
261        // No usable bound means no limit at all — the launch overhead is not a limit on
262        // its own, so the short-circuit stays disabled.
263        let bounds = Bounds {
264            bounds: vec![],
265            launch_overhead: Duration::from_millis(500),
266        };
267        assert_eq!(bounds.time_limit(), None);
268    }
269}
cubecl_runtime/tune/bounds_generator.rs

cubecl_runtime/tune/
bounds_generator.rs
