[RFC PATCH v3 bpf-next 0/5] BPF-based per-cgroup NUMA balancing mode selection

[RFC PATCH v3 bpf-next 0/5] 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 to address it.

Enabling NUMA balancing globally is still not acceptable since it 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 wants NUMA balancing disabled,
while a batch service wants its PMEM-resident memory promoted when hot
without paying for scheduler-driven migration. No single global setting
serves both.

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

This series introduces the fine-grained per-cgroup NUMA balancing
control. It grows out of Tejun's suggestion to store the existing
NUMA_BALANCING_* bitmask per cgroup and expose it through BPF kfuncs.
cgroup::numab_mode holds the request: NUMAB_FOLLOW_SYSTEM (-1) defers
to the global sysctl state, anything else is a bitmask.
bpf_numab_set_mode() writes a request, bpf_numab_get_mode() reads the
combined request of the cgroup and of every ancestor back. The hot
paths walk the same chain at read time behind a static key, so while no
cgroup makes a request the overhead is one patched-out branch.

An opt-out wins and binds the subtree: a DISABLED request anywhere in
the chain opts the cgroup's tasks out entirely. Otherwise the requested
bits OR together and the kernel reads them independently, the way
kernel.numa_balancing reads them, so a combined NORMAL|TIERING engages
both modes. A DISABLED request only speaks for the cgroup's tasks,
never for its memory, which follows the remaining requests of its
ancestors.

The cost of the additive combination is that an ancestor cannot cap
tiering short of a full opt-out: a descendant's TIERING request
re-engages promotion for its own memory, and promotion is the mode that
consumes the shared fast tier.

The readback contract is worth spelling out: bpf_numab_get_mode()
returns the combined request of the cgroup and its ancestors, and the
kernel reads it the same way it reads a single request, each bit
independently as kernel.numa_balancing does, so a combined
NORMAL|TIERING engages both modes. When any cgroup in the chain makes a
request, the value read back is the mode in effect for the cgroup's tasks;
when none does, the readback is NUMAB_FOLLOW_SYSTEM (-1) and the sysctl
decides. The answer speaks for the cgroup's tasks only: its memory follows
the remaining requests of its chain when it opts out, and that folio-side
view is not exposed by the kfuncs. These are deliberate simplifications;
both can be revisited if callers turn out to need them.

The sysctl's own documentation is untouched; the override is documented
where it lives, in the CONFIG_BPF_NUMAB Kconfig help and the kfunc docs.

The override was measured with schbench, 384 workers, on a two-node
384-core machine. Whatever kernel.numa_balancing is set to, a cgroup
request puts the cgroup at the throughput of the equivalent
system-wide setting, within a fraction of a percent: the override
itself costs nothing measurable, performance follows the effective
state alone. The two states themselves differ by less than a percent
under this load.

Changes
=======

RFC v2->RFC v3:
- The struct_ops interface is replaced by two kfuncs, following
  Tejun's suggestion.
- Improve on the function naming and simplify the code (David)

RFC v2: https://lwn.net/Articles/1099962/

RFC v1->RFC v2:
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, 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.

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

Yafang Shao (5):
  sched: fix stale get_pref_llc() stub for !CONFIG_SCHED_CACHE
  sched: use static_branch_likely() for the NUMA balancing tick check
  sched: add helpers for numa balancing
  sched: add BPF-based per-cgroup NUMA balancing mode selection
  selftests/bpf: add tests for the numab kfuncs

 include/linux/cgroup-defs.h                   |  10 +
 include/linux/sched/numa_balancing.h          | 149 ++++++++++++
 init/Kconfig                                  |  20 ++
 kernel/cgroup/cgroup.c                        |  14 ++
 kernel/sched/bpf_numab.c                      | 195 +++++++++++++++
 kernel/sched/build_policy.c                   |   4 +
 kernel/sched/fair.c                           |  33 +--
 kernel/sched/sched.h                          |   1 -
 mm/memory-tiers.c                             |   9 +-
 mm/mempolicy.c                                |   3 +-
 mm/migrate.c                                  |   7 +-
 mm/vmscan.c                                   |   6 +-
 .../selftests/bpf/prog_tests/bpf_numab.c      | 227 ++++++++++++++++++
 tools/testing/selftests/bpf/progs/bpf_numab.c |  32 +++
 14 files changed, 681 insertions(+), 29 deletions(-)
 create mode 100644 kernel/sched/bpf_numab.c
 create mode 100644 tools/testing/selftests/bpf/prog_tests/bpf_numab.c
 create mode 100644 tools/testing/selftests/bpf/progs/bpf_numab.c

-- 
2.52.0

[RFC PATCH v3 bpf-next 1/5] sched: fix stale get_pref_llc() stub for !CONFIG_SCHED_CACHE

[Yafang Shao]

Commit b636fef85bda ("sched/cache: Introduce task_struct->sched_cache_grp
to fix UAF") switched get_pref_llc() from taking a struct mm_struct * to
taking the task's struct sched_cache_group *, but only updated the
CONFIG_SCHED_CACHE definition and its caller: the stub in the
!CONFIG_SCHED_CACHE branch kept the old mm_struct parameter.

Fixes: b636fef85bda ("sched/cache: Introduce task_struct->sched_cache_grp to fix UAF")
Signed-off-by: Yafang Shao <laoar.shao@gmail.com>
Cc: Tim Chen <tim.c.chen@linux.intel.com>
Cc: Chen Yu <yu.c.chen@intel.com>
---
 kernel/sched/fair.c | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/kernel/sched/fair.c b/kernel/sched/fair.c
index 57360f5cdde4..973647ad6c2d 100644
--- a/kernel/sched/fair.c
+++ b/kernel/sched/fair.c
@@ -2194,7 +2194,7 @@ void init_sched_mm(struct task_struct *p) { }
 static void task_tick_cache(struct rq *rq, struct task_struct *p) { }
 
 static inline int get_pref_llc(struct task_struct *p,
-			       struct mm_struct *mm)
+			       struct sched_cache_group *grp)
 {
 	return -1;
 }
-- 
2.52.0

[RFC PATCH v3 bpf-next 2/5] sched: use static_branch_likely() for the NUMA balancing tick check

[Yafang Shao]

task_tick_fair() checks sched_numa_balancing on every tick. The
static_branch_unlikely() annotation reserves the nop for the disabled case,
so the default configuration (NUMA_BALANCING_DEFAULT_ENABLED=y) pays a
taken jump per tick. Flip the polarity to static_branch_likely() so that
the fast path matches the default.

Signed-off-by: Yafang Shao <laoar.shao@gmail.com>
---
 kernel/sched/fair.c | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/kernel/sched/fair.c b/kernel/sched/fair.c
index 973647ad6c2d..71dc47d6114d 100644
--- a/kernel/sched/fair.c
+++ b/kernel/sched/fair.c
@@ -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 (static_branch_likely(&sched_numa_balancing))
 		task_tick_numa(rq, curr);
 
 	task_tick_cache(rq, curr);
-- 
2.52.0

[RFC PATCH v3 bpf-next 3/5] sched: add helpers for numa balancing

[Yafang Shao]

Three new helpers task_numab_enabled(), numab_in_tiering_mode() and
numab_in_normal_mode() 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>
Cc: David Hildenbrand (Arm) <david@kernel.org>
---
 include/linux/sched/numa_balancing.h | 33 ++++++++++++++++++++++++++++
 kernel/sched/fair.c                  | 31 +++++++++++++-------------
 kernel/sched/sched.h                 |  1 -
 mm/memory-tiers.c                    |  3 ++-
 mm/mempolicy.c                       |  3 +--
 mm/migrate.c                         |  7 +++---
 mm/vmscan.c                          |  6 ++---
 7 files changed, 59 insertions(+), 25 deletions(-)

diff --git a/include/linux/sched/numa_balancing.h b/include/linux/sched/numa_balancing.h
index 52b22c5c396d..36de57f7fb06 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,26 @@ 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. */
+static inline bool task_numab_enabled(struct task_struct *p)
+{
+	return static_branch_likely(&sched_numa_balancing);
+}
+
+/* Whether NUMA balancing is in memory tiering mode. */
+static inline bool numab_in_tiering_mode(void)
+{
+	return sysctl_numa_balancing_mode & NUMA_BALANCING_MEMORY_TIERING;
+}
+
+/* Whether NUMA balancing is in normal mode. */
+static inline bool numab_in_normal_mode(void)
+{
+	return sysctl_numa_balancing_mode & NUMA_BALANCING_NORMAL;
+}
 #else
 static inline void task_numa_fault(int last_node, int node, int pages,
 				   int flags)
@@ -52,6 +73,18 @@ 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_in_tiering_mode(void)
+{
+	return false;
+}
+static inline bool numab_in_normal_mode(void)
+{
+	return false;
+}
 #endif
 
 #endif /* _LINUX_SCHED_NUMA_BALANCING_H */
diff --git a/kernel/sched/fair.c b/kernel/sched/fair.c
index 71dc47d6114d..066818498c76 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>
 
@@ -1478,7 +1479,8 @@ static inline int get_sched_cache_scale(int mul)
 	return (1 + (tol - 1) * mul);
 }
 
-static bool exceed_llc_capacity(struct sched_cache_group *grp, int cpu)
+static bool exceed_llc_capacity(struct sched_cache_group *grp, struct task_struct *p,
+				int cpu)
 {
 #ifdef CONFIG_NUMA_BALANCING
 	unsigned long llc, footprint;
@@ -1491,7 +1493,7 @@ static bool exceed_llc_capacity(struct sched_cache_group *grp, int cpu)
 	if (!sd)
 		return true;
 
-	if (static_branch_likely(&sched_numa_balancing)) {
+	if (task_numab_enabled(p)) {
 		/*
 		 * TBD: RDT exclusive LLC ways reserved should be
 		 * excluded.
@@ -1879,7 +1881,7 @@ static int get_pref_llc(struct task_struct *p, struct sched_cache_group *grp)
 		 * than sched_setnuma() at least -- and thus the
 		 * conflict only exists for a short period of time.
 		 */
-		if (static_branch_likely(&sched_numa_balancing) &&
+		if (task_numab_enabled(p) &&
 		    p->numa_preferred_nid >= 0 &&
 		    cpu_to_node(mm_sched_cpu) != p->numa_preferred_nid)
 			mm_sched_llc = -1;
@@ -1932,7 +1934,7 @@ void account_mm_sched(struct rq *rq, struct task_struct *p, s64 delta_exec)
 	 */
 	if ((long)(epoch - READ_ONCE(grp->epoch)) > llc_epoch_affinity_timeout ||
 	    invalid_llc_nr(grp, p, cpu_of(rq)) ||
-	    exceed_llc_capacity(grp, cpu_of(rq))) {
+	    exceed_llc_capacity(grp, p, cpu_of(rq))) {
 		if (READ_ONCE(grp->cpu) != -1)
 			WRITE_ONCE(grp->cpu, -1);
 	}
@@ -1980,7 +1982,7 @@ static void get_scan_cpumasks(cpumask_var_t cpus, struct task_struct *p,
 #ifdef CONFIG_NUMA_BALANCING
 	int cpu, curr_cpu, nid, pref_nid;
 
-	if (!static_branch_likely(&sched_numa_balancing))
+	if (!task_numab_enabled(p))
 		goto out;
 
 	cpu = READ_ONCE(grp->cpu);
@@ -2076,7 +2078,7 @@ static void task_cache_work(struct callback_head *work)
 
 	curr_cpu = task_cpu(p);
 	if (invalid_llc_nr(grp, p, curr_cpu) ||
-	    exceed_llc_capacity(grp, curr_cpu)) {
+	    exceed_llc_capacity(grp, p, curr_cpu)) {
 		if (READ_ONCE(grp->cpu) != -1)
 			WRITE_ONCE(grp->cpu, -1);
 
@@ -2963,8 +2965,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_in_tiering_mode()) && !node_is_toptier(src_nid) &&
+	    !cpupid_valid(last_cpupid))
 		return false;
 
 	/*
@@ -4194,7 +4196,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 +4208,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_in_tiering_mode() || !cpupid_valid(last_cpupid)))
 		return;
 
 	/* Allocate buffer to track faults on a per-node basis */
@@ -4665,7 +4666,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 +10698,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))
@@ -10960,7 +10961,7 @@ static enum llc_mig can_migrate_llc_task(struct lb_env *env,
 
 	/* skip cache aware load balance for too many threads */
 	if (invalid_llc_nr(grp, p, dst_cpu) ||
-	    exceed_llc_capacity(grp, dst_cpu)) {
+	    exceed_llc_capacity(grp, p, dst_cpu)) {
 		if (READ_ONCE(grp->cpu) != -1)
 			WRITE_ONCE(grp->cpu, -1);
 		return mig_unrestricted;
@@ -15351,7 +15352,7 @@ static void task_tick_fair(struct rq *rq, struct task_struct *curr, int queued)
 	if (queued)
 		return;
 
-	if (static_branch_likely(&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..9009b164edc2 100644
--- a/mm/memory-tiers.c
+++ b/mm/memory-tiers.c
@@ -6,6 +6,7 @@
 #include <linux/memory.h>
 #include <linux/memory-tiers.h>
 #include <linux/notifier.h>
+#include <linux/sched/numa_balancing.h>
 #include <linux/sched/sysctl.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_in_tiering_mode() &&
 	       !node_is_toptier(folio_nid(folio));
 }
 #endif
diff --git a/mm/mempolicy.c b/mm/mempolicy.c
index 79053ece02cd..442c93079fb0 100644
--- a/mm/mempolicy.c
+++ b/mm/mempolicy.c
@@ -871,8 +871,7 @@ bool folio_can_map_prot_numa(struct folio *folio, struct vm_area_struct *vma,
 	 * Skip scanning top tier node if normal numa
 	 * balancing is disabled
 	 */
-	if (!(sysctl_numa_balancing_mode & NUMA_BALANCING_NORMAL) &&
-	    node_is_toptier(nid))
+	if (!numab_in_normal_mode() && node_is_toptier(nid))
 		return false;
 
 	if (folio_use_access_time(folio))
diff --git a/mm/migrate.c b/mm/migrate.c
index 15b45832bcfa..a622df9a9b42 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_in_tiering_mode()) {
 		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_in_tiering_mode())
 			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_in_tiering_mode())
 		    && !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..17455255fc30 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_in_tiering_mode() ? 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_in_tiering_mode())
 			mark = promo_wmark_pages(zone);
 		else
 			mark = high_wmark_pages(zone);
-- 
2.52.0

[RFC PATCH v3 bpf-next 4/5] sched: add BPF-based per-cgroup NUMA balancing mode selection

[Yafang Shao]

Global NUMA balancing offers no per-workload choice: kernel.numa_balancing
is a single system-wide sysctl, where zero disables the machinery and the
set bits select normal balancing, memory tiering, or both for every
workload alike. However, a latency-sensitive service wants NUMA balancing
disabled, while a batch service wants its PMEM-resident memory promoted
when hot without paying for scheduler-driven migration. No single global
setting serves both.

Add a per-cgroup request. cgroup::numab_mode stores the mode a cgroup asks
for: NUMAB_FOLLOW_SYSTEM (-1) defers to the global sysctl state, anything
else is a NUMA_BALANCING_* bitmask. It is initialized to -1 at cgroup
creation since kzalloc's 0 would read as an explicit opt-out. Two kfuncs
manage it: bpf_numab_set_mode() writes the request, bpf_numab_get_mode()
reads the combined request back.

The hot paths combine the requests of the matching cgroup and of every
ancestor: task callsites walk the task's cgroup chain, folio callsites walk
the chain of the memcg the folio is charged to, so a memory answer is keyed
on the folio's owner rather than on whoever happens to touch it. A kswapd
or demotion pass decides the same way the fault path does. An opt-out
wins: a DISABLED request anywhere in the chain opts the cgroup's tasks
out entirely, and nothing below can re-enable it. Otherwise the
requested bits OR together and the kernel reads them independently, the
way kernel.numa_balancing reads them, so a combined NORMAL|TIERING
engages both modes. A DISABLED request only speaks for tasks, never for
the cgroup's memory. Demotion and the misplaced-migration fallback stay
global, keyed on node-level reclaim state rather than on any cgroup's
request. On the folio side the top-tier scanning exemption asks the
placement question, so a NORMAL|TIERING request also feeds the
placement statistics.

The combination is deliberate: a nearest-request-wins rule would let
any descendant re-enable machinery an ancestor turned off, and an
opt-out is only a subtree-wide guarantee if it binds. Requests add
instead, and the opt-out is what binds. The cost is that an ancestor
cannot cap tiering short of a full opt-out: a descendant's TIERING
request re-engages promotion for its own memory, and promotion is the
mode that consumes the shared fast tier.

The requests are combined at read time rather than propagated at write
time. The walk costs a few loads up a chain a handful of levels deep, a new
cgroup is correct by construction because it starts transparent and sees
its ancestors' current state, and bpf_numab_get_mode() runs the same walk
the kernel applies. Propagating a combined value down the subtree instead
would need hooks on cgroup creation to stay coherent with concurrent sets,
and cache a second field for the folio view, which reads opt-outs
differently.

The consult is behind the numab_cgroup_mode_set static key: while no cgroup
overrides the global state, the hot paths pay one patched-out branch and
nothing else. The readers run in atomic context (the scheduler tick,
hinting faults, load balancing, hardirq included), so the walk takes no
reference: rcu_read_lock() keeps the chain alive and the field is read
with READ_ONCE(). The kfunc accounts the static key transitions under a
mutex, and the cgroup free path drops the count of a request nobody
cleared, so the key goes off again once the last request is gone.
bpf_numab_set_mode() takes numab_mode_mutex and the jump label locks
inside static_branch_inc()/static_branch_dec(), so it is marked
KF_SLEEPABLE and only sleepable programs may call it.
bpf_numab_get_mode() never sleeps, so it carries no flag.
BPF_PROG_TYPE_SYSCALL programs run from user context, where sleeping is
allowed, so they qualify.

Suggested-by: Tejun Heo <tj@kernel.org>
Signed-off-by: Yafang Shao <laoar.shao@gmail.com>
---
 include/linux/cgroup-defs.h          |  10 ++
 include/linux/sched/numa_balancing.h | 118 +++++++++++++++-
 init/Kconfig                         |  20 +++
 kernel/cgroup/cgroup.c               |  14 ++
 kernel/sched/bpf_numab.c             | 195 +++++++++++++++++++++++++++
 kernel/sched/build_policy.c          |   4 +
 kernel/sched/fair.c                  |   4 +-
 mm/memory-tiers.c                    |   8 +-
 mm/mempolicy.c                       |   2 +-
 mm/migrate.c                         |   4 +-
 10 files changed, 369 insertions(+), 10 deletions(-)
 create mode 100644 kernel/sched/bpf_numab.c

diff --git a/include/linux/cgroup-defs.h b/include/linux/cgroup-defs.h
index 3754d697854b..5e97806634e4 100644
--- a/include/linux/cgroup-defs.h
+++ b/include/linux/cgroup-defs.h
@@ -639,6 +639,16 @@ struct cgroup {
 #ifdef CONFIG_EXT_SUB_SCHED
 	struct scx_sched __rcu *scx_sched;
 #endif
+#ifdef CONFIG_BPF_NUMAB
+	/*
+	 * The NUMA balancing mode requested for this cgroup:
+	 * NUMAB_FOLLOW_SYSTEM (-1) defers to the global sysctl state,
+	 * anything else is a NUMA_BALANCING_* bitmask. Written by
+	 * bpf_numab_set_mode() under numab_mode_mutex, read locklessly
+	 * under RCU.
+	 */
+	int numab_mode;
+#endif
 
 	/* All ancestors including self */
 	union {
diff --git a/include/linux/sched/numa_balancing.h b/include/linux/sched/numa_balancing.h
index 36de57f7fb06..d63891a8be01 100644
--- a/include/linux/sched/numa_balancing.h
+++ b/include/linux/sched/numa_balancing.h
@@ -36,9 +36,51 @@ 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. */
+#ifdef CONFIG_BPF_NUMAB
+/*
+ * A cgroup can override the global NUMA balancing state for the tasks it
+ * runs and the memory it owns: the bpf_numab_set_mode() kfunc stores a
+ * requested mode in struct cgroup, and the readers below combine the
+ * requests of the matching cgroup and of every ancestor.
+ * numab_cgroup_mode_set stays off while no cgroup makes a request, so
+ * the consult costs the hot paths nothing until then.
+ */
+extern struct static_key_false numab_cgroup_mode_set;
+
+/* The per-cgroup request domain: this, or a NUMA_BALANCING_* bitmask. */
+#define NUMAB_FOLLOW_SYSTEM (-1)
+
+/*
+ * The combined request for the chain of @p's cgroup:
+ * NUMA_BALANCING_DISABLED when any cgroup in the chain opts out, the OR
+ * of the requested bits otherwise, NUMAB_FOLLOW_SYSTEM when no cgroup
+ * makes a request.
+ */
+int task_numab_mode(struct task_struct *p);
+
+/*
+ * The combined request for the chain of the memcg @folio is charged to.
+ * An opt-out never speaks for memory, so it is not part of the answer;
+ * an uncharged folio makes no request.
+ */
+int folio_numab_mode(struct folio *folio);
+#endif
+
+/*
+ * Whether NUMA balancing applies to @p: the global sysctl state, possibly
+ * overridden by the requests of @p's cgroup and of every ancestor.
+ * A cgroup's request always takes priority over the system setting.
+ */
 static inline bool task_numab_enabled(struct task_struct *p)
 {
+#ifdef CONFIG_BPF_NUMAB
+	if (static_branch_unlikely(&numab_cgroup_mode_set)) {
+		int mode = task_numab_mode(p);
+
+		if (mode != NUMAB_FOLLOW_SYSTEM)
+			return mode != NUMA_BALANCING_DISABLED;
+	}
+#endif
 	return static_branch_likely(&sched_numa_balancing);
 }
 
@@ -53,6 +95,68 @@ static inline bool numab_in_normal_mode(void)
 {
 	return sysctl_numa_balancing_mode & NUMA_BALANCING_NORMAL;
 }
+
+/*
+ * 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; an opt-out wins, otherwise
+ * the bits are read independently, as the sysctl reads them. Without
+ * any request, the global sysctl decides.
+ */
+static inline bool task_numab_in_tiering_mode(struct task_struct *p)
+{
+#ifdef CONFIG_BPF_NUMAB
+	if (static_branch_unlikely(&numab_cgroup_mode_set)) {
+		int mode = task_numab_mode(p);
+
+		if (mode != NUMAB_FOLLOW_SYSTEM)
+			return mode & NUMA_BALANCING_MEMORY_TIERING;
+	}
+#endif
+	return numab_in_tiering_mode();
+}
+
+/*
+ * 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-side request), so only
+ * NORMAL and TIERING are consulted here; without either, the global
+ * sysctl decides.
+ */
+static inline bool folio_numab_in_tiering_mode(struct folio *folio)
+{
+#ifdef CONFIG_BPF_NUMAB
+	if (static_branch_unlikely(&numab_cgroup_mode_set)) {
+		int mode = folio_numab_mode(folio);
+
+		if (mode != NUMAB_FOLLOW_SYSTEM)
+			return mode & NUMA_BALANCING_MEMORY_TIERING;
+	}
+#endif
+	return numab_in_tiering_mode();
+}
+
+/*
+ * Whether @folio's owner asks for the normal mode, i.e. whether hinting
+ * faults on its top-tier pages should still feed the scheduler's
+ * placement statistics. The requests of the owner cgroup and of every
+ * ancestor are combined; without any request, the global sysctl
+ * decides. This is deliberately independent of tiering participation:
+ * a folio can record access times and still keep its top-tier scanning.
+ */
+static inline bool folio_numab_in_normal_mode(struct folio *folio)
+{
+#ifdef CONFIG_BPF_NUMAB
+	if (static_branch_unlikely(&numab_cgroup_mode_set)) {
+		int mode = folio_numab_mode(folio);
+
+		if (mode != NUMAB_FOLLOW_SYSTEM)
+			return mode & NUMA_BALANCING_NORMAL;
+	}
+#endif
+	return numab_in_normal_mode();
+}
 #else
 static inline void task_numa_fault(int last_node, int node, int pages,
 				   int flags)
@@ -85,6 +189,18 @@ static inline bool numab_in_normal_mode(void)
 {
 	return false;
 }
+static inline bool task_numab_in_tiering_mode(struct task_struct *p)
+{
+	return false;
+}
+static inline bool folio_numab_in_tiering_mode(struct folio *folio)
+{
+	return false;
+}
+static inline bool folio_numab_in_normal_mode(struct folio *folio)
+{
+	return false;
+}
 #endif
 
 #endif /* _LINUX_SCHED_NUMA_BALANCING_H */
diff --git a/init/Kconfig b/init/Kconfig
index ad592fdf29af..2db065d97eab 100644
--- a/init/Kconfig
+++ b/init/Kconfig
@@ -1036,6 +1036,26 @@ config NUMA_BALANCING
 
 	  This system will be inactive on UMA systems.
 
+config BPF_NUMAB
+	bool "Per-cgroup NUMA balancing mode selection"
+	default y
+	depends on NUMA_BALANCING && MEMCG && BPF_SYSCALL
+	help
+	  Allow BPF programs to select the NUMA balancing mode (disabled,
+	  normal, memory tiering, or both) on a per-cgroup basis, overriding the
+	  global NUMA balancing sysctl. A latency-sensitive service can
+	  opt its tasks out of balancing, or a batch service can have its
+	  PMEM-resident memory promoted when hot, without changing the
+	  mode for the rest of the system.
+
+	  The request is written and read back through the
+	  bpf_numab_set_mode() and bpf_numab_get_mode() kfuncs, callable
+	  from BPF_PROG_TYPE_SYSCALL programs. While no cgroup makes a
+	  request, the NUMA balancing hot paths pay a single patched-out
+	  branch and nothing else.
+
+	  If unsure, say Y.
+
 config SCHED_CACHE
 	bool "Cache aware load balance"
 	default y
diff --git a/kernel/cgroup/cgroup.c b/kernel/cgroup/cgroup.c
index 227d09704ca5..e9c01a017a1d 100644
--- a/kernel/cgroup/cgroup.c
+++ b/kernel/cgroup/cgroup.c
@@ -39,6 +39,7 @@
 #include <linux/proc_fs.h>
 #include <linux/rcupdate.h>
 #include <linux/sched.h>
+#include <linux/sched/numa_balancing.h>
 #include <linux/sched/task.h>
 #include <linux/slab.h>
 #include <linux/spinlock.h>
@@ -2122,6 +2123,9 @@ static void init_cgroup_housekeeping(struct cgroup *cgrp)
 	for (int i = 0; i < ARRAY_SIZE(cgrp->bpf.revisions); i++)
 		cgrp->bpf.revisions[i] = 1;
 #endif
+#ifdef CONFIG_BPF_NUMAB
+	cgrp->numab_mode = NUMAB_FOLLOW_SYSTEM;
+#endif
 
 	init_waitqueue_head(&cgrp->offline_waitq);
 	INIT_WORK(&cgrp->release_agent_work, cgroup1_release_agent);
@@ -5719,6 +5723,16 @@ static void css_free_rwork_fn(struct work_struct *work)
 			cgroup1_pidlist_destroy_all(cgrp);
 		cancel_work_sync(&cgrp->release_agent_work);
 		bpf_cgrp_storage_free(cgrp);
+#ifdef CONFIG_BPF_NUMAB
+		/*
+		 * Drop the static key count of a request the caller never
+		 * cleared. No setter can race: this path runs only after
+		 * the last css ref is gone, and the kfunc pins the cgroup
+		 * while it runs.
+		 */
+		if (cgrp->numab_mode != NUMAB_FOLLOW_SYSTEM)
+			static_branch_dec(&numab_cgroup_mode_set);
+#endif
 
 		if (cgroup_parent(cgrp)) {
 			/*
diff --git a/kernel/sched/bpf_numab.c b/kernel/sched/bpf_numab.c
new file mode 100644
index 000000000000..4a9da305ccb1
--- /dev/null
+++ b/kernel/sched/bpf_numab.c
@@ -0,0 +1,195 @@
+// SPDX-License-Identifier: GPL-2.0-or-later
+/*
+ * Per-cgroup NUMA balancing mode
+ *
+ * A cgroup can override the global NUMA balancing state for the tasks it
+ * runs and the memory it owns: the bpf_numab_set_mode() kfunc stores a
+ * requested mode in struct cgroup, and the NUMA balancing hot paths combine
+ * the requests of the matching cgroup and of every ancestor. For a task
+ * that is the chain of the task's cgroup; for a folio it is the chain of
+ * the memcg the folio is charged to. An opt-out wins and binds the
+ * subtree; otherwise the requested bits OR together and are read
+ * independently, the way the sysctl reads them.
+ *
+ * Author: Yafang Shao <laoar.shao@gmail.com>
+ */
+#include <linux/bpf.h>
+#include <linux/btf.h>
+#include <linux/cgroup.h>
+#include <linux/memcontrol.h>
+#include <linux/module.h>
+#include <linux/mutex.h>
+#include <linux/sched.h>
+#include <linux/sched/numa_balancing.h>
+
+/*
+ * Set while any cgroup overrides the global state: the hot paths skip the
+ * consult entirely while it stays off. Counted per cgroup that makes a
+ * request, so the key goes off again once the last request is gone.
+ */
+DEFINE_STATIC_KEY_FALSE(numab_cgroup_mode_set);
+
+/*
+ * Combine the requests of @cgrp and of every ancestor. A task chain
+ * stops at an explicit opt-out: a task can never be forced into
+ * balancing. A folio chain skips opt-outs instead (a cgroup's DISABLED
+ * request only speaks for its tasks, never for its memory) and keeps
+ * collecting the remaining requests.
+ */
+static int numab_cgroup_mode(struct cgroup *cgrp, bool folio_side)
+{
+	int req = 0;
+
+	for (; cgrp; cgrp = cgroup_parent(cgrp)) {
+		/* data_race: the field is written under numab_mode_mutex. */
+		int mode = data_race(READ_ONCE(cgrp->numab_mode));
+
+		if (mode == NUMAB_FOLLOW_SYSTEM)
+			continue;
+		if (mode == NUMA_BALANCING_DISABLED) {
+			if (!folio_side)
+				return NUMA_BALANCING_DISABLED;
+			continue;
+		}
+		req |= mode;
+	}
+	return req ?: NUMAB_FOLLOW_SYSTEM;
+}
+
+/*
+ * The combined request for @p. Kernel threads have no user memory to
+ * balance, so they make no request. The walk needs no reference: the
+ * cgroup chain is freed only after an RCU grace period, so
+ * rcu_read_lock() keeps it alive for the whole walk.
+ */
+int task_numab_mode(struct task_struct *p)
+{
+	int mode;
+
+	if (!p->mm)
+		return NUMAB_FOLLOW_SYSTEM;
+
+	rcu_read_lock();
+	mode = numab_cgroup_mode(task_dfl_cgroup(p), false);
+	rcu_read_unlock();
+	return mode;
+}
+
+/*
+ * The combined request for @folio's owner. An uncharged folio makes no
+ * request. A folio charged to a v1 memcg resolves to a cgroup chain the
+ * kfunc cannot reach, so it carries no request either and the folio
+ * follows the global state.
+ */
+int folio_numab_mode(struct folio *folio)
+{
+	struct mem_cgroup *memcg;
+	struct cgroup *cgrp;
+	int mode;
+
+	rcu_read_lock();
+	memcg = folio_memcg(folio);
+	cgrp = memcg ? memcg->css.cgroup : NULL;
+	mode = cgrp ? numab_cgroup_mode(cgrp, true) : NUMAB_FOLLOW_SYSTEM;
+	rcu_read_unlock();
+	return mode;
+}
+
+/*
+ * Serialize the static key accounting: without it, two concurrent
+ * setters on the same cgroup could both count one transition, and the
+ * key would stay on after the last request is gone.
+ */
+static DEFINE_MUTEX(numab_mode_mutex);
+
+__bpf_kfunc_start_defs();
+
+/**
+ * bpf_numab_set_mode() - Set the NUMA balancing mode requested for a cgroup.
+ * @cgid: kernfs id of the cgroup, e.g. from name_to_handle_at() on its
+ * directory or from bpf_get_current_cgroup_id()
+ * @mode: NUMAB_FOLLOW_SYSTEM to defer to the global sysctl state, or a
+ * NUMA_BALANCING_* bitmask. The kernel reads the bits independently,
+ * the way kernel.numa_balancing reads them, so NORMAL|TIERING engages
+ * both modes. A DISABLED request only speaks for the cgroup's tasks,
+ * never for its memory, which follows the remaining requests of its
+ * ancestors.
+ *
+ * Return: 0 on success, -EINVAL for a mode outside the domain, -ENOENT
+ * for an id that resolves to no cgroup.
+ */
+__bpf_kfunc int bpf_numab_set_mode(u64 cgid, s64 mode)
+{
+	struct cgroup *cgrp;
+	int old;
+
+	if (mode != NUMAB_FOLLOW_SYSTEM &&
+	    (mode < 0 || mode > (NUMA_BALANCING_NORMAL |
+				 NUMA_BALANCING_MEMORY_TIERING)))
+		return -EINVAL;
+
+	cgrp = __cgroup_get_from_id(cgid);
+	if (IS_ERR(cgrp))
+		return PTR_ERR(cgrp);
+
+	mutex_lock(&numab_mode_mutex);
+	old = cgrp->numab_mode;
+	if (old == NUMAB_FOLLOW_SYSTEM && mode != NUMAB_FOLLOW_SYSTEM)
+		static_branch_inc(&numab_cgroup_mode_set);
+	else if (old != NUMAB_FOLLOW_SYSTEM && mode == NUMAB_FOLLOW_SYSTEM)
+		static_branch_dec(&numab_cgroup_mode_set);
+	WRITE_ONCE(cgrp->numab_mode, (int)mode);
+	mutex_unlock(&numab_mode_mutex);
+	cgroup_put(cgrp);
+	return 0;
+}
+
+/**
+ * bpf_numab_get_mode() - Read the NUMA balancing mode that applies to a cgroup.
+ * @cgid: kernfs id of the cgroup
+ *
+ * The combined request of the cgroup and of every ancestor: an opt-out
+ * wins, otherwise the requested bits OR together, and the kernel reads
+ * them independently, as kernel.numa_balancing does, so a combined
+ * NORMAL|TIERING engages both modes. NUMAB_FOLLOW_SYSTEM when nobody
+ * makes a request. This is the view of the cgroup's tasks: its memory
+ * follows the remaining requests of the chain when it opts out, and
+ * that folio-side view stays internal to the kernel.
+ *
+ * Return: the combined mode, NUMAB_FOLLOW_SYSTEM (-1) included, or
+ * -ENOENT for an id that resolves to no cgroup.
+ */
+__bpf_kfunc s64 bpf_numab_get_mode(u64 cgid)
+{
+	struct cgroup *cgrp;
+	s64 mode;
+
+	cgrp = __cgroup_get_from_id(cgid);
+	if (IS_ERR(cgrp))
+		return PTR_ERR(cgrp);
+
+	rcu_read_lock();
+	mode = numab_cgroup_mode(cgrp, false);
+	rcu_read_unlock();
+
+	cgroup_put(cgrp);
+	return mode;
+}
+
+__bpf_kfunc_end_defs();
+
+BTF_KFUNCS_START(numab_kfunc_ids)
+BTF_ID_FLAGS(func, bpf_numab_set_mode, KF_SLEEPABLE)
+BTF_ID_FLAGS(func, bpf_numab_get_mode)
+BTF_KFUNCS_END(numab_kfunc_ids)
+
+static const struct btf_kfunc_id_set numab_kfunc_set = {
+	.owner = THIS_MODULE,
+	.set   = &numab_kfunc_ids,
+};
+
+static int __init numab_kfunc_init(void)
+{
+	return register_btf_kfunc_id_set(BPF_PROG_TYPE_SYSCALL, &numab_kfunc_set);
+}
+late_initcall(numab_kfunc_init);
diff --git a/kernel/sched/build_policy.c b/kernel/sched/build_policy.c
index 2a828725a7f9..3746c6675ff3 100644
--- a/kernel/sched/build_policy.c
+++ b/kernel/sched/build_policy.c
@@ -75,4 +75,8 @@
 # include "ext/sub.c"
 #endif
 
+#ifdef CONFIG_BPF_NUMAB
+#include "bpf_numab.c"
+#endif
+
 #include "syscalls.c"
diff --git a/kernel/sched/fair.c b/kernel/sched/fair.c
index 066818498c76..2a4ad93a262a 100644
--- a/kernel/sched/fair.c
+++ b/kernel/sched/fair.c
@@ -2965,7 +2965,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_in_tiering_mode()) && !node_is_toptier(src_nid) &&
+	if (!folio_numab_in_tiering_mode(folio) && !node_is_toptier(src_nid) &&
 	    !cpupid_valid(last_cpupid))
 		return false;
 
@@ -4208,7 +4208,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_in_tiering_mode() || !cpupid_valid(last_cpupid)))
+	    (task_numab_in_tiering_mode(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 9009b164edc2..c68ce7f9171e 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_in_tiering_mode() &&
+	return folio_numab_in_tiering_mode(folio) &&
 	       !node_is_toptier(folio_nid(folio));
 }
 #endif
diff --git a/mm/mempolicy.c b/mm/mempolicy.c
index 442c93079fb0..fe8ae6cf5c58 100644
--- a/mm/mempolicy.c
+++ b/mm/mempolicy.c
@@ -871,7 +871,7 @@ bool folio_can_map_prot_numa(struct folio *folio, struct vm_area_struct *vma,
 	 * Skip scanning top tier node if normal numa
 	 * balancing is disabled
 	 */
-	if (!numab_in_normal_mode() && node_is_toptier(nid))
+	if (!folio_numab_in_normal_mode(folio) && node_is_toptier(nid))
 		return false;
 
 	if (folio_use_access_time(folio))
diff --git a/mm/migrate.c b/mm/migrate.c
index a622df9a9b42..4b87d465c5fd 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_in_tiering_mode()) {
+	if (folio_numab_in_tiering_mode(folio)) {
 		bool f_toptier = node_is_toptier(folio_nid(folio));
 		bool t_toptier = node_is_toptier(folio_nid(newfolio));
 
@@ -2771,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 ((numab_in_tiering_mode())
+		if (folio_numab_in_tiering_mode(folio)
 		    && !node_is_toptier(folio_nid(folio))
 		    && node_is_toptier(node))
 			mod_lruvec_state(lruvec, PGPROMOTE_SUCCESS, nr_succeeded);
-- 
2.52.0

[RFC PATCH v3 bpf-next 5/5] selftests/bpf: add tests for the numab kfuncs

[Yafang Shao]

Exercise the bpf_numab_set_mode() and bpf_numab_get_mode() kfuncs
through BPF_PROG_TYPE_SYSCALL programs: the set/get roundtrip, the
mode domain (NUMAB_FOLLOW_SYSTEM and the NUMA_BALANCING_* bitmasks
accepted, anything else rejected with -EINVAL), unknown cgroup ids
rejected with -ENOENT, and the hierarchy combination the ancestor walk
produces (inheritance, OR-ing, and an opt-out winning over every
request in the chain).

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

diff --git a/tools/testing/selftests/bpf/prog_tests/bpf_numab.c b/tools/testing/selftests/bpf/prog_tests/bpf_numab.c
new file mode 100644
index 000000000000..7ce4b8153c42
--- /dev/null
+++ b/tools/testing/selftests/bpf/prog_tests/bpf_numab.c
@@ -0,0 +1,227 @@
+// SPDX-License-Identifier: GPL-2.0
+
+#include <test_progs.h>
+#include "cgroup_helpers.h"
+#include "bpf_numab.skel.h"
+
+#define CGROUP_PATH	"/bpf_numab"
+#define CHILD_PATH	CGROUP_PATH "/child"
+
+/* The per-cgroup request domain, mirroring the kernel constants. */
+#define NUMAB_FOLLOW_SYSTEM	(-1)
+#define NUMAB_DISABLED		0
+#define NUMAB_NORMAL		1
+#define NUMAB_TIERING		2
+#define NUMAB_NORMAL_TIERING	3
+
+static int numab_set(struct bpf_numab *skel, u64 cgid,
+		     long mode)
+{
+	LIBBPF_OPTS(bpf_test_run_opts, opts);
+	int ret;
+
+	skel->bss->set_cgid = cgid;
+	skel->bss->set_mode = mode;
+	ret = bpf_prog_test_run_opts(bpf_program__fd(skel->progs.numab_set),
+				     &opts);
+	if (!ASSERT_OK(ret, "set run"))
+		return -1;
+	return skel->bss->set_ret;
+}
+
+static long numab_get(struct bpf_numab *skel, u64 cgid)
+{
+	LIBBPF_OPTS(bpf_test_run_opts, opts);
+	int ret;
+
+	skel->bss->get_cgid = cgid;
+	ret = bpf_prog_test_run_opts(bpf_program__fd(skel->progs.numab_get),
+				     &opts);
+	if (!ASSERT_OK(ret, "get run"))
+		return -1;
+	return skel->bss->get_mode;
+}
+
+/*
+ * Roundtrip through the combined view: the test cgroup has no ancestor
+ * making a request, so the combined mode is the stored one.
+ */
+static void test_set_get(void)
+{
+	struct bpf_numab *skel;
+	int cgroup_fd;
+	u64 cgid;
+
+	cgroup_fd = test__join_cgroup(CGROUP_PATH);
+	if (!ASSERT_GE(cgroup_fd, 0, "join_cgroup"))
+		return;
+
+	skel = bpf_numab__open_and_load();
+	if (!ASSERT_OK_PTR(skel, "open_and_load"))
+		goto cleanup;
+
+	cgid = get_cgroup_id(CGROUP_PATH);
+	if (!ASSERT_GT(cgid, 0, "get_cgroup_id"))
+		goto cleanup;
+
+	ASSERT_EQ(numab_set(skel, cgid, NUMAB_NORMAL), 0, "set normal");
+	ASSERT_EQ(numab_get(skel, cgid), NUMAB_NORMAL, "get normal");
+
+	ASSERT_EQ(numab_set(skel, cgid, NUMAB_TIERING), 0, "set tiering");
+	ASSERT_EQ(numab_get(skel, cgid), NUMAB_TIERING, "get tiering");
+
+	ASSERT_EQ(numab_set(skel, cgid, NUMAB_NORMAL_TIERING), 0, "set both");
+	ASSERT_EQ(numab_get(skel, cgid), NUMAB_NORMAL_TIERING, "get both");
+
+	ASSERT_EQ(numab_set(skel, cgid, NUMAB_DISABLED), 0, "set disabled");
+	ASSERT_EQ(numab_get(skel, cgid), NUMAB_DISABLED, "get disabled");
+
+	/* Back to the system setting. */
+	ASSERT_EQ(numab_set(skel, cgid, NUMAB_FOLLOW_SYSTEM), 0, "unset");
+	ASSERT_EQ(numab_get(skel, cgid), NUMAB_FOLLOW_SYSTEM, "get unset");
+
+cleanup:
+	bpf_numab__destroy(skel);
+	close(cgroup_fd);
+}
+
+/*
+ * The mode domain: NUMAB_FOLLOW_SYSTEM and the NUMA_BALANCING_* bitmasks
+ * are accepted, anything else is rejected with -EINVAL.
+ */
+static void test_invalid_mode(void)
+{
+	struct bpf_numab *skel;
+	int cgroup_fd;
+	u64 cgid;
+
+	cgroup_fd = test__join_cgroup(CGROUP_PATH);
+	if (!ASSERT_GE(cgroup_fd, 0, "join_cgroup"))
+		return;
+
+	skel = bpf_numab__open_and_load();
+	if (!ASSERT_OK_PTR(skel, "open_and_load"))
+		goto cleanup;
+
+	cgid = get_cgroup_id(CGROUP_PATH);
+	if (!ASSERT_GT(cgid, 0, "get_cgroup_id"))
+		goto cleanup;
+
+	ASSERT_EQ(numab_set(skel, cgid, -2), -EINVAL, "set below domain");
+	ASSERT_EQ(numab_set(skel, cgid, 4), -EINVAL, "set above domain");
+	ASSERT_EQ(numab_set(skel, cgid, 0x1000), -EINVAL, "set far above");
+
+	/* The rejected requests must not have stuck. */
+	ASSERT_EQ(numab_get(skel, cgid), NUMAB_FOLLOW_SYSTEM, "get unchanged");
+
+cleanup:
+	bpf_numab__destroy(skel);
+	close(cgroup_fd);
+}
+
+/*
+ * An id that resolves to no cgroup is rejected with -ENOENT.
+ */
+static void test_bad_cgid(void)
+{
+	struct bpf_numab *skel;
+
+	skel = bpf_numab__open_and_load();
+	if (!ASSERT_OK_PTR(skel, "open_and_load"))
+		return;
+
+	ASSERT_EQ(numab_set(skel, 0, NUMAB_NORMAL), -ENOENT, "set cgid 0");
+	ASSERT_EQ(numab_get(skel, 0), -ENOENT, "get cgid 0");
+	ASSERT_EQ(numab_set(skel, 0xdeadbeef, NUMAB_NORMAL), -ENOENT,
+		  "set bad cgid");
+	ASSERT_EQ(numab_get(skel, 0xdeadbeef), -ENOENT, "get bad cgid");
+
+	bpf_numab__destroy(skel);
+}
+
+/*
+ * The requests of a cgroup and of every ancestor combine: an opt-out
+ * wins over every request in the chain, otherwise the requested bits OR
+ * together, and an unset cgroup inherits what its ancestors ask for.
+ */
+static void test_hierarchy(void)
+{
+	u64 parent_cgid, child_cgid;
+	int parent_fd, child_fd;
+	struct bpf_numab *skel;
+
+	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_parent;
+
+	skel = bpf_numab__open_and_load();
+	if (!ASSERT_OK_PTR(skel, "open_and_load"))
+		goto cleanup;
+
+	parent_cgid = get_cgroup_id(CGROUP_PATH);
+	child_cgid = get_cgroup_id(CHILD_PATH);
+	if (!ASSERT_GT(parent_cgid, 0, "get_parent_cgid") ||
+	    !ASSERT_GT(child_cgid, 0, "get_child_cgid"))
+		goto cleanup;
+
+	/* The child inherits an ancestor's request untouched. */
+	ASSERT_EQ(numab_set(skel, parent_cgid, NUMAB_TIERING), 0,
+		  "set parent tiering");
+	ASSERT_EQ(numab_get(skel, child_cgid), NUMAB_TIERING,
+		  "child inherits");
+
+	/* A request ORs with the ancestor's; NORMAL does not undo TIERING. */
+	ASSERT_EQ(numab_set(skel, child_cgid, NUMAB_NORMAL), 0,
+		  "set child normal");
+	ASSERT_EQ(numab_get(skel, child_cgid), NUMAB_NORMAL_TIERING,
+		  "or both");
+
+	/* The child's own opt-out wins. */
+	ASSERT_EQ(numab_set(skel, child_cgid, NUMAB_DISABLED), 0,
+		  "set child disabled");
+	ASSERT_EQ(numab_get(skel, child_cgid), NUMAB_DISABLED,
+		  "own opt-out wins");
+
+	/* So does an ancestor's: a task is never forced into balancing. */
+	ASSERT_EQ(numab_set(skel, child_cgid, NUMAB_TIERING), 0,
+		  "set child tiering");
+	ASSERT_EQ(numab_set(skel, parent_cgid, NUMAB_DISABLED), 0,
+		  "set parent disabled");
+	ASSERT_EQ(numab_get(skel, child_cgid), NUMAB_DISABLED,
+		  "ancestor opt-out wins");
+
+	/* Unsetting the parent re-exposes the child's own request. */
+	ASSERT_EQ(numab_set(skel, parent_cgid, NUMAB_FOLLOW_SYSTEM), 0,
+		  "unset parent");
+	ASSERT_EQ(numab_get(skel, child_cgid), NUMAB_TIERING,
+		  "child request again");
+
+	/* Unsetting the child leaves the chain with no request. */
+	ASSERT_EQ(numab_set(skel, child_cgid, NUMAB_FOLLOW_SYSTEM), 0,
+		  "unset child");
+	ASSERT_EQ(numab_get(skel, child_cgid), NUMAB_FOLLOW_SYSTEM,
+		  "no request");
+
+cleanup:
+	bpf_numab__destroy(skel);
+	close(child_fd);
+	remove_cgroup(CHILD_PATH);
+cleanup_parent:
+	close(parent_fd);
+}
+
+void test_bpf_numab(void)
+{
+	if (test__start_subtest("set_get"))
+		test_set_get();
+	if (test__start_subtest("invalid_mode"))
+		test_invalid_mode();
+	if (test__start_subtest("bad_cgid"))
+		test_bad_cgid();
+	if (test__start_subtest("hierarchy"))
+		test_hierarchy();
+}
diff --git a/tools/testing/selftests/bpf/progs/bpf_numab.c b/tools/testing/selftests/bpf/progs/bpf_numab.c
new file mode 100644
index 000000000000..b386cbbc19fa
--- /dev/null
+++ b/tools/testing/selftests/bpf/progs/bpf_numab.c
@@ -0,0 +1,32 @@
+// SPDX-License-Identifier: GPL-2.0
+
+#include <vmlinux.h>
+#include <bpf/bpf_helpers.h>
+
+/* Inputs, set from userspace before each run. */
+long set_cgid;
+long set_mode;
+long get_cgid;
+
+/* Results, read back after each run. */
+long set_ret;
+long get_mode;
+
+extern int bpf_numab_set_mode(u64 cgid, s64 mode) __ksym;
+extern s64 bpf_numab_get_mode(u64 cgid) __ksym;
+
+SEC("syscall")
+int numab_set(void *ctx)
+{
+	set_ret = bpf_numab_set_mode((u64)set_cgid, set_mode);
+	return 0;
+}
+
+SEC("syscall")
+int numab_get(void *ctx)
+{
+	get_mode = bpf_numab_get_mode((u64)get_cgid);
+	return 0;
+}
+
+char _license[] SEC("license") = "GPL";
-- 
2.52.0

Re: [RFC PATCH v3 bpf-next 5/5] selftests/bpf: add tests for the numab kfuncs

[sashiko-bot]

Thank you for your contribution! Sashiko AI review found 1 potential issue(s) to consider:
- [Medium] selftests/bpf: missing kfunc BTF skip check in test_bpf_numab()
--

commit 07e4c14213f29b11a7ce19b88e46e524dc48e78d
Author: Yafang Shao <laoar.shao@gmail.com>

selftests/bpf: add tests for the numab kfuncs

This commit adds BPF selftests exercising the bpf_numab_set_mode() and
bpf_numab_get_mode() kfuncs across set/get roundtrips, mode validation,
invalid cgroup IDs, and cgroup hierarchy combinations.

> diff --git a/tools/testing/selftests/bpf/prog_tests/bpf_numab.c b/tools/testing/selftests/bpf/prog_tests/bpf_numab.c
> new file mode 100644
> index 0000000000000..7ce4b8153c42f
> --- /dev/null
> +++ b/tools/testing/selftests/bpf/prog_tests/bpf_numab.c
> @@ -0,0 +1,227 @@
[ ... ]
> +static void test_set_get(void)
> +{
> +	struct bpf_numab *skel;
> +	int cgroup_fd;
> +	u64 cgid;
> +
> +	cgroup_fd = test__join_cgroup(CGROUP_PATH);
> +	if (!ASSERT_GE(cgroup_fd, 0, "join_cgroup"))
> +		return;
> +
> +	skel = bpf_numab__open_and_load();
> +	if (!ASSERT_OK_PTR(skel, "open_and_load"))
> +		goto cleanup;
[ ... ]
> +void test_bpf_numab(void)
> +{
> +	if (test__start_subtest("set_get"))
> +		test_set_get();
> +	if (test__start_subtest("invalid_mode"))
> +		test_invalid_mode();
> +	if (test__start_subtest("bad_cgid"))
> +		test_bad_cgid();
> +	if (test__start_subtest("hierarchy"))
> +		test_hierarchy();
> +}

[Severity: Medium]
What happens when test_bpf_numab() runs on a kernel built without
CONFIG_BPF_NUMAB=y (for example when CONFIG_NUMA_BALANCING is disabled or on
non-x86_64 selftest configs where CONFIG_NUMA_BALANCING is not enabled)?

Because tools/testing/selftests/bpf/progs/bpf_numab.c declares
bpf_numab_set_mode and bpf_numab_get_mode as strong __ksym externs:

    extern int bpf_numab_set_mode(u64 cgid, s64 mode) __ksym;
    extern s64 bpf_numab_get_mode(u64 cgid) __ksym;

calling test_bpf_numab() -> test_set_get() -> bpf_numab__open_and_load()
(as well as test_invalid_mode(), test_bad_cgid(), and test_hierarchy())
fails in bpf_object__resolve_ksym_func_btf_id() with -ESRCH and trips
ASSERT_OK_PTR(skel, "open_and_load").

Should test_bpf_numab() check vmlinux BTF for bpf_numab_set_mode and
bpf_numab_get_mode via btf__find_by_name_kind(..., BTF_KIND_FUNC) and call
test__skip() when the kfuncs are not present?

> diff --git a/tools/testing/selftests/bpf/progs/bpf_numab.c b/tools/testing/selftests/bpf/progs/bpf_numab.c
[ ... ]

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

Re: [RFC PATCH v3 bpf-next 5/5] selftests/bpf: add tests for the numab kfuncs

[Yafang Shao]

On Mon, Oct 19, 2026 at 2:51 PM <sashiko-bot@kernel.org> wrote:
>
> Thank you for your contribution! Sashiko AI review found 1 potential issue(s) to consider:
> - [Medium] selftests/bpf: missing kfunc BTF skip check in test_bpf_numab()
> --
>
> commit 07e4c14213f29b11a7ce19b88e46e524dc48e78d
> Author: Yafang Shao <laoar.shao@gmail.com>
>
> selftests/bpf: add tests for the numab kfuncs
>
> This commit adds BPF selftests exercising the bpf_numab_set_mode() and
> bpf_numab_get_mode() kfuncs across set/get roundtrips, mode validation,
> invalid cgroup IDs, and cgroup hierarchy combinations.
>
> > diff --git a/tools/testing/selftests/bpf/prog_tests/bpf_numab.c b/tools/testing/selftests/bpf/prog_tests/bpf_numab.c
> > new file mode 100644
> > index 0000000000000..7ce4b8153c42f
> > --- /dev/null
> > +++ b/tools/testing/selftests/bpf/prog_tests/bpf_numab.c
> > @@ -0,0 +1,227 @@
> [ ... ]
> > +static void test_set_get(void)
> > +{
> > +     struct bpf_numab *skel;
> > +     int cgroup_fd;
> > +     u64 cgid;
> > +
> > +     cgroup_fd = test__join_cgroup(CGROUP_PATH);
> > +     if (!ASSERT_GE(cgroup_fd, 0, "join_cgroup"))
> > +             return;
> > +
> > +     skel = bpf_numab__open_and_load();
> > +     if (!ASSERT_OK_PTR(skel, "open_and_load"))
> > +             goto cleanup;
> [ ... ]
> > +void test_bpf_numab(void)
> > +{
> > +     if (test__start_subtest("set_get"))
> > +             test_set_get();
> > +     if (test__start_subtest("invalid_mode"))
> > +             test_invalid_mode();
> > +     if (test__start_subtest("bad_cgid"))
> > +             test_bad_cgid();
> > +     if (test__start_subtest("hierarchy"))
> > +             test_hierarchy();
> > +}
>
> [Severity: Medium]
> What happens when test_bpf_numab() runs on a kernel built without
> CONFIG_BPF_NUMAB=y (for example when CONFIG_NUMA_BALANCING is disabled or on
> non-x86_64 selftest configs where CONFIG_NUMA_BALANCING is not enabled)?

Nice catch.
will add the configs to tools/testing/selftests/bpf/config.

>
> Because tools/testing/selftests/bpf/progs/bpf_numab.c declares
> bpf_numab_set_mode and bpf_numab_get_mode as strong __ksym externs:
>
>     extern int bpf_numab_set_mode(u64 cgid, s64 mode) __ksym;
>     extern s64 bpf_numab_get_mode(u64 cgid) __ksym;
>
> calling test_bpf_numab() -> test_set_get() -> bpf_numab__open_and_load()
> (as well as test_invalid_mode(), test_bad_cgid(), and test_hierarchy())
> fails in bpf_object__resolve_ksym_func_btf_id() with -ESRCH and trips
> ASSERT_OK_PTR(skel, "open_and_load").
>
> Should test_bpf_numab() check vmlinux BTF for bpf_numab_set_mode and
> bpf_numab_get_mode via btf__find_by_name_kind(..., BTF_KIND_FUNC) and call
> test__skip() when the kfuncs are not present?
>
> > diff --git a/tools/testing/selftests/bpf/progs/bpf_numab.c b/tools/testing/selftests/bpf/progs/bpf_numab.c
> [ ... ]
>
> --
> Sashiko AI review · https://sashiko.dev/#/patchset/20261019064157.29885-1-laoar.shao@gmail.com?part=5



-- 
Regards
Yafang

Re: [RFC PATCH v3 bpf-next 4/5] sched: add BPF-based per-cgroup NUMA balancing mode selection

[sashiko-bot]

Thank you for your contribution! Sashiko AI review found 1 potential issue(s) to consider:
- [Critical] mm: use-after-free on source folio in migrate_misplaced_folio()
--

commit 5de3ed8c82d3030544f5f025b8ae47dd8339bfeb
Author: Yafang Shao <laoar.shao@gmail.com>

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

This commit introduces per-cgroup NUMA balancing mode selection managed via
the bpf_numab_set_mode() and bpf_numab_get_mode() kfuncs. Hot paths combine
requests along the task or folio's memcg cgroup hierarchy behind the
numab_cgroup_mode_set static key to override the global sysctl setting.

> diff --git a/mm/migrate.c b/mm/migrate.c
> index 211c29bd2a283..c0692c7b27467 100644
> --- a/mm/migrate.c
> +++ b/mm/migrate.c

[ ... ]

> @@ -2770,7 +2770,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 ((numab_in_tiering_mode())
> +		if (folio_numab_in_tiering_mode(folio)
>  		    && !node_is_toptier(folio_nid(folio))
>  		    && node_is_toptier(node))
>  			mod_lruvec_state(lruvec, PGPROMOTE_SUCCESS, nr_succeeded);

[Severity: Critical]
Has the source folio already been uncharged and freed by the time
migrate_misplaced_folio() calls folio_numab_in_tiering_mode(folio) here?

Before calling migrate_pages(), migrate_misplaced_folio() snapshots the
folio's memcg because migration transfers the charge and releases the
isolation reference on the source folio:

migrate_misplaced_folio()
  memcg = get_mem_cgroup_from_folio(folio);
  migrate_pages()
    migrate_pages_batch()
      migrate_folios_move()
        migrate_folio_move()
          move_to_new_folio()
            folio_migrate_flags()
              mem_cgroup_migrate()
                old->memcg_data = 0;
          migrate_folio_done()
            folio_put(src);

When nr_succeeded is non-zero, calling folio_numab_in_tiering_mode(folio) ->
folio_numab_mode(folio) -> folio_memcg(folio) -> folio_objcg(folio) reads
folio->memcg_data on the freed source folio:

static inline struct obj_cgroup *folio_objcg(const struct folio *folio)
{
	unsigned long memcg_data = folio->memcg_data;

	VM_BUG_ON_FOLIO(folio_test_slab(folio), folio);
	VM_BUG_ON_FOLIO(memcg_data & MEMCG_DATA_OBJEXTS, folio);

	return (struct obj_cgroup *)(memcg_data & ~OBJEXTS_FLAGS_MASK);
}

If the freed page has not been reallocated yet, folio->memcg_data is 0, so
folio_memcg(folio) returns NULL and folio_numab_mode(folio) falls back to
NUMAB_FOLLOW_SYSTEM rather than using the owner cgroup's tiering mode for
PGPROMOTE_SUCCESS accounting.

If another CPU reallocates the freed page before this check runs (for example
as a slab page setting PG_slab and MEMCG_DATA_OBJEXTS, or as a compound tail
page), can folio_objcg(folio) trigger VM_BUG_ON_FOLIO() or dereference a
slabobj_ext pointer as a struct obj_cgroup *?

Would it make sense to evaluate folio_numab_in_tiering_mode(folio) before
migrate_pages(), or derive the mode from the already-pinned memcg pointer?

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

Re: [RFC PATCH v3 bpf-next 3/5] sched: add helpers for numa balancing

[bot+bpf-ci]

[-- Attachment #1: Type: text/plain, Size: 2044 bytes --]

> sched: add helpers for numa balancing

TL;DR: Nits only: redundant parentheses around the new numab_in_tiering_mode()
helper calls in should_numa_migrate_memory() and migrate_misplaced_folio()
could be dropped.

> diff --git a/kernel/sched/fair.c b/kernel/sched/fair.c
> index 71dc47d6114d..066818498c76 100644
> --- a/kernel/sched/fair.c
> +++ b/kernel/sched/fair.c

[ ... ]

> @@ -2963,8 +2965,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_in_tiering_mode()) && !node_is_toptier(src_nid) &&
> +	    !cpupid_valid(last_cpupid))
>  		return false;

this isn't a bug, but could the extra parentheses in
!(numab_in_tiering_mode()) be dropped here?  They were needed around the
old bitwise test, and the other converted sites such as
task_numa_fault() call the helper without them.

[ ... ]

> diff --git a/mm/migrate.c b/mm/migrate.c
> index 15b45832bcfa..a622df9a9b42 100644
> --- a/mm/migrate.c
> +++ b/mm/migrate.c

[ ... ]

> @@ -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_in_tiering_mode())
>  		    && !node_is_toptier(folio_nid(folio))
>  		    && node_is_toptier(node))
>  			mod_lruvec_state(lruvec, PGPROMOTE_SUCCESS, nr_succeeded);

The same applies to the parentheses in
if ((numab_in_tiering_mode()) here.


---
AI reviewed your patch. Please fix the bug or email reply why it's not a bug.
See: https://github.com/kernel-patches/vmtest/blob/master/ci/claude/README.md

CI run summary: https://github.com/kernel-patches/bpf/actions/runs/39213018427

Re: [RFC PATCH v3 bpf-next 4/5] sched: add BPF-based per-cgroup NUMA balancing mode selection

[bot+bpf-ci]

[-- Attachment #1: Type: text/plain, Size: 2148 bytes --]

> sched: add BPF-based per-cgroup NUMA balancing mode selection

TL;DR: Comment only: the comment on task_numab_in_tiering_mode() may describe
the return value backwards, since task_numa_fault() skips non-top-tier
statistics when it returns true.

> diff --git a/include/linux/sched/numa_balancing.h b/include/linux/sched/numa_balancing.h
> index 36de57f7fb06..d63891a8be01 100644
> --- a/include/linux/sched/numa_balancing.h
> +++ b/include/linux/sched/numa_balancing.h

[ ... ]

> @@ -53,6 +95,68 @@ static inline bool numab_in_normal_mode(void)
>  {
>  	return sysctl_numa_balancing_mode & NUMA_BALANCING_NORMAL;
>  }
> +
> +/*
> + * 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; an opt-out wins, otherwise
> + * the bits are read independently, as the sysctl reads them. Without
> + * any request, the global sysctl decides.
> + */
> +static inline bool task_numab_in_tiering_mode(struct task_struct *p)
> +{
> +#ifdef CONFIG_BPF_NUMAB
> +	if (static_branch_unlikely(&numab_cgroup_mode_set)) {
> +		int mode = task_numab_mode(p);
> +
> +		if (mode != NUMAB_FOLLOW_SYSTEM)
> +			return mode & NUMA_BALANCING_MEMORY_TIERING;
> +	}
> +#endif
> +	return numab_in_tiering_mode();
> +}

The comment above task_numab_in_tiering_mode() says it reports "whether
faults on non-top-tier memory still feed its NUMA statistics".  Is that
the right way around?

The only caller, task_numa_fault() in kernel/sched/fair.c, uses it like
this:

	if (!node_is_toptier(mem_node) &&
	    (task_numab_in_tiering_mode(p) || !cpupid_valid(last_cpupid)))
		return;

so a true return makes task_numa_fault() bail out before recording any
statistics for non-top-tier memory.  Should the comment say that faults on
non-top-tier memory are kept out of the task's NUMA statistics instead?


---
AI reviewed your patch. Please fix the bug or email reply why it's not a bug.
See: https://github.com/kernel-patches/vmtest/blob/master/ci/claude/README.md

CI run summary: https://github.com/kernel-patches/bpf/actions/runs/39213018427

Re: [RFC PATCH v3 bpf-next 4/5] sched: add BPF-based per-cgroup NUMA balancing mode selection

[Yafang Shao]

On Mon, Oct 19, 2026 at 2:58 PM <sashiko-bot@kernel.org> wrote:
>
> Thank you for your contribution! Sashiko AI review found 1 potential issue(s) to consider:
> - [Critical] mm: use-after-free on source folio in migrate_misplaced_folio()
> --
>
> commit 5de3ed8c82d3030544f5f025b8ae47dd8339bfeb
> Author: Yafang Shao <laoar.shao@gmail.com>
>
> sched: add BPF-based per-cgroup NUMA balancing mode selection
>
> This commit introduces per-cgroup NUMA balancing mode selection managed via
> the bpf_numab_set_mode() and bpf_numab_get_mode() kfuncs. Hot paths combine
> requests along the task or folio's memcg cgroup hierarchy behind the
> numab_cgroup_mode_set static key to override the global sysctl setting.
>
> > diff --git a/mm/migrate.c b/mm/migrate.c
> > index 211c29bd2a283..c0692c7b27467 100644
> > --- a/mm/migrate.c
> > +++ b/mm/migrate.c
>
> [ ... ]
>
> > @@ -2770,7 +2770,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 ((numab_in_tiering_mode())
> > +             if (folio_numab_in_tiering_mode(folio)
> >                   && !node_is_toptier(folio_nid(folio))
> >                   && node_is_toptier(node))
> >                       mod_lruvec_state(lruvec, PGPROMOTE_SUCCESS, nr_succeeded);
>
> [Severity: Critical]
> Has the source folio already been uncharged and freed by the time
> migrate_misplaced_folio() calls folio_numab_in_tiering_mode(folio) here?
>
> Before calling migrate_pages(), migrate_misplaced_folio() snapshots the
> folio's memcg because migration transfers the charge and releases the
> isolation reference on the source folio:
>
> migrate_misplaced_folio()
>   memcg = get_mem_cgroup_from_folio(folio);
>   migrate_pages()
>     migrate_pages_batch()
>       migrate_folios_move()
>         migrate_folio_move()
>           move_to_new_folio()
>             folio_migrate_flags()
>               mem_cgroup_migrate()
>                 old->memcg_data = 0;
>           migrate_folio_done()
>             folio_put(src);
>
> When nr_succeeded is non-zero, calling folio_numab_in_tiering_mode(folio) ->
> folio_numab_mode(folio) -> folio_memcg(folio) -> folio_objcg(folio) reads
> folio->memcg_data on the freed source folio:
>
> static inline struct obj_cgroup *folio_objcg(const struct folio *folio)
> {
>         unsigned long memcg_data = folio->memcg_data;
>
>         VM_BUG_ON_FOLIO(folio_test_slab(folio), folio);
>         VM_BUG_ON_FOLIO(memcg_data & MEMCG_DATA_OBJEXTS, folio);
>
>         return (struct obj_cgroup *)(memcg_data & ~OBJEXTS_FLAGS_MASK);
> }
>
> If the freed page has not been reallocated yet, folio->memcg_data is 0, so
> folio_memcg(folio) returns NULL and folio_numab_mode(folio) falls back to
> NUMAB_FOLLOW_SYSTEM rather than using the owner cgroup's tiering mode for
> PGPROMOTE_SUCCESS accounting.
>
> If another CPU reallocates the freed page before this check runs (for example
> as a slab page setting PG_slab and MEMCG_DATA_OBJEXTS, or as a compound tail
> page), can folio_objcg(folio) trigger VM_BUG_ON_FOLIO() or dereference a
> slabobj_ext pointer as a struct obj_cgroup *?
>
> Would it make sense to evaluate folio_numab_in_tiering_mode(folio) before
> migrate_pages(), or derive the mode from the already-pinned memcg pointer?

Nice catch. The check has to run before migrate_pages().

diff --git a/mm/migrate.c b/mm/migrate.c
index 4b87d465c5fd..84f224c42ab2 100644
--- a/mm/migrate.c
+++ b/mm/migrate.c
@@ -2761,6 +2761,7 @@ int migrate_misplaced_folio(struct folio *folio, int node)
        LIST_HEAD(migratepages);
        struct mem_cgroup *memcg = get_mem_cgroup_from_folio(folio);
        struct lruvec *lruvec = mem_cgroup_lruvec(memcg, pgdat);
+       bool folio_tiering = folio_numab_in_tiering_mode(folio);

        list_add(&folio->lru, &migratepages);
        nr_remaining = migrate_pages(&migratepages, alloc_misplaced_dst_folio,
@@ -2771,7 +2772,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 (folio_numab_in_tiering_mode(folio)
+               if (folio_tiering
                    && !node_is_toptier(folio_nid(folio))
                    && node_is_toptier(node))
                        mod_lruvec_state(lruvec, PGPROMOTE_SUCCESS,
nr_succeeded);


BTW, the folio_nid(folio) next to it has the same problem: after a
successful migration the source folio is uncharged and its isolation
reference dropped, so it may already be back in the allocator, and
a reallocated page's flags give a wrong nid. That one predates this
series (it goes back to e39bb6be9f2b), so I will fix it separately:

diff --git a/mm/migrate.c b/mm/migrate.c
index 15b45832bcfa..685c06a472ea 100644
--- a/mm/migrate.c
+++ b/mm/migrate.c
@@ -2760,6 +2760,7 @@ int migrate_misplaced_folio(struct folio *folio, int node)
        LIST_HEAD(migratepages);
        struct mem_cgroup *memcg = get_mem_cgroup_from_folio(folio);
        struct lruvec *lruvec = mem_cgroup_lruvec(memcg, pgdat);
+       int src_nid = folio_nid(folio);

        list_add(&folio->lru, &migratepages);
        nr_remaining = migrate_pages(&migratepages, alloc_misplaced_dst_folio,
@@ -2771,7 +2772,7 @@ int migrate_misplaced_folio(struct folio *folio, int node)
                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)
-                   && !node_is_toptier(folio_nid(folio))
+                   && !node_is_toptier(src_node))
                    && node_is_toptier(node))
                        mod_lruvec_state(lruvec, PGPROMOTE_SUCCESS,
nr_succeeded);
        }

Pls correct me if I'm wrong.

-- 
Regards
Yafang