[PATCH 2/3] rust: kernel: add bench

[Daniel Gomez]

From: Daniel Gomez <da.gomez@xxxxxxxxxxx>

Rust bench for sampled benchmarking with statistics.

The XArray benchmark will be the first user. The runner is a module so
other Rust benchmarks such as find_bit_benchmark_rust can share it. It
can also be extended with percentiles in the future.

Assisted-by: LLM
Signed-off-by: Daniel Gomez <da.gomez@xxxxxxxxxxx>
---
 MAINTAINERS          |   1 +
 rust/kernel/bench.rs | 173 +++++++++++++++++++++++++++++++++++++++++++++++++++
 rust/kernel/lib.rs   |   1 +
 3 files changed, 175 insertions(+)

diff --git a/MAINTAINERS b/MAINTAINERS
index c40a254c35d3d..b8bdfe9e22226 100644
--- a/MAINTAINERS
+++ b/MAINTAINERS
@@ -29690,6 +29690,7 @@ W:	https://rust-for-linux.com
 B:	https://github.com/Rust-for-Linux/linux/issues
 C:	https://rust-for-linux.zulipchat.com
 T:	git git://git.kernel.org/pub/scm/linux/kernel/git/da.gomez/linux.git rxarray-next
+F:	rust/kernel/bench.rs
 F:	rust/kernel/rxarray.rs
 
 XBOX DVD IR REMOTE
diff --git a/rust/kernel/bench.rs b/rust/kernel/bench.rs
new file mode 100644
index 0000000000000..d1c85c49c7969
--- /dev/null
+++ b/rust/kernel/bench.rs
@@ -0,0 +1,173 @@
+// SPDX-License-Identifier: GPL-2.0
+
+//! Sampled benchmarks with in-kernel statistics.
+//!
+//! Each sample times `iterations` of a workload. A [`Row`] prints the minimum, median, maximum and
+//! mean of the sample times in nanoseconds. The caller owns the timer and the printing:
+//!
+//! ```ignore
+//! let mut bench = Bencher::new(samples, entries)?;
+//! pr_info!("{samples} samples x {entries} entries, ns per sample:\n");
+//! pr_info!("{}\n", bench::Heading);
+//! pr_info!("{}\n", bench.run("store", XArray::new, store));
+//! pr_info!("total runtime {}\n", bench.runtime());
+//! ```
+
+use crate::{
+    fmt,
+    prelude::*,
+    time::{Delta, Instant, Monotonic}, //
+};
+
+/// The column headings of a table of [`Row`]s, in the same columns.
+pub struct Heading;
+
+impl fmt::Display for Heading {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.write_fmt(fmt!(
+            "{:<18} {:>12} {:>12} {:>12} {:>12} {:>12}",
+            "benchmark",
+            "min",
+            "median",
+            "max",
+            "mean",
+            "runtime"
+        ))
+    }
+}
+
+/// A wall time.
+pub struct Runtime(pub Delta);
+
+impl fmt::Display for Runtime {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let ms = self.0.as_millis();
+        let secs = ms / 1000;
+        let mins = secs / 60;
+        let hours = mins / 60;
+        let width = f.width().unwrap_or(0);
+        if hours > 0 {
+            let w = width.saturating_sub(4);
+            write!(f, "{hours:w$}h{:02}m", mins % 60)
+        } else if mins > 0 {
+            let w = width.saturating_sub(4);
+            write!(f, "{mins:w$}m{:02}s", secs % 60)
+        } else {
+            let w = width.saturating_sub(5);
+            write!(f, "{secs:w$}.{:03}s", ms % 1000)
+        }
+    }
+}
+
+/// Stats across samples, in nanoseconds.
+pub struct Stats {
+    /// The fastest sample.
+    pub min: i64,
+    /// The middle sample.
+    pub median: i64,
+    /// The slowest sample.
+    pub max: i64,
+    /// The mean of the samples, rounded down.
+    pub mean: i64,
+}
+
+impl Stats {
+    /// Computes the statistics of the non-empty `samples`, sorting them in place.
+    pub fn new(samples: &mut [i64]) -> Self {
+        samples.sort_unstable();
+        let len = samples.len();
+        Self {
+            min: samples[0],
+            median: (samples[(len - 1) / 2] + samples[len / 2]) / 2,
+            max: samples[len - 1],
+            mean: samples.iter().sum::<i64>() / len as i64,
+        }
+    }
+}
+
+/// One row of the table: the benchmark's name, its statistics and its wall time, in the columns
+/// of [`Heading`].
+pub struct Row<'a> {
+    name: &'a str,
+    stats: Stats,
+    runtime: Runtime,
+}
+
+impl fmt::Display for Row<'_> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.write_fmt(fmt!(
+            "{:<18} {:>12} {:>12} {:>12} {:>12} {:>12}",
+            self.name,
+            self.stats.min,
+            self.stats.median,
+            self.stats.max,
+            self.stats.mean,
+            self.runtime
+        ))
+    }
+}
+
+/// Benchmark runner.
+pub struct Bencher {
+    iterations: usize,
+    timings: KVVec<i64>,
+    runtime: Delta,
+}
+
+impl Bencher {
+    /// Allocates one timing slot per sample up front, outside the timed loops.
+    ///
+    /// Returns `EINVAL` if `samples` or `iterations` is 0.
+    pub fn new(samples: usize, iterations: usize) -> Result<Self> {
+        if samples == 0 || iterations == 0 {
+            return Err(EINVAL);
+        }
+        Ok(Self {
+            iterations,
+            timings: KVVec::from_elem(0, samples, GFP_KERNEL)?,
+            runtime: Delta::ZERO,
+        })
+    }
+
+    /// Runs `bench` on a fresh `setup` value once per sample and returns the table row of `name`.
+    ///
+    /// `bench` returns the [`Delta`] of the window it timed.
+    pub fn run<'a, A>(
+        &mut self,
+        name: &'a str,
+        setup: impl Fn() -> A,
+        bench: impl Fn(A, usize) -> Delta,
+    ) -> Row<'a> {
+        let start = Instant::<Monotonic>::now();
+        for ns in &mut self.timings {
+            *ns = bench(setup(), self.iterations).as_nanos();
+        }
+        let elapsed = start.elapsed();
+        self.runtime += elapsed;
+
+        Row {
+            name,
+            stats: Stats::new(&mut self.timings),
+            runtime: Runtime(elapsed),
+        }
+    }
+
+    /// The wall time of every run so far.
+    pub fn runtime(&self) -> Runtime {
+        Runtime(self.runtime)
+    }
+}
+
+#[macros::kunit_tests(rust_bench)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn stats() {
+        let odd = Stats::new(&mut [5, 1, 9]);
+        assert_eq!((odd.min, odd.median, odd.max, odd.mean), (1, 5, 9, 5));
+        // Median check for even samples.
+        let med = Stats::new(&mut [4, 1, 9, 5]);
+        assert_eq!(med.median, 4);
+    }
+}
diff --git a/rust/kernel/lib.rs b/rust/kernel/lib.rs
index 1e3c8d3051e53..95162ab0c3f13 100644
--- a/rust/kernel/lib.rs
+++ b/rust/kernel/lib.rs
@@ -48,6 +48,7 @@
 #[cfg(CONFIG_AUXILIARY_BUS)]
 pub mod auxiliary;
 pub mod bitfield;
+pub mod bench;
 pub mod bitmap;
 pub mod bits;
 #[cfg(CONFIG_BLOCK)]

-- 
2.55.0