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