Add a _from flavor for bitmap_weight(). It's useful when traversing
bitmaps in a loop to not count bits from the beginning every time.

The test for bitmap_weight{,_from}() is added as well.

Signed-off-by: Yury Norov <yury.norov@gmail.com>
---
 include/linux/bitmap.h | 14 ++++++++++++++
 lib/bitmap.c           | 25 +++++++++++++++++++++++++
 lib/test_bitmap.c      | 18 ++++++++++++++++++
 3 files changed, 57 insertions(+)

diff --git a/include/linux/bitmap.h b/include/linux/bitmap.h
index 692d0a1dbe92..6acbdd2abd0c 100644
--- a/include/linux/bitmap.h
+++ b/include/linux/bitmap.h
@@ -168,6 +168,8 @@ bool __bitmap_subset(const unsigned long *bitmap1,
 unsigned int __bitmap_weight(const unsigned long *bitmap, unsigned int nbits);
 unsigned int __bitmap_weight_and(const unsigned long *bitmap1,
 				 const unsigned long *bitmap2, unsigned int nbits);
+unsigned int __bitmap_weight_from(const unsigned long *bitmap,
+					unsigned int bits, unsigned int start);
 void __bitmap_set(unsigned long *map, unsigned int start, int len);
 void __bitmap_clear(unsigned long *map, unsigned int start, int len);
 
@@ -446,6 +448,18 @@ unsigned long bitmap_weight_and(const unsigned long *src1,
 	return __bitmap_weight_and(src1, src2, nbits);
 }
 
+static __always_inline
+unsigned int bitmap_weight_from(const unsigned long *src, unsigned int nbits, unsigned int start)
+{
+	if (unlikely(start >= nbits))
+		return 0;
+
+	if (small_const_nbits(nbits - start) && nbits / BITS_PER_LONG == start / BITS_PER_LONG)
+		return hweight_long(*src & GENMASK(nbits-1, start));
+
+	return __bitmap_weight_from(src, nbits, start);
+}
+
 static __always_inline void bitmap_set(unsigned long *map, unsigned int start,
 		unsigned int nbits)
 {
diff --git a/lib/bitmap.c b/lib/bitmap.c
index cf77d56cf223..65c64911c92f 100644
--- a/lib/bitmap.c
+++ b/lib/bitmap.c
@@ -345,6 +345,24 @@ EXPORT_SYMBOL(__bitmap_subset);
 	w;									\
 })
 
+#define BITMAP_WEIGHT_FROM(FETCH, bits, start)					\
+({										\
+	unsigned long __bits = (bits), val, idx, w = 0, __start = (start), mask;\
+										\
+	mask = BITMAP_FIRST_WORD_MASK(__start);					\
+	idx = __start / BITS_PER_LONG;						\
+										\
+	for (val = (FETCH) & mask; idx < __bits / BITS_PER_LONG; val = (FETCH))	{\
+		w += hweight_long(val);						\
+		idx++;								\
+	}									\
+										\
+	if (__bits % BITS_PER_LONG)						\
+		w += hweight_long((val) & BITMAP_LAST_WORD_MASK(__bits));	\
+										\
+	w;									\
+})
+
 unsigned int __bitmap_weight(const unsigned long *bitmap, unsigned int bits)
 {
 	return BITMAP_WEIGHT(bitmap[idx], bits);
@@ -358,6 +376,13 @@ unsigned int __bitmap_weight_and(const unsigned long *bitmap1,
 }
 EXPORT_SYMBOL(__bitmap_weight_and);
 
+unsigned int __bitmap_weight_from(const unsigned long *bitmap,
+					unsigned int bits, unsigned int start)
+{
+	return BITMAP_WEIGHT_FROM(bitmap[idx], bits, start);
+}
+EXPORT_SYMBOL(__bitmap_weight_from);
+
 void __bitmap_set(unsigned long *map, unsigned int start, int len)
 {
 	unsigned long *p = map + BIT_WORD(start);
diff --git a/lib/test_bitmap.c b/lib/test_bitmap.c
index 3248ed58a817..a5d823f7589d 100644
--- a/lib/test_bitmap.c
+++ b/lib/test_bitmap.c
@@ -361,6 +361,23 @@ static void __init test_bitmap_region(void)
 	expect_eq_uint(bitmap_weight(bmap, 1000), 0);
 }
 
+static void __init test_weight(void)
+{
+	DECLARE_BITMAP(bmap, 1024);
+	unsigned int idx, w;
+
+	for (idx = 0; idx < 1024; idx++)
+		__assign_bit(idx, bmap, idx);
+
+	w = bitmap_weight(bmap, 1024);
+	for (idx = 0; idx < 1024; idx++) {
+		unsigned int w1 = bitmap_weight(bmap, idx);
+		unsigned int w2 = bitmap_weight_from(bmap, 1024, idx);
+
+		expect_eq_uint(w1 + w2, w);
+	}
+}
+
 #define EXP2_IN_BITS	(sizeof(exp2) * 8)
 
 static void __init test_replace(void)
@@ -1260,6 +1277,7 @@ static void __init selftest(void)
 	test_copy();
 	test_bitmap_region();
 	test_replace();
+	test_weight();
 	test_bitmap_arr32();
 	test_bitmap_arr64();
 	test_bitmap_parse();
-- 
2.39.2

On 28/08/2023 20.43, Yury Norov wrote:
> Add a _from flavor for bitmap_weight(). It's useful when traversing
> bitmaps in a loop to not count bits from the beginning every time.
> 
> The test for bitmap_weight{,_from}() is added as well.
> 
> Signed-off-by: Yury Norov <yury.norov@gmail.com>
> ---
>  include/linux/bitmap.h | 14 ++++++++++++++
>  lib/bitmap.c           | 25 +++++++++++++++++++++++++
>  lib/test_bitmap.c      | 18 ++++++++++++++++++
>  3 files changed, 57 insertions(+)
> 
> diff --git a/include/linux/bitmap.h b/include/linux/bitmap.h
> index 692d0a1dbe92..6acbdd2abd0c 100644
> --- a/include/linux/bitmap.h
> +++ b/include/linux/bitmap.h
> @@ -168,6 +168,8 @@ bool __bitmap_subset(const unsigned long *bitmap1,
>  unsigned int __bitmap_weight(const unsigned long *bitmap, unsigned int nbits);
>  unsigned int __bitmap_weight_and(const unsigned long *bitmap1,
>  				 const unsigned long *bitmap2, unsigned int nbits);
> +unsigned int __bitmap_weight_from(const unsigned long *bitmap,
> +					unsigned int bits, unsigned int start);
>  void __bitmap_set(unsigned long *map, unsigned int start, int len);
>  void __bitmap_clear(unsigned long *map, unsigned int start, int len);
>  
> @@ -446,6 +448,18 @@ unsigned long bitmap_weight_and(const unsigned long *src1,
>  	return __bitmap_weight_and(src1, src2, nbits);
>  }
>  
> +static __always_inline
> +unsigned int bitmap_weight_from(const unsigned long *src, unsigned int nbits, unsigned int start)
> +{
> +	if (unlikely(start >= nbits))
> +		return 0;
> +
> +	if (small_const_nbits(nbits - start) && nbits / BITS_PER_LONG == start / BITS_PER_LONG)
> +		return hweight_long(*src & GENMASK(nbits-1, start));

This must be buggy. If I pass compile-time constants nbits==96 and
start==64, the whole if() is true, and we call GENMASK with total
garbage arguments, and access src[0] and not src[1] as that start-value
would suggest.

Don't optimize for irrelevant cases. The expected use of this function
must be that start is some runtime thing. So just make the whole if()
just test whether nbits is small_const, and if so, I think the
hweight_long() line is actually right as-is (though I can never remember
GENMASK argument order).
>  
> +#define BITMAP_WEIGHT_FROM(FETCH, bits, start)					\
> +({										\
> +	unsigned long __bits = (bits), val, idx, w = 0, __start = (start), mask;\
> +										\
> +	mask = BITMAP_FIRST_WORD_MASK(__start);					\
> +	idx = __start / BITS_PER_LONG;						\
> +										\
> +	for (val = (FETCH) & mask; idx < __bits / BITS_PER_LONG; val = (FETCH))	{\
> +		w += hweight_long(val);						\
> +		idx++;								\
> +	}									\
> +										\
> +	if (__bits % BITS_PER_LONG)						\
> +		w += hweight_long((val) & BITMAP_LAST_WORD_MASK(__bits));	\
> +										\
> +	w;									\
> +})
>

Why does the entire function body need to be an unholy statement expression?

Rasmus