xref: /linux/kernel/sched/ext/sub.h (revision 75c268ed57f2197a4f1fd9259f42cb6065dd311b)
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 void scx_process_sync_ecaps(struct rq *rq, struct task_struct *prev);
33 void scx_unbypass_replay_ecaps(struct rq *rq, struct scx_sched *sch);
34 void scx_online_ecaps(struct rq *rq);
35 void scx_offline_ecaps(struct rq *rq);
36 void scx_discard_ecaps_to_sync(s32 cpu, struct scx_sched_pcpu *pcpu);
37 void scx_discard_stale_ecaps_syncs(void);
38 struct scx_dispatch_q *scx_local_or_reject_dsq(struct scx_sched *sch, struct rq *rq,
39 					       struct task_struct *p, u64 *enq_flags);
40 bool scx_task_reenq_on_cap_revoke(struct rq *rq, struct task_struct *p);
41 void scx_reenq_reject(struct rq *rq);
42 
43 static inline const char *sch_cgrp_path(struct scx_sched *sch)
44 {
45 	return sch->cgrp_path;
46 }
47 
48 #else	/* CONFIG_EXT_SUB_SCHED */
49 
50 static inline struct scx_sched *scx_next_descendant_pre(struct scx_sched *pos, struct scx_sched *root) { return pos ? NULL : root; }
51 static inline struct scx_sched *scx_skip_subtree_pre(struct scx_sched *pos, struct scx_sched *root) { return NULL; }
52 static inline void scx_set_task_sched(struct task_struct *p, struct scx_sched *sch) {}
53 static inline struct cgroup *sch_cgroup(struct scx_sched *sch) { return NULL; }
54 static inline const char *sch_cgrp_path(struct scx_sched *sch) { return "/"; }
55 static inline void set_cgroup_sched(struct cgroup *cgrp, struct scx_sched *sch) {}
56 static inline void drain_descendants(struct scx_sched *sch) { }
57 static inline void scx_sub_disable(struct scx_sched *sch) { }
58 static inline void scx_free_pshards(struct scx_sched *sch) {}
59 static inline s32 scx_alloc_pshards(struct scx_sched *sch) { return 0; }
60 static inline void scx_init_root_caps(struct scx_sched *sch) {}
61 static inline void scx_process_sync_ecaps(struct rq *rq, struct task_struct *prev) {}
62 static inline void scx_unbypass_replay_ecaps(struct rq *rq, struct scx_sched *sch) {}
63 static inline void scx_online_ecaps(struct rq *rq) {}
64 static inline void scx_offline_ecaps(struct rq *rq) {}
65 static inline void scx_discard_ecaps_to_sync(s32 cpu, struct scx_sched_pcpu *pcpu) {}
66 static inline void scx_discard_stale_ecaps_syncs(void) {}
67 static inline struct scx_dispatch_q *scx_local_or_reject_dsq(struct scx_sched *sch, struct rq *rq, struct task_struct *p, u64 *enq_flags) { return &rq->scx.local_dsq; }
68 static inline bool scx_task_reenq_on_cap_revoke(struct rq *rq, struct task_struct *p) { return false; }
69 static inline void scx_reenq_reject(struct rq *rq) {}
70 
71 #endif	/* CONFIG_EXT_SUB_SCHED */
72 
73 /**
74  * scx_for_each_descendant_pre - pre-order walk of a sched's descendants
75  * @pos: iteration cursor
76  * @root: sched to walk the descendants of
77  *
78  * Walk @root's descendants. @root is included in the iteration and the first
79  * node to be visited. Must be called with scx_enable_mutex, scx_sched_lock, or
80  * RCU read lock.
81  */
82 #define scx_for_each_descendant_pre(pos, root)					\
83 	for ((pos) = scx_next_descendant_pre(NULL, (root)); (pos);		\
84 	     (pos) = scx_next_descendant_pre((pos), (root)))
85 
86 #ifdef CONFIG_EXT_SUB_SCHED
87 
88 /**
89  * scx_missing_caps - The caps in @needed that @sch lacks on @cpu
90  * @sch: sched to test
91  * @cpu: cpu to test on
92  * @needed: bitmask of SCX_CAP_* values
93  *
94  * Return the caps in @needed that @sch lacks for @cpu, 0 if it holds them all.
95  */
96 static inline u64 scx_missing_caps(struct scx_sched *sch, s32 cpu, u64 needed)
97 {
98 	u64 ecaps;
99 
100 	/* root holds every cap on every cpu */
101 	if (!sch->level)
102 		return 0;
103 
104 	ecaps = READ_ONCE(per_cpu_ptr(sch->pcpu, cpu)->ecaps);
105 
106 	return needed & ~ecaps;
107 }
108 
109 /*
110  * Cap semantics: which caps an action requires, and which caps a cap implies.
111  * Keep all such mappings collected here.
112  */
113 
114 /* map @enq_flags to the SCX_CAP_* bit required for the local-DSQ insert */
115 static inline u64 scx_caps_for_enq(u64 enq_flags)
116 {
117 	/* a restored task must be put into the local DSQ regardless of caps */
118 	if (enq_flags & SCX_ENQ_IGNORE_CAPS)
119 		return 0;
120 	if (enq_flags & SCX_ENQ_IMMED)
121 		return SCX_CAP_ENQ_IMMED;
122 	return SCX_CAP_ENQ;
123 }
124 
125 /* map queued @p to the SCX_CAP_* bit required to stay on its local DSQ */
126 static inline u64 scx_caps_for_task(struct task_struct *p)
127 {
128 	if (p->scx.flags & SCX_TASK_IMMED)
129 		return SCX_CAP_ENQ_IMMED;
130 	return SCX_CAP_ENQ;
131 }
132 
133 /* the cap @sch needs to preempt @rq's current task, 0 if none */
134 static inline u64 scx_caps_for_preempt(struct scx_sched *sch, struct rq *rq)
135 {
136 	struct task_struct *curr = rq->curr;
137 
138 	/* a non-ext task can't be preempted by ext, own-subtree needs no cap */
139 	if (curr->sched_class != &ext_sched_class ||
140 	    scx_is_descendant(scx_task_sched(curr), sch))
141 		return 0;
142 	return SCX_CAP_PREEMPT;
143 }
144 
145 /* caps implied by holding @cap */
146 static inline u64 scx_caps_implied(u64 cap)
147 {
148 	switch (cap) {
149 	case SCX_CAP_PREEMPT:
150 		return SCX_CAP_ENQ | SCX_CAP_ENQ_IMMED;
151 	case SCX_CAP_ENQ:
152 		return SCX_CAP_ENQ_IMMED;
153 	}
154 	return 0;
155 }
156 
157 /* may @p keep running on @rq's cpu? requires baseline cpu access */
158 static inline bool scx_task_can_stay_on_cpu(struct rq *rq, struct task_struct *p)
159 {
160 	/* a migration-disabled task is let in without caps, keep it likewise */
161 	if (unlikely(is_migration_disabled(p)))
162 		return true;
163 
164 	return likely(!scx_missing_caps(scx_task_sched(p), cpu_of(rq), SCX_CAP_BASE));
165 }
166 
167 #else	/* CONFIG_EXT_SUB_SCHED */
168 
169 static inline u64 scx_missing_caps(struct scx_sched *sch, s32 cpu, u64 needed) { return 0; }
170 static inline u64 scx_caps_for_preempt(struct scx_sched *sch, struct rq *rq) { return 0; }
171 static inline bool scx_task_can_stay_on_cpu(struct rq *rq, struct task_struct *p) { return true; }
172 
173 #endif	/* CONFIG_EXT_SUB_SCHED */
174 
175 /*
176  * One user of this function is scx_bpf_dispatch() which can be called
177  * recursively as sub-sched dispatches nest. Always inline to reduce stack usage
178  * from the call frame.
179  */
180 static __always_inline bool
181 scx_dispatch_sched(struct scx_sched *sch, struct rq *rq,
182 		   struct task_struct *prev, bool nested)
183 {
184 	struct scx_dsp_ctx *dspc = &this_cpu_ptr(sch->pcpu)->dsp_ctx;
185 	int nr_loops = SCX_DSP_MAX_LOOPS;
186 	s32 cpu = cpu_of(rq);
187 	bool prev_on_sch = (prev->sched_class == &ext_sched_class) &&
188 		scx_task_on_sched(sch, prev);
189 
190 	if (scx_consume_global_dsq(sch, rq))
191 		return true;
192 
193 	if (scx_bypass_dsp_enabled(sch)) {
194 		/* if @sch is bypassing, only the bypass DSQs are active */
195 		if (scx_bypassing(sch, cpu))
196 			return scx_consume_dispatch_q(sch, rq, scx_bypass_dsq(sch, cpu), 0);
197 
198 #ifdef CONFIG_EXT_SUB_SCHED
199 		/*
200 		 * If @sch isn't bypassing but its children are, @sch is
201 		 * responsible for making forward progress for both its own
202 		 * tasks that aren't bypassing and the bypassing descendants'
203 		 * tasks. The following implements a simple built-in behavior -
204 		 * let each CPU try to run the bypass DSQ every Nth time.
205 		 *
206 		 * Later, if necessary, we can add an ops flag to suppress the
207 		 * auto-consumption and a kfunc to consume the bypass DSQ and,
208 		 * so that the BPF scheduler can fully control scheduling of
209 		 * bypassed tasks.
210 		 */
211 		struct scx_sched_pcpu *pcpu = per_cpu_ptr(sch->pcpu, cpu);
212 
213 		if (!(pcpu->bypass_host_seq++ % SCX_BYPASS_HOST_NTH) &&
214 		    scx_consume_dispatch_q(sch, rq, scx_bypass_dsq(sch, cpu), 0)) {
215 			__scx_add_event(sch, SCX_EV_SUB_BYPASS_DISPATCH, 1);
216 			return true;
217 		}
218 #endif	/* CONFIG_EXT_SUB_SCHED */
219 	}
220 
221 	if (unlikely(!SCX_HAS_OP(sch, dispatch)) || !scx_rq_online(rq))
222 		return false;
223 
224 	dspc->rq = rq;
225 
226 	/*
227 	 * The dispatch loop. Because scx_flush_dispatch_buf() may drop the rq
228 	 * lock, the local DSQ might still end up empty after a successful
229 	 * ops.dispatch(). If the local DSQ is empty even after ops.dispatch()
230 	 * produced some tasks, retry. The BPF scheduler may depend on this
231 	 * looping behavior to simplify its implementation.
232 	 */
233 	do {
234 		dspc->nr_tasks = 0;
235 
236 		if (nested) {
237 			SCX_CALL_OP(sch, dispatch, rq, scx_cpu_arg(cpu),
238 				    prev_on_sch ? prev : NULL);
239 		} else {
240 			/* stash @prev so that nested invocations can access it */
241 			rq->scx.sub_dispatch_prev = prev;
242 			SCX_CALL_OP(sch, dispatch, rq, scx_cpu_arg(cpu),
243 				    prev_on_sch ? prev : NULL);
244 			rq->scx.sub_dispatch_prev = NULL;
245 		}
246 
247 		scx_flush_dispatch_buf(sch, rq);
248 
249 		if ((prev->scx.flags & SCX_TASK_QUEUED) && prev->scx.slice) {
250 			rq->scx.flags |= SCX_RQ_BAL_KEEP;
251 			return true;
252 		}
253 		if (rq->scx.local_dsq.nr)
254 			return true;
255 		if (scx_consume_global_dsq(sch, rq))
256 			return true;
257 
258 		/*
259 		 * ops.dispatch() can trap us in this loop by repeatedly
260 		 * dispatching ineligible tasks. Break out once in a while to
261 		 * allow the watchdog to run. As IRQ can't be enabled in
262 		 * balance(), we want to complete this scheduling cycle and then
263 		 * start a new one. IOW, we want to call resched_curr() on the
264 		 * next, most likely idle, task, not the current one. Use
265 		 * __scx_bpf_kick_cpu() for deferred kicking.
266 		 */
267 		if (unlikely(!--nr_loops)) {
268 			scx_kick_cpu(sch, cpu, 0);
269 			break;
270 		}
271 	} while (dspc->nr_tasks);
272 
273 	/*
274 	 * Prevent the CPU from going idle while bypassed descendants have tasks
275 	 * queued. Without this fallback, bypassed tasks could stall if the host
276 	 * scheduler's ops.dispatch() doesn't yield any tasks.
277 	 */
278 	if (scx_bypass_dsp_enabled(sch))
279 		return scx_consume_dispatch_q(sch, rq, scx_bypass_dsq(sch, cpu), 0);
280 
281 	return false;
282 }
283 
284 #endif /* _KERNEL_SCHED_EXT_SUB_H */
285