History

sched_ext: Split sub-scheduler implementation into sub.c

The sub-scheduler implementation has grown and will continue to expand. Move the sub-scheduler functions from ext.c into a new kernel/sched/ext/sub.c. sub.h holds the prototypes and the !CONFIG_EXT_SUB_SCHED no-op stubs. scx_dispatch_sched() is shared: balance_one() in ext.c and the scx_bpf_sub_dispatch() kfunc in sub.c both call it, and the latter re-enters it as sub-scheduler dispatch nests. It moves into sub.h as a static __always_inline so both callers keep it inlined and per-level stack stays bounded across the recursion. The event macros it uses move to internal.h. No functional change.

Tejun Heo committed 4 months ago

Commit daf8e161 parent 305f86e5 files · +860 −810Browse files

Unified · Split

5 files
  1. kernel/sched/build_policy.c +2 −0
  2. kernel/sched/ext/ext.c +1 −810
  3. kernel/sched/ext/internal.h +28 −0
  4. kernel/sched/ext/sub.c +668 −0
  5. kernel/sched/ext/sub.h +161 −0
kernel/sched/build_policy.c+2−0
@@ -66,10 +66,12 @@ · 65 unmodified lines
6666 # include "ext/cid.h"
6767 # include "ext/arena.h"
6868 # include "ext/idle.h"
69+# include "ext/sub.h"
6970 # include "ext/ext.c"
7071 # include "ext/cid.c"
7172 # include "ext/arena.c"
7273 # include "ext/idle.c"
74+# include "ext/sub.c"
7375 #endif
7476
7577 #include "syscalls.c"
kernel/sched/ext/ext.c+1−810
@@ -19,6 +19,7 @@ · 18 unmodified lines
1919 #include "cid.h"
2020 #include "arena.h"
2121 #include "idle.h"
22+#include "sub.h"
2223
2324 DEFINE_RAW_SPINLOCK(scx_sched_lock);
2425
@@ -271,58 +272,6 @@ static bool u32_before(u32 a, u32 b) · 246 unmodified lines
271272 return (s32)(a - b) < 0;
272273 }
273274
274−#ifdef CONFIG_EXT_SUB_SCHED
275−/**
276− * scx_next_descendant_pre - find the next descendant for pre-order walk
277− * @pos: the current position (%NULL to initiate traversal)
278− * @root: sched whose descendants to walk
279− *
280− * To be used by scx_for_each_descendant_pre(). Find the next descendant to
281− * visit for pre-order traversal of @root's descendants. @root is included in
282− * the iteration and the first node to be visited.
283− */
284−static struct scx_sched *scx_next_descendant_pre(struct scx_sched *pos,
285− struct scx_sched *root)
286−{
287− struct scx_sched *next;
288−
289− lockdep_assert(lockdep_is_held(&scx_enable_mutex) ||
290− lockdep_is_held(&scx_sched_lock));
291−
292− /* if first iteration, visit @root */
293− if (!pos)
294− return root;
295−
296− /* visit the first child if exists */
297− next = list_first_entry_or_null(&pos->children, struct scx_sched, sibling);
298− if (next)
299− return next;
300−
301− /* no child, visit my or the closest ancestor's next sibling */
302− while (pos != root) {
303− if (!list_is_last(&pos->sibling, &scx_parent(pos)->children))
304− return list_next_entry(pos, sibling);
305− pos = scx_parent(pos);
306− }
307−
308− return NULL;
309−}
310−
311−static struct scx_sched *scx_find_sub_sched(u64 cgroup_id)
312−{
313− return rhashtable_lookup(&scx_sched_hash, &cgroup_id,
314− scx_sched_hash_params);
315−}
316−
317−static void scx_set_task_sched(struct task_struct *p, struct scx_sched *sch)
318−{
319− rcu_assign_pointer(p->scx.sched, sch);
320−}
321−#else /* CONFIG_EXT_SUB_SCHED */
322−static inline struct scx_sched *scx_next_descendant_pre(struct scx_sched *pos, struct scx_sched *root) { return pos ? NULL : root; }
323−static inline void scx_set_task_sched(struct task_struct *p, struct scx_sched *sch) {}
324−#endif /* CONFIG_EXT_SUB_SCHED */
325−
326275 /**
327276 * scx_is_descendant - Test whether sched is a descendant
328277 * @sch: sched to test
@@ -337,19 +286,6 @@ static bool scx_is_descendant(struct scx_sched *sch, struct scx_sched *ancestor) · 8 unmodified lines
337286 return sch->ancestors[ancestor->level] == ancestor;
338287 }
339288
340−/**
341− * scx_for_each_descendant_pre - pre-order walk of a sched's descendants
342− * @pos: iteration cursor
343− * @root: sched to walk the descendants of
344− *
345− * Walk @root's descendants. @root is included in the iteration and the first
346− * node to be visited. Must be called with either scx_enable_mutex or
347− * scx_sched_lock held.
348− */
349−#define scx_for_each_descendant_pre(pos, root) \
350− for ((pos) = scx_next_descendant_pre(NULL, (root)); (pos); \
351− (pos) = scx_next_descendant_pre((pos), (root)))
352−
353289 static struct scx_dispatch_q *find_global_dsq(struct scx_sched *sch, s32 cpu)
354290 {
355291 return &sch->pnode[cpu_to_node(cpu)]->global_dsq;
@@ -935,32 +871,6 @@ struct task_struct *scx_task_iter_next_locked(struct scx_task_iter *iter) · 579 unmodified lines
935871 return NULL;
936872 }
937873
938−/**
939− * scx_add_event - Increase an event counter for 'name' by 'cnt'
940− * @sch: scx_sched to account events for
941− * @name: an event name defined in struct scx_event_stats
942− * @cnt: the number of the event occurred
943− *
944− * This can be used when preemption is not disabled.
945− */
946−#define scx_add_event(sch, name, cnt) do { \
947− this_cpu_add((sch)->pcpu->event_stats.name, (cnt)); \
948− trace_sched_ext_event(#name, (cnt)); \
949−} while(0)
950−
951−/**
952− * __scx_add_event - Increase an event counter for 'name' by 'cnt'
953− * @sch: scx_sched to account events for
954− * @name: an event name defined in struct scx_event_stats
955− * @cnt: the number of the event occurred
956− *
957− * This should be used only when preemption is disabled.
958− */
959−#define __scx_add_event(sch, name, cnt) do { \
960− __this_cpu_add((sch)->pcpu->event_stats.name, (cnt)); \
961− trace_sched_ext_event(#name, cnt); \
962−} while(0)
963−
964874 /**
965875 * scx_dump_event - Dump an event 'kind' in 'events' to 's'
966876 * @s: output seq_buf
@@ -2681,115 +2591,6 @@ static inline void maybe_queue_balance_callback(struct rq *rq) · 1714 unmodified lines
26812591 rq->scx.flags &= ~SCX_RQ_BAL_CB_PENDING;
26822592 }
26832593
2684−/*
2685− * One user of this function is scx_bpf_dispatch() which can be called
2686− * recursively as sub-sched dispatches nest. Always inline to reduce stack usage
2687− * from the call frame.
2688− */
2689−static __always_inline bool
2690−scx_dispatch_sched(struct scx_sched *sch, struct rq *rq,
2691− struct task_struct *prev, bool nested)
2692−{
2693− struct scx_dsp_ctx *dspc = &this_cpu_ptr(sch->pcpu)->dsp_ctx;
2694− int nr_loops = SCX_DSP_MAX_LOOPS;
2695− s32 cpu = cpu_of(rq);
2696− bool prev_on_sch = (prev->sched_class == &ext_sched_class) &&
2697− scx_task_on_sched(sch, prev);
2698−
2699− if (scx_consume_global_dsq(sch, rq))
2700− return true;
2701−
2702− if (scx_bypass_dsp_enabled(sch)) {
2703− /* if @sch is bypassing, only the bypass DSQs are active */
2704− if (scx_bypassing(sch, cpu))
2705− return scx_consume_dispatch_q(sch, rq, scx_bypass_dsq(sch, cpu), 0);
2706−
2707−#ifdef CONFIG_EXT_SUB_SCHED
2708− /*
2709− * If @sch isn't bypassing but its children are, @sch is
2710− * responsible for making forward progress for both its own
2711− * tasks that aren't bypassing and the bypassing descendants'
2712− * tasks. The following implements a simple built-in behavior -
2713− * let each CPU try to run the bypass DSQ every Nth time.
2714− *
2715− * Later, if necessary, we can add an ops flag to suppress the
2716− * auto-consumption and a kfunc to consume the bypass DSQ and,
2717− * so that the BPF scheduler can fully control scheduling of
2718− * bypassed tasks.
2719− */
2720− struct scx_sched_pcpu *pcpu = per_cpu_ptr(sch->pcpu, cpu);
2721−
2722− if (!(pcpu->bypass_host_seq++ % SCX_BYPASS_HOST_NTH) &&
2723− scx_consume_dispatch_q(sch, rq, scx_bypass_dsq(sch, cpu), 0)) {
2724− __scx_add_event(sch, SCX_EV_SUB_BYPASS_DISPATCH, 1);
2725− return true;
2726− }
2727−#endif /* CONFIG_EXT_SUB_SCHED */
2728− }
2729−
2730− if (unlikely(!SCX_HAS_OP(sch, dispatch)) || !scx_rq_online(rq))
2731− return false;
2732−
2733− dspc->rq = rq;
2734−
2735− /*
2736− * The dispatch loop. Because scx_flush_dispatch_buf() may drop the rq
2737− * lock, the local DSQ might still end up empty after a successful
2738− * ops.dispatch(). If the local DSQ is empty even after ops.dispatch()
2739− * produced some tasks, retry. The BPF scheduler may depend on this
2740− * looping behavior to simplify its implementation.
2741− */
2742− do {
2743− dspc->nr_tasks = 0;
2744−
2745− if (nested) {
2746− SCX_CALL_OP(sch, dispatch, rq, scx_cpu_arg(cpu),
2747− prev_on_sch ? prev : NULL);
2748− } else {
2749− /* stash @prev so that nested invocations can access it */
2750− rq->scx.sub_dispatch_prev = prev;
2751− SCX_CALL_OP(sch, dispatch, rq, scx_cpu_arg(cpu),
2752− prev_on_sch ? prev : NULL);
2753− rq->scx.sub_dispatch_prev = NULL;
2754− }
2755−
2756− scx_flush_dispatch_buf(sch, rq);
2757−
2758− if ((prev->scx.flags & SCX_TASK_QUEUED) && prev->scx.slice) {
2759− rq->scx.flags |= SCX_RQ_BAL_KEEP;
2760− return true;
2761− }
2762− if (rq->scx.local_dsq.nr)
2763− return true;
2764− if (scx_consume_global_dsq(sch, rq))
2765− return true;
2766−
2767− /*
2768− * ops.dispatch() can trap us in this loop by repeatedly
2769− * dispatching ineligible tasks. Break out once in a while to
2770− * allow the watchdog to run. As IRQ can't be enabled in
2771− * balance(), we want to complete this scheduling cycle and then
2772− * start a new one. IOW, we want to call resched_curr() on the
2773− * next, most likely idle, task, not the current one. Use
2774− * __scx_bpf_kick_cpu() for deferred kicking.
2775− */
2776− if (unlikely(!--nr_loops)) {
2777− scx_kick_cpu(sch, cpu, 0);
2778− break;
2779− }
2780− } while (dspc->nr_tasks);
2781−
2782− /*
2783− * Prevent the CPU from going idle while bypassed descendants have tasks
2784− * queued. Without this fallback, bypassed tasks could stall if the host
2785− * scheduler's ops.dispatch() doesn't yield any tasks.
2786− */
2787− if (scx_bypass_dsp_enabled(sch))
2788− return scx_consume_dispatch_q(sch, rq, scx_bypass_dsq(sch, cpu), 0);
2789−
2790− return false;
2791−}
2792−
27932594 static int balance_one(struct rq *rq, struct task_struct *prev)
27942595 {
27952596 struct scx_sched *sch = scx_root;
@@ -4469,26 +4270,6 @@ static inline void scx_cgroup_lock(void) {} · 1673 unmodified lines
44694270 static inline void scx_cgroup_unlock(void) {}
44704271 #endif /* CONFIG_EXT_GROUP_SCHED || CONFIG_EXT_SUB_SCHED */
44714272
4472−#ifdef CONFIG_EXT_SUB_SCHED
4473−static struct cgroup *sch_cgroup(struct scx_sched *sch)
4474−{
4475− return sch->cgrp;
4476−}
4477−
4478−/* for each descendant of @cgrp including self, set ->scx_sched to @sch */
4479−static void set_cgroup_sched(struct cgroup *cgrp, struct scx_sched *sch)
4480−{
4481− struct cgroup *pos;
4482− struct cgroup_subsys_state *css;
4483−
4484− cgroup_for_each_live_descendant_pre(pos, css, cgrp)
4485− rcu_assign_pointer(pos->scx_sched, sch);
4486−}
4487−#else /* CONFIG_EXT_SUB_SCHED */
4488−static inline struct cgroup *sch_cgroup(struct scx_sched *sch) { return NULL; }
4489−static inline void set_cgroup_sched(struct cgroup *cgrp, struct scx_sched *sch) {}
4490−#endif /* CONFIG_EXT_SUB_SCHED */
4491−
44924273 /*
44934274 * Omitted operations:
44944275 *
@@ -5765,202 +5546,6 @@ void scx_log_sched_disable(struct scx_sched *sch) · 1270 unmodified lines
57655546 }
57665547 }
57675548
5768−#ifdef CONFIG_EXT_SUB_SCHED
5769−static DECLARE_WAIT_QUEUE_HEAD(scx_unlink_waitq);
5770−
5771−static void drain_descendants(struct scx_sched *sch)
5772−{
5773− /*
5774− * Child scheds that finished the critical part of disabling will take
5775− * themselves off @sch->children. Wait for it to drain. As propagation
5776− * is recursive, empty @sch->children means that all proper descendant
5777− * scheds reached unlinking stage.
5778− */
5779− wait_event(scx_unlink_waitq, list_empty(&sch->children));
5780−}
5781−
5782−static void scx_fail_parent(struct scx_sched *sch,
5783− struct task_struct *failed, s32 fail_code)
5784−{
5785− struct scx_sched *parent = scx_parent(sch);
5786− struct scx_task_iter sti;
5787− struct task_struct *p;
5788−
5789− scx_error(parent, "ops.init_task() failed (%d) for %s[%d] while disabling a sub-scheduler",
5790− fail_code, failed->comm, failed->pid);
5791−
5792− /*
5793− * Once $parent is bypassed, it's safe to put SCX_TASK_NONE tasks into
5794− * it. This may cause downstream failures on the BPF side but $parent is
5795− * dying anyway.
5796− */
5797− scx_bypass(parent, true);
5798−
5799− scx_task_iter_start(&sti, sch->cgrp);
5800− while ((p = scx_task_iter_next_locked(&sti))) {
5801− if (scx_task_on_sched(parent, p))
5802− continue;
5803−
5804− scoped_guard (sched_change, p, DEQUEUE_SAVE | DEQUEUE_MOVE) {
5805− scx_disable_and_exit_task(sch, p);
5806− scx_set_task_sched(p, parent);
5807− }
5808− }
5809− scx_task_iter_stop(&sti);
5810−}
5811−
5812−static void scx_sub_disable(struct scx_sched *sch)
5813−{
5814− struct scx_sched *parent = scx_parent(sch);
5815− struct scx_task_iter sti;
5816− struct task_struct *p;
5817− int ret;
5818−
5819− /*
5820− * Guarantee forward progress and wait for descendants to be disabled.
5821− * To limit disruptions, $parent is not bypassed. Tasks are fully
5822− * prepped and then inserted back into $parent.
5823− */
5824− scx_bypass(sch, true);
5825− drain_descendants(sch);
5826−
5827− /*
5828− * Here, every runnable task is guaranteed to make forward progress and
5829− * we can safely use blocking synchronization constructs. Actually
5830− * disable ops.
5831− */
5832− mutex_lock(&scx_enable_mutex);
5833− percpu_down_write(&scx_fork_rwsem);
5834− scx_cgroup_lock();
5835−
5836− set_cgroup_sched(sch_cgroup(sch), parent);
5837−
5838− scx_task_iter_start(&sti, sch->cgrp);
5839− while ((p = scx_task_iter_next_locked(&sti))) {
5840− struct rq *rq;
5841− struct rq_flags rf;
5842−
5843− /* filter out duplicate visits */
5844− if (scx_task_on_sched(parent, p))
5845− continue;
5846−
5847− /*
5848− * By the time control reaches here, all descendant schedulers
5849− * should already have been disabled.
5850− */
5851− WARN_ON_ONCE(!scx_task_on_sched(sch, p));
5852−
5853− /*
5854− * @p is pinned by the iter: css_task_iter_next() takes a
5855− * reference and holds it until the next iter_next() call, so
5856− * @p->usage is guaranteed > 0.
5857− */
5858− get_task_struct(p);
5859−
5860− scx_task_iter_unlock(&sti);
5861−
5862− /*
5863− * $p is READY or ENABLED on @sch. Initialize for $parent,
5864− * disable and exit from @sch, and then switch over to $parent.
5865− *
5866− * If a task fails to initialize for $parent, the only available
5867− * action is disabling $parent too. While this allows disabling
5868− * of a child sched to cause the parent scheduler to fail, the
5869− * failure can only originate from ops.init_task() of the
5870− * parent. A child can't directly affect the parent through its
5871− * own failures.
5872− */
5873− ret = __scx_init_task(parent, p, false);
5874− if (ret) {
5875− scx_fail_parent(sch, p, ret);
5876− put_task_struct(p);
5877− break;
5878− }
5879−
5880− rq = task_rq_lock(p, &rf);
5881−
5882− if (scx_get_task_state(p) == SCX_TASK_DEAD) {
5883− /*
5884− * sched_ext_dead() raced us between __scx_init_task()
5885− * and this rq lock and ran exit_task() on @sch (the
5886− * sched @p was on at that point), not on $parent.
5887− * $parent's just-completed init is owed an exit_task()
5888− * and we issue it here.
5889− */
5890− scx_sub_init_cancel_task(parent, p);
5891− task_rq_unlock(rq, p, &rf);
5892− put_task_struct(p);
5893− continue;
5894− }
5895−
5896− scoped_guard (sched_change, p, DEQUEUE_SAVE | DEQUEUE_MOVE) {
5897− /*
5898− * $p is initialized for $parent and still attached to
5899− * @sch. Disable and exit for @sch, switch over to
5900− * $parent, override the state to READY to account for
5901− * $p having already been initialized, and then enable.
5902− */
5903− scx_disable_and_exit_task(sch, p);
5904− scx_set_task_state(p, SCX_TASK_INIT_BEGIN);
5905− scx_set_task_state(p, SCX_TASK_INIT);
5906− scx_set_task_sched(p, parent);
5907− scx_set_task_state(p, SCX_TASK_READY);
5908− scx_enable_task(parent, p);
5909− }
5910−
5911− task_rq_unlock(rq, p, &rf);
5912− put_task_struct(p);
5913− }
5914− scx_task_iter_stop(&sti);
5915−
5916− scx_disable_dump(sch);
5917−
5918− scx_cgroup_unlock();
5919− percpu_up_write(&scx_fork_rwsem);
5920−
5921− /*
5922− * All tasks are moved off of @sch but there may still be on-going
5923− * operations (e.g. ops.select_cpu()). Drain them by flushing RCU. Use
5924− * the expedited version as ancestors may be waiting in bypass mode.
5925− * Also, tell the parent that there is no need to keep running bypass
5926− * DSQs for us.
5927− */
5928− synchronize_rcu_expedited();
5929− scx_disable_bypass_dsp(sch);
5930−
5931− scx_unlink_sched(sch);
5932−
5933− mutex_unlock(&scx_enable_mutex);
5934−
5935− /*
5936− * @sch is now unlinked from the parent's children list. Notify and call
5937− * ops.sub_detach/exit(). Note that ops.sub_detach/exit() must be called
5938− * after unlinking and releasing all locks. See scx_claim_exit().
5939− */
5940− wake_up_all(&scx_unlink_waitq);
5941−
5942− if (parent->ops.sub_detach && sch->sub_attached) {
5943− struct scx_sub_detach_args sub_detach_args = {
5944− .ops = &sch->ops,
5945− .cgroup_path = sch->cgrp_path,
5946− };
5947− SCX_CALL_OP(parent, sub_detach, NULL,
5948− &sub_detach_args);
5949− }
5950−
5951− scx_log_sched_disable(sch);
5952−
5953− if (sch->ops.exit)
5954− SCX_CALL_OP(sch, exit, NULL, sch->exit_info);
5955− if (sch->sub_kset)
5956− kobject_del(&sch->sub_kset->kobj);
5957− kobject_del(&sch->kobj);
5958−}
5959−#else /* CONFIG_EXT_SUB_SCHED */
5960−static inline void drain_descendants(struct scx_sched *sch) { }
5961−static inline void scx_sub_disable(struct scx_sched *sch) { }
5962−#endif /* CONFIG_EXT_SUB_SCHED */
5963−
59645549 static void scx_root_disable(struct scx_sched *sch)
59655550 {
59665551 struct scx_task_iter sti;
@@ -7350,347 +6935,6 @@ err_disable: · 1383 unmodified lines
73506935 cmd->ret = 0;
73516936 }
73526937
7353−#ifdef CONFIG_EXT_SUB_SCHED
7354−/* verify that a scheduler can be attached to @cgrp and return the parent */
7355−static struct scx_sched *find_parent_sched(struct cgroup *cgrp)
7356−{
7357− struct scx_sched *parent = cgrp->scx_sched;
7358− struct scx_sched *pos;
7359−
7360− lockdep_assert_held(&scx_sched_lock);
7361−
7362− /* can't attach twice to the same cgroup */
7363− if (parent->cgrp == cgrp)
7364− return ERR_PTR(-EBUSY);
7365−
7366− /* does $parent allow sub-scheds? */
7367− if (!parent->ops.sub_attach)
7368− return ERR_PTR(-EOPNOTSUPP);
7369−
7370− /* can't insert between $parent and its exiting children */
7371− list_for_each_entry(pos, &parent->children, sibling)
7372− if (cgroup_is_descendant(pos->cgrp, cgrp))
7373− return ERR_PTR(-EBUSY);
7374−
7375− return parent;
7376−}
7377−
7378−static bool assert_task_ready_or_enabled(struct task_struct *p)
7379−{
7380− u32 state = scx_get_task_state(p);
7381−
7382− switch (state) {
7383− case SCX_TASK_READY:
7384− case SCX_TASK_ENABLED:
7385− return true;
7386− default:
7387− WARN_ONCE(true, "sched_ext: Invalid task state %d for %s[%d] during enabling sub sched",
7388− state, p->comm, p->pid);
7389− return false;
7390− }
7391−}
7392−
7393−static void scx_sub_enable_workfn(struct kthread_work *work)
7394−{
7395− struct scx_enable_cmd *cmd = container_of(work, struct scx_enable_cmd, work);
7396− struct sched_ext_ops *ops = cmd->ops;
7397− struct cgroup *cgrp;
7398− struct scx_sched *parent, *sch;
7399− struct scx_task_iter sti;
7400− struct task_struct *p;
7401− s32 i, ret;
7402−
7403− mutex_lock(&scx_enable_mutex);
7404−
7405− if (!scx_enabled()) {
7406− ret = -ENODEV;
7407− goto out_unlock;
7408− }
7409−
7410− /* See scx_root_enable_workfn() for the @ops->priv check. */
7411− if (rcu_access_pointer(ops->priv)) {
7412− ret = -EBUSY;
7413− goto out_unlock;
7414− }
7415−
7416− cgrp = cgroup_get_from_id(ops->sub_cgroup_id);
7417− if (IS_ERR(cgrp)) {
7418− ret = PTR_ERR(cgrp);
7419− goto out_unlock;
7420− }
7421−
7422− raw_spin_lock_irq(&scx_sched_lock);
7423− parent = find_parent_sched(cgrp);
7424− if (IS_ERR(parent)) {
7425− raw_spin_unlock_irq(&scx_sched_lock);
7426− ret = PTR_ERR(parent);
7427− goto out_put_cgrp;
7428− }
7429− kobject_get(&parent->kobj);
7430− raw_spin_unlock_irq(&scx_sched_lock);
7431−
7432− /* scx_alloc_and_add_sched() consumes @cgrp whether it succeeds or not */
7433− sch = scx_alloc_and_add_sched(cmd, cgrp, parent);
7434− kobject_put(&parent->kobj);
7435− if (IS_ERR(sch)) {
7436− ret = PTR_ERR(sch);
7437− goto out_unlock;
7438− }
7439−
7440− ret = scx_link_sched(sch);
7441− if (ret)
7442− goto err_disable;
7443−
7444− if (sch->level >= SCX_SUB_MAX_DEPTH) {
7445− scx_error(sch, "max nesting depth %d violated",
7446− SCX_SUB_MAX_DEPTH);
7447− goto err_disable;
7448− }
7449−
7450− if (sch->ops.init) {
7451− ret = SCX_CALL_OP_RET(sch, init, NULL);
7452− if (ret) {
7453− ret = scx_ops_sanitize_err(sch, "init", ret);
7454− scx_error(sch, "ops.init() failed (%d)", ret);
7455− goto err_disable;
7456− }
7457− sch->exit_info->flags |= SCX_EFLAG_INITIALIZED;
7458− }
7459−
7460− ret = scx_arena_pool_init(sch);
7461− if (ret)
7462− goto err_disable;
7463−
7464− ret = scx_set_cmask_scratch_alloc(sch);
7465− if (ret)
7466− goto err_disable;
7467−
7468− if (scx_validate_ops(sch, ops))
7469− goto err_disable;
7470−
7471− struct scx_sub_attach_args sub_attach_args = {
7472− .ops = &sch->ops,
7473− .cgroup_path = sch->cgrp_path,
7474− };
7475−
7476− ret = SCX_CALL_OP_RET(parent, sub_attach, NULL,
7477− &sub_attach_args);
7478− if (ret) {
7479− ret = scx_ops_sanitize_err(sch, "sub_attach", ret);
7480− scx_error(sch, "parent rejected (%d)", ret);
7481− goto err_disable;
7482− }
7483− sch->sub_attached = true;
7484−
7485− scx_bypass(sch, true);
7486−
7487− for (i = SCX_OPI_BEGIN; i < SCX_OPI_END; i++)
7488− if (((void (**)(void))ops)[i])
7489− set_bit(i, sch->has_op);
7490−
7491− percpu_down_write(&scx_fork_rwsem);
7492− scx_cgroup_lock();
7493−
7494− /*
7495− * Set cgroup->scx_sched's and check CSS_ONLINE. Either we see
7496− * !CSS_ONLINE or scx_cgroup_lifetime_notify() sees and shoots us down.
7497− */
7498− set_cgroup_sched(sch_cgroup(sch), sch);
7499− if (!(cgrp->self.flags & CSS_ONLINE)) {
7500− scx_error(sch, "cgroup is not online");
7501− goto err_unlock_and_disable;
7502− }
7503−
7504− /*
7505− * Initialize tasks for the new child $sch without exiting them for
7506− * $parent so that the tasks can always be reverted back to $parent
7507− * sched on child init failure.
7508− */
7509− WARN_ON_ONCE(scx_enabling_sub_sched);
7510− scx_enabling_sub_sched = sch;
7511−
7512− scx_task_iter_start(&sti, sch->cgrp);
7513− while ((p = scx_task_iter_next_locked(&sti))) {
7514− struct rq *rq;
7515− struct rq_flags rf;
7516−
7517− /*
7518− * Task iteration may visit the same task twice when racing
7519− * against exiting. Use %SCX_TASK_SUB_INIT to mark tasks which
7520− * finished __scx_init_task() and skip if set.
7521− *
7522− * A task may exit and get freed between __scx_init_task()
7523− * completion and scx_enable_task(). In such cases,
7524− * scx_disable_and_exit_task() must exit the task for both the
7525− * parent and child scheds.
7526− */
7527− if (p->scx.flags & SCX_TASK_SUB_INIT)
7528− continue;
7529−
7530− /* @p is pinned by the iter; see scx_sub_disable() */
7531− get_task_struct(p);
7532−
7533− if (!assert_task_ready_or_enabled(p)) {
7534− ret = -EINVAL;
7535− goto abort;
7536− }
7537−
7538− scx_task_iter_unlock(&sti);
7539−
7540− /*
7541− * As $p is still on $parent, it can't be transitioned to INIT.
7542− * Let's worry about task state later. Use __scx_init_task().
7543− */
7544− ret = __scx_init_task(sch, p, false);
7545− if (ret)
7546− goto abort;
7547−
7548− rq = task_rq_lock(p, &rf);
7549−
7550− if (scx_get_task_state(p) == SCX_TASK_DEAD) {
7551− /*
7552− * sched_ext_dead() raced us between __scx_init_task()
7553− * and this rq lock and ran exit_task() on $parent (the
7554− * sched @p was on at that point), not on @sch. @sch's
7555− * just-completed init is owed an exit_task() and we
7556− * issue it here.
7557− */
7558− scx_sub_init_cancel_task(sch, p);
7559− task_rq_unlock(rq, p, &rf);
7560− put_task_struct(p);
7561− continue;
7562− }
7563−
7564− p->scx.flags |= SCX_TASK_SUB_INIT;
7565− task_rq_unlock(rq, p, &rf);
7566−
7567− put_task_struct(p);
7568− }
7569− scx_task_iter_stop(&sti);
7570−
7571− /*
7572− * All tasks are prepped. Disable/exit tasks for $parent and enable for
7573− * the new @sch.
7574− */
7575− scx_task_iter_start(&sti, sch->cgrp);
7576− while ((p = scx_task_iter_next_locked(&sti))) {
7577− /*
7578− * Use clearing of %SCX_TASK_SUB_INIT to detect and skip
7579− * duplicate iterations.
7580− */
7581− if (!(p->scx.flags & SCX_TASK_SUB_INIT))
7582− continue;
7583−
7584− scoped_guard (sched_change, p, DEQUEUE_SAVE | DEQUEUE_MOVE) {
7585− /*
7586− * $p must be either READY or ENABLED. If ENABLED,
7587− * __scx_disabled_and_exit_task() first disables and
7588− * makes it READY. However, after exiting $p, it will
7589− * leave $p as READY.
7590− */
7591− assert_task_ready_or_enabled(p);
7592− __scx_disable_and_exit_task(parent, p);
7593−
7594− /*
7595− * $p is now only initialized for @sch and READY, which
7596− * is what we want. Assign it to @sch and enable.
7597− */
7598− scx_set_task_sched(p, sch);
7599− scx_enable_task(sch, p);
7600−
7601− p->scx.flags &= ~SCX_TASK_SUB_INIT;
7602− }
7603− }
7604− scx_task_iter_stop(&sti);
7605−
7606− scx_enabling_sub_sched = NULL;
7607−
7608− scx_cgroup_unlock();
7609− percpu_up_write(&scx_fork_rwsem);
7610−
7611− scx_bypass(sch, false);
7612−
7613− pr_info("sched_ext: BPF sub-scheduler \"%s\" enabled\n", sch->ops.name);
7614− kobject_uevent(&sch->kobj, KOBJ_ADD);
7615− ret = 0;
7616− goto out_unlock;
7617−
7618−out_put_cgrp:
7619− cgroup_put(cgrp);
7620−out_unlock:
7621− mutex_unlock(&scx_enable_mutex);
7622− cmd->ret = ret;
7623− return;
7624−
7625−abort:
7626− put_task_struct(p);
7627− scx_task_iter_stop(&sti);
7628−
7629− /*
7630− * Undo __scx_init_task() for tasks we marked. scx_enable_task() never
7631− * ran for @sch on them, so calling scx_disable_task() here would invoke
7632− * ops.disable() without a matching ops.enable(). scx_enabling_sub_sched
7633− * must stay set until SUB_INIT is cleared from every marked task -
7634− * scx_disable_and_exit_task() reads it when a task exits concurrently.
7635− */
7636− scx_task_iter_start(&sti, sch->cgrp);
7637− while ((p = scx_task_iter_next_locked(&sti))) {
7638− if (p->scx.flags & SCX_TASK_SUB_INIT) {
7639− scx_sub_init_cancel_task(sch, p);
7640− p->scx.flags &= ~SCX_TASK_SUB_INIT;
7641− }
7642− }
7643− scx_task_iter_stop(&sti);
7644− scx_enabling_sub_sched = NULL;
7645−err_unlock_and_disable:
7646− /* we'll soon enter disable path, keep bypass on */
7647− scx_cgroup_unlock();
7648− percpu_up_write(&scx_fork_rwsem);
7649−err_disable:
7650− mutex_unlock(&scx_enable_mutex);
7651− scx_flush_disable_work(sch);
7652− cmd->ret = 0;
7653−}
7654−
7655−static s32 scx_cgroup_lifetime_notify(struct notifier_block *nb,
7656− unsigned long action, void *data)
7657−{
7658− struct cgroup *cgrp = data;
7659− struct cgroup *parent = cgroup_parent(cgrp);
7660−
7661− if (!cgroup_on_dfl(cgrp))
7662− return NOTIFY_OK;
7663−
7664− switch (action) {
7665− case CGROUP_LIFETIME_ONLINE:
7666− /* inherit ->scx_sched from $parent */
7667− if (parent)
7668− rcu_assign_pointer(cgrp->scx_sched, parent->scx_sched);
7669− break;
7670− case CGROUP_LIFETIME_OFFLINE:
7671− /* if there is a sched attached, shoot it down */
7672− if (cgrp->scx_sched && cgrp->scx_sched->cgrp == cgrp)
7673− scx_exit(cgrp->scx_sched, SCX_EXIT_UNREG_KERN,
7674− SCX_ECODE_RSN_CGROUP_OFFLINE,
7675− "cgroup %llu going offline", cgroup_id(cgrp));
7676− break;
7677− }
7678−
7679− return NOTIFY_OK;
7680−}
7681−
7682−static struct notifier_block scx_cgroup_lifetime_nb = {
7683− .notifier_call = scx_cgroup_lifetime_notify,
7684−};
7685−
7686−static s32 __init scx_cgroup_lifetime_notifier_init(void)
7687−{
7688− return blocking_notifier_chain_register(&cgroup_lifetime_notifier,
7689− &scx_cgroup_lifetime_nb);
7690−}
7691−core_initcall(scx_cgroup_lifetime_notifier_init);
7692−#endif /* CONFIG_EXT_SUB_SCHED */
7693−
76946938 static s32 scx_enable(struct scx_enable_cmd *cmd, struct bpf_link *link)
76956939 {
76966940 static struct kthread_worker *helper;
@@ -7837,20 +7081,6 @@ static int bpf_scx_init_member(const struct btf_type *t, · 140 unmodified lines
78377081 return 0;
78387082 }
78397083
7840−#ifdef CONFIG_EXT_SUB_SCHED
7841−static void scx_pstack_recursion_on_dispatch(struct bpf_prog *prog)
7842−{
7843− struct scx_sched *sch;
7844−
7845− guard(rcu)();
7846− sch = scx_prog_sched(prog->aux);
7847− if (unlikely(!sch))
7848− return;
7849−
7850− scx_error(sch, "dispatch recursion detected");
7851−}
7852−#endif /* CONFIG_EXT_SUB_SCHED */
7853−
78547084 static int bpf_scx_check_member(const struct btf_type *t,
78557085 const struct btf_member *member,
78567086 const struct bpf_prog *prog)
@@ -9021,45 +8251,6 @@ __bpf_kfunc bool scx_bpf_dsq_move_vtime(struct bpf_iter_scx_dsq *it__iter, · 1164 unmodified lines
90218251 p, dsq_id, enq_flags | SCX_ENQ_DSQ_PRIQ);
90228252 }
90238253
9024−#ifdef CONFIG_EXT_SUB_SCHED
9025−/**
9026− * scx_bpf_sub_dispatch - Trigger dispatching on a child scheduler
9027− * @cgroup_id: cgroup ID of the child scheduler to dispatch
9028− * @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
9029− *
9030− * Allows a parent scheduler to trigger dispatching on one of its direct
9031− * child schedulers. The child scheduler runs its dispatch operation to
9032− * move tasks from dispatch queues to the local runqueue.
9033− *
9034− * Returns: true on success, false if cgroup_id is invalid, not a direct
9035− * child, or caller lacks dispatch permission.
9036− */
9037−__bpf_kfunc bool scx_bpf_sub_dispatch(u64 cgroup_id, const struct bpf_prog_aux *aux)
9038−{
9039− struct rq *this_rq = this_rq();
9040− struct scx_sched *parent, *child;
9041−
9042− guard(rcu)();
9043− parent = scx_prog_sched(aux);
9044− if (unlikely(!parent))
9045− return false;
9046−
9047− child = scx_find_sub_sched(cgroup_id);
9048−
9049− if (unlikely(!child))
9050− return false;
9051−
9052− if (unlikely(scx_parent(child) != parent)) {
9053− scx_error(parent, "trying to dispatch a distant sub-sched on cgroup %llu",
9054− cgroup_id);
9055− return false;
9056− }
9057−
9058− return scx_dispatch_sched(child, this_rq, this_rq->scx.sub_dispatch_prev,
9059− true);
9060−}
9061−#endif /* CONFIG_EXT_SUB_SCHED */
9062−
90638254 __bpf_kfunc_end_defs();
90648255
90658256 BTF_KFUNCS_START(scx_kfunc_ids_dispatch)
kernel/sched/ext/internal.h+28−0
@@ -11,6 +11,34 @@ · 10 unmodified lines
1111 #include "../sched.h"
1212 #include "types.h"
1313
14+#include <trace/events/sched_ext.h>
15+
16+/**
17+ * scx_add_event - Increase an event counter for 'name' by 'cnt'
18+ * @sch: scx_sched to account events for
19+ * @name: an event name defined in struct scx_event_stats
20+ * @cnt: the number of the event occurred
21+ *
22+ * This can be used when preemption is not disabled.
23+ */
24+#define scx_add_event(sch, name, cnt) do { \
25+ this_cpu_add((sch)->pcpu->event_stats.name, (cnt)); \
26+ trace_sched_ext_event(#name, (cnt)); \
27+} while(0)
28+
29+/**
30+ * __scx_add_event - Increase an event counter for 'name' by 'cnt'
31+ * @sch: scx_sched to account events for
32+ * @name: an event name defined in struct scx_event_stats
33+ * @cnt: the number of the event occurred
34+ *
35+ * This should be used only when preemption is disabled.
36+ */
37+#define __scx_add_event(sch, name, cnt) do { \
38+ __this_cpu_add((sch)->pcpu->event_stats.name, (cnt)); \
39+ trace_sched_ext_event(#name, cnt); \
40+} while(0)
41+
1442 #define SCX_OP_IDX(op) (offsetof(struct sched_ext_ops, op) / sizeof(void (*)(void)))
1543 #define SCX_MOFF_IDX(moff) ((moff) / sizeof(void (*)(void)))
1644
kernel/sched/ext/sub.c+668−0
@@ -0,0 +1,668 @@
1+// SPDX-License-Identifier: GPL-2.0
2+/*
3+ * BPF extensible scheduler class: Documentation/scheduler/sched-ext.rst
4+ *
5+ * Sub-scheduler hierarchy support.
6+ *
7+ * A sub-scheduler is an scx_sched attached to a cgroup subtree under another
8+ * scx_sched. This file holds the sub-scheduler implementation: the scheduler
9+ * tree walk, capability delegation, per-shard cap state and its sync, and the
10+ * sub-scheduler enable/disable paths. The core dispatch/enqueue machinery it
11+ * builds on lives in ext.c.
12+ *
13+ * Copyright (c) 2026 Meta Platforms, Inc. and affiliates.
14+ * Copyright (c) 2026 Tejun Heo <tj@kernel.org>
15+ */
16+#include <linux/rhashtable.h>
17+#include "internal.h"
18+#include "cid.h"
19+#include "arena.h"
20+#include "sub.h"
21+
22+#ifdef CONFIG_EXT_SUB_SCHED
23+
24+/**
25+ * scx_next_descendant_pre - find the next descendant for pre-order walk
26+ * @pos: the current position (%NULL to initiate traversal)
27+ * @root: sched whose descendants to walk
28+ *
29+ * To be used by scx_for_each_descendant_pre(). Find the next descendant to
30+ * visit for pre-order traversal of @root's descendants. @root is included in
31+ * the iteration and the first node to be visited.
32+ */
33+struct scx_sched *scx_next_descendant_pre(struct scx_sched *pos, struct scx_sched *root)
34+{
35+ struct scx_sched *next;
36+
37+ lockdep_assert(lockdep_is_held(&scx_enable_mutex) ||
38+ lockdep_is_held(&scx_sched_lock));
39+
40+ /* if first iteration, visit @root */
41+ if (!pos)
42+ return root;
43+
44+ /* visit the first child if exists */
45+ next = list_first_entry_or_null(&pos->children, struct scx_sched, sibling);
46+ if (next)
47+ return next;
48+
49+ /* no child, visit my or the closest ancestor's next sibling */
50+ while (pos != root) {
51+ if (!list_is_last(&pos->sibling, &scx_parent(pos)->children))
52+ return list_next_entry(pos, sibling);
53+ pos = scx_parent(pos);
54+ }
55+
56+ return NULL;
57+}
58+
59+static struct scx_sched *scx_find_sub_sched(u64 cgroup_id)
60+{
61+ return rhashtable_lookup(&scx_sched_hash, &cgroup_id,
62+ scx_sched_hash_params);
63+}
64+
65+void scx_set_task_sched(struct task_struct *p, struct scx_sched *sch)
66+{
67+ rcu_assign_pointer(p->scx.sched, sch);
68+}
69+
70+struct cgroup *sch_cgroup(struct scx_sched *sch)
71+{
72+ return sch->cgrp;
73+}
74+
75+/* for each descendant of @cgrp including self, set ->scx_sched to @sch */
76+void set_cgroup_sched(struct cgroup *cgrp, struct scx_sched *sch)
77+{
78+ struct cgroup *pos;
79+ struct cgroup_subsys_state *css;
80+
81+ cgroup_for_each_live_descendant_pre(pos, css, cgrp)
82+ rcu_assign_pointer(pos->scx_sched, sch);
83+}
84+
85+static DECLARE_WAIT_QUEUE_HEAD(scx_unlink_waitq);
86+
87+void drain_descendants(struct scx_sched *sch)
88+{
89+ /*
90+ * Child scheds that finished the critical part of disabling will take
91+ * themselves off @sch->children. Wait for it to drain. As propagation
92+ * is recursive, empty @sch->children means that all proper descendant
93+ * scheds reached unlinking stage.
94+ */
95+ wait_event(scx_unlink_waitq, list_empty(&sch->children));
96+}
97+
98+static void scx_fail_parent(struct scx_sched *sch,
99+ struct task_struct *failed, s32 fail_code)
100+{
101+ struct scx_sched *parent = scx_parent(sch);
102+ struct scx_task_iter sti;
103+ struct task_struct *p;
104+
105+ scx_error(parent, "ops.init_task() failed (%d) for %s[%d] while disabling a sub-scheduler",
106+ fail_code, failed->comm, failed->pid);
107+
108+ /*
109+ * Once $parent is bypassed, it's safe to put SCX_TASK_NONE tasks into
110+ * it. This may cause downstream failures on the BPF side but $parent is
111+ * dying anyway.
112+ */
113+ scx_bypass(parent, true);
114+
115+ scx_task_iter_start(&sti, sch->cgrp);
116+ while ((p = scx_task_iter_next_locked(&sti))) {
117+ if (scx_task_on_sched(parent, p))
118+ continue;
119+
120+ scoped_guard (sched_change, p, DEQUEUE_SAVE | DEQUEUE_MOVE) {
121+ scx_disable_and_exit_task(sch, p);
122+ scx_set_task_sched(p, parent);
123+ }
124+ }
125+ scx_task_iter_stop(&sti);
126+}
127+
128+void scx_sub_disable(struct scx_sched *sch)
129+{
130+ struct scx_sched *parent = scx_parent(sch);
131+ struct scx_task_iter sti;
132+ struct task_struct *p;
133+ int ret;
134+
135+ /*
136+ * Guarantee forward progress and wait for descendants to be disabled.
137+ * To limit disruptions, $parent is not bypassed. Tasks are fully
138+ * prepped and then inserted back into $parent.
139+ */
140+ scx_bypass(sch, true);
141+ drain_descendants(sch);
142+
143+ /*
144+ * Here, every runnable task is guaranteed to make forward progress and
145+ * we can safely use blocking synchronization constructs. Actually
146+ * disable ops.
147+ */
148+ mutex_lock(&scx_enable_mutex);
149+ percpu_down_write(&scx_fork_rwsem);
150+ scx_cgroup_lock();
151+
152+ set_cgroup_sched(sch_cgroup(sch), parent);
153+
154+ scx_task_iter_start(&sti, sch->cgrp);
155+ while ((p = scx_task_iter_next_locked(&sti))) {
156+ struct rq *rq;
157+ struct rq_flags rf;
158+
159+ /* filter out duplicate visits */
160+ if (scx_task_on_sched(parent, p))
161+ continue;
162+
163+ /*
164+ * By the time control reaches here, all descendant schedulers
165+ * should already have been disabled.
166+ */
167+ WARN_ON_ONCE(!scx_task_on_sched(sch, p));
168+
169+ /*
170+ * @p is pinned by the iter: css_task_iter_next() takes a
171+ * reference and holds it until the next iter_next() call, so
172+ * @p->usage is guaranteed > 0.
173+ */
174+ get_task_struct(p);
175+
176+ scx_task_iter_unlock(&sti);
177+
178+ /*
179+ * $p is READY or ENABLED on @sch. Initialize for $parent,
180+ * disable and exit from @sch, and then switch over to $parent.
181+ *
182+ * If a task fails to initialize for $parent, the only available
183+ * action is disabling $parent too. While this allows disabling
184+ * of a child sched to cause the parent scheduler to fail, the
185+ * failure can only originate from ops.init_task() of the
186+ * parent. A child can't directly affect the parent through its
187+ * own failures.
188+ */
189+ ret = __scx_init_task(parent, p, false);
190+ if (ret) {
191+ scx_fail_parent(sch, p, ret);
192+ put_task_struct(p);
193+ break;
194+ }
195+
196+ rq = task_rq_lock(p, &rf);
197+
198+ if (scx_get_task_state(p) == SCX_TASK_DEAD) {
199+ /*
200+ * sched_ext_dead() raced us between __scx_init_task()
201+ * and this rq lock and ran exit_task() on @sch (the
202+ * sched @p was on at that point), not on $parent.
203+ * $parent's just-completed init is owed an exit_task()
204+ * and we issue it here.
205+ */
206+ scx_sub_init_cancel_task(parent, p);
207+ task_rq_unlock(rq, p, &rf);
208+ put_task_struct(p);
209+ continue;
210+ }
211+
212+ scoped_guard (sched_change, p, DEQUEUE_SAVE | DEQUEUE_MOVE) {
213+ /*
214+ * $p is initialized for $parent and still attached to
215+ * @sch. Disable and exit for @sch, switch over to
216+ * $parent, override the state to READY to account for
217+ * $p having already been initialized, and then enable.
218+ */
219+ scx_disable_and_exit_task(sch, p);
220+ scx_set_task_state(p, SCX_TASK_INIT_BEGIN);
221+ scx_set_task_state(p, SCX_TASK_INIT);
222+ scx_set_task_sched(p, parent);
223+ scx_set_task_state(p, SCX_TASK_READY);
224+ scx_enable_task(parent, p);
225+ }
226+
227+ task_rq_unlock(rq, p, &rf);
228+ put_task_struct(p);
229+ }
230+ scx_task_iter_stop(&sti);
231+
232+ scx_disable_dump(sch);
233+
234+ scx_cgroup_unlock();
235+ percpu_up_write(&scx_fork_rwsem);
236+
237+ /*
238+ * All tasks are moved off of @sch but there may still be on-going
239+ * operations (e.g. ops.select_cpu()). Drain them by flushing RCU. Use
240+ * the expedited version as ancestors may be waiting in bypass mode.
241+ * Also, tell the parent that there is no need to keep running bypass
242+ * DSQs for us.
243+ */
244+ synchronize_rcu_expedited();
245+ scx_disable_bypass_dsp(sch);
246+
247+ scx_unlink_sched(sch);
248+
249+ mutex_unlock(&scx_enable_mutex);
250+
251+ /*
252+ * @sch is now unlinked from the parent's children list. Notify and call
253+ * ops.sub_detach/exit(). Note that ops.sub_detach/exit() must be called
254+ * after unlinking and releasing all locks. See scx_claim_exit().
255+ */
256+ wake_up_all(&scx_unlink_waitq);
257+
258+ if (parent->ops.sub_detach && sch->sub_attached) {
259+ struct scx_sub_detach_args sub_detach_args = {
260+ .ops = &sch->ops,
261+ .cgroup_path = sch->cgrp_path,
262+ };
263+ SCX_CALL_OP(parent, sub_detach, NULL,
264+ &sub_detach_args);
265+ }
266+
267+ scx_log_sched_disable(sch);
268+
269+ if (sch->ops.exit)
270+ SCX_CALL_OP(sch, exit, NULL, sch->exit_info);
271+ if (sch->sub_kset)
272+ kobject_del(&sch->sub_kset->kobj);
273+ kobject_del(&sch->kobj);
274+}
275+
276+/* verify that a scheduler can be attached to @cgrp and return the parent */
277+static struct scx_sched *find_parent_sched(struct cgroup *cgrp)
278+{
279+ struct scx_sched *parent = cgrp->scx_sched;
280+ struct scx_sched *pos;
281+
282+ lockdep_assert_held(&scx_sched_lock);
283+
284+ /* can't attach twice to the same cgroup */
285+ if (parent->cgrp == cgrp)
286+ return ERR_PTR(-EBUSY);
287+
288+ /* does $parent allow sub-scheds? */
289+ if (!parent->ops.sub_attach)
290+ return ERR_PTR(-EOPNOTSUPP);
291+
292+ /* can't insert between $parent and its exiting children */
293+ list_for_each_entry(pos, &parent->children, sibling)
294+ if (cgroup_is_descendant(pos->cgrp, cgrp))
295+ return ERR_PTR(-EBUSY);
296+
297+ return parent;
298+}
299+
300+static bool assert_task_ready_or_enabled(struct task_struct *p)
301+{
302+ u32 state = scx_get_task_state(p);
303+
304+ switch (state) {
305+ case SCX_TASK_READY:
306+ case SCX_TASK_ENABLED:
307+ return true;
308+ default:
309+ WARN_ONCE(true, "sched_ext: Invalid task state %d for %s[%d] during enabling sub sched",
310+ state, p->comm, p->pid);
311+ return false;
312+ }
313+}
314+
315+void scx_sub_enable_workfn(struct kthread_work *work)
316+{
317+ struct scx_enable_cmd *cmd = container_of(work, struct scx_enable_cmd, work);
318+ struct sched_ext_ops *ops = cmd->ops;
319+ struct cgroup *cgrp;
320+ struct scx_sched *parent, *sch;
321+ struct scx_task_iter sti;
322+ struct task_struct *p;
323+ s32 i, ret;
324+
325+ mutex_lock(&scx_enable_mutex);
326+
327+ if (!scx_enabled()) {
328+ ret = -ENODEV;
329+ goto out_unlock;
330+ }
331+
332+ /* See scx_root_enable_workfn() for the @ops->priv check. */
333+ if (rcu_access_pointer(ops->priv)) {
334+ ret = -EBUSY;
335+ goto out_unlock;
336+ }
337+
338+ cgrp = cgroup_get_from_id(ops->sub_cgroup_id);
339+ if (IS_ERR(cgrp)) {
340+ ret = PTR_ERR(cgrp);
341+ goto out_unlock;
342+ }
343+
344+ raw_spin_lock_irq(&scx_sched_lock);
345+ parent = find_parent_sched(cgrp);
346+ if (IS_ERR(parent)) {
347+ raw_spin_unlock_irq(&scx_sched_lock);
348+ ret = PTR_ERR(parent);
349+ goto out_put_cgrp;
350+ }
351+ kobject_get(&parent->kobj);
352+ raw_spin_unlock_irq(&scx_sched_lock);
353+
354+ /* scx_alloc_and_add_sched() consumes @cgrp whether it succeeds or not */
355+ sch = scx_alloc_and_add_sched(cmd, cgrp, parent);
356+ kobject_put(&parent->kobj);
357+ if (IS_ERR(sch)) {
358+ ret = PTR_ERR(sch);
359+ goto out_unlock;
360+ }
361+
362+ ret = scx_link_sched(sch);
363+ if (ret)
364+ goto err_disable;
365+
366+ if (sch->level >= SCX_SUB_MAX_DEPTH) {
367+ scx_error(sch, "max nesting depth %d violated",
368+ SCX_SUB_MAX_DEPTH);
369+ goto err_disable;
370+ }
371+
372+ if (sch->ops.init) {
373+ ret = SCX_CALL_OP_RET(sch, init, NULL);
374+ if (ret) {
375+ ret = scx_ops_sanitize_err(sch, "init", ret);
376+ scx_error(sch, "ops.init() failed (%d)", ret);
377+ goto err_disable;
378+ }
379+ sch->exit_info->flags |= SCX_EFLAG_INITIALIZED;
380+ }
381+
382+ ret = scx_arena_pool_init(sch);
383+ if (ret)
384+ goto err_disable;
385+
386+ ret = scx_set_cmask_scratch_alloc(sch);
387+ if (ret)
388+ goto err_disable;
389+
390+ if (scx_validate_ops(sch, ops))
391+ goto err_disable;
392+
393+ struct scx_sub_attach_args sub_attach_args = {
394+ .ops = &sch->ops,
395+ .cgroup_path = sch->cgrp_path,
396+ };
397+
398+ ret = SCX_CALL_OP_RET(parent, sub_attach, NULL,
399+ &sub_attach_args);
400+ if (ret) {
401+ ret = scx_ops_sanitize_err(sch, "sub_attach", ret);
402+ scx_error(sch, "parent rejected (%d)", ret);
403+ goto err_disable;
404+ }
405+ sch->sub_attached = true;
406+
407+ scx_bypass(sch, true);
408+
409+ for (i = SCX_OPI_BEGIN; i < SCX_OPI_END; i++)
410+ if (((void (**)(void))ops)[i])
411+ set_bit(i, sch->has_op);
412+
413+ percpu_down_write(&scx_fork_rwsem);
414+ scx_cgroup_lock();
415+
416+ /*
417+ * Set cgroup->scx_sched's and check CSS_ONLINE. Either we see
418+ * !CSS_ONLINE or scx_cgroup_lifetime_notify() sees and shoots us down.
419+ */
420+ set_cgroup_sched(sch_cgroup(sch), sch);
421+ if (!(cgrp->self.flags & CSS_ONLINE)) {
422+ scx_error(sch, "cgroup is not online");
423+ goto err_unlock_and_disable;
424+ }
425+
426+ /*
427+ * Initialize tasks for the new child $sch without exiting them for
428+ * $parent so that the tasks can always be reverted back to $parent
429+ * sched on child init failure.
430+ */
431+ WARN_ON_ONCE(scx_enabling_sub_sched);
432+ scx_enabling_sub_sched = sch;
433+
434+ scx_task_iter_start(&sti, sch->cgrp);
435+ while ((p = scx_task_iter_next_locked(&sti))) {
436+ struct rq *rq;
437+ struct rq_flags rf;
438+
439+ /*
440+ * Task iteration may visit the same task twice when racing
441+ * against exiting. Use %SCX_TASK_SUB_INIT to mark tasks which
442+ * finished __scx_init_task() and skip if set.
443+ *
444+ * A task may exit and get freed between __scx_init_task()
445+ * completion and scx_enable_task(). In such cases,
446+ * scx_disable_and_exit_task() must exit the task for both the
447+ * parent and child scheds.
448+ */
449+ if (p->scx.flags & SCX_TASK_SUB_INIT)
450+ continue;
451+
452+ /* @p is pinned by the iter; see scx_sub_disable() */
453+ get_task_struct(p);
454+
455+ if (!assert_task_ready_or_enabled(p)) {
456+ ret = -EINVAL;
457+ goto abort;
458+ }
459+
460+ scx_task_iter_unlock(&sti);
461+
462+ /*
463+ * As $p is still on $parent, it can't be transitioned to INIT.
464+ * Let's worry about task state later. Use __scx_init_task().
465+ */
466+ ret = __scx_init_task(sch, p, false);
467+ if (ret)
468+ goto abort;
469+
470+ rq = task_rq_lock(p, &rf);
471+
472+ if (scx_get_task_state(p) == SCX_TASK_DEAD) {
473+ /*
474+ * sched_ext_dead() raced us between __scx_init_task()
475+ * and this rq lock and ran exit_task() on $parent (the
476+ * sched @p was on at that point), not on @sch. @sch's
477+ * just-completed init is owed an exit_task() and we
478+ * issue it here.
479+ */
480+ scx_sub_init_cancel_task(sch, p);
481+ task_rq_unlock(rq, p, &rf);
482+ put_task_struct(p);
483+ continue;
484+ }
485+
486+ p->scx.flags |= SCX_TASK_SUB_INIT;
487+ task_rq_unlock(rq, p, &rf);
488+
489+ put_task_struct(p);
490+ }
491+ scx_task_iter_stop(&sti);
492+
493+ /*
494+ * All tasks are prepped. Disable/exit tasks for $parent and enable for
495+ * the new @sch.
496+ */
497+ scx_task_iter_start(&sti, sch->cgrp);
498+ while ((p = scx_task_iter_next_locked(&sti))) {
499+ /*
500+ * Use clearing of %SCX_TASK_SUB_INIT to detect and skip
501+ * duplicate iterations.
502+ */
503+ if (!(p->scx.flags & SCX_TASK_SUB_INIT))
504+ continue;
505+
506+ scoped_guard (sched_change, p, DEQUEUE_SAVE | DEQUEUE_MOVE) {
507+ /*
508+ * $p must be either READY or ENABLED. If ENABLED,
509+ * __scx_disabled_and_exit_task() first disables and
510+ * makes it READY. However, after exiting $p, it will
511+ * leave $p as READY.
512+ */
513+ assert_task_ready_or_enabled(p);
514+ __scx_disable_and_exit_task(parent, p);
515+
516+ /*
517+ * $p is now only initialized for @sch and READY, which
518+ * is what we want. Assign it to @sch and enable.
519+ */
520+ scx_set_task_sched(p, sch);
521+ scx_enable_task(sch, p);
522+
523+ p->scx.flags &= ~SCX_TASK_SUB_INIT;
524+ }
525+ }
526+ scx_task_iter_stop(&sti);
527+
528+ scx_enabling_sub_sched = NULL;
529+
530+ scx_cgroup_unlock();
531+ percpu_up_write(&scx_fork_rwsem);
532+
533+ scx_bypass(sch, false);
534+
535+ pr_info("sched_ext: BPF sub-scheduler \"%s\" enabled\n", sch->ops.name);
536+ kobject_uevent(&sch->kobj, KOBJ_ADD);
537+ ret = 0;
538+ goto out_unlock;
539+
540+out_put_cgrp:
541+ cgroup_put(cgrp);
542+out_unlock:
543+ mutex_unlock(&scx_enable_mutex);
544+ cmd->ret = ret;
545+ return;
546+
547+abort:
548+ put_task_struct(p);
549+ scx_task_iter_stop(&sti);
550+
551+ /*
552+ * Undo __scx_init_task() for tasks we marked. scx_enable_task() never
553+ * ran for @sch on them, so calling scx_disable_task() here would invoke
554+ * ops.disable() without a matching ops.enable(). scx_enabling_sub_sched
555+ * must stay set until SUB_INIT is cleared from every marked task -
556+ * scx_disable_and_exit_task() reads it when a task exits concurrently.
557+ */
558+ scx_task_iter_start(&sti, sch->cgrp);
559+ while ((p = scx_task_iter_next_locked(&sti))) {
560+ if (p->scx.flags & SCX_TASK_SUB_INIT) {
561+ scx_sub_init_cancel_task(sch, p);
562+ p->scx.flags &= ~SCX_TASK_SUB_INIT;
563+ }
564+ }
565+ scx_task_iter_stop(&sti);
566+ scx_enabling_sub_sched = NULL;
567+err_unlock_and_disable:
568+ /* we'll soon enter disable path, keep bypass on */
569+ scx_cgroup_unlock();
570+ percpu_up_write(&scx_fork_rwsem);
571+err_disable:
572+ mutex_unlock(&scx_enable_mutex);
573+ scx_flush_disable_work(sch);
574+ cmd->ret = 0;
575+}
576+
577+static s32 scx_cgroup_lifetime_notify(struct notifier_block *nb,
578+ unsigned long action, void *data)
579+{
580+ struct cgroup *cgrp = data;
581+ struct cgroup *parent = cgroup_parent(cgrp);
582+
583+ if (!cgroup_on_dfl(cgrp))
584+ return NOTIFY_OK;
585+
586+ switch (action) {
587+ case CGROUP_LIFETIME_ONLINE:
588+ /* inherit ->scx_sched from $parent */
589+ if (parent)
590+ rcu_assign_pointer(cgrp->scx_sched, parent->scx_sched);
591+ break;
592+ case CGROUP_LIFETIME_OFFLINE:
593+ /* if there is a sched attached, shoot it down */
594+ if (cgrp->scx_sched && cgrp->scx_sched->cgrp == cgrp)
595+ scx_exit(cgrp->scx_sched, SCX_EXIT_UNREG_KERN,
596+ SCX_ECODE_RSN_CGROUP_OFFLINE,
597+ "cgroup %llu going offline", cgroup_id(cgrp));
598+ break;
599+ }
600+
601+ return NOTIFY_OK;
602+}
603+
604+static struct notifier_block scx_cgroup_lifetime_nb = {
605+ .notifier_call = scx_cgroup_lifetime_notify,
606+};
607+
608+static s32 __init scx_cgroup_lifetime_notifier_init(void)
609+{
610+ return blocking_notifier_chain_register(&cgroup_lifetime_notifier,
611+ &scx_cgroup_lifetime_nb);
612+}
613+core_initcall(scx_cgroup_lifetime_notifier_init);
614+
615+void scx_pstack_recursion_on_dispatch(struct bpf_prog *prog)
616+{
617+ struct scx_sched *sch;
618+
619+ guard(rcu)();
620+ sch = scx_prog_sched(prog->aux);
621+ if (unlikely(!sch))
622+ return;
623+
624+ scx_error(sch, "dispatch recursion detected");
625+}
626+
627+__bpf_kfunc_start_defs();
628+
629+/**
630+ * scx_bpf_sub_dispatch - Trigger dispatching on a child scheduler
631+ * @cgroup_id: cgroup ID of the child scheduler to dispatch
632+ * @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
633+ *
634+ * Allows a parent scheduler to trigger dispatching on one of its direct
635+ * child schedulers. The child scheduler runs its dispatch operation to
636+ * move tasks from dispatch queues to the local runqueue.
637+ *
638+ * Returns: true on success, false if cgroup_id is invalid, not a direct
639+ * child, or caller lacks dispatch permission.
640+ */
641+__bpf_kfunc bool scx_bpf_sub_dispatch(u64 cgroup_id, const struct bpf_prog_aux *aux)
642+{
643+ struct rq *this_rq = this_rq();
644+ struct scx_sched *parent, *child;
645+
646+ guard(rcu)();
647+ parent = scx_prog_sched(aux);
648+ if (unlikely(!parent))
649+ return false;
650+
651+ child = scx_find_sub_sched(cgroup_id);
652+
653+ if (unlikely(!child))
654+ return false;
655+
656+ if (unlikely(scx_parent(child) != parent)) {
657+ scx_error(parent, "trying to dispatch a distant sub-sched on cgroup %llu",
658+ cgroup_id);
659+ return false;
660+ }
661+
662+ return scx_dispatch_sched(child, this_rq, this_rq->scx.sub_dispatch_prev,
663+ true);
664+}
665+
666+__bpf_kfunc_end_defs();
667+
668+#endif /* CONFIG_EXT_SUB_SCHED */
kernel/sched/ext/sub.h+161−0
@@ -0,0 +1,161 @@
1+/* SPDX-License-Identifier: GPL-2.0 */
2+/*
3+ * BPF extensible scheduler class: Documentation/scheduler/sched-ext.rst
4+ *
5+ * Sub-scheduler hierarchy support.
6+ *
7+ * Copyright (c) 2026 Meta Platforms, Inc. and affiliates.
8+ * Copyright (c) 2026 Tejun Heo <tj@kernel.org>
9+ */
10+#ifndef _KERNEL_SCHED_EXT_SUB_H
11+#define _KERNEL_SCHED_EXT_SUB_H
12+
13+#include "internal.h"
14+#include "cid.h"
15+
16+#ifdef CONFIG_EXT_SUB_SCHED
17+
18+struct scx_sched *scx_next_descendant_pre(struct scx_sched *pos, struct scx_sched *root);
19+void scx_set_task_sched(struct task_struct *p, struct scx_sched *sch);
20+struct cgroup *sch_cgroup(struct scx_sched *sch);
21+void set_cgroup_sched(struct cgroup *cgrp, struct scx_sched *sch);
22+void scx_pstack_recursion_on_dispatch(struct bpf_prog *prog);
23+void drain_descendants(struct scx_sched *sch);
24+void scx_sub_disable(struct scx_sched *sch);
25+void scx_sub_enable_workfn(struct kthread_work *work);
26+bool scx_bpf_sub_dispatch(u64 cgroup_id, const struct bpf_prog_aux *aux);
27+
28+#else /* CONFIG_EXT_SUB_SCHED */
29+
30+static inline struct scx_sched *scx_next_descendant_pre(struct scx_sched *pos, struct scx_sched *root) { return pos ? NULL : root; }
31+static inline void scx_set_task_sched(struct task_struct *p, struct scx_sched *sch) {}
32+static inline struct cgroup *sch_cgroup(struct scx_sched *sch) { return NULL; }
33+static inline void set_cgroup_sched(struct cgroup *cgrp, struct scx_sched *sch) {}
34+static inline void drain_descendants(struct scx_sched *sch) { }
35+static inline void scx_sub_disable(struct scx_sched *sch) { }
36+
37+#endif /* CONFIG_EXT_SUB_SCHED */
38+
39+/**
40+ * scx_for_each_descendant_pre - pre-order walk of a sched's descendants
41+ * @pos: iteration cursor
42+ * @root: sched to walk the descendants of
43+ *
44+ * Walk @root's descendants. @root is included in the iteration and the first
45+ * node to be visited. Must be called with either scx_enable_mutex or
46+ * scx_sched_lock held.
47+ */
48+#define scx_for_each_descendant_pre(pos, root) \
49+ for ((pos) = scx_next_descendant_pre(NULL, (root)); (pos); \
50+ (pos) = scx_next_descendant_pre((pos), (root)))
51+
52+/*
53+ * One user of this function is scx_bpf_dispatch() which can be called
54+ * recursively as sub-sched dispatches nest. Always inline to reduce stack usage
55+ * from the call frame.
56+ */
57+static __always_inline bool
58+scx_dispatch_sched(struct scx_sched *sch, struct rq *rq,
59+ struct task_struct *prev, bool nested)
60+{
61+ struct scx_dsp_ctx *dspc = &this_cpu_ptr(sch->pcpu)->dsp_ctx;
62+ int nr_loops = SCX_DSP_MAX_LOOPS;
63+ s32 cpu = cpu_of(rq);
64+ bool prev_on_sch = (prev->sched_class == &ext_sched_class) &&
65+ scx_task_on_sched(sch, prev);
66+
67+ if (scx_consume_global_dsq(sch, rq))
68+ return true;
69+
70+ if (scx_bypass_dsp_enabled(sch)) {
71+ /* if @sch is bypassing, only the bypass DSQs are active */
72+ if (scx_bypassing(sch, cpu))
73+ return scx_consume_dispatch_q(sch, rq, scx_bypass_dsq(sch, cpu), 0);
74+
75+#ifdef CONFIG_EXT_SUB_SCHED
76+ /*
77+ * If @sch isn't bypassing but its children are, @sch is
78+ * responsible for making forward progress for both its own
79+ * tasks that aren't bypassing and the bypassing descendants'
80+ * tasks. The following implements a simple built-in behavior -
81+ * let each CPU try to run the bypass DSQ every Nth time.
82+ *
83+ * Later, if necessary, we can add an ops flag to suppress the
84+ * auto-consumption and a kfunc to consume the bypass DSQ and,
85+ * so that the BPF scheduler can fully control scheduling of
86+ * bypassed tasks.
87+ */
88+ struct scx_sched_pcpu *pcpu = per_cpu_ptr(sch->pcpu, cpu);
89+
90+ if (!(pcpu->bypass_host_seq++ % SCX_BYPASS_HOST_NTH) &&
91+ scx_consume_dispatch_q(sch, rq, scx_bypass_dsq(sch, cpu), 0)) {
92+ __scx_add_event(sch, SCX_EV_SUB_BYPASS_DISPATCH, 1);
93+ return true;
94+ }
95+#endif /* CONFIG_EXT_SUB_SCHED */
96+ }
97+
98+ if (unlikely(!SCX_HAS_OP(sch, dispatch)) || !scx_rq_online(rq))
99+ return false;
100+
101+ dspc->rq = rq;
102+
103+ /*
104+ * The dispatch loop. Because scx_flush_dispatch_buf() may drop the rq
105+ * lock, the local DSQ might still end up empty after a successful
106+ * ops.dispatch(). If the local DSQ is empty even after ops.dispatch()
107+ * produced some tasks, retry. The BPF scheduler may depend on this
108+ * looping behavior to simplify its implementation.
109+ */
110+ do {
111+ dspc->nr_tasks = 0;
112+
113+ if (nested) {
114+ SCX_CALL_OP(sch, dispatch, rq, scx_cpu_arg(cpu),
115+ prev_on_sch ? prev : NULL);
116+ } else {
117+ /* stash @prev so that nested invocations can access it */
118+ rq->scx.sub_dispatch_prev = prev;
119+ SCX_CALL_OP(sch, dispatch, rq, scx_cpu_arg(cpu),
120+ prev_on_sch ? prev : NULL);
121+ rq->scx.sub_dispatch_prev = NULL;
122+ }
123+
124+ scx_flush_dispatch_buf(sch, rq);
125+
126+ if ((prev->scx.flags & SCX_TASK_QUEUED) && prev->scx.slice) {
127+ rq->scx.flags |= SCX_RQ_BAL_KEEP;
128+ return true;
129+ }
130+ if (rq->scx.local_dsq.nr)
131+ return true;
132+ if (scx_consume_global_dsq(sch, rq))
133+ return true;
134+
135+ /*
136+ * ops.dispatch() can trap us in this loop by repeatedly
137+ * dispatching ineligible tasks. Break out once in a while to
138+ * allow the watchdog to run. As IRQ can't be enabled in
139+ * balance(), we want to complete this scheduling cycle and then
140+ * start a new one. IOW, we want to call resched_curr() on the
141+ * next, most likely idle, task, not the current one. Use
142+ * __scx_bpf_kick_cpu() for deferred kicking.
143+ */
144+ if (unlikely(!--nr_loops)) {
145+ scx_kick_cpu(sch, cpu, 0);
146+ break;
147+ }
148+ } while (dspc->nr_tasks);
149+
150+ /*
151+ * Prevent the CPU from going idle while bypassed descendants have tasks
152+ * queued. Without this fallback, bypassed tasks could stall if the host
153+ * scheduler's ops.dispatch() doesn't yield any tasks.
154+ */
155+ if (scx_bypass_dsp_enabled(sch))
156+ return scx_consume_dispatch_q(sch, rq, scx_bypass_dsq(sch, cpu), 0);
157+
158+ return false;
159+}
160+
161+#endif /* _KERNEL_SCHED_EXT_SUB_H */