The function calculates a Hamming weight of a bitmap starting from an
arbitrary bit.

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

diff --git a/include/linux/bitmap.h b/include/linux/bitmap.h
index b0395e4ccf90..0f4789e1f7cb 100644
--- a/include/linux/bitmap.h
+++ b/include/linux/bitmap.h
@@ -57,6 +57,7 @@ struct device;
  *  bitmap_weight(src, nbits)                   Hamming Weight: number set bits
  *  bitmap_weight_and(src1, src2, nbits)        Hamming Weight of and'ed bitmap
  *  bitmap_weight_andnot(src1, src2, nbits)     Hamming Weight of andnot'ed bitmap
+ *  bitmap_weight_from(src, start, nbits)       Hamming Weight starting from @start
  *  bitmap_set(dst, pos, nbits)                 Set specified bit area
  *  bitmap_clear(dst, pos, nbits)               Clear specified bit area
  *  bitmap_find_next_zero_area(buf, len, pos, n, mask)  Find bit free area
@@ -184,6 +185,8 @@ unsigned int __bitmap_weight_and(const unsigned long *bitmap1,
 				 const unsigned long *bitmap2, unsigned int nbits);
 unsigned int __bitmap_weight_andnot(const unsigned long *bitmap1,
 				    const unsigned long *bitmap2, unsigned int nbits);
+unsigned long __bitmap_weight_from(const unsigned long *bitmap,
+				   unsigned int start, unsigned int nbits);
 void __bitmap_set(unsigned long *map, unsigned int start, int len);
 void __bitmap_clear(unsigned long *map, unsigned int start, int len);
 
@@ -479,6 +482,28 @@ unsigned long bitmap_weight_andnot(const unsigned long *src1,
 	return __bitmap_weight_andnot(src1, src2, nbits);
 }
 
+/**
+ * bitmap_weight_from - Hamming weight for a memory region
+ * @bitmap: The base address
+ * @start: The bitnumber to starts weighting
+ * @nbits: the bitmap size in bits
+ *
+ * Returns the number of set bits in the region, or > @nbits in case of error.
+ */
+static __always_inline
+unsigned long bitmap_weight_from(const unsigned long *bitmap,
+				   unsigned int start, unsigned int nbits)
+{
+	if (small_const_nbits(nbits)) {
+		if (unlikely(start >= nbits))
+			return nbits + 1;
+
+		return hweight_long(*bitmap & GENMASK(nbits - 1, start));
+	}
+
+	return __bitmap_weight_from(bitmap, start, nbits);
+}
+
 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 9dc526507875..698d15933c84 100644
--- a/lib/bitmap.c
+++ b/lib/bitmap.c
@@ -335,6 +335,27 @@ EXPORT_SYMBOL(__bitmap_subset);
 	w;									\
 })
 
+#define BITMAP_WEIGHT_FROM(FETCH, start, nbits)	\
+({										\
+	unsigned int __start = (start), __end = (nbits), idx, w;		\
+										\
+	if (unlikely(__start >= __end)) {					\
+		w = __end + 1;							\
+		goto out;							\
+	}									\
+										\
+	idx = __start / BITS_PER_LONG;						\
+	w = hweight_long((FETCH) & BITMAP_FIRST_WORD_MASK(__start));		\
+										\
+	for (idx++; idx * BITS_PER_LONG < __end; idx++)				\
+		w += hweight_long(FETCH);					\
+										\
+	if (__end % BITS_PER_LONG)						\
+		w += hweight_long((FETCH) & BITMAP_LAST_WORD_MASK(__end));	\
+out:										\
+	w;									\
+})
+
 unsigned int __bitmap_weight(const unsigned long *bitmap, unsigned int bits)
 {
 	return BITMAP_WEIGHT(bitmap[idx], bits);
@@ -361,6 +382,13 @@ unsigned int __bitmap_weighted_or(unsigned long *dst, const unsigned long *bitma
 	return BITMAP_WEIGHT(({dst[idx] = bitmap1[idx] | bitmap2[idx]; dst[idx]; }), bits);
 }
 
+unsigned long __bitmap_weight_from(const unsigned long *bitmap,
+				   unsigned int start, unsigned int nbits)
+{
+	return BITMAP_WEIGHT_FROM(bitmap[idx], start, nbits);
+}
+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 c83829ef557f..0fead99bf867 100644
--- a/lib/test_bitmap.c
+++ b/lib/test_bitmap.c
@@ -850,6 +850,34 @@ static void __init test_for_each_set_bit_from(void)
 		expect_eq_bitmap(tmp, copy, 500);
 	}
 }
+static void __init test_bitmap_weight(void)
+{
+	unsigned int bit, w1, w2, w;
+	DECLARE_BITMAP(b, 30);
+
+	bitmap_parselist("all:1/2", b, 30);
+
+	/* Test inline implementation */
+	w = bitmap_weight(b, 30);
+	w1 = bitmap_weight(b, 15);
+	w2 = bitmap_weight_from(b, 15, 30);
+
+	expect_eq_uint(15, w);
+	expect_eq_uint(8, w1);
+	expect_eq_uint(7, w2);
+
+	/* Test outline implementation */
+	w = bitmap_weight(exp1, EXP1_IN_BITS);
+	for (bit = 0; bit < EXP1_IN_BITS; bit++) {
+		w1 = bitmap_weight(exp1, bit);
+		w2 = bitmap_weight_from(exp1, bit, EXP1_IN_BITS);
+		expect_eq_uint(w1 + w2, w);
+	}
+
+	/* Test out-of-range */
+	w = bitmap_weight_from(b, 31, 30);
+	expect_eq_uint(0, !!(w < 30));
+}
 
 static void __init test_for_each_clear_bit(void)
 {
@@ -1441,6 +1469,7 @@ static void __init selftest(void)
 	test_bitmap_const_eval();
 	test_bitmap_read_write();
 	test_bitmap_read_perf();
+	test_bitmap_weight();
 	test_bitmap_write_perf();
 
 	test_find_nth_bit();
-- 
2.43.0

* Yury Norov (NVIDIA) <yury.norov@gmail.com> wrote:

> The function calculates a Hamming weight of a bitmap starting from an
> arbitrary bit.
> 
> Signed-off-by: Yury Norov (NVIDIA) <yury.norov@gmail.com>
> ---
>  include/linux/bitmap.h | 25 +++++++++++++++++++++++++
>  lib/bitmap.c           | 28 ++++++++++++++++++++++++++++
>  lib/test_bitmap.c      | 29 +++++++++++++++++++++++++++++
>  3 files changed, 82 insertions(+)
> 
> diff --git a/include/linux/bitmap.h b/include/linux/bitmap.h
> index b0395e4ccf90..0f4789e1f7cb 100644
> --- a/include/linux/bitmap.h
> +++ b/include/linux/bitmap.h
> @@ -57,6 +57,7 @@ struct device;
>   *  bitmap_weight(src, nbits)                   Hamming Weight: number set bits
>   *  bitmap_weight_and(src1, src2, nbits)        Hamming Weight of and'ed bitmap
>   *  bitmap_weight_andnot(src1, src2, nbits)     Hamming Weight of andnot'ed bitmap
> + *  bitmap_weight_from(src, start, nbits)       Hamming Weight starting from @start
>   *  bitmap_set(dst, pos, nbits)                 Set specified bit area
>   *  bitmap_clear(dst, pos, nbits)               Clear specified bit area
>   *  bitmap_find_next_zero_area(buf, len, pos, n, mask)  Find bit free area
> @@ -184,6 +185,8 @@ unsigned int __bitmap_weight_and(const unsigned long *bitmap1,
>  				 const unsigned long *bitmap2, unsigned int nbits);
>  unsigned int __bitmap_weight_andnot(const unsigned long *bitmap1,
>  				    const unsigned long *bitmap2, unsigned int nbits);
> +unsigned long __bitmap_weight_from(const unsigned long *bitmap,
> +				   unsigned int start, unsigned int nbits);
>  void __bitmap_set(unsigned long *map, unsigned int start, int len);
>  void __bitmap_clear(unsigned long *map, unsigned int start, int len);
>  
> @@ -479,6 +482,28 @@ unsigned long bitmap_weight_andnot(const unsigned long *src1,
>  	return __bitmap_weight_andnot(src1, src2, nbits);
>  }
>  
> +/**
> + * bitmap_weight_from - Hamming weight for a memory region
> + * @bitmap: The base address
> + * @start: The bitnumber to starts weighting
> + * @nbits: the bitmap size in bits
> + *
> + * Returns the number of set bits in the region, or > @nbits in case of error.
> + */
> +static __always_inline
> +unsigned long bitmap_weight_from(const unsigned long *bitmap,
> +				   unsigned int start, unsigned int nbits)
> +{
> +	if (small_const_nbits(nbits)) {
> +		if (unlikely(start >= nbits))
> +			return nbits + 1;
> +
> +		return hweight_long(*bitmap & GENMASK(nbits - 1, start));
> +	}
> +
> +	return __bitmap_weight_from(bitmap, start, nbits);

The 'nbits' name and description is actively misleading: it suggests
the number of bits searched, like bitmap_write() has nbits for
the number of bits written, but in reality it's the *end* index,
not the size of the area to search.

Note how it contrasts with how bitmap_write(..,start,nbits)
works.

So please rename it to something more suitable, like 'end', or so.

Thanks,

	Ingo