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

[Andrea Righi <arighi@nvidia.com>]

Hi Cheng-Yang,

this has been sitting in my TODO list for a while, so thanks for looking
at it. :)

Comments below.

On Mon, Mar 16, 2026 at 02:10:14AM +0800, Cheng-Yang Chou wrote:
> 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;

Do we need this array? Isn't this always nr_node_ids - 1 (all other
possible nodes)?

Also, if nodes are sparse we should probably initialize the array with
NUMA_NO_NODE items and skip them when we scan the array (in addition to the
node_online() check).

> +
>  /*
>   * 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);

Considering the comment above, we probably need only nr_node_ids - 1 items.

> +	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
> 

Thanks,
-Andrea