[PATCH] sched/psi: use for_each_set_bit() in psi_group_change() task-count walk

[Usama Arif <usama.arif@linux.dev>]

psi_group_change() walks the @clear and @set bitmasks to
decrement/increment groupc->tasks[t]. Both masks are at most
NR_PSI_TASK_COUNTS (=4) wide, dense at [0, 4), and typically
sparse. Today's form visits every position up to the highest set
bit:

	for (t = 0, m = clear; m; m &= ~(1 << t), t++) {
		if (!(m & (1 << t)))
			continue;
		...
	}

so a mask with only bit 3 set still spins four times; the same
open-coded shape repeats for @set. The code is also unnecessarily
hard to read.

Switch both walks to for_each_set_bit() which is easier to read
and also more efficient. As NR_PSI_TASK_COUNTS is a compile-time
constant <= BITS_PER_LONG, find_next_bit() folds into its
small_const_nbits() fast path (single load + GENMASK + __ffs), lowering
to a bit-scan where one exists (x86 TZCNT/BSF, arm64 RBIT+CLZ).

psi_group_change() runs from psi_task_switch() and psi_task_change()
once per ancestor psi_group per event, so the saved iterations
multiply out on any hot scheduler workload.

No functional change intended.

Signed-off-by: Usama Arif <usama.arif@linux.dev>
---
 kernel/sched/psi.c | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/kernel/sched/psi.c b/kernel/sched/psi.c
index d9c9d9480a45..f5ae1ceb21a7 100644
--- a/kernel/sched/psi.c
+++ b/kernel/sched/psi.c
@@ -798,7 +798,8 @@ static void psi_group_change(struct psi_group *group, int cpu,
 			     u64 now, bool wake_clock)
 {
 	struct psi_group_cpu *groupc;
-	unsigned int t, m;
+	unsigned long clear_bits, set_bits;
+	unsigned int t;
 	u32 state_mask;
 
 	lockdep_assert_rq_held(cpu_rq(cpu));
@@ -824,9 +825,8 @@ static void psi_group_change(struct psi_group *group, int cpu,
 	 * The rest of the state mask is calculated based on the task
 	 * counts. Update those first, then construct the mask.
 	 */
-	for (t = 0, m = clear; m; m &= ~(1 << t), t++) {
-		if (!(m & (1 << t)))
-			continue;
+	clear_bits = clear;
+	for_each_set_bit(t, &clear_bits, NR_PSI_TASK_COUNTS) {
 		if (groupc->tasks[t]) {
 			groupc->tasks[t]--;
 		} else if (!psi_bug) {
@@ -838,9 +838,9 @@ static void psi_group_change(struct psi_group *group, int cpu,
 		}
 	}
 
-	for (t = 0; set; set &= ~(1 << t), t++)
-		if (set & (1 << t))
-			groupc->tasks[t]++;
+	set_bits = set;
+	for_each_set_bit(t, &set_bits, NR_PSI_TASK_COUNTS)
+		groupc->tasks[t]++;
 
 	if (!group->enabled) {
 		/*
-- 
2.53.0-Meta