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

[Usama Arif]

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@xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx/
> 
> 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?