[RFC PATCH v3 15/17] sched/cache: Include recently used preferred nodes in LLC scans

[Jianyong Wu]

get_scan_cpumasks() includes the current task's preferred NUMA node,
but other tasks sharing the mm may prefer different nodes. This makes
the LLC scan scope depend on which task performs the scan.

Record preferred nodes observed while accounting tasks that share the
mm, and include recently observed nodes in the scan mask. Expire stale
records using the LLC affinity timeout and clear them when the mm's
preferred cache state is invalidated.

Signed-off-by: Jianyong Wu <wujianyong@xxxxxxxx>
---
 include/linux/mm_types.h |  6 +-----
 include/linux/sched.h    |  1 +
 kernel/sched/fair.c      | 44 ++++++++++++++++++++++++++++++++++++++++
 3 files changed, 46 insertions(+), 5 deletions(-)

diff --git a/include/linux/mm_types.h b/include/linux/mm_types.h
index b18c2b2e7d2c..7fef34c6f33a 100644
--- a/include/linux/mm_types.h
+++ b/include/linux/mm_types.h
@@ -1626,11 +1626,7 @@ static inline int mm_alloc_sched_noprof(struct mm_struct *mm)
 
 #define mm_alloc_sched(...)	alloc_hooks(mm_alloc_sched_noprof(__VA_ARGS__))
 
-static inline void mm_destroy_sched(struct mm_struct *mm)
-{
-	free_percpu(mm->sc_stat.pcpu_sched);
-	mm->sc_stat.pcpu_sched = NULL;
-}
+void mm_destroy_sched(struct mm_struct *mm);
 #else /* !CONFIG_SCHED_CACHE */
 
 static inline int mm_alloc_sched(struct mm_struct *mm) { return 0; }
diff --git a/include/linux/sched.h b/include/linux/sched.h
index 1e4136c2b2a3..7fd4ea8c9037 100644
--- a/include/linux/sched.h
+++ b/include/linux/sched.h
@@ -2393,6 +2393,7 @@ struct sched_cache_time {
 
 struct sched_cache_stat {
 	struct sched_cache_time __percpu *pcpu_sched;
+	unsigned long *node_epoch;
 	raw_spinlock_t lock;
 	unsigned long epoch;
 	u64 nr_running_avg;
diff --git a/kernel/sched/fair.c b/kernel/sched/fair.c
index 20a6a7f5c2f7..b93b6e6c9dea 100644
--- a/kernel/sched/fair.c
+++ b/kernel/sched/fair.c
@@ -1594,6 +1594,9 @@ void mm_init_sched(struct mm_struct *mm,
 	mm->sc_stat.next_scan = jiffies;
 	mm->sc_stat.nr_running_avg = 0;
 	mm->sc_stat.footprint = 0;
+	mm->sc_stat.node_epoch = kcalloc(num_possible_nodes(),
+					 sizeof(*mm->sc_stat.node_epoch),
+					 GFP_KERNEL);
 	/*
 	 * The update to mm->sc_stat should not be reordered
 	 * before initialization to mm's other fields, in case
@@ -1602,6 +1605,14 @@ void mm_init_sched(struct mm_struct *mm,
 	smp_store_release(&mm->sc_stat.pcpu_sched, _pcpu_sched);
 }
 
+void mm_destroy_sched(struct mm_struct *mm)
+{
+	free_percpu(mm->sc_stat.pcpu_sched);
+	mm->sc_stat.pcpu_sched = NULL;
+	kfree(mm->sc_stat.node_epoch);
+	mm->sc_stat.node_epoch = NULL;
+}
+
 /* because why would C be fully specified */
 static __always_inline void __shr_u64(u64 *val, unsigned int n)
 {
@@ -1689,7 +1700,9 @@ void account_mm_sched(struct rq *rq, struct task_struct *p, s64 delta_exec)
 {
 	struct sched_cache_time *pcpu_sched;
 	struct mm_struct *mm = p->mm;
+	unsigned long *node_epoch;
 	int mm_sched_llc = -1;
+	int node, pref_nid = NUMA_NO_NODE;
 	unsigned long epoch;
 
 	if (!sched_cache_enabled())
@@ -1713,6 +1726,13 @@ void account_mm_sched(struct rq *rq, struct task_struct *p, s64 delta_exec)
 		epoch = rq->cpu_epoch;
 	}
 
+	node_epoch = mm->sc_stat.node_epoch;
+#ifdef CONFIG_NUMA_BALANCING
+	pref_nid = READ_ONCE(p->numa_preferred_nid);
+#endif
+	if (pref_nid != NUMA_NO_NODE && mm->sc_stat.node_epoch)
+		mm->sc_stat.node_epoch[pref_nid] = epoch;
+
 	/*
 	 * If this process hasn't hit task_cache_work() for a while invalidate
 	 * its preferred state.
@@ -1722,6 +1742,16 @@ void account_mm_sched(struct rq *rq, struct task_struct *p, s64 delta_exec)
 	    exceed_llc_capacity(mm, cpu_of(rq))) {
 		if (READ_ONCE(mm->sc_stat.cpu) != -1)
 			WRITE_ONCE(mm->sc_stat.cpu, -1);
+
+		if (node_epoch)
+			memset(mm->sc_stat.node_epoch, 0,
+				sizeof(*node_epoch) * num_possible_nodes());
+	} else if (node_epoch) {
+		for_each_node(node) {
+			if ((long)(epoch - READ_ONCE(node_epoch[node])) >
+			    llc_epoch_affinity_timeout)
+				WRITE_ONCE(node_epoch[node], 0);
+		}
 	}
 
 	mm_sched_llc = get_pref_llc(p, mm);
@@ -1765,10 +1795,12 @@ static void get_scan_cpumasks(cpumask_var_t cpus, struct task_struct *p)
 {
 #ifdef CONFIG_NUMA_BALANCING
 	int cpu, curr_cpu, nid, pref_nid;
+	unsigned long *node_epoch;
 
 	if (!static_branch_likely(&sched_numa_balancing))
 		goto out;
 
+	node_epoch = p->mm->sc_stat.node_epoch;
 	cpu = READ_ONCE(p->mm->sc_stat.cpu);
 	if (cpu != -1)
 		nid = cpu_to_node(cpu);
@@ -1789,6 +1821,18 @@ static void get_scan_cpumasks(cpumask_var_t cpus, struct task_struct *p)
 	if (pref_nid == NUMA_NO_NODE)
 		goto out;
 
+	if (node_epoch) {
+		int node;
+
+		for_each_node(node) {
+			if (node == pref_nid)
+				continue;
+
+			if (node_epoch[node] > 0)
+				cpumask_or(cpus, cpus, cpumask_of_node(node));
+		}
+	}
+
 	cpumask_or(cpus, cpus, cpumask_of_node(pref_nid));
 
 	/* honor the task's preferred LLC CPU */
-- 
2.34.1