[PATCH 1/2] sched: psi: loosen clock sync between scheduler and aggregator

[PATCH 1/2] sched: psi: loosen clock sync between scheduler and aggregator

[Johannes Weiner]

In the aggregator, catch races between state snooping and task state
conclusions explicitly by checking for sample underflows; then move
the clock reads out of the reader's seqcount protection.

This shrinks the critical section and allows switching the scheduler
side to looser (cheaper) clock sourcing in the next patch.

Suggested-by: Chengming Zhou <chengming.zhou@linux.dev>
Signed-off-by: Johannes Weiner <hannes@cmpxchg.org>
---
 kernel/sched/psi.c | 34 +++++++++++++++++++++++++++-------
 1 file changed, 27 insertions(+), 7 deletions(-)

diff --git a/kernel/sched/psi.c b/kernel/sched/psi.c
index 59fdb7ebbf22..4b7bf8eb46c2 100644
--- a/kernel/sched/psi.c
+++ b/kernel/sched/psi.c
@@ -285,7 +285,6 @@ static void get_recent_times(struct psi_group *group, int cpu,
 	/* Snapshot a coherent view of the CPU state */
 	do {
 		seq = psi_read_begin(cpu);
-		now = cpu_clock(cpu);
 		memcpy(times, groupc->times, sizeof(groupc->times));
 		state_mask = groupc->state_mask;
 		state_start = groupc->state_start;
@@ -293,6 +292,9 @@ static void get_recent_times(struct psi_group *group, int cpu,
 			memcpy(tasks, groupc->tasks, sizeof(groupc->tasks));
 	} while (psi_read_retry(cpu, seq));
 
+	if (state_mask)
+		now = cpu_clock(cpu);
+
 	/* Calculate state time deltas against the previous snapshot */
 	for (s = 0; s < NR_PSI_STATES; s++) {
 		u32 delta;
@@ -308,7 +310,22 @@ static void get_recent_times(struct psi_group *group, int cpu,
 		if (state_mask & (1 << s))
 			times[s] += now - state_start;
 
+		/*
+		 * This snooping ahead can obviously race with the
+		 * state concluding on the cpu. If we previously
+		 * snooped to a time past where the state concludes,
+		 * times[s] can now be behind times_prev[s].
+		 *
+		 * time_after32() would be the obvious choice, but
+		 * S32_MAX is right around two seconds, which is the
+		 * aggregation interval; if the aggregator gets
+		 * delayed, there would be a risk of dismissing
+		 * genuinely large samples. Use a larger margin.
+		 */
 		delta = times[s] - groupc->times_prev[aggregator][s];
+		if (delta > psi_period + (psi_period >> 1))
+			delta = 0;
+
 		groupc->times_prev[aggregator][s] = times[s];
 
 		times[s] = delta;
@@ -908,16 +925,18 @@ static void psi_flags_change(struct task_struct *task, int clear, int set)
 
 void psi_task_change(struct task_struct *task, int clear, int set)
 {
-	int cpu = task_cpu(task);
+	int cpu;
 	u64 now;
 
 	if (!task->pid)
 		return;
 
+	cpu = task_cpu(task);
+	now = cpu_clock(cpu);
+
 	psi_flags_change(task, clear, set);
 
 	psi_write_begin(cpu);
-	now = cpu_clock(cpu);
 	for_each_group(group, task_psi_group(task))
 		psi_group_change(group, cpu, clear, set, now, true);
 	psi_write_end(cpu);
@@ -928,10 +947,9 @@ void psi_task_switch(struct task_struct *prev, struct task_struct *next,
 {
 	struct psi_group *common = NULL;
 	int cpu = task_cpu(prev);
-	u64 now;
+	u64 now = cpu_clock(cpu);
 
 	psi_write_begin(cpu);
-	now = cpu_clock(cpu);
 
 	if (next->pid) {
 		psi_flags_change(next, 0, TSK_ONCPU);
@@ -999,6 +1017,7 @@ void psi_task_switch(struct task_struct *prev, struct task_struct *next,
 				psi_group_change(group, cpu, clear, set, now, wake_clock);
 		}
 	}
+
 	psi_write_end(cpu);
 }
 
@@ -1027,9 +1046,9 @@ void psi_account_irqtime(struct rq *rq, struct task_struct *curr, struct task_st
 		return;
 	rq->psi_irq_time = irq;
 
-	psi_write_begin(cpu);
 	now = cpu_clock(cpu);
 
+	psi_write_begin(cpu);
 	for_each_group(group, task_psi_group(curr)) {
 		if (!group->enabled)
 			continue;
@@ -1234,8 +1253,9 @@ void psi_cgroup_restart(struct psi_group *group)
 
 		guard(rq_lock_irq)(cpu_rq(cpu));
 
-		psi_write_begin(cpu);
 		now = cpu_clock(cpu);
+
+		psi_write_begin(cpu);
 		psi_group_change(group, cpu, 0, 0, now, true);
 		psi_write_end(cpu);
 	}
-- 
2.52.0

[PATCH 2/2] sched: psi: use rq_clock() during task state changes

[Johannes Weiner]

In the hottest psi paths, the scheduler already caches the cpu_clock()
call for the event in rq->clock. Now that the clocks between state
changes and pressure aggregation don't need to be synchronized inside
the seqcount section anymore, use the cheaper rq_clock().

Add update_rq_clock() calls to the few places where psi is entered
without the rq already locked.

schbench -n 0 (no think ops):

Before: average rps: 204408.50
 After: average rps: 204755.90

     2.67%     -0.54%  [kernel.kallsyms]         [k] sched_clock_noinstr

Signed-off-by: Johannes Weiner <hannes@cmpxchg.org>
---
 kernel/sched/core.c  |  3 ++-
 kernel/sched/psi.c   | 12 +++++++-----
 kernel/sched/stats.h |  1 +
 3 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index 41ba0be16911..cc66415b85a1 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -5519,9 +5519,10 @@ void sched_tick(void)
 	rq_lock(rq, &rf);
 	donor = rq->donor;
 
+	update_rq_clock(rq);
+
 	psi_account_irqtime(rq, donor, NULL);
 
-	update_rq_clock(rq);
 	hw_pressure = arch_scale_hw_pressure(cpu_of(rq));
 	update_hw_load_avg(rq_clock_task(rq), rq, hw_pressure);
 
diff --git a/kernel/sched/psi.c b/kernel/sched/psi.c
index 4b7bf8eb46c2..4a0a83f1d1dc 100644
--- a/kernel/sched/psi.c
+++ b/kernel/sched/psi.c
@@ -932,7 +932,7 @@ void psi_task_change(struct task_struct *task, int clear, int set)
 		return;
 
 	cpu = task_cpu(task);
-	now = cpu_clock(cpu);
+	now = rq_clock(cpu_rq(cpu));
 
 	psi_flags_change(task, clear, set);
 
@@ -947,7 +947,7 @@ void psi_task_switch(struct task_struct *prev, struct task_struct *next,
 {
 	struct psi_group *common = NULL;
 	int cpu = task_cpu(prev);
-	u64 now = cpu_clock(cpu);
+	u64 now = rq_clock(cpu_rq(cpu));
 
 	psi_write_begin(cpu);
 
@@ -1026,9 +1026,8 @@ void psi_account_irqtime(struct rq *rq, struct task_struct *curr, struct task_st
 {
 	int cpu = task_cpu(curr);
 	struct psi_group_cpu *groupc;
+	u64 irq, now;
 	s64 delta;
-	u64 irq;
-	u64 now;
 
 	if (static_branch_likely(&psi_disabled) || !irqtime_enabled())
 		return;
@@ -1046,7 +1045,7 @@ void psi_account_irqtime(struct rq *rq, struct task_struct *curr, struct task_st
 		return;
 	rq->psi_irq_time = irq;
 
-	now = cpu_clock(cpu);
+	now = rq_clock(rq);
 
 	psi_write_begin(cpu);
 	for_each_group(group, task_psi_group(curr)) {
@@ -1089,6 +1088,7 @@ void psi_memstall_enter(unsigned long *flags)
 	 * race with CPU migration.
 	 */
 	rq = this_rq_lock_irq(&rf);
+	update_rq_clock(rq);
 
 	current->in_memstall = 1;
 	psi_task_change(current, 0, TSK_MEMSTALL | TSK_MEMSTALL_RUNNING);
@@ -1119,6 +1119,7 @@ void psi_memstall_leave(unsigned long *flags)
 	 * race with CPU migration.
 	 */
 	rq = this_rq_lock_irq(&rf);
+	update_rq_clock(rq);
 
 	current->in_memstall = 0;
 	psi_task_change(current, TSK_MEMSTALL | TSK_MEMSTALL_RUNNING, 0);
@@ -1187,6 +1188,7 @@ void cgroup_move_task(struct task_struct *task, struct css_set *to)
 	}
 
 	rq = task_rq_lock(task, &rf);
+	update_rq_clock(rq);
 
 	/*
 	 * We may race with schedule() dropping the rq lock between
diff --git a/kernel/sched/stats.h b/kernel/sched/stats.h
index c903f1a42891..cb67a81e92b6 100644
--- a/kernel/sched/stats.h
+++ b/kernel/sched/stats.h
@@ -210,6 +210,7 @@ static inline void psi_ttwu_dequeue(struct task_struct *p)
 		struct rq *rq;
 
 		rq = __task_rq_lock(p, &rf);
+		update_rq_clock(rq);
 		psi_task_change(p, p->psi_flags, 0);
 		__task_rq_unlock(rq, p, &rf);
 	}
-- 
2.52.0

Re: [PATCH 1/2] sched: psi: loosen clock sync between scheduler and aggregator

[Chengming Zhou]

On 2025/12/10 23:58, Johannes Weiner wrote:
> In the aggregator, catch races between state snooping and task state
> conclusions explicitly by checking for sample underflows; then move
> the clock reads out of the reader's seqcount protection.
> 
> This shrinks the critical section and allows switching the scheduler
> side to looser (cheaper) clock sourcing in the next patch.
> 
> Suggested-by: Chengming Zhou <chengming.zhou@linux.dev>
> Signed-off-by: Johannes Weiner <hannes@cmpxchg.org>

LGTM!

Reviewed-by: Chengming Zhou <chengming.zhou@linux.dev>

Thanks.

> ---
>   kernel/sched/psi.c | 34 +++++++++++++++++++++++++++-------
>   1 file changed, 27 insertions(+), 7 deletions(-)
> 
> diff --git a/kernel/sched/psi.c b/kernel/sched/psi.c
> index 59fdb7ebbf22..4b7bf8eb46c2 100644
> --- a/kernel/sched/psi.c
> +++ b/kernel/sched/psi.c
> @@ -285,7 +285,6 @@ static void get_recent_times(struct psi_group *group, int cpu,
>   	/* Snapshot a coherent view of the CPU state */
>   	do {
>   		seq = psi_read_begin(cpu);
> -		now = cpu_clock(cpu);
>   		memcpy(times, groupc->times, sizeof(groupc->times));
>   		state_mask = groupc->state_mask;
>   		state_start = groupc->state_start;
> @@ -293,6 +292,9 @@ static void get_recent_times(struct psi_group *group, int cpu,
>   			memcpy(tasks, groupc->tasks, sizeof(groupc->tasks));
>   	} while (psi_read_retry(cpu, seq));
>   
> +	if (state_mask)
> +		now = cpu_clock(cpu);
> +
>   	/* Calculate state time deltas against the previous snapshot */
>   	for (s = 0; s < NR_PSI_STATES; s++) {
>   		u32 delta;
> @@ -308,7 +310,22 @@ static void get_recent_times(struct psi_group *group, int cpu,
>   		if (state_mask & (1 << s))
>   			times[s] += now - state_start;
>   
> +		/*
> +		 * This snooping ahead can obviously race with the
> +		 * state concluding on the cpu. If we previously
> +		 * snooped to a time past where the state concludes,
> +		 * times[s] can now be behind times_prev[s].
> +		 *
> +		 * time_after32() would be the obvious choice, but
> +		 * S32_MAX is right around two seconds, which is the
> +		 * aggregation interval; if the aggregator gets
> +		 * delayed, there would be a risk of dismissing
> +		 * genuinely large samples. Use a larger margin.
> +		 */
>   		delta = times[s] - groupc->times_prev[aggregator][s];
> +		if (delta > psi_period + (psi_period >> 1))
> +			delta = 0;
> +
>   		groupc->times_prev[aggregator][s] = times[s];
>   
>   		times[s] = delta;
> @@ -908,16 +925,18 @@ static void psi_flags_change(struct task_struct *task, int clear, int set)
>   
>   void psi_task_change(struct task_struct *task, int clear, int set)
>   {
> -	int cpu = task_cpu(task);
> +	int cpu;
>   	u64 now;
>   
>   	if (!task->pid)
>   		return;
>   
> +	cpu = task_cpu(task);
> +	now = cpu_clock(cpu);
> +
>   	psi_flags_change(task, clear, set);
>   
>   	psi_write_begin(cpu);
> -	now = cpu_clock(cpu);
>   	for_each_group(group, task_psi_group(task))
>   		psi_group_change(group, cpu, clear, set, now, true);
>   	psi_write_end(cpu);
> @@ -928,10 +947,9 @@ void psi_task_switch(struct task_struct *prev, struct task_struct *next,
>   {
>   	struct psi_group *common = NULL;
>   	int cpu = task_cpu(prev);
> -	u64 now;
> +	u64 now = cpu_clock(cpu);
>   
>   	psi_write_begin(cpu);
> -	now = cpu_clock(cpu);
>   
>   	if (next->pid) {
>   		psi_flags_change(next, 0, TSK_ONCPU);
> @@ -999,6 +1017,7 @@ void psi_task_switch(struct task_struct *prev, struct task_struct *next,
>   				psi_group_change(group, cpu, clear, set, now, wake_clock);
>   		}
>   	}
> +
>   	psi_write_end(cpu);
>   }
>   
> @@ -1027,9 +1046,9 @@ void psi_account_irqtime(struct rq *rq, struct task_struct *curr, struct task_st
>   		return;
>   	rq->psi_irq_time = irq;
>   
> -	psi_write_begin(cpu);
>   	now = cpu_clock(cpu);
>   
> +	psi_write_begin(cpu);
>   	for_each_group(group, task_psi_group(curr)) {
>   		if (!group->enabled)
>   			continue;
> @@ -1234,8 +1253,9 @@ void psi_cgroup_restart(struct psi_group *group)
>   
>   		guard(rq_lock_irq)(cpu_rq(cpu));
>   
> -		psi_write_begin(cpu);
>   		now = cpu_clock(cpu);
> +
> +		psi_write_begin(cpu);
>   		psi_group_change(group, cpu, 0, 0, now, true);
>   		psi_write_end(cpu);
>   	}

Re: [PATCH 2/2] sched: psi: use rq_clock() during task state changes

[Chengming Zhou]

On 2025/12/10 23:58, Johannes Weiner wrote:
> In the hottest psi paths, the scheduler already caches the cpu_clock()
> call for the event in rq->clock. Now that the clocks between state
> changes and pressure aggregation don't need to be synchronized inside
> the seqcount section anymore, use the cheaper rq_clock().
> 
> Add update_rq_clock() calls to the few places where psi is entered
> without the rq already locked.
> 
> schbench -n 0 (no think ops):
> 
> Before: average rps: 204408.50
>   After: average rps: 204755.90
> 
>       2.67%     -0.54%  [kernel.kallsyms]         [k] sched_clock_noinstr
> 
> Signed-off-by: Johannes Weiner <hannes@cmpxchg.org>

Reviewed-by: Chengming Zhou <chengming.zhou@linux.dev>

Thanks.

> ---
>   kernel/sched/core.c  |  3 ++-
>   kernel/sched/psi.c   | 12 +++++++-----
>   kernel/sched/stats.h |  1 +
>   3 files changed, 10 insertions(+), 6 deletions(-)
> 
> diff --git a/kernel/sched/core.c b/kernel/sched/core.c
> index 41ba0be16911..cc66415b85a1 100644
> --- a/kernel/sched/core.c
> +++ b/kernel/sched/core.c
> @@ -5519,9 +5519,10 @@ void sched_tick(void)
>   	rq_lock(rq, &rf);
>   	donor = rq->donor;
>   
> +	update_rq_clock(rq);
> +
>   	psi_account_irqtime(rq, donor, NULL);
>   
> -	update_rq_clock(rq);
>   	hw_pressure = arch_scale_hw_pressure(cpu_of(rq));
>   	update_hw_load_avg(rq_clock_task(rq), rq, hw_pressure);
>   
> diff --git a/kernel/sched/psi.c b/kernel/sched/psi.c
> index 4b7bf8eb46c2..4a0a83f1d1dc 100644
> --- a/kernel/sched/psi.c
> +++ b/kernel/sched/psi.c
> @@ -932,7 +932,7 @@ void psi_task_change(struct task_struct *task, int clear, int set)
>   		return;
>   
>   	cpu = task_cpu(task);
> -	now = cpu_clock(cpu);
> +	now = rq_clock(cpu_rq(cpu));
>   
>   	psi_flags_change(task, clear, set);
>   
> @@ -947,7 +947,7 @@ void psi_task_switch(struct task_struct *prev, struct task_struct *next,
>   {
>   	struct psi_group *common = NULL;
>   	int cpu = task_cpu(prev);
> -	u64 now = cpu_clock(cpu);
> +	u64 now = rq_clock(cpu_rq(cpu));
>   
>   	psi_write_begin(cpu);
>   
> @@ -1026,9 +1026,8 @@ void psi_account_irqtime(struct rq *rq, struct task_struct *curr, struct task_st
>   {
>   	int cpu = task_cpu(curr);
>   	struct psi_group_cpu *groupc;
> +	u64 irq, now;
>   	s64 delta;
> -	u64 irq;
> -	u64 now;
>   
>   	if (static_branch_likely(&psi_disabled) || !irqtime_enabled())
>   		return;
> @@ -1046,7 +1045,7 @@ void psi_account_irqtime(struct rq *rq, struct task_struct *curr, struct task_st
>   		return;
>   	rq->psi_irq_time = irq;
>   
> -	now = cpu_clock(cpu);
> +	now = rq_clock(rq);
>   
>   	psi_write_begin(cpu);
>   	for_each_group(group, task_psi_group(curr)) {
> @@ -1089,6 +1088,7 @@ void psi_memstall_enter(unsigned long *flags)
>   	 * race with CPU migration.
>   	 */
>   	rq = this_rq_lock_irq(&rf);
> +	update_rq_clock(rq);
>   
>   	current->in_memstall = 1;
>   	psi_task_change(current, 0, TSK_MEMSTALL | TSK_MEMSTALL_RUNNING);
> @@ -1119,6 +1119,7 @@ void psi_memstall_leave(unsigned long *flags)
>   	 * race with CPU migration.
>   	 */
>   	rq = this_rq_lock_irq(&rf);
> +	update_rq_clock(rq);
>   
>   	current->in_memstall = 0;
>   	psi_task_change(current, TSK_MEMSTALL | TSK_MEMSTALL_RUNNING, 0);
> @@ -1187,6 +1188,7 @@ void cgroup_move_task(struct task_struct *task, struct css_set *to)
>   	}
>   
>   	rq = task_rq_lock(task, &rf);
> +	update_rq_clock(rq);
>   
>   	/*
>   	 * We may race with schedule() dropping the rq lock between
> diff --git a/kernel/sched/stats.h b/kernel/sched/stats.h
> index c903f1a42891..cb67a81e92b6 100644
> --- a/kernel/sched/stats.h
> +++ b/kernel/sched/stats.h
> @@ -210,6 +210,7 @@ static inline void psi_ttwu_dequeue(struct task_struct *p)
>   		struct rq *rq;
>   
>   		rq = __task_rq_lock(p, &rf);
> +		update_rq_clock(rq);
>   		psi_task_change(p, p->psi_flags, 0);
>   		__task_rq_unlock(rq, p, &rf);
>   	}