[RFC PATCH sched_ext/for-7.4 0/5] sched_ext: Make ops.cpu_acquire/release() reliable

[RFC PATCH sched_ext/for-7.4 0/5] sched_ext: Make ops.cpu_acquire/release() reliable

[Kuba Piecuch]

Hello,

ops.cpu_acquire() and ops.cpu_release() are meant to tell the BPF
scheduler when a CPU becomes available to it and when it's taken away.
As they are now, they can't be used to reliably track that:

- An RT task waking up on an idle CPU doesn't trigger ops.cpu_release().
  The release is only emitted from put_prev_task_scx(), i.e. when an SCX
  task is switched out.

- A running SCX task moving to a higher class through
  sched_setscheduler() or a PI boost doesn't trigger it either, as no
  context switch takes place.

- ops.cpu_acquire() is emitted speculatively from dispatch_one(). If a
  higher class task wins the pick afterwards (dispatch dropped the rq
  lock, RETRY_TASK, core-sched picking for a sibling), the CPU stays
  acquired while an RT task runs on it.

- rq->scx.cpu_released carries over from one scheduler to the next and
  CPU hotplug isn't handled.

Commit a3f5d4822253 ("sched_ext: Allow scx_bpf_reenqueue_local() to be
called from anywhere") deprecated the callbacks, as closing these gaps
was thought to require a hook invoked regardless of sched class. This
series shows that it doesn't.

Whether a CPU is available to the BPF scheduler only depends on the
class of rq->donor: the CPU is available while an SCX task or the idle
task runs on it and released while a stop, DL, RT or fair (partial
mode) task does. Every change of availability has an SCX task or the
idle task on one side of the switch, so the ext and idle classes' own
put_prev_task() and set_next_task() see all of them. Switches between
higher class tasks never change availability. The one remaining case, a
running task leaving SCX without a context switch, is caught by
switched_from_scx().

kernel/sched/core.c is left untouched. The idle class gets two calls
next to its existing scx_update_idle() calls and sched.h gets one new
scx_rq flag.

Points worth discussing:

- Core-sched forced idle counts as available. Whether a task can be
  placed on a forced idle CPU depends on that task's cookie, which is
  for the BPF scheduler to judge.

- ops.cpu_acquire() moves from before ops.dispatch() to the point where
  the CPU actually switches to an SCX task or idle, so ops.dispatch()
  may now run on a released CPU. A speculative acquire before dispatch
  can't be made reliable without hooks in every class.

- The callbacks may be invoked for a remote CPU, with that CPU's rq
  locked: sched_setscheduler() or a PI boost of a task running on
  another CPU, and scheduler enable. Patch 1 makes the legacy
  scx_bpf_reenqueue_local() re-enqueue the locked rq's local DSQ instead
  of assuming smp_processor_id().

- Every CPU starts out released, and each available online CPU is
  acquired when the scheduler is enabled. The callbacks are only
  invoked between ops.cpu_online() and ops.cpu_offline(). Because
  rq_online_scx() runs before ops.cpu_online() when a CPU comes up,
  this uses a new SCX_RQ_OPS_ONLINE flag rather than SCX_RQ_ONLINE.

- Schedulers which don't implement the callbacks only pay for a static
  branch.

- The callbacks stay cpu-form only, which makes them root-only by
  construction. cid-form variants and fan-out to sub-schedulers (gated
  on SCX_CAP_BASE, like ops.update_idle()) could follow if there's
  interest.

Patch 1 makes the legacy scx_bpf_reenqueue_local() target the locked
rq. Patch 2 emits the callbacks on actual transitions and handles
scheduler enable/disable and CPU hotplug. Patch 3 releases the CPU when
its running task leaves SCX. Patch 4 un-deprecates the callbacks and
documents the contract. Patch 5 adds selftests.

Testing was done in virtme-ng with PROVE_LOCKING, DEBUG_ATOMIC_SLEEP,
SCHED_CORE and EXT_SUB_SCHED enabled. The full sched_ext selftest suite
passes without splats. The new cpu_avail test checks the reported state
on every context switch through a sched_switch tracepoint and covers
idle -> RT, remote class changes, the legacy re-enqueue from a remote
release, PI boosting, hotplug, scheduler reloads and partial mode. It
fails on the base kernel (idle -> RT) and with only patches 1-2 applied
(remote class change). cpu_avail_core, run in an SMT VM, checks that
core-sched forced idle doesn't release a CPU while an RT task with a
matching cookie does.

Not included:

- tools/sched_ext/include/scx/*autogen* isn't regenerated for the new
  SCX_CPU_PREEMPT_FAIR and SCX_CPU_PREEMPT_CLASS_CHANGE values. I'll add
  that as a separate patch in a non-RFC posting.

- Unrelated to this series: the comment on SCX_RQ_ONLINE in sched.h says
  it tracks the ops.cpu_online()/cpu_offline() state, but as noted above
  rq_online_scx() runs before ops.cpu_online() on hotplug. I can send a
  separate fix if that's useful.

Kuba Piecuch (5):
  sched_ext: Make legacy scx_bpf_reenqueue_local() target the locked rq
  sched_ext: Emit ops.cpu_acquire/release() on actual CPU transitions
  sched_ext: Release the CPU when its running task leaves SCX
  sched_ext: Un-deprecate ops.cpu_acquire/release()
  selftests/sched_ext: Add cpu_avail tests

 kernel/sched/ext/ext.c                        | 265 +++++--
 kernel/sched/ext/ext.h                        |  19 +
 kernel/sched/ext/internal.h                   |  69 +-
 kernel/sched/idle.c                           |   4 +
 kernel/sched/sched.h                          |   6 +
 tools/testing/selftests/sched_ext/Makefile    |   2 +
 .../selftests/sched_ext/cpu_avail.bpf.c       | 218 ++++++
 tools/testing/selftests/sched_ext/cpu_avail.c | 717 ++++++++++++++++++
 tools/testing/selftests/sched_ext/cpu_avail.h | 237 ++++++
 .../selftests/sched_ext/cpu_avail_core.c      | 256 +++++++
 .../testing/selftests/sched_ext/maximal.bpf.c |  10 +
 11 files changed, 1706 insertions(+), 97 deletions(-)
 create mode 100644 tools/testing/selftests/sched_ext/cpu_avail.bpf.c
 create mode 100644 tools/testing/selftests/sched_ext/cpu_avail.c
 create mode 100644 tools/testing/selftests/sched_ext/cpu_avail.h
 create mode 100644 tools/testing/selftests/sched_ext/cpu_avail_core.c


base-commit: 2f2f7eeef8b0a1a9b9ca998231419873ccae8d4e
-- 
2.56.0.176.g3575e65682-goog

[RFC PATCH 1/5] sched_ext: Make legacy scx_bpf_reenqueue_local() target the locked rq

[Kuba Piecuch]

The legacy, synchronous scx_bpf_reenqueue_local() can only be called
from ops.cpu_release(). It assumes that ops.cpu_release() always runs on
the CPU being released and re-enqueues the tasks on the local DSQ of
smp_processor_id().

That assumption is about to stop holding. A following patch emits
ops.cpu_release() when the task running on a CPU leaves SCX through
sched_setscheduler() or a PI boost. That can happen from any CPU, with
the target CPU's rq locked. Picking the rq from smp_processor_id() would
then trip lockdep_assert_rq_held() and walk another CPU's local DSQ
without holding that CPU's rq lock.

SCX_CALL_OP() already records the rq the operation is invoked for in
scx_locked_rq(), and scx_bpf_dsq_reenq() /
scx_bpf_reenqueue_local___v2() already resolve SCX_DSQ_LOCAL against it.
Do the same in the legacy kfunc. Nothing else in reenq_local() depends
on running on the target CPU: the local DSQ is protected by the rq lock,
and direct dispatches to a third CPU's local DSQ are already deferred
through ddsp_deferred_locals.

No functional change for the existing callers, which all run on the CPU
whose rq is locked.

Also update the kerneldoc of scx_bpf_reenqueue_local___v2(), which still
says that it works on the local DSQ of the caller's CPU, to match
scx_bpf_dsq_reenq().

Signed-off-by: Kuba Piecuch <jpiecuch@google.com>
---
 kernel/sched/ext/ext.c | 19 ++++++++++++++-----
 1 file changed, 14 insertions(+), 5 deletions(-)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index f1b9d52f624a..f64438290218 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -9501,9 +9501,10 @@ __bpf_kfunc_start_defs();
  * scx_bpf_reenqueue_local - Re-enqueue tasks on a local DSQ
  * @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
  *
- * Iterate over all of the tasks currently enqueued on the local DSQ of the
- * caller's CPU, and re-enqueue them in the BPF scheduler. Returns the number of
- * processed tasks. Can only be called from ops.cpu_release().
+ * Iterate over all of the tasks currently enqueued on the local DSQ of the CPU
+ * ops.cpu_release() is being invoked for, and re-enqueue them in the BPF
+ * scheduler. Returns the number of processed tasks. Can only be called from
+ * ops.cpu_release().
  */
 __bpf_kfunc u32 scx_bpf_reenqueue_local(const struct bpf_prog_aux *aux)
 {
@@ -9515,7 +9516,13 @@ __bpf_kfunc u32 scx_bpf_reenqueue_local(const struct bpf_prog_aux *aux)
 	if (unlikely(!sch))
 		return 0;
 
-	rq = cpu_rq(smp_processor_id());
+	/*
+	 * Use the rq ops.cpu_release() was invoked for rather than the local
+	 * one. They are not necessarily the same CPU.
+	 */
+	rq = scx_locked_rq();
+	if (WARN_ON_ONCE(!rq))
+		return 0;
 	lockdep_assert_rq_held(rq);
 
 	return reenq_local(sch, rq, SCX_REENQ_ANY);
@@ -10057,7 +10064,9 @@ __bpf_kfunc void scx_bpf_dsq_reenq(u64 dsq_id, u64 reenq_flags,
  * @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
  *
  * Iterate over all of the tasks currently enqueued on the local DSQ of the
- * caller's CPU, and re-enqueue them in the BPF scheduler. Can be called from
+ * rq the current scheduler operation is locked to - e.g. the CPU
+ * ops.cpu_release() is being invoked for - or the calling CPU's when no rq is
+ * locked, and re-enqueue them in the BPF scheduler. Can be called from
  * anywhere.
  *
  * This is now a special case of scx_bpf_dsq_reenq() and may be removed in the
-- 
2.56.0.176.g3575e65682-goog

[RFC PATCH 2/5] sched_ext: Emit ops.cpu_acquire/release() on actual CPU transitions

[Kuba Piecuch]

ops.cpu_acquire() and ops.cpu_release() are meant to tell the BPF
scheduler when a CPU becomes available to it and when it's taken away.
In their current form they miss events and can get stuck in the wrong
state:

- ops.cpu_release() is only emitted from put_prev_task_scx(), so a CPU
  going from idle straight to an RT, DL or stop task is never released.
  This also covers the case where ttwu() direct-dispatched a task to an
  idle CPU's local DSQ but a higher class task wins the next pick:
  nothing on the SCX side runs at all.

- ops.cpu_acquire() is emitted speculatively from dispatch_one(), before
  the pick is final. If dispatch drops the rq lock and a higher class
  task shows up, or the pick is retried, the CPU stays "acquired" while
  an RT task runs. Under core scheduling, dispatch_one() also runs for
  the siblings, so a CPU can be acquired even though its final pick is
  an RT task.

- rq->scx.cpu_released is never reset. It carries over from one
  scheduler to the next, and CPU hotplug isn't handled.

Whether a CPU is available to the BPF scheduler only depends on the
class of its scheduling context, rq->donor: the CPU is available while
an SCX task or the idle task runs on it and released while a stop, DL,
RT or fair (partial mode) task does. Every change of availability has
an SCX task or the idle task on one side, so the ext and idle classes'
own put_prev_task() and set_next_task() see all of them. No hook in the
core or in the other classes is needed.

Replace switch_class() and the acquire in dispatch_one() with a small,
idempotent state machine, __scx_update_cpu_avail(), which looks at the
task about to become rq->donor and only invokes a callback on an actual
edge. Call it from:

- set_next_task_scx(), before ops.running(): acquire.
- put_prev_task_scx(), after the outgoing task went back to the local
  DSQ, so that ops.cpu_release() can still re-enqueue it: release if the
  next task is from a higher class.
- put_prev_task_idle(), after ops.update_idle(cpu, false): release if
  the next task is from a higher class.
- set_next_task_idle(), before ops.update_idle(cpu, true): acquire.

Core-sched forced idle counts as available. Whether the BPF scheduler
can place a task on a forced idle CPU depends on that task's cookie,
which is for the BPF scheduler to decide. Idle-to-idle re-picks don't
change availability and pick_task_idle() is left alone.

Every CPU starts released. When a root scheduler implementing either
callback is enabled, reset all CPUs to released, turn on the static key
gating the hooks and acquire every available online CPU. The key is
turned off under scx_enable_mutex when the scheduler is disabled, before
tasks are switched out of SCX, so the BPF scheduler doesn't see a storm
of releases. In-flight callbacks run with an rq lock held and IRQs
disabled and are flushed by the synchronize_rcu() before ops.exit().
Schedulers which don't implement the callbacks only pay for a static
branch.

For hotplug, the callbacks are only invoked between ops.cpu_online() and
ops.cpu_offline(): ops.cpu_offline() implies a release and the CPU is
acquired right after ops.cpu_online() if it's available. This can't be
keyed off SCX_RQ_ONLINE as rq_online_scx() runs before ops.cpu_online()
when a CPU comes up: cpuset_cpu_active() rebuilds the sched domains
before sched_cpu_activate() calls scx_rq_activate(). Add
SCX_RQ_OPS_ONLINE, which scx_rq_activate() sets right after
ops.cpu_online() and scx_rq_deactivate() clears right before
ops.cpu_offline().

Add SCX_CPU_PREEMPT_FAIR for releases to a fair task in partial mode,
which used to be reported as SCX_CPU_PREEMPT_UNKNOWN. It's appended to
keep the existing values unchanged.

This changes when ops.cpu_acquire() is called: it's now called when the
CPU actually switches to an SCX task or idle rather than before
ops.dispatch(). ops.dispatch() may thus be called while the CPU is still
released.

A running SCX task changing class without a context switch is not
handled yet and is addressed by the next patch.

Signed-off-by: Kuba Piecuch <jpiecuch@google.com>
---
 kernel/sched/ext/ext.c      | 220 +++++++++++++++++++++++++-----------
 kernel/sched/ext/ext.h      |  19 ++++
 kernel/sched/ext/internal.h |   2 +
 kernel/sched/idle.c         |   4 +
 kernel/sched/sched.h        |   6 +
 5 files changed, 185 insertions(+), 66 deletions(-)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index f64438290218..4d02d2d6e6d0 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -80,6 +80,9 @@ static bool scx_switching_all;
 DEFINE_STATIC_KEY_FALSE(__scx_switched_all);
 static DEFINE_STATIC_KEY_FALSE(__scx_tid_to_task_enabled);
 
+/* the root scheduler implements ops.cpu_acquire() or ops.cpu_release() */
+DEFINE_STATIC_KEY_FALSE(__scx_cpu_avail_enabled);
+
 /*
  * Gates cgroup ops delivery. Set at the end of the cgroup init phase of root
  * enable and cleared before root disable starts tearing down tasks, both under
@@ -2964,19 +2967,6 @@ static enum scx_dsp_verdict dispatch_one(struct rq *rq, struct task_struct *prev
 
 	scx_process_sync_ecaps(rq, prev);
 
-	if ((root_sch->ops.flags & SCX_OPS_HAS_CPU_PREEMPT) &&
-	    unlikely(rq->scx.cpu_released)) {
-		/*
-		 * If the previous sched_class for the current CPU was not SCX,
-		 * notify the BPF scheduler that it again has control of the
-		 * core. This callback complements ->cpu_release(), which is
-		 * emitted in switch_class().
-		 */
-		if (root_sch->ops.cpu_acquire)
-			SCX_CALL_OP(root_sch, cpu_acquire, rq, cpu, NULL);
-		rq->scx.cpu_released = false;
-	}
-
 	if (prev->sched_class == &ext_sched_class) {
 		update_curr_scx(rq);
 
@@ -3044,10 +3034,101 @@ static enum scx_dsp_verdict dispatch_one(struct rq *rq, struct task_struct *prev
 	return verdict;
 }
 
+/*
+ * A CPU is available to the BPF scheduler while its scheduling context is an
+ * SCX task or the idle task. Core-sched forced idle counts as available:
+ * whether a task can be placed on such a CPU depends on that task's cookie,
+ * which is the BPF scheduler's call.
+ */
+static bool scx_cpu_avail(struct rq *rq, struct task_struct *next)
+{
+	return next->sched_class == &ext_sched_class || next == rq->idle;
+}
+
+static enum scx_cpu_preempt_reason cpu_release_reason(struct task_struct *next)
+{
+	if (next->sched_class == &stop_sched_class)
+		return SCX_CPU_PREEMPT_STOP;
+	if (next->sched_class == &dl_sched_class)
+		return SCX_CPU_PREEMPT_DL;
+	if (next->sched_class == &rt_sched_class)
+		return SCX_CPU_PREEMPT_RT;
+	if (next->sched_class == &fair_sched_class)
+		return SCX_CPU_PREEMPT_FAIR;
+	return SCX_CPU_PREEMPT_UNKNOWN;
+}
+
+/*
+ * ops.cpu_acquire/release() are edge-triggered on @rq->scx.cpu_released and
+ * only fire between ops.cpu_online() and ops.cpu_offline(), see
+ * scx_rq_activate() and scx_rq_deactivate().
+ */
+static void scx_cpu_release(struct rq *rq, enum scx_cpu_preempt_reason reason,
+			    struct task_struct *task)
+{
+	struct scx_sched *sch = scx_root_protected_live();
+
+	lockdep_assert_rq_held(rq);
+
+	if (rq->scx.cpu_released || !(rq->scx.flags & SCX_RQ_OPS_ONLINE))
+		return;
+
+	/* update before calling out so that a nested re-evaluation is a no-op */
+	rq->scx.cpu_released = true;
+
+	if (sch->ops.cpu_release) {
+		struct scx_cpu_release_args args = {
+			.reason = reason,
+			.task = task,
+		};
+
+		SCX_CALL_OP(sch, cpu_release, rq, cpu_of(rq), &args);
+	}
+}
+
+static void scx_cpu_acquire(struct rq *rq)
+{
+	struct scx_sched *sch = scx_root_protected_live();
+
+	lockdep_assert_rq_held(rq);
+
+	if (!rq->scx.cpu_released || !(rq->scx.flags & SCX_RQ_OPS_ONLINE))
+		return;
+
+	rq->scx.cpu_released = false;
+
+	if (sch->ops.cpu_acquire)
+		SCX_CALL_OP(sch, cpu_acquire, rq, cpu_of(rq), NULL);
+}
+
+/**
+ * __scx_update_cpu_avail - Re-evaluate whether a CPU is available
+ * @rq: rq of the CPU, locked
+ * @next: task that is (about to become) @rq->donor
+ *
+ * Availability only depends on the class of @rq->donor and every change has
+ * an SCX task or the idle task on one side. Callers re-evaluate whenever that
+ * may have happened and a callback is only invoked on an actual transition.
+ */
+void __scx_update_cpu_avail(struct rq *rq, struct task_struct *next)
+{
+	if (scx_cpu_avail(rq, next))
+		scx_cpu_acquire(rq);
+	else
+		scx_cpu_release(rq, cpu_release_reason(next), next);
+}
+
 static void set_next_task_scx(struct rq *rq, struct task_struct *p, bool first)
 {
 	struct scx_sched *sch = scx_task_sched(p);
 
+	/*
+	 * Acquire before ops.dequeue(SCX_DEQ_CORE_SCHED_EXEC) and
+	 * ops.running(). This also covers a running task switching to SCX
+	 * through sched_change_end().
+	 */
+	scx_update_cpu_avail(rq, p);
+
 	if (p->scx.flags & SCX_TASK_QUEUED) {
 		/*
 		 * Core-sched might decide to execute @p before it is
@@ -3105,57 +3186,6 @@ static void set_next_task_scx(struct rq *rq, struct task_struct *p, bool first)
 	}
 }
 
-static enum scx_cpu_preempt_reason
-preempt_reason_from_class(const struct sched_class *class)
-{
-	if (class == &stop_sched_class)
-		return SCX_CPU_PREEMPT_STOP;
-	if (class == &dl_sched_class)
-		return SCX_CPU_PREEMPT_DL;
-	if (class == &rt_sched_class)
-		return SCX_CPU_PREEMPT_RT;
-	return SCX_CPU_PREEMPT_UNKNOWN;
-}
-
-static void switch_class(struct rq *rq, struct task_struct *next)
-{
-	struct scx_sched *sch = scx_root_protected_live();
-	const struct sched_class *next_class = next->sched_class;
-
-	if (!(sch->ops.flags & SCX_OPS_HAS_CPU_PREEMPT))
-		return;
-
-	/*
-	 * The callback is conceptually meant to convey that the CPU is no
-	 * longer under the control of SCX. Therefore, don't invoke the callback
-	 * if the next class is below SCX (in which case the BPF scheduler has
-	 * actively decided not to schedule any tasks on the CPU).
-	 */
-	if (sched_class_above(&ext_sched_class, next_class))
-		return;
-
-	/*
-	 * At this point we know that SCX was preempted by a higher priority
-	 * sched_class, so invoke the ->cpu_release() callback if we have not
-	 * done so already. We only send the callback once between SCX being
-	 * preempted, and it regaining control of the CPU.
-	 *
-	 * ->cpu_release() complements ->cpu_acquire(), which is emitted the
-	 *  next time that dispatch_one() is invoked.
-	 */
-	if (!rq->scx.cpu_released) {
-		if (sch->ops.cpu_release) {
-			struct scx_cpu_release_args args = {
-				.reason = preempt_reason_from_class(next_class),
-				.task = next,
-			};
-
-			SCX_CALL_OP(sch, cpu_release, rq, cpu_of(rq), &args);
-		}
-		rq->scx.cpu_released = true;
-	}
-}
-
 static void put_prev_task_scx(struct rq *rq, struct task_struct *p,
 			      struct task_struct *next)
 {
@@ -3254,8 +3284,14 @@ static void put_prev_task_scx(struct rq *rq, struct task_struct *p,
 	}
 
 switch_class:
+	/*
+	 * Release the CPU if @next is from a higher class. This is done after
+	 * @p went back to the local DSQ so that ops.cpu_release() can
+	 * re-enqueue it. SCX -> idle (including core-sched forced idle) is
+	 * not a transition as both are available.
+	 */
 	if (next && next->sched_class != &ext_sched_class)
-		switch_class(rq, next);
+		scx_update_cpu_avail(rq, next);
 }
 
 static void kick_sync_wait_bal_cb(struct rq *rq)
@@ -3749,10 +3785,29 @@ static void handle_hotplug(struct rq *rq, bool online)
 void scx_rq_activate(struct rq *rq)
 {
 	handle_hotplug(rq, true);
+
+	/*
+	 * The CPU comes up released. Now that ops.cpu_online() has been
+	 * called, let ops.cpu_acquire/release() through and acquire the CPU if
+	 * it's available. This can't be done from rq_online_scx() which may
+	 * run before ops.cpu_online() when the sched domains are rebuilt.
+	 */
+	guard(rq_lock_irqsave)(rq);
+	rq->scx.flags |= SCX_RQ_OPS_ONLINE;
+	scx_update_cpu_avail(rq, rq->donor);
 }
 
 void scx_rq_deactivate(struct rq *rq)
 {
+	/*
+	 * ops.cpu_offline() implies release. No ops.cpu_acquire/release()
+	 * until ops.cpu_online().
+	 */
+	scoped_guard(rq_lock_irqsave, rq) {
+		rq->scx.flags &= ~SCX_RQ_OPS_ONLINE;
+		rq->scx.cpu_released = true;
+	}
+
 	handle_hotplug(rq, false);
 }
 
@@ -6579,6 +6634,15 @@ static void scx_root_disable(struct scx_sched *sch)
 	 */
 	mutex_lock(&scx_enable_mutex);
 
+	/*
+	 * Stop ops.cpu_acquire/release() before switching tasks out so that
+	 * the BPF scheduler doesn't see a storm of releases. This must be under
+	 * scx_enable_mutex so that it's ordered after the enable path turning
+	 * the key on. In-flight callbacks hold an rq lock with IRQs disabled
+	 * and are flushed by synchronize_rcu() before ops.exit().
+	 */
+	static_branch_disable(&__scx_cpu_avail_enabled);
+
 	was_switched_all = scx_switched_all();
 
 	static_branch_disable(&__scx_switched_all);
@@ -7991,6 +8055,30 @@ static void scx_root_enable_workfn(struct kthread_work *work)
 		goto err_disable;
 	}
 
+	/*
+	 * Every CPU starts released. Once the hooks are live, acquire the
+	 * available CPUs. CPUs which don't context switch (e.g. idle ones)
+	 * wouldn't get there otherwise. scx_root_disable() turns the key off
+	 * under scx_enable_mutex, so it can't race this.
+	 */
+	if (sch->ops.flags & SCX_OPS_HAS_CPU_PREEMPT) {
+		for_each_possible_cpu(cpu) {
+			struct rq *rq = cpu_rq(cpu);
+
+			guard(rq_lock_irqsave)(rq);
+			rq->scx.cpu_released = true;
+		}
+
+		static_branch_enable(&__scx_cpu_avail_enabled);
+
+		for_each_possible_cpu(cpu) {
+			struct rq *rq = cpu_rq(cpu);
+
+			guard(rq_lock_irqsave)(rq);
+			__scx_update_cpu_avail(rq, rq->donor);
+		}
+	}
+
 	if (!(ops->flags & SCX_OPS_SWITCH_PARTIAL))
 		static_branch_enable(&__scx_switched_all);
 
@@ -8871,7 +8959,7 @@ void __init init_sched_ext_class(void)
 		rq->scx.kick_cpus_irq_work = IRQ_WORK_INIT_HARD(kick_cpus_irq_workfn);
 
 		if (cpu_online(cpu))
-			cpu_rq(cpu)->scx.flags |= SCX_RQ_ONLINE;
+			cpu_rq(cpu)->scx.flags |= SCX_RQ_ONLINE | SCX_RQ_OPS_ONLINE;
 	}
 
 	register_sysrq_key('S', &sysrq_sched_ext_reset_op);
diff --git a/kernel/sched/ext/ext.h b/kernel/sched/ext/ext.h
index 0b7fc46aee08..c03fd8bc0d66 100644
--- a/kernel/sched/ext/ext.h
+++ b/kernel/sched/ext/ext.h
@@ -70,6 +70,25 @@ static inline void scx_update_idle(struct rq *rq, bool idle, bool do_notify)
 static inline void scx_update_idle(struct rq *rq, bool idle, bool do_notify) {}
 #endif
 
+#ifdef CONFIG_SCHED_CLASS_EXT
+DECLARE_STATIC_KEY_FALSE(__scx_cpu_avail_enabled);
+void __scx_update_cpu_avail(struct rq *rq, struct task_struct *next);
+
+/*
+ * @next is (about to become) @rq->donor. Notify the root BPF scheduler if this
+ * changes whether @rq's CPU is available to it. Only enabled if the root
+ * scheduler implements ops.cpu_acquire() or ops.cpu_release(). A NULL @next,
+ * from put_prev_task() without a next task, is ignored.
+ */
+static inline void scx_update_cpu_avail(struct rq *rq, struct task_struct *next)
+{
+	if (static_branch_unlikely(&__scx_cpu_avail_enabled) && next)
+		__scx_update_cpu_avail(rq, next);
+}
+#else
+static inline void scx_update_cpu_avail(struct rq *rq, struct task_struct *next) {}
+#endif
+
 #ifdef CONFIG_CGROUP_SCHED
 #ifdef CONFIG_EXT_GROUP_SCHED
 void scx_tg_init(struct task_group *tg);
diff --git a/kernel/sched/ext/internal.h b/kernel/sched/ext/internal.h
index 352547aedbb7..2d8c789c1e76 100644
--- a/kernel/sched/ext/internal.h
+++ b/kernel/sched/ext/internal.h
@@ -298,6 +298,8 @@ enum scx_cpu_preempt_reason {
 	SCX_CPU_PREEMPT_STOP,
 	/* unknown reason for SCX being preempted */
 	SCX_CPU_PREEMPT_UNKNOWN,
+	/* next task is being scheduled by &fair_sched_class (partial mode) */
+	SCX_CPU_PREEMPT_FAIR,
 };
 
 /*
diff --git a/kernel/sched/idle.c b/kernel/sched/idle.c
index eb73b65ce6c4..89313f697124 100644
--- a/kernel/sched/idle.c
+++ b/kernel/sched/idle.c
@@ -484,12 +484,16 @@ static void put_prev_task_idle(struct rq *rq, struct task_struct *prev, struct t
 {
 	update_curr_idle(rq);
 	scx_update_idle(rq, false, true);
+	/* idle -> stop/DL/RT/fair releases the CPU from the BPF scheduler */
+	scx_update_cpu_avail(rq, next);
 	update_rq_avg_idle(rq);
 }
 
 static void set_next_task_idle(struct rq *rq, struct task_struct *next, bool first)
 {
 	update_idle_core(rq);
+	/* stop/DL/RT/fair -> idle acquires the CPU for the BPF scheduler */
+	scx_update_cpu_avail(rq, next);
 	scx_update_idle(rq, true, true);
 	schedstat_inc(rq->sched_goidle);
 	next->se.exec_start = rq_clock_task(rq);
diff --git a/kernel/sched/sched.h b/kernel/sched/sched.h
index 4c25fbe84fb5..7a93ee3d442d 100644
--- a/kernel/sched/sched.h
+++ b/kernel/sched/sched.h
@@ -787,6 +787,12 @@ enum scx_rq_flags {
 	 */
 	SCX_RQ_ONLINE		= 1 << 0,
 	SCX_RQ_CAN_STOP_TICK	= 1 << 1,
+	/*
+	 * Set between ops.cpu_online() and ops.cpu_offline(). Unlike
+	 * %SCX_RQ_ONLINE, this is never set before ops.cpu_online() has been
+	 * called. Gates ops.cpu_acquire/release().
+	 */
+	SCX_RQ_OPS_ONLINE	= 1 << 2,
 	SCX_RQ_CLK_VALID	= 1 << 5, /* RQ clock is fresh and valid */
 	SCX_RQ_BAL_CB_PENDING	= 1 << 6, /* must queue a cb after dispatching */
 	SCX_RQ_SUB_IDLE_RENOTIFY	= 1 << 7, /* sub-scheds are owed update_idle() */
-- 
2.56.0.176.g3575e65682-goog

[RFC PATCH 3/5] sched_ext: Release the CPU when its running task leaves SCX

[Kuba Piecuch]

When the task running on a CPU switches from SCX to another sched_class
through sched_setscheduler() or a PI boost, the CPU is taken away from
the BPF scheduler without a context switch. sched_change_begin() calls
put_prev_task() with a NULL @next, so put_prev_task_scx() can't tell
where the task is going, and the new class's set_next_task() doesn't
notify SCX. The BPF scheduler is never told that the CPU has been
released.

Release the CPU from switched_from_scx() if @p is the CPU's scheduling
context. Every class @p can switch to is above SCX, so the CPU is no
longer available. As the new class isn't known yet, report the new
SCX_CPU_PREEMPT_CLASS_CHANGE reason, with @args->task pointing to @p.
The release happens before ops.disable() for @p.

switched_from_scx() is the only reliable place for this:
put_prev_task_scx() with a NULL @next is also called for changes which
keep @p in SCX, and the wakeup_preempt() call at the end of
sched_change_end() is skipped if rq->next_class has already been raised.

The reverse direction needs nothing new: when a running task switches to
SCX, sched_change_end() calls set_next_task_scx(), which acquires the
CPU.

sched_setscheduler() and PI boosting can target a task running on
another CPU, so ops.cpu_release() can now be invoked for a remote CPU,
with that CPU's rq locked. The legacy scx_bpf_reenqueue_local() was
already switched to the locked rq to handle this.

Signed-off-by: Kuba Piecuch <jpiecuch@google.com>
---
 kernel/sched/ext/ext.c      | 11 +++++++++++
 kernel/sched/ext/internal.h |  2 ++
 2 files changed, 13 insertions(+)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index 4d02d2d6e6d0..d56c692f3dfa 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -4433,6 +4433,17 @@ static void switching_to_scx(struct rq *rq, struct task_struct *p)
 
 static void switched_from_scx(struct rq *rq, struct task_struct *p)
 {
+	/*
+	 * The running task changing class (sched_setscheduler(), PI boost)
+	 * keeps the CPU without a context switch, so neither
+	 * put_prev_task_scx() nor the idle class sees it. Every class @p can
+	 * switch to is above SCX, release the CPU. The new class isn't known
+	 * yet. The scheduler being disabled is excluded by the static key.
+	 */
+	if (static_branch_unlikely(&__scx_cpu_avail_enabled) &&
+	    task_current_donor(rq, p))
+		scx_cpu_release(rq, SCX_CPU_PREEMPT_CLASS_CHANGE, p);
+
 	if (task_dead_and_done(p))
 		return;
 
diff --git a/kernel/sched/ext/internal.h b/kernel/sched/ext/internal.h
index 2d8c789c1e76..65b6ec7dc2f3 100644
--- a/kernel/sched/ext/internal.h
+++ b/kernel/sched/ext/internal.h
@@ -300,6 +300,8 @@ enum scx_cpu_preempt_reason {
 	SCX_CPU_PREEMPT_UNKNOWN,
 	/* next task is being scheduled by &fair_sched_class (partial mode) */
 	SCX_CPU_PREEMPT_FAIR,
+	/* the running task switched from SCX to another sched_class */
+	SCX_CPU_PREEMPT_CLASS_CHANGE,
 };
 
 /*
-- 
2.56.0.176.g3575e65682-goog

[RFC PATCH 4/5] sched_ext: Un-deprecate ops.cpu_acquire/release()

[Kuba Piecuch]

Commit a3f5d4822253 ("sched_ext: Allow scx_bpf_reenqueue_local() to be
called from anywhere") deprecated ops.cpu_acquire() and
ops.cpu_release() because they missed events, and fixing that was
thought to require a hook invoked regardless of sched class. The
previous patches make the callbacks reliable using only hooks in the ext
and idle classes: they are now emitted on every transition of a CPU
between being available to the BPF scheduler and being released, and
only then.

A reliable, edge-triggered view of CPU availability is useful to BPF
schedulers beyond re-enqueueing preempted tasks, e.g. to keep track of
which CPUs can be used for placement without tracing every context
switch through the sched_switch tracepoint.

Drop the deprecation warning from scx_validate_ops() and the deprecation
notes from the kerneldoc. Document the contract in the kerneldoc: what
available means, the order relative to ops.running() and
ops.update_idle(), that ops.dispatch() may run on a released CPU, that
the callbacks may be invoked for a remote CPU, and the rules for
scheduler load and CPU hotplug. Also note that for
SCX_CPU_PREEMPT_CLASS_CHANGE, the task in the release arguments is the
one which left SCX.

The callbacks stay cpu-form only. They remain at the tail of struct
sched_ext_ops, past SCX_OPI_END, and struct sched_ext_ops_cid doesn't
have them. Since cpu-form schedulers can neither be sub-schedulers nor
have sub-schedulers, they are root-only by construction.

Signed-off-by: Kuba Piecuch <jpiecuch@google.com>
---
 kernel/sched/ext/ext.c      | 15 +++------
 kernel/sched/ext/internal.h | 65 ++++++++++++++++++++++++++++---------
 2 files changed, 54 insertions(+), 26 deletions(-)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index d56c692f3dfa..7f4496667b1a 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -7681,14 +7681,6 @@ int scx_validate_ops(struct scx_sched *sch, const struct sched_ext_ops *ops)
 		return -EINVAL;
 	}
 
-	/*
-	 * cid-form's struct is shorter and doesn't include the cpu_acquire /
-	 * cpu_release tail; reading those fields off a cid-form @ops would
-	 * run past the BPF allocation. Skip for cid-form.
-	 */
-	if (!sch->is_cid_type && (ops->cpu_acquire || ops->cpu_release))
-		pr_warn_ratelimited("ops->cpu_acquire/release() are deprecated, use sched_switch TP instead\n");
-
 	/*
 	 * Sub-scheduler support is tied to the cid-form struct_ops. A sub-sched
 	 * attaches through a cid-form-only interface (sub_attach/sub_detach),
@@ -11283,9 +11275,10 @@ static int __init scx_init(void)
 	/* @priv tail must align since both share the same data block */
 	CID_OFFSET_MATCH(priv, priv);
 	/*
-	 * cid-form must end exactly at @priv - scx_validate_ops() skips
-	 * cpu_acquire/cpu_release for cid-form because reading those fields
-	 * past the BPF allocation would be UB.
+	 * cid-form must end exactly at @priv. It doesn't have the cpu-form only
+	 * cpu_acquire/cpu_release and reading them off a cid-form @ops would
+	 * run past the BPF allocation. The kernel copy in sch->ops keeps them
+	 * zero for cid-form, see scx_alloc_and_add_sched().
 	 */
 	BUILD_BUG_ON(offsetof(struct sched_ext_ops_cid, __end) !=
 		     offsetofend(struct sched_ext_ops, priv));
diff --git a/kernel/sched/ext/internal.h b/kernel/sched/ext/internal.h
index 65b6ec7dc2f3..415fc0cd314c 100644
--- a/kernel/sched/ext/internal.h
+++ b/kernel/sched/ext/internal.h
@@ -315,7 +315,11 @@ struct scx_cpu_release_args {
 	/* the reason the CPU was preempted */
 	enum scx_cpu_preempt_reason reason;
 
-	/* the task that's going to be scheduled on the CPU */
+	/*
+	 * The task that's going to be scheduled on the CPU. For
+	 * %SCX_CPU_PREEMPT_CLASS_CHANGE, the task which left SCX and keeps
+	 * running on the CPU.
+	 */
 	struct task_struct	*task;
 };
 
@@ -1019,12 +1023,33 @@ struct sched_ext_ops {
 	void __rcu *priv;
 
 	/*
-	 * Deprecated callbacks. Kept at the end of the struct so the cid-form
-	 * struct (sched_ext_ops_cid) can omit them without affecting the
-	 * shared field offsets. Use SCX_ENQ_IMMED instead. Sitting past
-	 * SCX_OPI_END means has_op doesn't cover them, so SCX_HAS_OP() cannot
-	 * be used; callers must test sch->ops.cpu_acquire / cpu_release
-	 * directly.
+	 * cpu-form only callbacks. Kept at the end of the struct so the
+	 * cid-form struct (sched_ext_ops_cid) can omit them without affecting
+	 * the shared field offsets. Sitting past SCX_OPI_END means has_op
+	 * doesn't cover them, so SCX_HAS_OP() cannot be used; callers must test
+	 * sch->ops.cpu_acquire / cpu_release directly.
+	 *
+	 * A CPU is available to the BPF scheduler while its scheduling context
+	 * is an SCX task or the idle task, including core-sched forced idle:
+	 * whether a task can run on a forced idle CPU depends on the task's
+	 * core cookie. A CPU is released while a stop, deadline, RT or fair
+	 * (%SCX_OPS_SWITCH_PARTIAL) task runs on it, and while it's offline.
+	 *
+	 * ops.cpu_acquire() and ops.cpu_release() are called on every
+	 * transition between the two states and only then, so they always
+	 * alternate for a given CPU. Every CPU starts released when the
+	 * scheduler is loaded and ops.cpu_acquire() is called for each online
+	 * CPU which is available at that point. ops.cpu_offline() implies a
+	 * release and neither callback is called for a CPU between
+	 * ops.cpu_offline() and ops.cpu_online(). ops.cpu_acquire() is called
+	 * right after ops.cpu_online() if the CPU is available.
+	 *
+	 * Both callbacks are called with the rq of @cpu locked and may be
+	 * called from a CPU other than @cpu. Don't assume that @cpu equals
+	 * bpf_get_smp_processor_id().
+	 *
+	 * Both callbacks are only effective for the root scheduler: cpu-form
+	 * schedulers can't be sub-schedulers or have sub-schedulers.
 	 */
 
 	/**
@@ -1033,7 +1058,11 @@ struct sched_ext_ops {
 	 * @args: Acquire arguments, see the struct definition.
 	 *
 	 * A CPU that was previously released from the BPF scheduler is now once
-	 * again under its control. Deprecated; use SCX_ENQ_IMMED instead.
+	 * again available to it. Called when the CPU switches to an SCX task
+	 * or the idle task, before ops.running() and ops.update_idle(@cpu,
+	 * %true) for it. ops.dispatch() is called when the CPU is looking for
+	 * a task to run and may thus be called on a CPU which is still
+	 * released.
 	 */
 	void (*cpu_acquire)(s32 cpu, struct scx_cpu_acquire_args *args);
 
@@ -1042,11 +1071,17 @@ struct sched_ext_ops {
 	 * @cpu: The CPU being released by the BPF scheduler.
 	 * @args: Release arguments, see the struct definition.
 	 *
-	 * The specified CPU is no longer under the control of the BPF
-	 * scheduler. This could be because it was preempted by a higher
-	 * priority sched_class, though there may be other reasons as well. The
-	 * caller should consult @args->reason to determine the cause.
-	 * Deprecated; use SCX_ENQ_IMMED instead.
+	 * The specified CPU is no longer available to the BPF scheduler. This
+	 * could be because a task of a higher priority sched_class is about
+	 * to run on it or because the task running on it switched to another
+	 * sched_class. The callee should consult @args->reason to determine
+	 * the cause. When an SCX task is preempted, this is called after the
+	 * task has been put back on the local DSQ, or re-enqueued to the BPF
+	 * scheduler if it's an %SCX_TASK_IMMED task. When the CPU leaves idle,
+	 * this is called after ops.update_idle(@cpu, %false).
+	 *
+	 * scx_bpf_reenqueue_local() can be used to re-enqueue the tasks queued
+	 * on @cpu's local DSQ.
 	 */
 	void (*cpu_release)(s32 cpu, struct scx_cpu_release_args *args);
 };
@@ -1058,7 +1093,7 @@ struct sched_ext_ops {
  * where applicable. Layout up to and including @priv matches sched_ext_ops
  * byte-for-byte (verified by BUILD_BUG_ON checks at scx_init() time) so
  * shared field offsets work for both struct types in bpf_scx_init_member()
- * and bpf_scx_check_member(). The deprecated cpu_acquire/cpu_release
+ * and bpf_scx_check_member(). The cpu-form only cpu_acquire/cpu_release
  * callbacks at the tail of sched_ext_ops are omitted here entirely.
  *
  * Differences from sched_ext_ops:
@@ -1071,7 +1106,7 @@ struct sched_ext_ops {
  *   - cpu_offline      -> cid_offline
  *   - dump_cpu         -> dump_cid
  *   - cgroup_*         -> cpuctl_* (they track the cgroup cpu controller)
- *   - cpu_acquire/cpu_release  -> not present (deprecated in sched_ext_ops)
+ *   - cpu_acquire/cpu_release  -> not present (cpu-form only)
  *
  * BPF schedulers using this type cannot call cpu-form scx_bpf_* kfuncs;
  * use the cid-form variants instead. Enforced at BPF verifier time via
-- 
2.56.0.176.g3575e65682-goog

[RFC PATCH 5/5] selftests/sched_ext: Add cpu_avail tests

[Kuba Piecuch]

Add tests verifying that ops.cpu_acquire() and ops.cpu_release() are
invoked on every change of a CPU's availability to the BPF scheduler,
and only then.

The cpu_avail BPF scheduler tracks the state reported by the callbacks
per CPU, flagging double acquires and releases as well as callbacks
between ops.cpu_offline() and ops.cpu_online(). A sched_switch
tracepoint program checks after every pick that the reported state
matches the incoming task. In full switch mode, a CPU must be released
iff the task about to run has an RT or deadline priority, which covers
the stop task and PI-boosted tasks too. In partial mode, tasks without
the SCHED_EXT policy release it as well. The checker skips a CPU
while userspace hotplugs it, as the kernel releases the CPU before
ops.cpu_offline() updates the BPF side.

cpu_avail runs the following scenarios on top of the invariant checker:

- idle_to_rt: an RT task waking up on an idle CPU releases it, with the
  callback invoked from the idle task, and the CPU is acquired again
  when the RT task sleeps.
- class_change: switching a running SCX task to SCHED_FIFO from another
  CPU releases its CPU remotely with SCX_CPU_PREEMPT_CLASS_CHANGE and
  switching it back acquires the CPU. If the task happened not to be
  running at that moment, e.g. because the stopper preempted it, the
  switch is retried once.
- remote_reenq: same as class_change with tasks queued on the target
  CPU's local DSQ and ops.cpu_release() calling the legacy, synchronous
  scx_bpf_reenqueue_local(). The queued tasks must be re-enqueued.
- pi_boost: an SCX task holding a PI mutex is boosted by an RT waiter on
  another CPU, which releases its CPU until it's deboosted.
- churn: RT tasks, SCX tasks and class changes on all CPUs at once.
- hotplug: no callbacks while the CPU is offline and an acquire after it
  comes back online.
- partial: in partial mode, idle -> fair and a fair task preempting an
  SCX task release the CPU with SCX_CPU_PREEMPT_FAIR, and
  a running SCHED_EXT task switching to SCHED_OTHER releases it with
  SCX_CPU_PREEMPT_CLASS_CHANGE.
- reload: every scheduler instance gets an initial acquire for every
  online CPU.

cpu_avail_core needs SMT and runs SCX tasks with different core-sched
cookies on two siblings. The resulting forced idle must not release the
CPUs. An RT task sharing the other sibling's cookie must release and
re-acquire its CPU.

Also add the cpu_acquire/release() stubs back to the maximal scheduler
now that they are no longer deprecated.

Signed-off-by: Kuba Piecuch <jpiecuch@google.com>
---
 tools/testing/selftests/sched_ext/Makefile    |   2 +
 .../selftests/sched_ext/cpu_avail.bpf.c       | 218 ++++++
 tools/testing/selftests/sched_ext/cpu_avail.c | 717 ++++++++++++++++++
 tools/testing/selftests/sched_ext/cpu_avail.h | 237 ++++++
 .../selftests/sched_ext/cpu_avail_core.c      | 256 +++++++
 .../testing/selftests/sched_ext/maximal.bpf.c |  10 +
 6 files changed, 1440 insertions(+)
 create mode 100644 tools/testing/selftests/sched_ext/cpu_avail.bpf.c
 create mode 100644 tools/testing/selftests/sched_ext/cpu_avail.c
 create mode 100644 tools/testing/selftests/sched_ext/cpu_avail.h
 create mode 100644 tools/testing/selftests/sched_ext/cpu_avail_core.c

diff --git a/tools/testing/selftests/sched_ext/Makefile b/tools/testing/selftests/sched_ext/Makefile
index af919a56a4f4..c0861bd33d3a 100644
--- a/tools/testing/selftests/sched_ext/Makefile
+++ b/tools/testing/selftests/sched_ext/Makefile
@@ -162,6 +162,8 @@ endef
 all_test_bpfprogs := $(foreach prog,$(wildcard *.bpf.c),$(INCLUDE_DIR)/$(patsubst %.c,%.skel.h,$(prog)))
 
 auto-test-targets :=			\
+	cpu_avail			\
+	cpu_avail_core			\
 	create_dsq			\
 	dequeue				\
 	dequeue_iter			\
diff --git a/tools/testing/selftests/sched_ext/cpu_avail.bpf.c b/tools/testing/selftests/sched_ext/cpu_avail.bpf.c
new file mode 100644
index 000000000000..d0052d1a18a0
--- /dev/null
+++ b/tools/testing/selftests/sched_ext/cpu_avail.bpf.c
@@ -0,0 +1,218 @@
+// SPDX-License-Identifier: GPL-2.0
+/*
+ * Track ops.cpu_acquire/release() and verify that they reliably reflect
+ * whether each CPU is available to the BPF scheduler.
+ *
+ * A CPU is available while it runs an SCX task or the idle task. In full
+ * switch mode, every task with a normal priority is an SCX task and a CPU must
+ * be released iff the task running on it has an RT or deadline priority. This
+ * includes the stop task and PI-boosted tasks. In partial mode, only
+ * tasks with a normal priority and the SCHED_EXT policy are SCX tasks.
+ *
+ * The sched_switch tracepoint fires after the pick with the rq locked and
+ * checks that the state reported through the callbacks matches the incoming
+ * task.
+ */
+
+#include <scx/common.bpf.h>
+
+char _license[] SEC("license") = "GPL";
+
+#define MAX_CPUS	1024
+#define NR_REASONS	8
+#define MAX_RT_PRIO	100
+#define SCHED_EXT	7
+
+/* set by userspace */
+const volatile u32 nr_cpus;
+u32 hog_tgid;			/* tasks of this tgid get an infinite slice */
+bool reenq_enabled;		/* cpu_release() of reenq_cpu calls v1 reenqueue */
+s32 reenq_cpu;
+bool check_enabled;		/* enables the sched_switch invariant checker */
+bool partial;			/* cpu_avail_partial_ops is attached */
+bool hp_busy;			/* hp_cpu is being hotplugged, don't check it */
+s32 hp_cpu;
+
+/* state as reported by the callbacks, every CPU starts released */
+bool acquired[MAX_CPUS];
+bool offline[MAX_CPUS];
+/* don't check the invariant until the first acquire after ops.cpu_online() */
+bool online_pending[MAX_CPUS];
+
+u64 nr_acquires[MAX_CPUS];
+u64 nr_releases[MAX_CPUS];
+u64 nr_reasons[MAX_CPUS][NR_REASONS];
+u64 nr_remote_releases[MAX_CPUS];	/* invoked from another CPU */
+u64 nr_idle_releases[MAX_CPUS];		/* invoked from the idle task */
+u64 nr_idle_switches[MAX_CPUS];		/* context switches to idle */
+u64 nr_hog_fair_releases[MAX_CPUS];	/* FAIR releases preempting the hog */
+
+u64 nr_dbl_acquires;
+u64 nr_dbl_releases;
+u64 nr_offline_cbs;			/* between cpu_offline() and cpu_online() */
+u64 nr_checks;
+u64 nr_violations;
+s32 violation_cpu;
+s32 violation_pid;
+s32 violation_prio;
+bool violation_acquired;
+
+/* legacy v1 scx_bpf_reenqueue_local() results */
+u64 nr_reenq_calls;
+u64 nr_reenq_remote;			/* invoked from another CPU */
+u64 nr_reenq_tasks;
+u64 nr_reenq_remote_tasks;		/* tasks re-enqueued by remote calls */
+u64 nr_reenq_enqueues;			/* ops.enqueue() with SCX_ENQ_REENQ */
+
+static __always_inline bool cpu_ok(s32 cpu)
+{
+	return cpu >= 0 && cpu < MAX_CPUS && cpu < nr_cpus;
+}
+
+void BPF_STRUCT_OPS(cpu_avail_enqueue, struct task_struct *p, u64 enq_flags)
+{
+	u64 slice = SCX_SLICE_DFL;
+
+	if (enq_flags & SCX_ENQ_REENQ)
+		__sync_fetch_and_add(&nr_reenq_enqueues, 1);
+
+	if (hog_tgid && p->tgid == hog_tgid)
+		slice = SCX_SLICE_INF;
+
+	/* queue pinned tasks on the local DSQ of the CPU they're enqueued on */
+	if (p->nr_cpus_allowed == 1)
+		scx_bpf_dsq_insert(p, SCX_DSQ_LOCAL, slice, enq_flags);
+	else
+		scx_bpf_dsq_insert(p, SCX_DSQ_GLOBAL, slice, enq_flags);
+}
+
+void BPF_STRUCT_OPS(cpu_avail_cpu_acquire, s32 cpu,
+		    struct scx_cpu_acquire_args *args)
+{
+	if (!cpu_ok(cpu))
+		return;
+
+	if (offline[cpu])
+		__sync_fetch_and_add(&nr_offline_cbs, 1);
+	if (acquired[cpu])
+		__sync_fetch_and_add(&nr_dbl_acquires, 1);
+
+	acquired[cpu] = true;
+	online_pending[cpu] = false;
+	__sync_fetch_and_add(&nr_acquires[cpu], 1);
+}
+
+void BPF_STRUCT_OPS(cpu_avail_cpu_release, s32 cpu,
+		    struct scx_cpu_release_args *args)
+{
+	struct task_struct *cur = bpf_get_current_task_btf();
+	bool remote = bpf_get_smp_processor_id() != cpu;
+	u32 reason = args->reason;
+
+	if (!cpu_ok(cpu))
+		return;
+
+	if (offline[cpu])
+		__sync_fetch_and_add(&nr_offline_cbs, 1);
+	if (!acquired[cpu])
+		__sync_fetch_and_add(&nr_dbl_releases, 1);
+
+	acquired[cpu] = false;
+	__sync_fetch_and_add(&nr_releases[cpu], 1);
+	if (reason < NR_REASONS)
+		__sync_fetch_and_add(&nr_reasons[cpu][reason], 1);
+
+	if (remote)
+		__sync_fetch_and_add(&nr_remote_releases[cpu], 1);
+	else if (cur->pid == 0)
+		__sync_fetch_and_add(&nr_idle_releases[cpu], 1);
+	else if (hog_tgid && cur->tgid == hog_tgid &&
+		 reason == SCX_CPU_PREEMPT_FAIR)
+		__sync_fetch_and_add(&nr_hog_fair_releases[cpu], 1);
+
+	if (reenq_enabled && cpu == reenq_cpu) {
+		u32 nr = scx_bpf_reenqueue_local___v1();
+
+		__sync_fetch_and_add(&nr_reenq_calls, 1);
+		__sync_fetch_and_add(&nr_reenq_tasks, nr);
+		if (remote) {
+			__sync_fetch_and_add(&nr_reenq_remote, 1);
+			__sync_fetch_and_add(&nr_reenq_remote_tasks, nr);
+		}
+	}
+}
+
+void BPF_STRUCT_OPS(cpu_avail_cpu_online, s32 cpu)
+{
+	if (!cpu_ok(cpu))
+		return;
+
+	offline[cpu] = false;
+	online_pending[cpu] = true;
+}
+
+void BPF_STRUCT_OPS(cpu_avail_cpu_offline, s32 cpu)
+{
+	if (!cpu_ok(cpu))
+		return;
+
+	/* ops.cpu_offline() implies release */
+	acquired[cpu] = false;
+	offline[cpu] = true;
+}
+
+SEC("tp_btf/sched_switch")
+int BPF_PROG(cpu_avail_sched_switch, bool preempt, struct task_struct *prev,
+	     struct task_struct *next, unsigned int prev_state)
+{
+	s32 cpu = bpf_get_smp_processor_id();
+	bool avail;
+
+	if (!cpu_ok(cpu))
+		return 0;
+
+	if (next->pid == 0)
+		__sync_fetch_and_add(&nr_idle_switches[cpu], 1);
+
+	if (!check_enabled || offline[cpu] || online_pending[cpu] ||
+	    (hp_busy && cpu == hp_cpu))
+		return 0;
+
+	avail = next->pid == 0 ||
+		(next->prio >= MAX_RT_PRIO &&
+		 (!partial || next->policy == SCHED_EXT));
+
+	__sync_fetch_and_add(&nr_checks, 1);
+	if (acquired[cpu] != avail &&
+	    __sync_fetch_and_add(&nr_violations, 1) == 0) {
+		violation_cpu = cpu;
+		violation_pid = next->pid;
+		violation_prio = next->prio;
+		violation_acquired = acquired[cpu];
+	}
+
+	return 0;
+}
+
+SEC(".struct_ops.link")
+struct sched_ext_ops cpu_avail_ops = {
+	.enqueue		= (void *)cpu_avail_enqueue,
+	.cpu_acquire		= (void *)cpu_avail_cpu_acquire,
+	.cpu_release		= (void *)cpu_avail_cpu_release,
+	.cpu_online		= (void *)cpu_avail_cpu_online,
+	.cpu_offline		= (void *)cpu_avail_cpu_offline,
+	.name			= "cpu_avail",
+	.timeout_ms		= 10000,
+};
+
+SEC(".struct_ops.link")
+struct sched_ext_ops cpu_avail_partial_ops = {
+	.enqueue		= (void *)cpu_avail_enqueue,
+	.cpu_acquire		= (void *)cpu_avail_cpu_acquire,
+	.cpu_release		= (void *)cpu_avail_cpu_release,
+	.cpu_online		= (void *)cpu_avail_cpu_online,
+	.cpu_offline		= (void *)cpu_avail_cpu_offline,
+	.flags			= SCX_OPS_SWITCH_PARTIAL,
+	.name			= "cpu_avail_partial",
+	.timeout_ms		= 10000,
+};
diff --git a/tools/testing/selftests/sched_ext/cpu_avail.c b/tools/testing/selftests/sched_ext/cpu_avail.c
new file mode 100644
index 000000000000..3796f94b05a9
--- /dev/null
+++ b/tools/testing/selftests/sched_ext/cpu_avail.c
@@ -0,0 +1,717 @@
+// SPDX-License-Identifier: GPL-2.0
+/*
+ * Verify that ops.cpu_acquire/release() are invoked on every change of a
+ * CPU's availability to the BPF scheduler, and only then.
+ */
+#define _GNU_SOURCE
+#include <pthread.h>
+#include <stdatomic.h>
+#include "cpu_avail.h"
+
+#define NR_RELOADS	5
+
+struct cpu_avail_ctx {
+	struct cpu_avail *skel;
+	struct bpf_link *link;
+	cpu_set_t orig_mask;
+	int nr_cpus;
+	int ctl_cpu;		/* where the test itself runs */
+	int tgt_cpu;		/* the CPU being tested */
+};
+
+/* poll @cond for up to @ms milliseconds */
+#define WAIT_FOR(cond, ms)						\
+({									\
+	u64 __end = now_ns() + (ms) * 1000000ULL;			\
+	bool __ok;							\
+									\
+	while (!(__ok = (cond)) && now_ns() < __end)			\
+		sleep_ms(1);						\
+	__ok;								\
+})
+
+/*
+ * An RT task waking up on an idle CPU must release it and the CPU must be
+ * acquired again when the RT task goes back to sleep. This never goes
+ * through the SCX class.
+ */
+static enum scx_test_status test_idle_to_rt(struct cpu_avail_ctx *ctx)
+{
+	struct cpu_avail__bss *bss = ctx->skel->bss;
+	struct task_spec spec = {
+		.cpu = ctx->tgt_cpu, .policy = SCHED_FIFO, .prio = 10,
+		.period_us = 5000, .run_us = 500, .duration_ms = 500,
+	};
+	int cpu = ctx->tgt_cpu;
+	u64 rt = bss->nr_reasons[cpu][REASON(RT)];
+	u64 idle = bss->nr_idle_releases[cpu];
+	u64 acq = bss->nr_acquires[cpu];
+	pid_t pid;
+
+	pid = spawn_task(&spec);
+	SCX_FAIL_IF(pid < 0, "Failed to fork (%d)", errno);
+	SCX_FAIL_IF(!wait_task(pid), "RT task failed");
+
+	SCX_GT(bss->nr_reasons[cpu][REASON(RT)], rt);
+	SCX_GT(bss->nr_idle_releases[cpu], idle);
+	SCX_GT(bss->nr_acquires[cpu], acq);
+	SCX_ASSERT(WAIT_FOR(bss->acquired[cpu], 100));
+
+	return SCX_TEST_PASS;
+}
+
+/* @policy is the hog's SCX policy, SCHED_OTHER or SCHED_EXT */
+static pid_t spawn_hog(struct cpu_avail_ctx *ctx, int policy)
+{
+	struct task_spec spec = { .cpu = ctx->tgt_cpu, .policy = policy };
+	int go_fd;
+	pid_t pid;
+
+	pid = fork_task(&spec, &go_fd);
+	if (pid < 0)
+		return pid;
+
+	/* infinite slice so that it stays current until it changes class */
+	ctx->skel->bss->hog_tgid = pid;
+	start_task(go_fd);
+	return pid;
+}
+
+static void stop_hog(struct cpu_avail_ctx *ctx, pid_t pid)
+{
+	kill_task(pid);
+	ctx->skel->bss->hog_tgid = 0;
+}
+
+/*
+ * Switch @hog, which runs on the target CPU, from @scx_policy to @policy from
+ * the control CPU and wait for the resulting remote CLASS_CHANGE release. The
+ * hog is normally current but something else, e.g. the stopper, may have
+ * preempted it at that moment, making the release local. If so, switch it
+ * back and retry once.
+ */
+static bool hog_leave_scx(struct cpu_avail_ctx *ctx, pid_t hog, int scx_policy,
+			  int policy, int prio)
+{
+	struct cpu_avail__bss *bss = ctx->skel->bss;
+	int cpu = ctx->tgt_cpu, i;
+	u64 cc, remote;
+
+	for (i = 0; i < 2; i++) {
+		if (!WAIT_FOR(bss->acquired[cpu], 500)) {
+			SCX_ERR("CPU %d not acquired", cpu);
+			return false;
+		}
+		sleep_ms(50);
+
+		cc = bss->nr_reasons[cpu][REASON(CLASS_CHANGE)];
+		remote = bss->nr_remote_releases[cpu];
+		if (set_policy(hog, policy, prio)) {
+			SCX_ERR("Failed to set policy %d (%d)", policy, errno);
+			return false;
+		}
+		if (WAIT_FOR(bss->nr_reasons[cpu][REASON(CLASS_CHANGE)] > cc, 100) &&
+		    bss->nr_remote_releases[cpu] > remote)
+			return true;
+
+		fprintf(stderr, "No remote CLASS_CHANGE release (%lu -> %lu, %lu -> %lu)%s\n",
+			cc, bss->nr_reasons[cpu][REASON(CLASS_CHANGE)],
+			remote, bss->nr_remote_releases[cpu],
+			i ? "" : ", retrying");
+		if (set_policy(hog, scx_policy, 0)) {
+			SCX_ERR("Failed to set policy %d (%d)", scx_policy, errno);
+			return false;
+		}
+	}
+
+	return false;
+}
+
+/*
+ * The running SCX task (policy @scx_policy) switching to @policy through
+ * sched_setscheduler() from another CPU must release the CPU without a
+ * context switch, and the switch back must acquire it.
+ */
+static enum scx_test_status class_change(struct cpu_avail_ctx *ctx,
+					 int scx_policy, int policy, int prio)
+{
+	struct cpu_avail__bss *bss = ctx->skel->bss;
+	int cpu = ctx->tgt_cpu;
+	enum scx_test_status ret = SCX_TEST_FAIL;
+	u64 acq;
+	pid_t hog;
+
+	hog = spawn_hog(ctx, scx_policy);
+	SCX_FAIL_IF(hog < 0, "Failed to fork (%d)", errno);
+
+	if (!hog_leave_scx(ctx, hog, scx_policy, policy, prio))
+		goto out;
+	sleep_ms(20);
+
+	acq = bss->nr_acquires[cpu];
+	if (set_policy(hog, scx_policy, 0)) {
+		SCX_ERR("Failed to set policy %d (%d)", scx_policy, errno);
+		goto out;
+	}
+	if (!WAIT_FOR(bss->nr_acquires[cpu] > acq && bss->acquired[cpu], 100)) {
+		SCX_ERR("CPU %d not acquired after switching back", cpu);
+		goto out;
+	}
+
+	ret = SCX_TEST_PASS;
+out:
+	stop_hog(ctx, hog);
+	return ret;
+}
+
+static enum scx_test_status test_class_change(struct cpu_avail_ctx *ctx)
+{
+	return class_change(ctx, SCHED_OTHER, SCHED_FIFO, 1);
+}
+
+/*
+ * Same as test_class_change() but with tasks queued on the target CPU's local
+ * DSQ and ops.cpu_release() calling the legacy scx_bpf_reenqueue_local().
+ * The tasks must be re-enqueued from the target CPU's local DSQ even though
+ * ops.cpu_release() runs on another CPU.
+ */
+static enum scx_test_status test_remote_reenq(struct cpu_avail_ctx *ctx)
+{
+	struct cpu_avail__bss *bss = ctx->skel->bss;
+	struct task_spec spec = { .cpu = ctx->tgt_cpu, .policy = SCHED_OTHER };
+	enum scx_test_status ret = SCX_TEST_FAIL;
+	pid_t hog, waiters[2] = { -1, -1 };
+	u64 remote, tasks, enqs;
+	int i;
+
+	hog = spawn_hog(ctx, SCHED_OTHER);
+	SCX_FAIL_IF(hog < 0, "Failed to fork (%d)", errno);
+
+	if (!WAIT_FOR(bss->acquired[ctx->tgt_cpu], 500)) {
+		SCX_ERR("CPU %d not acquired", ctx->tgt_cpu);
+		goto out;
+	}
+	sleep_ms(50);
+
+	/* these queue up behind the hog on the target CPU's local DSQ */
+	for (i = 0; i < 2; i++) {
+		waiters[i] = spawn_task(&spec);
+		if (waiters[i] < 0) {
+			SCX_ERR("Failed to fork (%d)", errno);
+			goto out;
+		}
+	}
+	sleep_ms(100);
+
+	/* a first attempt may release and re-enqueue locally, count remote only */
+	remote = bss->nr_reenq_remote;
+	tasks = bss->nr_reenq_remote_tasks;
+	enqs = bss->nr_reenq_enqueues;
+
+	bss->reenq_cpu = ctx->tgt_cpu;
+	bss->reenq_enabled = true;
+	if (!hog_leave_scx(ctx, hog, SCHED_OTHER, SCHED_FIFO, 1))
+		goto out;
+	bss->reenq_enabled = false;
+
+	if (bss->nr_reenq_remote <= remote ||
+	    bss->nr_reenq_remote_tasks <= tasks ||
+	    bss->nr_reenq_enqueues <= enqs) {
+		SCX_ERR("Bad reenq: remote %lu->%lu tasks %lu->%lu enqs %lu->%lu",
+			remote, bss->nr_reenq_remote, tasks,
+			bss->nr_reenq_remote_tasks, enqs, bss->nr_reenq_enqueues);
+		goto out;
+	}
+
+	ret = SCX_TEST_PASS;
+out:
+	bss->reenq_enabled = false;
+	for (i = 0; i < 2; i++)
+		kill_task(waiters[i]);
+	stop_hog(ctx, hog);
+	return ret;
+}
+
+struct pi_ctx {
+	pthread_mutex_t mutex;
+	atomic_bool pinned;
+	atomic_bool held;
+	int owner_cpu;
+	int waiter_cpu;
+	bool failed;
+};
+
+static void *pi_owner_fn(void *arg)
+{
+	struct pi_ctx *pi = arg;
+
+	if (pin_to_cpu(0, pi->owner_cpu)) {
+		pi->failed = true;
+		atomic_store(&pi->pinned, true);
+		atomic_store(&pi->held, true);
+		return NULL;
+	}
+	atomic_store(&pi->pinned, true);
+
+	pthread_mutex_lock(&pi->mutex);
+	atomic_store(&pi->held, true);
+	spin_us(300 * 1000);
+	pthread_mutex_unlock(&pi->mutex);
+	return NULL;
+}
+
+static void *pi_waiter_fn(void *arg)
+{
+	struct pi_ctx *pi = arg;
+	struct sched_param param = { .sched_priority = 10 };
+
+	if (pin_to_cpu(0, pi->waiter_cpu) ||
+	    pthread_setschedparam(pthread_self(), SCHED_FIFO, &param)) {
+		pi->failed = true;
+		return NULL;
+	}
+
+	while (!atomic_load(&pi->held))
+		;
+	/* boosts the owner, which runs on another CPU */
+	pthread_mutex_lock(&pi->mutex);
+	pthread_mutex_unlock(&pi->mutex);
+	return NULL;
+}
+
+static int run_pi(int owner_cpu, int waiter_cpu)
+{
+	struct pi_ctx pi = { .owner_cpu = owner_cpu, .waiter_cpu = waiter_cpu };
+	pthread_mutexattr_t attr;
+	pthread_t owner, waiter;
+
+	pthread_mutexattr_init(&attr);
+	if (pthread_mutexattr_setprotocol(&attr, PTHREAD_PRIO_INHERIT) ||
+	    pthread_mutex_init(&pi.mutex, &attr))
+		return EXIT_FAILURE;
+
+	if (pthread_create(&owner, NULL, pi_owner_fn, &pi))
+		return EXIT_FAILURE;
+
+	/*
+	 * The threads start on the waiter's CPU. Don't let the busy-waiting RT
+	 * waiter starve the owner before it has moved to its own CPU.
+	 */
+	while (!atomic_load(&pi.pinned))
+		usleep(1000);
+
+	if (pthread_create(&waiter, NULL, pi_waiter_fn, &pi)) {
+		pthread_join(owner, NULL);
+		return EXIT_FAILURE;
+	}
+	pthread_join(waiter, NULL);
+	pthread_join(owner, NULL);
+
+	return pi.failed ? EXIT_FAILURE : EXIT_SUCCESS;
+}
+
+/*
+ * An SCX task holding a PI mutex is boosted to RT by an RT waiter on another
+ * CPU. Its CPU must be released while it's boosted and acquired again once
+ * it's deboosted.
+ */
+static enum scx_test_status test_pi_boost(struct cpu_avail_ctx *ctx)
+{
+	struct cpu_avail__bss *bss = ctx->skel->bss;
+	int cpu = ctx->tgt_cpu;
+	u64 cc = bss->nr_reasons[cpu][REASON(CLASS_CHANGE)];
+	u64 rt = bss->nr_reasons[cpu][REASON(RT)];
+	u64 acq = bss->nr_acquires[cpu];
+	pid_t pid;
+
+	pid = fork();
+	SCX_FAIL_IF(pid < 0, "Failed to fork (%d)", errno);
+	if (!pid)
+		exit(run_pi(ctx->tgt_cpu, ctx->ctl_cpu));
+
+	SCX_FAIL_IF(!wait_task(pid), "PI test failed");
+
+	/* boosting a running owner releases with CLASS_CHANGE, a queued one RT */
+	SCX_GT(bss->nr_reasons[cpu][REASON(CLASS_CHANGE)] - cc +
+	       bss->nr_reasons[cpu][REASON(RT)] - rt, 0);
+	SCX_GT(bss->nr_acquires[cpu], acq);
+	SCX_ASSERT(WAIT_FOR(bss->acquired[cpu], 100));
+
+	return SCX_TEST_PASS;
+}
+
+/*
+ * Run RT tasks, SCX tasks and sched_setscheduler() class changes on all CPUs
+ * at once. The sched_switch checker must not find any inconsistencies.
+ */
+static enum scx_test_status test_churn(struct cpu_avail_ctx *ctx)
+{
+	struct cpu_avail__bss *bss = ctx->skel->bss;
+	struct task_spec spinner = { .cpu = -1, .policy = SCHED_OTHER };
+	struct task_spec rt = {
+		.policy = SCHED_FIFO, .prio = 20,
+		.period_us = 1000, .run_us = 200, .duration_ms = 1500,
+	};
+	enum scx_test_status ret = SCX_TEST_FAIL;
+	int nr_spinners = 0, nr_rts = 0, cpu, i;
+	pid_t *spinners, *rts;
+	u64 checks = bss->nr_checks;
+	u64 end;
+
+	spinners = calloc(ctx->nr_cpus, sizeof(pid_t));
+	rts = calloc(ctx->nr_cpus, sizeof(pid_t));
+	if (!spinners || !rts) {
+		SCX_ERR("Failed to allocate");
+		goto out;
+	}
+
+	for (cpu = 0; cpu < ctx->nr_cpus; cpu++) {
+		if (!cpu_is_online(cpu))
+			continue;
+
+		spinners[nr_spinners] = spawn_task(&spinner);
+		if (spinners[nr_spinners] < 0) {
+			SCX_ERR("Failed to fork (%d)", errno);
+			goto out;
+		}
+		nr_spinners++;
+
+		rt.cpu = cpu;
+		rts[nr_rts] = spawn_task(&rt);
+		if (rts[nr_rts] < 0) {
+			SCX_ERR("Failed to fork (%d)", errno);
+			goto out;
+		}
+		nr_rts++;
+	}
+
+	/* flip a spinner between SCX and RT while the RT tasks run */
+	end = now_ns() + 1500 * 1000000ULL;
+	for (i = 0; now_ns() < end; i++) {
+		if (set_policy(spinners[0], i % 2 ? SCHED_OTHER : SCHED_FIFO,
+			       i % 2 ? 0 : 1)) {
+			SCX_ERR("Failed to set policy (%d)", errno);
+			goto out;
+		}
+		sleep_ms(2);
+	}
+
+	for (i = 0; i < nr_rts; i++) {
+		if (!wait_task(rts[i])) {
+			SCX_ERR("RT task failed");
+			rts[i] = -1;
+			goto out;
+		}
+		rts[i] = -1;
+	}
+
+	if (bss->nr_checks <= checks || bss->nr_violations) {
+		SCX_ERR("checks %lu -> %lu, violations %lu", checks,
+			bss->nr_checks, bss->nr_violations);
+		goto out;
+	}
+
+	ret = SCX_TEST_PASS;
+out:
+	for (i = 0; i < nr_rts; i++)
+		kill_task(rts[i]);
+	for (i = 0; i < nr_spinners; i++)
+		kill_task(spinners[i]);
+	free(spinners);
+	free(rts);
+	return ret;
+}
+
+/*
+ * Neither callback may be invoked between ops.cpu_offline() and
+ * ops.cpu_online(), and the CPU must be acquired after ops.cpu_online().
+ */
+static enum scx_test_status test_hotplug(struct cpu_avail_ctx *ctx)
+{
+	struct cpu_avail__bss *bss = ctx->skel->bss;
+	enum scx_test_status ret = SCX_TEST_FAIL;
+	int cpu = ctx->tgt_cpu;
+	char path[64];
+	u64 acq;
+
+	snprintf(path, sizeof(path), "/sys/devices/system/cpu/cpu%d/online", cpu);
+	if (access(path, W_OK)) {
+		fprintf(stderr, "CPU %d can't be offlined, skipping hotplug\n", cpu);
+		return SCX_TEST_PASS;
+	}
+
+	/*
+	 * The kernel releases the CPU right before ops.cpu_offline(), which
+	 * updates the BPF side without holding the rq lock. Don't check the
+	 * CPU while it's being hotplugged.
+	 */
+	bss->hp_cpu = cpu;
+	bss->hp_busy = true;
+
+	acq = bss->nr_acquires[cpu];
+	if (file_write_long(path, 0)) {
+		SCX_ERR("Failed to offline CPU %d", cpu);
+		goto out;
+	}
+	if (!bss->offline[cpu]) {
+		SCX_ERR("ops.cpu_offline() not called for CPU %d", cpu);
+		goto out;
+	}
+	sleep_ms(100);
+
+	if (file_write_long(path, 1)) {
+		SCX_ERR("Failed to online CPU %d", cpu);
+		goto out;
+	}
+	if (bss->offline[cpu] ||
+	    !WAIT_FOR(bss->nr_acquires[cpu] > acq && bss->acquired[cpu], 500)) {
+		SCX_ERR("CPU %d not acquired after onlining", cpu);
+		goto out;
+	}
+	ret = SCX_TEST_PASS;
+out:
+	bss->hp_busy = false;
+	if (ret != SCX_TEST_PASS)
+		return ret;
+	SCX_EQ(bss->nr_offline_cbs, 0);
+	return SCX_TEST_PASS;
+}
+
+/*
+ * In partial mode, only SCHED_EXT tasks are SCX tasks and fair tasks
+ * take the CPU away from the BPF scheduler: idle -> fair, SCX -> fair
+ * preemption and an SCX task switching to SCHED_OTHER must release the CPU
+ * with SCX_CPU_PREEMPT_FAIR or SCX_CPU_PREEMPT_CLASS_CHANGE.
+ */
+static enum scx_test_status partial_scenarios(struct cpu_avail_ctx *ctx)
+{
+	struct cpu_avail__bss *bss = ctx->skel->bss;
+	struct task_spec fair = {
+		.cpu = ctx->tgt_cpu, .policy = SCHED_OTHER,
+		.period_us = 5000, .run_us = 500, .duration_ms = 300,
+	};
+	int cpu = ctx->tgt_cpu;
+	enum scx_test_status ret;
+	u64 rels, idle, acq;
+	pid_t pid, hog;
+
+	/* idle -> fair -> idle */
+	fprintf(stderr, "cpu_avail: partial idle_to_fair on CPU %d\n", cpu);
+	rels = bss->nr_reasons[cpu][REASON(FAIR)];
+	idle = bss->nr_idle_releases[cpu];
+	acq = bss->nr_acquires[cpu];
+	pid = spawn_task(&fair);
+	SCX_FAIL_IF(pid < 0, "Failed to fork (%d)", errno);
+	SCX_FAIL_IF(!wait_task(pid), "fair task failed");
+	SCX_GT(bss->nr_reasons[cpu][REASON(FAIR)], rels);
+	SCX_GT(bss->nr_idle_releases[cpu], idle);
+	SCX_GT(bss->nr_acquires[cpu], acq);
+	SCX_ASSERT(WAIT_FOR(bss->acquired[cpu], 100));
+
+	/* a fair task preempting an SCX task */
+	fprintf(stderr, "cpu_avail: partial scx_to_fair on CPU %d\n", cpu);
+	hog = spawn_hog(ctx, SCHED_EXT);
+	SCX_FAIL_IF(hog < 0, "Failed to fork (%d)", errno);
+	if (!WAIT_FOR(bss->acquired[cpu], 500)) {
+		stop_hog(ctx, hog);
+		SCX_FAIL("CPU %d not acquired", cpu);
+	}
+	/* only count releases from put_prev_task_scx() of the hog */
+	rels = bss->nr_hog_fair_releases[cpu];
+	acq = bss->nr_acquires[cpu];
+	pid = spawn_task(&fair);
+	if (pid < 0 || !wait_task(pid)) {
+		stop_hog(ctx, hog);
+		SCX_FAIL("fair task failed");
+	}
+	stop_hog(ctx, hog);
+	SCX_GT(bss->nr_hog_fair_releases[cpu], rels);
+	SCX_GT(bss->nr_acquires[cpu], acq);
+
+	/* SCHED_EXT -> SCHED_OTHER on a running task */
+	fprintf(stderr, "cpu_avail: partial class_change on CPU %d\n", cpu);
+	ret = class_change(ctx, SCHED_EXT, SCHED_OTHER, 0);
+	if (ret != SCX_TEST_PASS)
+		return ret;
+
+	return check_consistency(ctx->skel);
+}
+
+static enum scx_test_status test_partial(struct cpu_avail_ctx *ctx)
+{
+	struct cpu_avail *main_skel = ctx->skel;
+	enum scx_test_status ret;
+	struct bpf_link *link;
+
+	ctx->skel = cpu_avail_open_load();
+	if (!ctx->skel) {
+		ctx->skel = main_skel;
+		SCX_FAIL("Failed to open and load skel");
+	}
+
+	ctx->skel->bss->partial = true;
+	link = bpf_map__attach_struct_ops(ctx->skel->maps.cpu_avail_partial_ops);
+	if (!link) {
+		ret = SCX_TEST_FAIL;
+		SCX_ERR("Failed to attach scheduler");
+	} else {
+		ctx->skel->bss->check_enabled = true;
+		ret = partial_scenarios(ctx);
+		bpf_link__destroy(link);
+	}
+
+	cpu_avail__destroy(ctx->skel);
+	ctx->skel = main_skel;
+	return ret;
+}
+
+/*
+ * Every scheduler instance starts with all CPUs released and must see an
+ * initial ops.cpu_acquire() for every online CPU, regardless of what the
+ * previous instance left behind.
+ */
+static enum scx_test_status test_reload(struct cpu_avail_ctx *ctx)
+{
+	enum scx_test_status ret;
+	int i, cpu;
+
+	for (i = 0; i < NR_RELOADS; i++) {
+		struct cpu_avail *skel;
+		struct bpf_link *link;
+
+		skel = cpu_avail_open_load();
+		SCX_FAIL_IF(!skel, "Failed to open and load skel");
+
+		link = bpf_map__attach_struct_ops(skel->maps.cpu_avail_ops);
+		if (!link) {
+			cpu_avail__destroy(skel);
+			SCX_FAIL("Failed to attach scheduler");
+		}
+		skel->bss->check_enabled = true;
+		sleep_ms(100);
+
+		ret = SCX_TEST_PASS;
+		for (cpu = 0; cpu < ctx->nr_cpus; cpu++) {
+			if (cpu_is_online(cpu) && !skel->bss->nr_acquires[cpu]) {
+				SCX_ERR("reload %d: CPU %d not acquired", i, cpu);
+				ret = SCX_TEST_FAIL;
+			}
+		}
+		if (ret == SCX_TEST_PASS)
+			ret = check_consistency(skel);
+
+		bpf_link__destroy(link);
+		cpu_avail__destroy(skel);
+		if (ret != SCX_TEST_PASS)
+			return ret;
+	}
+
+	return SCX_TEST_PASS;
+}
+
+static enum scx_test_status setup(void **ctxp)
+{
+	struct cpu_avail_ctx *ctx;
+	int cpu;
+
+	ctx = calloc(1, sizeof(*ctx));
+	SCX_FAIL_IF(!ctx, "Failed to allocate");
+
+	ctx->nr_cpus = libbpf_num_possible_cpus();
+	ctx->ctl_cpu = -1;
+	ctx->tgt_cpu = -1;
+	for (cpu = 0; cpu < ctx->nr_cpus; cpu++) {
+		if (!cpu_is_online(cpu))
+			continue;
+		if (ctx->ctl_cpu < 0)
+			ctx->ctl_cpu = cpu;
+		else
+			ctx->tgt_cpu = cpu;
+	}
+	if (ctx->tgt_cpu < 0) {
+		fprintf(stderr, "Needs at least 2 online CPUs\n");
+		free(ctx);
+		return SCX_TEST_SKIP;
+	}
+
+	if (sched_getaffinity(0, sizeof(ctx->orig_mask), &ctx->orig_mask) ||
+	    pin_to_cpu(0, ctx->ctl_cpu)) {
+		SCX_ERR("Failed to pin to CPU %d", ctx->ctl_cpu);
+		free(ctx);
+		return SCX_TEST_FAIL;
+	}
+
+	ctx->skel = cpu_avail_open_load();
+	if (!ctx->skel) {
+		sched_setaffinity(0, sizeof(ctx->orig_mask), &ctx->orig_mask);
+		free(ctx);
+		SCX_FAIL("Failed to open and load skel");
+	}
+
+	*ctxp = ctx;
+	return SCX_TEST_PASS;
+}
+
+static enum scx_test_status run(void *arg)
+{
+	struct cpu_avail_ctx *ctx = arg;
+	static const struct {
+		const char *name;
+		enum scx_test_status (*fn)(struct cpu_avail_ctx *ctx);
+	} tests[] = {
+		{ "idle_to_rt", test_idle_to_rt },
+		{ "class_change", test_class_change },
+		{ "remote_reenq", test_remote_reenq },
+		{ "pi_boost", test_pi_boost },
+		{ "churn", test_churn },
+		{ "hotplug", test_hotplug },
+	};
+	enum scx_test_status ret;
+	int i;
+
+	ctx->link = bpf_map__attach_struct_ops(ctx->skel->maps.cpu_avail_ops);
+	SCX_FAIL_IF(!ctx->link, "Failed to attach scheduler");
+	ctx->skel->bss->check_enabled = true;
+
+	for (i = 0; i < ARRAY_SIZE(tests); i++) {
+		fprintf(stderr, "cpu_avail: %s on CPU %d\n", tests[i].name,
+			ctx->tgt_cpu);
+		ret = tests[i].fn(ctx);
+		if (ret != SCX_TEST_PASS)
+			return ret;
+	}
+
+	ret = check_consistency(ctx->skel);
+	if (ret != SCX_TEST_PASS)
+		return ret;
+
+	bpf_link__destroy(ctx->link);
+	ctx->link = NULL;
+	ctx->skel->bss->check_enabled = false;
+
+	ret = test_partial(ctx);
+	if (ret != SCX_TEST_PASS)
+		return ret;
+
+	fprintf(stderr, "cpu_avail: reload\n");
+	return test_reload(ctx);
+}
+
+static void cleanup(void *arg)
+{
+	struct cpu_avail_ctx *ctx = arg;
+
+	bpf_link__destroy(ctx->link);
+	cpu_avail__destroy(ctx->skel);
+	sched_setaffinity(0, sizeof(ctx->orig_mask), &ctx->orig_mask);
+	free(ctx);
+}
+
+struct scx_test cpu_avail = {
+	.name = "cpu_avail",
+	.description = "Verify that ops.cpu_acquire/release() track CPU availability",
+	.setup = setup,
+	.run = run,
+	.cleanup = cleanup,
+};
+REGISTER_SCX_TEST(&cpu_avail)
diff --git a/tools/testing/selftests/sched_ext/cpu_avail.h b/tools/testing/selftests/sched_ext/cpu_avail.h
new file mode 100644
index 000000000000..97659e1061b8
--- /dev/null
+++ b/tools/testing/selftests/sched_ext/cpu_avail.h
@@ -0,0 +1,237 @@
+/* SPDX-License-Identifier: GPL-2.0 */
+/*
+ * Helpers shared by the cpu_avail and cpu_avail_core tests.
+ */
+#ifndef __CPU_AVAIL_H__
+#define __CPU_AVAIL_H__
+
+#include <bpf/bpf.h>
+#include <errno.h>
+#include <sched.h>
+#include <signal.h>
+#include <stdbool.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <sys/wait.h>
+#include <time.h>
+#include <unistd.h>
+#include <scx/common.h>
+#include "cpu_avail.bpf.skel.h"
+#include "scx_test.h"
+#include "util.h"
+
+#ifndef ARRAY_SIZE
+#define ARRAY_SIZE(x)	(sizeof(x) / sizeof((x)[0]))
+#endif
+
+#ifndef SCHED_EXT
+#define SCHED_EXT	7
+#endif
+
+#define REASON(__name)							\
+	__COMPAT_ENUM_OR_ZERO("scx_cpu_preempt_reason", "SCX_CPU_PREEMPT_" #__name)
+
+static inline u64 now_ns(void)
+{
+	struct timespec ts;
+
+	clock_gettime(CLOCK_MONOTONIC, &ts);
+	return ts.tv_sec * 1000000000ULL + ts.tv_nsec;
+}
+
+static inline void spin_us(u64 us)
+{
+	u64 end = now_ns() + us * 1000;
+
+	while (now_ns() < end)
+		;
+}
+
+static inline void sleep_ms(u64 ms)
+{
+	usleep(ms * 1000);
+}
+
+static inline int pin_to_cpu(pid_t pid, int cpu)
+{
+	cpu_set_t mask;
+
+	CPU_ZERO(&mask);
+	CPU_SET(cpu, &mask);
+	return sched_setaffinity(pid, sizeof(mask), &mask);
+}
+
+static inline int set_policy(pid_t pid, int policy, int prio)
+{
+	struct sched_param param = { .sched_priority = prio };
+
+	return sched_setscheduler(pid, policy, &param);
+}
+
+static inline bool cpu_is_online(int cpu)
+{
+	char path[64];
+
+	snprintf(path, sizeof(path), "/sys/devices/system/cpu/cpu%d", cpu);
+	if (access(path, F_OK))
+		return false;
+
+	/* CPUs which can't be offlined don't have the file */
+	snprintf(path, sizeof(path), "/sys/devices/system/cpu/cpu%d/online", cpu);
+	if (access(path, F_OK))
+		return true;
+
+	return file_read_long(path) > 0;
+}
+
+/*
+ * What a forked task does once it's been released. @period_us == 0 means spin
+ * forever. Otherwise, sleep for @period_us - @run_us and then spin for @run_us
+ * in a loop for @duration_ms.
+ */
+struct task_spec {
+	int cpu;		/* -1 for no affinity */
+	int policy;
+	int prio;
+	u64 period_us;
+	u64 run_us;
+	u64 duration_ms;
+};
+
+static inline void task_body(const struct task_spec *spec)
+{
+	u64 end;
+
+	if (!spec->period_us) {
+		while (true)
+			spin_us(1000);
+	}
+
+	end = now_ns() + spec->duration_ms * 1000000ULL;
+	while (now_ns() < end) {
+		usleep(spec->period_us - spec->run_us);
+		spin_us(spec->run_us);
+	}
+}
+
+/*
+ * Fork a task according to @spec. The child blocks until start_task() is
+ * called on the returned pid so that the caller can set things up for it
+ * first. Returns the child's pid, or -1 on failure. @go_fd receives the write
+ * end of the pipe that releases the child.
+ */
+static inline pid_t fork_task(const struct task_spec *spec, int *go_fd)
+{
+	int fds[2];
+	pid_t pid;
+	char c;
+
+	if (pipe(fds))
+		return -1;
+
+	pid = fork();
+	if (pid < 0) {
+		close(fds[0]);
+		close(fds[1]);
+		return -1;
+	}
+
+	if (!pid) {
+		close(fds[1]);
+		if (read(fds[0], &c, 1) != 1)
+			exit(EXIT_FAILURE);
+		close(fds[0]);
+		if (spec->cpu >= 0 && pin_to_cpu(0, spec->cpu))
+			exit(EXIT_FAILURE);
+		if (spec->policy != SCHED_OTHER &&
+		    set_policy(0, spec->policy, spec->prio))
+			exit(EXIT_FAILURE);
+		task_body(spec);
+		exit(EXIT_SUCCESS);
+	}
+
+	close(fds[0]);
+	*go_fd = fds[1];
+	return pid;
+}
+
+static inline void start_task(int go_fd)
+{
+	char c = 1;
+
+	if (write(go_fd, &c, 1) != 1)
+		perror("write");
+	close(go_fd);
+}
+
+static inline pid_t spawn_task(const struct task_spec *spec)
+{
+	int go_fd;
+	pid_t pid;
+
+	pid = fork_task(spec, &go_fd);
+	if (pid > 0)
+		start_task(go_fd);
+	return pid;
+}
+
+static inline bool wait_task(pid_t pid)
+{
+	int status;
+
+	if (waitpid(pid, &status, 0) != pid)
+		return false;
+	return WIFEXITED(status) && WEXITSTATUS(status) == EXIT_SUCCESS;
+}
+
+static inline void kill_task(pid_t pid)
+{
+	if (pid <= 0)
+		return;
+	kill(pid, SIGKILL);
+	waitpid(pid, NULL, 0);
+}
+
+static inline struct cpu_avail *cpu_avail_open_load(void)
+{
+	struct cpu_avail *skel;
+
+	skel = cpu_avail__open();
+	if (!skel)
+		return NULL;
+	SCX_ENUM_INIT(skel);
+	skel->rodata->nr_cpus = libbpf_num_possible_cpus();
+	if (cpu_avail__load(skel))
+		goto err;
+
+	/* attach the sched_switch checker, a struct_ops is attached later */
+	bpf_map__set_autoattach(skel->maps.cpu_avail_ops, false);
+	bpf_map__set_autoattach(skel->maps.cpu_avail_partial_ops, false);
+	if (cpu_avail__attach(skel))
+		goto err;
+
+	return skel;
+err:
+	cpu_avail__destroy(skel);
+	return NULL;
+}
+
+/* common checks on the state accumulated by the BPF side */
+static inline enum scx_test_status check_consistency(struct cpu_avail *skel)
+{
+	struct cpu_avail__bss *bss = skel->bss;
+
+	if (bss->nr_violations)
+		SCX_ERR("first violation: cpu=%d pid=%d prio=%d acquired=%d",
+			bss->violation_cpu, bss->violation_pid,
+			bss->violation_prio, bss->violation_acquired);
+	SCX_EQ(bss->nr_violations, 0);
+	SCX_GT(bss->nr_checks, 0);
+	SCX_EQ(bss->nr_dbl_acquires, 0);
+	SCX_EQ(bss->nr_dbl_releases, 0);
+	SCX_EQ(bss->nr_offline_cbs, 0);
+	return SCX_TEST_PASS;
+}
+
+#endif /* __CPU_AVAIL_H__ */
diff --git a/tools/testing/selftests/sched_ext/cpu_avail_core.c b/tools/testing/selftests/sched_ext/cpu_avail_core.c
new file mode 100644
index 000000000000..971c4d37eb78
--- /dev/null
+++ b/tools/testing/selftests/sched_ext/cpu_avail_core.c
@@ -0,0 +1,256 @@
+// SPDX-License-Identifier: GPL-2.0
+/*
+ * Verify ops.cpu_acquire/release() under core scheduling. A CPU which is
+ * forced idle stays available to the BPF scheduler, while an RT task with a
+ * matching cookie running on it releases it.
+ */
+#define _GNU_SOURCE
+#include <sys/prctl.h>
+#include "cpu_avail.h"
+
+#ifndef PR_SCHED_CORE
+#define PR_SCHED_CORE			62
+#endif
+#ifndef PR_SCHED_CORE_CREATE
+#define PR_SCHED_CORE_CREATE		1
+#endif
+#ifndef PR_SCHED_CORE_SHARE_FROM
+#define PR_SCHED_CORE_SHARE_FROM	3
+#endif
+#ifndef PR_SCHED_CORE_SCOPE_THREAD
+#define PR_SCHED_CORE_SCOPE_THREAD	0
+#endif
+
+struct core_ctx {
+	struct cpu_avail *skel;
+	struct bpf_link *link;
+	cpu_set_t orig_mask;
+	int sib[2];		/* SMT siblings under test */
+	int ctl_cpu;		/* where the test itself runs, not a sibling */
+};
+
+static bool find_siblings(int cpu, int sib[2])
+{
+	char path[128], buf[64];
+	FILE *f;
+	bool ret;
+
+	snprintf(path, sizeof(path),
+		 "/sys/devices/system/cpu/cpu%d/topology/thread_siblings_list", cpu);
+	f = fopen(path, "r");
+	if (!f)
+		return false;
+	ret = fgets(buf, sizeof(buf), f) &&
+	      sscanf(buf, "%d%*[,-]%d", &sib[0], &sib[1]) == 2 &&
+	      sib[0] != sib[1];
+	fclose(f);
+	return ret;
+}
+
+static int core_cookie_create(pid_t pid)
+{
+	return prctl(PR_SCHED_CORE, PR_SCHED_CORE_CREATE, pid,
+		     PR_SCHED_CORE_SCOPE_THREAD, 0);
+}
+
+/* fork a spinner pinned to @cpu with a cookie of its own */
+static pid_t spawn_cookie_spinner(int cpu)
+{
+	struct task_spec spec = { .cpu = cpu, .policy = SCHED_OTHER };
+	int go_fd;
+	pid_t pid;
+
+	pid = fork_task(&spec, &go_fd);
+	if (pid < 0)
+		return pid;
+
+	if (core_cookie_create(pid)) {
+		int err = errno;
+
+		close(go_fd);
+		kill_task(pid);
+		errno = err;
+		return -1;
+	}
+
+	start_task(go_fd);
+	return pid;
+}
+
+/* fork a periodic RT task pinned to @cpu sharing @cookie_pid's cookie */
+static pid_t spawn_cookie_rt(int cpu, pid_t cookie_pid)
+{
+	struct task_spec spec = {
+		.cpu = cpu, .policy = SCHED_FIFO, .prio = 10,
+		.period_us = 5000, .run_us = 1000, .duration_ms = 1000,
+	};
+	pid_t pid;
+
+	pid = fork();
+	if (pid)
+		return pid;
+
+	if (prctl(PR_SCHED_CORE, PR_SCHED_CORE_SHARE_FROM, cookie_pid,
+		  PR_SCHED_CORE_SCOPE_THREAD, 0))
+		exit(EXIT_FAILURE);
+	if (pin_to_cpu(0, spec.cpu) || set_policy(0, spec.policy, spec.prio))
+		exit(EXIT_FAILURE);
+	task_body(&spec);
+	exit(EXIT_SUCCESS);
+}
+
+static enum scx_test_status setup(void **ctxp)
+{
+	struct core_ctx *ctx;
+	int nr_cpus = libbpf_num_possible_cpus();
+	int cpu;
+
+	ctx = calloc(1, sizeof(*ctx));
+	SCX_FAIL_IF(!ctx, "Failed to allocate");
+
+	ctx->ctl_cpu = -1;
+	for (cpu = 0; cpu < nr_cpus; cpu++) {
+		if (cpu_is_online(cpu) && find_siblings(cpu, ctx->sib))
+			break;
+	}
+	if (cpu == nr_cpus || !cpu_is_online(ctx->sib[0]) ||
+	    !cpu_is_online(ctx->sib[1])) {
+		fprintf(stderr, "No online SMT siblings\n");
+		free(ctx);
+		return SCX_TEST_SKIP;
+	}
+
+	for (cpu = 0; cpu < nr_cpus; cpu++) {
+		if (cpu_is_online(cpu) && cpu != ctx->sib[0] && cpu != ctx->sib[1]) {
+			ctx->ctl_cpu = cpu;
+			break;
+		}
+	}
+	if (ctx->ctl_cpu < 0) {
+		fprintf(stderr, "No online CPU outside of the SMT siblings\n");
+		free(ctx);
+		return SCX_TEST_SKIP;
+	}
+
+	/* the test itself shouldn't share the core */
+	if (sched_getaffinity(0, sizeof(ctx->orig_mask), &ctx->orig_mask) ||
+	    pin_to_cpu(0, ctx->ctl_cpu)) {
+		SCX_ERR("Failed to pin to CPU %d", ctx->ctl_cpu);
+		free(ctx);
+		return SCX_TEST_FAIL;
+	}
+
+	ctx->skel = cpu_avail_open_load();
+	if (!ctx->skel) {
+		sched_setaffinity(0, sizeof(ctx->orig_mask), &ctx->orig_mask);
+		free(ctx);
+		SCX_FAIL("Failed to open and load skel");
+	}
+
+	*ctxp = ctx;
+	return SCX_TEST_PASS;
+}
+
+static enum scx_test_status run(void *arg)
+{
+	struct core_ctx *ctx = arg;
+	struct cpu_avail__bss *bss = ctx->skel->bss;
+	enum scx_test_status ret = SCX_TEST_FAIL;
+	pid_t spinners[2] = { -1, -1 }, rt = -1;
+	u64 unknown[2], idle[2], rels[2], rt_rels, acqs;
+	int *sib = ctx->sib, i;
+
+	ctx->link = bpf_map__attach_struct_ops(ctx->skel->maps.cpu_avail_ops);
+	SCX_FAIL_IF(!ctx->link, "Failed to attach scheduler");
+	bss->check_enabled = true;
+
+	fprintf(stderr, "cpu_avail_core: siblings %d,%d\n", sib[0], sib[1]);
+
+	/* SCX tasks with different cookies on the siblings force idle */
+	for (i = 0; i < 2; i++) {
+		spinners[i] = spawn_cookie_spinner(sib[i]);
+		if (spinners[i] < 0) {
+			if (errno == EINVAL || errno == ENODEV) {
+				fprintf(stderr, "Core scheduling not supported (%d)\n",
+					errno);
+				ret = SCX_TEST_SKIP;
+			} else {
+				SCX_ERR("Failed to fork (%d)", errno);
+			}
+			goto out;
+		}
+	}
+	sleep_ms(100);
+
+	for (i = 0; i < 2; i++) {
+		unknown[i] = bss->nr_reasons[sib[i]][REASON(UNKNOWN)];
+		idle[i] = bss->nr_idle_switches[sib[i]];
+		rels[i] = bss->nr_releases[sib[i]];
+	}
+	sleep_ms(1000);
+
+	for (i = 0; i < 2; i++) {
+		fprintf(stderr, "cpu_avail_core: CPU %d: %lu idle switches, %lu releases\n",
+			sib[i], bss->nr_idle_switches[sib[i]] - idle[i],
+			bss->nr_releases[sib[i]] - rels[i]);
+		if (bss->nr_reasons[sib[i]][REASON(UNKNOWN)] != unknown[i]) {
+			SCX_ERR("CPU %d released while forced idle", sib[i]);
+			goto out;
+		}
+	}
+	if (bss->nr_idle_switches[sib[0]] == idle[0] &&
+	    bss->nr_idle_switches[sib[1]] == idle[1]) {
+		SCX_ERR("No forced idle observed");
+		goto out;
+	}
+
+	/* an RT task matching sib[0]'s cookie runs on sib[1] */
+	rt_rels = bss->nr_reasons[sib[1]][REASON(RT)];
+	acqs = bss->nr_acquires[sib[1]];
+	rt = spawn_cookie_rt(sib[1], spinners[0]);
+	if (rt < 0) {
+		SCX_ERR("Failed to fork (%d)", errno);
+		goto out;
+	}
+	if (!wait_task(rt)) {
+		rt = -1;
+		SCX_ERR("RT task failed");
+		goto out;
+	}
+	rt = -1;
+
+	fprintf(stderr, "cpu_avail_core: CPU %d: %lu RT releases, %lu acquires\n",
+		sib[1], bss->nr_reasons[sib[1]][REASON(RT)] - rt_rels,
+		bss->nr_acquires[sib[1]] - acqs);
+	if (bss->nr_reasons[sib[1]][REASON(RT)] == rt_rels ||
+	    bss->nr_acquires[sib[1]] == acqs) {
+		SCX_ERR("CPU %d not released and acquired around the RT task", sib[1]);
+		goto out;
+	}
+
+	ret = check_consistency(ctx->skel);
+out:
+	kill_task(rt);
+	for (i = 0; i < 2; i++)
+		kill_task(spinners[i]);
+	return ret;
+}
+
+static void cleanup(void *arg)
+{
+	struct core_ctx *ctx = arg;
+
+	bpf_link__destroy(ctx->link);
+	cpu_avail__destroy(ctx->skel);
+	sched_setaffinity(0, sizeof(ctx->orig_mask), &ctx->orig_mask);
+	free(ctx);
+}
+
+struct scx_test cpu_avail_core = {
+	.name = "cpu_avail_core",
+	.description = "Verify ops.cpu_acquire/release() under core scheduling",
+	.setup = setup,
+	.run = run,
+	.cleanup = cleanup,
+};
+REGISTER_SCX_TEST(&cpu_avail_core)
diff --git a/tools/testing/selftests/sched_ext/maximal.bpf.c b/tools/testing/selftests/sched_ext/maximal.bpf.c
index 04a369078aac..c2ee5437df65 100644
--- a/tools/testing/selftests/sched_ext/maximal.bpf.c
+++ b/tools/testing/selftests/sched_ext/maximal.bpf.c
@@ -67,6 +67,14 @@ void BPF_STRUCT_OPS(maximal_set_cpumask, struct task_struct *p,
 void BPF_STRUCT_OPS(maximal_update_idle, s32 cpu, bool idle)
 {}
 
+void BPF_STRUCT_OPS(maximal_cpu_acquire, s32 cpu,
+		    struct scx_cpu_acquire_args *args)
+{}
+
+void BPF_STRUCT_OPS(maximal_cpu_release, s32 cpu,
+		    struct scx_cpu_release_args *args)
+{}
+
 SEC("tp_btf/sched_switch")
 int BPF_PROG(maximal_sched_switch, bool preempt, struct task_struct *prev,
 	     struct task_struct *next, unsigned int prev_state)
@@ -149,6 +157,8 @@ struct sched_ext_ops maximal_ops = {
 	.set_weight		= (void *) maximal_set_weight,
 	.set_cpumask		= (void *) maximal_set_cpumask,
 	.update_idle		= (void *) maximal_update_idle,
+	.cpu_acquire		= (void *) maximal_cpu_acquire,
+	.cpu_release		= (void *) maximal_cpu_release,
 	.cpu_online		= (void *) maximal_cpu_online,
 	.cpu_offline		= (void *) maximal_cpu_offline,
 	.init_task		= (void *) maximal_init_task,
-- 
2.56.0.176.g3575e65682-goog

Re: [RFC PATCH sched_ext/for-7.4 0/5] sched_ext: Make ops.cpu_acquire/release() reliable

[Kuba Piecuch]

On Sun, Oct 18, 2026 at 12:14:44PM +0000, Kuba Piecuch wrote:
> - ops.cpu_acquire() moves from before ops.dispatch() to the point where
>   the CPU actually switches to an SCX task or idle, so ops.dispatch()
>   may now run on a released CPU. A speculative acquire before dispatch
>   can't be made reliable without hooks in every class.

I'd like to reconsider this part of the approach, for two reasons.

First, it changes the semantics of ops.cpu_acquire() relative to the
original implementation. Schedulers could so far assume that the CPU
is acquired whenever ops.dispatch() runs. With this patchset they can't.

Second, it loses a case ops.cpu_release() should cover. Say an RT task
blocks, ops.dispatch() inserts T into the local DSQ, the rq lock gets
dropped, another RT task wakes up on the CPU, the SCX pick returns
RETRY_TASK and the RT task wins. That's an RT -> RT switch, so the CPU
never becomes available, no ops.cpu_release() is invoked and T sits on
the local DSQ until the CPU picks from SCX again.

The standard answer to that would be SCX_ENQ_IMMED, but IMMED isn't
bulletproof either. It relies on rq->next_class being raised through
wakeup_preempt() whenever a higher-class task is about to take the CPU,
and there are several ways around that:

- sched_setscheduler() or a PI boost promoting a queued task to RT/DL:
  sched_change_end() doesn't call wakeup_preempt() for non-running
  tasks, switched_to_rt() only does resched_curr() and switched_to_dl()
  calls wakeup_preempt_dl() directly.

- DL replenishment (dl_task_timer()), CFS bandwidth unthrottling
  (partial mode) and RT group unthrottling enqueue and resched_curr()
  without going through wakeup_preempt().

- When the ext dl_server picks an SCX task, next_class becomes ext
  although RT tasks are queued, so rq_is_open() accepts IMMED inserts.
  Once the server runs out of budget, RT takes over without any wakeup.

- With core scheduling, a forced idle pick resets next_class to idle
  while tasks remain on the local DSQ, so a later RT wakeup goes to
  wakeup_preempt_idle() rather than wakeup_preempt_scx().

In all of these, IMMED tasks end up stranded on the local DSQ behind
the higher-class task.

Making next_class exact would mean touching rt, dl, fair and core, and
still wouldn't cover the dl_server and forced idle cases. What I think
would fix it is a single SCX hook in __schedule(), right after the pick:

	next = pick_next_task(rq, &rf);
	rq->next_class = next->sched_class;
	scx_settle_pick(rq, next);

gated on scx_enabled() and next's class being above ext. At that point
the pick is final, the hook runs on the CPU that's going to run next
(core-sched siblings settle in their own __schedule()), and a preempted
prev as well as anything ops.dispatch() inserted are already on the
local DSQ. If rq->scx.nr_immed is non-zero, the hook schedules a
re-enqueue of the IMMED tasks, whatever made the higher-class task
runnable.

The same hook also allows fixing ops.cpu_acquire/release() while
preserving the original semantics. ops.cpu_acquire() stays in
dispatch_one(), before ops.dispatch(), and the hook releases the CPU
if it's still acquired when the pick lands above ext, with the task
actually picked in args->task. That correctly handles every abandoned
pick (RETRY_TASK, the ext dl_server returning NULL, core-sched
restarts), and in the example with an RT task above ops.cpu_release()
gets to re-enqueue T. The release calls v1 adds to put_prev_task_scx()
and put_prev_task_idle() go away. The rest of v1 (class changes,
hotplug, enable/disable) stays.

Tejun, Peter, would a hook like this be acceptable? The core.c
footprint is one line, it's a static branch when SCX isn't enabled, and
the out-of-line call only happens when the pick goes to a class above
SCX. If you're OK with it, I'd send the IMMED fix first, separately
(Fixes: 98d709cba319 ("sched_ext: Implement SCX_ENQ_IMMED")), and
rebase v2 of this series on top of it.

Thanks,
Kuba