[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

[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

[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 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