1 // SPDX-License-Identifier: GPL-2.0
2 /*
3 * BPF extensible scheduler class: Documentation/scheduler/sched-ext.rst
4 *
5 * Built-in idle CPU tracking policy.
6 *
7 * Copyright (c) 2022 Meta Platforms, Inc. and affiliates.
8 * Copyright (c) 2022 Tejun Heo <tj@kernel.org>
9 * Copyright (c) 2022 David Vernet <dvernet@meta.com>
10 * Copyright (c) 2024 Andrea Righi <arighi@nvidia.com>
11 */
12 #include "internal.h"
13 #include "cid.h"
14 #include "idle.h"
15 #include "sub.h"
16
17 /* Enable/disable built-in idle CPU selection policy */
18 static DEFINE_STATIC_KEY_FALSE(scx_builtin_idle_enabled);
19
20 /* Enable/disable per-node idle cpumasks */
21 static DEFINE_STATIC_KEY_FALSE(scx_builtin_idle_per_node);
22
23 /* Enable/disable LLC aware optimizations */
24 static DEFINE_STATIC_KEY_FALSE(scx_selcpu_topo_llc);
25
26 /* Enable/disable NUMA aware optimizations */
27 static DEFINE_STATIC_KEY_FALSE(scx_selcpu_topo_numa);
28
29 /*
30 * cpumasks to track idle CPUs within each NUMA node.
31 *
32 * If SCX_OPS_BUILTIN_IDLE_PER_NODE is not enabled, a single global cpumask
33 * from is used to track all the idle CPUs in the system.
34 */
35 struct scx_idle_cpus {
36 cpumask_var_t cpu;
37 cpumask_var_t smt;
38 };
39
40 /*
41 * Global host-wide idle cpumasks (used when SCX_OPS_BUILTIN_IDLE_PER_NODE
42 * is not enabled).
43 */
44 static struct scx_idle_cpus scx_idle_global_masks;
45
46 /*
47 * Per-node idle cpumasks.
48 */
49 static struct scx_idle_cpus **scx_idle_node_masks;
50
51 /*
52 * Local per-CPU cpumasks (used to generate temporary idle cpumasks).
53 */
54 static DEFINE_PER_CPU(cpumask_var_t, local_idle_cpumask);
55 static DEFINE_PER_CPU(cpumask_var_t, local_llc_idle_cpumask);
56 static DEFINE_PER_CPU(cpumask_var_t, local_numa_idle_cpumask);
57
58 /*
59 * Return the idle masks associated to a target @node.
60 *
61 * NUMA_NO_NODE identifies the global idle cpumask.
62 */
idle_cpumask(int node)63 static struct scx_idle_cpus *idle_cpumask(int node)
64 {
65 return node == NUMA_NO_NODE ? &scx_idle_global_masks : scx_idle_node_masks[node];
66 }
67
68 /*
69 * Returns the NUMA node ID associated with a @cpu, or NUMA_NO_NODE if
70 * per-node idle cpumasks are disabled.
71 */
scx_cpu_node_if_enabled(int cpu)72 static int scx_cpu_node_if_enabled(int cpu)
73 {
74 if (!static_branch_maybe(CONFIG_NUMA, &scx_builtin_idle_per_node))
75 return NUMA_NO_NODE;
76
77 return cpu_to_node(cpu);
78 }
79
scx_idle_test_and_clear_cpu(int cpu)80 static bool scx_idle_test_and_clear_cpu(int cpu)
81 {
82 int node = scx_cpu_node_if_enabled(cpu);
83 struct cpumask *idle_cpus = idle_cpumask(node)->cpu;
84
85 /*
86 * SMT mask should be cleared whether we can claim @cpu or not. The SMT
87 * cluster is not wholly idle either way. This also prevents
88 * scx_pick_idle_cpu() from getting caught in an infinite loop.
89 */
90 if (sched_smt_active()) {
91 const struct cpumask *smt = cpu_smt_mask(cpu);
92 struct cpumask *idle_smts = idle_cpumask(node)->smt;
93
94 /*
95 * If offline, @cpu is not its own sibling and
96 * scx_pick_idle_cpu() can get caught in an infinite loop as
97 * @cpu is never cleared from the idle SMT mask. Ensure that
98 * @cpu is eventually cleared.
99 *
100 * NOTE: Use cpumask_intersects() and cpumask_test_cpu() to
101 * reduce memory writes, which may help alleviate cache
102 * coherence pressure.
103 */
104 if (cpumask_intersects(smt, idle_smts))
105 cpumask_andnot(idle_smts, idle_smts, smt);
106 else if (cpumask_test_cpu(cpu, idle_smts))
107 __cpumask_clear_cpu(cpu, idle_smts);
108 }
109
110 return cpumask_test_and_clear_cpu(cpu, idle_cpus);
111 }
112
113 /*
114 * Pick an idle CPU in a specific NUMA node.
115 */
pick_idle_cpu_in_node(const struct cpumask * cpus_allowed,int node,u64 flags)116 static s32 pick_idle_cpu_in_node(const struct cpumask *cpus_allowed, int node, u64 flags)
117 {
118 int cpu;
119
120 retry:
121 if (sched_smt_active()) {
122 cpu = cpumask_any_and_distribute(idle_cpumask(node)->smt, cpus_allowed);
123 if (cpu < nr_cpu_ids)
124 goto found;
125
126 if (flags & SCX_PICK_IDLE_CORE)
127 return -EBUSY;
128 }
129
130 cpu = cpumask_any_and_distribute(idle_cpumask(node)->cpu, cpus_allowed);
131 if (cpu >= nr_cpu_ids)
132 return -EBUSY;
133
134 found:
135 if (scx_idle_test_and_clear_cpu(cpu))
136 return cpu;
137 else
138 goto retry;
139 }
140
141 #ifdef CONFIG_NUMA
142 /*
143 * Tracks nodes that have not yet been visited when searching for an idle
144 * CPU across all available nodes.
145 */
146 static DEFINE_PER_CPU(nodemask_t, per_cpu_unvisited);
147
148 /*
149 * Search for an idle CPU across all nodes, excluding @node.
150 */
pick_idle_cpu_from_online_nodes(const struct cpumask * cpus_allowed,int node,u64 flags)151 static s32 pick_idle_cpu_from_online_nodes(const struct cpumask *cpus_allowed, int node, u64 flags)
152 {
153 nodemask_t *unvisited;
154 s32 cpu = -EBUSY;
155
156 preempt_disable();
157 unvisited = this_cpu_ptr(&per_cpu_unvisited);
158
159 /*
160 * Restrict the search to the online nodes (excluding the current
161 * node that has been visited already).
162 */
163 nodes_copy(*unvisited, node_states[N_ONLINE]);
164 node_clear(node, *unvisited);
165
166 /*
167 * Traverse all nodes in order of increasing distance, starting
168 * from @node.
169 *
170 * This loop is O(N^2), with N being the amount of NUMA nodes,
171 * which might be quite expensive in large NUMA systems. However,
172 * this complexity comes into play only when a scheduler enables
173 * SCX_OPS_BUILTIN_IDLE_PER_NODE and it's requesting an idle CPU
174 * without specifying a target NUMA node, so it shouldn't be a
175 * bottleneck is most cases.
176 *
177 * As a future optimization we may want to cache the list of nodes
178 * in a per-node array, instead of actually traversing them every
179 * time.
180 */
181 for_each_node_numadist(node, *unvisited) {
182 cpu = pick_idle_cpu_in_node(cpus_allowed, node, flags);
183 if (cpu >= 0)
184 break;
185 }
186 preempt_enable();
187
188 return cpu;
189 }
190 #else
191 static inline s32
pick_idle_cpu_from_online_nodes(const struct cpumask * cpus_allowed,int node,u64 flags)192 pick_idle_cpu_from_online_nodes(const struct cpumask *cpus_allowed, int node, u64 flags)
193 {
194 return -EBUSY;
195 }
196 #endif
197
198 /*
199 * Find an idle CPU in the system, starting from @node.
200 */
scx_pick_idle_cpu(const struct cpumask * cpus_allowed,int node,u64 flags)201 static s32 scx_pick_idle_cpu(const struct cpumask *cpus_allowed, int node, u64 flags)
202 {
203 s32 cpu;
204
205 /*
206 * Always search in the starting node first (this is an
207 * optimization that can save some cycles even when the search is
208 * not limited to a single node).
209 */
210 cpu = pick_idle_cpu_in_node(cpus_allowed, node, flags);
211 if (cpu >= 0)
212 return cpu;
213
214 /*
215 * Stop the search if we are using only a single global cpumask
216 * (NUMA_NO_NODE) or if the search is restricted to the first node
217 * only.
218 */
219 if (node == NUMA_NO_NODE || flags & SCX_PICK_IDLE_IN_NODE)
220 return -EBUSY;
221
222 /*
223 * Extend the search to the other online nodes.
224 */
225 return pick_idle_cpu_from_online_nodes(cpus_allowed, node, flags);
226 }
227
228 /*
229 * Return the amount of CPUs in the same LLC domain of @cpu (or zero if the LLC
230 * domain is not defined).
231 */
llc_weight(s32 cpu)232 static unsigned int llc_weight(s32 cpu)
233 {
234 struct sched_domain *sd;
235
236 sd = rcu_dereference(per_cpu(sd_llc, cpu));
237 if (!sd)
238 return 0;
239
240 return sd->span_weight;
241 }
242
243 /*
244 * Return the cpumask representing the LLC domain of @cpu (or NULL if the LLC
245 * domain is not defined).
246 */
llc_span(s32 cpu)247 static struct cpumask *llc_span(s32 cpu)
248 {
249 struct sched_domain *sd;
250
251 sd = rcu_dereference(per_cpu(sd_llc, cpu));
252 if (!sd)
253 return NULL;
254
255 return sched_domain_span(sd);
256 }
257
258 /*
259 * Return the amount of CPUs in the same NUMA domain of @cpu (or zero if the
260 * NUMA domain is not defined).
261 */
numa_weight(s32 cpu)262 static unsigned int numa_weight(s32 cpu)
263 {
264 struct sched_domain *sd;
265 struct sched_group *sg;
266
267 sd = rcu_dereference(per_cpu(sd_numa, cpu));
268 if (!sd)
269 return 0;
270 sg = sd->groups;
271 if (!sg)
272 return 0;
273
274 return sg->group_weight;
275 }
276
277 /*
278 * Return the cpumask representing the NUMA domain of @cpu (or NULL if the NUMA
279 * domain is not defined).
280 */
numa_span(s32 cpu)281 static struct cpumask *numa_span(s32 cpu)
282 {
283 struct sched_domain *sd;
284 struct sched_group *sg;
285
286 sd = rcu_dereference(per_cpu(sd_numa, cpu));
287 if (!sd)
288 return NULL;
289 sg = sd->groups;
290 if (!sg)
291 return NULL;
292
293 return sched_group_span(sg);
294 }
295
296 /*
297 * Return true if the LLC domains do not perfectly overlap with the NUMA
298 * domains, false otherwise.
299 */
llc_numa_mismatch(void)300 static bool llc_numa_mismatch(void)
301 {
302 int cpu;
303
304 /*
305 * We need to scan all online CPUs to verify whether their scheduling
306 * domains overlap.
307 *
308 * While it is rare to encounter architectures with asymmetric NUMA
309 * topologies, CPU hotplugging or virtualized environments can result
310 * in asymmetric configurations.
311 *
312 * For example:
313 *
314 * NUMA 0:
315 * - LLC 0: cpu0..cpu7
316 * - LLC 1: cpu8..cpu15 [offline]
317 *
318 * NUMA 1:
319 * - LLC 0: cpu16..cpu23
320 * - LLC 1: cpu24..cpu31
321 *
322 * In this case, if we only check the first online CPU (cpu0), we might
323 * incorrectly assume that the LLC and NUMA domains are fully
324 * overlapping, which is incorrect (as NUMA 1 has two distinct LLC
325 * domains).
326 */
327 for_each_online_cpu(cpu)
328 if (llc_weight(cpu) != numa_weight(cpu))
329 return true;
330
331 return false;
332 }
333
334 /*
335 * Initialize topology-aware scheduling.
336 *
337 * Detect if the system has multiple LLC or multiple NUMA domains and enable
338 * cache-aware / NUMA-aware scheduling optimizations in the default CPU idle
339 * selection policy.
340 *
341 * Assumption: the kernel's internal topology representation assumes that each
342 * CPU belongs to a single LLC domain, and that each LLC domain is entirely
343 * contained within a single NUMA node.
344 */
scx_idle_update_selcpu_topology(struct sched_ext_ops * ops)345 void scx_idle_update_selcpu_topology(struct sched_ext_ops *ops)
346 {
347 bool enable_llc = false, enable_numa = false;
348 unsigned int nr_cpus;
349 s32 cpu = cpumask_first(cpu_online_mask);
350
351 /*
352 * Enable LLC domain optimization only when there are multiple LLC
353 * domains among the online CPUs. If all online CPUs are part of a
354 * single LLC domain, the idle CPU selection logic can choose any
355 * online CPU without bias.
356 *
357 * Note that it is sufficient to check the LLC domain of the first
358 * online CPU to determine whether a single LLC domain includes all
359 * CPUs.
360 */
361 rcu_read_lock();
362 nr_cpus = llc_weight(cpu);
363 if (nr_cpus > 0) {
364 if (nr_cpus < num_online_cpus())
365 enable_llc = true;
366 pr_debug("sched_ext: LLC=%*pb weight=%u\n",
367 cpumask_pr_args(llc_span(cpu)), llc_weight(cpu));
368 }
369
370 /*
371 * Enable NUMA optimization only when there are multiple NUMA domains
372 * among the online CPUs and the NUMA domains don't perfectly overlap
373 * with the LLC domains.
374 *
375 * If all CPUs belong to the same NUMA node and the same LLC domain,
376 * enabling both NUMA and LLC optimizations is unnecessary, as checking
377 * for an idle CPU in the same domain twice is redundant.
378 *
379 * If SCX_OPS_BUILTIN_IDLE_PER_NODE is enabled ignore the NUMA
380 * optimization, as we would naturally select idle CPUs within
381 * specific NUMA nodes querying the corresponding per-node cpumask.
382 */
383 if (!(ops->flags & SCX_OPS_BUILTIN_IDLE_PER_NODE)) {
384 nr_cpus = numa_weight(cpu);
385 if (nr_cpus > 0) {
386 if (nr_cpus < num_online_cpus() && llc_numa_mismatch())
387 enable_numa = true;
388 pr_debug("sched_ext: NUMA=%*pb weight=%u\n",
389 cpumask_pr_args(numa_span(cpu)), nr_cpus);
390 }
391 }
392 rcu_read_unlock();
393
394 pr_debug("sched_ext: LLC idle selection %s\n",
395 str_enabled_disabled(enable_llc));
396 pr_debug("sched_ext: NUMA idle selection %s\n",
397 str_enabled_disabled(enable_numa));
398
399 if (enable_llc)
400 static_branch_enable_cpuslocked(&scx_selcpu_topo_llc);
401 else
402 static_branch_disable_cpuslocked(&scx_selcpu_topo_llc);
403 if (enable_numa)
404 static_branch_enable_cpuslocked(&scx_selcpu_topo_numa);
405 else
406 static_branch_disable_cpuslocked(&scx_selcpu_topo_numa);
407 }
408
409 /*
410 * Return true if @p can run on all possible CPUs, false otherwise.
411 */
task_affinity_all(const struct task_struct * p)412 static inline bool task_affinity_all(const struct task_struct *p)
413 {
414 return p->nr_cpus_allowed >= num_possible_cpus();
415 }
416
417 /*
418 * Built-in CPU idle selection policy:
419 *
420 * 1. Prioritize full-idle cores:
421 * - always prioritize CPUs from fully idle cores (both logical CPUs are
422 * idle) to avoid interference caused by SMT.
423 *
424 * 2. Reuse the same CPU:
425 * - prefer the last used CPU to take advantage of cached data (L1, L2) and
426 * branch prediction optimizations.
427 *
428 * 3. Prefer @prev_cpu's SMT sibling:
429 * - if @prev_cpu is busy and no fully idle core is available, try to
430 * place the task on an idle SMT sibling of @prev_cpu; keeping the
431 * task on the same core makes migration cheaper, preserves L1 cache
432 * locality and reduces wakeup latency.
433 *
434 * 4. Pick a CPU within the same LLC (Last-Level Cache):
435 * - if the above conditions aren't met, pick a CPU that shares the same
436 * LLC, if the LLC domain is a subset of @cpus_allowed, to maintain
437 * cache locality.
438 *
439 * 5. Pick a CPU within the same NUMA node, if enabled:
440 * - choose a CPU from the same NUMA node, if the node cpumask is a
441 * subset of @cpus_allowed, to reduce memory access latency.
442 *
443 * 6. Pick any idle CPU within the @cpus_allowed domain.
444 *
445 * Step 4 and 5 are performed only if the system has, respectively,
446 * multiple LLCs / multiple NUMA nodes (see scx_selcpu_topo_llc and
447 * scx_selcpu_topo_numa) and they don't contain the same subset of CPUs.
448 *
449 * If %SCX_OPS_BUILTIN_IDLE_PER_NODE is enabled, the search will always
450 * begin in @prev_cpu's node and proceed to other nodes in order of
451 * increasing distance.
452 *
453 * Return the picked CPU if idle, or a negative value otherwise.
454 *
455 * NOTE: tasks that can only run on 1 CPU are excluded by this logic, because
456 * we never call ops.select_cpu() for them, see select_task_rq().
457 */
scx_select_cpu_dfl(struct task_struct * p,s32 prev_cpu,u64 wake_flags,const struct cpumask * cpus_allowed,u64 flags)458 s32 scx_select_cpu_dfl(struct task_struct *p, s32 prev_cpu, u64 wake_flags,
459 const struct cpumask *cpus_allowed, u64 flags)
460 {
461 const struct cpumask *llc_cpus = NULL, *numa_cpus = NULL;
462 const struct cpumask *allowed = cpus_allowed ?: p->cpus_ptr;
463 int node = scx_cpu_node_if_enabled(prev_cpu);
464 bool is_prev_allowed;
465 s32 cpu;
466
467 preempt_disable();
468
469 /*
470 * Determine the subset of CPUs usable by @p within @cpus_allowed.
471 */
472 if (allowed != p->cpus_ptr) {
473 struct cpumask *local_cpus = this_cpu_cpumask_var_ptr(local_idle_cpumask);
474
475 if (task_affinity_all(p)) {
476 allowed = cpus_allowed;
477 } else if (cpumask_and(local_cpus, cpus_allowed, p->cpus_ptr)) {
478 allowed = local_cpus;
479 } else {
480 cpu = -EBUSY;
481 goto out_enable;
482 }
483 }
484
485 /*
486 * Check whether @prev_cpu is still within the allowed set. If not,
487 * we can still try selecting a nearby CPU.
488 */
489 is_prev_allowed = cpumask_test_cpu(prev_cpu, allowed);
490
491 /*
492 * This is necessary to protect llc_cpus.
493 */
494 rcu_read_lock();
495
496 /*
497 * Determine the subset of CPUs that the task can use in its
498 * current LLC and node.
499 *
500 * If the task can run on all CPUs, use the node and LLC cpumasks
501 * directly.
502 */
503 if (static_branch_maybe(CONFIG_NUMA, &scx_selcpu_topo_numa)) {
504 struct cpumask *local_cpus = this_cpu_cpumask_var_ptr(local_numa_idle_cpumask);
505 const struct cpumask *cpus = numa_span(prev_cpu);
506
507 if (allowed == p->cpus_ptr && task_affinity_all(p))
508 numa_cpus = cpus;
509 else if (cpus && cpumask_and(local_cpus, allowed, cpus))
510 numa_cpus = local_cpus;
511 }
512
513 if (static_branch_maybe(CONFIG_SCHED_MC, &scx_selcpu_topo_llc)) {
514 struct cpumask *local_cpus = this_cpu_cpumask_var_ptr(local_llc_idle_cpumask);
515 const struct cpumask *cpus = llc_span(prev_cpu);
516
517 if (allowed == p->cpus_ptr && task_affinity_all(p))
518 llc_cpus = cpus;
519 else if (cpus && cpumask_and(local_cpus, allowed, cpus))
520 llc_cpus = local_cpus;
521 }
522
523 /*
524 * If WAKE_SYNC, try to migrate the wakee to the waker's CPU.
525 */
526 if (wake_flags & SCX_WAKE_SYNC) {
527 int waker_node;
528
529 /*
530 * If the waker's CPU is cache affine and prev_cpu is idle,
531 * then avoid a migration.
532 */
533 cpu = smp_processor_id();
534 if (is_prev_allowed && cpus_share_cache(cpu, prev_cpu) &&
535 scx_idle_test_and_clear_cpu(prev_cpu)) {
536 cpu = prev_cpu;
537 goto out_unlock;
538 }
539
540 /*
541 * If the waker's local DSQ is empty, and the system is under
542 * utilized, try to wake up @p to the local DSQ of the waker.
543 *
544 * Checking only for an empty local DSQ is insufficient as it
545 * could give the wakee an unfair advantage when the system is
546 * oversaturated.
547 *
548 * Checking only for the presence of idle CPUs is also
549 * insufficient as the local DSQ of the waker could have tasks
550 * piled up on it even if there is an idle core elsewhere on
551 * the system.
552 */
553 waker_node = scx_cpu_node_if_enabled(cpu);
554 if (!(current->flags & PF_EXITING) &&
555 cpu_rq(cpu)->scx.local_dsq.nr == 0 &&
556 (!(flags & SCX_PICK_IDLE_IN_NODE) || (waker_node == node)) &&
557 !cpumask_empty(idle_cpumask(waker_node)->cpu)) {
558 if (cpumask_test_cpu(cpu, allowed)) {
559 scx_idle_test_and_clear_cpu(cpu);
560 goto out_unlock;
561 }
562 }
563 }
564
565 /*
566 * If CPU has SMT, any wholly idle CPU is likely a better pick than
567 * partially idle @prev_cpu.
568 */
569 if (sched_smt_active()) {
570 /*
571 * Keep using @prev_cpu if it's part of a fully idle core.
572 */
573 if (is_prev_allowed &&
574 cpumask_test_cpu(prev_cpu, idle_cpumask(node)->smt) &&
575 scx_idle_test_and_clear_cpu(prev_cpu)) {
576 cpu = prev_cpu;
577 goto out_unlock;
578 }
579
580 /*
581 * Search for any fully idle core in the same LLC domain.
582 */
583 if (llc_cpus) {
584 cpu = pick_idle_cpu_in_node(llc_cpus, node, SCX_PICK_IDLE_CORE);
585 if (cpu >= 0)
586 goto out_unlock;
587 }
588
589 /*
590 * Search for any fully idle core in the same NUMA node.
591 */
592 if (numa_cpus) {
593 cpu = pick_idle_cpu_in_node(numa_cpus, node, SCX_PICK_IDLE_CORE);
594 if (cpu >= 0)
595 goto out_unlock;
596 }
597
598 /*
599 * Search for any full-idle core usable by the task.
600 *
601 * If the node-aware idle CPU selection policy is enabled
602 * (%SCX_OPS_BUILTIN_IDLE_PER_NODE), the search will always
603 * begin in prev_cpu's node and proceed to other nodes in
604 * order of increasing distance.
605 */
606 cpu = scx_pick_idle_cpu(allowed, node, flags | SCX_PICK_IDLE_CORE);
607 if (cpu >= 0)
608 goto out_unlock;
609
610 /*
611 * Give up if we're strictly looking for a full-idle SMT
612 * core.
613 */
614 if (flags & SCX_PICK_IDLE_CORE) {
615 cpu = -EBUSY;
616 goto out_unlock;
617 }
618 }
619
620 /*
621 * Use @prev_cpu if it's idle.
622 */
623 if (is_prev_allowed && scx_idle_test_and_clear_cpu(prev_cpu)) {
624 cpu = prev_cpu;
625 goto out_unlock;
626 }
627
628 /*
629 * Use @prev_cpu's sibling if it's idle.
630 */
631 if (sched_smt_active()) {
632 for_each_cpu_and(cpu, cpu_smt_mask(prev_cpu), allowed) {
633 if (cpu == prev_cpu)
634 continue;
635 if (scx_idle_test_and_clear_cpu(cpu))
636 goto out_unlock;
637 }
638 }
639
640 /*
641 * Search for any idle CPU in the same LLC domain.
642 */
643 if (llc_cpus) {
644 cpu = pick_idle_cpu_in_node(llc_cpus, node, 0);
645 if (cpu >= 0)
646 goto out_unlock;
647 }
648
649 /*
650 * Search for any idle CPU in the same NUMA node.
651 */
652 if (numa_cpus) {
653 cpu = pick_idle_cpu_in_node(numa_cpus, node, 0);
654 if (cpu >= 0)
655 goto out_unlock;
656 }
657
658 /*
659 * Search for any idle CPU usable by the task.
660 *
661 * If the node-aware idle CPU selection policy is enabled
662 * (%SCX_OPS_BUILTIN_IDLE_PER_NODE), the search will always begin
663 * in prev_cpu's node and proceed to other nodes in order of
664 * increasing distance.
665 */
666 cpu = scx_pick_idle_cpu(allowed, node, flags);
667
668 out_unlock:
669 rcu_read_unlock();
670 out_enable:
671 preempt_enable();
672
673 return cpu;
674 }
675
676 /*
677 * Initialize global and per-node idle cpumasks.
678 */
scx_idle_init_masks(void)679 void scx_idle_init_masks(void)
680 {
681 int i;
682
683 /* Allocate global idle cpumasks */
684 BUG_ON(!alloc_cpumask_var(&scx_idle_global_masks.cpu, GFP_KERNEL));
685 BUG_ON(!alloc_cpumask_var(&scx_idle_global_masks.smt, GFP_KERNEL));
686
687 /* Allocate per-node idle cpumasks (use nr_node_ids for non-contiguous NUMA nodes) */
688 scx_idle_node_masks = kzalloc_objs(*scx_idle_node_masks, nr_node_ids);
689 BUG_ON(!scx_idle_node_masks);
690
691 for_each_node(i) {
692 scx_idle_node_masks[i] = kzalloc_node(sizeof(**scx_idle_node_masks),
693 GFP_KERNEL, i);
694 BUG_ON(!scx_idle_node_masks[i]);
695
696 BUG_ON(!alloc_cpumask_var_node(&scx_idle_node_masks[i]->cpu, GFP_KERNEL, i));
697 BUG_ON(!alloc_cpumask_var_node(&scx_idle_node_masks[i]->smt, GFP_KERNEL, i));
698 }
699
700 /* Allocate local per-cpu idle cpumasks */
701 for_each_possible_cpu(i) {
702 BUG_ON(!alloc_cpumask_var_node(&per_cpu(local_idle_cpumask, i),
703 GFP_KERNEL, cpu_to_node(i)));
704 BUG_ON(!alloc_cpumask_var_node(&per_cpu(local_llc_idle_cpumask, i),
705 GFP_KERNEL, cpu_to_node(i)));
706 BUG_ON(!alloc_cpumask_var_node(&per_cpu(local_numa_idle_cpumask, i),
707 GFP_KERNEL, cpu_to_node(i)));
708 }
709 }
710
update_builtin_idle(int cpu,bool idle)711 static void update_builtin_idle(int cpu, bool idle)
712 {
713 int node = scx_cpu_node_if_enabled(cpu);
714 struct cpumask *idle_cpus = idle_cpumask(node)->cpu;
715
716 assign_cpu(cpu, idle_cpus, idle);
717
718 if (sched_smt_active()) {
719 const struct cpumask *smt = cpu_smt_mask(cpu);
720 struct cpumask *idle_smts = idle_cpumask(node)->smt;
721
722 if (idle) {
723 /*
724 * idle_smt handling is racy but that's fine as it's
725 * only for optimization and self-correcting.
726 */
727 if (!cpumask_subset(smt, idle_cpus))
728 return;
729 cpumask_or(idle_smts, idle_smts, smt);
730 } else {
731 cpumask_andnot(idle_smts, idle_smts, smt);
732 }
733 }
734 }
735
736 /*
737 * Notify schedulers of an idle transition on @cpu's cid, delivering to every
738 * sched that holds %SCX_CAP_BASE on the cid (the root holds every cap). A real
739 * transition (@do_notify) reaches all holders. A forced one (@root_renotify for
740 * the root, a sub-sched's idle_renotify marker for a sub) reaches only the owed
741 * scheds.
742 */
scx_idle_notify(struct rq * rq,bool idle,bool do_notify,bool root_renotify)743 static void scx_idle_notify(struct rq *rq, bool idle, bool do_notify, bool root_renotify)
744 {
745 s32 cpu = cpu_of(rq);
746 s32 cid = scx_cpu_arg(cpu);
747 struct scx_sched *root = scx_root_protected_live();
748 struct scx_sched *pos;
749
750 lockdep_assert_rq_held(rq);
751
752 /* with no sub-sched, only the root can be owed a notification */
753 if (!scx_has_subs()) {
754 if ((do_notify || root_renotify) &&
755 SCX_HAS_OP(root, update_idle) && !scx_bypassing(root, cpu))
756 SCX_CALL_OP(root, update_idle, rq, cid, idle);
757 return;
758 }
759
760 pos = scx_next_descendant_pre(NULL, root);
761 while (pos) {
762 bool forced = false;
763
764 if (unlikely(scx_missing_caps(pos, cpu, SCX_CAP_BASE))) {
765 pos = scx_skip_subtree_pre(pos, root);
766 continue;
767 }
768
769 if (!pos->level) {
770 forced = root_renotify;
771 }
772 #ifdef CONFIG_EXT_SUB_SCHED
773 else if (per_cpu_ptr(pos->pcpu, cpu)->idle_renotify) {
774 per_cpu_ptr(pos->pcpu, cpu)->idle_renotify = false;
775 forced = true;
776 }
777 #endif
778 if ((do_notify || forced) && SCX_HAS_OP(pos, update_idle) &&
779 !scx_bypassing(pos, cpu))
780 SCX_CALL_OP(pos, update_idle, rq, cid, idle);
781 pos = scx_next_descendant_pre(pos, root);
782 }
783 }
784
785 /*
786 * Update the idle state of a CPU to @idle.
787 *
788 * If @do_notify is true, ops.update_idle() is invoked to notify the scx
789 * scheduler of an actual idle state transition (idle to busy or vice
790 * versa). If @do_notify is false, only the idle state in the idle masks is
791 * refreshed without invoking ops.update_idle().
792 *
793 * This distinction is necessary, because an idle CPU can be "reserved" and
794 * awakened via scx_bpf_pick_idle_cpu() + scx_bpf_kick_cpu(), marking it as
795 * busy even if no tasks are dispatched. In this case, the CPU may return
796 * to idle without a true state transition. Refreshing the idle masks
797 * without invoking ops.update_idle() ensures accurate idle state tracking
798 * while avoiding unnecessary updates and maintaining balanced state
799 * transitions.
800 */
__scx_update_idle(struct rq * rq,bool idle,bool do_notify)801 void __scx_update_idle(struct rq *rq, bool idle, bool do_notify)
802 {
803 int cpu = cpu_of(rq);
804
805 lockdep_assert_rq_held(rq);
806
807 /*
808 * pick_task_idle() calls here only on an idle-to-idle re-pick and the
809 * transitions call with @do_notify, so every reaching call updates the
810 * masks.
811 */
812 if (static_branch_likely(&scx_builtin_idle_enabled))
813 update_builtin_idle(cpu, idle);
814
815 /*
816 * ops.update_idle() fires on real idle transitions, indicated by
817 * @do_notify and managed by put_prev_task_idle()/set_next_task_idle().
818 * An idle pick also fires it to flush a forced notify owed to a sched
819 * that missed transitions while bypassed or on a cid it just gained.
820 * unbypass_renotify_idle() and scx_process_sync_ecaps() arm the per-rq
821 * gates, and scx_idle_notify() targets the owed scheds.
822 *
823 * This must come after the builtin idle update so that BPF schedulers
824 * can create interlocking between ops.update_idle() and ops.enqueue() -
825 * either enqueue() sees the idle bit or update_idle() sees the task
826 * that enqueue() queued.
827 */
828 if (do_notify ||
829 (idle && (rq->scx.flags &
830 (SCX_RQ_SUB_IDLE_RENOTIFY | SCX_RQ_ROOT_IDLE_RENOTIFY)))) {
831 bool root_renotify = rq->scx.flags & SCX_RQ_ROOT_IDLE_RENOTIFY;
832
833 rq->scx.flags &= ~(SCX_RQ_SUB_IDLE_RENOTIFY | SCX_RQ_ROOT_IDLE_RENOTIFY);
834 scx_idle_notify(rq, idle, do_notify, root_renotify);
835 }
836 }
837
reset_idle_masks(struct sched_ext_ops * ops)838 static void reset_idle_masks(struct sched_ext_ops *ops)
839 {
840 int node;
841
842 /*
843 * Start with all CPUs marked busy. The idle masks are populated when
844 * bypass is lifted and each idle CPU is forced through an idle re-pick.
845 * This may temporarily omit idle CPUs but never advertises a busy CPU as
846 * idle.
847 */
848 if (!(ops->flags & SCX_OPS_BUILTIN_IDLE_PER_NODE)) {
849 cpumask_clear(idle_cpumask(NUMA_NO_NODE)->cpu);
850 cpumask_clear(idle_cpumask(NUMA_NO_NODE)->smt);
851 return;
852 }
853
854 for_each_node(node) {
855 cpumask_clear(idle_cpumask(node)->cpu);
856 cpumask_clear(idle_cpumask(node)->smt);
857 }
858 }
859
scx_idle_enable(struct sched_ext_ops * ops)860 void scx_idle_enable(struct sched_ext_ops *ops)
861 {
862 if (!ops->update_idle || (ops->flags & SCX_OPS_KEEP_BUILTIN_IDLE))
863 static_branch_enable_cpuslocked(&scx_builtin_idle_enabled);
864 else
865 static_branch_disable_cpuslocked(&scx_builtin_idle_enabled);
866
867 if (ops->flags & SCX_OPS_BUILTIN_IDLE_PER_NODE)
868 static_branch_enable_cpuslocked(&scx_builtin_idle_per_node);
869 else
870 static_branch_disable_cpuslocked(&scx_builtin_idle_per_node);
871
872 reset_idle_masks(ops);
873 }
874
scx_idle_disable(void)875 void scx_idle_disable(void)
876 {
877 static_branch_disable(&scx_builtin_idle_enabled);
878 static_branch_disable(&scx_builtin_idle_per_node);
879 }
880
881 /********************************************************************************
882 * Helpers that can be called from the BPF scheduler.
883 */
884
validate_node(struct scx_sched * sch,int node)885 static int validate_node(struct scx_sched *sch, int node)
886 {
887 if (!static_branch_likely(&scx_builtin_idle_per_node)) {
888 scx_error(sch, "per-node idle tracking is disabled");
889 return -EOPNOTSUPP;
890 }
891
892 /* Return no entry for NUMA_NO_NODE (not a critical scx error) */
893 if (node == NUMA_NO_NODE)
894 return -ENOENT;
895
896 /* Make sure node is in a valid range */
897 if (node < 0 || node >= nr_node_ids) {
898 scx_error(sch, "invalid node %d", node);
899 return -EINVAL;
900 }
901
902 /* Make sure the node is part of the set of possible nodes */
903 if (!node_possible(node)) {
904 scx_error(sch, "unavailable node %d", node);
905 return -EINVAL;
906 }
907
908 return node;
909 }
910
911 __bpf_kfunc_start_defs();
912
check_builtin_idle_enabled(struct scx_sched * sch)913 static bool check_builtin_idle_enabled(struct scx_sched *sch)
914 {
915 if (static_branch_likely(&scx_builtin_idle_enabled))
916 return true;
917
918 scx_error(sch, "built-in idle tracking is disabled");
919 return false;
920 }
921
922 /*
923 * Determine whether @p is a migration-disabled task in the context of BPF
924 * code.
925 *
926 * We can't simply check whether @p->migration_disabled is set in a
927 * sched_ext callback, because the BPF prolog (__bpf_prog_enter) may disable
928 * migration for the current task while running BPF code.
929 *
930 * Since the BPF prolog calls migrate_disable() only when CONFIG_PREEMPT_RCU
931 * is enabled (via rcu_read_lock_dont_migrate()), migration_disabled == 1 for
932 * the current task is ambiguous only in that case: it could be from the BPF
933 * prolog rather than a real migrate_disable() call.
934 *
935 * Without CONFIG_PREEMPT_RCU, the BPF prolog never calls migrate_disable(),
936 * so migration_disabled == 1 always means the task is truly
937 * migration-disabled.
938 *
939 * Therefore, when migration_disabled == 1 and CONFIG_PREEMPT_RCU is enabled,
940 * check whether @p is the current task or not: if it is, then migration was
941 * not disabled before entering the callback, otherwise migration was disabled.
942 *
943 * Returns true if @p is migration-disabled, false otherwise.
944 */
is_bpf_migration_disabled(const struct task_struct * p)945 static bool is_bpf_migration_disabled(const struct task_struct *p)
946 {
947 if (p->migration_disabled == 1) {
948 if (IS_ENABLED(CONFIG_PREEMPT_RCU))
949 return p != current;
950 return true;
951 }
952 return p->migration_disabled;
953 }
954
select_cpu_from_kfunc(struct scx_sched * sch,struct task_struct * p,s32 prev_cpu,u64 wake_flags,const struct cpumask * allowed,u64 flags)955 static s32 select_cpu_from_kfunc(struct scx_sched *sch, struct task_struct *p,
956 s32 prev_cpu, u64 wake_flags,
957 const struct cpumask *allowed, u64 flags)
958 {
959 unsigned long irq_flags;
960 bool we_locked = false;
961 s32 cpu;
962
963 if (!scx_cpu_valid(sch, prev_cpu, NULL))
964 return -EINVAL;
965
966 if (!check_builtin_idle_enabled(sch))
967 return -EBUSY;
968
969 /*
970 * Accessing p->cpus_ptr / p->nr_cpus_allowed needs either @p's rq
971 * lock or @p's pi_lock. Three cases:
972 *
973 * - inside ops.select_cpu(): try_to_wake_up() holds the wake-up
974 * task's pi_lock; the wake-up task is recorded in kf_tasks[0]
975 * by SCX_CALL_OP_TASK_RET().
976 * - other rq-locked SCX op: scx_locked_rq() points at the held rq.
977 * - truly unlocked (UNLOCKED ops, SYSCALL, non-SCX struct_ops):
978 * nothing held, take pi_lock ourselves.
979 *
980 * In the first two cases, BPF schedulers may pass an arbitrary task
981 * that the held lock doesn't cover. Refuse those.
982 */
983 if (this_rq()->scx.in_select_cpu) {
984 if (!scx_kf_arg_task_ok(sch, p))
985 return -EINVAL;
986 lockdep_assert_held(&p->pi_lock);
987 } else if (scx_locked_rq()) {
988 if (task_rq(p) != scx_locked_rq())
989 goto cross_task;
990 } else {
991 raw_spin_lock_irqsave(&p->pi_lock, irq_flags);
992 we_locked = true;
993 }
994
995 /*
996 * This may also be called from ops.enqueue(), so we need to handle
997 * per-CPU tasks as well. For these tasks, we can skip all idle CPU
998 * selection optimizations and simply check whether the previously
999 * used CPU is idle and within the allowed cpumask.
1000 */
1001 if (p->nr_cpus_allowed == 1 || is_bpf_migration_disabled(p)) {
1002 if (cpumask_test_cpu(prev_cpu, allowed ?: p->cpus_ptr) &&
1003 scx_idle_test_and_clear_cpu(prev_cpu))
1004 cpu = prev_cpu;
1005 else
1006 cpu = -EBUSY;
1007 } else {
1008 cpu = scx_select_cpu_dfl(p, prev_cpu, wake_flags,
1009 allowed ?: p->cpus_ptr, flags);
1010 }
1011
1012 if (we_locked)
1013 raw_spin_unlock_irqrestore(&p->pi_lock, irq_flags);
1014
1015 return cpu;
1016
1017 cross_task:
1018 scx_error(sch, "select_cpu kfunc called cross-task on %s[%d]",
1019 p->comm, p->pid);
1020 return -EINVAL;
1021 }
1022
1023 /**
1024 * scx_bpf_cpu_node - Return the NUMA node the given @cpu belongs to, or
1025 * trigger an error if @cpu is invalid
1026 * @cpu: target CPU
1027 * @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
1028 */
scx_bpf_cpu_node(s32 cpu,const struct bpf_prog_aux * aux)1029 __bpf_kfunc s32 scx_bpf_cpu_node(s32 cpu, const struct bpf_prog_aux *aux)
1030 {
1031 struct scx_sched *sch;
1032
1033 guard(rcu)();
1034
1035 sch = scx_prog_sched(aux);
1036 if (unlikely(!sch) || !scx_cpu_valid(sch, cpu, NULL))
1037 return NUMA_NO_NODE;
1038 return cpu_to_node(cpu);
1039 }
1040
1041 /**
1042 * scx_bpf_select_cpu_dfl - The default implementation of ops.select_cpu()
1043 * @p: task_struct to select a CPU for
1044 * @prev_cpu: CPU @p was on previously
1045 * @wake_flags: %SCX_WAKE_* flags
1046 * @is_idle: out parameter indicating whether the returned CPU is idle
1047 * @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
1048 *
1049 * Can be called from ops.select_cpu(), ops.enqueue(), or from an unlocked
1050 * context such as a BPF test_run() call, as long as built-in CPU selection
1051 * is enabled: ops.update_idle() is missing or %SCX_OPS_KEEP_BUILTIN_IDLE
1052 * is set.
1053 *
1054 * Returns the picked CPU with *@is_idle indicating whether the picked CPU is
1055 * currently idle and thus a good candidate for direct dispatching.
1056 */
scx_bpf_select_cpu_dfl(struct task_struct * p,s32 prev_cpu,u64 wake_flags,bool * is_idle,const struct bpf_prog_aux * aux)1057 __bpf_kfunc s32 scx_bpf_select_cpu_dfl(struct task_struct *p, s32 prev_cpu,
1058 u64 wake_flags, bool *is_idle,
1059 const struct bpf_prog_aux *aux)
1060 {
1061 struct scx_sched *sch;
1062 s32 cpu;
1063
1064 guard(rcu)();
1065
1066 sch = scx_prog_sched(aux);
1067 if (unlikely(!sch))
1068 return -ENODEV;
1069
1070 cpu = select_cpu_from_kfunc(sch, p, prev_cpu, wake_flags, NULL, 0);
1071 if (cpu >= 0) {
1072 *is_idle = true;
1073 return cpu;
1074 }
1075 *is_idle = false;
1076 return prev_cpu;
1077 }
1078
1079 struct scx_bpf_select_cpu_and_args {
1080 /* @p and @cpus_allowed can't be packed together as KF_RCU is not transitive */
1081 s32 prev_cpu;
1082 u64 wake_flags;
1083 u64 flags;
1084 };
1085
1086 /**
1087 * __scx_bpf_select_cpu_and - Arg-wrapped CPU selection with cpumask
1088 * @p: task_struct to select a CPU for
1089 * @cpus_allowed: cpumask of allowed CPUs
1090 * @args: struct containing the rest of the arguments
1091 * @args->prev_cpu: CPU @p was on previously
1092 * @args->wake_flags: %SCX_WAKE_* flags
1093 * @args->flags: %SCX_PICK_IDLE* flags
1094 * @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
1095 *
1096 * Wrapper kfunc that takes arguments via struct to work around BPF's 5 argument
1097 * limit. BPF programs should use scx_bpf_select_cpu_and() which is provided
1098 * as an inline wrapper in common.bpf.h.
1099 *
1100 * Can be called from ops.select_cpu(), ops.enqueue(), or from an unlocked
1101 * context such as a BPF test_run() call, as long as built-in CPU selection
1102 * is enabled: ops.update_idle() is missing or %SCX_OPS_KEEP_BUILTIN_IDLE
1103 * is set.
1104 *
1105 * @p, @args->prev_cpu and @args->wake_flags match ops.select_cpu().
1106 *
1107 * Returns the selected idle CPU, which will be automatically awakened upon
1108 * returning from ops.select_cpu() and can be used for direct dispatch, or
1109 * a negative value if no idle CPU is available.
1110 */
1111 __bpf_kfunc s32
__scx_bpf_select_cpu_and(struct task_struct * p,const struct cpumask * cpus_allowed,struct scx_bpf_select_cpu_and_args * args,const struct bpf_prog_aux * aux)1112 __scx_bpf_select_cpu_and(struct task_struct *p, const struct cpumask *cpus_allowed,
1113 struct scx_bpf_select_cpu_and_args *args,
1114 const struct bpf_prog_aux *aux)
1115 {
1116 struct scx_sched *sch;
1117
1118 guard(rcu)();
1119
1120 sch = scx_prog_sched(aux);
1121 if (unlikely(!sch))
1122 return -ENODEV;
1123
1124 return select_cpu_from_kfunc(sch, p, args->prev_cpu, args->wake_flags,
1125 cpus_allowed, args->flags);
1126 }
1127
1128 /*
1129 * COMPAT: Will be removed in v6.22.
1130 */
scx_bpf_select_cpu_and(struct task_struct * p,s32 prev_cpu,u64 wake_flags,const struct cpumask * cpus_allowed,u64 flags)1131 __bpf_kfunc s32 scx_bpf_select_cpu_and(struct task_struct *p, s32 prev_cpu, u64 wake_flags,
1132 const struct cpumask *cpus_allowed, u64 flags)
1133 {
1134 struct scx_sched *sch;
1135
1136 guard(rcu)();
1137
1138 sch = rcu_dereference(scx_root);
1139 if (unlikely(!sch))
1140 return -ENODEV;
1141
1142 #ifdef CONFIG_EXT_SUB_SCHED
1143 /*
1144 * Disallow if any sub-scheds are attached. There is no way to tell
1145 * which scheduler called us, just error out @p's scheduler.
1146 */
1147 if (unlikely(!list_empty(&sch->children))) {
1148 scx_error(scx_task_sched(p), "__scx_bpf_select_cpu_and() must be used");
1149 return -EINVAL;
1150 }
1151 #endif
1152
1153 return select_cpu_from_kfunc(sch, p, prev_cpu, wake_flags,
1154 cpus_allowed, flags);
1155 }
1156
1157 /**
1158 * scx_bpf_get_idle_cpumask_node - Get a referenced kptr to the
1159 * idle-tracking per-CPU cpumask of a target NUMA node.
1160 * @node: target NUMA node
1161 * @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
1162 *
1163 * Returns an empty cpumask if idle tracking is not enabled, if @node is
1164 * not valid, or running on a UP kernel. In this case the actual error will
1165 * be reported to the BPF scheduler via scx_error().
1166 */
1167 __bpf_kfunc const struct cpumask *
scx_bpf_get_idle_cpumask_node(s32 node,const struct bpf_prog_aux * aux)1168 scx_bpf_get_idle_cpumask_node(s32 node, const struct bpf_prog_aux *aux)
1169 {
1170 struct scx_sched *sch;
1171
1172 guard(rcu)();
1173
1174 sch = scx_prog_sched(aux);
1175 if (unlikely(!sch))
1176 return cpu_none_mask;
1177
1178 node = validate_node(sch, node);
1179 if (node < 0)
1180 return cpu_none_mask;
1181
1182 return idle_cpumask(node)->cpu;
1183 }
1184
1185 /**
1186 * scx_bpf_get_idle_cpumask - Get a referenced kptr to the idle-tracking
1187 * per-CPU cpumask.
1188 * @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
1189 *
1190 * Returns an empty mask if idle tracking is not enabled, or running on a
1191 * UP kernel.
1192 */
scx_bpf_get_idle_cpumask(const struct bpf_prog_aux * aux)1193 __bpf_kfunc const struct cpumask *scx_bpf_get_idle_cpumask(const struct bpf_prog_aux *aux)
1194 {
1195 struct scx_sched *sch;
1196
1197 guard(rcu)();
1198
1199 sch = scx_prog_sched(aux);
1200 if (unlikely(!sch))
1201 return cpu_none_mask;
1202
1203 if (static_branch_unlikely(&scx_builtin_idle_per_node)) {
1204 scx_error(sch, "SCX_OPS_BUILTIN_IDLE_PER_NODE enabled");
1205 return cpu_none_mask;
1206 }
1207
1208 if (!check_builtin_idle_enabled(sch))
1209 return cpu_none_mask;
1210
1211 return idle_cpumask(NUMA_NO_NODE)->cpu;
1212 }
1213
1214 /**
1215 * scx_bpf_get_idle_smtmask_node - Get a referenced kptr to the
1216 * idle-tracking, per-physical-core cpumask of a target NUMA node. Can be
1217 * used to determine if an entire physical core is free.
1218 * @node: target NUMA node
1219 * @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
1220 *
1221 * Returns an empty cpumask if idle tracking is not enabled, if @node is
1222 * not valid, or running on a UP kernel. In this case the actual error will
1223 * be reported to the BPF scheduler via scx_error().
1224 */
1225 __bpf_kfunc const struct cpumask *
scx_bpf_get_idle_smtmask_node(s32 node,const struct bpf_prog_aux * aux)1226 scx_bpf_get_idle_smtmask_node(s32 node, const struct bpf_prog_aux *aux)
1227 {
1228 struct scx_sched *sch;
1229
1230 guard(rcu)();
1231
1232 sch = scx_prog_sched(aux);
1233 if (unlikely(!sch))
1234 return cpu_none_mask;
1235
1236 node = validate_node(sch, node);
1237 if (node < 0)
1238 return cpu_none_mask;
1239
1240 if (sched_smt_active())
1241 return idle_cpumask(node)->smt;
1242 else
1243 return idle_cpumask(node)->cpu;
1244 }
1245
1246 /**
1247 * scx_bpf_get_idle_smtmask - Get a referenced kptr to the idle-tracking,
1248 * per-physical-core cpumask. Can be used to determine if an entire physical
1249 * core is free.
1250 * @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
1251 *
1252 * Returns an empty mask if idle tracking is not enabled, or running on a
1253 * UP kernel.
1254 */
scx_bpf_get_idle_smtmask(const struct bpf_prog_aux * aux)1255 __bpf_kfunc const struct cpumask *scx_bpf_get_idle_smtmask(const struct bpf_prog_aux *aux)
1256 {
1257 struct scx_sched *sch;
1258
1259 guard(rcu)();
1260
1261 sch = scx_prog_sched(aux);
1262 if (unlikely(!sch))
1263 return cpu_none_mask;
1264
1265 if (static_branch_unlikely(&scx_builtin_idle_per_node)) {
1266 scx_error(sch, "SCX_OPS_BUILTIN_IDLE_PER_NODE enabled");
1267 return cpu_none_mask;
1268 }
1269
1270 if (!check_builtin_idle_enabled(sch))
1271 return cpu_none_mask;
1272
1273 if (sched_smt_active())
1274 return idle_cpumask(NUMA_NO_NODE)->smt;
1275 else
1276 return idle_cpumask(NUMA_NO_NODE)->cpu;
1277 }
1278
1279 /**
1280 * scx_bpf_put_idle_cpumask - Release a previously acquired referenced kptr to
1281 * either the percpu, or SMT idle-tracking cpumask.
1282 * @idle_mask: &cpumask to use
1283 */
scx_bpf_put_idle_cpumask(const struct cpumask * idle_mask)1284 __bpf_kfunc void scx_bpf_put_idle_cpumask(const struct cpumask *idle_mask)
1285 {
1286 /*
1287 * Empty function body because we aren't actually acquiring or releasing
1288 * a reference to a global idle cpumask, which is read-only in the
1289 * caller and is never released. The acquire / release semantics here
1290 * are just used to make the cpumask a trusted pointer in the caller.
1291 */
1292 }
1293
1294 /**
1295 * scx_bpf_test_and_clear_cpu_idle - Test and clear @cpu's idle state
1296 * @cpu: cpu to test and clear idle for
1297 * @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
1298 *
1299 * Returns %true if @cpu was idle and its idle state was successfully cleared.
1300 * %false otherwise.
1301 *
1302 * Unavailable if ops.update_idle() is implemented and
1303 * %SCX_OPS_KEEP_BUILTIN_IDLE is not set.
1304 */
scx_bpf_test_and_clear_cpu_idle(s32 cpu,const struct bpf_prog_aux * aux)1305 __bpf_kfunc bool scx_bpf_test_and_clear_cpu_idle(s32 cpu, const struct bpf_prog_aux *aux)
1306 {
1307 struct scx_sched *sch;
1308
1309 guard(rcu)();
1310
1311 sch = scx_prog_sched(aux);
1312 if (unlikely(!sch))
1313 return false;
1314
1315 if (!check_builtin_idle_enabled(sch))
1316 return false;
1317
1318 if (!scx_cpu_valid(sch, cpu, NULL))
1319 return false;
1320
1321 return scx_idle_test_and_clear_cpu(cpu);
1322 }
1323
1324 /**
1325 * scx_bpf_pick_idle_cpu_node - Pick and claim an idle cpu from @node
1326 * @cpus_allowed: Allowed cpumask
1327 * @node: target NUMA node
1328 * @flags: %SCX_PICK_IDLE_* flags
1329 * @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
1330 *
1331 * Pick and claim an idle cpu in @cpus_allowed from the NUMA node @node.
1332 *
1333 * Returns the picked idle cpu number on success, or -%EBUSY if no matching
1334 * cpu was found.
1335 *
1336 * The search starts from @node and proceeds to other online NUMA nodes in
1337 * order of increasing distance (unless SCX_PICK_IDLE_IN_NODE is specified,
1338 * in which case the search is limited to the target @node).
1339 *
1340 * Always returns an error if ops.update_idle() is implemented and
1341 * %SCX_OPS_KEEP_BUILTIN_IDLE is not set, or if
1342 * %SCX_OPS_BUILTIN_IDLE_PER_NODE is not set.
1343 */
scx_bpf_pick_idle_cpu_node(const struct cpumask * cpus_allowed,s32 node,u64 flags,const struct bpf_prog_aux * aux)1344 __bpf_kfunc s32 scx_bpf_pick_idle_cpu_node(const struct cpumask *cpus_allowed,
1345 s32 node, u64 flags,
1346 const struct bpf_prog_aux *aux)
1347 {
1348 struct scx_sched *sch;
1349
1350 guard(rcu)();
1351
1352 sch = scx_prog_sched(aux);
1353 if (unlikely(!sch))
1354 return -ENODEV;
1355
1356 node = validate_node(sch, node);
1357 if (node < 0)
1358 return node;
1359
1360 return scx_pick_idle_cpu(cpus_allowed, node, flags);
1361 }
1362
1363 /**
1364 * scx_bpf_pick_idle_cpu - Pick and claim an idle cpu
1365 * @cpus_allowed: Allowed cpumask
1366 * @flags: %SCX_PICK_IDLE_* flags
1367 * @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
1368 *
1369 * Pick and claim an idle cpu in @cpus_allowed. Returns the picked idle cpu
1370 * number on success. -%EBUSY if no matching cpu was found.
1371 *
1372 * Idle CPU tracking may race against CPU scheduling state transitions. For
1373 * example, this function may return -%EBUSY as CPUs are transitioning into the
1374 * idle state. If the caller then assumes that there will be dispatch events on
1375 * the CPUs as they were all busy, the scheduler may end up stalling with CPUs
1376 * idling while there are pending tasks. Use scx_bpf_pick_any_cpu() and
1377 * scx_bpf_kick_cpu() to guarantee that there will be at least one dispatch
1378 * event in the near future.
1379 *
1380 * Unavailable if ops.update_idle() is implemented and
1381 * %SCX_OPS_KEEP_BUILTIN_IDLE is not set.
1382 *
1383 * Always returns an error if %SCX_OPS_BUILTIN_IDLE_PER_NODE is set, use
1384 * scx_bpf_pick_idle_cpu_node() instead.
1385 */
scx_bpf_pick_idle_cpu(const struct cpumask * cpus_allowed,u64 flags,const struct bpf_prog_aux * aux)1386 __bpf_kfunc s32 scx_bpf_pick_idle_cpu(const struct cpumask *cpus_allowed,
1387 u64 flags, const struct bpf_prog_aux *aux)
1388 {
1389 struct scx_sched *sch;
1390
1391 guard(rcu)();
1392
1393 sch = scx_prog_sched(aux);
1394 if (unlikely(!sch))
1395 return -ENODEV;
1396
1397 if (static_branch_maybe(CONFIG_NUMA, &scx_builtin_idle_per_node)) {
1398 scx_error(sch, "per-node idle tracking is enabled");
1399 return -EBUSY;
1400 }
1401
1402 if (!check_builtin_idle_enabled(sch))
1403 return -EBUSY;
1404
1405 return scx_pick_idle_cpu(cpus_allowed, NUMA_NO_NODE, flags);
1406 }
1407
1408 /**
1409 * scx_bpf_pick_any_cpu_node - Pick and claim an idle cpu if available
1410 * or pick any CPU from @node
1411 * @cpus_allowed: Allowed cpumask
1412 * @node: target NUMA node
1413 * @flags: %SCX_PICK_IDLE_* flags
1414 * @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
1415 *
1416 * Pick and claim an idle cpu in @cpus_allowed. If none is available, pick any
1417 * CPU in @cpus_allowed. Guaranteed to succeed and returns the picked idle cpu
1418 * number if @cpus_allowed is not empty. -%EBUSY is returned if @cpus_allowed is
1419 * empty.
1420 *
1421 * The search starts from @node and proceeds to other online NUMA nodes in
1422 * order of increasing distance (unless %SCX_PICK_IDLE_IN_NODE is specified,
1423 * in which case the search is limited to the target @node, regardless of
1424 * the CPU idle state).
1425 *
1426 * If ops.update_idle() is implemented and %SCX_OPS_KEEP_BUILTIN_IDLE is not
1427 * set, this function can't tell which CPUs are idle and will always pick any
1428 * CPU.
1429 */
scx_bpf_pick_any_cpu_node(const struct cpumask * cpus_allowed,s32 node,u64 flags,const struct bpf_prog_aux * aux)1430 __bpf_kfunc s32 scx_bpf_pick_any_cpu_node(const struct cpumask *cpus_allowed,
1431 s32 node, u64 flags,
1432 const struct bpf_prog_aux *aux)
1433 {
1434 struct scx_sched *sch;
1435 s32 cpu;
1436
1437 guard(rcu)();
1438
1439 sch = scx_prog_sched(aux);
1440 if (unlikely(!sch))
1441 return -ENODEV;
1442
1443 node = validate_node(sch, node);
1444 if (node < 0)
1445 return node;
1446
1447 cpu = scx_pick_idle_cpu(cpus_allowed, node, flags);
1448 if (cpu >= 0)
1449 return cpu;
1450
1451 if (flags & SCX_PICK_IDLE_IN_NODE)
1452 cpu = cpumask_any_and_distribute(cpumask_of_node(node), cpus_allowed);
1453 else
1454 cpu = cpumask_any_distribute(cpus_allowed);
1455 if (cpu < nr_cpu_ids)
1456 return cpu;
1457 else
1458 return -EBUSY;
1459 }
1460
1461 /**
1462 * scx_bpf_pick_any_cpu - Pick and claim an idle cpu if available or pick any CPU
1463 * @cpus_allowed: Allowed cpumask
1464 * @flags: %SCX_PICK_IDLE_* flags
1465 * @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
1466 *
1467 * Pick and claim an idle cpu in @cpus_allowed. If none is available, pick any
1468 * CPU in @cpus_allowed. Guaranteed to succeed and returns the picked idle cpu
1469 * number if @cpus_allowed is not empty. -%EBUSY is returned if @cpus_allowed is
1470 * empty.
1471 *
1472 * If ops.update_idle() is implemented and %SCX_OPS_KEEP_BUILTIN_IDLE is not
1473 * set, this function can't tell which CPUs are idle and will always pick any
1474 * CPU.
1475 *
1476 * Always returns an error if %SCX_OPS_BUILTIN_IDLE_PER_NODE is set, use
1477 * scx_bpf_pick_any_cpu_node() instead.
1478 */
scx_bpf_pick_any_cpu(const struct cpumask * cpus_allowed,u64 flags,const struct bpf_prog_aux * aux)1479 __bpf_kfunc s32 scx_bpf_pick_any_cpu(const struct cpumask *cpus_allowed,
1480 u64 flags, const struct bpf_prog_aux *aux)
1481 {
1482 struct scx_sched *sch;
1483 s32 cpu;
1484
1485 guard(rcu)();
1486
1487 sch = scx_prog_sched(aux);
1488 if (unlikely(!sch))
1489 return -ENODEV;
1490
1491 if (static_branch_maybe(CONFIG_NUMA, &scx_builtin_idle_per_node)) {
1492 scx_error(sch, "per-node idle tracking is enabled");
1493 return -EBUSY;
1494 }
1495
1496 if (static_branch_likely(&scx_builtin_idle_enabled)) {
1497 cpu = scx_pick_idle_cpu(cpus_allowed, NUMA_NO_NODE, flags);
1498 if (cpu >= 0)
1499 return cpu;
1500 }
1501
1502 cpu = cpumask_any_distribute(cpus_allowed);
1503 if (cpu < nr_cpu_ids)
1504 return cpu;
1505 else
1506 return -EBUSY;
1507 }
1508
1509 __bpf_kfunc_end_defs();
1510
1511 BTF_KFUNCS_START(scx_kfunc_ids_idle)
1512 BTF_ID_FLAGS(func, scx_bpf_cpu_node, KF_IMPLICIT_ARGS)
1513 BTF_ID_FLAGS(func, scx_bpf_get_idle_cpumask_node, KF_IMPLICIT_ARGS | KF_ACQUIRE)
1514 BTF_ID_FLAGS(func, scx_bpf_get_idle_cpumask, KF_IMPLICIT_ARGS | KF_ACQUIRE)
1515 BTF_ID_FLAGS(func, scx_bpf_get_idle_smtmask_node, KF_IMPLICIT_ARGS | KF_ACQUIRE)
1516 BTF_ID_FLAGS(func, scx_bpf_get_idle_smtmask, KF_IMPLICIT_ARGS | KF_ACQUIRE)
1517 BTF_ID_FLAGS(func, scx_bpf_put_idle_cpumask, KF_RELEASE)
1518 BTF_ID_FLAGS(func, scx_bpf_test_and_clear_cpu_idle, KF_IMPLICIT_ARGS)
1519 BTF_ID_FLAGS(func, scx_bpf_pick_idle_cpu_node, KF_IMPLICIT_ARGS | KF_RCU)
1520 BTF_ID_FLAGS(func, scx_bpf_pick_idle_cpu, KF_IMPLICIT_ARGS | KF_RCU)
1521 BTF_ID_FLAGS(func, scx_bpf_pick_any_cpu_node, KF_IMPLICIT_ARGS | KF_RCU)
1522 BTF_ID_FLAGS(func, scx_bpf_pick_any_cpu, KF_IMPLICIT_ARGS | KF_RCU)
1523 BTF_KFUNCS_END(scx_kfunc_ids_idle)
1524
1525 static const struct btf_kfunc_id_set scx_kfunc_set_idle = {
1526 .owner = THIS_MODULE,
1527 .set = &scx_kfunc_ids_idle,
1528 .filter = scx_kfunc_context_filter,
1529 };
1530
1531 /*
1532 * The select_cpu kfuncs internally call task_rq_lock() when invoked from an
1533 * rq-unlocked context, and thus cannot be safely called from arbitrary tracing
1534 * contexts where @p's pi_lock state is unknown. Keep them out of
1535 * BPF_PROG_TYPE_TRACING by registering them in their own set which is exposed
1536 * only to STRUCT_OPS and SYSCALL programs.
1537 *
1538 * These kfuncs are also members of scx_kfunc_ids_unlocked (see ext.c) because
1539 * they're callable from unlocked contexts in addition to ops.select_cpu() and
1540 * ops.enqueue().
1541 */
1542 BTF_KFUNCS_START(scx_kfunc_ids_select_cpu)
1543 BTF_ID_FLAGS(func, __scx_bpf_select_cpu_and, KF_IMPLICIT_ARGS | KF_RCU)
1544 BTF_ID_FLAGS(func, scx_bpf_select_cpu_and, KF_RCU)
1545 BTF_ID_FLAGS(func, scx_bpf_select_cpu_dfl, KF_IMPLICIT_ARGS | KF_RCU)
1546 BTF_KFUNCS_END(scx_kfunc_ids_select_cpu)
1547
1548 static const struct btf_kfunc_id_set scx_kfunc_set_select_cpu = {
1549 .owner = THIS_MODULE,
1550 .set = &scx_kfunc_ids_select_cpu,
1551 .filter = scx_kfunc_context_filter,
1552 };
1553
scx_idle_init(void)1554 int scx_idle_init(void)
1555 {
1556 return register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS, &scx_kfunc_set_idle) ?:
1557 register_btf_kfunc_id_set(BPF_PROG_TYPE_TRACING, &scx_kfunc_set_idle) ?:
1558 register_btf_kfunc_id_set(BPF_PROG_TYPE_SYSCALL, &scx_kfunc_set_idle) ?:
1559 register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS, &scx_kfunc_set_select_cpu) ?:
1560 register_btf_kfunc_id_set(BPF_PROG_TYPE_SYSCALL, &scx_kfunc_set_select_cpu);
1561 }
1562