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

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

[Usama Arif]

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

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

[Usama Arif]

On Fri, 17 Jul 2026 03:59:39 -0700 Usama Arif <usama.arif@linux.dev> wrote:

> 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>


Hello Peter!

Just wanted to check if there were any comments or feedback on this patch?

I checked sched/core branch and didn't see the patch there.

Thanks!
Usama