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

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

[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, 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.

Signed-off-by: Usama Arif <usama.arif@linux.dev>
---
 kernel/sched/psi.c | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/kernel/sched/psi.c b/kernel/sched/psi.c
index d9c9d9480a45..f5ae1ceb21a7 100644
--- a/kernel/sched/psi.c
+++ b/kernel/sched/psi.c
@@ -798,7 +798,8 @@ 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 long clear_bits, set_bits;
+	unsigned int t;
 	u32 state_mask;
 
 	lockdep_assert_rq_held(cpu_rq(cpu));
@@ -824,9 +825,8 @@ static void psi_group_change(struct psi_group *group, int cpu,
 	 * 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;
+	clear_bits = clear;
+	for_each_set_bit(t, &clear_bits, NR_PSI_TASK_COUNTS) {
 		if (groupc->tasks[t]) {
 			groupc->tasks[t]--;
 		} else if (!psi_bug) {
@@ -838,9 +838,9 @@ static void psi_group_change(struct psi_group *group, int cpu,
 		}
 	}
 
-	for (t = 0; set; set &= ~(1 << t), t++)
-		if (set & (1 << t))
-			groupc->tasks[t]++;
+	set_bits = set;
+	for_each_set_bit(t, &set_bits, NR_PSI_TASK_COUNTS)
+		groupc->tasks[t]++;
 
 	if (!group->enabled) {
 		/*
-- 
2.53.0-Meta

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

[K Prateek Nayak]

Hello Usama,

On 7/14/2026 7:50 PM, Usama Arif wrote:
> diff --git a/kernel/sched/psi.c b/kernel/sched/psi.c
> index d9c9d9480a45..f5ae1ceb21a7 100644
> --- a/kernel/sched/psi.c
> +++ b/kernel/sched/psi.c
> @@ -798,7 +798,8 @@ 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 long clear_bits, set_bits;
> +	unsigned int t;
>  	u32 state_mask;
>  
>  	lockdep_assert_rq_held(cpu_rq(cpu));
> @@ -824,9 +825,8 @@ static void psi_group_change(struct psi_group *group, int cpu,
>  	 * 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;
> +	clear_bits = clear;

nit.

Can't we convert the arguments to unsigned long instead of assigning
them to these local variables?

Apart form that, for_each_set_bit() is indeed much cleaner. Feel free
to include:

Reviewed-by: K Prateek Nayak <kprateek.nayak@amd.com>

> +	for_each_set_bit(t, &clear_bits, NR_PSI_TASK_COUNTS) {
>  		if (groupc->tasks[t]) {
>  			groupc->tasks[t]--;
>  		} else if (!psi_bug) {

-- 
Thanks and Regards,
Prateek

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

[Usama Arif]

On 15/07/2026 04:49, K Prateek Nayak wrote:
> Hello Usama,
> 
> On 7/14/2026 7:50 PM, Usama Arif wrote:
>> diff --git a/kernel/sched/psi.c b/kernel/sched/psi.c
>> index d9c9d9480a45..f5ae1ceb21a7 100644
>> --- a/kernel/sched/psi.c
>> +++ b/kernel/sched/psi.c
>> @@ -798,7 +798,8 @@ 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 long clear_bits, set_bits;
>> +	unsigned int t;
>>  	u32 state_mask;
>>  
>>  	lockdep_assert_rq_held(cpu_rq(cpu));
>> @@ -824,9 +825,8 @@ static void psi_group_change(struct psi_group *group, int cpu,
>>  	 * 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;
>> +	clear_bits = clear;
> 
> nit.
> 
> Can't we convert the arguments to unsigned long instead of assigning
> them to these local variables?

I kind of prefer the explicit conversion. Also all the callers take int
instead of long. The printk would also change from %x to %lx.

The assignment itself should hopefully be free? The compiler hopefully
optimizes it away.

> 
> Apart form that, for_each_set_bit() is indeed much cleaner. Feel free
> to include:
> 
> Reviewed-by: K Prateek Nayak <kprateek.nayak@amd.com>
> 

Thanks for the review!

>> +	for_each_set_bit(t, &clear_bits, NR_PSI_TASK_COUNTS) {
>>  		if (groupc->tasks[t]) {
>>  			groupc->tasks[t]--;
>>  		} else if (!psi_bug) {
> 

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

[Johannes Weiner]

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.

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