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

[Johannes Weiner <hannes@cmpxchg.org>]

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