xref: /linux/kernel/sched/ext/sub.h (revision 78f8d726e62e0b72a4b11e5778d2d7e252b076c0)
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 
15 #ifdef CONFIG_EXT_SUB_SCHED
16 
17 struct scx_sched *scx_skip_subtree_pre(struct scx_sched *pos, struct scx_sched *root);
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 scx_pstack_recursion_on_caps_updated(struct bpf_prog *prog);
24 void drain_descendants(struct scx_sched *sch);
25 void scx_sub_disable(struct scx_sched *sch);
26 void scx_sub_enable_workfn(struct kthread_work *work);
27 bool scx_bpf_sub_dispatch(u64 cgroup_id, const struct bpf_prog_aux *aux);
28 void scx_free_pshards(struct scx_sched *sch);
29 s32 scx_alloc_pshards(struct scx_sched *sch);
30 void scx_init_root_caps(struct scx_sched *sch);
31 void scx_process_sync_ecaps(struct rq *rq, struct task_struct *prev);
32 void scx_unbypass_replay_ecaps(struct rq *rq, struct scx_sched *sch);
33 void scx_online_ecaps(struct rq *rq);
34 void scx_offline_ecaps(struct rq *rq);
35 void scx_discard_ecaps_to_sync(s32 cpu, struct scx_sched_pcpu *pcpu);
36 void scx_discard_stale_ecaps_syncs(void);
37 struct scx_dispatch_q *scx_resolve_local_dsq(struct scx_sched *sch, struct rq *rq,
38 					     struct task_struct *p, u64 *enq_flags);
39 bool scx_task_reenq_on_cap_revoke(struct rq *rq, struct task_struct *p);
40 void scx_reenq_reject(struct rq *rq);
41 
42 /*
43  * cgrp->scx_sched is written by root/sub enable/disable under all of
44  * scx_enable_mutex, scx_fork_rwsem and cgroup_mutex. A new cgroup inherits the
45  * parent's sched under just cgroup_mutex but is not yet reachable by the other
46  * two lock holders. Any one of the three locks stabilizes the association.
47  */
48 static inline struct scx_sched *scx_cgroup_sched(struct cgroup *cgrp)
49 {
50 	return rcu_dereference_check(cgrp->scx_sched,
51 				     lockdep_is_held(&cgroup_mutex) ||
52 				     percpu_rwsem_is_held(&scx_fork_rwsem) ||
53 				     lockdep_is_held(&scx_enable_mutex));
54 }
55 
56 static inline const char *sch_cgrp_path(struct scx_sched *sch)
57 {
58 	return sch->cgrp_path;
59 }
60 
61 /* a dying sub's hot-path influence ends in scx_sched_free_rcu_work() */
62 static inline void scx_dec_has_subs(struct scx_sched *sch)
63 {
64 	if (sch->level)
65 		static_branch_dec(&__scx_has_subs);
66 }
67 
68 #else	/* CONFIG_EXT_SUB_SCHED */
69 
70 static inline struct scx_sched *scx_next_descendant_pre(struct scx_sched *pos, struct scx_sched *root) { return pos ? NULL : root; }
71 static inline struct scx_sched *scx_skip_subtree_pre(struct scx_sched *pos, struct scx_sched *root) { return NULL; }
72 static inline void scx_set_task_sched(struct task_struct *p, struct scx_sched *sch) {}
73 static inline struct cgroup *sch_cgroup(struct scx_sched *sch) { return NULL; }
74 static inline const char *sch_cgrp_path(struct scx_sched *sch) { return "/"; }
75 static inline void set_cgroup_sched(struct cgroup *cgrp, struct scx_sched *sch) {}
76 static inline void drain_descendants(struct scx_sched *sch) { }
77 static inline void scx_sub_disable(struct scx_sched *sch) { }
78 static inline void scx_free_pshards(struct scx_sched *sch) {}
79 static inline s32 scx_alloc_pshards(struct scx_sched *sch) { return 0; }
80 static inline void scx_init_root_caps(struct scx_sched *sch) {}
81 static inline void scx_process_sync_ecaps(struct rq *rq, struct task_struct *prev) {}
82 static inline void scx_unbypass_replay_ecaps(struct rq *rq, struct scx_sched *sch) {}
83 static inline void scx_online_ecaps(struct rq *rq) {}
84 static inline void scx_offline_ecaps(struct rq *rq) {}
85 static inline void scx_discard_ecaps_to_sync(s32 cpu, struct scx_sched_pcpu *pcpu) {}
86 static inline void scx_discard_stale_ecaps_syncs(void) {}
87 static inline struct scx_dispatch_q *scx_resolve_local_dsq(struct scx_sched *sch, struct rq *rq, struct task_struct *p, u64 *enq_flags) { return &rq->scx.local_dsq; }
88 static inline bool scx_task_reenq_on_cap_revoke(struct rq *rq, struct task_struct *p) { return false; }
89 static inline void scx_reenq_reject(struct rq *rq) {}
90 static inline void scx_dec_has_subs(struct scx_sched *sch) {}
91 
92 #endif	/* CONFIG_EXT_SUB_SCHED */
93 
94 /**
95  * scx_for_each_descendant_pre - pre-order walk of a sched's descendants
96  * @pos: iteration cursor
97  * @root: sched to walk the descendants of
98  *
99  * Walk @root's descendants. @root is included in the iteration and the first
100  * node to be visited. Must be called with scx_enable_mutex, scx_sched_lock, or
101  * RCU read lock.
102  */
103 #define scx_for_each_descendant_pre(pos, root)					\
104 	for ((pos) = scx_next_descendant_pre(NULL, (root)); (pos);		\
105 	     (pos) = scx_next_descendant_pre((pos), (root)))
106 
107 #ifdef CONFIG_EXT_SUB_SCHED
108 
109 /**
110  * scx_missing_caps - The caps in @needed that @sch lacks on @cpu
111  * @sch: sched to test
112  * @cpu: cpu to test on
113  * @needed: bitmask of SCX_CAP_* values
114  *
115  * Return the caps in @needed that @sch lacks for @cpu, 0 if it holds them all.
116  */
117 static inline u64 scx_missing_caps(struct scx_sched *sch, s32 cpu, u64 needed)
118 {
119 	u64 ecaps;
120 
121 	/* no sub-scheds, no missing caps */
122 	if (!scx_has_subs())
123 		return 0;
124 
125 	/* root holds every cap on every cpu */
126 	if (!sch->level)
127 		return 0;
128 
129 	ecaps = READ_ONCE(per_cpu_ptr(sch->pcpu, cpu)->ecaps);
130 
131 	return needed & ~ecaps;
132 }
133 
134 /*
135  * Cap semantics: which caps an action requires, and which caps a cap implies.
136  * Keep all such mappings collected here.
137  */
138 
139 /* map @enq_flags to the SCX_CAP_* bit required for the local-DSQ insert */
140 static inline u64 scx_caps_for_enq(u64 enq_flags)
141 {
142 	/* a restored task must be put into the local DSQ regardless of caps */
143 	if (unlikely(enq_flags & SCX_ENQ_IGNORE_CAPS))
144 		return 0;
145 	if (enq_flags & SCX_ENQ_IMMED)
146 		return SCX_CAP_ENQ_IMMED;
147 	return SCX_CAP_ENQ;
148 }
149 
150 /* map queued @p to the SCX_CAP_* bit required to stay on its local DSQ */
151 static inline u64 scx_caps_for_task(struct task_struct *p)
152 {
153 	if (p->scx.flags & SCX_TASK_IMMED)
154 		return SCX_CAP_ENQ_IMMED;
155 	return SCX_CAP_ENQ;
156 }
157 
158 /* the cap @sch needs to preempt @rq's current task, 0 if none */
159 static inline u64 scx_caps_for_preempt(struct scx_sched *sch, struct rq *rq, u64 enq_flags)
160 {
161 	struct task_struct *curr = rq->curr;
162 
163 	/* a kernel-forced placement preempts regardless of caps */
164 	if (unlikely(enq_flags & SCX_ENQ_IGNORE_CAPS))
165 		return 0;
166 	/* a non-ext task can't be preempted by ext, own-subtree needs no cap */
167 	if (curr->sched_class != &ext_sched_class ||
168 	    scx_is_descendant(scx_task_sched(curr), sch))
169 		return 0;
170 	return SCX_CAP_PREEMPT;
171 }
172 
173 /* caps implied by holding @cap */
174 static inline u64 scx_caps_implied(u64 cap)
175 {
176 	switch (cap) {
177 	case SCX_CAP_PREEMPT:
178 		return SCX_CAP_ENQ | SCX_CAP_ENQ_IMMED;
179 	case SCX_CAP_ENQ:
180 		return SCX_CAP_ENQ_IMMED;
181 	}
182 	return 0;
183 }
184 
185 /* may @p keep running on @rq's cpu? requires baseline cpu access */
186 static inline bool scx_task_can_stay_on_cpu(struct rq *rq, struct task_struct *p)
187 {
188 	if (!scx_has_subs())
189 		return true;
190 
191 	/* a migration-disabled task is let in without caps, keep it likewise */
192 	if (unlikely(is_migration_disabled(p)))
193 		return true;
194 
195 	return likely(!scx_missing_caps(scx_task_sched(p), cpu_of(rq), SCX_CAP_BASE));
196 }
197 
198 #else	/* CONFIG_EXT_SUB_SCHED */
199 
200 static inline u64 scx_missing_caps(struct scx_sched *sch, s32 cpu, u64 needed) { return 0; }
201 static inline u64 scx_caps_for_preempt(struct scx_sched *sch, struct rq *rq, u64 enq_flags) { return 0; }
202 static inline bool scx_task_can_stay_on_cpu(struct rq *rq, struct task_struct *p) { return true; }
203 
204 #endif	/* CONFIG_EXT_SUB_SCHED */
205 
206 #endif /* _KERNEL_SCHED_EXT_SUB_H */
207