[PATCH v2] lib: test bitmap vs IDA 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, IDA and a Maple Tree. All 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 approaches
scale. Report allocation and free times separately because bitmap,
IDA and Maple Tree removal have substantially different costs.

On x86/kvm, the output example is:

type       alloc (ns)      free (ns)   capacity   memory (B)
bitmap      179573071         342105    1000000       125000
IDA          46555636       33931498    1000000       134864
maple        18629665       19741396    1000000      1548304
bitmap        1630912          30933     100000        12504
IDA           6144785        3354590     100000        14288
maple         1745026        1825032     100000       155408
bitmap          28448           3374      10000         1256
IDA            418978         333641      10000         1872
maple          185398         211138      10000        15632
bitmap           2253            610       1000          128
IDA             42755          36432       1000          144
maple           19728          23474       1000         1552

Reported IDA and Maple Tree memory figures exclude slab overhead
and transient allocations. The Maple Tree figure is additionally
a lower-bound estimate that assumes fully occupied leaf nodes and
excludes internal nodes.

IDA has no region-allocation API, so each region is implemented as
a sequence of single-ID allocations. The IDs remain contiguous
because this benchmark fills an initially empty IDA monotonically.

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           3907            831       2048          256
IDA             83545          71007       2048          848
maple           34632          39710       2048         3600

Link: https://lore.kernel.org/all/20260710-chid-maple-v1-1-4ee869055268@nvidia.com/
Signed-off-by: Yury Norov <ynorov@nvidia.com>
---
v2: add IDA benchmark (Matthew)
v1: https://lore.kernel.org/all/20260711013910.349586-1-ynorov@nvidia.com/

 MAINTAINERS                  |   3 +
 lib/Kconfig.debug            |   9 ++
 lib/Makefile                 |   1 +
 lib/region_alloc_benchmark.c | 174 +++++++++++++++++++++++++++++++++++
 4 files changed, 187 insertions(+)
 create mode 100644 lib/region_alloc_benchmark.c

diff --git a/MAINTAINERS b/MAINTAINERS
index 7cc4bca5a2c5..9e487a94aba4 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
@@ -29323,6 +29325,7 @@ F:	Documentation/core-api/xarray.rst
 F:	include/linux/idr.h
 F:	include/linux/xarray.h
 F:	lib/idr.c
+F:	lib/region_alloc_benchmark.c
 F:	lib/test_xarray.c
 F:	lib/xarray.c
 F:	tools/testing/radix-tree
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, IDA and Maple Tree region allocation"
+	help
+	  This builds a microbenchmark comparing variable-sized region
+	  allocation using bitmaps, IDA and Maple Tree. 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..0e000c4f9926
--- /dev/null
+++ b/lib/region_alloc_benchmark.c
@@ -0,0 +1,174 @@
+// SPDX-License-Identifier: GPL-2.0-only
+/* Benchmark bitmap, IDA and Maple Tree allocation of variable-sized regions. */
+
+#include <linux/bitmap.h>
+#include <linux/idr.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 size_t __init ida_memory(unsigned long nr_ids)
+{
+	unsigned long entries = DIV_ROUND_UP(nr_ids, IDA_BITMAP_BITS);
+	unsigned long bitmaps = nr_ids / IDA_BITMAP_BITS;
+	unsigned long nodes = 0;
+
+	if (nr_ids % IDA_BITMAP_BITS > BITS_PER_XA_VALUE)
+		bitmaps++;
+
+	while (entries > 1) {
+		entries = DIV_ROUND_UP(entries, XA_CHUNK_SIZE);
+		nodes += entries;
+	}
+
+	return sizeof(struct ida) + bitmaps * sizeof(struct ida_bitmap) +
+		nodes * sizeof(struct xa_node);
+}
+
+static unsigned long __init benchmark_ida(unsigned long capacity)
+{
+	struct ida ida = IDA_INIT(ida);
+	unsigned long count, index, offset, nr_ids = 0;
+	ktime_t alloc_time, free_time;
+	int id = -ENOSPC;
+
+	alloc_time = ktime_get();
+	for (count = 0; count < ARRAY_SIZE(region_sizes); count++) {
+		for (offset = 0; offset < region_sizes[count]; offset++) {
+			id = ida_alloc_max(&ida, capacity - 1, GFP_KERNEL);
+			if (id < 0)
+				break;
+			if (!offset)
+				region_indexes[count] = id;
+		}
+		if (id < 0) {
+			while (offset--)
+				ida_free(&ida, region_indexes[count] + offset);
+			break;
+		}
+		WARN_ON(id != region_indexes[count] + region_sizes[count] - 1);
+		nr_ids += region_sizes[count];
+	}
+	alloc_time = ktime_get() - alloc_time;
+	WARN_ON(id != -ENOSPC);
+
+	index = count;
+	free_time = ktime_get();
+	while (index--)
+		for (offset = 0; offset < region_sizes[index]; offset++)
+			ida_free(&ida, region_indexes[index] + offset);
+	free_time = ktime_get() - free_time;
+	pr_err("%-6s   %12llu   %12llu   %8lu   %10zu\n", "IDA", alloc_time,
+		free_time, capacity, ida_memory(nr_ids));
+	ida_destroy(&ida);
+
+	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, ida_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 IDA 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]);
+		ida_count = benchmark_ida(capacities[i]);
+		maple_count = benchmark_maple_tree(capacities[i]);
+		WARN_ON(bitmap_count != ida_count);
+		WARN_ON(bitmap_count != maple_count);
+	}
+
+	/* Let the benchmark be loaded and run repeatedly without rmmod. */
+	return -EINVAL;
+}
+module_init(region_alloc_benchmark);
+
+MODULE_AUTHOR("Yury Norov <ynorov@nvidia.com>");
+MODULE_DESCRIPTION("Benchmark bitmap, IDA and Maple Tree region allocation");
+MODULE_LICENSE("GPL");
-- 
2.53.0