[PATCH 5/6] sched_ext: Add a kfunc to read a task's preferred NUMA node

[Andrea Righi <arighi@nvidia.com>]

NUMA balancing derives a preferred node for each task from its hinting
faults, p->numa_preferred_nid, but a BPF scheduler has no reliable way
to read it.

Add scx_bpf_task_numa_nid() to return the task's preferred node, or
NUMA_NO_NODE when the task has none. The value is a hint the scheduler
can weigh against its own placement rules. It is read without the task's
rq lock, so it is a snapshot.

This complements ops.set_numa_node(), which reports changes as they
happen, by letting a scheduler read the current preference from any of
its other callbacks.

Only report the preference of a task on a BPF scheduler if the scheduler
that owns the task implements ops.set_numa_node(), and NUMA_NO_NODE
otherwise: such a task isn't scanned, so its preference would be left
over from its previous scheduling class. Check the owner, not the
caller, as with sub-schedulers they can differ and only the owner
decides whether the task is scanned. A task that isn't on sched_ext
reports its preference as is, which NUMA balancing keeps current for
fair tasks.

Add a compat wrapper that returns NUMA_NO_NODE on kernels without the
kfunc.

Signed-off-by: Andrea Righi <arighi@nvidia.com>
---
 kernel/sched/ext/ext.c                   | 42 ++++++++++++++++++++++++
 tools/sched_ext/include/scx/compat.bpf.h | 10 ++++++
 2 files changed, 52 insertions(+)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index 076e3757c0634..fba2c83451ebc 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -11530,6 +11530,47 @@ __bpf_kfunc s32 scx_bpf_task_cid(const struct task_struct *p)
 	return tbl[task_cpu(p)];
 }
 
+/**
+ * scx_bpf_task_numa_nid - Node NUMA balancing would prefer a task to run on
+ * @p: task of interest
+ *
+ * Return p->numa_preferred_nid: the node NUMA balancing has derived from @p's
+ * hinting faults as the best one to run @p on, which is the node with the most
+ * faults, adjusted for CPU-less nodes and for the NUMA group @p may be part of.
+ * Return NUMA_NO_NODE when @p has no preference yet. A task has none until its
+ * address space has been scanned a few times. Disabling kernel.numa_balancing
+ * stops the scan, which leaves the last preference in place.
+ *
+ * A task on a BPF scheduler that doesn't implement ops.set_numa_node() isn't
+ * scanned, so its preference would be left over from its previous scheduling
+ * class. Return NUMA_NO_NODE for such a task. For a task that isn't on
+ * sched_ext, return its preference as is.
+ *
+ * The value is a hint, not an instruction: nothing in the kernel moves a task
+ * under a BPF scheduler to this node and the scheduler is free to weigh it
+ * against its own placement rules. The value is read without @p's rq lock, so
+ * it is a snapshot. In particular, a task that just changed scheduling class
+ * may report a preference from before the change until its new class has
+ * rescanned its address space.
+ */
+__bpf_kfunc s32 scx_bpf_task_numa_nid(const struct task_struct *p)
+{
+#ifdef CONFIG_NUMA_BALANCING
+	if (task_on_scx(p)) {
+		struct scx_sched *sch;
+
+		guard(rcu)();
+		sch = scx_task_sched_rcu(p);
+		if (!sch || !SCX_HAS_OP(sch, set_numa_node))
+			return NUMA_NO_NODE;
+	}
+
+	return READ_ONCE(p->numa_preferred_nid);
+#else
+	return NUMA_NO_NODE;
+#endif
+}
+
 /**
  * scx_bpf_locked_rq - Return the rq currently locked by SCX
  * @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
@@ -11836,6 +11877,7 @@ BTF_ID_FLAGS(func, scx_bpf_put_cpumask, KF_RELEASE)
 BTF_ID_FLAGS(func, scx_bpf_task_running, KF_RCU)
 BTF_ID_FLAGS(func, scx_bpf_task_cpu, KF_RCU)
 BTF_ID_FLAGS(func, scx_bpf_task_cid, KF_RCU)
+BTF_ID_FLAGS(func, scx_bpf_task_numa_nid, KF_RCU)
 BTF_ID_FLAGS(func, scx_bpf_locked_rq, KF_IMPLICIT_ARGS | KF_RET_NULL)
 BTF_ID_FLAGS(func, scx_bpf_cpu_curr, KF_IMPLICIT_ARGS | KF_RET_NULL | KF_RCU_PROTECTED)
 BTF_ID_FLAGS(func, scx_bpf_cid_curr, KF_IMPLICIT_ARGS | KF_RET_NULL | KF_RCU_PROTECTED)
diff --git a/tools/sched_ext/include/scx/compat.bpf.h b/tools/sched_ext/include/scx/compat.bpf.h
index 3520a749f8ce6..7dc94e54db167 100644
--- a/tools/sched_ext/include/scx/compat.bpf.h
+++ b/tools/sched_ext/include/scx/compat.bpf.h
@@ -478,6 +478,16 @@ static inline bool scx_bpf_task_set_dsq_vtime(struct task_struct *p, u64 vtime)
 	return true;
 }
 
+/* v7.5: Add scx_bpf_task_numa_nid(). */
+s32 scx_bpf_task_numa_nid___new(const struct task_struct *p) __ksym __weak;
+
+static inline s32 scx_bpf_task_numa_nid(const struct task_struct *p)
+{
+	if (bpf_ksym_exists(scx_bpf_task_numa_nid___new))
+		return scx_bpf_task_numa_nid___new(p);
+	return NUMA_NO_NODE;
+}
+
 /* v7.4: Add scx_bpf_task_set_lazy_resched(). */
 bool scx_bpf_task_set_lazy_resched___new(struct task_struct *p, bool lazy) __ksym __weak;
 
-- 
2.56.0