[PATCHSET v2 sched_ext/for-7.5] sched_ext: Add NUMA balancing support

[Andrea Righi <arighi@nvidia.com>]

Tasks running under a BPF scheduler are currently invisible to NUMA
balancing: their address space is never scanned for hinting faults, so
they never record NUMA fault statistics and never acquire a preferred
node. A BPF scheduler therefore has no information about where a task's
memory is when deciding where to run it.

This series lets a BPF scheduler opt into the NUMA hinting-fault scan by
implementing a new callback, ops.set_numa_node(). The existing scan is
driven from the sched_ext tick for the tasks of such a scheduler, and
the hinting faults then work for them as they do for fair tasks: the
kernel maintains their NUMA fault statistics and preferred node, and
migrates the memory a task accesses toward the node it runs on. The
tasks themselves are not migrated: their placement is left to the BPF
scheduler.

ops.set_numa_node() reports the task's preferred node: once after
ops.enable(), as ops.set_weight() does for the weight, and then every
time the preference changes. Changes are reported from sched_setnuma(),
under the task's rq lock and before the task is re-enqueued, so the
scheduler can update its placement state before ops.enqueue() sees the
task.

To keep sched core free of class-specific code, sched_setnuma() notifies
the task's class through a new optional sched_class->numa_node_changed()
method, which only sched_ext implements at the moment.

scx_bpf_task_numa_nid() lets a scheduler read the current preference
from any of its other callbacks. For a task on a BPF scheduler that
doesn't implement ops.set_numa_node() it returns NUMA_NO_NODE, as such a
task isn't scanned and its preference would be left over from its
previous scheduling class. The value is advisory; nothing in the kernel
acts on it for these tasks.

Ideally the ultimate goal is to give BPF schedulers both controls that
NUMA balancing combines: moving a task to the node where its memory is
and moving the memory a task uses to the node where the task runs. This
series provides the first: with the preferred node available, a
scheduler can decide where to run the task. A second series, to follow
once this one lands, will provide the second: a kfunc that lets a
scheduler request the node a task's memory should migrate to.

Vladimir Vdovin recently posted an RFC [1] that drives the scan from the
sched_ext tick in the same way and skips the fair-class task migration
for tasks under a BPF scheduler. This series covers what that RFC
proposed and adds the callback and the kfunc on top.

This topic was also discussed at the sched_ext microconference at LPC
2026 [2], in the context of placing accelerator workloads close to their
devices and memory.

Patches 1 and 2 are small preparatory changes to the generic NUMA
balancing code: the first makes the scan available to other scheduling
classes, the second skips the fair-class task migration for tasks owned
by a BPF scheduler and never picks them as swap partners of a migrating
fair task. Patch 3 adds the sched_class method. Apart from that swap
exclusion, which only matters when fair and BPF-scheduled tasks share
CPUs (scx partial mode enabled), none of them changes the behavior of
fair, RT or deadline tasks. Patch 4 adds ops.set_numa_node() and drives
the scan for the schedulers that implement it. Patch 5 adds
scx_bpf_task_numa_nid(), which lets a BPF scheduler read a task's
preferred node from any of its callbacks and returns NUMA_NO_NODE for
the tasks of a scheduler that doesn't implement ops.set_numa_node().
Patch 6 adds a selftest.

Changes in v2:
 - Replace the SCX_OPS_NUMA_BALANCING flag with the ops.set_numa_node()
   callback, which reports preferred node changes as events and enables the
   hinting-fault scan when it is implemented (Tejun Heo)
 - Keep driving the scan from the donor under proxy execution, consistent with
   task_tick_fair() (Vladimir Vdovin)
 - Never pick a task owned by a BPF scheduler as the swap partner of a
   fair task migrated by NUMA balancing
 - Return NUMA_NO_NODE from scx_bpf_task_numa_nid() for tasks whose scheduler
   doesn't implement ops.set_numa_node()
 - Extend the selftest to check that every enabled task gets an initial
   notification before it runs, that changes are reported without duplicates
   and that the notifications agree with scx_bpf_task_numa_nid(), and require
   a preferred node change on a machine with NUMA balancing enabled
 - Link to v1:
   https://lore.kernel.org/r/20261004072901.3579967-1-arighi@nvidia.com

[1] https://lore.kernel.org/r/20261002124559.10367-1-deliran@verdict.gg
[2] https://lpc.events/event/20/sessions/270/

Git tree:
git://git.kernel.org/pub/scm/linux/kernel/git/arighi/linux.git scx-numa-preferred-node

Andrea Righi (6):
  sched/numa: Let other scheduling classes drive NUMA scanning
  sched/numa: Leave the placement of a BPF-scheduled task to its
    scheduler
  sched/numa: Notify the scheduling class of preferred node changes
  sched_ext: Report preferred NUMA nodes through ops.set_numa_node()
  sched_ext: Add a kfunc to read a task's preferred NUMA node
  selftests/sched_ext: Test preferred NUMA node notifications

 kernel/sched/core.c                           |   7 +-
 kernel/sched/ext/ext.c                        |  88 +++++++
 kernel/sched/ext/internal.h                   |  20 ++
 kernel/sched/fair.c                           |  21 +-
 kernel/sched/sched.h                          |  12 +
 tools/sched_ext/include/scx/compat.bpf.h      |  10 +
 tools/testing/selftests/sched_ext/Makefile    |   1 +
 .../selftests/sched_ext/numa_nid.bpf.c        | 142 +++++++++++
 tools/testing/selftests/sched_ext/numa_nid.c  | 233 ++++++++++++++++++
 9 files changed, 527 insertions(+), 7 deletions(-)
 create mode 100644 tools/testing/selftests/sched_ext/numa_nid.bpf.c
 create mode 100644 tools/testing/selftests/sched_ext/numa_nid.c