[RFC PATCH v2 0/3] BPF-based per-cgroup NUMA balancing mode selection

[RFC PATCH v2 0/3] BPF-based per-cgroup NUMA balancing mode selection

[Yafang Shao]

The Problem
===========

In our large fleet of Kubernetes-managed servers, NUMA balancing has
been historically disabled globally on each server. With the increasing
deployment of multi-socket AMD EPYC servers, cross-NUMA access has
become a critical performance issue, which made us reconsider.

Enabling NUMA balancing globally is still not acceptable: it increases
the overall system overhead and can introduce latency spikes for
latency-sensitive workloads. What we need is to choose per workload
whether NUMA balancing runs at all, and which mode it runs in. A
latency-sensitive service may want the scheduler's NUMA statistics
without tiering's promotion machinery, while a batch service wants its
PMEM-resident memory promoted when hot without paying for
scheduler-driven migration. No single global switch serves both.

The Design
==========

This series adds a bpf_numab_ops struct_ops which attaches to a
cgroup. The kernel consults the policies of the matching cgroup and of
every ancestor and combines what they return. One policy question, two
surfaces: the kernel passes the task when it acts on a task and the
folio when it acts on a folio; exactly one of the two is set, so a
memory answer cannot be keyed on whoever happens to touch the folio.
A policy can always opt a workload out, so no cgroup is ever forced
into heavier machinery than its own policy asks for. The demotion and
migration-fallback machinery stays global, as it is keyed on
node-level reclaim state rather than on any cgroup's policy.

The hook is designed to be free for everyone who does not use it.
Every call site (scheduler tick, NUMA hinting faults, folio
migration) first tests the cgroup_bpf_enabled() static key of the
new attach type. The key stays off until the first policy is
attached anywhere in the system, so with no policy attached the
added cost is a single patched branch.

The BPF consult comes before the global sched_numa_balancing static key
rather than after it, because the primary use case is a fleet which runs
with global NUMA balancing disabled and enables it per workload: a policy
must be able to turn balancing on for its cgroup while the global switch
stays off.

Once any policy is attached, the key is on for the whole system and
every task pays the walk: a cgroup lookup, a percpu refcount tryget
and an iteration of the pre-aggregated effective array, so the cost
is proportional to the number of policies in the chain rather than
to its depth, plus one trampoline call per policy. This is also why
patch #2 rejects a policy which implements nothing: attaching one
would flip the static key for the whole system and make every task
pay the walk for no effect. The key turns back off when the last
policy is detached.

Measured with schbench (64 worker threads, 30 s per run) on the test
box. Wakeup latency percentiles from the final cumulative window, in
usec:

                            p50  p99
  balancing=0, no policy    15   207
  balancing=0, policy       15   205
  balancing=1, no policy    15   199
  balancing=1, policy       14   201

With a policy attached, wakeup percentiles stay within a few percent
of the no-policy baseline whether the global switch is on or off, and
throughput stays within the benchmark's own run-to-run variance.

On this box the global on and off states are themselves
indistinguishable for this workload, so these numbers isolate the
cost of the hook, not the benefit of the selection; the benefit side
needs multi-node, memory-heavy workloads.

Changes since v1
================

v1 keyed the hook on the memory cgroup and temporarily embedded the cgroup
id into the struct-ops value, because the cgroup struct_ops integration was
still under development. v2 attaches through the cgroup struct_ops support
that has since landed in bpf-next [0], keys the task surface on the task's
cgroup and the memory surface on the folio's memcg, and adds the
memory-tiering mode alongside enable and disable.

This is still an RFC with limited testing. Feedback is welcome.

[0] https://lore.kernel.org/bpf/20260706171918.317102-1-ameryhung@gmail.com/

RFC v1: https://lwn.net/Articles/1054030/

Yafang Shao (3):
  sched: add helpers for numa balancing
  sched: add per-cgroup NUMA balancing mode selection via BPF
  selftests/bpf: add tests for the bpf_numab_ops struct_ops

 include/linux/bpf-cgroup-defs.h               |   1 +
 include/linux/bpf-cgroup.h                    |   2 +-
 include/linux/bpf_numab.h                     |  77 ++++++
 include/linux/sched/numa_balancing.h          |  82 ++++++
 kernel/sched/Makefile                         |   5 +
 kernel/sched/bpf_numab.c                      | 164 ++++++++++++
 kernel/sched/fair.c                           |  16 +-
 kernel/sched/sched.h                          |   1 -
 mm/memory-tiers.c                             |   9 +-
 mm/mempolicy.c                                |   7 +-
 mm/migrate.c                                  |   7 +-
 mm/vmscan.c                                   |   6 +-
 .../selftests/bpf/prog_tests/bpf_numab_ops.c  | 250 ++++++++++++++++++
 .../selftests/bpf/progs/bpf_numab_ops.c       |  57 ++++
 14 files changed, 660 insertions(+), 24 deletions(-)
 create mode 100644 include/linux/bpf_numab.h
 create mode 100644 kernel/sched/bpf_numab.c
 create mode 100644 tools/testing/selftests/bpf/prog_tests/bpf_numab_ops.c
 create mode 100644 tools/testing/selftests/bpf/progs/bpf_numab_ops.c

-- 
2.52.0

[RFC PATCH v2 1/3] sched: add helpers for numa balancing

[Yafang Shao]

Two new helpers task_numab_enabled() and numab_mode_tiering() are
introduced to replace the scattered direct accesses to the
sched_numa_balancing static key and the sysctl_numa_balancing_mode
bitmask.

No functional change.

Signed-off-by: Yafang Shao <laoar.shao@gmail.com>
---
 include/linux/sched/numa_balancing.h | 31 ++++++++++++++++++++++++++++
 kernel/sched/fair.c                  | 16 +++++++-------
 kernel/sched/sched.h                 |  1 -
 mm/memory-tiers.c                    |  3 ++-
 mm/migrate.c                         |  7 ++++---
 mm/vmscan.c                          |  6 +++---
 6 files changed, 48 insertions(+), 16 deletions(-)

diff --git a/include/linux/sched/numa_balancing.h b/include/linux/sched/numa_balancing.h
index 52b22c5c396d..353772b49c5f 100644
--- a/include/linux/sched/numa_balancing.h
+++ b/include/linux/sched/numa_balancing.h
@@ -8,6 +8,7 @@
  */
 
 #include <linux/sched.h>
+#include <linux/sched/sysctl.h>
 
 #define TNF_MIGRATED	0x01
 #define TNF_NO_GROUP	0x02
@@ -32,6 +33,28 @@ extern void set_numabalancing_state(bool enabled);
 extern void task_numa_free(struct task_struct *p, bool final);
 bool should_numa_migrate_memory(struct task_struct *p, struct folio *folio,
 				int src_nid, int dst_cpu);
+
+extern struct static_key_false sched_numa_balancing;
+
+/*
+ * Whether NUMA balancing applies to @p.  The @p argument is unused for
+ * now: a later patch will consult a per-cgroup BPF policy with it.
+ */
+static inline bool task_numab_enabled(struct task_struct *p)
+{
+	return static_branch_unlikely(&sched_numa_balancing);
+}
+
+/*
+ * The global sysctl state, for the system-wide tiering machinery: this
+ * reflects kernel.numa_balancing_mode only and takes no task.
+ */
+static inline bool numab_mode_tiering(void)
+{
+	if (sysctl_numa_balancing_mode & NUMA_BALANCING_MEMORY_TIERING)
+		return true;
+	return false;
+}
 #else
 static inline void task_numa_fault(int last_node, int node, int pages,
 				   int flags)
@@ -52,6 +75,14 @@ static inline bool should_numa_migrate_memory(struct task_struct *p,
 {
 	return true;
 }
+static inline bool task_numab_enabled(struct task_struct *p)
+{
+	return false;
+}
+static inline bool numab_mode_tiering(void)
+{
+	return false;
+}
 #endif
 
 #endif /* _LINUX_SCHED_NUMA_BALANCING_H */
diff --git a/kernel/sched/fair.c b/kernel/sched/fair.c
index 57360f5cdde4..45086e6f246c 100644
--- a/kernel/sched/fair.c
+++ b/kernel/sched/fair.c
@@ -39,6 +39,7 @@
 #include <linux/sched/cputime.h>
 #include <linux/sched/isolation.h>
 #include <linux/sched/nohz.h>
+#include <linux/sched/numa_balancing.h>
 #include <linux/sched/prio.h>
 #include <linux/static_call.h>
 
@@ -2963,8 +2964,8 @@ bool should_numa_migrate_memory(struct task_struct *p, struct folio *folio,
 	this_cpupid = cpu_pid_to_cpupid(dst_cpu, current->pid);
 	last_cpupid = folio_xchg_last_cpupid(folio, this_cpupid);
 
-	if (!(sysctl_numa_balancing_mode & NUMA_BALANCING_MEMORY_TIERING) &&
-	    !node_is_toptier(src_nid) && !cpupid_valid(last_cpupid))
+	if (!(numab_mode_tiering()) && !node_is_toptier(src_nid) &&
+	    !cpupid_valid(last_cpupid))
 		return false;
 
 	/*
@@ -4194,7 +4195,7 @@ void task_numa_fault(int last_cpupid, int mem_node, int pages, int flags)
 	struct numa_group *ng;
 	int priv;
 
-	if (!static_branch_likely(&sched_numa_balancing))
+	if (!task_numab_enabled(p))
 		return;
 
 	/* for example, ksmd faulting in a user's mm */
@@ -4206,8 +4207,7 @@ void task_numa_fault(int last_cpupid, int mem_node, int pages, int flags)
 	 * node for memory tiering mode.
 	 */
 	if (!node_is_toptier(mem_node) &&
-	    (sysctl_numa_balancing_mode & NUMA_BALANCING_MEMORY_TIERING ||
-	     !cpupid_valid(last_cpupid)))
+	    (numab_mode_tiering() || !cpupid_valid(last_cpupid)))
 		return;
 
 	/* Allocate buffer to track faults on a per-node basis */
@@ -4665,7 +4665,7 @@ static void update_scan_period(struct task_struct *p, int new_cpu)
 	int src_nid = cpu_to_node(task_cpu(p));
 	int dst_nid = cpu_to_node(new_cpu);
 
-	if (!static_branch_likely(&sched_numa_balancing))
+	if (!task_numab_enabled(p))
 		return;
 
 	if (!p->mm || !p->numa_faults || (p->flags & PF_EXITING))
@@ -10697,7 +10697,7 @@ static long migrate_degrades_locality(struct task_struct *p, struct lb_env *env)
 	unsigned long src_weight, dst_weight;
 	int src_nid, dst_nid, dist;
 
-	if (!static_branch_likely(&sched_numa_balancing))
+	if (!task_numab_enabled(p))
 		return 0;
 
 	if (!p->numa_faults || !(env->sd->flags & SD_NUMA))
@@ -15351,7 +15351,7 @@ static void task_tick_fair(struct rq *rq, struct task_struct *curr, int queued)
 	if (queued)
 		return;
 
-	if (static_branch_unlikely(&sched_numa_balancing))
+	if (task_numab_enabled(curr))
 		task_tick_numa(rq, curr);
 
 	task_tick_cache(rq, curr);
diff --git a/kernel/sched/sched.h b/kernel/sched/sched.h
index e656c7059bf8..38e21c8c57dd 100644
--- a/kernel/sched/sched.h
+++ b/kernel/sched/sched.h
@@ -2465,7 +2465,6 @@ extern struct static_key sched_feat_keys[__SCHED_FEAT_NR];
 
 #endif /* !CONFIG_JUMP_LABEL */
 
-extern struct static_key_false sched_numa_balancing;
 extern struct static_key_false sched_schedstats;
 
 static inline u64 global_rt_period(void)
diff --git a/mm/memory-tiers.c b/mm/memory-tiers.c
index 54851d8a195b..ec7b9aadda98 100644
--- a/mm/memory-tiers.c
+++ b/mm/memory-tiers.c
@@ -7,6 +7,7 @@
 #include <linux/memory-tiers.h>
 #include <linux/notifier.h>
 #include <linux/sched/sysctl.h>
+#include <linux/sched/numa_balancing.h>
 
 #include "internal.h"
 
@@ -64,7 +65,7 @@ static const struct bus_type memory_tier_subsys = {
  */
 bool folio_use_access_time(struct folio *folio)
 {
-	return (sysctl_numa_balancing_mode & NUMA_BALANCING_MEMORY_TIERING) &&
+	return numab_mode_tiering() &&
 	       !node_is_toptier(folio_nid(folio));
 }
 #endif
diff --git a/mm/migrate.c b/mm/migrate.c
index 15b45832bcfa..4d2ba87eaa34 100644
--- a/mm/migrate.c
+++ b/mm/migrate.c
@@ -41,6 +41,7 @@
 #include <linux/ptrace.h>
 #include <linux/memory.h>
 #include <linux/sched/sysctl.h>
+#include <linux/sched/numa_balancing.h>
 #include <linux/memory-tiers.h>
 #include <linux/pagewalk.h>
 
@@ -818,7 +819,7 @@ void folio_migrate_flags(struct folio *newfolio, struct folio *folio)
 	 * memory node, reset cpupid, because that is used to record
 	 * page access time in slow memory node.
 	 */
-	if (sysctl_numa_balancing_mode & NUMA_BALANCING_MEMORY_TIERING) {
+	if (numab_mode_tiering()) {
 		bool f_toptier = node_is_toptier(folio_nid(folio));
 		bool t_toptier = node_is_toptier(folio_nid(newfolio));
 
@@ -2718,7 +2719,7 @@ int migrate_misplaced_folio_prepare(struct folio *folio,
 	if (!migrate_balanced_pgdat(pgdat, nr_pages)) {
 		int z;
 
-		if (!(sysctl_numa_balancing_mode & NUMA_BALANCING_MEMORY_TIERING))
+		if (!(numab_mode_tiering()))
 			return -EAGAIN;
 		for (z = pgdat->nr_zones - 1; z >= 0; z--) {
 			if (managed_zone(pgdat->node_zones + z))
@@ -2770,7 +2771,7 @@ int migrate_misplaced_folio(struct folio *folio, int node)
 	if (nr_succeeded) {
 		count_vm_numa_events(NUMA_PAGE_MIGRATE, nr_succeeded);
 		count_memcg_events(memcg, NUMA_PAGE_MIGRATE, nr_succeeded);
-		if ((sysctl_numa_balancing_mode & NUMA_BALANCING_MEMORY_TIERING)
+		if ((numab_mode_tiering())
 		    && !node_is_toptier(folio_nid(folio))
 		    && node_is_toptier(node))
 			mod_lruvec_state(lruvec, PGPROMOTE_SUCCESS, nr_succeeded);
diff --git a/mm/vmscan.c b/mm/vmscan.c
index f11491ee9ed5..af46869bb15e 100644
--- a/mm/vmscan.c
+++ b/mm/vmscan.c
@@ -65,6 +65,7 @@
 
 #include <linux/swapops.h>
 #include <linux/sched/sysctl.h>
+#include <linux/sched/numa_balancing.h>
 
 #include "internal.h"
 #include "page_alloc.h"
@@ -5029,8 +5030,7 @@ static bool should_abort_scan(struct lruvec *lruvec, struct scan_control *sc)
 	if (!current_is_kswapd() || sc->order)
 		return false;
 
-	mark = sysctl_numa_balancing_mode & NUMA_BALANCING_MEMORY_TIERING ?
-	       WMARK_PROMO : WMARK_HIGH;
+	mark = numab_mode_tiering() ? WMARK_PROMO : WMARK_HIGH;
 
 	for (i = 0; i <= sc->reclaim_idx; i++) {
 		struct zone *zone = lruvec_pgdat(lruvec)->node_zones + i;
@@ -6974,7 +6974,7 @@ static bool pgdat_balanced(pg_data_t *pgdat, int order, int highest_zoneidx)
 		enum zone_stat_item item;
 		unsigned long free_pages;
 
-		if (sysctl_numa_balancing_mode & NUMA_BALANCING_MEMORY_TIERING)
+		if (numab_mode_tiering())
 			mark = promo_wmark_pages(zone);
 		else
 			mark = high_wmark_pages(zone);
-- 
2.52.0

Re: [RFC PATCH v2 1/3] sched: add helpers for numa balancing

[sashiko-bot]

Thank you for your contribution! Sashiko AI review found 2 potential issue(s) to consider:
- [Medium] sched: Incomplete refactoring leaves direct sched_numa_balancing accesses
- [Low] sched: performance regression in task_numa_fault() due to incorrect static branch hint
--

commit 6333e53c9f6ee67bda34cc209d9f45e01a0af5fd
Author: Yafang Shao <laoar.shao@gmail.com>

    sched: add helpers for numa balancing
    
    Two new helpers task_numab_enabled() and numab_mode_tiering() are
    introduced to replace the scattered direct accesses to the
    sched_numa_balancing static key and the sysctl_numa_balancing_mode
    bitmask.

> diff --git a/include/linux/sched/numa_balancing.h b/include/linux/sched/numa_balancing.h
> index 52b22c5c396db..353772b49c5f9 100644
> --- a/include/linux/sched/numa_balancing.h
> +++ b/include/linux/sched/numa_balancing.h
> @@ -32,6 +33,28 @@ extern void set_numabalancing_state(bool enabled);
>  extern void task_numa_free(struct task_struct *p, bool final);
>  bool should_numa_migrate_memory(struct task_struct *p, struct folio *folio,
>  				int src_nid, int dst_cpu);
> +
> +extern struct static_key_false sched_numa_balancing;
> +
> +/*
> + * Whether NUMA balancing applies to @p.  The @p argument is unused for
> + * now: a later patch will consult a per-cgroup BPF policy with it.
> + */
> +static inline bool task_numab_enabled(struct task_struct *p)
> +{
> +	return static_branch_unlikely(&sched_numa_balancing);
> +}

[Severity: Low]
Does this wrapper unintentionally flip the branch prediction hint from
likely to unlikely on hot paths?

For instance, task_numa_fault() is almost exclusively called from NUMA
hinting fault handlers, which only trigger when NUMA balancing is active.
Before this patch, it used static_branch_likely():

-	if (!static_branch_likely(&sched_numa_balancing))
+	if (!task_numab_enabled(p))

Because the new task_numab_enabled() wrapper hardcodes
static_branch_unlikely(), could the compiler now outline the body of
task_numa_fault() and other functions that expect the condition to usually
be true?

[ ... ]

> diff --git a/kernel/sched/fair.c b/kernel/sched/fair.c
> index c52c29dcb6a2e..451e26da5af24 100644
> --- a/kernel/sched/fair.c
> +++ b/kernel/sched/fair.c
> @@ -4199,7 +4200,7 @@ void task_numa_fault(int last_cpupid, int mem_node, int pages, int flags)
>  	struct numa_group *ng;
>  	int priv;
>  
> -	if (!static_branch_likely(&sched_numa_balancing))
> +	if (!task_numab_enabled(p))
>  		return;

[Severity: Medium]
Are there a few direct accesses to the sched_numa_balancing static key in
kernel/sched/fair.c that might have been missed by this refactoring?

For example:

kernel/sched/fair.c:exceed_llc_capacity() {
    ...
	if (static_branch_likely(&sched_numa_balancing)) {
    ...
}

kernel/sched/fair.c:get_pref_llc() {
    ...
		if (static_branch_likely(&sched_numa_balancing) &&
		    p->numa_preferred_nid >= 0 &&
    ...
}

kernel/sched/fair.c:get_scan_cpumasks() {
    ...
	if (!static_branch_likely(&sched_numa_balancing))
		goto out;
    ...
}

Since get_pref_llc() and get_scan_cpumasks() already take a task_struct
pointer 'p', shouldn't they also be updated to use the new
task_numab_enabled(p) wrapper?

If these direct checks remain, might they bypass the per-cgroup BPF policy
that the commit message notes will be consulted in a later patch?

-- 
Sashiko AI review · https://sashiko.dev/#/patchset/20261012120438.52849-1-laoar.shao@gmail.com?part=1

Re: [RFC PATCH v2 1/3] sched: add helpers for numa balancing

[Yafang Shao]

On Mon, Oct 12, 2026 at 8:19 PM <sashiko-bot@kernel.org> wrote:
>
> Thank you for your contribution! Sashiko AI review found 2 potential issue(s) to consider:
> - [Medium] sched: Incomplete refactoring leaves direct sched_numa_balancing accesses
> - [Low] sched: performance regression in task_numa_fault() due to incorrect static branch hint
> --
>
> commit 6333e53c9f6ee67bda34cc209d9f45e01a0af5fd
> Author: Yafang Shao <laoar.shao@gmail.com>
>
>     sched: add helpers for numa balancing
>
>     Two new helpers task_numab_enabled() and numab_mode_tiering() are
>     introduced to replace the scattered direct accesses to the
>     sched_numa_balancing static key and the sysctl_numa_balancing_mode
>     bitmask.
>
> > diff --git a/include/linux/sched/numa_balancing.h b/include/linux/sched/numa_balancing.h
> > index 52b22c5c396db..353772b49c5f9 100644
> > --- a/include/linux/sched/numa_balancing.h
> > +++ b/include/linux/sched/numa_balancing.h
> > @@ -32,6 +33,28 @@ extern void set_numabalancing_state(bool enabled);
> >  extern void task_numa_free(struct task_struct *p, bool final);
> >  bool should_numa_migrate_memory(struct task_struct *p, struct folio *folio,
> >                               int src_nid, int dst_cpu);
> > +
> > +extern struct static_key_false sched_numa_balancing;
> > +
> > +/*
> > + * Whether NUMA balancing applies to @p.  The @p argument is unused for
> > + * now: a later patch will consult a per-cgroup BPF policy with it.
> > + */
> > +static inline bool task_numab_enabled(struct task_struct *p)
> > +{
> > +     return static_branch_unlikely(&sched_numa_balancing);
> > +}
>
> [Severity: Low]
> Does this wrapper unintentionally flip the branch prediction hint from
> likely to unlikely on hot paths?
>
> For instance, task_numa_fault() is almost exclusively called from NUMA
> hinting fault handlers, which only trigger when NUMA balancing is active.
> Before this patch, it used static_branch_likely():
>
> -       if (!static_branch_likely(&sched_numa_balancing))
> +       if (!task_numab_enabled(p))
>
> Because the new task_numab_enabled() wrapper hardcodes
> static_branch_unlikely(), could the compiler now outline the body of
> task_numa_fault() and other functions that expect the condition to usually
> be true?

Nice catch.

They are all static_branch_likely(&sched_numa_balancing) except
the one in task_tick_fair().  Since NUMA balancing is enabled by
default:

config NUMA_BALANCING_DEFAULT_ENABLED
        bool "Automatically enable NUMA aware memory/task placement"
        default y
        depends on NUMA_BALANCING
        help
          If set, automatic NUMA balancing will be enabled if running on a NUMA
          machine.

I believe we'd better make the change below first and then switch
the helper to static_branch_likely().

--- a/kernel/sched/fair.c
+++ b/kernel/sched/fair.c
@@ -15082,7 +15082,7 @@ static void task_tick_fair(struct rq *rq,
struct task_struct *curr, int queued)
        if (queued)
                return;

-       if (static_branch_unlikely(&sched_numa_balancing))
+       if (static_branch_likely(&sched_numa_balancing))
                task_tick_numa(rq, curr);

        task_tick_cache(rq, curr);

>
> [ ... ]
>
> > diff --git a/kernel/sched/fair.c b/kernel/sched/fair.c
> > index c52c29dcb6a2e..451e26da5af24 100644
> > --- a/kernel/sched/fair.c
> > +++ b/kernel/sched/fair.c
> > @@ -4199,7 +4200,7 @@ void task_numa_fault(int last_cpupid, int mem_node, int pages, int flags)
> >       struct numa_group *ng;
> >       int priv;
> >
> > -     if (!static_branch_likely(&sched_numa_balancing))
> > +     if (!task_numab_enabled(p))
> >               return;
>
> [Severity: Medium]
> Are there a few direct accesses to the sched_numa_balancing static key in
> kernel/sched/fair.c that might have been missed by this refactoring?
>
> For example:
>
> kernel/sched/fair.c:exceed_llc_capacity() {
>     ...
>         if (static_branch_likely(&sched_numa_balancing)) {
>     ...
> }
>
> kernel/sched/fair.c:get_pref_llc() {
>     ...
>                 if (static_branch_likely(&sched_numa_balancing) &&
>                     p->numa_preferred_nid >= 0 &&
>     ...
> }
>
> kernel/sched/fair.c:get_scan_cpumasks() {
>     ...
>         if (!static_branch_likely(&sched_numa_balancing))
>                 goto out;
>     ...
> }
>
> Since get_pref_llc() and get_scan_cpumasks() already take a task_struct
> pointer 'p', shouldn't they also be updated to use the new
> task_numab_enabled(p) wrapper?
>
> If these direct checks remain, might they bypass the per-cgroup BPF policy
> that the commit message notes will be consulted in a later patch?

I will update it.

-- 
Regards
Yafang

Re: [RFC PATCH v2 1/3] sched: add helpers for numa balancing

[David Hildenbrand (Arm)]

On 10/12/26 14:04, Yafang Shao wrote:
> Two new helpers task_numab_enabled() and numab_mode_tiering() are
> introduced to replace the scattered direct accesses to the
> sched_numa_balancing static key and the sysctl_numa_balancing_mode
> bitmask.
> 

I'll only comment on the new exposed helpers. Some of these might make sense
even without the bigger series.

(obviously, minus the task change)

> No functional change.
> 
> Signed-off-by: Yafang Shao <laoar.shao@gmail.com>
> ---
>  include/linux/sched/numa_balancing.h | 31 ++++++++++++++++++++++++++++
>  kernel/sched/fair.c                  | 16 +++++++-------
>  kernel/sched/sched.h                 |  1 -
>  mm/memory-tiers.c                    |  3 ++-
>  mm/migrate.c                         |  7 ++++---
>  mm/vmscan.c                          |  6 +++---
>  6 files changed, 48 insertions(+), 16 deletions(-)
> 
> diff --git a/include/linux/sched/numa_balancing.h b/include/linux/sched/numa_balancing.h
> index 52b22c5c396d..353772b49c5f 100644
> --- a/include/linux/sched/numa_balancing.h
> +++ b/include/linux/sched/numa_balancing.h
> @@ -8,6 +8,7 @@
>   */
>  
>  #include <linux/sched.h>
> +#include <linux/sched/sysctl.h>
>  
>  #define TNF_MIGRATED	0x01
>  #define TNF_NO_GROUP	0x02
> @@ -32,6 +33,28 @@ extern void set_numabalancing_state(bool enabled);
>  extern void task_numa_free(struct task_struct *p, bool final);
>  bool should_numa_migrate_memory(struct task_struct *p, struct folio *folio,
>  				int src_nid, int dst_cpu);
> +
> +extern struct static_key_false sched_numa_balancing;
> +
> +/*
> + * Whether NUMA balancing applies to @p.  The @p argument is unused for
> + * now: a later patch will consult a per-cgroup BPF policy with it.

Talking about later patches in a function doc is ... wrong.

> + */
> +static inline bool task_numab_enabled(struct task_struct *p)
> +{
> +	return static_branch_unlikely(&sched_numa_balancing);
> +}
> +
> +/*
> + * The global sysctl state, for the system-wide tiering machinery: this
> + * reflects kernel.numa_balancing_mode only and takes no task.

The latter part doesn't make a lot of sense in this documentation.

> + */
> +static inline bool numab_mode_tiering(void)

The function name isn't really intuitive. Should this be

	numab_in_tiering_mode

	
	numab_tiering_enabled

or something like this.

> +{
> +	if (sysctl_numa_balancing_mode & NUMA_BALANCING_MEMORY_TIERING)
> +		return true;
> +	return false;


return sysctl_numa_balancing_mode & NUMA_BALANCING_MEMORY_TIERING;

?

Both functions could directly have proper kerneldoc if we are already
documenting these.

> +}
>  #else
>  static inline void task_numa_fault(int last_node, int node, int pages,
>  				   int flags)
> @@ -52,6 +75,14 @@ static inline bool should_numa_migrate_memory(struct task_struct *p,
>  {
>  	return true;
>  }
> +static inline bool task_numab_enabled(struct task_struct *p)
> +{
> +	return false;
> +}
> +static inline bool numab_mode_tiering(void)
> +{
> +	return false;
> +}


-- 
Cheers,

David

Re: [RFC PATCH v2 1/3] sched: add helpers for numa balancing

[Yafang Shao]

On Tue, Oct 13, 2026 at 3:38 PM David Hildenbrand (Arm)
<david@kernel.org> wrote:
>
> On 10/12/26 14:04, Yafang Shao wrote:
> > Two new helpers task_numab_enabled() and numab_mode_tiering() are
> > introduced to replace the scattered direct accesses to the
> > sched_numa_balancing static key and the sysctl_numa_balancing_mode
> > bitmask.
> >
>
> I'll only comment on the new exposed helpers. Some of these might make sense
> even without the bigger series.
>
> (obviously, minus the task change)
>
> > No functional change.
> >
> > Signed-off-by: Yafang Shao <laoar.shao@gmail.com>
> > ---
> >  include/linux/sched/numa_balancing.h | 31 ++++++++++++++++++++++++++++
> >  kernel/sched/fair.c                  | 16 +++++++-------
> >  kernel/sched/sched.h                 |  1 -
> >  mm/memory-tiers.c                    |  3 ++-
> >  mm/migrate.c                         |  7 ++++---
> >  mm/vmscan.c                          |  6 +++---
> >  6 files changed, 48 insertions(+), 16 deletions(-)
> >
> > diff --git a/include/linux/sched/numa_balancing.h b/include/linux/sched/numa_balancing.h
> > index 52b22c5c396d..353772b49c5f 100644
> > --- a/include/linux/sched/numa_balancing.h
> > +++ b/include/linux/sched/numa_balancing.h
> > @@ -8,6 +8,7 @@
> >   */
> >
> >  #include <linux/sched.h>
> > +#include <linux/sched/sysctl.h>
> >
> >  #define TNF_MIGRATED 0x01
> >  #define TNF_NO_GROUP 0x02
> > @@ -32,6 +33,28 @@ extern void set_numabalancing_state(bool enabled);
> >  extern void task_numa_free(struct task_struct *p, bool final);
> >  bool should_numa_migrate_memory(struct task_struct *p, struct folio *folio,
> >                               int src_nid, int dst_cpu);
> > +
> > +extern struct static_key_false sched_numa_balancing;
> > +
> > +/*
> > + * Whether NUMA balancing applies to @p.  The @p argument is unused for
> > + * now: a later patch will consult a per-cgroup BPF policy with it.
>
> Talking about later patches in a function doc is ... wrong.

Thanks for your guidance.
I will update it.

>
> > + */
> > +static inline bool task_numab_enabled(struct task_struct *p)
> > +{
> > +     return static_branch_unlikely(&sched_numa_balancing);
> > +}
> > +
> > +/*
> > + * The global sysctl state, for the system-wide tiering machinery: this
> > + * reflects kernel.numa_balancing_mode only and takes no task.
>
> The latter part doesn't make a lot of sense in this documentation.

I will update it.

>
> > + */
> > +static inline bool numab_mode_tiering(void)
>
> The function name isn't really intuitive. Should this be
>
>         numab_in_tiering_mode
>
>
>         numab_tiering_enabled
>
> or something like this.

Agreed, thanks.  I will rename it to numab_in_tiering_mode().

>
> > +{
> > +     if (sysctl_numa_balancing_mode & NUMA_BALANCING_MEMORY_TIERING)
> > +             return true;
> > +     return false;
>
>
> return sysctl_numa_balancing_mode & NUMA_BALANCING_MEMORY_TIERING;
>
> ?

Agreed, will simplify as suggested.

>
> Both functions could directly have proper kerneldoc if we are already
> documenting these.

Sure. I will update it.

-- 
Regards
Yafang

[RFC PATCH v2 2/3] sched: add per-cgroup NUMA balancing mode selection via BPF

[Yafang Shao]

Global NUMA balancing is all-or-nothing: kernel.numa_balancing and
kernel.numa_balancing_mode pick the mode for every workload.  A
latency-sensitive service wants the scheduler's NUMA statistics but
not tiering's promotion machinery; a batch service wants its
PMEM-resident memory promoted when hot without paying for scheduler
driven migration.  No global switch serves both.

Add a bpf_numab_ops struct_ops. A policy attached to a cgroup picks a
mode for that cgroup; the kernel runs the policies of the matching
cgroup and of every ancestor and combines what they return.  One
policy question, two surfaces: the kernel passes the task when it acts
on a task (scanner, fault statistics) and the folio when it acts on a
folio (cpupid vs access-time bookkeeping, scanning exemptions);
exactly one of the two is set, so a memory answer cannot be keyed on
whoever happens to touch the folio.

The memory surface converges on folio_numab_mode_tiering(): it decides
folio_use_access_time(), the top-tier scanning exemption, the
cross-tier cpupid reset on migration and the cpupid migration guard.
The task surface: task_numab_enabled() reads the combined request
directly, and task_numab_mode_tiering() decides whether faults on
non-top-tier memory feed the statistics. Both compose with the global
fallback numab_mode_tiering(), so the series is a no-op until a policy
is attached.

DISABLED beats NORMAL beats TIERING: a policy can always opt out, so
no cgroup is ever forced into heavier machinery than its own policy
asks for.  The demotion and migration-fallback machinery stays global:
it is keyed on node-level reclaim state, not on any cgroup's policy; a
TIERING cgroup under a normal-mode system gets promotion without
demotion.

A policy which implements nothing makes every task pay for the enabled
static key, so bpf_numab_validate() rejects it.

Signed-off-by: Yafang Shao <laoar.shao@gmail.com>
---
 include/linux/bpf-cgroup-defs.h      |   1 +
 include/linux/bpf-cgroup.h           |   2 +-
 include/linux/bpf_numab.h            |  77 +++++++++++++
 include/linux/sched/numa_balancing.h |  65 +++++++++--
 kernel/sched/Makefile                |   5 +
 kernel/sched/bpf_numab.c             | 164 +++++++++++++++++++++++++++
 kernel/sched/fair.c                  |   4 +-
 mm/memory-tiers.c                    |   8 +-
 mm/mempolicy.c                       |   7 +-
 mm/migrate.c                         |   4 +-
 10 files changed, 317 insertions(+), 20 deletions(-)
 create mode 100644 include/linux/bpf_numab.h
 create mode 100644 kernel/sched/bpf_numab.c

diff --git a/include/linux/bpf-cgroup-defs.h b/include/linux/bpf-cgroup-defs.h
index 0147b8bec973..c0056d170a1b 100644
--- a/include/linux/bpf-cgroup-defs.h
+++ b/include/linux/bpf-cgroup-defs.h
@@ -48,6 +48,7 @@ enum cgroup_bpf_attach_type {
 	CGROUP_UNIX_GETSOCKNAME,
 	CGROUP_INET_SOCK_RELEASE,
 	CGROUP_TCP_SOCK_OPS,
+	CGROUP_NUMAB_OPS,
 	CGROUP_LSM_START,
 	CGROUP_LSM_END = CGROUP_LSM_START + CGROUP_LSM_NUM - 1,
 	MAX_CGROUP_BPF_ATTACH_TYPE
diff --git a/include/linux/bpf-cgroup.h b/include/linux/bpf-cgroup.h
index 8a75a6cd7309..bbfa171e449c 100644
--- a/include/linux/bpf-cgroup.h
+++ b/include/linux/bpf-cgroup.h
@@ -119,7 +119,7 @@ struct bpf_prog_list {
 
 static inline bool cgroup_bpf_is_struct_ops_atype(enum cgroup_bpf_attach_type atype)
 {
-	return atype == CGROUP_TCP_SOCK_OPS;
+	return atype == CGROUP_TCP_SOCK_OPS || atype == CGROUP_NUMAB_OPS;
 }
 void cgroup_bpf_struct_ops_register(int atype, u32 type_id, void *cfi_stubs, bool mult_trace);
 int cgroup_bpf_struct_ops_attach(struct bpf_map *map, const union bpf_attr *attr);
diff --git a/include/linux/bpf_numab.h b/include/linux/bpf_numab.h
new file mode 100644
index 000000000000..7f78eef1c159
--- /dev/null
+++ b/include/linux/bpf_numab.h
@@ -0,0 +1,77 @@
+/* SPDX-License-Identifier: GPL-2.0 */
+/*
+ * BPF policy hooks for NUMA balancing.
+ *
+ * A bpf_numab_ops is attached to a cgroup. A NUMA balancing decision runs
+ * the policies of that cgroup and of every ancestor, and the kernel combines
+ * what they return. BPF only picks between things the kernel already does:
+ * with global NUMA balancing disabled, a policy can enable it per workload;
+ * with global NUMA balancing enabled, a policy can disable it for
+ * latency-sensitive workloads or opt their memory into tiering.
+ */
+#ifndef _LINUX_BPF_NUMAB_H
+#define _LINUX_BPF_NUMAB_H
+
+#include <linux/types.h>
+
+struct task_struct;
+struct folio;
+
+/*
+ * What a policy can ask for NUMA balancing. The kernel ORs the requests
+ * of all matching policies, so one policy cannot undo another's opt-in.
+ * When several requests combine, the conservative choice wins:
+ * BPF_NUMAB_DISABLED beats BPF_NUMAB_NORMAL, which beats BPF_NUMAB_TIERING.
+ * A policy can always opt a workload out, so no cgroup is ever forced
+ * into heavier machinery than its own policy asks for.
+ *
+ * The mode bits mirror the global sysctl_numa_balancing_mode semantics:
+ * BPF_NUMAB_NORMAL asks for the normal mode, where hinting faults feed
+ * the scheduler's NUMA statistics. BPF_NUMAB_TIERING asks for the
+ * memory tiering mode, where the cgroup's non-top-tier memory records
+ * access times in place of cpupids and is promoted when hot.
+ * BPF_NUMAB_GLOBAL (0) makes no request and defers to the global sysctl
+ * state.
+ */
+enum bpf_numab_mode {
+	BPF_NUMAB_GLOBAL	= 0,
+	BPF_NUMAB_DISABLED	= 1U << 0,
+	BPF_NUMAB_NORMAL	= 1U << 1,
+	BPF_NUMAB_TIERING	= 1U << 2,
+};
+
+#define BPF_NUMAB_VALID_MASK	(BPF_NUMAB_DISABLED | BPF_NUMAB_NORMAL | \
+				 BPF_NUMAB_TIERING)
+
+struct bpf_numab_ops {
+	/**
+	 * numab_mode - request a NUMA balancing mode
+	 * @p: task the NUMA balancing machinery is about to act on, or NULL
+	 * @folio: folio whose NUMA balancing semantics are being decided,
+	 *	   or NULL
+	 *
+	 * Exactly one of @p and @folio is set. The kernel passes @p when
+	 * it acts on a task (whether its scanner runs, whether a fault
+	 * feeds its statistics) and @folio when it acts on a folio
+	 * (whether its memory carries cpupids or access times, whether its
+	 * top-tier pages are worth scanning). A policy must NULL-check
+	 * before dereferencing.
+	 *
+	 * Called from the NUMA balancing hot paths (hinting faults,
+	 * scanner passes, load balancing). It must not sleep.
+	 *
+	 * Return: bits from enum bpf_numab_mode.  Other bits are dropped.
+	 */
+	u32 (*numab_mode)(struct task_struct *p, struct folio *folio);
+};
+
+#if defined(CONFIG_CGROUP_BPF) && defined(CONFIG_NUMA_BALANCING)
+u32 bpf_numab_mode(struct task_struct *p, struct folio *folio);
+#else
+static inline u32 bpf_numab_mode(struct task_struct *p, struct folio *folio)
+{
+	return BPF_NUMAB_GLOBAL;
+}
+#endif
+
+#endif /* _LINUX_BPF_NUMAB_H */
diff --git a/include/linux/sched/numa_balancing.h b/include/linux/sched/numa_balancing.h
index 353772b49c5f..fb85e968b1f3 100644
--- a/include/linux/sched/numa_balancing.h
+++ b/include/linux/sched/numa_balancing.h
@@ -7,6 +7,7 @@
  * implements memory access pattern based NUMA-balancing:
  */
 
+#include <linux/bpf_numab.h>
 #include <linux/sched.h>
 #include <linux/sched/sysctl.h>
 
@@ -36,25 +37,67 @@ bool should_numa_migrate_memory(struct task_struct *p, struct folio *folio,
 
 extern struct static_key_false sched_numa_balancing;
 
-/*
- * Whether NUMA balancing applies to @p.  The @p argument is unused for
- * now: a later patch will consult a per-cgroup BPF policy with it.
+/* Whether NUMA balancing applies to @p: the global sysctl state, possibly
+ * overridden by a BPF policy attached to @p's cgroup. The requests of
+ * the task's cgroup and of every ancestor are combined; a policy can
+ * always opt out, so a latency-sensitive workload is never forced in.
  */
 static inline bool task_numab_enabled(struct task_struct *p)
 {
+	u32 state = bpf_numab_mode(p, NULL);
+
+	/* A policy can always opt out, regardless of the global state. */
+	if (state & BPF_NUMAB_DISABLED)
+		return false;
+	if (state & (BPF_NUMAB_NORMAL | BPF_NUMAB_TIERING))
+		return true;
 	return static_branch_unlikely(&sched_numa_balancing);
 }
 
-/*
- * The global sysctl state, for the system-wide tiering machinery: this
- * reflects kernel.numa_balancing_mode only and takes no task.
- */
 static inline bool numab_mode_tiering(void)
 {
 	if (sysctl_numa_balancing_mode & NUMA_BALANCING_MEMORY_TIERING)
 		return true;
 	return false;
 }
+
+/*
+ * Whether @p runs in tiering mode, i.e. whether faults on non-top-tier
+ * memory still feed its NUMA statistics. The requests of the task's
+ * cgroup and of every ancestor are combined; DISABLED and NORMAL both
+ * opt out of tiering (a policy can always opt out), so only an
+ * unopposed TIERING request engages it. Without any request, the
+ * global sysctl decides.
+ */
+static inline bool task_numab_mode_tiering(struct task_struct *p)
+{
+	u32 state = bpf_numab_mode(p, NULL);
+
+	if (state & (BPF_NUMAB_DISABLED | BPF_NUMAB_NORMAL))
+		return false;
+	if (state & BPF_NUMAB_TIERING)
+		return true;
+	return numab_mode_tiering();
+}
+
+/*
+ * Whether @folio's memory participates in tiering, i.e. whether its
+ * _last_cpupid carries an access time instead of a cpupid. The requests
+ * of the owner cgroup and of every ancestor are combined. DISABLED
+ * never speaks for memory (it is a task-surface request), so only
+ * NORMAL and TIERING are consulted here; without either, the global
+ * sysctl decides.
+ */
+static inline bool folio_numab_mode_tiering(struct folio *folio)
+{
+	u32 state = bpf_numab_mode(NULL, folio);
+
+	if (state & BPF_NUMAB_NORMAL)
+		return false;
+	if (state & BPF_NUMAB_TIERING)
+		return true;
+	return numab_mode_tiering();
+}
 #else
 static inline void task_numa_fault(int last_node, int node, int pages,
 				   int flags)
@@ -83,6 +126,14 @@ static inline bool numab_mode_tiering(void)
 {
 	return false;
 }
+static inline bool task_numab_mode_tiering(struct task_struct *p)
+{
+	return false;
+}
+static inline bool folio_numab_mode_tiering(struct folio *folio)
+{
+	return false;
+}
 #endif
 
 #endif /* _LINUX_SCHED_NUMA_BALANCING_H */
diff --git a/kernel/sched/Makefile b/kernel/sched/Makefile
index b1f1a367034f..688d4cd12c4f 100644
--- a/kernel/sched/Makefile
+++ b/kernel/sched/Makefile
@@ -40,3 +40,8 @@ obj-y += core.o
 obj-y += fair.o
 obj-y += build_policy.o
 obj-y += build_utility.o
+
+# BPF-based per-cgroup NUMA balancing
+ifdef CONFIG_NUMA_BALANCING
+obj-$(CONFIG_CGROUP_BPF) += bpf_numab.o
+endif
diff --git a/kernel/sched/bpf_numab.c b/kernel/sched/bpf_numab.c
new file mode 100644
index 000000000000..924adf30889e
--- /dev/null
+++ b/kernel/sched/bpf_numab.c
@@ -0,0 +1,164 @@
+// SPDX-License-Identifier: GPL-2.0-or-later
+/*
+ * BPF-based per-cgroup NUMA balancing
+ *
+ * A bpf_numab_ops is attached to a cgroup via a struct_ops link. Whenever
+ * the kernel is about to act on NUMA balancing, it runs the policies of the
+ * matching cgroup and of every ancestor, and combines what they return: for
+ * a task, the chain of the task's cgroup; for a folio, the chain of the
+ * memcg the folio is charged to.
+ *
+ * Author: Yafang Shao <laoar.shao@gmail.com>
+ */
+#include <linux/bpf.h>
+#include <linux/bpf-cgroup.h>
+#include <linux/bpf_numab.h>
+#include <linux/bpf_verifier.h>
+#include <linux/cgroup.h>
+#include <linux/memcontrol.h>
+#include <linux/sched.h>
+
+/*
+ * CFI stubs. These really run: a slot points at them while its policy is
+ * being detached. Return BPF_NUMAB_GLOBAL, the identity for the
+ * kernel's OR.
+ */
+static u32 numab_mode_stub(struct task_struct *p, struct folio *folio)
+{
+	return BPF_NUMAB_GLOBAL;
+}
+
+static struct bpf_numab_ops __bpf_numab_ops = {
+	.numab_mode = numab_mode_stub,
+};
+
+static const struct bpf_func_proto *
+bpf_numab_get_func_proto(enum bpf_func_id func_id, const struct bpf_prog *prog)
+{
+	return bpf_base_func_proto(func_id, prog);
+}
+
+static bool bpf_numab_is_valid_access(int off, int size,
+				      enum bpf_access_type type,
+				      const struct bpf_prog *prog,
+				      struct bpf_insn_access_aux *info)
+{
+	if (type != BPF_READ)
+		return false;
+	return bpf_tracing_btf_ctx_access(off, size, type, prog, info);
+}
+
+static int bpf_numab_init_member(const struct btf_type *t,
+				 const struct btf_member *member,
+				 void *kdata, const void *udata)
+{
+	/* Mandatory: the core calls it without a NULL check. */
+	return 0;
+}
+
+static int bpf_numab_check_member(const struct btf_type *t,
+				  const struct btf_member *member,
+				  const struct bpf_prog *prog)
+{
+	/*
+	 * Members also run from the scheduler tick and from load
+	 * balancing, which hold the rq lock, so they can never sleep.
+	 */
+	if (prog->sleepable)
+		return -EINVAL;
+	return 0;
+}
+
+static int bpf_numab_init(struct btf *btf)
+{
+	return 0;
+}
+
+static int bpf_numab_validate(void *kdata)
+{
+	struct bpf_numab_ops *ops = kdata;
+
+	/*
+	 * A policy which implements nothing has no effect but still makes
+	 * every task pay for the enabled static key. Reject it.
+	 */
+	if (!ops->numab_mode)
+		return -EINVAL;
+	return 0;
+}
+
+static const struct bpf_verifier_ops bpf_numab_verifier_ops = {
+	.get_func_proto		= bpf_numab_get_func_proto,
+	.is_valid_access	= bpf_numab_is_valid_access,
+};
+
+static struct bpf_struct_ops bpf_numab_ops_desc = {
+	.verifier_ops	= &bpf_numab_verifier_ops,
+	.init		= bpf_numab_init,
+	.init_member	= bpf_numab_init_member,
+	.check_member	= bpf_numab_check_member,
+	.validate	= bpf_numab_validate,
+	.name		= "bpf_numab_ops",
+	.cgroup_atype	= CGROUP_NUMAB_OPS,
+	.cfi_stubs	= &__bpf_numab_ops,
+	.owner		= THIS_MODULE,
+};
+
+/*
+ * Run every numab_mode on the matching cgroup and its ancestors, and
+ * return the combined request for the caller to act on. A task question
+ * walks @p's cgroup; a folio question walks the chain of the memcg @folio
+ * is charged to, and an uncharged folio makes no request.
+ */
+u32 bpf_numab_mode(struct task_struct *p, struct folio *folio)
+{
+	const struct bpf_prog_array_item *item;
+	const struct bpf_numab_ops *ops;
+	struct mem_cgroup *memcg;
+	struct cgroup *cgrp;
+	u32 acc = BPF_NUMAB_GLOBAL;
+
+	if (!cgroup_bpf_enabled(CGROUP_NUMAB_OPS))
+		return acc;
+
+	/* Kernel threads have no user memory to balance. */
+	if (p && !p->mm)
+		return acc;
+
+	rcu_read_lock();
+
+	if (p) {
+		cgrp = task_dfl_cgroup(p);
+	} else {
+		memcg = folio_memcg(folio);
+		cgrp = memcg ? memcg->css.cgroup : NULL;
+	}
+
+	/*
+	 * A cgroup outlives its attachments while it is live, and
+	 * cgroup_bpf_release() frees the arrays when the cgroup goes.
+	 */
+	if (!cgrp || !percpu_ref_tryget_live_rcu(&cgrp->bpf.refcnt))
+		goto out;
+
+	bpf_cgroup_struct_ops_foreach(ops, item, cgrp, CGROUP_NUMAB_OPS) {
+		if (ops->numab_mode)
+			acc |= ops->numab_mode(p, folio) & BPF_NUMAB_VALID_MASK;
+	}
+
+	cgroup_bpf_put(cgrp);
+out:
+	rcu_read_unlock();
+	return acc;
+}
+
+static int __init bpf_numab_ops_register(void)
+{
+	int err;
+
+	err = register_bpf_struct_ops(&bpf_numab_ops_desc, bpf_numab_ops);
+	if (err)
+		pr_err("bpf_numab: Failed to register struct_ops (%d)\n", err);
+	return err;
+}
+late_initcall(bpf_numab_ops_register);
diff --git a/kernel/sched/fair.c b/kernel/sched/fair.c
index 45086e6f246c..2cabe8c3ec1e 100644
--- a/kernel/sched/fair.c
+++ b/kernel/sched/fair.c
@@ -2964,7 +2964,7 @@ bool should_numa_migrate_memory(struct task_struct *p, struct folio *folio,
 	this_cpupid = cpu_pid_to_cpupid(dst_cpu, current->pid);
 	last_cpupid = folio_xchg_last_cpupid(folio, this_cpupid);
 
-	if (!(numab_mode_tiering()) && !node_is_toptier(src_nid) &&
+	if (!folio_numab_mode_tiering(folio) && !node_is_toptier(src_nid) &&
 	    !cpupid_valid(last_cpupid))
 		return false;
 
@@ -4207,7 +4207,7 @@ void task_numa_fault(int last_cpupid, int mem_node, int pages, int flags)
 	 * node for memory tiering mode.
 	 */
 	if (!node_is_toptier(mem_node) &&
-	    (numab_mode_tiering() || !cpupid_valid(last_cpupid)))
+	    (task_numab_mode_tiering(p) || !cpupid_valid(last_cpupid)))
 		return;
 
 	/* Allocate buffer to track faults on a per-node basis */
diff --git a/mm/memory-tiers.c b/mm/memory-tiers.c
index ec7b9aadda98..b401f2e36d99 100644
--- a/mm/memory-tiers.c
+++ b/mm/memory-tiers.c
@@ -57,15 +57,15 @@ static const struct bus_type memory_tier_subsys = {
  * folio_use_access_time - check if a folio reuses cpupid for page access time
  * @folio: folio to check
  *
- * folio's _last_cpupid field is repurposed by memory tiering. In memory
- * tiering mode, cpupid of slow memory folio (not toptier memory) is used to
- * record page access time.
+ * folio's _last_cpupid field is repurposed by memory tiering. When the
+ * owner of @folio participates in tiering, cpupid of slow memory folio
+ * (not toptier memory) is used to record page access time.
  *
  * Return: the folio _last_cpupid is used to record page access time
  */
 bool folio_use_access_time(struct folio *folio)
 {
-	return numab_mode_tiering() &&
+	return folio_numab_mode_tiering(folio) &&
 	       !node_is_toptier(folio_nid(folio));
 }
 #endif
diff --git a/mm/mempolicy.c b/mm/mempolicy.c
index 79053ece02cd..6b8359ce16f8 100644
--- a/mm/mempolicy.c
+++ b/mm/mempolicy.c
@@ -868,11 +868,10 @@ bool folio_can_map_prot_numa(struct folio *folio, struct vm_area_struct *vma,
 		return false;
 
 	/*
-	 * Skip scanning top tier node if normal numa
-	 * balancing is disabled
+	 * Skip scanning top tier memory whose owner participates in
+	 * tiering: hinting faults on it would buy nothing.
 	 */
-	if (!(sysctl_numa_balancing_mode & NUMA_BALANCING_NORMAL) &&
-	    node_is_toptier(nid))
+	if (folio_numab_mode_tiering(folio) && node_is_toptier(nid))
 		return false;
 
 	if (folio_use_access_time(folio))
diff --git a/mm/migrate.c b/mm/migrate.c
index 4d2ba87eaa34..dd4e370ae9f8 100644
--- a/mm/migrate.c
+++ b/mm/migrate.c
@@ -819,7 +819,7 @@ void folio_migrate_flags(struct folio *newfolio, struct folio *folio)
 	 * memory node, reset cpupid, because that is used to record
 	 * page access time in slow memory node.
 	 */
-	if (numab_mode_tiering()) {
+	if (folio_numab_mode_tiering(folio)) {
 		bool f_toptier = node_is_toptier(folio_nid(folio));
 		bool t_toptier = node_is_toptier(folio_nid(newfolio));
 
@@ -2719,7 +2719,7 @@ int migrate_misplaced_folio_prepare(struct folio *folio,
 	if (!migrate_balanced_pgdat(pgdat, nr_pages)) {
 		int z;
 
-		if (!(numab_mode_tiering()))
+		if (!(folio_numab_mode_tiering(folio)))
 			return -EAGAIN;
 		for (z = pgdat->nr_zones - 1; z >= 0; z--) {
 			if (managed_zone(pgdat->node_zones + z))
-- 
2.52.0

Re: [RFC PATCH v2 2/3] sched: add per-cgroup NUMA balancing mode selection via BPF

[sashiko-bot]

Thank you for your contribution! Sashiko AI review found 4 potential issue(s) to consider:
- [High] mm: folio_can_map_prot_numa() disables NORMAL NUMA balancing on top-tier nodes when TIERING is active
- [High] sched: AB-BA deadlock in bpf_numab_mode (rq->lock -> pi_lock)
- [High] bpf: missing __nullable annotation in bpf_numab_ops leads to kernel panic
- [Medium] sched: BPF NUMA balancing opt-in ignored by LLC scheduling
--

commit 3975a69bbe7653d443a34c5ce6984dadc110d558
Author: Yafang Shao <laoar.shao@gmail.com>

sched: add per-cgroup NUMA balancing mode selection via BPF

Global NUMA balancing is all-or-nothing: kernel.numa_balancing and
kernel.numa_balancing_mode pick the mode for every workload. A
latency-sensitive service wants the scheduler's NUMA statistics but
not tiering's promotion machinery. This patch adds a bpf_numab_ops
struct_ops to allow policies attached to a cgroup to pick a mode.

> diff --git a/include/linux/bpf_numab.h b/include/linux/bpf_numab.h
> --- /dev/null
> +++ b/include/linux/bpf_numab.h
[ ... ]
> +struct bpf_numab_ops {
> +	/**
> +	 * numab_mode - request a NUMA balancing mode
> +	 * @p: task the NUMA balancing machinery is about to act on, or NULL
> +	 * @folio: folio whose NUMA balancing semantics are being decided,
> +	 *	   or NULL
> +	 *
> +	 * Exactly one of @p and @folio is set. The kernel passes @p when
[ ... ]
> +	 * Return: bits from enum bpf_numab_mode.  Other bits are dropped.
> +	 */
> +	u32 (*numab_mode)(struct task_struct *p, struct folio *folio);

[Severity: High]
Could the lack of a __nullable annotation here lead to a NULL pointer
dereference in user BPF programs? 

The BPF verifier relies on the __nullable suffix on struct_ops stub
arguments to mark them as PTR_MAYBE_NULL. Without it, the verifier assumes
they are strictly non-NULL and optimizes away any user-provided NULL checks.
When the kernel passes a NULL argument (as designed, since only one of p or
folio is set), the executing BPF program will crash the system.

> +};
[ ... ]
> diff --git a/include/linux/sched/numa_balancing.h b/include/linux/sched/numa_balancing.h
> --- a/include/linux/sched/numa_balancing.h
> +++ b/include/linux/sched/numa_balancing.h
[ ... ]
>  static inline bool task_numab_enabled(struct task_struct *p)
>  {
> +	u32 state = bpf_numab_mode(p, NULL);
> +
> +	/* A policy can always opt out, regardless of the global state. */
> +	if (state & BPF_NUMAB_DISABLED)
> +		return false;
> +	if (state & (BPF_NUMAB_NORMAL | BPF_NUMAB_TIERING))
> +		return true;
>  	return static_branch_unlikely(&sched_numa_balancing);
>  }

[Severity: Medium]
Is this BPF opt-in overlooked by the LLC scheduling logic?

This new helper allows workloads to opt into NUMA balancing even when global
NUMA balancing is disabled. However, functions like get_pref_llc() and
get_scan_cpumasks() still check the global static key sched_numa_balancing
directly instead of using task_numab_enabled(p). For a workload that
explicitly opts in via BPF, its numa_preferred_nid and CPU scan masks will
be skipped.

[ ... ]
> diff --git a/kernel/sched/bpf_numab.c b/kernel/sched/bpf_numab.c
> --- /dev/null
> +++ b/kernel/sched/bpf_numab.c
[ ... ]
> +static u32 numab_mode_stub(struct task_struct *p, struct folio *folio)
> +{
> +	return BPF_NUMAB_GLOBAL;
> +}

[Severity: High]
Should this stub function also include the __nullable suffix on its
arguments?

As mentioned in include/linux/bpf_numab.h, omitting the __nullable suffix
(e.g., p__nullable) prevents the verifier from recognizing the arguments as
optional. This leads to dead-code elimination of user NULL checks and a
subsequent system crash.

[ ... ]
> +u32 bpf_numab_mode(struct task_struct *p, struct folio *folio)
> +{
[ ... ]
> +	bpf_cgroup_struct_ops_foreach(ops, item, cgrp, CGROUP_NUMAB_OPS) {
> +		if (ops->numab_mode)
> +			acc |= ops->numab_mode(p, folio) & BPF_NUMAB_VALID_MASK;
> +	}
> +
> +	cgroup_bpf_put(cgrp);

[Severity: High]
Could dropping the cgroup BPF reference here trigger an AB-BA deadlock?

The scheduler path calls task_tick_fair() under rq->lock, which invokes
bpf_numab_mode(). Dropping the cgroup BPF reference here via
cgroup_bpf_put() can trigger queue_work() if this is the final reference
while the cgroup is being offlined. 

queue_work() acquires pool->lock and wakes up a worker by acquiring its
pi_lock. Because standard wakeup paths acquire pi_lock before rq->lock, this
inverted rq->lock -> pool->lock -> pi_lock sequence leads to a deadlock.

Since this code is already protected by rcu_read_lock(), is the reference
counting strictly necessary here?

> +out:
> +	rcu_read_unlock();
> +	return acc;
> +}
[ ... ]
> diff --git a/mm/mempolicy.c b/mm/mempolicy.c
> --- a/mm/mempolicy.c
> +++ b/mm/mempolicy.c
> @@ -884,11 +884,10 @@ bool folio_can_map_prot_numa(struct folio *folio, struct vm_area_struct *vma,
>  		return false;
>  
>  	/*
> -	 * Skip scanning top tier node if normal numa
> -	 * balancing is disabled
> +	 * Skip scanning top tier memory whose owner participates in
> +	 * tiering: hinting faults on it would buy nothing.
>  	 */
> -	if (!(sysctl_numa_balancing_mode & NUMA_BALANCING_NORMAL) &&
> -	    node_is_toptier(nid))
> +	if (folio_numab_mode_tiering(folio) && node_is_toptier(nid))
>  		return false;

[Severity: High]
Does this unconditionally disable NORMAL NUMA balancing on top-tier nodes
when TIERING is active?

Previously, this skipped hinting faults on top-tier nodes only if
NUMA_BALANCING_NORMAL was explicitly disabled. The new logic skips top-tier
scanning entirely if tiering is enabled.

When the global sysctl_numa_balancing_mode is 3 (both NORMAL and TIERING
enabled), folio_numab_mode_tiering() returns true. This breaks standard NUMA
page placement optimization for fast memory on tiered systems by dropping
all top-tier scanning.

>  
>  	if (folio_use_access_time(folio))

-- 
Sashiko AI review · https://sashiko.dev/#/patchset/20261012120438.52849-1-laoar.shao@gmail.com?part=2

Re: [RFC PATCH v2 2/3] sched: add per-cgroup NUMA balancing mode selection via BPF

[Yafang Shao]

On Mon, Oct 12, 2026 at 8:25 PM <sashiko-bot@kernel.org> wrote:
>
> Thank you for your contribution! Sashiko AI review found 4 potential issue(s) to consider:
> - [High] mm: folio_can_map_prot_numa() disables NORMAL NUMA balancing on top-tier nodes when TIERING is active
> - [High] sched: AB-BA deadlock in bpf_numab_mode (rq->lock -> pi_lock)
> - [High] bpf: missing __nullable annotation in bpf_numab_ops leads to kernel panic
> - [Medium] sched: BPF NUMA balancing opt-in ignored by LLC scheduling
> --
>
> commit 3975a69bbe7653d443a34c5ce6984dadc110d558
> Author: Yafang Shao <laoar.shao@gmail.com>
>
> sched: add per-cgroup NUMA balancing mode selection via BPF
>
> Global NUMA balancing is all-or-nothing: kernel.numa_balancing and
> kernel.numa_balancing_mode pick the mode for every workload. A
> latency-sensitive service wants the scheduler's NUMA statistics but
> not tiering's promotion machinery. This patch adds a bpf_numab_ops
> struct_ops to allow policies attached to a cgroup to pick a mode.
>
> > diff --git a/include/linux/bpf_numab.h b/include/linux/bpf_numab.h
> > --- /dev/null
> > +++ b/include/linux/bpf_numab.h
> [ ... ]
> > +struct bpf_numab_ops {
> > +     /**
> > +      * numab_mode - request a NUMA balancing mode
> > +      * @p: task the NUMA balancing machinery is about to act on, or NULL
> > +      * @folio: folio whose NUMA balancing semantics are being decided,
> > +      *         or NULL
> > +      *
> > +      * Exactly one of @p and @folio is set. The kernel passes @p when
> [ ... ]
> > +      * Return: bits from enum bpf_numab_mode.  Other bits are dropped.
> > +      */
> > +     u32 (*numab_mode)(struct task_struct *p, struct folio *folio);
>
> [Severity: High]
> Could the lack of a __nullable annotation here lead to a NULL pointer
> dereference in user BPF programs?

Nice catch.

It should be
u32 (*numab_mode)(struct task_struct *p__nullable, struct folio
*folio__nullable);

>
> The BPF verifier relies on the __nullable suffix on struct_ops stub
> arguments to mark them as PTR_MAYBE_NULL. Without it, the verifier assumes
> they are strictly non-NULL and optimizes away any user-provided NULL checks.
> When the kernel passes a NULL argument (as designed, since only one of p or
> folio is set), the executing BPF program will crash the system.
>
> > +};
> [ ... ]
> > diff --git a/include/linux/sched/numa_balancing.h b/include/linux/sched/numa_balancing.h
> > --- a/include/linux/sched/numa_balancing.h
> > +++ b/include/linux/sched/numa_balancing.h
> [ ... ]
> >  static inline bool task_numab_enabled(struct task_struct *p)
> >  {
> > +     u32 state = bpf_numab_mode(p, NULL);
> > +
> > +     /* A policy can always opt out, regardless of the global state. */
> > +     if (state & BPF_NUMAB_DISABLED)
> > +             return false;
> > +     if (state & (BPF_NUMAB_NORMAL | BPF_NUMAB_TIERING))
> > +             return true;
> >       return static_branch_unlikely(&sched_numa_balancing);
> >  }
>
> [Severity: Medium]
> Is this BPF opt-in overlooked by the LLC scheduling logic?

Right, LLC scheduling was overlooked on my end. Will fix in v2.

>
> This new helper allows workloads to opt into NUMA balancing even when global
> NUMA balancing is disabled. However, functions like get_pref_llc() and
> get_scan_cpumasks() still check the global static key sched_numa_balancing
> directly instead of using task_numab_enabled(p). For a workload that
> explicitly opts in via BPF, its numa_preferred_nid and CPU scan masks will
> be skipped.
>
> [ ... ]
> > diff --git a/kernel/sched/bpf_numab.c b/kernel/sched/bpf_numab.c
> > --- /dev/null
> > +++ b/kernel/sched/bpf_numab.c
> [ ... ]
> > +static u32 numab_mode_stub(struct task_struct *p, struct folio *folio)
> > +{
> > +     return BPF_NUMAB_GLOBAL;
> > +}
>
> [Severity: High]
> Should this stub function also include the __nullable suffix on its
> arguments?

I will add the __nullable here as well.

>
> As mentioned in include/linux/bpf_numab.h, omitting the __nullable suffix
> (e.g., p__nullable) prevents the verifier from recognizing the arguments as
> optional. This leads to dead-code elimination of user NULL checks and a
> subsequent system crash.
>
> [ ... ]
> > +u32 bpf_numab_mode(struct task_struct *p, struct folio *folio)
> > +{
> [ ... ]
> > +     bpf_cgroup_struct_ops_foreach(ops, item, cgrp, CGROUP_NUMAB_OPS) {
> > +             if (ops->numab_mode)
> > +                     acc |= ops->numab_mode(p, folio) & BPF_NUMAB_VALID_MASK;
> > +     }
> > +
> > +     cgroup_bpf_put(cgrp);
>
> [Severity: High]
> Could dropping the cgroup BPF reference here trigger an AB-BA deadlock?
>
> The scheduler path calls task_tick_fair() under rq->lock, which invokes
> bpf_numab_mode(). Dropping the cgroup BPF reference here via
> cgroup_bpf_put() can trigger queue_work() if this is the final reference
> while the cgroup is being offlined.
>
> queue_work() acquires pool->lock and wakes up a worker by acquiring its
> pi_lock. Because standard wakeup paths acquire pi_lock before rq->lock, this
> inverted rq->lock -> pool->lock -> pi_lock sequence leads to a deadlock.
>
> Since this code is already protected by rcu_read_lock(), is the reference
> counting strictly necessary here?

Nice catch.

The walk needs no refcount on cgrp->bpf: the cgroup itself is only
freed after an RCU grace period (css_free_rwork_fn() runs through
queue_rcu_work()), and the effective prog array is freed through
bpf_prog_array_free() (kfree_rcu), so rcu_read_lock() keeps both
alive for the whole walk -- the same protection bpf_tcp_ops_call()
relies on.

While removing it I noticed the copy from the memcg_ext series
missed a warning its walk carries: it is only v1-safe because
memory.high happens to be v2-only, and a hook on a path v1 can
reach needs its own cgroup_on_dfl() test -- a v1 cgroup has an
uninitialised bpf.refcnt and a NULL effective array.  The folio
surface here can reach v1 memcgs on a hybrid system, so the walk
now reads:

    rcu_read_lock();

    if (p) {
            cgrp = task_dfl_cgroup(p);
    } else {
            memcg = folio_memcg(folio);
            cgrp = memcg ? memcg->css.cgroup : NULL;
    }

    /* Only the default hierarchy has a cgroup_bpf. */
    if (!cgrp || !cgroup_on_dfl(cgrp))
            goto out;

    bpf_cgroup_struct_ops_foreach(ops, item, cgrp, CGROUP_NUMAB_OPS) {
            if (ops->numab_mode)
                    acc |= ops->numab_mode(p, folio) &
                           BPF_NUMAB_VALID_MASK;
    }
out:
    rcu_read_unlock();
    return acc;

>
> > +out:
> > +     rcu_read_unlock();
> > +     return acc;
> > +}
> [ ... ]
> > diff --git a/mm/mempolicy.c b/mm/mempolicy.c
> > --- a/mm/mempolicy.c
> > +++ b/mm/mempolicy.c
> > @@ -884,11 +884,10 @@ bool folio_can_map_prot_numa(struct folio *folio, struct vm_area_struct *vma,
> >               return false;
> >
> >       /*
> > -      * Skip scanning top tier node if normal numa
> > -      * balancing is disabled
> > +      * Skip scanning top tier memory whose owner participates in
> > +      * tiering: hinting faults on it would buy nothing.
> >        */
> > -     if (!(sysctl_numa_balancing_mode & NUMA_BALANCING_NORMAL) &&
> > -         node_is_toptier(nid))
> > +     if (folio_numab_mode_tiering(folio) && node_is_toptier(nid))
> >               return false;
>
> [Severity: High]
> Does this unconditionally disable NORMAL NUMA balancing on top-tier nodes
> when TIERING is active?
>
> Previously, this skipped hinting faults on top-tier nodes only if
> NUMA_BALANCING_NORMAL was explicitly disabled. The new logic skips top-tier
> scanning entirely if tiering is enabled.

You are right, this is a regression.
I will fix it by introducing a new folio_numab_mode_normal().

>
> When the global sysctl_numa_balancing_mode is 3 (both NORMAL and TIERING
> enabled), folio_numab_mode_tiering() returns true. This breaks standard NUMA
> page placement optimization for fast memory on tiered systems by dropping
> all top-tier scanning.
>
> >
> >       if (folio_use_access_time(folio))
>
> --
> Sashiko AI review · https://sashiko.dev/#/patchset/20261012120438.52849-1-laoar.shao@gmail.com?part=2



-- 
Regards
Yafang

[RFC PATCH v2 3/3] selftests/bpf: add tests for the bpf_numab_ops struct_ops

[Yafang Shao]

Five cases:
- attach/detach with counter proof
- hierarchy aggregation
- multi-policy combination
- inheritance
- a surface-dispatching policy

Assisted-by: LLM
Signed-off-by: Yafang Shao <laoar.shao@gmail.com>
---
 .../selftests/bpf/prog_tests/bpf_numab_ops.c  | 250 ++++++++++++++++++
 .../selftests/bpf/progs/bpf_numab_ops.c       |  57 ++++
 2 files changed, 307 insertions(+)
 create mode 100644 tools/testing/selftests/bpf/prog_tests/bpf_numab_ops.c
 create mode 100644 tools/testing/selftests/bpf/progs/bpf_numab_ops.c

diff --git a/tools/testing/selftests/bpf/prog_tests/bpf_numab_ops.c b/tools/testing/selftests/bpf/prog_tests/bpf_numab_ops.c
new file mode 100644
index 000000000000..3e5dac8e3131
--- /dev/null
+++ b/tools/testing/selftests/bpf/prog_tests/bpf_numab_ops.c
@@ -0,0 +1,250 @@
+// SPDX-License-Identifier: GPL-2.0
+
+#include <test_progs.h>
+#include <time.h>
+#include "cgroup_helpers.h"
+#include "bpf_numab_ops.skel.h"
+
+#define CGROUP_PATH	"/bpf_numab_ops"
+#define CHILD_PATH	CGROUP_PATH "/child"
+
+/*
+ * task_tick_fair() consults the per-cgroup NUMA balancing policy on every
+ * tick of a running task, so a short busy loop is enough to prove that the
+ * hook runs. This works regardless of the global numa_balancing sysctl
+ * state and regardless of the machine's NUMA topology: both the enabled
+ * and disabled code paths end up calling every attached policy.
+ */
+static void burn_cpu_ms(int ms)
+{
+	struct timespec now;
+	u64 deadline;
+
+	clock_gettime(CLOCK_MONOTONIC, &now);
+	deadline = now.tv_sec * 1000ULL + now.tv_nsec / 1000000ULL + ms;
+
+	do {
+		clock_gettime(CLOCK_MONOTONIC, &now);
+	} while (now.tv_sec * 1000ULL + now.tv_nsec / 1000000ULL < deadline);
+}
+
+static void test_attach_detach(void)
+{
+	struct bpf_numab_ops *skel;
+	struct bpf_link *link = NULL;
+	int cgroup_fd, before;
+
+	cgroup_fd = test__join_cgroup(CGROUP_PATH);
+	if (!ASSERT_GE(cgroup_fd, 0, "join_cgroup"))
+		return;
+
+	skel = bpf_numab_ops__open_and_load();
+	if (!ASSERT_OK_PTR(skel, "open_and_load"))
+		goto cleanup;
+
+	link = bpf_map__attach_cgroup_opts(skel->maps.numab_enable_ops,
+					   cgroup_fd, NULL);
+	if (!ASSERT_OK_PTR(link, "attach_enable_ops"))
+		goto cleanup;
+
+	skel->bss->enable_cnt = 0;
+	burn_cpu_ms(300);
+	ASSERT_GT(skel->bss->enable_cnt, 0, "enable_cnt");
+
+	/*
+	 * Detach and let a grace period pass so that no reader is still
+	 * walking the old effective array, then verify the hook stopped
+	 * running.
+	 */
+	before = skel->bss->enable_cnt;
+	bpf_link__destroy(link);
+	link = NULL;
+	sleep(1);
+	burn_cpu_ms(300);
+	ASSERT_EQ(skel->bss->enable_cnt, before, "enable_cnt_after_detach");
+
+cleanup:
+	bpf_link__destroy(link);
+	bpf_numab_ops__destroy(skel);
+	close(cgroup_fd);
+}
+
+/*
+ * Attach to the parent cgroup while the task lives in the child: the
+ * child's effective array must aggregate the parent's policy.
+ */
+static void test_hierarchy(void)
+{
+	struct bpf_numab_ops *skel;
+	struct bpf_link *plink = NULL;
+	int parent_fd, child_fd = -1;
+
+	parent_fd = test__join_cgroup(CGROUP_PATH);
+	if (!ASSERT_GE(parent_fd, 0, "join_parent_cgroup"))
+		return;
+
+	child_fd = create_and_get_cgroup(CHILD_PATH);
+	if (!ASSERT_GE(child_fd, 0, "create_child_cgroup"))
+		goto cleanup;
+
+	if (!ASSERT_OK(join_cgroup(CHILD_PATH), "join_child_cgroup"))
+		goto cleanup;
+
+	skel = bpf_numab_ops__open_and_load();
+	if (!ASSERT_OK_PTR(skel, "open_and_load"))
+		goto cleanup;
+
+	plink = bpf_map__attach_cgroup_opts(skel->maps.numab_enable_ops,
+					    parent_fd, NULL);
+	if (!ASSERT_OK_PTR(plink, "attach_to_parent"))
+		goto cleanup;
+
+	skel->bss->enable_cnt = 0;
+	burn_cpu_ms(300);
+	ASSERT_GT(skel->bss->enable_cnt, 0, "enable_cnt_from_parent_policy");
+
+cleanup:
+	bpf_link__destroy(plink);
+	bpf_numab_ops__destroy(skel);
+	if (child_fd >= 0) {
+		close(child_fd);
+		join_cgroup(CGROUP_PATH);
+		remove_cgroup(CHILD_PATH);
+	}
+	close(parent_fd);
+}
+
+/*
+ * Attach two policies to the same cgroup. Both must run: the kernel
+ * combines what they return.
+ */
+static void test_multi(void)
+{
+	struct bpf_numab_ops *skel;
+	struct bpf_link *link1 = NULL, *link2 = NULL;
+	int cgroup_fd;
+
+	cgroup_fd = test__join_cgroup(CGROUP_PATH);
+	if (!ASSERT_GE(cgroup_fd, 0, "join_cgroup"))
+		return;
+
+	skel = bpf_numab_ops__open_and_load();
+	if (!ASSERT_OK_PTR(skel, "open_and_load"))
+		goto cleanup;
+
+	link1 = bpf_map__attach_cgroup_opts(skel->maps.numab_enable_ops,
+					    cgroup_fd, NULL);
+	if (!ASSERT_OK_PTR(link1, "attach_enable_ops"))
+		goto cleanup;
+
+	link2 = bpf_map__attach_cgroup_opts(skel->maps.numab_disable_ops,
+					    cgroup_fd, NULL);
+	if (!ASSERT_OK_PTR(link2, "attach_disable_ops"))
+		goto cleanup;
+
+	skel->bss->enable_cnt = 0;
+	skel->bss->disable_cnt = 0;
+	burn_cpu_ms(300);
+	ASSERT_GT(skel->bss->enable_cnt, 0, "enable_cnt");
+	ASSERT_GT(skel->bss->disable_cnt, 0, "disable_cnt");
+
+cleanup:
+	bpf_link__destroy(link2);
+	bpf_link__destroy(link1);
+	bpf_numab_ops__destroy(skel);
+	close(cgroup_fd);
+}
+
+/*
+ * Attach to the parent first, then create and join the child cgroup. The
+ * child must inherit the parent's effective array via cgroup_bpf_inherit().
+ */
+static void test_inherit(void)
+{
+	struct bpf_numab_ops *skel;
+	struct bpf_link *plink = NULL;
+	int parent_fd, child_fd = -1;
+
+	parent_fd = test__join_cgroup(CGROUP_PATH);
+	if (!ASSERT_GE(parent_fd, 0, "join_parent_cgroup"))
+		return;
+
+	skel = bpf_numab_ops__open_and_load();
+	if (!ASSERT_OK_PTR(skel, "open_and_load"))
+		goto cleanup;
+
+	plink = bpf_map__attach_cgroup_opts(skel->maps.numab_enable_ops,
+					    parent_fd, NULL);
+	if (!ASSERT_OK_PTR(plink, "attach_to_parent"))
+		goto cleanup;
+
+	child_fd = create_and_get_cgroup(CHILD_PATH);
+	if (!ASSERT_GE(child_fd, 0, "create_child_cgroup"))
+		goto cleanup;
+
+	if (!ASSERT_OK(join_cgroup(CHILD_PATH), "join_child_cgroup"))
+		goto cleanup;
+
+	skel->bss->enable_cnt = 0;
+	burn_cpu_ms(300);
+	ASSERT_GT(skel->bss->enable_cnt, 0, "enable_cnt_inherited");
+
+cleanup:
+	bpf_link__destroy(plink);
+	bpf_numab_ops__destroy(skel);
+	if (child_fd >= 0) {
+		close(child_fd);
+		join_cgroup(CGROUP_PATH);
+		remove_cgroup(CHILD_PATH);
+	}
+	close(parent_fd);
+}
+
+/*
+ * A surface-dispatching policy must load, attach and run from the task
+ * surface.  The folio surface runs from the scanner and fault paths,
+ * which need a scan pass and foreign-node folios to consult, so folio_cnt
+ * is left for debugging rather than asserted here.
+ */
+static void test_mixed(void)
+{
+	struct bpf_numab_ops *skel;
+	struct bpf_link *link = NULL;
+	int cgroup_fd;
+
+	cgroup_fd = test__join_cgroup(CGROUP_PATH);
+	if (!ASSERT_GE(cgroup_fd, 0, "join_cgroup"))
+		return;
+
+	skel = bpf_numab_ops__open_and_load();
+	if (!ASSERT_OK_PTR(skel, "open_and_load"))
+		goto cleanup;
+
+	link = bpf_map__attach_cgroup_opts(skel->maps.numab_mixed_ops,
+					   cgroup_fd, NULL);
+	if (!ASSERT_OK_PTR(link, "attach_mixed_ops"))
+		goto cleanup;
+
+	skel->bss->task_cnt = 0;
+	burn_cpu_ms(300);
+	ASSERT_GT(skel->bss->task_cnt, 0, "task_cnt");
+
+cleanup:
+	bpf_link__destroy(link);
+	bpf_numab_ops__destroy(skel);
+	close(cgroup_fd);
+}
+
+void serial_test_bpf_numab_ops(void)
+{
+	if (test__start_subtest("attach_detach"))
+		test_attach_detach();
+	if (test__start_subtest("hierarchy"))
+		test_hierarchy();
+	if (test__start_subtest("multi"))
+		test_multi();
+	if (test__start_subtest("mixed"))
+		test_mixed();
+	if (test__start_subtest("inherit"))
+		test_inherit();
+}
diff --git a/tools/testing/selftests/bpf/progs/bpf_numab_ops.c b/tools/testing/selftests/bpf/progs/bpf_numab_ops.c
new file mode 100644
index 000000000000..c9abff2cd694
--- /dev/null
+++ b/tools/testing/selftests/bpf/progs/bpf_numab_ops.c
@@ -0,0 +1,57 @@
+// SPDX-License-Identifier: GPL-2.0
+
+#include "vmlinux.h"
+#include <bpf/bpf_helpers.h>
+#include <bpf/bpf_tracing.h>
+
+/* Counters, bumped from the NUMA balancing hot paths. */
+int enable_cnt;
+int disable_cnt;
+int task_cnt;
+int folio_cnt;
+
+SEC("struct_ops")
+u32 BPF_PROG(numab_mode_enable, struct task_struct *p, struct folio *folio)
+{
+	__sync_fetch_and_add(&enable_cnt, 1);
+	return BPF_NUMAB_NORMAL;
+}
+
+SEC(".struct_ops.link")
+struct bpf_numab_ops numab_enable_ops = {
+	.numab_mode = (void *)numab_mode_enable,
+};
+
+SEC("struct_ops")
+u32 BPF_PROG(numab_mode_disable, struct task_struct *p, struct folio *folio)
+{
+	__sync_fetch_and_add(&disable_cnt, 1);
+	return BPF_NUMAB_DISABLED;
+}
+
+SEC(".struct_ops.link")
+struct bpf_numab_ops numab_disable_ops = {
+	.numab_mode = (void *)numab_mode_disable,
+};
+
+/*
+ * A surface-dispatching policy: exactly one of @p and @folio is set, so
+ * the prog answers the task surface and the memory surface separately.
+ */
+SEC("struct_ops")
+u32 BPF_PROG(numab_mode_mixed, struct task_struct *p, struct folio *folio)
+{
+	if (p) {
+		__sync_fetch_and_add(&task_cnt, 1);
+		return BPF_NUMAB_NORMAL;
+	}
+	__sync_fetch_and_add(&folio_cnt, 1);
+	return BPF_NUMAB_TIERING;
+}
+
+SEC(".struct_ops.link")
+struct bpf_numab_ops numab_mixed_ops = {
+	.numab_mode = (void *)numab_mode_mixed,
+};
+
+char _license[] SEC("license") = "GPL";
-- 
2.52.0

Re: [RFC PATCH v2 3/3] selftests/bpf: add tests for the bpf_numab_ops struct_ops

[sashiko-bot]

Thank you for your contribution! Sashiko AI review found 2 potential issue(s) to consider:
- [High] selftests/bpf: uninitialized skel pointer dereference on test_hierarchy error path
- [Medium] selftests/bpf: race condition leading to flaky test in test_attach_detach
--

commit 0d0b1887ad5f6f3d52adb583711d13aa30078932
Author: Yafang Shao <laoar.shao@gmail.com>

selftests/bpf: add tests for the bpf_numab_ops struct_ops

This commit adds selftests for the bpf_numab_ops struct_ops. It covers five
scenarios including attach/detach behavior, hierarchy aggregation, multi-policy
combination, inheritance, and surface-dispatching.

> diff --git a/tools/testing/selftests/bpf/prog_tests/bpf_numab_ops.c b/tools/testing/selftests/bpf/prog_tests/bpf_numab_ops.c
> new file mode 100644
> index 0000000000000..3e5dac8e31313
> --- /dev/null
> +++ b/tools/testing/selftests/bpf/prog_tests/bpf_numab_ops.c

[ ... ]

> +static void test_attach_detach(void)
> +{

[ ... ]

> +	/*
> +	 * Detach and let a grace period pass so that no reader is still
> +	 * walking the old effective array, then verify the hook stopped
> +	 * running.
> +	 */
> +	before = skel->bss->enable_cnt;
> +	bpf_link__destroy(link);
> +	link = NULL;
> +	sleep(1);
> +	burn_cpu_ms(300);
> +	ASSERT_EQ(skel->bss->enable_cnt, before, "enable_cnt_after_detach");

[Severity: Medium]
Could this snapshot timing cause the test to be flaky?

Because the task is running during bpf_link__destroy() and BPF hooks are
attached via RCU, scheduler ticks can still trigger the hook and increment
the counter after the before snapshot is taken but before the detachment is
complete or during the grace period.

If any ticks hit the task during teardown, the counts will not match and the
test will spuriously fail. Should the baseline counter be snapshotted after
the link is destroyed and the RCU grace period (sleep(1)) has fully elapsed?

[ ... ]

> +static void test_hierarchy(void)
> +{
> +	struct bpf_numab_ops *skel;
> +	struct bpf_link *plink = NULL;
> +	int parent_fd, child_fd = -1;
> +
> +	parent_fd = test__join_cgroup(CGROUP_PATH);
> +	if (!ASSERT_GE(parent_fd, 0, "join_parent_cgroup"))
> +		return;
> +
> +	child_fd = create_and_get_cgroup(CHILD_PATH);
> +	if (!ASSERT_GE(child_fd, 0, "create_child_cgroup"))
> +		goto cleanup;
> +
> +	if (!ASSERT_OK(join_cgroup(CHILD_PATH), "join_child_cgroup"))
> +		goto cleanup;
> +
> +	skel = bpf_numab_ops__open_and_load();

[ ... ]

> +cleanup:
> +	bpf_link__destroy(plink);
> +	bpf_numab_ops__destroy(skel);

[Severity: High]
Does this code dereference an uninitialized pointer if cgroup creation fails?

The skel pointer is declared without initialization. If either
create_and_get_cgroup() or join_cgroup() fails, the code jumps to the cleanup
label before skel is assigned a value.

At the cleanup label, bpf_numab_ops__destroy(skel) is invoked with the
uninitialized skel pointer containing garbage stack data. This bypasses the
internal NULL check in the destructor and can cause a segmentation fault.

Would initializing struct bpf_numab_ops *skel = NULL; prevent this?

-- 
Sashiko AI review · https://sashiko.dev/#/patchset/20261012120438.52849-1-laoar.shao@gmail.com?part=3

Re: [RFC PATCH v2 0/3] BPF-based per-cgroup NUMA balancing mode selection

[Tejun Heo]

Hello,

On Mon, Oct 12, 2026 at 08:04:35PM +0800, Yafang Shao wrote:
> This series adds a bpf_numab_ops struct_ops which attaches to a
> cgroup. The kernel consults the policies of the matching cgroup and of
> every ancestor and combines what they return.

The op doesn't compute anything. It returns a fixed mode per cgroup and
that's all the series uses it for. BPF is unnecessary here. This is a
prctl() feature, not a struct_ops one.

Thanks.

-- 
tejun

Re: [RFC PATCH v2 0/3] BPF-based per-cgroup NUMA balancing mode selection

[Yafang Shao]

On Tue, Oct 13, 2026 at 1:51 AM Tejun Heo <tj@kernel.org> wrote:
>
> Hello,
>
> On Mon, Oct 12, 2026 at 08:04:35PM +0800, Yafang Shao wrote:
> > This series adds a bpf_numab_ops struct_ops which attaches to a
> > cgroup. The kernel consults the policies of the matching cgroup and of
> > every ancestor and combines what they return.
>
> The op doesn't compute anything. It returns a fixed mode per cgroup and
> that's all the series uses it for. BPF is unnecessary here. This is a
> prctl() feature, not a struct_ops one.

Part of the consults are keyed on struct folio rather than struct
task_struct: the folio surface answers from the folio's memcg, the
owner's request, not the touching task's. I don't see how to key
these callsites to a per-task setting.  Please correct me if I'm
wrong.

prctl() is also a stable uAPI rather than a kernel-internal
interface: a new option is frozen the day it merges, which is a
poor fit for a feature whose semantics are still moving.  The
struct_ops/kfunc surface is explicitly exempt from the uAPI
stability rules, which is what lets us experiment with the kernel
feature rapidly.

-- 
Regards
Yafang