[PATCH v2] sched/psi: use __ffs() to walk task-count bitmasks in psi_group_change()

[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 and typically have one or two bits
set. 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 through the
skip path. Switch both walks to __ffs() + m &= m-1 form:

	while (clear) {
		t = __ffs(clear);
		clear &= clear - 1;
		...
	}

which iterates only over the set bits and terminates naturally on
m == 0. m & (m - 1) clears the lowest set bit. This code is easier
to read as well.

An in-kernel microbench (noinline, same body, IRQs off, pinned CPU
on Zen4c, min-of-10 cyc/call) over mask distributions produced by
common scheduler PSI paths:

  mask pattern                        old    new    delta
  empty          (clear=0x0, set=0x0)  3.68   3.68    +0%
  sleep          (clear=0x4, set=0x0)  9.60   3.74   -61%
  iowait-sleep   (clear=0x4, set=0x1) 12.32   4.45   -63%
  memstall-sleep (clear=0xc, set=0x0) 11.87   5.54   -53%
  wake           (clear=0x0, set=0x4)  7.10   3.70   -47%
  iowait-wake    (clear=0x1, set=0x4) 10.83   4.48   -58%

Every non-empty case wins 47-63%: old cost tracks the highest set bit
(linear walk), new cost tracks the count of set bits (skip zeros via
TZCNT). Single-bit patterns run at the empty-case floor.

The generated psi_group_change() text also shrinks by 67 bytes under
-O2 -march=x86-64 (756 -> 689): no scratch register for a "constant 1"
(only __ffs's operand is needed), simpler bit-clear (LEA+AND vs
SHL+NOT+AND after the test), and no skip-if-unset check per position.

Acked-by: Johannes Weiner <hannes@cmpxchg.org>
Signed-off-by: Usama Arif <usama.arif@linux.dev>
---
v1 -> v2
- Switch from using for_each_set_bit() to __ffs(). Based on the
  benchmark: https://gist.github.com/uarif1/e1bf78b54f50099b354b84684f880fda
  and code size reduction for psi_group_change().
---
 kernel/sched/psi.c | 20 ++++++++++++--------
 1 file changed, 12 insertions(+), 8 deletions(-)

diff --git a/kernel/sched/psi.c b/kernel/sched/psi.c
index d9c9d9480a45..2951614cae17 100644
--- a/kernel/sched/psi.c
+++ b/kernel/sched/psi.c
@@ -798,7 +798,7 @@ 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 int t, clear_orig;
 	u32 state_mask;
 
 	lockdep_assert_rq_held(cpu_rq(cpu));
@@ -820,27 +820,31 @@ static void psi_group_change(struct psi_group *group, int cpu,
 		state_mask = groupc->state_mask & PSI_ONCPU;
 	}
 
+	clear_orig = clear;
+
 	/*
 	 * 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;
+	while (clear) {
+		t = __ffs(clear);
+		clear &= clear - 1;
 		if (groupc->tasks[t]) {
 			groupc->tasks[t]--;
 		} else if (!psi_bug) {
 			printk_deferred(KERN_ERR "psi: task underflow! cpu=%d t=%d tasks=[%u %u %u %u] clear=%x set=%x\n",
 					cpu, t, groupc->tasks[0],
 					groupc->tasks[1], groupc->tasks[2],
-					groupc->tasks[3], clear, set);
+					groupc->tasks[3], clear_orig, set);
 			psi_bug = 1;
 		}
 	}
 
-	for (t = 0; set; set &= ~(1 << t), t++)
-		if (set & (1 << t))
-			groupc->tasks[t]++;
+	while (set) {
+		t = __ffs(set);
+		set &= set - 1;
+		groupc->tasks[t]++;
+	}
 
 	if (!group->enabled) {
 		/*
-- 
2.53.0-Meta