[RFC PATCH v3 08/17] sched/cache: Introduce helpers for task migration decisions

[Jianyong Wu]

Cache-aware aggregation must accommodate workloads that need more than
one LLC. Use the node and intra-node LLC affinity sequences established
by earlier patches to distinguish moves toward the preferred region
from moves away from it, and consult utilization to permit expansion.

Add get_node_stats() to accumulate a node's utilization, capacity and LLC
count without allocating memory in the load balance path.

Extend fits_llc_capacity() to support both <= and < comparisons against
the aggregation threshold. Reaching the threshold after arrival is
allowed; departure checks reject a move if the remaining utilization
falls below the threshold.

For moves toward the preferred region, permit a move if the destination
LLC fits the task, or if the source is noticeably busier and remains no
lighter than the destination after the move.

For moves away, walk the cached affinity sequence from the source node
to the destination. Within one node, leave the source LLC at or above
the threshold, and permit moves that reduce LLC imbalance without
reversing the source and destination utilization ordering.

Across nodes, require the source node to remain at or above the
threshold after departure. An intervening node vetoes expansion if its
aggregate utilization still fits after charging one task per LLC, or
per pair of LLCs when its raw distance to the cache-preferred node is
less than the destination's. This estimates average LLC headroom rather
than guaranteeing room in an individual LLC.

At the destination, permit a fitting or balancing move; otherwise leave
the decision to generic load balancing. Nodes preceding the source or
following the destination do not participate in the decision.

Signed-off-by: Jianyong Wu <wujianyong@xxxxxxxx>
---
 kernel/sched/fair.c | 185 ++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 180 insertions(+), 5 deletions(-)

diff --git a/kernel/sched/fair.c b/kernel/sched/fair.c
index fe357d79b6e0..4eb7cede9080 100644
--- a/kernel/sched/fair.c
+++ b/kernel/sched/fair.c
@@ -10478,7 +10478,7 @@ static inline int task_is_ineligible_on_dst_cpu(struct task_struct *p, int dest_
  *
  * (default: ~50%, tunable via debugfs)
  */
-static bool fits_llc_capacity(unsigned long util, unsigned long max)
+static bool fits_llc_capacity(unsigned long util, unsigned long max, bool inclusive)
 {
 	u32 aggr_pct = llc_overaggr_pct;
 
@@ -10489,6 +10489,9 @@ static bool fits_llc_capacity(unsigned long util, unsigned long max)
 	if (cpu_smt_num_threads == 1)
 		aggr_pct = (aggr_pct * 3 / 2);
 
+	if (inclusive)
+		return util * 100 <= max * aggr_pct;
+
 	return util * 100 < max * aggr_pct;
 }
 
@@ -10592,8 +10595,8 @@ static enum llc_mig can_migrate_llc(int src_cpu, int dst_cpu,
 	src_util = src_util < tsk_util ? 0 : src_util - tsk_util;
 	dst_util = dst_util + tsk_util;
 
-	if (!fits_llc_capacity(dst_util, dst_cap) &&
-	    !fits_llc_capacity(src_util, src_cap))
+	if (!fits_llc_capacity(dst_util, dst_cap, false) &&
+	    !fits_llc_capacity(src_util, src_cap, false))
 		return mig_unrestricted;
 
 	if (to_pref) {
@@ -10603,7 +10606,7 @@ static enum llc_mig can_migrate_llc(int src_cpu, int dst_cpu,
 		 * than the src, in which case migration will
 		 * increase the imbalance too much.
 		 */
-		if (!fits_llc_capacity(dst_util, dst_cap) &&
+		if (!fits_llc_capacity(dst_util, dst_cap, false) &&
 		    util_greater(dst_util, src_util))
 			return mig_forbid;
 	} else {
@@ -10614,13 +10617,185 @@ static enum llc_mig can_migrate_llc(int src_cpu, int dst_cpu,
 		 * of preferred LLC, leading to migration again
 		 * back to preferred LLC.
 		 */
-		if (fits_llc_capacity(src_util, src_cap) ||
+		if (fits_llc_capacity(src_util, src_cap, false) ||
 		    !util_greater(src_util, dst_util))
 			return mig_forbid;
 	}
 	return mig_llc;
 }
 
+/* Accumulate a NUMA node's utilization and capacity at LLC granularity. */
+static bool get_node_stats(int node, unsigned long *util_out,
+			   unsigned long *cap_out, int *nr_llcs)
+{
+	unsigned long util = 0, cap = 0;
+	const struct cpumask *span;
+	int nr = 0;
+
+	if (!util_out || !cap_out || !nr_llcs)
+		return false;
+
+	for_each_llc_node_span(node, span) {
+		unsigned long util_tmp, cap_tmp;
+
+		if (!get_llc_stats(cpumask_first(span), &util_tmp, &cap_tmp))
+			return false;
+
+		util += util_tmp;
+		cap += cap_tmp;
+		nr++;
+	}
+
+	*util_out = util;
+	*cap_out = cap;
+	*nr_llcs = nr;
+
+	return true;
+}
+
+static enum llc_mig __maybe_unused can_migrate_llcs(int src_cpu, int dst_cpu,
+				     struct task_struct *p, bool to_pref)
+{
+	unsigned long src_util, src_pre, src_cap, dst_util, dst_cap, dst_pre;
+	int target_node, src_node, dst_node, dst_distance;
+	int node, target_cpu = src_cpu, n_llc;
+	unsigned long tsk_util = 0;
+	bool got_src = false;
+
+	if (!get_llc_stats(src_cpu, &src_util, &src_cap) ||
+	    !get_llc_stats(dst_cpu, &dst_util, &dst_cap))
+		return mig_unrestricted;
+
+	if (p) {
+		struct mm_struct *mm = p->mm;
+
+		if (mm)
+			target_cpu = READ_ONCE(mm->sc_stat.cpu);
+		target_cpu = target_cpu >= 0 ? target_cpu : src_cpu;
+
+		tsk_util = task_util(p);
+	}
+
+	/* Utilization of the destination LLC before the task moves in. */
+	dst_pre = dst_util;
+	src_pre = src_util;
+	src_util = src_util < tsk_util ? 0 : src_util - tsk_util;
+	dst_util = dst_util + tsk_util;
+
+	if (to_pref) {
+		if (fits_llc_capacity(dst_util, dst_cap, true))
+			return mig_llc;
+
+		if (util_greater(src_pre, dst_pre) &&
+		    src_util >= dst_util)
+			return mig_llc;
+
+		return mig_forbid;
+	}
+
+	src_node = cpu_to_node(src_cpu);
+	dst_node = cpu_to_node(dst_cpu);
+	target_node = cpu_to_node(target_cpu);
+	dst_distance = node_distance(target_node, dst_node);
+
+	/*
+	 * For moves away from the cache preferred node, examine the source,
+	 * intervening nodes and destination in the cached affinity order.
+	 */
+	for_each_affi_node(target_cpu, node) {
+		unsigned long nu = 0, nc = 0;
+
+		if (!got_src) {
+			if (node != src_node)
+				continue;
+			got_src = true;
+
+			if (src_node == dst_node) {
+				/* Leave the source LLC at or above the threshold. */
+				if (fits_llc_capacity(src_util, src_cap, false))
+					return mig_forbid;
+
+				if (util_greater(src_pre, dst_pre) &&
+				    src_util >= dst_util)
+					return mig_llc;
+
+				/* Nodes beyond the destination cannot veto this move. */
+				break;
+			}
+
+			/* Cross-node moves must leave the source node filled. */
+			if (get_node_stats(src_node, &nu, &nc, &n_llc) && n_llc) {
+				nu = nu < tsk_util ? 0 : nu - tsk_util;
+				if (fits_llc_capacity(nu, nc, false))
+					return mig_forbid;
+			}
+			continue;
+		}
+
+		if (node != dst_node) {
+			if (!get_node_stats(node, &nu, &nc, &n_llc) || !n_llc)
+				continue;
+
+			/*
+			 * Estimate average LLC headroom in intervening nodes.
+			 * Charge one task per LLC, or per pair of LLCs when
+			 * this node is closer to the target than the destination.
+			 */
+			if (node_distance(target_node, node) < dst_distance)
+				nu += DIV_ROUND_UP(n_llc, 2) * tsk_util;
+			else
+				nu += n_llc * tsk_util;
+
+			if (fits_llc_capacity(nu, nc, true))
+				return mig_forbid;
+			continue;
+		}
+
+		if (fits_llc_capacity(dst_util, dst_cap, true))
+			return mig_llc;
+
+		if (util_greater(src_pre, dst_pre) && src_util >= dst_util)
+			return mig_llc;
+
+		break;
+	}
+
+	return mig_unrestricted;
+}
+
+/* Decide whether the destination precedes the source in the affinity order. */
+static bool __maybe_unused if_to_prefer(int src_cpu, int dst_cpu, int pref_llc)
+{
+	int src_rank, dst_rank, src_node, dst_node, pref_node;
+
+	if (pref_llc < 0)
+		return false;
+
+	src_node = cpu_to_node(src_cpu);
+	dst_node = cpu_to_node(dst_cpu);
+	pref_node = llc_to_node(pref_llc);
+	if (pref_node < 0)
+		return false;
+
+	if (src_node == dst_node) {
+		if (src_node == pref_node) {
+			src_rank = affi_llc_rank_of(pref_llc, llc_id(src_cpu),
+						 pref_node);
+			dst_rank = affi_llc_rank_of(pref_llc, llc_id(dst_cpu),
+						 pref_node);
+			if (src_rank < 0 || dst_rank < 0)
+				return false;
+			return src_rank > dst_rank;
+		}
+		/* Outside the cache preferred node, use ascending LLC IDs. */
+		return llc_id(src_cpu) > llc_id(dst_cpu);
+	}
+
+	src_rank = affi_node_rank_of(pref_node, src_node);
+	dst_rank = affi_node_rank_of(pref_node, dst_node);
+	return src_rank >= 0 && dst_rank >= 0 && src_rank > dst_rank;
+}
+
 /*
  * Check if task p can migrate from source LLC to
  * destination LLC in terms of cache aware load balance.
-- 
2.34.1