[RFC PATCH v3 11/13] lib/sbm: Add helpers to allocate, set, clear, and traverse the bits on sbm

[K Prateek Nayak <kprateek.nayak@amd.com>]

From: Peter Zijlstra <peterz@infradead.org>

Introduce helpers to allocate a sparsebitmap (sbm) of arch configured
length, set a bit on the sbm, clear a bit from the sbm, and iterate all
the set indices on a sbm structure.

  [ yu.c.chen: Fixes for sbm implementation. ]
  [ kprateek: Adapting sbm implementation to a flat array implementation. ]

(Not-yet-)Signed-off-by: Peter Zijlstra <peterz@infradead.org>
(Not-yet-)Signed-off-by: Chen Yu <yu.c.chen@intel.com>
Signed-off-by: K Prateek Nayak <kprateek.nayak@amd.com>
---
 include/linux/sbm.h | 85 +++++++++++++++++++++++++++++++++++++++++++++
 lib/sbm.c           | 55 +++++++++++++++++++++++++++++
 2 files changed, 140 insertions(+)

diff --git a/include/linux/sbm.h b/include/linux/sbm.h
index 232b0076bb3f..63b116e52e6c 100644
--- a/include/linux/sbm.h
+++ b/include/linux/sbm.h
@@ -2,6 +2,8 @@
 #ifndef _LINUX_SBM_H
 #define _LINUX_SBM_H
 
+#include <linux/bitmap.h>
+
 /*
  * Masks and shifts for sbm index to translate
  * a sbm leaf to CPU.
@@ -9,12 +11,95 @@
 extern int __sbm_shift;
 extern int __sbm_mask;
 
+struct sbm {
+	unsigned long   bitmap;
+} ____cacheline_aligned;
+
 int arch_sbm_cpu_instance_id(int cpu);
 void sbm_set_topology(int num_instances, int max_threads_per_instance);
 
 int sbm_cpu_to_idx(int cpu);
 int sbm_idx_to_cpu(int idx);
 
+struct sbm *sbm_alloc(void);
+bool sbm_empty(struct sbm *sbm);
+int sbm_find_next_bit(struct sbm *sbm, int start);
+
+#define __sbm_op(sbm, func)				\
+({							\
+	int idx = sbm_cpu_to_idx(cpu);			\
+	int nr = idx >> __sbm_shift;			\
+	int bit = idx & __sbm_mask;			\
+							\
+	func(bit, &sbm[nr].bitmap);			\
+})
+
+static inline void sbm_cpu_set(struct sbm *sbm, int cpu)
+{
+	__sbm_op(sbm, set_bit);
+}
+
+static inline void sbm_cpu_clear(struct sbm *sbm, int cpu)
+{
+	__sbm_op(sbm, clear_bit);
+}
+
+static inline void __sbm_cpu_set(struct sbm *sbm, int cpu)
+{
+	__sbm_op(sbm, __set_bit);
+}
+
+static inline void __sbm_cpu_clear(struct sbm *sbm, int cpu)
+{
+	__sbm_op(sbm, __clear_bit);
+}
+
+static inline bool sbm_cpu_test(struct sbm *sbm, int cpu)
+{
+	return __sbm_op(sbm, test_bit);
+}
+
+static __always_inline
+unsigned int sbm_find_next_bit_wrap(struct sbm *sbm, int start)
+{
+	int bit = sbm_find_next_bit(sbm, start);
+
+	if (bit >= 0 || start == 0)
+		return bit;
+
+	bit = sbm_find_next_bit(sbm, 0);
+	return bit < start ? bit : -1;
+}
+
+static __always_inline
+unsigned int __sbm_for_each_wrap(struct sbm *sbm, int start, int n)
+{
+	int bit;
+
+	/* If not wrapped around */
+	if (n > start) {
+		/* and have a bit, just return it. */
+		bit = sbm_find_next_bit(sbm, n);
+		if (bit >= 0)
+			return bit;
+
+		/* Otherwise, wrap around and ... */
+		n = 0;
+	}
+
+	/* Search the other part. */
+	bit = sbm_find_next_bit(sbm, n);
+	return bit < start ? bit : -1;
+}
+
+#define sbm_for_each_set_bit(sbm, idx) \
+	for (int idx = sbm_find_next_bit(sbm, 0); \
+	     idx >= 0; idx = sbm_find_next_bit(sbm, idx+1))
+
+#define sbm_for_each_set_bit_wrap(sbm, idx, start) \
+	for (int idx = sbm_find_next_bit_wrap(sbm, start); \
+	     idx >= 0; idx = __sbm_for_each_wrap(sbm, start, idx+1))
+
 int alloc_sbm_index(int cpu);
 void free_sbm_index(int cpu);
 int sbm_init(void);
diff --git a/lib/sbm.c b/lib/sbm.c
index e5b0508b6825..eeca5ce06d50 100644
--- a/lib/sbm.c
+++ b/lib/sbm.c
@@ -12,6 +12,7 @@
 
 static int sbm_max_threads_per_instance __ro_after_init = -1;
 static int sbm_num_instance __ro_after_init = -1;
+static int sbm_max_populated_index;
 
 int __sbm_shift __ro_after_init;
 int __sbm_mask __ro_after_init;
@@ -35,6 +36,11 @@ static __always_inline int *_sbm_idx_to_cpu(void)
 	return runtime_const_ptr(__sbm_idx_to_cpu);
 }
 
+static int sbm_max_index(void)
+{
+	return READ_ONCE(sbm_max_populated_index);
+}
+
 int sbm_cpu_to_idx(int cpu)
 {
 	return _sbm_cpu_to_idx()[cpu];
@@ -45,6 +51,44 @@ int sbm_idx_to_cpu(int idx)
 	return _sbm_idx_to_cpu()[idx];
 }
 
+struct sbm *sbm_alloc(void)
+{
+	return kzalloc_objs(struct sbm, sbm_max_threads_per_instance * sbm_num_instance);
+}
+
+bool sbm_empty(struct sbm *sbm)
+{
+	int i;
+
+	for (i = 0; i <= sbm_max_index(); ++i) {
+		if (sbm[i].bitmap)
+			return false;
+	}
+
+	return true;
+}
+
+int sbm_find_next_bit(struct sbm *sbm, int start)
+{
+	u32 nr = runtime_const_shift_right_32(start, __sbm_shift);
+	u32 bit = runtime_const_mask_32(start, __sbm_mask);
+	unsigned long tmp = 0, mask = (~0UL) << bit;
+
+	for (; nr <= sbm_max_index(); nr++) {
+		tmp = sbm[nr].bitmap & mask;
+		if (tmp)
+			break;
+		/*
+		 * Consider full bitmask from
+		 * second iteration.
+		 */
+		mask = ~0UL;
+	}
+	if (!tmp)
+		return -1;
+	return (nr << __sbm_shift) | __ffs(tmp);
+}
+
 /*
  * Certain architectures may skip initializing sbm propoerties
  * while having an arch_sbm_cpu_instance_id() definition.
@@ -105,6 +149,8 @@ int alloc_sbm_index(int cpu)
 	_sbm_idx_to_cpu()[idx] = cpu;
 	_sbm_cpu_to_idx()[cpu] = idx;
 
+	WRITE_ONCE(sbm_max_populated_index, max(sbm_max_populated_index, i));
+
 	return 0;
 }
 
@@ -127,6 +173,15 @@ void free_sbm_index(int cpu)
 	if (find_first_bit(&__sbm_idx_metadata[leaf].allocated_mask, BITS_PER_LONG) ==
 	    BITS_PER_LONG)
 		__sbm_idx_metadata[leaf].instance_id = -1;
+
+	if (leaf == sbm_max_populated_index) {
+		for (idx = leaf - 1; idx > -1; idx--) {
+			if (__sbm_idx_metadata[idx].instance_id != -1)
+				break;
+		}
+
+		WRITE_ONCE(sbm_max_populated_index, max(idx, 0));
+	}
 }
 
 void __init sbm_set_topology(int num_instances, int max_threads_per_instance)
-- 
2.34.1