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

[Johannes Weiner <hannes@cmpxchg.org>]

On Wed, Jul 15, 2026 at 06:03:57PM +0100, Usama Arif wrote:
> 
> 
> 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.

This looks great, thanks for digging deeper into it.

> No functional change intended.

... he says right after large performance deltas lol. Drop that line?

> Signed-off-by: Usama Arif <usama.arif@linux.dev>

Acked-by: Johannes Weiner <hannes@cmpxchg.org>