[PATCH sched_ext/for-7.4] sched_ext: Add a CID NUMA node lookup kfunc

[Andrea Righi <arighi@nvidia.com>]

CID-form schedulers may need the physical NUMA node for a CID, for
example to allocate arena memory on that node.

The CID topology's node index is not a physical NUMA node ID, and
scx_bpf_cpu_node() is restricted to CPU-form schedulers.

Add scx_bpf_cid_node() to return the physical NUMA node for a CID and
provide a compatibility helper that returns NUMA_NO_NODE on kernels
without this kfunc.

Signed-off-by: Andrea Righi <arighi@nvidia.com>
---
 kernel/sched/ext/cid.c                   | 24 ++++++++++++++++++++++++
 tools/sched_ext/include/scx/common.bpf.h |  1 +
 tools/sched_ext/include/scx/compat.bpf.h |  5 +++++
 3 files changed, 30 insertions(+)

diff --git a/kernel/sched/ext/cid.c b/kernel/sched/ext/cid.c
index 39f88deb94bc4..932f640ccc2d8 100644
--- a/kernel/sched/ext/cid.c
+++ b/kernel/sched/ext/cid.c
@@ -640,6 +640,29 @@ __bpf_kfunc s32 scx_bpf_cid_to_cpu(s32 cid, const struct bpf_prog_aux *aux)
 	return scx_cid_to_cpu(sch, cid);
 }
 
+/**
+ * scx_bpf_cid_node - Return the NUMA node the given @cid belongs to
+ * @cid: cid to look up
+ * @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
+ *
+ * Return NUMA_NO_NODE if @cid is invalid or there is no scheduler. Unlike
+ * scx_bpf_cpu_node(), this kfunc is available only with CID support, so BPF
+ * schedulers can use its presence to detect NUMA-aware CID lookups.
+ */
+__bpf_kfunc s32 scx_bpf_cid_node(s32 cid, const struct bpf_prog_aux *aux)
+{
+	struct scx_sched *sch;
+	s32 cpu;
+
+	guard(rcu)();
+
+	sch = scx_prog_sched(aux);
+	if (unlikely(!sch))
+		return NUMA_NO_NODE;
+	cpu = scx_cid_to_cpu(sch, cid);
+	return cpu < 0 ? NUMA_NO_NODE : cpu_to_node(cpu);
+}
+
 /**
  * scx_bpf_cpu_to_cid - Return the cid for @cpu
  * @cpu: cpu to look up
@@ -952,6 +975,7 @@ static const struct btf_kfunc_id_set scx_kfunc_set_init_cids = {
 
 BTF_KFUNCS_START(scx_kfunc_ids_cid)
 BTF_ID_FLAGS(func, scx_bpf_cid_to_cpu, KF_IMPLICIT_ARGS)
+BTF_ID_FLAGS(func, scx_bpf_cid_node, KF_IMPLICIT_ARGS)
 BTF_ID_FLAGS(func, scx_bpf_cpu_to_cid, KF_IMPLICIT_ARGS)
 BTF_ID_FLAGS(func, scx_bpf_cid_topo, KF_IMPLICIT_ARGS)
 BTF_KFUNCS_END(scx_kfunc_ids_cid)
diff --git a/tools/sched_ext/include/scx/common.bpf.h b/tools/sched_ext/include/scx/common.bpf.h
index 9660b43d612ce..41342f1bdd311 100644
--- a/tools/sched_ext/include/scx/common.bpf.h
+++ b/tools/sched_ext/include/scx/common.bpf.h
@@ -106,6 +106,7 @@ u64 scx_bpf_now(void) __ksym __weak;
 void scx_bpf_events(struct scx_event_stats *events, size_t events__sz) __ksym __weak;
 s32 scx_bpf_cpu_to_cid(s32 cpu) __ksym __weak;
 s32 scx_bpf_cid_to_cpu(s32 cid) __ksym __weak;
+s32 scx_bpf_cid_node(s32 cid) __ksym __weak;
 void scx_bpf_cid_topo(s32 cid, struct scx_cid_topo *out) __ksym __weak;
 void scx_bpf_kick_cid(s32 cid, u64 flags) __ksym __weak;
 s32 scx_bpf_task_cid(const struct task_struct *p) __ksym __weak;
diff --git a/tools/sched_ext/include/scx/compat.bpf.h b/tools/sched_ext/include/scx/compat.bpf.h
index 7d870331b2a6b..c5c90d1d81e04 100644
--- a/tools/sched_ext/include/scx/compat.bpf.h
+++ b/tools/sched_ext/include/scx/compat.bpf.h
@@ -243,6 +243,11 @@ static inline bool __COMPAT_is_enq_cpu_selected(u64 enq_flags)
 	 scx_bpf_pick_any_cpu_node(cpus_allowed, node, flags) :			\
 	 scx_bpf_pick_any_cpu(cpus_allowed, flags))
 
+/* v7.4: Add scx_bpf_cid_node(). */
+#define __COMPAT_scx_bpf_cid_node(cid)					\
+	(bpf_ksym_exists(scx_bpf_cid_node) ?				\
+	 scx_bpf_cid_node(cid) : NUMA_NO_NODE)
+
 /*
  * v6.18: Add a helper to retrieve the current task running on a CPU.
  *
-- 
2.55.0