[PATCH 1/2] sched_ext: Cache per-node NUMA distance order in scx_idle_init_masks()

[Cheng-Yang Chou <yphbchou0911@gmail.com>]

Add scx_numa_node_order[] and scx_numa_node_order_cnt[], per-node
arrays that store NUMA nodes sorted by increasing distance from each
node. They are allocated and populated once during boot in
scx_idle_init_masks() using the existing for_each_node_numadist()
O(N^2) traversal, so the cost is paid only at init time.

pick_idle_cpu_from_online_nodes() will use these cached arrays to
reduce its per-call complexity from O(N^2) to O(N).

Signed-off-by: Cheng-Yang Chou <yphbchou0911@gmail.com>
---
 kernel/sched/ext_idle.c | 38 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/kernel/sched/ext_idle.c b/kernel/sched/ext_idle.c
index 03be4d664267..623a0356d6c3 100644
--- a/kernel/sched/ext_idle.c
+++ b/kernel/sched/ext_idle.c
@@ -144,6 +144,14 @@ static s32 pick_idle_cpu_in_node(const struct cpumask *cpus_allowed, int node, u
  */
 static DEFINE_PER_CPU(nodemask_t, per_cpu_unvisited);
 
+/*
+ * Per-node sorted arrays of NUMA nodes in increasing distance order,
+ * pre-computed at init time to avoid the O(N^2) traversal on each call
+ * to pick_idle_cpu_from_online_nodes().
+ */
+static int **scx_numa_node_order;
+static int *scx_numa_node_order_cnt;
+
 /*
  * Search for an idle CPU across all nodes, excluding @node.
  */
@@ -685,6 +693,36 @@ void scx_idle_init_masks(void)
 		BUG_ON(!alloc_cpumask_var_node(&per_cpu(local_numa_idle_cpumask, i),
 					       GFP_KERNEL, cpu_to_node(i)));
 	}
+
+#ifdef CONFIG_NUMA
+	/*
+	 * Pre-compute per-node NUMA distance order so that
+	 * pick_idle_cpu_from_online_nodes() can iterate in O(N) instead of
+	 * the O(N^2) traversal required by for_each_node_numadist().
+	 */
+	scx_numa_node_order = kcalloc(nr_node_ids, sizeof(*scx_numa_node_order), GFP_KERNEL);
+	BUG_ON(!scx_numa_node_order);
+	scx_numa_node_order_cnt = kcalloc(nr_node_ids, sizeof(*scx_numa_node_order_cnt), GFP_KERNEL);
+	BUG_ON(!scx_numa_node_order_cnt);
+
+	for_each_node(i) {
+		scx_numa_node_order[i] = kcalloc(nr_node_ids, sizeof(int), GFP_KERNEL);
+		BUG_ON(!scx_numa_node_order[i]);
+	}
+
+	rcu_read_lock();
+	for_each_node(i) {
+		nodemask_t unvisited;
+		int node = i, j = 0;
+
+		nodes_copy(unvisited, node_states[N_POSSIBLE]);
+		node_clear(i, unvisited);
+		for_each_node_numadist(node, unvisited)
+			scx_numa_node_order[i][j++] = node;
+		scx_numa_node_order_cnt[i] = j;
+	}
+	rcu_read_unlock();
+#endif
 }
 
 static void update_builtin_idle(int cpu, bool idle)
-- 
2.48.1