[PATCH 3/3] lib/xarray_benchmark_rust: add module

[Daniel Gomez <da.gomez@kernel.org>]

From: Daniel Gomez <da.gomez@samsung.com>

Benchmark XArray Rust APIs using `kernel::bench`.

Assisted-by: LLM
Signed-off-by: Daniel Gomez <da.gomez@samsung.com>
---
 MAINTAINERS                  |   1 +
 lib/Kconfig.debug            |  10 ++++
 lib/Makefile                 |   1 +
 lib/xarray_benchmark_rust.rs | 117 +++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 129 insertions(+)

diff --git a/MAINTAINERS b/MAINTAINERS
index b8bdfe9e22226..f6e45a88475b0 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:	lib/xarray_benchmark_rust.rs
 F:	rust/kernel/bench.rs
 F:	rust/kernel/rxarray.rs
 
diff --git a/lib/Kconfig.debug b/lib/Kconfig.debug
index 134b15a44625e..6e008f657cd4d 100644
--- a/lib/Kconfig.debug
+++ b/lib/Kconfig.debug
@@ -2667,6 +2667,16 @@ config FIND_BIT_BENCHMARK_RUST
 
 	  If unsure, say N.
 
+config XARRAY_BENCHMARK_RUST
+	tristate "Benchmark the XArray Rust APIs"
+	depends on RUST
+	help
+	  This builds the "xarray_benchmark_rust" module. It runs the same
+	  workloads through the Rust abstraction of the C XArray (kernel::xarray),
+	  and through the Rust XArray implementation (kernel::rxarray).
+
+	  If unsure, say N.
+
 config TEST_FIRMWARE
 	tristate "Test firmware loading via userspace interface"
 	depends on FW_LOADER
diff --git a/lib/Makefile b/lib/Makefile
index dfab958327c5c..9401f5146592e 100644
--- a/lib/Makefile
+++ b/lib/Makefile
@@ -66,6 +66,7 @@ obj-y += kstrtox.o
 obj-$(CONFIG_FIND_BIT_BENCHMARK) += find_bit_benchmark.o
 obj-$(CONFIG_REGION_ALLOC_BENCHMARK) += region_alloc_benchmark.o
 obj-$(CONFIG_FIND_BIT_BENCHMARK_RUST) += find_bit_benchmark_rust.o
+obj-$(CONFIG_XARRAY_BENCHMARK_RUST) += xarray_benchmark_rust.o
 obj-$(CONFIG_TEST_BPF) += test_bpf.o
 test_dhry-objs := dhry_1.o dhry_2.o dhry_run.o
 obj-$(CONFIG_TEST_DHRY) += test_dhry.o
diff --git a/lib/xarray_benchmark_rust.rs b/lib/xarray_benchmark_rust.rs
new file mode 100644
index 0000000000000..fa69b975d8d9f
--- /dev/null
+++ b/lib/xarray_benchmark_rust.rs
@@ -0,0 +1,117 @@
+// SPDX-License-Identifier: GPL-2.0
+//
+//! Benchmark for the XArray Rust APIs.
+
+use kernel::{
+    bench::{self, Bencher},
+    prelude::*,
+    rxarray::{self, XArray4, XArray6},
+    time::{Delta, Instant, Monotonic},
+    xarray::{self, AllocKind}, //
+};
+
+/// Stores integer `entries` at `0..entries` in the empty `xa` and returns the time taken.
+fn store_int<const SHIFT: usize, const SIZE: usize>(
+    mut xa: rxarray::XArray<KBox<u64>, SHIFT, SIZE>,
+    entries: usize,
+) -> Delta {
+    let time = Instant::<Monotonic>::now();
+    for i in 0..entries {
+        let entry = rxarray::Entry::try_int(i * 10).expect("the value fits in an integer entry");
+        xa.store(i, entry, GFP_KERNEL).expect("store");
+    }
+    time.elapsed()
+}
+
+/// Stores pointer `entries` at `0..entries` in the empty `xa` and returns the time taken.
+fn store_ptr<const SHIFT: usize, const SIZE: usize>(
+    mut xa: rxarray::XArray<KBox<u64>, SHIFT, SIZE>,
+    entries: usize,
+) -> Delta {
+    let time = Instant::<Monotonic>::now();
+    for i in 0..entries {
+        let entry = rxarray::Entry::Pointer(KBox::new(i as u64, GFP_KERNEL).expect("allocation"));
+        xa.store(i, entry, GFP_KERNEL).expect("store");
+    }
+    time.elapsed()
+}
+
+/// Allocates the empty C XArray that [`store_ptr_xarray`] stores into.
+fn new_ptr_xarray() -> Pin<KBox<xarray::XArray<KBox<u64>>>> {
+    KBox::pin_init(xarray::XArray::new(AllocKind::Alloc), GFP_KERNEL).expect("allocation")
+}
+
+/// Stores pointer `entries` at `0..entries` in the empty C XArray `xa` and returns the
+/// time taken.
+fn store_ptr_xarray(xa: Pin<KBox<xarray::XArray<KBox<u64>>>>, entries: usize) -> Delta {
+    let time = Instant::<Monotonic>::now();
+    for i in 0..entries {
+        let value = KBox::new(i as u64, GFP_KERNEL).expect("allocation");
+        xa.lock()
+            .store(i, value, GFP_KERNEL)
+            .map_err(|e| e.error)
+            .expect("store");
+    }
+    time.elapsed()
+}
+
+/// Runs every benchmark.
+fn benchmark(samples: usize, entries: usize) -> Result {
+    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_int_rxarray4", XArray4::<KBox<u64>>::new, store_int)
+    );
+    pr_info!(
+        "{}\n",
+        bench.run("store_int_rxarray6", XArray6::<KBox<u64>>::new, store_int)
+    );
+    pr_info!(
+        "{}\n",
+        bench.run("store_ptr_rxarray4", XArray4::<KBox<u64>>::new, store_ptr)
+    );
+    pr_info!(
+        "{}\n",
+        bench.run("store_ptr_rxarray6", XArray6::<KBox<u64>>::new, store_ptr)
+    );
+    pr_info!(
+        "{}\n",
+        bench.run("store_ptr_xarray", new_ptr_xarray, store_ptr_xarray)
+    );
+    pr_info!("total runtime {}\n", bench.runtime());
+    Ok(())
+}
+
+/// The benchmark module.
+struct Benchmark;
+
+impl kernel::Module for Benchmark {
+    fn init(_module: &'static ThisModule) -> Result<Self> {
+        let samples = module_parameters::samples.value();
+        let entries = module_parameters::entries.value();
+        benchmark(samples, entries)?;
+
+        Ok(Benchmark)
+    }
+}
+
+module! {
+    type: Benchmark,
+    name: "xarray_benchmark_rust",
+    authors: ["Daniel Gomez <da.gomez@samsung.com>"],
+    description: "Benchmark: XArray",
+    license: "GPL v2",
+    params: {
+        samples: usize {
+            default: 100,
+            description: "Timed runs per benchmark",
+        },
+        entries: usize {
+            default: 100_000,
+            description: "Entries stored per run",
+        },
+    },
+}

-- 
2.55.0