[PATCH v7 linux 1/2] sched/cache: Reduce the overhead of task_cache_work by only scan the visisted cpus

[Luo Gengkun <luogengkun2@huawei.com>]

The overhead of task_cache_work() is high, especially in multi-NUMA systems.
Currently, task_cache_work() tries to find the pref_llc by scanning all CPUs
in the system. However, most of these scans are meaningless, such as those
for CPUs that have never been visited or were accessed a long time ago.

To address this problem, introduce visited_cpus to track the visited CPUs
and evict them once they have not been accessed for a duration exceeding
llc_epoch_affinity_timeout.

Signed-off-by: Luo Gengkun <luogengkun2@huawei.com>
---
 include/linux/mm_types.h |  6 +++++
 include/linux/sched.h    |  2 ++
 kernel/sched/fair.c      | 48 +++++++++++++++++++++++++++++++---------
 3 files changed, 46 insertions(+), 10 deletions(-)

diff --git a/include/linux/mm_types.h b/include/linux/mm_types.h
index b18c2b2e7d2c..35559079e4d4 100644
--- a/include/linux/mm_types.h
+++ b/include/linux/mm_types.h
@@ -1620,6 +1620,11 @@ static inline int mm_alloc_sched_noprof(struct mm_struct *mm)
 	if (!pcpu_sched)
 		return -ENOMEM;
 
+	if (!zalloc_cpumask_var(&mm->sc_stat.visited_cpus, GFP_KERNEL)) {
+		free_percpu(pcpu_sched);
+		return -ENOMEM;
+	}
+
 	mm_init_sched(mm, pcpu_sched);
 	return 0;
 }
@@ -1630,6 +1635,7 @@ static inline void mm_destroy_sched(struct mm_struct *mm)
 {
 	free_percpu(mm->sc_stat.pcpu_sched);
 	mm->sc_stat.pcpu_sched = NULL;
+	free_cpumask_var(mm->sc_stat.visited_cpus);
 }
 #else /* !CONFIG_SCHED_CACHE */
 
diff --git a/include/linux/sched.h b/include/linux/sched.h
index 373bcc0598d1..b461a71a65da 100644
--- a/include/linux/sched.h
+++ b/include/linux/sched.h
@@ -2388,6 +2388,7 @@ static __always_inline int task_mm_cid(struct task_struct *t)
 struct sched_cache_time {
 	u64 runtime;
 	unsigned long epoch;
+	unsigned long epoch_last_visit;
 };
 
 struct sched_cache_stat {
@@ -2398,6 +2399,7 @@ struct sched_cache_stat {
 	unsigned long next_scan;
 	unsigned long footprint;
 	int cpu;
+	cpumask_var_t visited_cpus;
 } ____cacheline_aligned_in_smp;
 
 #else
diff --git a/kernel/sched/fair.c b/kernel/sched/fair.c
index d78467ec6ee1..c56b2df66de8 100644
--- a/kernel/sched/fair.c
+++ b/kernel/sched/fair.c
@@ -1585,6 +1585,7 @@ void mm_init_sched(struct mm_struct *mm,
 		pcpu_sched->runtime = 0;
 		/* a slightly stale cpu epoch is acceptible */
 		pcpu_sched->epoch = rq->cpu_epoch;
+		pcpu_sched->epoch_last_visit = rq->cpu_epoch;
 		epoch = rq->cpu_epoch;
 	}
 
@@ -1635,13 +1636,23 @@ static inline void __update_mm_sched(struct rq *rq,
 	}
 }
 
-static unsigned long fraction_mm_sched(struct rq *rq,
-				       struct sched_cache_time *pcpu_sched)
+static unsigned long fraction_mm_sched(int cpu,
+				       struct mm_struct *mm)
 {
+	struct sched_cache_time *pcpu_sched =
+		per_cpu_ptr(mm->sc_stat.pcpu_sched, cpu);
+	struct rq *rq = cpu_rq(cpu);
+
 	guard(raw_spinlock_irqsave)(&rq->cpu_epoch_lock);
 
 	__update_mm_sched(rq, pcpu_sched);
 
+	/* Skip the rq that has not been hit for a long time */
+	if ((rq->cpu_epoch - pcpu_sched->epoch_last_visit) > llc_epoch_affinity_timeout) {
+		cpumask_clear_cpu(cpu, mm->sc_stat.visited_cpus);
+		return 0;
+	}
+
 	/*
 	 * Runtime is a geometric series (r=0.5) and as such will sum to twice
 	 * the accumulation period, this means the multiplcation here should
@@ -1711,6 +1722,9 @@ void account_mm_sched(struct rq *rq, struct task_struct *p, s64 delta_exec)
 		pcpu_sched->runtime += delta_exec;
 		rq->cpu_runtime += delta_exec;
 		epoch = rq->cpu_epoch;
+		pcpu_sched->epoch_last_visit = epoch;
+		if (!cpumask_test_cpu(cpu_of(rq), mm->sc_stat.visited_cpus))
+			cpumask_set_cpu(cpu_of(rq), mm->sc_stat.visited_cpus);
 	}
 
 	/*
@@ -1867,6 +1881,18 @@ static void task_cache_work(struct callback_head *work)
 		guard(rcu)();
 
 		get_scan_cpumasks(cpus, p);
+		/*
+		 * Data race: While evaluating the visited_cpus without
+		 * a lock, a CPU could be concurrently set by
+		 * account_mm_sched(), meaning the scan might skip the newly
+		 * visited CPU if the bit changes during the scan. This is
+		 * a deliberate trade-off between accuracy and efficiency:
+		 * locking would prevent this race but incur extra overhead.
+		 * The missed runtime contribution is negligible because it
+		 * implies this process hasn't run on that CPU for a long
+		 * time, and will be captured in the next cycle.
+		 */
+		cpumask_and(cpus, cpus, mm->sc_stat.visited_cpus);
 
 		for_each_cpu(cpu, cpus) {
 			/* XXX sched_cluster_active */
@@ -1877,19 +1903,21 @@ static void task_cache_work(struct callback_head *work)
 			if (!sd)
 				continue;
 
-			for_each_cpu(i, sched_domain_span(sd)) {
-				occ = fraction_mm_sched(cpu_rq(i),
-							per_cpu_ptr(mm->sc_stat.pcpu_sched, i));
+			for_each_cpu_and(i, sched_domain_span(sd), mm->sc_stat.visited_cpus) {
+				cur = rcu_dereference_all(cpu_rq(i)->curr);
+				if (cur && !(cur->flags & (PF_EXITING | PF_KTHREAD)) &&
+				    cur->mm == mm)
+					nr_running++;
+
+				occ = fraction_mm_sched(i, mm);
+				if (occ == 0)
+					continue;
+
 				a_occ += occ;
 				if (occ > m_occ) {
 					m_occ = occ;
 					m_cpu = i;
 				}
-
-				cur = rcu_dereference_all(cpu_rq(i)->curr);
-				if (cur && !(cur->flags & (PF_EXITING | PF_KTHREAD)) &&
-				    cur->mm == mm)
-					nr_running++;
 			}
 
 			/*
-- 
2.34.1