xref: /linux/kernel/sched/ext/sub.h (revision 5f2a9a4c2e6c827eafa8491764125934c5466a84)
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_skip_subtree_pre(struct scx_sched *pos, struct scx_sched *root);
19 struct scx_sched *scx_next_descendant_pre(struct scx_sched *pos, struct scx_sched *root);
20 void scx_set_task_sched(struct task_struct *p, struct scx_sched *sch);
21 struct cgroup *sch_cgroup(struct scx_sched *sch);
22 void set_cgroup_sched(struct cgroup *cgrp, struct scx_sched *sch);
23 void scx_pstack_recursion_on_dispatch(struct bpf_prog *prog);
24 void scx_pstack_recursion_on_caps_updated(struct bpf_prog *prog);
25 void drain_descendants(struct scx_sched *sch);
26 void scx_sub_disable(struct scx_sched *sch);
27 void scx_sub_enable_workfn(struct kthread_work *work);
28 bool scx_bpf_sub_dispatch(u64 cgroup_id, const struct bpf_prog_aux *aux);
29 void scx_free_pshards(struct scx_sched *sch);
30 s32 scx_alloc_pshards(struct scx_sched *sch);
31 void scx_init_root_caps(struct scx_sched *sch);
32 
33 static inline const char *sch_cgrp_path(struct scx_sched *sch)
34 {
35 	return sch->cgrp_path;
36 }
37 
38 #else	/* CONFIG_EXT_SUB_SCHED */
39 
40 static inline struct scx_sched *scx_next_descendant_pre(struct scx_sched *pos, struct scx_sched *root) { return pos ? NULL : root; }
41 static inline struct scx_sched *scx_skip_subtree_pre(struct scx_sched *pos, struct scx_sched *root) { return NULL; }
42 static inline void scx_set_task_sched(struct task_struct *p, struct scx_sched *sch) {}
43 static inline struct cgroup *sch_cgroup(struct scx_sched *sch) { return NULL; }
44 static inline const char *sch_cgrp_path(struct scx_sched *sch) { return "/"; }
45 static inline void set_cgroup_sched(struct cgroup *cgrp, struct scx_sched *sch) {}
46 static inline void drain_descendants(struct scx_sched *sch) { }
47 static inline void scx_sub_disable(struct scx_sched *sch) { }
48 static inline void scx_free_pshards(struct scx_sched *sch) {}
49 static inline s32 scx_alloc_pshards(struct scx_sched *sch) { return 0; }
50 static inline void scx_init_root_caps(struct scx_sched *sch) {}
51 
52 #endif	/* CONFIG_EXT_SUB_SCHED */
53 
54 /**
55  * scx_for_each_descendant_pre - pre-order walk of a sched's descendants
56  * @pos: iteration cursor
57  * @root: sched to walk the descendants of
58  *
59  * Walk @root's descendants. @root is included in the iteration and the first
60  * node to be visited. Must be called with scx_enable_mutex, scx_sched_lock, or
61  * RCU read lock.
62  */
63 #define scx_for_each_descendant_pre(pos, root)					\
64 	for ((pos) = scx_next_descendant_pre(NULL, (root)); (pos);		\
65 	     (pos) = scx_next_descendant_pre((pos), (root)))
66 
67 /*
68  * One user of this function is scx_bpf_dispatch() which can be called
69  * recursively as sub-sched dispatches nest. Always inline to reduce stack usage
70  * from the call frame.
71  */
72 static __always_inline bool
73 scx_dispatch_sched(struct scx_sched *sch, struct rq *rq,
74 		   struct task_struct *prev, bool nested)
75 {
76 	struct scx_dsp_ctx *dspc = &this_cpu_ptr(sch->pcpu)->dsp_ctx;
77 	int nr_loops = SCX_DSP_MAX_LOOPS;
78 	s32 cpu = cpu_of(rq);
79 	bool prev_on_sch = (prev->sched_class == &ext_sched_class) &&
80 		scx_task_on_sched(sch, prev);
81 
82 	if (scx_consume_global_dsq(sch, rq))
83 		return true;
84 
85 	if (scx_bypass_dsp_enabled(sch)) {
86 		/* if @sch is bypassing, only the bypass DSQs are active */
87 		if (scx_bypassing(sch, cpu))
88 			return scx_consume_dispatch_q(sch, rq, scx_bypass_dsq(sch, cpu), 0);
89 
90 #ifdef CONFIG_EXT_SUB_SCHED
91 		/*
92 		 * If @sch isn't bypassing but its children are, @sch is
93 		 * responsible for making forward progress for both its own
94 		 * tasks that aren't bypassing and the bypassing descendants'
95 		 * tasks. The following implements a simple built-in behavior -
96 		 * let each CPU try to run the bypass DSQ every Nth time.
97 		 *
98 		 * Later, if necessary, we can add an ops flag to suppress the
99 		 * auto-consumption and a kfunc to consume the bypass DSQ and,
100 		 * so that the BPF scheduler can fully control scheduling of
101 		 * bypassed tasks.
102 		 */
103 		struct scx_sched_pcpu *pcpu = per_cpu_ptr(sch->pcpu, cpu);
104 
105 		if (!(pcpu->bypass_host_seq++ % SCX_BYPASS_HOST_NTH) &&
106 		    scx_consume_dispatch_q(sch, rq, scx_bypass_dsq(sch, cpu), 0)) {
107 			__scx_add_event(sch, SCX_EV_SUB_BYPASS_DISPATCH, 1);
108 			return true;
109 		}
110 #endif	/* CONFIG_EXT_SUB_SCHED */
111 	}
112 
113 	if (unlikely(!SCX_HAS_OP(sch, dispatch)) || !scx_rq_online(rq))
114 		return false;
115 
116 	dspc->rq = rq;
117 
118 	/*
119 	 * The dispatch loop. Because scx_flush_dispatch_buf() may drop the rq
120 	 * lock, the local DSQ might still end up empty after a successful
121 	 * ops.dispatch(). If the local DSQ is empty even after ops.dispatch()
122 	 * produced some tasks, retry. The BPF scheduler may depend on this
123 	 * looping behavior to simplify its implementation.
124 	 */
125 	do {
126 		dspc->nr_tasks = 0;
127 
128 		if (nested) {
129 			SCX_CALL_OP(sch, dispatch, rq, scx_cpu_arg(cpu),
130 				    prev_on_sch ? prev : NULL);
131 		} else {
132 			/* stash @prev so that nested invocations can access it */
133 			rq->scx.sub_dispatch_prev = prev;
134 			SCX_CALL_OP(sch, dispatch, rq, scx_cpu_arg(cpu),
135 				    prev_on_sch ? prev : NULL);
136 			rq->scx.sub_dispatch_prev = NULL;
137 		}
138 
139 		scx_flush_dispatch_buf(sch, rq);
140 
141 		if ((prev->scx.flags & SCX_TASK_QUEUED) && prev->scx.slice) {
142 			rq->scx.flags |= SCX_RQ_BAL_KEEP;
143 			return true;
144 		}
145 		if (rq->scx.local_dsq.nr)
146 			return true;
147 		if (scx_consume_global_dsq(sch, rq))
148 			return true;
149 
150 		/*
151 		 * ops.dispatch() can trap us in this loop by repeatedly
152 		 * dispatching ineligible tasks. Break out once in a while to
153 		 * allow the watchdog to run. As IRQ can't be enabled in
154 		 * balance(), we want to complete this scheduling cycle and then
155 		 * start a new one. IOW, we want to call resched_curr() on the
156 		 * next, most likely idle, task, not the current one. Use
157 		 * __scx_bpf_kick_cpu() for deferred kicking.
158 		 */
159 		if (unlikely(!--nr_loops)) {
160 			scx_kick_cpu(sch, cpu, 0);
161 			break;
162 		}
163 	} while (dspc->nr_tasks);
164 
165 	/*
166 	 * Prevent the CPU from going idle while bypassed descendants have tasks
167 	 * queued. Without this fallback, bypassed tasks could stall if the host
168 	 * scheduler's ops.dispatch() doesn't yield any tasks.
169 	 */
170 	if (scx_bypass_dsp_enabled(sch))
171 		return scx_consume_dispatch_q(sch, rq, scx_bypass_dsq(sch, cpu), 0);
172 
173 	return false;
174 }
175 
176 #endif /* _KERNEL_SCHED_EXT_SUB_H */
177