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

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

On 15/07/2026 14:54, Johannes Weiner wrote:
> On Tue, Jul 14, 2026 at 07:20:57AM -0700, Usama Arif 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, 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.
> 
> This actually started out using ffs. Because the performance is so
> sensitive in this path, this was handtuned to scheduler benchmarks.
> 
> https://lore.kernel.org/all/20180718120318.GC2476@hirez.programming.kicks-ass.net/
> 
> Not the worst idea to revisit this, but you have to be careful, look
> at the asm, and benchmark it. gcc is producing more code for me with
> your patch.


So I tried with claude to create a kernel module that benchmarks the 3
implementations, the current one, for_each_set_bit and __ffs [1]:
__ffs actually performs best. 

psi_bench: running on cpu 1, iters=2000000, trials=10
psi_bench empty            (clear=0x0, set=0x0)     original=3.68  foreach=3.81  ffs=3.68   (for=+3%  ffs=+0%)
psi_bench sleep            (clear=0x4, set=0x0)     original=9.60  foreach=5.18  ffs=3.74   (for=-46%  ffs=-61%)
psi_bench iowait-sleep     (clear=0x4, set=0x1)     original=12.32  foreach=6.90  ffs=4.45   (for=-43%  ffs=-63%)
psi_bench memstall-sleep   (clear=0xc, set=0x0)     original=11.87  foreach=5.98  ffs=5.54   (for=-49%  ffs=-53%)
psi_bench wake             (clear=0x0, set=0x4)     original=7.10  foreach=5.17  ffs=3.70   (for=-27%  ffs=-47%)
psi_bench iowait-wake      (clear=0x1, set=0x4)     original=10.83  foreach=6.91  ffs=4.48   (for=-36%  ffs=-58%)
psi_bench: done
psi_bench: unloaded

[1] https://gist.github.com/uarif1/e1bf78b54f50099b354b84684f880fda


I built the below patch and the code size reduces as well by
67 bytes from 756B to 689B. If it looks ok, I can send it as v2?


From 052c845707f65d58d5cf8cb1da6515e268159e83 Mon Sep 17 00:00:00 2001
From: Usama Arif <usama.arif@linux.dev>
Date: Sat, 11 Jul 2026 03:57:01 -0700
Subject: [PATCH] sched/psi: use __ffs() to walk task-count bitmasks in
 psi_group_change()

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/Bergamo, 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.

No functional change intended.

Signed-off-by: Usama Arif <usama.arif@linux.dev>
---
 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