[PATCHSET v3 sched_ext/for-7.3-fixes] sched_ext: Pass the initial cmask to cid-form ops.enable()

[PATCHSET v3 sched_ext/for-7.3-fixes] sched_ext: Pass the initial cmask to cid-form ops.enable()

[Tejun Heo]

Hello,

v3: - The initial set_cmask() is delivered before set_weight() so that the
      mask is in place when weight-dependent state is derived (Andrea Righi).
    - Added a selftest checking the cmask enable() receives, the set_cmask()
      that follows it and their ordering (Andrea Righi).

v2: https://lore.kernel.org/r/2f0eb5c762d55edba5365f1006dbb395@kernel.org
v1: https://lore.kernel.org/r/c8343c2e457508a6cd3fcab66b0823b6@kernel.org

A cid-form scheduler has no p->cpus_ptr equivalent to read. A task's cid
mask is only visible through ops.set_cmask(), which doesn't fire when a task
enters a scheduler through fork, sub-sched enable or re-home, so schedulers
seed the mask from p->cpus_ptr in ops.init_task(), and that seed goes stale
when affinity changes between init_task() and enable().

The first patch passes the initial cmask to cid-form ops.enable() through a
new args struct and has set_cmask() repeat it right after, so a scheduler can
track affinity in set_cmask() alone. The second adds a selftest for it.

Verified by building the kernel, tools/sched_ext and the selftests. The new
selftest passes on a 16-CPU VM, standalone three times and inside the full
sched_ext suite. scx_qmap ran under fork churn with a throwaway enable()
check comparing the delivered mask against p->cpus_ptr on about 55k enables
with no mismatch.

Based on sched_ext/for-7.3-fixes (9ec7ba20c97d).

This patchset contains the following 2 patches.

 0001 sched_ext: Pass the initial cmask to cid-form ops.enable()
 0002 selftests/sched_ext: Check the cmask cid-form ops.enable() receives

The patchset is also available in the following git branch:

 git://git.kernel.org/pub/scm/linux/kernel/git/tj/sched_ext.git cid-enable-args-v3

diffstat follows. Thanks.

 kernel/sched/ext/ext.c                               |  92 ++++++++++++++-------
 kernel/sched/ext/internal.h                          |  56 +++++++++++--
 tools/sched_ext/scx_qmap.bpf.c                       |   3 -
 tools/testing/selftests/sched_ext/Makefile           |   1 +
 tools/testing/selftests/sched_ext/enable_cmask.bpf.c | 213 +++++++++++++++++++++++++++++++++++++++++++++++++
 tools/testing/selftests/sched_ext/enable_cmask.c     | 138 ++++++++++++++++++++++++++++++++
 6 files changed, 462 insertions(+), 41 deletions(-)

--
tejun

[PATCH 1/2] sched_ext: Pass the initial cmask to cid-form ops.enable()

[Tejun Heo]

The cid-form API has an obvious hole. A task's cid mask is only visible
through ops.set_cmask(), which fires on affinity changes and class switches
but not when a task enters a scheduler through fork, sub-sched enable or
re-home, and there is no p->cpus_ptr equivalent to fall back on. Schedulers
work around it by seeding the mask in ops.init_task() from p->cpus_ptr cid
by cid, which is subtly wrong: on sub-sched enable and re-home, an affinity
change between init_task() and enable() is delivered to the sched the task
is still on, and nothing corrects the new sched's copy afterwards.

Fix it by adding struct scx_enable_args to cid-form ops.enable() carrying
the task's cmask, built in the per-cpu scratch under the rq lock as the task
enters the scheduler, and calling set_cmask() with the same mask right after
enable(), ahead of set_weight(). A scheduler can then track affinity in
set_cmask() alone, and scx_qmap drops its init_task() seed. set_cmask() no
longer fires for a cid-form task before it is enabled, and the class-switch
republish in switching_to_scx() is limited to the cpu form.

This changes the cid-form ops.enable() signature, which is fine as the
cid-form API is still considered unreleased. An args struct rather than a
bare cmask argument leaves room for more initial state without another
signature change, and the cmask travels as a plain arena address because BTF
can't type arena struct members yet.

v2: The cmask travels as a u64 arena address, cmask_arena_addr, instead of a
kernel-typed pointer, with the typing limitation and the planned typed alias
documented (Sashiko review).

v3: The initial set_cmask() is delivered before set_weight() so that the
mask is in place when weight-dependent state is derived (Andrea Righi).
Selftest added.

Signed-off-by: Tejun Heo <tj@kernel.org>
---
 kernel/sched/ext/ext.c         | 92 +++++++++++++++++++++++-----------
 kernel/sched/ext/internal.h    | 56 ++++++++++++++++++---
 tools/sched_ext/scx_qmap.bpf.c |  3 --
 3 files changed, 110 insertions(+), 41 deletions(-)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index f568fd9973f6..ad391a8cbd05 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -447,37 +447,45 @@ static void switch_rq_lock(struct rq *from, struct rq *to)
 DEFINE_STATIC_KEY_FALSE(__scx_is_cid_type);
 
 /**
- * scx_call_op_set_cpumask - invoke ops.set_cpumask / ops_cid.set_cmask for @task
+ * scx_fill_cmask_scratch - Build this cpu's arena cmask from @cpumask
+ * @sch: scx_sched whose scratch to fill
+ * @cpumask: cpus to translate into cids
+ *
+ * The scratch lives in BPF-writable arena memory and its header can't be
+ * trusted, so it is rewritten from kernel geometry rather than read. Caller
+ * must hold an rq lock so this cpu is the sole kernel writer for as long as the
+ * returned address is in use.
+ */
+static struct scx_cmask *scx_fill_cmask_scratch(struct scx_sched *sch,
+						const struct cpumask *cpumask)
+{
+	struct scx_cmask *kern_va = *this_cpu_ptr(sch->set_cmask_scratch);
+	struct scx_cmask_ref ref;
+
+	scx_cmask_ref_init_kern(sch, kern_va, 0, num_possible_cpus(), &ref);
+	scx_cmask_ref_from_cpumask(&ref, cpumask);
+	return kern_va;
+}
+
+/**
+ * scx_call_op_set_cpumask - Invoke the set_cpumask or set_cmask op for @task
  * @sch: scx_sched being invoked
  * @rq: rq to update as the currently-locked rq, or NULL
  * @task: task whose affinity is changing
  * @cpumask: new cpumask
  *
- * For cid-form schedulers, translate @cpumask to a cmask via the per-cpu
- * scratch in cid.c and dispatch through the ops_cid union view. Caller
- * must hold @rq's rq lock so this_cpu_ptr is stable across the call.
+ * For cid-form schedulers, translate @cpumask to a cmask in the per-cpu scratch
+ * and dispatch through the ops_cid union view. Caller must hold @rq's rq lock.
  */
 static inline void scx_call_op_set_cpumask(struct scx_sched *sch, struct rq *rq,
 					   struct task_struct *task,
 					   const struct cpumask *cpumask)
 {
-	if (scx_is_cid_type()) {
-		struct scx_cmask *kern_va = *this_cpu_ptr(sch->set_cmask_scratch);
-		struct scx_cmask_ref ref;
-
-		/*
-		 * Build the per-cpu arena cmask from kernel geometry via @ref,
-		 * never reading its BPF-writable header. set_cmask()'s __arena
-		 * argument takes the kernel address and the struct_ops
-		 * trampoline rebases it into BPF's arena pointer form. The rq
-		 * lock makes this cpu the sole kernel writer.
-		 */
-		scx_cmask_ref_init_kern(sch, kern_va, 0, num_possible_cpus(), &ref);
-		scx_cmask_ref_from_cpumask(&ref, cpumask);
-		SCX_CALL_CID_OP_TASK(sch, set_cmask, rq, task, kern_va);
-	} else {
+	if (scx_is_cid_type())
+		SCX_CALL_CID_OP_TASK(sch, set_cmask, rq, task,
+				     scx_fill_cmask_scratch(sch, cpumask));
+	else
 		SCX_CALL_OP_TASK(sch, set_cpumask, rq, task, cpumask);
-	}
 }
 
 enum scx_dsq_iter_flags {
@@ -3634,8 +3642,12 @@ static void set_cpus_allowed_scx(struct task_struct *p,
 	 *
 	 * Fine-grained memory write control is enforced by BPF making the const
 	 * designation pointless. Cast it away when calling the operation.
+	 *
+	 * The cid form receives the initial mask when the task is enabled and
+	 * hears about changes only afterwards, see struct scx_enable_args.
 	 */
-	if (SCX_HAS_OP(sch, set_cpumask))
+	if (SCX_HAS_OP(sch, set_cpumask) &&
+	    (!scx_is_cid_type() || scx_get_task_state(p) == SCX_TASK_ENABLED))
 		scx_call_op_set_cpumask(sch, task_rq(p), p, (struct cpumask *)p->cpus_ptr);
 }
 
@@ -3944,8 +3956,27 @@ static void __scx_enable_task(struct scx_sched *sch, struct task_struct *p)
 
 	p->scx.weight = sched_weight_to_cgroup(weight);
 
-	if (SCX_HAS_OP(sch, enable))
-		SCX_CALL_OP_TASK(sch, enable, rq, p);
+	if (SCX_HAS_OP(sch, enable)) {
+		if (scx_is_cid_type()) {
+			struct scx_cmask *cmask = scx_fill_cmask_scratch(sch, p->cpus_ptr);
+			struct scx_enable_args args = {
+				.cmask_arena_addr = scx_kaddr_to_arena(sch, cmask),
+			};
+
+			SCX_CALL_CID_OP_TASK(sch, enable, rq, p, &args);
+		} else {
+			SCX_CALL_OP_TASK(sch, enable, rq, p);
+		}
+	}
+
+	/*
+	 * The initial mask also goes out through set_cmask() so a scheduler can
+	 * track affinity there alone, and before set_weight() so that the mask
+	 * is in place when weight-dependent state is derived, see struct
+	 * scx_enable_args.
+	 */
+	if (scx_is_cid_type() && SCX_HAS_OP(sch, set_cpumask))
+		scx_call_op_set_cpumask(sch, rq, p, p->cpus_ptr);
 
 	if (SCX_HAS_OP(sch, set_weight))
 		SCX_CALL_OP_TASK(sch, set_weight, rq, p, p->scx.weight);
@@ -4288,9 +4319,10 @@ static void switching_to_scx(struct rq *rq, struct task_struct *p)
 
 	/*
 	 * set_cpus_allowed_scx() is not called while @p is associated with a
-	 * different scheduler class. Keep the BPF scheduler up-to-date.
+	 * different scheduler class. Keep the BPF scheduler up-to-date. The cid
+	 * form gets its mask from scx_enable_task().
 	 */
-	if (SCX_HAS_OP(sch, set_cpumask))
+	if (!scx_is_cid_type() && SCX_HAS_OP(sch, set_cpumask))
 		scx_call_op_set_cpumask(sch, rq, p, (struct cpumask *)p->cpus_ptr);
 }
 
@@ -8374,10 +8406,11 @@ static struct bpf_struct_ops bpf_sched_ext_ops = {
 /*
  * cid-form cfi stubs. Stubs whose signatures match the cpu-form (param types
  * identical, only param names differ across structs) are reused. Some need
- * fresh stubs, set_cmask due to an argument type difference and the sub-sched
- * notifiers because no cpu-form stub exists to reuse.
+ * fresh stubs, set_cmask and enable due to argument differences and the
+ * sub-sched notifiers because no cpu-form stub exists to reuse.
  */
 static void sched_ext_ops_cid__set_cmask(struct task_struct *p, const struct scx_cmask *cmask__arena) {}
+static void sched_ext_ops_cid__enable(struct task_struct *p, struct scx_enable_args *args) {}
 static void sched_ext_ops__sub_caps_updated(const struct scx_cmask *cmask__arena, u64 caps) {}
 static void sched_ext_ops__sub_ecaps_updated(s32 cid, u64 before, u64 after) {}
 
@@ -8398,7 +8431,7 @@ static struct sched_ext_ops_cid __bpf_ops_sched_ext_ops_cid = {
 	.update_idle		= sched_ext_ops__update_idle,
 	.init_task		= sched_ext_ops__init_task,
 	.exit_task		= sched_ext_ops__exit_task,
-	.enable			= sched_ext_ops__enable,
+	.enable			= sched_ext_ops_cid__enable,
 	.disable		= sched_ext_ops__disable,
 #ifdef CONFIG_EXT_GROUP_SCHED
 	.cpuctl_init		= sched_ext_ops__cgroup_init,
@@ -10421,8 +10454,7 @@ __bpf_kfunc const void *scx_bpf_online_cmask(const struct bpf_prog_aux *aux)
 	if (unlikely(!online))
 		return NULL;
 
-	/* BPF rebases by the low 32 bits, like __arena callback args */
-	return (void *)((unsigned long)online - sch->arena_kern_base);
+	return (void *)scx_kaddr_to_arena(sch, online);
 }
 
 /**
diff --git a/kernel/sched/ext/internal.h b/kernel/sched/ext/internal.h
index d150de10a5c9..1df8f583b0ec 100644
--- a/kernel/sched/ext/internal.h
+++ b/kernel/sched/ext/internal.h
@@ -250,6 +250,31 @@ struct scx_exit_task_args {
 	bool cancelled;
 };
 
+/**
+ * struct scx_enable_args - Argument container for cid-form ops.enable()
+ * @cmask_arena_addr: BPF arena address of the cmask of cids the task may run on
+ *
+ * @cmask_arena_addr is the task's affinity as it enters the scheduler.
+ * set_cmask() delivers the same mask right after enable(), before set_weight()
+ * and the first enqueue, then every affinity change afterwards, and is never
+ * called before enable(). A scheduler may therefore track affinity in
+ * set_cmask() alone.
+ *
+ * The kernel builds the mask in the scheduler arena from its own geometry, so
+ * the header is valid regardless of what the scheduler last wrote there. The
+ * memory is per-cpu scratch reused once the callback returns: copy the bits
+ * out, don't keep the address. The set_cmask() argument follows the same rules.
+ *
+ * The address is a plain value rather than a typed pointer because BTF can't
+ * mark a struct member as an arena pointer yet and a pointer member would reach
+ * the program typed as a kernel pointer. Cast it to struct scx_cmask __arena *
+ * before use. Once arena members can be typed, a typed alias will join this
+ * field in an anonymous union at the same offset.
+ */
+struct scx_enable_args {
+	u64	cmask_arena_addr;
+};
+
 /* argument container for ops.cgroup_init() */
 struct scx_cgroup_init_args {
 	/* the weight of the cgroup [1..10000] */
@@ -1037,6 +1062,7 @@ struct sched_ext_ops {
  *   - dispatch         -> dispatch (cpu arg is now cid)
  *   - update_idle      -> update_idle (cpu arg is now cid)
  *   - set_cpumask      -> set_cmask (cmask instead of cpumask)
+ *   - enable           -> enable (takes struct scx_enable_args)
  *   - cpu_online       -> cid_online
  *   - cpu_offline      -> cid_offline
  *   - dump_cpu         -> dump_cid
@@ -1070,7 +1096,7 @@ struct sched_ext_ops_cid {
 			  struct scx_init_task_args *args);
 	void (*exit_task)(struct task_struct *p,
 			   struct scx_exit_task_args *args);
-	void (*enable)(struct task_struct *p);
+	void (*enable)(struct task_struct *p, struct scx_enable_args *args);
 	void (*disable)(struct task_struct *p);
 	void (*dump)(struct scx_dump_ctx *ctx);
 	void (*dump_cid)(struct scx_dump_ctx *ctx, s32 cid, bool idle);
@@ -1533,7 +1559,8 @@ struct scx_sched {
 	 * by BUILD_BUG_ON in scx_init()). The anonymous union lets the kernel
 	 * access either view of the same storage without function-pointer
 	 * casts: use .ops for cpu-form and shared fields, .ops_cid for the
-	 * cid-renamed callbacks (set_cmask, select_cid, cid_online, ...).
+	 * callbacks whose cid-form signature differs (set_cmask, enable,
+	 * select_cid, cid_online, ...).
 	 */
 	union {
 		struct sched_ext_ops		ops;
@@ -1556,9 +1583,9 @@ struct scx_sched {
 	uintptr_t		arena_kern_base;
 
 	/*
-	 * Per-CPU arena cmask used by scx_call_op_set_cpumask() to hand a cmask
-	 * to ops_cid.set_cmask(). The kernel writes through the stored kern_va
-	 * and passes it to the callback's __arena argument.
+	 * Per-CPU arena cmask the kernel fills from a task's cpumask and hands
+	 * to ops_cid.enable() and ops_cid.set_cmask(). The stored pointers are
+	 * the kernel addresses.
 	 */
 	struct scx_cmask * __percpu *set_cmask_scratch;
 	struct scx_cmask *online_cmask;
@@ -1669,6 +1696,19 @@ static inline void *scx_arena_to_kaddr(struct scx_sched *sch, const void *bpf_pt
 	return (void *)(sch->arena_kern_base + (u32)(uintptr_t)bpf_ptr);
 }
 
+/**
+ * scx_kaddr_to_arena - Translate a kernel arena address to the BPF form
+ * @sch: scheduler whose arena hosts @kaddr
+ * @kaddr: kernel address inside @sch's arena
+ *
+ * __arena callback arguments need no translation. Addresses handed to BPF any
+ * other way, such as struct fields and kfunc return values, go through this.
+ */
+static inline uintptr_t scx_kaddr_to_arena(struct scx_sched *sch, const void *kaddr)
+{
+	return (uintptr_t)kaddr - sch->arena_kern_base;
+}
+
 enum scx_wake_flags {
 	/* expose select WF_* flags as enums */
 	SCX_WAKE_FORK		= WF_FORK,
@@ -2302,9 +2342,9 @@ do {										\
 } while (0)
 
 /*
- * Dispatch a task op through the cid-form ops_cid table. Only set_cmask() needs
- * this: it takes an arena cmask address instead of a cpumask, so it cannot be
- * invoked via its cpu-form set_cpumask() slot.
+ * Dispatch a task op through the cid-form ops_cid table, for the ops whose
+ * cid-form signature differs from the cpu-form slot: set_cmask() takes an arena
+ * cmask instead of a cpumask and enable() takes scx_enable_args.
  */
 #define SCX_CALL_CID_OP_TASK(sch, op, locked_rq, task, args...)			\
 	__SCX_CALL_OP_TASK(sch, ops_cid, op, locked_rq, task, ##args)
diff --git a/tools/sched_ext/scx_qmap.bpf.c b/tools/sched_ext/scx_qmap.bpf.c
index 67b7c01cae55..2f3653199842 100644
--- a/tools/sched_ext/scx_qmap.bpf.c
+++ b/tools/sched_ext/scx_qmap.bpf.c
@@ -961,9 +961,6 @@ s32 BPF_STRUCT_OPS_SLEEPABLE(qmap_init_task, struct task_struct *p,
 	taskc->highpri = false;
 	taskc->core_sched_seq = 0;
 	cmask_init(&taskc->cpus_allowed, 0, scx_bpf_nr_cids());
-	bpf_rcu_read_lock();
-	cmask_from_cpumask(&taskc->cpus_allowed, p->cpus_ptr);
-	bpf_rcu_read_unlock();
 
 	v = bpf_task_storage_get(&task_ctx_stor, p, NULL,
 				 BPF_LOCAL_STORAGE_GET_F_CREATE);
-- 
2.55.0

[PATCH 2/2] selftests/sched_ext: Check the cmask cid-form ops.enable() receives

[Tejun Heo]

cid-form ops.enable() now hands the task's cmask to the scheduler and
set_cmask() repeats it right after. Add a cid-form selftest that checks both
against p->cpus_ptr, that they match each other, that the initial
set_cmask() lands before set_weight() and before the task first becomes
runnable, and that set_cmask() never precedes enable(), across class-switch
enables, fork-path enables and live affinity changes.

Signed-off-by: Tejun Heo <tj@kernel.org>
---
 tools/testing/selftests/sched_ext/Makefile    |   1 +
 .../selftests/sched_ext/enable_cmask.bpf.c    | 213 ++++++++++++++++++
 .../selftests/sched_ext/enable_cmask.c        | 138 ++++++++++++
 3 files changed, 352 insertions(+)
 create mode 100644 tools/testing/selftests/sched_ext/enable_cmask.bpf.c
 create mode 100644 tools/testing/selftests/sched_ext/enable_cmask.c

diff --git a/tools/testing/selftests/sched_ext/Makefile b/tools/testing/selftests/sched_ext/Makefile
index 49897727f535..4e06d0baaeec 100644
--- a/tools/testing/selftests/sched_ext/Makefile
+++ b/tools/testing/selftests/sched_ext/Makefile
@@ -169,6 +169,7 @@ auto-test-targets :=			\
 	ddsp_bogus_dsq_fail		\
 	ddsp_vtimelocal_fail		\
 	dsp_local_on			\
+	enable_cmask			\
 	enq_select_cpu			\
 	exit				\
 	hotplug				\
diff --git a/tools/testing/selftests/sched_ext/enable_cmask.bpf.c b/tools/testing/selftests/sched_ext/enable_cmask.bpf.c
new file mode 100644
index 000000000000..bbbf989d2c69
--- /dev/null
+++ b/tools/testing/selftests/sched_ext/enable_cmask.bpf.c
@@ -0,0 +1,213 @@
+// SPDX-License-Identifier: GPL-2.0
+/*
+ * A cid-form scheduler checking the cmask cid-form ops.enable() receives: the
+ * header, every cid bit against p->cpus_ptr, and that set_cmask() follows with
+ * the same mask before set_weight() and before the task first becomes runnable,
+ * and never runs before enable().
+ *
+ * Copyright (c) 2026 Tejun Heo <tj@kernel.org>
+ */
+#include <scx/common.bpf.h>
+
+#define MAX_CPUS 1024
+
+char _license[] SEC("license") = "GPL";
+
+struct {
+	__uint(type, BPF_MAP_TYPE_ARENA);
+	__uint(map_flags, BPF_F_MMAPABLE);
+	__uint(max_entries, 1 << 16);
+} arena SEC(".maps");
+
+struct task_ctx {
+	u64	enable_fp;	/* fingerprint of the mask enable() received */
+	bool	enabled;
+	bool	pending;	/* enable() ran, the initial set_cmask() hasn't */
+};
+
+struct {
+	__uint(type, BPF_MAP_TYPE_TASK_STORAGE);
+	__uint(map_flags, BPF_F_NO_PREALLOC);
+	__type(key, int);
+	__type(value, struct task_ctx);
+} task_ctx_stor SEC(".maps");
+
+u64 nr_enable, nr_initial_set_cmask, nr_set_cmask, nr_set_weight;
+s32 bad_cid;
+bool bad_want, bad_got;
+
+UEI_DEFINE(uei);
+
+static struct task_ctx *lookup_task_ctx(struct task_struct *p)
+{
+	struct task_ctx *tctx;
+
+	tctx = bpf_task_storage_get(&task_ctx_stor, p, 0, 0);
+	if (!tctx)
+		scx_bpf_error("task_ctx lookup failed for %s[%d]", p->comm, p->pid);
+	return tctx;
+}
+
+/*
+ * Verify @m's header and every cid bit against @p's cpumask and fingerprint the
+ * bits into @fp. Return 0 on success, -EINVAL on a bad header, -ENOENT on a cid
+ * without a cpu and -EIO on a bit mismatch with the details in @bad_*.
+ */
+static int check_mask(struct task_struct *p, const struct scx_cmask __arena *m, u64 *fp)
+{
+	u32 nr_cids = scx_bpf_nr_cids();
+	u64 h = 0;
+	s32 cid;
+
+	if (m->base || m->nr_cids != nr_cids)
+		return -EINVAL;
+
+	bpf_for(cid, 0, MAX_CPUS) {
+		bool want, got;
+		s32 cpu;
+
+		if (cid >= nr_cids)
+			break;
+		cpu = scx_bpf_cid_to_cpu(cid);
+		if (cpu < 0)
+			return -ENOENT;
+		want = bpf_cpumask_test_cpu(cpu, p->cpus_ptr);
+		got = cmask_test(cid, m);
+		if (want != got) {
+			bad_cid = cid;
+			bad_want = want;
+			bad_got = got;
+			return -EIO;
+		}
+		h = h * 31 + got;
+	}
+
+	*fp = h;
+	return 0;
+}
+
+s32 BPF_STRUCT_OPS_SLEEPABLE(enable_cmask_init_task, struct task_struct *p,
+			     struct scx_init_task_args *args)
+{
+	if (!bpf_task_storage_get(&task_ctx_stor, p, 0, BPF_LOCAL_STORAGE_GET_F_CREATE))
+		return -ENOMEM;
+	return 0;
+}
+
+void BPF_STRUCT_OPS(enable_cmask_enable, struct task_struct *p, struct scx_enable_args *args)
+{
+	struct scx_cmask __arena *m = (struct scx_cmask __arena *)args->cmask_arena_addr;
+	struct task_ctx *tctx;
+	int ret;
+
+	asm volatile("" :: "r"(&arena));
+	tctx = lookup_task_ctx(p);
+	if (!tctx)
+		return;
+
+	__sync_fetch_and_add(&nr_enable, 1);
+	if (tctx->enabled || tctx->pending) {
+		scx_bpf_error("enable: %s[%d] enabled twice", p->comm, p->pid);
+		return;
+	}
+
+	ret = check_mask(p, m, &tctx->enable_fp);
+	if (ret) {
+		scx_bpf_error("enable: %s[%d] cmask check failed %d cid=%d want=%d got=%d",
+			      p->comm, p->pid, ret, bad_cid, bad_want, bad_got);
+		return;
+	}
+	tctx->enabled = true;
+	tctx->pending = true;
+}
+
+void BPF_STRUCT_OPS(enable_cmask_set_cmask, struct task_struct *p,
+		    struct scx_cmask __arena *m)
+{
+	struct task_ctx *tctx;
+	u64 fp;
+	int ret;
+
+	asm volatile("" :: "r"(&arena));
+	tctx = lookup_task_ctx(p);
+	if (!tctx)
+		return;
+
+	__sync_fetch_and_add(&nr_set_cmask, 1);
+	if (!tctx->enabled) {
+		scx_bpf_error("set_cmask: %s[%d] not enabled", p->comm, p->pid);
+		return;
+	}
+
+	ret = check_mask(p, m, &fp);
+	if (ret) {
+		scx_bpf_error("set_cmask: %s[%d] cmask check failed %d cid=%d want=%d got=%d",
+			      p->comm, p->pid, ret, bad_cid, bad_want, bad_got);
+		return;
+	}
+
+	if (tctx->pending) {
+		if (fp != tctx->enable_fp) {
+			scx_bpf_error("set_cmask: %s[%d] initial mask differs from enable()",
+				      p->comm, p->pid);
+			return;
+		}
+		tctx->pending = false;
+		__sync_fetch_and_add(&nr_initial_set_cmask, 1);
+	}
+}
+
+void BPF_STRUCT_OPS(enable_cmask_set_weight, struct task_struct *p, u32 weight)
+{
+	struct task_ctx *tctx;
+
+	tctx = lookup_task_ctx(p);
+	if (!tctx)
+		return;
+
+	__sync_fetch_and_add(&nr_set_weight, 1);
+	if (tctx->pending)
+		scx_bpf_error("set_weight: %s[%d] before the initial set_cmask()", p->comm,
+			      p->pid);
+}
+
+void BPF_STRUCT_OPS(enable_cmask_runnable, struct task_struct *p, u64 enq_flags)
+{
+	struct task_ctx *tctx;
+
+	tctx = lookup_task_ctx(p);
+	if (!tctx)
+		return;
+
+	if (tctx->pending)
+		scx_bpf_error("runnable: %s[%d] before the initial set_cmask()", p->comm,
+			      p->pid);
+}
+
+void BPF_STRUCT_OPS(enable_cmask_disable, struct task_struct *p)
+{
+	struct task_ctx *tctx;
+
+	tctx = lookup_task_ctx(p);
+	if (!tctx)
+		return;
+
+	tctx->enabled = false;
+	tctx->pending = false;
+}
+
+void BPF_STRUCT_OPS(enable_cmask_exit, struct scx_exit_info *ei)
+{
+	UEI_RECORD(uei, ei);
+}
+
+SCX_OPS_CID_DEFINE(enable_cmask_ops,
+		   .init_task	= (void *)enable_cmask_init_task,
+		   .enable	= (void *)enable_cmask_enable,
+		   .set_cmask	= (void *)enable_cmask_set_cmask,
+		   .set_weight	= (void *)enable_cmask_set_weight,
+		   .runnable	= (void *)enable_cmask_runnable,
+		   .disable	= (void *)enable_cmask_disable,
+		   .exit	= (void *)enable_cmask_exit,
+		   .flags	= SCX_OPS_SWITCH_PARTIAL,
+		   .name	= "enable_cmask");
diff --git a/tools/testing/selftests/sched_ext/enable_cmask.c b/tools/testing/selftests/sched_ext/enable_cmask.c
new file mode 100644
index 000000000000..556bcad4431d
--- /dev/null
+++ b/tools/testing/selftests/sched_ext/enable_cmask.c
@@ -0,0 +1,138 @@
+// SPDX-License-Identifier: GPL-2.0
+/* Copyright (c) 2026 Tejun Heo <tj@kernel.org> */
+#define _GNU_SOURCE
+#include <sched.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <time.h>
+#include <unistd.h>
+#include <sys/wait.h>
+#include <bpf/bpf.h>
+#include <scx/common.h>
+#include "enable_cmask.bpf.skel.h"
+#include "scx_test.h"
+
+#define SCHED_EXT 7
+#define NR_CHILDREN 8
+#define MAX_CPUS 1024
+
+static int cpus[MAX_CPUS];
+static int nr_cpus;
+
+static void spin_ms(int ms)
+{
+	struct timespec start, now;
+
+	clock_gettime(CLOCK_MONOTONIC, &start);
+	do {
+		clock_gettime(CLOCK_MONOTONIC, &now);
+	} while ((now.tv_sec - start.tv_sec) * 1000 +
+		 (now.tv_nsec - start.tv_nsec) / 1000000 < ms);
+}
+
+static int pin(pid_t pid, int idx)
+{
+	cpu_set_t set;
+
+	CPU_ZERO(&set);
+	CPU_SET(cpus[idx % nr_cpus], &set);
+	return sched_setaffinity(pid, sizeof(set), &set);
+}
+
+/*
+ * Pin, switch to SCHED_EXT for a class-switch enable, fork a grandchild that
+ * inherits the policy for a fork-path enable, then change affinity a few times
+ * while running for set_cmask() on live tasks.
+ */
+static int child(int idx)
+{
+	struct sched_param param = {};
+	int i, status;
+	pid_t pid;
+
+	if (pin(0, idx) || sched_setscheduler(0, SCHED_EXT, &param))
+		return 1;
+
+	pid = fork();
+	if (pid < 0)
+		return 1;
+	if (!pid) {
+		spin_ms(20);
+		return 0;
+	}
+
+	for (i = 1; i <= 4; i++) {
+		if (pin(0, idx + i))
+			return 1;
+		spin_ms(5);
+	}
+
+	return waitpid(pid, &status, 0) == pid && !status ? 0 : 1;
+}
+
+static enum scx_test_status run(void *ctx)
+{
+	struct enable_cmask *skel;
+	struct bpf_link *link;
+	pid_t pids[NR_CHILDREN];
+	cpu_set_t set;
+	int i, status, failed = 0;
+
+	if (!__COMPAT_struct_has_field("scx_enable_args", "cmask_arena_addr"))
+		return SCX_TEST_SKIP;
+
+	SCX_FAIL_IF(sched_getaffinity(0, sizeof(set), &set), "Failed to read affinity");
+	for (i = 0; i < MAX_CPUS && i < CPU_SETSIZE; i++)
+		if (CPU_ISSET(i, &set))
+			cpus[nr_cpus++] = i;
+	if (nr_cpus < 2)
+		return SCX_TEST_SKIP;
+
+	skel = enable_cmask__open();
+	SCX_FAIL_IF(!skel, "Failed to open");
+	SCX_ENUM_INIT(skel);
+	SCX_FAIL_IF(enable_cmask__load(skel), "Failed to load skel");
+
+	link = bpf_map__attach_struct_ops(skel->maps.enable_cmask_ops);
+	SCX_FAIL_IF(!link, "Failed to attach struct_ops");
+
+	for (i = 0; i < NR_CHILDREN; i++) {
+		pids[i] = fork();
+		SCX_FAIL_IF(pids[i] < 0, "Failed to fork");
+		if (!pids[i])
+			exit(child(i));
+	}
+
+	/* affinity changes from the outside race with the children's own */
+	for (i = 0; i < NR_CHILDREN; i++)
+		pin(pids[i], i + NR_CHILDREN);
+
+	for (i = 0; i < NR_CHILDREN; i++) {
+		if (waitpid(pids[i], &status, 0) != pids[i] || status)
+			failed++;
+	}
+
+	bpf_link__destroy(link);
+
+	SCX_EQ(skel->data->uei.kind, EXIT_KIND(SCX_EXIT_UNREG));
+	SCX_EQ(failed, 0);
+	SCX_GE(skel->bss->nr_enable, 2 * NR_CHILDREN);
+	SCX_EQ(skel->bss->nr_initial_set_cmask, skel->bss->nr_enable);
+	SCX_GT(skel->bss->nr_set_cmask, skel->bss->nr_initial_set_cmask);
+	SCX_GE(skel->bss->nr_set_weight, skel->bss->nr_enable);
+	printf("enable=%lu initial_set_cmask=%lu set_cmask=%lu set_weight=%lu\n",
+	       (unsigned long)skel->bss->nr_enable,
+	       (unsigned long)skel->bss->nr_initial_set_cmask,
+	       (unsigned long)skel->bss->nr_set_cmask,
+	       (unsigned long)skel->bss->nr_set_weight);
+
+	enable_cmask__destroy(skel);
+	return SCX_TEST_PASS;
+}
+
+struct scx_test enable_cmask = {
+	.name = "enable_cmask",
+	.description = "Check the cid-form ops.enable() cmask and the set_cmask() after it",
+	.run = run,
+};
+REGISTER_SCX_TEST(&enable_cmask)
-- 
2.55.0

Re: [PATCH 1/2] sched_ext: Pass the initial cmask to cid-form ops.enable()

[Andrea Righi]

Hi Tejun,

On Fri, Sep 18, 2026 at 02:28:37PM -1000, Tejun Heo wrote:
> The cid-form API has an obvious hole. A task's cid mask is only visible
> through ops.set_cmask(), which fires on affinity changes and class switches
> but not when a task enters a scheduler through fork, sub-sched enable or
> re-home, and there is no p->cpus_ptr equivalent to fall back on. Schedulers
> work around it by seeding the mask in ops.init_task() from p->cpus_ptr cid
> by cid, which is subtly wrong: on sub-sched enable and re-home, an affinity
> change between init_task() and enable() is delivered to the sched the task
> is still on, and nothing corrects the new sched's copy afterwards.
> 
> Fix it by adding struct scx_enable_args to cid-form ops.enable() carrying
> the task's cmask, built in the per-cpu scratch under the rq lock as the task
> enters the scheduler, and calling set_cmask() with the same mask right after
> enable(), ahead of set_weight(). A scheduler can then track affinity in
> set_cmask() alone, and scx_qmap drops its init_task() seed. set_cmask() no
> longer fires for a cid-form task before it is enabled, and the class-switch
> republish in switching_to_scx() is limited to the cpu form.
> 
> This changes the cid-form ops.enable() signature, which is fine as the
> cid-form API is still considered unreleased. An args struct rather than a
> bare cmask argument leaves room for more initial state without another
> signature change, and the cmask travels as a plain arena address because BTF
> can't type arena struct members yet.
> 
> v2: The cmask travels as a u64 arena address, cmask_arena_addr, instead of a
> kernel-typed pointer, with the typing limitation and the planned typed alias
> documented (Sashiko review).
> 
> v3: The initial set_cmask() is delivered before set_weight() so that the
> mask is in place when weight-dependent state is derived (Andrea Righi).
> Selftest added.
> 
> Signed-off-by: Tejun Heo <tj@kernel.org>

Looks good now.

Reviewed-by: Andrea Righi <arighi@nvidia.com>

Thanks,
-Andrea

> ---
>  kernel/sched/ext/ext.c         | 92 +++++++++++++++++++++++-----------
>  kernel/sched/ext/internal.h    | 56 ++++++++++++++++++---
>  tools/sched_ext/scx_qmap.bpf.c |  3 --
>  3 files changed, 110 insertions(+), 41 deletions(-)
> 
> diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
> index f568fd9973f6..ad391a8cbd05 100644
> --- a/kernel/sched/ext/ext.c
> +++ b/kernel/sched/ext/ext.c
> @@ -447,37 +447,45 @@ static void switch_rq_lock(struct rq *from, struct rq *to)
>  DEFINE_STATIC_KEY_FALSE(__scx_is_cid_type);
>  
>  /**
> - * scx_call_op_set_cpumask - invoke ops.set_cpumask / ops_cid.set_cmask for @task
> + * scx_fill_cmask_scratch - Build this cpu's arena cmask from @cpumask
> + * @sch: scx_sched whose scratch to fill
> + * @cpumask: cpus to translate into cids
> + *
> + * The scratch lives in BPF-writable arena memory and its header can't be
> + * trusted, so it is rewritten from kernel geometry rather than read. Caller
> + * must hold an rq lock so this cpu is the sole kernel writer for as long as the
> + * returned address is in use.
> + */
> +static struct scx_cmask *scx_fill_cmask_scratch(struct scx_sched *sch,
> +						const struct cpumask *cpumask)
> +{
> +	struct scx_cmask *kern_va = *this_cpu_ptr(sch->set_cmask_scratch);
> +	struct scx_cmask_ref ref;
> +
> +	scx_cmask_ref_init_kern(sch, kern_va, 0, num_possible_cpus(), &ref);
> +	scx_cmask_ref_from_cpumask(&ref, cpumask);
> +	return kern_va;
> +}
> +
> +/**
> + * scx_call_op_set_cpumask - Invoke the set_cpumask or set_cmask op for @task
>   * @sch: scx_sched being invoked
>   * @rq: rq to update as the currently-locked rq, or NULL
>   * @task: task whose affinity is changing
>   * @cpumask: new cpumask
>   *
> - * For cid-form schedulers, translate @cpumask to a cmask via the per-cpu
> - * scratch in cid.c and dispatch through the ops_cid union view. Caller
> - * must hold @rq's rq lock so this_cpu_ptr is stable across the call.
> + * For cid-form schedulers, translate @cpumask to a cmask in the per-cpu scratch
> + * and dispatch through the ops_cid union view. Caller must hold @rq's rq lock.
>   */
>  static inline void scx_call_op_set_cpumask(struct scx_sched *sch, struct rq *rq,
>  					   struct task_struct *task,
>  					   const struct cpumask *cpumask)
>  {
> -	if (scx_is_cid_type()) {
> -		struct scx_cmask *kern_va = *this_cpu_ptr(sch->set_cmask_scratch);
> -		struct scx_cmask_ref ref;
> -
> -		/*
> -		 * Build the per-cpu arena cmask from kernel geometry via @ref,
> -		 * never reading its BPF-writable header. set_cmask()'s __arena
> -		 * argument takes the kernel address and the struct_ops
> -		 * trampoline rebases it into BPF's arena pointer form. The rq
> -		 * lock makes this cpu the sole kernel writer.
> -		 */
> -		scx_cmask_ref_init_kern(sch, kern_va, 0, num_possible_cpus(), &ref);
> -		scx_cmask_ref_from_cpumask(&ref, cpumask);
> -		SCX_CALL_CID_OP_TASK(sch, set_cmask, rq, task, kern_va);
> -	} else {
> +	if (scx_is_cid_type())
> +		SCX_CALL_CID_OP_TASK(sch, set_cmask, rq, task,
> +				     scx_fill_cmask_scratch(sch, cpumask));
> +	else
>  		SCX_CALL_OP_TASK(sch, set_cpumask, rq, task, cpumask);
> -	}
>  }
>  
>  enum scx_dsq_iter_flags {
> @@ -3634,8 +3642,12 @@ static void set_cpus_allowed_scx(struct task_struct *p,
>  	 *
>  	 * Fine-grained memory write control is enforced by BPF making the const
>  	 * designation pointless. Cast it away when calling the operation.
> +	 *
> +	 * The cid form receives the initial mask when the task is enabled and
> +	 * hears about changes only afterwards, see struct scx_enable_args.
>  	 */
> -	if (SCX_HAS_OP(sch, set_cpumask))
> +	if (SCX_HAS_OP(sch, set_cpumask) &&
> +	    (!scx_is_cid_type() || scx_get_task_state(p) == SCX_TASK_ENABLED))
>  		scx_call_op_set_cpumask(sch, task_rq(p), p, (struct cpumask *)p->cpus_ptr);
>  }
>  
> @@ -3944,8 +3956,27 @@ static void __scx_enable_task(struct scx_sched *sch, struct task_struct *p)
>  
>  	p->scx.weight = sched_weight_to_cgroup(weight);
>  
> -	if (SCX_HAS_OP(sch, enable))
> -		SCX_CALL_OP_TASK(sch, enable, rq, p);
> +	if (SCX_HAS_OP(sch, enable)) {
> +		if (scx_is_cid_type()) {
> +			struct scx_cmask *cmask = scx_fill_cmask_scratch(sch, p->cpus_ptr);
> +			struct scx_enable_args args = {
> +				.cmask_arena_addr = scx_kaddr_to_arena(sch, cmask),
> +			};
> +
> +			SCX_CALL_CID_OP_TASK(sch, enable, rq, p, &args);
> +		} else {
> +			SCX_CALL_OP_TASK(sch, enable, rq, p);
> +		}
> +	}
> +
> +	/*
> +	 * The initial mask also goes out through set_cmask() so a scheduler can
> +	 * track affinity there alone, and before set_weight() so that the mask
> +	 * is in place when weight-dependent state is derived, see struct
> +	 * scx_enable_args.
> +	 */
> +	if (scx_is_cid_type() && SCX_HAS_OP(sch, set_cpumask))
> +		scx_call_op_set_cpumask(sch, rq, p, p->cpus_ptr);
>  
>  	if (SCX_HAS_OP(sch, set_weight))
>  		SCX_CALL_OP_TASK(sch, set_weight, rq, p, p->scx.weight);
> @@ -4288,9 +4319,10 @@ static void switching_to_scx(struct rq *rq, struct task_struct *p)
>  
>  	/*
>  	 * set_cpus_allowed_scx() is not called while @p is associated with a
> -	 * different scheduler class. Keep the BPF scheduler up-to-date.
> +	 * different scheduler class. Keep the BPF scheduler up-to-date. The cid
> +	 * form gets its mask from scx_enable_task().
>  	 */
> -	if (SCX_HAS_OP(sch, set_cpumask))
> +	if (!scx_is_cid_type() && SCX_HAS_OP(sch, set_cpumask))
>  		scx_call_op_set_cpumask(sch, rq, p, (struct cpumask *)p->cpus_ptr);
>  }
>  
> @@ -8374,10 +8406,11 @@ static struct bpf_struct_ops bpf_sched_ext_ops = {
>  /*
>   * cid-form cfi stubs. Stubs whose signatures match the cpu-form (param types
>   * identical, only param names differ across structs) are reused. Some need
> - * fresh stubs, set_cmask due to an argument type difference and the sub-sched
> - * notifiers because no cpu-form stub exists to reuse.
> + * fresh stubs, set_cmask and enable due to argument differences and the
> + * sub-sched notifiers because no cpu-form stub exists to reuse.
>   */
>  static void sched_ext_ops_cid__set_cmask(struct task_struct *p, const struct scx_cmask *cmask__arena) {}
> +static void sched_ext_ops_cid__enable(struct task_struct *p, struct scx_enable_args *args) {}
>  static void sched_ext_ops__sub_caps_updated(const struct scx_cmask *cmask__arena, u64 caps) {}
>  static void sched_ext_ops__sub_ecaps_updated(s32 cid, u64 before, u64 after) {}
>  
> @@ -8398,7 +8431,7 @@ static struct sched_ext_ops_cid __bpf_ops_sched_ext_ops_cid = {
>  	.update_idle		= sched_ext_ops__update_idle,
>  	.init_task		= sched_ext_ops__init_task,
>  	.exit_task		= sched_ext_ops__exit_task,
> -	.enable			= sched_ext_ops__enable,
> +	.enable			= sched_ext_ops_cid__enable,
>  	.disable		= sched_ext_ops__disable,
>  #ifdef CONFIG_EXT_GROUP_SCHED
>  	.cpuctl_init		= sched_ext_ops__cgroup_init,
> @@ -10421,8 +10454,7 @@ __bpf_kfunc const void *scx_bpf_online_cmask(const struct bpf_prog_aux *aux)
>  	if (unlikely(!online))
>  		return NULL;
>  
> -	/* BPF rebases by the low 32 bits, like __arena callback args */
> -	return (void *)((unsigned long)online - sch->arena_kern_base);
> +	return (void *)scx_kaddr_to_arena(sch, online);
>  }
>  
>  /**
> diff --git a/kernel/sched/ext/internal.h b/kernel/sched/ext/internal.h
> index d150de10a5c9..1df8f583b0ec 100644
> --- a/kernel/sched/ext/internal.h
> +++ b/kernel/sched/ext/internal.h
> @@ -250,6 +250,31 @@ struct scx_exit_task_args {
>  	bool cancelled;
>  };
>  
> +/**
> + * struct scx_enable_args - Argument container for cid-form ops.enable()
> + * @cmask_arena_addr: BPF arena address of the cmask of cids the task may run on
> + *
> + * @cmask_arena_addr is the task's affinity as it enters the scheduler.
> + * set_cmask() delivers the same mask right after enable(), before set_weight()
> + * and the first enqueue, then every affinity change afterwards, and is never
> + * called before enable(). A scheduler may therefore track affinity in
> + * set_cmask() alone.
> + *
> + * The kernel builds the mask in the scheduler arena from its own geometry, so
> + * the header is valid regardless of what the scheduler last wrote there. The
> + * memory is per-cpu scratch reused once the callback returns: copy the bits
> + * out, don't keep the address. The set_cmask() argument follows the same rules.
> + *
> + * The address is a plain value rather than a typed pointer because BTF can't
> + * mark a struct member as an arena pointer yet and a pointer member would reach
> + * the program typed as a kernel pointer. Cast it to struct scx_cmask __arena *
> + * before use. Once arena members can be typed, a typed alias will join this
> + * field in an anonymous union at the same offset.
> + */
> +struct scx_enable_args {
> +	u64	cmask_arena_addr;
> +};
> +
>  /* argument container for ops.cgroup_init() */
>  struct scx_cgroup_init_args {
>  	/* the weight of the cgroup [1..10000] */
> @@ -1037,6 +1062,7 @@ struct sched_ext_ops {
>   *   - dispatch         -> dispatch (cpu arg is now cid)
>   *   - update_idle      -> update_idle (cpu arg is now cid)
>   *   - set_cpumask      -> set_cmask (cmask instead of cpumask)
> + *   - enable           -> enable (takes struct scx_enable_args)
>   *   - cpu_online       -> cid_online
>   *   - cpu_offline      -> cid_offline
>   *   - dump_cpu         -> dump_cid
> @@ -1070,7 +1096,7 @@ struct sched_ext_ops_cid {
>  			  struct scx_init_task_args *args);
>  	void (*exit_task)(struct task_struct *p,
>  			   struct scx_exit_task_args *args);
> -	void (*enable)(struct task_struct *p);
> +	void (*enable)(struct task_struct *p, struct scx_enable_args *args);
>  	void (*disable)(struct task_struct *p);
>  	void (*dump)(struct scx_dump_ctx *ctx);
>  	void (*dump_cid)(struct scx_dump_ctx *ctx, s32 cid, bool idle);
> @@ -1533,7 +1559,8 @@ struct scx_sched {
>  	 * by BUILD_BUG_ON in scx_init()). The anonymous union lets the kernel
>  	 * access either view of the same storage without function-pointer
>  	 * casts: use .ops for cpu-form and shared fields, .ops_cid for the
> -	 * cid-renamed callbacks (set_cmask, select_cid, cid_online, ...).
> +	 * callbacks whose cid-form signature differs (set_cmask, enable,
> +	 * select_cid, cid_online, ...).
>  	 */
>  	union {
>  		struct sched_ext_ops		ops;
> @@ -1556,9 +1583,9 @@ struct scx_sched {
>  	uintptr_t		arena_kern_base;
>  
>  	/*
> -	 * Per-CPU arena cmask used by scx_call_op_set_cpumask() to hand a cmask
> -	 * to ops_cid.set_cmask(). The kernel writes through the stored kern_va
> -	 * and passes it to the callback's __arena argument.
> +	 * Per-CPU arena cmask the kernel fills from a task's cpumask and hands
> +	 * to ops_cid.enable() and ops_cid.set_cmask(). The stored pointers are
> +	 * the kernel addresses.
>  	 */
>  	struct scx_cmask * __percpu *set_cmask_scratch;
>  	struct scx_cmask *online_cmask;
> @@ -1669,6 +1696,19 @@ static inline void *scx_arena_to_kaddr(struct scx_sched *sch, const void *bpf_pt
>  	return (void *)(sch->arena_kern_base + (u32)(uintptr_t)bpf_ptr);
>  }
>  
> +/**
> + * scx_kaddr_to_arena - Translate a kernel arena address to the BPF form
> + * @sch: scheduler whose arena hosts @kaddr
> + * @kaddr: kernel address inside @sch's arena
> + *
> + * __arena callback arguments need no translation. Addresses handed to BPF any
> + * other way, such as struct fields and kfunc return values, go through this.
> + */
> +static inline uintptr_t scx_kaddr_to_arena(struct scx_sched *sch, const void *kaddr)
> +{
> +	return (uintptr_t)kaddr - sch->arena_kern_base;
> +}
> +
>  enum scx_wake_flags {
>  	/* expose select WF_* flags as enums */
>  	SCX_WAKE_FORK		= WF_FORK,
> @@ -2302,9 +2342,9 @@ do {										\
>  } while (0)
>  
>  /*
> - * Dispatch a task op through the cid-form ops_cid table. Only set_cmask() needs
> - * this: it takes an arena cmask address instead of a cpumask, so it cannot be
> - * invoked via its cpu-form set_cpumask() slot.
> + * Dispatch a task op through the cid-form ops_cid table, for the ops whose
> + * cid-form signature differs from the cpu-form slot: set_cmask() takes an arena
> + * cmask instead of a cpumask and enable() takes scx_enable_args.
>   */
>  #define SCX_CALL_CID_OP_TASK(sch, op, locked_rq, task, args...)			\
>  	__SCX_CALL_OP_TASK(sch, ops_cid, op, locked_rq, task, ##args)
> diff --git a/tools/sched_ext/scx_qmap.bpf.c b/tools/sched_ext/scx_qmap.bpf.c
> index 67b7c01cae55..2f3653199842 100644
> --- a/tools/sched_ext/scx_qmap.bpf.c
> +++ b/tools/sched_ext/scx_qmap.bpf.c
> @@ -961,9 +961,6 @@ s32 BPF_STRUCT_OPS_SLEEPABLE(qmap_init_task, struct task_struct *p,
>  	taskc->highpri = false;
>  	taskc->core_sched_seq = 0;
>  	cmask_init(&taskc->cpus_allowed, 0, scx_bpf_nr_cids());
> -	bpf_rcu_read_lock();
> -	cmask_from_cpumask(&taskc->cpus_allowed, p->cpus_ptr);
> -	bpf_rcu_read_unlock();
>  
>  	v = bpf_task_storage_get(&task_ctx_stor, p, NULL,
>  				 BPF_LOCAL_STORAGE_GET_F_CREATE);
> -- 
> 2.55.0
> 

Re: [PATCH 2/2] selftests/sched_ext: Check the cmask cid-form ops.enable() receives

[Andrea Righi]

Hi Tejun,

nits below.

On Fri, Sep 18, 2026 at 02:28:38PM -1000, Tejun Heo wrote:
> cid-form ops.enable() now hands the task's cmask to the scheduler and
> set_cmask() repeats it right after. Add a cid-form selftest that checks both
> against p->cpus_ptr, that they match each other, that the initial
> set_cmask() lands before set_weight() and before the task first becomes
> runnable, and that set_cmask() never precedes enable(), across class-switch
> enables, fork-path enables and live affinity changes.
> 
> Signed-off-by: Tejun Heo <tj@kernel.org>
> ---
>  tools/testing/selftests/sched_ext/Makefile    |   1 +
>  .../selftests/sched_ext/enable_cmask.bpf.c    | 213 ++++++++++++++++++
>  .../selftests/sched_ext/enable_cmask.c        | 138 ++++++++++++
>  3 files changed, 352 insertions(+)
>  create mode 100644 tools/testing/selftests/sched_ext/enable_cmask.bpf.c
>  create mode 100644 tools/testing/selftests/sched_ext/enable_cmask.c
> 
> diff --git a/tools/testing/selftests/sched_ext/Makefile b/tools/testing/selftests/sched_ext/Makefile
> index 49897727f535..4e06d0baaeec 100644
> --- a/tools/testing/selftests/sched_ext/Makefile
> +++ b/tools/testing/selftests/sched_ext/Makefile
> @@ -169,6 +169,7 @@ auto-test-targets :=			\
>  	ddsp_bogus_dsq_fail		\
>  	ddsp_vtimelocal_fail		\
>  	dsp_local_on			\
> +	enable_cmask			\
>  	enq_select_cpu			\
>  	exit				\
>  	hotplug				\
> diff --git a/tools/testing/selftests/sched_ext/enable_cmask.bpf.c b/tools/testing/selftests/sched_ext/enable_cmask.bpf.c
> new file mode 100644
> index 000000000000..bbbf989d2c69
> --- /dev/null
> +++ b/tools/testing/selftests/sched_ext/enable_cmask.bpf.c
> @@ -0,0 +1,213 @@
> +// SPDX-License-Identifier: GPL-2.0
> +/*
> + * A cid-form scheduler checking the cmask cid-form ops.enable() receives: the
> + * header, every cid bit against p->cpus_ptr, and that set_cmask() follows with
> + * the same mask before set_weight() and before the task first becomes runnable,
> + * and never runs before enable().
> + *
> + * Copyright (c) 2026 Tejun Heo <tj@kernel.org>
> + */
> +#include <scx/common.bpf.h>
> +
> +#define MAX_CPUS 1024
> +
> +char _license[] SEC("license") = "GPL";
> +
> +struct {
> +	__uint(type, BPF_MAP_TYPE_ARENA);
> +	__uint(map_flags, BPF_F_MMAPABLE);
> +	__uint(max_entries, 1 << 16);
> +} arena SEC(".maps");
> +
> +struct task_ctx {
> +	u64	enable_fp;	/* fingerprint of the mask enable() received */
> +	bool	enabled;
> +	bool	pending;	/* enable() ran, the initial set_cmask() hasn't */
> +};
> +
> +struct {
> +	__uint(type, BPF_MAP_TYPE_TASK_STORAGE);
> +	__uint(map_flags, BPF_F_NO_PREALLOC);
> +	__type(key, int);
> +	__type(value, struct task_ctx);
> +} task_ctx_stor SEC(".maps");
> +
> +u64 nr_enable, nr_initial_set_cmask, nr_set_cmask, nr_set_weight;
> +s32 bad_cid;
> +bool bad_want, bad_got;

As Sashiko also reported, these globals can be written concurrently by
ops.enable() and ops.set_cmask() and could make a failure report misleading. It
should be possible to move these in a caller-local stack struct.

> +
> +UEI_DEFINE(uei);
> +
> +static struct task_ctx *lookup_task_ctx(struct task_struct *p)
> +{
> +	struct task_ctx *tctx;
> +
> +	tctx = bpf_task_storage_get(&task_ctx_stor, p, 0, 0);
> +	if (!tctx)
> +		scx_bpf_error("task_ctx lookup failed for %s[%d]", p->comm, p->pid);
> +	return tctx;
> +}
> +
> +/*
> + * Verify @m's header and every cid bit against @p's cpumask and fingerprint the
> + * bits into @fp. Return 0 on success, -EINVAL on a bad header, -ENOENT on a cid
> + * without a cpu and -EIO on a bit mismatch with the details in @bad_*.
> + */
> +static int check_mask(struct task_struct *p, const struct scx_cmask __arena *m, u64 *fp)
> +{
> +	u32 nr_cids = scx_bpf_nr_cids();
> +	u64 h = 0;
> +	s32 cid;
> +
> +	if (m->base || m->nr_cids != nr_cids)
> +		return -EINVAL;

Since this is checking the mask header, should this validate alloc_words too?

	if (m->base || m->nr_cids != nr_cids ||
	    m->alloc_words != CMASK_NR_WORDS(nr_cids))
		return -EINVAL;

> +
> +	bpf_for(cid, 0, MAX_CPUS) {
> +		bool want, got;
> +		s32 cpu;
> +
> +		if (cid >= nr_cids)
> +			break;

Can we iterate directly to nr_cids?

> +		cpu = scx_bpf_cid_to_cpu(cid);
> +		if (cpu < 0)
> +			return -ENOENT;
> +		want = bpf_cpumask_test_cpu(cpu, p->cpus_ptr);
> +		got = cmask_test(cid, m);
> +		if (want != got) {
> +			bad_cid = cid;
> +			bad_want = want;
> +			bad_got = got;
> +			return -EIO;
> +		}
> +		h = h * 31 + got;
> +	}
> +
> +	*fp = h;
> +	return 0;
> +}

Thanks,
-Andrea