regmap: Add maple tree based register cache

The current state of the art for sparse register maps is the
rbtree cache.  This works well for most applications but isn't
always ideal for sparser register maps since the rbtree can get
deep, requiring a lot of walking.  Fortunately the kernel has a
data structure intended to address this very problem, the maple
tree.  Provide an initial implementation of a register cache
based on the maple tree to start taking advantage of it.

The entries stored in the maple tree are arrays of register
values, with the maple tree keys holding the register addresses.
We store data in host native format rather than device native
format as we do for rbtree, this will be a benefit for devices
where we don't marshal data within regmap and simplifies the code
but will result in additional CPU overhead when syncing the cache
on devices where we do marshal data in regmap.

This should work well for a lot of devices, though there's some
additional areas that could be looked at such as caching the
last accessed entry like we do for rbtree and trying to minimise
the maple tree level locking. We should also use bulk writes
rather than single register writes when resyncing the cache where
possible, even if we don't store in device native format.

Very small register maps may continue to to better with rbtree
longer term.

Signed-off-by: Mark Brown <broonie@kernel.org>
Link: https://lore.kernel.org/r/20230325-regcache-maple-v3-2-23e271f93dc7@kernel.org
Signed-off-by: Mark Brown <broonie@kernel.org>

drivers/base/regmap/Makefile

@@ -3,7 +3,7 @@
 CFLAGS_regmap.o := -I$(src)
 
 obj-$(CONFIG_REGMAP) += regmap.o regcache.o
-obj-$(CONFIG_REGMAP) += regcache-rbtree.o regcache-flat.o
+obj-$(CONFIG_REGMAP) += regcache-rbtree.o regcache-flat.o regcache-maple.o
 obj-$(CONFIG_DEBUG_FS) += regmap-debugfs.o
 obj-$(CONFIG_REGMAP_KUNIT) += regmap-kunit.o
 obj-$(CONFIG_REGMAP_AC97) += regmap-ac97.o

drivers/base/regmap/internal.h

@@ -282,6 +282,7 @@ enum regmap_endian regmap_get_val_endian(struct device *dev,
 					 const struct regmap_config *config);
 
 extern struct regcache_ops regcache_rbtree_ops;
+extern struct regcache_ops regcache_maple_ops;
 extern struct regcache_ops regcache_flat_ops;
 
 static inline const char *regmap_name(const struct regmap *map)

drivers/base/regmap/regcache-maple.c

@@ -0,0 +1,278 @@
+// SPDX-License-Identifier: GPL-2.0
+//
+// Register cache access API - maple tree based cache
+//
+// Copyright 2023 Arm, Ltd
+//
+// Author: Mark Brown <broonie@kernel.org>
+
+#include <linux/debugfs.h>
+#include <linux/device.h>
+#include <linux/maple_tree.h>
+#include <linux/slab.h>
+
+#include "internal.h"
+
+static int regcache_maple_read(struct regmap *map,
+			       unsigned int reg, unsigned int *value)
+{
+	struct maple_tree *mt = map->cache;
+	MA_STATE(mas, mt, reg, reg);
+	unsigned long *entry;
+
+	rcu_read_lock();
+
+	entry = mas_find(&mas, reg);
+	if (!entry) {
+		rcu_read_unlock();
+		return -ENOENT;
+	}
+
+	*value = entry[reg - mas.index];
+
+	rcu_read_unlock();
+
+	return 0;
+}
+
+static int regcache_maple_write(struct regmap *map, unsigned int reg,
+				unsigned int val)
+{
+	struct maple_tree *mt = map->cache;
+	MA_STATE(mas, mt, reg, reg);
+	unsigned long *entry, *upper, *lower;
+	unsigned long index, last;
+	size_t lower_sz, upper_sz;
+	int ret;
+
+	rcu_read_lock();
+
+	entry = mas_find(&mas, reg);
+	if (entry) {
+		entry[reg - mas.index] = val;
+		rcu_read_unlock();
+		return 0;
+	}
+
+	/* Any adjacent entries to extend/merge? */
+	mas_set_range(&mas, reg - 1, reg + 1);
+	index = reg;
+	last = reg;
+
+	lower = mas_find(&mas, reg - 1);
+	if (lower) {
+		index = mas.index;
+		lower_sz = (mas.last - mas.index + 1) * sizeof(unsigned long);
+	}
+
+	upper = mas_find(&mas, reg + 1);
+	if (upper) {
+		last = mas.last;
+		upper_sz = (mas.last - mas.index + 1) * sizeof(unsigned long);
+	}
+
+	rcu_read_unlock();
+
+	entry = kmalloc((last - index + 1) * sizeof(unsigned long),
+			GFP_KERNEL);
+	if (!entry)
+		return -ENOMEM;
+
+	if (lower)
+		memcpy(entry, lower, lower_sz);
+	entry[reg - index] = val;
+	if (upper)
+		memcpy(&entry[reg - index + 1], upper, upper_sz);
+
+	/*
+	 * This is safe because the regmap lock means the Maple lock
+	 * is redundant, but we need to take it due to lockdep asserts
+	 * in the maple tree code.
+	 */
+	mas_lock(&mas);
+
+	mas_set_range(&mas, index, last);
+	ret = mas_store_gfp(&mas, entry, GFP_KERNEL);
+
+	mas_unlock(&mas);
+
+	if (ret == 0) {
+		kfree(lower);
+		kfree(upper);
+	}
+
+	return ret;
+}
+
+static int regcache_maple_drop(struct regmap *map, unsigned int min,
+			       unsigned int max)
+{
+	struct maple_tree *mt = map->cache;
+	MA_STATE(mas, mt, min, max);
+	unsigned long *entry, *lower, *upper;
+	unsigned long lower_index, lower_last;
+	unsigned long upper_index, upper_last;
+	int ret;
+
+	lower = NULL;
+	upper = NULL;
+
+	mas_lock(&mas);
+
+	mas_for_each(&mas, entry, max) {
+		/*
+		 * This is safe because the regmap lock means the
+		 * Maple lock is redundant, but we need to take it due
+		 * to lockdep asserts in the maple tree code.
+		 */
+		mas_unlock(&mas);
+
+		/* Do we need to save any of this entry? */
+		if (mas.index < min) {
+			lower_index = mas.index;
+			lower_last = min -1;
+
+			lower = kmemdup(entry, ((min - mas.index) *
+						sizeof(unsigned long)),
+					GFP_KERNEL);
+			if (!lower) {
+				ret = -ENOMEM;
+				goto out;
+			}
+		}
+
+		if (mas.last > max) {
+			upper_index = max + 1;
+			upper_last = mas.last;
+
+			upper = kmemdup(&entry[max + 1],
+					((mas.last - max) *
+					 sizeof(unsigned long)),
+					GFP_KERNEL);
+			if (!upper) {
+				ret = -ENOMEM;
+				goto out;
+			}
+		}
+
+		kfree(entry);
+		mas_lock(&mas);
+		mas_erase(&mas);
+
+		/* Insert new nodes with the saved data */
+		if (lower) {
+			mas_set_range(&mas, lower_index, lower_last);
+			ret = mas_store_gfp(&mas, lower, GFP_KERNEL);
+			if (ret != 0)
+				goto out;
+			lower = NULL;
+		}
+
+		if (upper) {
+			mas_set_range(&mas, upper_index, upper_last);
+			ret = mas_store_gfp(&mas, upper, GFP_KERNEL);
+			if (ret != 0)
+				goto out;
+			upper = NULL;
+		}
+	}
+
+out:
+	mas_unlock(&mas);
+	kfree(lower);
+	kfree(upper);
+
+	return ret;
+}
+
+static int regcache_maple_sync(struct regmap *map, unsigned int min,
+			       unsigned int max)
+{
+	struct maple_tree *mt = map->cache;
+	unsigned long *entry;
+	MA_STATE(mas, mt, min, max);
+	unsigned long lmin = min;
+	unsigned long lmax = max;
+	unsigned int r;
+	int ret;
+
+	map->cache_bypass = true;
+
+	rcu_read_lock();
+
+	mas_for_each(&mas, entry, max) {
+		for (r = max(mas.index, lmin); r <= min(mas.last, lmax); r++) {
+			ret = regcache_sync_val(map, r, entry[r - mas.index]);
+			if (ret != 0)
+				goto out;
+		}
+	}
+
+out:
+	rcu_read_unlock();
+
+	map->cache_bypass = false;
+
+	return ret;
+}
+
+static int regcache_maple_exit(struct regmap *map)
+{
+	struct maple_tree *mt = map->cache;
+	MA_STATE(mas, mt, 0, UINT_MAX);
+	unsigned int *entry;;
+
+	/* if we've already been called then just return */
+	if (!mt)
+		return 0;
+
+	mas_lock(&mas);
+	mas_for_each(&mas, entry, UINT_MAX)
+		kfree(entry);
+	__mt_destroy(mt);
+	mas_unlock(&mas);
+
+	kfree(mt);
+	map->cache = NULL;
+
+	return 0;
+}
+
+static int regcache_maple_init(struct regmap *map)
+{
+	struct maple_tree *mt;
+	int i;
+	int ret;
+
+	mt = kmalloc(sizeof(*mt), GFP_KERNEL);
+	if (!mt)
+		return -ENOMEM;
+	map->cache = mt;
+
+	mt_init(mt);
+
+	for (i = 0; i < map->num_reg_defaults; i++) {
+		ret = regcache_maple_write(map,
+					   map->reg_defaults[i].reg,
+					   map->reg_defaults[i].def);
+		if (ret)
+			goto err;
+	}
+
+	return 0;
+
+err:
+	regcache_maple_exit(map);
+	return ret;
+}
+
+struct regcache_ops regcache_maple_ops = {
+	.type = REGCACHE_MAPLE,
+	.name = "maple",
+	.init = regcache_maple_init,
+	.exit = regcache_maple_exit,
+	.read = regcache_maple_read,
+	.write = regcache_maple_write,
+	.drop = regcache_maple_drop,
+	.sync = regcache_maple_sync,
+};

drivers/base/regmap/regcache.c

@@ -17,6 +17,7 @@
 
 static const struct regcache_ops *cache_types[] = {
 	&regcache_rbtree_ops,
+	&regcache_maple_ops,
 	&regcache_flat_ops,
 };
 

drivers/base/regmap/regmap-kunit.c

@@ -29,19 +29,22 @@ static const struct regcache_types regcache_types_list[] = {
 	{ REGCACHE_NONE, "none" },
 	{ REGCACHE_FLAT, "flat" },
 	{ REGCACHE_RBTREE, "rbtree" },
+	{ REGCACHE_MAPLE, "maple" },
 };
 
 KUNIT_ARRAY_PARAM(regcache_types, regcache_types_list, case_to_desc);
 
 static const struct regcache_types real_cache_types_list[] = {
 	{ REGCACHE_FLAT, "flat" },
 	{ REGCACHE_RBTREE, "rbtree" },
+	{ REGCACHE_MAPLE, "maple" },
 };
 
 KUNIT_ARRAY_PARAM(real_cache_types, real_cache_types_list, case_to_desc);
 
 static const struct regcache_types sparse_cache_types_list[] = {
 	{ REGCACHE_RBTREE, "rbtree" },
+	{ REGCACHE_MAPLE, "maple" },
 };
 
 KUNIT_ARRAY_PARAM(sparse_cache_types, sparse_cache_types_list, case_to_desc);

include/linux/regmap.h

@@ -51,6 +51,7 @@ enum regcache_type {
 	REGCACHE_NONE,
 	REGCACHE_RBTREE,
 	REGCACHE_FLAT,
+	REGCACHE_MAPLE,
 };
 
 /**