|
| 1 | +use alloc::vec::Vec; |
| 2 | +use core::time::Duration; |
| 3 | +use cubecl_common::benchmark::{BenchmarkComputations, BenchmarkDurations}; |
| 4 | +use cubecl_common::profile::TimingMethod; |
| 5 | + |
| 6 | +/// The timings collected for one candidate, plus the small amount of state needed to decide |
| 7 | +/// whether it is still worth sampling. |
| 8 | +/// |
| 9 | +/// `BenchmarkComputations::score` is not usable while tuning is in flight: at a single sample the |
| 10 | +/// variance is zero and reads as perfect stability. This tracks the sample count explicitly so |
| 11 | +/// elimination can require evidence instead of inferring it from a degenerate variance. |
| 12 | +#[derive(Debug, Default)] |
| 13 | +pub(crate) struct SampleSet { |
| 14 | + durations: Vec<Duration>, |
| 15 | + stalled: u8, |
| 16 | +} |
| 17 | + |
| 18 | +impl SampleSet { |
| 19 | + pub(crate) fn push(&mut self, duration: Duration) { |
| 20 | + let before = self.best(); |
| 21 | + self.durations.push(duration); |
| 22 | + |
| 23 | + // This push ages the biased first sample out, so the two bests are taken over different |
| 24 | + // samples and their difference says nothing about progress. |
| 25 | + if self.durations.len() == DISCARD_THRESHOLD { |
| 26 | + self.stalled = 0; |
| 27 | + return; |
| 28 | + } |
| 29 | + |
| 30 | + let improved = match (before, self.best()) { |
| 31 | + (Some(before), Some(after)) => after < before.mul_f64(1.0 - CONVERGENCE_EPSILON), |
| 32 | + _ => true, |
| 33 | + }; |
| 34 | + |
| 35 | + self.stalled = if improved { |
| 36 | + 0 |
| 37 | + } else { |
| 38 | + self.stalled.saturating_add(1) |
| 39 | + }; |
| 40 | + } |
| 41 | + |
| 42 | + pub(crate) fn len(&self) -> usize { |
| 43 | + self.durations.len() |
| 44 | + } |
| 45 | + |
| 46 | + pub(crate) fn is_empty(&self) -> bool { |
| 47 | + self.durations.is_empty() |
| 48 | + } |
| 49 | + |
| 50 | + /// The first sample follows a single warmup, so it still carries allocation and clock ramp |
| 51 | + /// costs. It is dropped once enough samples remain without it. |
| 52 | + fn reliable(&self) -> &[Duration] { |
| 53 | + if self.durations.len() >= DISCARD_THRESHOLD { |
| 54 | + &self.durations[1..] |
| 55 | + } else { |
| 56 | + &self.durations |
| 57 | + } |
| 58 | + } |
| 59 | + |
| 60 | + pub(crate) fn best(&self) -> Option<Duration> { |
| 61 | + self.reliable().iter().min().copied() |
| 62 | + } |
| 63 | + |
| 64 | + pub(crate) fn converged(&self) -> bool { |
| 65 | + self.stalled >= CONVERGENCE_ROUNDS |
| 66 | + } |
| 67 | + |
| 68 | + /// Whether enough samples independently landed under `limit` to trust a short circuit. |
| 69 | + /// |
| 70 | + /// Unlike [`Self::best`] this counts the warmup-biased first sample, deliberately: the bias |
| 71 | + /// is toward being slower, so clearing the limit despite it is the conservative direction. |
| 72 | + /// Excluding it would also make a short circuit impossible during the first pass, where it |
| 73 | + /// is the only sample there is. |
| 74 | + pub(crate) fn confirmed_under(&self, limit: Duration, required: usize) -> bool { |
| 75 | + self.durations.iter().filter(|d| **d <= limit).count() >= required |
| 76 | + } |
| 77 | + |
| 78 | + /// Whether any sample at all landed under `limit` — the cheap gate before a candidate is |
| 79 | + /// worth spending confirmation samples on. |
| 80 | + pub(crate) fn any_under(&self, limit: Duration) -> bool { |
| 81 | + self.durations.iter().any(|d| *d <= limit) |
| 82 | + } |
| 83 | + |
| 84 | + pub(crate) fn computation(&self, method: TimingMethod) -> BenchmarkComputations { |
| 85 | + BenchmarkComputations::new(&BenchmarkDurations::from_durations( |
| 86 | + method, |
| 87 | + self.reliable().to_vec(), |
| 88 | + )) |
| 89 | + } |
| 90 | +} |
| 91 | + |
| 92 | +/// Total samples required before the warmup-biased first sample is discarded. |
| 93 | +const DISCARD_THRESHOLD: usize = 3; |
| 94 | +/// Relative improvement below which a new sample counts as no progress. |
| 95 | +const CONVERGENCE_EPSILON: f64 = 0.02; |
| 96 | +/// Consecutive non-improving samples before a candidate is considered converged. |
| 97 | +const CONVERGENCE_ROUNDS: u8 = 2; |
| 98 | + |
| 99 | +#[cfg(test)] |
| 100 | +mod tests { |
| 101 | + use super::*; |
| 102 | + |
| 103 | + fn set(durations: impl IntoIterator<Item = u64>) -> SampleSet { |
| 104 | + let mut set = SampleSet::default(); |
| 105 | + for millis in durations { |
| 106 | + set.push(Duration::from_millis(millis)); |
| 107 | + } |
| 108 | + set |
| 109 | + } |
| 110 | + |
| 111 | + #[test] |
| 112 | + fn keeps_every_sample_below_the_discard_threshold() { |
| 113 | + assert_eq!( |
| 114 | + set([10, 20]).reliable(), |
| 115 | + &[Duration::from_millis(10), Duration::from_millis(20)] |
| 116 | + ); |
| 117 | + } |
| 118 | + |
| 119 | + #[test] |
| 120 | + fn drops_the_warmup_biased_first_sample_once_enough_remain() { |
| 121 | + // The first sample is the fastest here, so dropping it has to move `best` upward: |
| 122 | + // that is the bias being removed, not data being lost. |
| 123 | + let set = set([5, 20, 21]); |
| 124 | + assert_eq!(set.reliable().len(), 2); |
| 125 | + assert_eq!(set.best(), Some(Duration::from_millis(20))); |
| 126 | + } |
| 127 | + |
| 128 | + #[test] |
| 129 | + fn converges_after_consecutive_non_improving_samples() { |
| 130 | + // Gains under the 2% epsilon do not count as progress. The third push only shifts the |
| 131 | + // window, so two further flat samples are what trips convergence. |
| 132 | + let mut set = set([100, 20, 20]); |
| 133 | + assert!(!set.converged()); |
| 134 | + set.push(Duration::from_millis(20)); |
| 135 | + assert!(!set.converged()); |
| 136 | + set.push(Duration::from_millis(20)); |
| 137 | + assert!(set.converged()); |
| 138 | + } |
| 139 | + |
| 140 | + #[test] |
| 141 | + fn aging_out_the_biased_sample_does_not_count_as_a_stall() { |
| 142 | + // The first sample is the fastest, so dropping it raises `best`. That is the window |
| 143 | + // moving, not the candidate failing to improve, and must not count toward convergence. |
| 144 | + let mut set = set([5, 20, 21]); |
| 145 | + assert!(!set.converged()); |
| 146 | + set.push(Duration::from_millis(20)); |
| 147 | + assert!(!set.converged()); |
| 148 | + set.push(Duration::from_millis(20)); |
| 149 | + assert!(set.converged()); |
| 150 | + } |
| 151 | + |
| 152 | + #[test] |
| 153 | + fn a_real_improvement_resets_convergence() { |
| 154 | + let mut set = set([20, 20, 20, 20, 20]); |
| 155 | + assert!(set.converged()); |
| 156 | + set.push(Duration::from_millis(5)); |
| 157 | + assert!(!set.converged()); |
| 158 | + } |
| 159 | + |
| 160 | + #[test] |
| 161 | + fn short_circuit_needs_independent_confirmations() { |
| 162 | + let limit = Duration::from_millis(10); |
| 163 | + // One fast sample among slow ones is exactly the lucky measurement that must not |
| 164 | + // commit a decision to the persistent cache on its own. |
| 165 | + assert!(!set([5, 50]).confirmed_under(limit, 2)); |
| 166 | + assert!(set([5, 50, 6]).confirmed_under(limit, 2)); |
| 167 | + } |
| 168 | + |
| 169 | + #[test] |
| 170 | + fn computation_is_built_from_the_reliable_samples_only() { |
| 171 | + let set = set([1, 20, 30]); |
| 172 | + let computation = set.computation(TimingMethod::System); |
| 173 | + assert_eq!(computation.min, Duration::from_millis(20)); |
| 174 | + assert_eq!(computation.max, Duration::from_millis(30)); |
| 175 | + } |
| 176 | + |
| 177 | + #[test] |
| 178 | + fn empty_set_has_no_best() { |
| 179 | + let set = SampleSet::default(); |
| 180 | + assert!(set.is_empty()); |
| 181 | + assert_eq!(set.best(), None); |
| 182 | + assert_eq!(set.len(), 0); |
| 183 | + } |
| 184 | +} |
0 commit comments