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