[PATCH 4/6] sched_ext: Report preferred NUMA nodes through ops.set_numa_node()

[Andrea Righi <arighi@nvidia.com>]

BPF schedulers need NUMA access information to place tasks close to
their memory, while retaining control over task placement. Deliver the
preferred node through ops.set_numa_node() so schedulers can update
their placement state when the preference changes instead of polling.

Report the current preference after ops.enable(), including
NUMA_NO_NODE when no preference is available. Notify subsequent changes
through sched_class->numa_node_changed(), which sched_setnuma() calls
under the task's rq lock, before re-enqueueing the task. Support both
CPU and CID schedulers. Kernels without NUMA balancing deliver only the
initial NUMA_NO_NODE notification.

Drive NUMA hinting-fault scans from the sched_ext tick when the callback
is implemented and global NUMA balancing is enabled. The faults maintain
NUMA statistics and can migrate accessed memory toward the accessing
CPU, while task placement remains the BPF scheduler's responsibility.

Suggested-by: Tejun Heo <tj@kernel.org>
Signed-off-by: Andrea Righi <arighi@nvidia.com>
---
 kernel/sched/ext/ext.c      | 46 +++++++++++++++++++++++++++++++++++++
 kernel/sched/ext/internal.h | 20 ++++++++++++++++
 2 files changed, 66 insertions(+)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index 2a2b56c6074c6..076e3757c0634 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -4379,6 +4379,21 @@ static void task_tick_scx(struct rq *rq, struct task_struct *donor, int queued)
 	else if (SCX_HAS_OP(sch, tick))
 		SCX_CALL_OP_TASK(sch, tick, rq, donor);
 
+	/*
+	 * If requested by the scheduler, drive the NUMA hinting-fault scan as
+	 * task_tick_fair() does. Neither the scan nor the statistics it feeds
+	 * depend on the scheduling class. Where the task then runs is left to
+	 * that scheduler: NUMA balancing neither migrates a task it owns nor
+	 * picks one as a swap partner, see numa_migrate_preferred() and
+	 * task_numa_compare().
+	 *
+	 * The tick that only refreshes an already queued task does not scan,
+	 * matching the @queued check in task_tick_fair().
+	 */
+	if (!queued && SCX_HAS_OP(sch, set_numa_node) &&
+	    static_branch_unlikely(&sched_numa_balancing))
+		task_tick_numa(rq, donor);
+
 	if (!donor->scx.slice) {
 		/* the slice can't be trusted while bypassing */
 		if (READ_ONCE(donor->scx.lazy_resched) &&
@@ -4528,6 +4543,15 @@ static void __scx_enable_task(struct scx_sched *sch, struct task_struct *p)
 
 	if (SCX_HAS_OP(sch, set_weight))
 		SCX_CALL_OP_TASK(sch, set_weight, rq, p, p->scx.weight);
+
+	if (SCX_HAS_OP(sch, set_numa_node)) {
+		s32 nid = NUMA_NO_NODE;
+
+#ifdef CONFIG_NUMA_BALANCING
+		nid = p->numa_preferred_nid;
+#endif
+		SCX_CALL_OP_TASK(sch, set_numa_node, rq, p, nid);
+	}
 }
 
 void scx_enable_task(struct scx_sched *sch, struct task_struct *p)
@@ -4834,6 +4858,20 @@ void sched_ext_dead(struct task_struct *p)
 	}
 }
 
+#ifdef CONFIG_NUMA_BALANCING
+static void numa_node_changed_scx(struct rq *rq, struct task_struct *p, int nid)
+{
+	struct scx_sched *sch = scx_task_sched(p);
+
+	if (task_dead_and_done(p) ||
+	    scx_get_task_state(p) != SCX_TASK_ENABLED)
+		return;
+
+	if (SCX_HAS_OP(sch, set_numa_node))
+		SCX_CALL_OP_TASK(sch, set_numa_node, rq, p, nid);
+}
+#endif
+
 static void reweight_task_scx(struct rq *rq, struct task_struct *p,
 			      const struct load_weight *lw)
 {
@@ -5745,6 +5783,10 @@ DEFINE_SCHED_CLASS(ext) = {
 
 	.update_curr		= update_curr_scx,
 
+#ifdef CONFIG_NUMA_BALANCING
+	.numa_node_changed	= numa_node_changed_scx,
+#endif
+
 #ifdef CONFIG_UCLAMP_TASK
 	.uclamp_enabled		= 1,
 #endif
@@ -9069,6 +9111,7 @@ static void sched_ext_ops__quiescent(struct task_struct *p, u64 deq_flags) {}
 static bool sched_ext_ops__yield(struct task_struct *from, struct task_struct *to__nullable) { return false; }
 static bool sched_ext_ops__core_sched_before(struct task_struct *a, struct task_struct *b) { return false; }
 static void sched_ext_ops__set_weight(struct task_struct *p, u32 weight) {}
+static void sched_ext_ops__set_numa_node(struct task_struct *p, s32 nid) {}
 static void sched_ext_ops__set_cpumask(struct task_struct *p, const struct cpumask *mask) {}
 static void sched_ext_ops__update_idle(s32 cpu, bool idle) {}
 static void sched_ext_ops__cpu_acquire(s32 cpu, struct scx_cpu_acquire_args *args) {}
@@ -9111,6 +9154,7 @@ static struct sched_ext_ops __bpf_ops_sched_ext_ops = {
 	.yield			= sched_ext_ops__yield,
 	.core_sched_before	= sched_ext_ops__core_sched_before,
 	.set_weight		= sched_ext_ops__set_weight,
+	.set_numa_node		= sched_ext_ops__set_numa_node,
 	.set_cpumask		= sched_ext_ops__set_cpumask,
 	.update_idle		= sched_ext_ops__update_idle,
 	.cpu_acquire		= sched_ext_ops__cpu_acquire,
@@ -9181,6 +9225,7 @@ static struct sched_ext_ops_cid __bpf_ops_sched_ext_ops_cid = {
 	.yield			= sched_ext_ops__yield,
 	.core_sched_before	= sched_ext_ops__core_sched_before,
 	.set_weight		= sched_ext_ops__set_weight,
+	.set_numa_node		= sched_ext_ops__set_numa_node,
 	.set_cmask		= sched_ext_ops_cid__set_cmask,
 	.update_idle		= sched_ext_ops__update_idle,
 	.init_task		= sched_ext_ops__init_task,
@@ -12026,6 +12071,7 @@ static int __init scx_init(void)
 	CID_OFFSET_MATCH(yield, yield);
 	CID_OFFSET_MATCH(core_sched_before, core_sched_before);
 	CID_OFFSET_MATCH(set_weight, set_weight);
+	CID_OFFSET_MATCH(set_numa_node, set_numa_node);
 	CID_OFFSET_MATCH(update_idle, update_idle);
 	CID_OFFSET_MATCH(init_task, init_task);
 	CID_OFFSET_MATCH(exit_task, exit_task);
diff --git a/kernel/sched/ext/internal.h b/kernel/sched/ext/internal.h
index 80c746054ce7a..a5a8044d3c5a2 100644
--- a/kernel/sched/ext/internal.h
+++ b/kernel/sched/ext/internal.h
@@ -616,6 +616,25 @@ struct sched_ext_ops {
 	 */
 	void (*set_weight)(struct task_struct *p, u32 weight);
 
+	/**
+	 * @set_numa_node: Report a task's preferred NUMA node
+	 * @p: task whose preferred node is being reported
+	 * @nid: preferred node, or NUMA_NO_NODE if there is no preference
+	 *
+	 * Called after ops.enable() to report the current preference and when
+	 * the preference changes. The task's rq lock is held, and updates are
+	 * reported before the task is re-enqueued, if it is runnable.
+	 *
+	 * Implementing this op enables NUMA hinting-fault scans for the
+	 * scheduler's tasks when CONFIG_NUMA_BALANCING and kernel.numa_balancing
+	 * are enabled. The faults maintain NUMA statistics and the preferred
+	 * node, and can migrate accessed memory toward the accessing CPU.
+	 * Task placement remains the BPF scheduler's responsibility; @nid is
+	 * advisory. Without CONFIG_NUMA_BALANCING, only the initial
+	 * NUMA_NO_NODE notification is delivered.
+	 */
+	void (*set_numa_node)(struct task_struct *p, s32 nid);
+
 	/**
 	 * @set_cpumask: Set CPU affinity
 	 * @p: task to set CPU affinity for
@@ -1214,6 +1233,7 @@ struct sched_ext_ops_cid {
 	bool (*core_sched_before)(struct task_struct *a,
 				   struct task_struct *b);
 	void (*set_weight)(struct task_struct *p, u32 weight);
+	void (*set_numa_node)(struct task_struct *p, s32 nid);
 	void (*set_cmask)(struct task_struct *p,
 			   const struct scx_cmask *cmask__arena);
 	void (*update_idle)(s32 cid, bool idle);
-- 
2.56.0