[PATCH] lib: test bitmap vs Maple Tree performance for region allocations

[Yury Norov <ynorov@nvidia.com>]

Compare the cost of allocating and freeing variable-sized regions using
a bitmap and a Maple Tree. Both implementations process the same randomly
generated sequence of region sizes, ranging from 1 to 32 entries, until
the configured capacity is exhausted.

Run the benchmark at several capacities to show how the two approaches
scale. Report allocation and free times separately because bitmap
clearing and Maple Tree removal have substantially different costs.

On x86/kvm, the output example is:

Start testing bitmap vs Maple Tree region allocation
type       alloc (ns)      free (ns)   capacity   memory (B)
bitmap      100191761         187248    1000000       125000
maple        12256712       12665044    1000000      1555472
bitmap        1094833          22182     100000        12504
maple         1348996        1120168     100000       157456
bitmap          19566           2230      10000         1256
maple          109854         111542      10000        16144
bitmap            918            222       1000          128
maple            8993           7986       1000         1552

Maple Tree memory consumption is a lower-bound estimate because it
does not account for internal nodes, partially occupied nodes, slab
overhead, and transient allocations.

The benchmark is motivated by the discussion at the link below about
choosing the best data structure for the channel ID pool with the
capacity of 2048 IDs for the nova GPU driver.

Specifically for 2048 IDs the result is:

bitmap           2261            472       2048          256
maple           21743          24055       2048         3344

Link: https://lore.kernel.org/all/20260710-chid-maple-v1-1-4ee869055268@nvidia.com/
Signed-off-by: Yury Norov <ynorov@nvidia.com>
---

The numbers are collected on top of bitmap_find_next_zero_area_off()
optimization:

https://lore.kernel.org/all/tencent_FBA5367F4AA8C57DE2198975767E416FC90A@qq.com/

It's in -next and scheduled for the next merge window.

 MAINTAINERS                  |   2 +
 lib/Kconfig.debug            |   9 +++
 lib/Makefile                 |   1 +
 lib/region_alloc_benchmark.c | 112 +++++++++++++++++++++++++++++++++++
 4 files changed, 124 insertions(+)
 create mode 100644 lib/region_alloc_benchmark.c

diff --git a/MAINTAINERS b/MAINTAINERS
index 7cc4bca5a2c5..9fc60489a0f0 100644
--- a/MAINTAINERS
+++ b/MAINTAINERS
@@ -4615,6 +4615,7 @@ F:	lib/bitmap.c
 F:	lib/cpumask.c
 F:	lib/find_bit.c
 F:	lib/find_bit_benchmark.c
+F:	lib/region_alloc_benchmark.c
 F:	lib/test_bitmap.c
 F:	lib/tests/cpumask_kunit.c
 F:	tools/include/linux/bitfield.h
@@ -15581,6 +15582,7 @@ F:	Documentation/core-api/maple_tree.rst
 F:	include/linux/maple_tree.h
 F:	include/trace/events/maple_tree.h
 F:	lib/maple_tree.c
+F:	lib/region_alloc_benchmark.c
 F:	lib/test_maple_tree.c
 F:	rust/helpers/maple_tree.c
 F:	rust/kernel/maple_tree.rs
diff --git a/lib/Kconfig.debug b/lib/Kconfig.debug
index 1244dcac2294..40872263cff1 100644
--- a/lib/Kconfig.debug
+++ b/lib/Kconfig.debug
@@ -2683,6 +2683,15 @@ config FIND_BIT_BENCHMARK
 
 	  If unsure, say N.
 
+config REGION_ALLOC_BENCHMARK
+	tristate "Benchmark bitmap and Maple Tree region allocation"
+	help
+	  This builds a microbenchmark comparing variable-sized region
+	  allocation using bitmaps and Maple Trees. The benchmark runs at
+	  module load time.
+
+	  If unsure, say N.
+
 config FIND_BIT_BENCHMARK_RUST
 	tristate "Test find_bit functions in Rust"
 	depends on RUST
diff --git a/lib/Makefile b/lib/Makefile
index 7f75cc6edf94..adb18810e3f7 100644
--- a/lib/Makefile
+++ b/lib/Makefile
@@ -64,6 +64,7 @@ obj-y += hexdump.o
 obj-$(CONFIG_TEST_HEXDUMP) += test_hexdump.o
 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_TEST_BPF) += test_bpf.o
 test_dhry-objs := dhry_1.o dhry_2.o dhry_run.o
diff --git a/lib/region_alloc_benchmark.c b/lib/region_alloc_benchmark.c
new file mode 100644
index 000000000000..7459e37608bc
--- /dev/null
+++ b/lib/region_alloc_benchmark.c
@@ -0,0 +1,112 @@
+// SPDX-License-Identifier: GPL-2.0-only
+/* Benchmark bitmap and Maple Tree allocation of variable-sized regions. */
+
+#include <linux/bitmap.h>
+#include <linux/kernel.h>
+#include <linux/maple_tree.h>
+#include <linux/module.h>
+#include <linux/printk.h>
+#include <linux/random.h>
+#include <linux/xarray.h>
+
+#define MAP_SIZE	(1000000UL)
+#define REGION_MAX_SIZE	32
+
+static DECLARE_BITMAP(alloc_bitmap, MAP_SIZE) __initdata;
+/* One more request guarantees that even an all-ones trace reaches ENOSPC. */
+static u8 region_sizes[MAP_SIZE + 1] __initdata;
+static unsigned long region_indexes[MAP_SIZE] __initdata;
+
+static unsigned long __init benchmark_bitmap(unsigned long capacity)
+{
+	unsigned long count, index;
+	ktime_t alloc_time, free_time;
+	size_t memory;
+
+	bitmap_zero(alloc_bitmap, MAP_SIZE);
+	alloc_time = ktime_get();
+	for (count = 0; count < ARRAY_SIZE(region_sizes); count++) {
+		index = bitmap_find_next_zero_area(alloc_bitmap, capacity, 0,
+						   region_sizes[count], 0);
+		if (index >= capacity)
+			break;
+
+		region_indexes[count] = index;
+		bitmap_set(alloc_bitmap, index, region_sizes[count]);
+	}
+	alloc_time = ktime_get() - alloc_time;
+
+	index = count;
+	free_time = ktime_get();
+	while (index--)
+		bitmap_clear(alloc_bitmap, region_indexes[index],
+			     region_sizes[index]);
+	free_time = ktime_get() - free_time;
+	memory = BITS_TO_LONGS(capacity) * sizeof(unsigned long);
+	pr_err("%-6s   %12llu   %12llu   %8lu   %10zu\n", "bitmap", alloc_time,
+		free_time, capacity, memory);
+
+	return count;
+}
+
+static unsigned long __init benchmark_maple_tree(unsigned long capacity)
+{
+	struct maple_tree mt = MTREE_INIT(mt, MT_FLAGS_ALLOC_RANGE);
+	unsigned long count, index;
+	ktime_t alloc_time, free_time;
+	size_t memory;
+	int ret;
+
+	alloc_time = ktime_get();
+	for (count = 0; count < ARRAY_SIZE(region_sizes); count++) {
+		ret = mtree_alloc_range(&mt, &index, xa_mk_value(count + 1),
+					region_sizes[count], 0, capacity - 1,
+					GFP_KERNEL);
+		if (ret)
+			break;
+
+		region_indexes[count] = index;
+	}
+	alloc_time = ktime_get() - alloc_time;
+	WARN_ON(ret != -EBUSY);
+
+	index = count;
+	free_time = ktime_get();
+	while (index--)
+		mtree_erase(&mt, region_indexes[index]);
+	free_time = ktime_get() - free_time;
+	/* Minimum storage assuming fully occupied allocation-range leaf nodes. */
+	memory = sizeof(mt) + DIV_ROUND_UP(count, MAPLE_ARANGE64_SLOTS) *
+		 sizeof(struct maple_node);
+	pr_err("%-6s   %12llu   %12llu   %8lu   %10zu\n", "maple", alloc_time,
+		free_time, capacity, memory);
+	mtree_destroy(&mt);
+
+	return count;
+}
+
+static int __init region_alloc_benchmark(void)
+{
+	static const unsigned long capacities[] = { 1000000, 100000, 10000, 1000 };
+	unsigned long bitmap_count, maple_count;
+	unsigned long i;
+
+	for (i = 0; i < ARRAY_SIZE(region_sizes); i++)
+		region_sizes[i] = get_random_u32_below(REGION_MAX_SIZE) + 1;
+
+	pr_err("\nStart testing bitmap vs Maple Tree region allocation\n");
+	pr_err("%-6s   %12s   %12s   %8s   %10s\n", "type", "alloc (ns)", "free (ns)",
+		"capacity", "memory (B)");
+	for (i = 0; i < ARRAY_SIZE(capacities); i++) {
+		bitmap_count = benchmark_bitmap(capacities[i]);
+		maple_count = benchmark_maple_tree(capacities[i]);
+		WARN_ON(bitmap_count != maple_count);
+	}
+
+	/* Let the benchmark be loaded and run repeatedly without rmmod. */
+	return -EINVAL;
+}
+module_init(region_alloc_benchmark);
+
+MODULE_DESCRIPTION("Benchmark bitmap and Maple Tree region allocation");
+MODULE_LICENSE("GPL");
-- 
2.53.0