Skip to main content

hrx/
benchmark_statistics.rs

1//! Reusable benchmark distributions and synchronized stage timing.
2
3use crate::{Error, Result, Stream};
4use std::{collections::BTreeMap, time::Instant};
5
6/// Select a nearest-rank percentile from sorted, nonempty samples.
7///
8/// `percent` must be in `1..=100`. This helper intentionally asserts its
9/// preconditions: benchmark collection owns both sorting and sample counts.
10#[must_use]
11pub fn percentile(sorted: &[f64], percent: usize) -> f64 {
12    assert!(!sorted.is_empty() && (1..=100).contains(&percent));
13    sorted[(sorted.len() * percent).div_ceil(100) - 1]
14}
15
16/// Millisecond distribution measured around submission and completion.
17#[derive(Clone, Copy, Debug, serde::Serialize)]
18pub struct Distribution {
19    /// Nearest-rank 50th percentile.
20    pub median_ms: f64,
21    /// Nearest-rank 95th percentile.
22    pub p95_ms: f64,
23}
24
25impl Distribution {
26    /// Summarize finite, nonempty millisecond samples.
27    pub fn from_samples(mut samples: Vec<f64>) -> Result<Self> {
28        if samples.is_empty()
29            || samples
30                .iter()
31                .any(|sample| !sample.is_finite() || *sample < 0.0)
32        {
33            return Err(Error::Message(
34                "benchmark samples must be finite, nonnegative, and nonempty".into(),
35            ));
36        }
37        samples.sort_by(f64::total_cmp);
38        Ok(Self {
39            median_ms: percentile(&samples, 50),
40            p95_ms: percentile(&samples, 95),
41        })
42    }
43}
44
45/// Named synchronized GPU stage samples collected by a benchmark.
46#[derive(Default)]
47pub struct StageTimings {
48    samples: BTreeMap<String, Vec<f64>>,
49}
50
51impl StageTimings {
52    /// Create an empty collector.
53    #[must_use]
54    pub fn new() -> Self {
55        Self::default()
56    }
57
58    /// Synchronize, run one stage, synchronize again, and record elapsed milliseconds.
59    pub fn measure<T>(
60        &mut self,
61        stream: &mut Stream,
62        name: impl Into<String>,
63        stage: impl FnOnce(&mut Stream) -> Result<T>,
64    ) -> Result<T> {
65        stream.synchronize()?;
66        let start = Instant::now();
67        let value = stage(stream)?;
68        stream.synchronize()?;
69        self.samples
70            .entry(name.into())
71            .or_default()
72            .push(start.elapsed().as_secs_f64() * 1_000.0);
73        Ok(value)
74    }
75
76    /// Add an already measured nonnegative millisecond sample.
77    pub fn push(&mut self, name: impl Into<String>, milliseconds: f64) -> Result<()> {
78        if !milliseconds.is_finite() || milliseconds < 0.0 {
79            return Err(Error::Message("invalid benchmark duration".into()));
80        }
81        self.samples
82            .entry(name.into())
83            .or_default()
84            .push(milliseconds);
85        Ok(())
86    }
87
88    /// Summarize each named stage in lexical order.
89    pub fn distributions(&self) -> Result<BTreeMap<String, Distribution>> {
90        self.samples
91            .iter()
92            .map(|(name, samples)| Ok((name.clone(), Distribution::from_samples(samples.clone())?)))
93            .collect()
94    }
95
96    /// Sum the samples for every stage, in milliseconds.
97    #[must_use]
98    pub fn totals_ms(&self) -> BTreeMap<String, f64> {
99        self.samples
100            .iter()
101            .map(|(name, samples)| (name.clone(), samples.iter().sum()))
102            .collect()
103    }
104
105    /// Drain the samples and return each stage's total milliseconds.
106    pub fn take_totals_ms(&mut self) -> BTreeMap<String, f64> {
107        std::mem::take(&mut self.samples)
108            .into_iter()
109            .map(|(name, samples)| (name, samples.into_iter().sum()))
110            .collect()
111    }
112
113    /// Whether no stage samples have been recorded.
114    #[must_use]
115    pub fn is_empty(&self) -> bool {
116        self.samples.is_empty()
117    }
118}
119
120#[cfg(test)]
121mod tests {
122    use super::*;
123    #[test]
124    fn nearest_rank_includes_the_correct_order_statistic() {
125        let values: Vec<_> = (1..=100).map(f64::from).collect();
126        assert_eq!(percentile(&values, 50), 50.0);
127        assert_eq!(percentile(&values, 95), 95.0);
128        assert_eq!(percentile(&values[..50], 50), 25.0);
129        assert_eq!(percentile(&values[..50], 95), 48.0);
130        assert_eq!(percentile(&[7.0], 95), 7.0);
131    }
132
133    #[test]
134    fn distributions_reject_bad_samples() {
135        assert!(Distribution::from_samples(vec![]).is_err());
136        assert!(Distribution::from_samples(vec![f64::NAN]).is_err());
137        let value = Distribution::from_samples((1..=100).map(f64::from).collect()).unwrap();
138        assert_eq!(value.median_ms, 50.0);
139        assert_eq!(value.p95_ms, 95.0);
140    }
141
142    #[test]
143    fn stage_totals_can_be_drained() {
144        let mut timings = StageTimings::new();
145        timings.push("load", 1.25).unwrap();
146        timings.push("load", 2.75).unwrap();
147        timings.push("store", 3.0).unwrap();
148        assert_eq!(timings.totals_ms()["load"], 4.0);
149        assert_eq!(timings.take_totals_ms()["store"], 3.0);
150        assert!(timings.is_empty());
151    }
152}