xref: /linux/kernel/sched/ext/cid.c (revision 98aec86f8c6bf2e1091fefb99eaba8b910e8eef1)
1 /* SPDX-License-Identifier: GPL-2.0 */
2 /*
3  * BPF extensible scheduler class: Documentation/scheduler/sched-ext.rst
4  *
5  * Copyright (c) 2026 Meta Platforms, Inc. and affiliates.
6  * Copyright (c) 2026 Tejun Heo <tj@kernel.org>
7  */
8 #include <linux/cacheinfo.h>
9 
10 #include "internal.h"
11 #include "cid.h"
12 
13 /*
14  * cid tables. The cid kfuncs are available whether the root scheduler is
15  * cid-form or cpu-form, the latter to allow gradual migration to cids, so every
16  * root builds a default mapping. Each root enable allocates a fresh set, builds
17  * it privately and publishes the __rcu globals below once the layout is final.
18  * Root disable unpublishes and RCU-frees the set. kfuncs may run before the
19  * tables are published and must check for NULL.
20  */
21 u32 scx_nr_cid_shards;
22 s16 __rcu *scx_cid_to_cpu_tbl;
23 s16 __rcu *scx_cpu_to_cid_tbl;
24 s32 __rcu *scx_cid_to_shard;
25 s32 __rcu *scx_shard_node;
26 struct scx_cid_shard __rcu *scx_cid_shard_ranges;
27 struct scx_cid_topo __rcu *scx_cid_topo;
28 
29 static struct scx_cid_tables *scx_cid_tables;	/* used only during alloc/free */
30 
31 #define SCX_CID_TOPO_NEG	(struct scx_cid_topo) {				\
32 	.core_cid = -1, .core_idx = -1, .llc_cid = -1, .llc_idx = -1,		\
33 	.node_cid = -1, .node_idx = -1, .shard_cid = -1, .shard_idx = -1,	\
34 	.cluster_cid = -1, .cluster_idx = -1,					\
35 }
36 
37 /*
38  * Return @cpu's LLC shared_cpu_map. If cacheinfo isn't populated (offline or
39  * !present), record @cpu in @fallbacks and return its node mask instead - the
40  * worst that can happen is that the cpu's LLC becomes coarser than reality.
41  */
42 static const struct cpumask *cpu_llc_mask(int cpu, struct cpumask *fallbacks)
43 {
44 	struct cpu_cacheinfo *ci = get_cpu_cacheinfo(cpu);
45 
46 	if (!ci || !ci->info_list || !ci->num_leaves) {
47 		cpumask_set_cpu(cpu, fallbacks);
48 		return cpumask_of_node(cpu_to_node(cpu));
49 	}
50 	return &ci->info_list[ci->num_leaves - 1].shared_cpu_map;
51 }
52 
53 /*
54  * Compute per-LLC shard layout. Each shard holds at most @shard_size cids, and
55  * in any case no more than SCX_CID_SHARD_MAX_CPUS. Cores are spread as evenly
56  * as possible across shards so cpu count is balanced: the first *@nr_large_p
57  * shards get (*@cores_per_shard_p + 1) cores, the rest get *@cores_per_shard_p.
58  */
59 static void calc_shard_layout(const struct cpumask *llc_cpus, u32 shard_size,
60 			      u32 *cores_per_shard_p, u32 *nr_large_p)
61 {
62 	u32 nr_cores = 0, nr_cpus = 0, nr_shards;
63 	int cpu;
64 
65 	for_each_cpu(cpu, llc_cpus) {
66 		nr_cpus++;
67 		if (cpumask_first(topology_sibling_cpumask(cpu)) == cpu)
68 			nr_cores++;
69 	}
70 
71 	nr_shards = max_t(u32, 1, DIV_ROUND_UP(nr_cpus, shard_size));
72 	nr_shards = max_t(u32, nr_shards,
73 			  DIV_ROUND_UP(nr_cpus, SCX_CID_SHARD_MAX_CPUS));
74 
75 	*cores_per_shard_p = nr_cores / nr_shards;
76 	*nr_large_p = nr_cores % nr_shards;
77 }
78 
79 static void scx_cid_tables_free(struct scx_cid_tables *tbls)
80 {
81 	if (!tbls)
82 		return;
83 	kvfree(tbls->cid_to_cpu);
84 	kvfree(tbls->cpu_to_cid);
85 	kvfree(tbls->cid_to_shard);
86 	kvfree(tbls->shard_node);
87 	kvfree(tbls->shard_ranges);
88 	kvfree(tbls->topo);
89 	kfree(tbls);
90 }
91 
92 static void scx_cid_tables_free_rcufn(struct rcu_head *rcu)
93 {
94 	scx_cid_tables_free(container_of(rcu, struct scx_cid_tables, rcu));
95 }
96 
97 static struct scx_cid_tables *scx_cid_alloc_tables(void)
98 {
99 	u32 npossible = num_possible_cpus();
100 	struct scx_cid_tables *tbls;
101 
102 	tbls = kzalloc_obj(*tbls);
103 	if (!tbls)
104 		return NULL;
105 
106 	tbls->cid_to_cpu = kvzalloc_objs(*tbls->cid_to_cpu, npossible);
107 	tbls->cpu_to_cid = kvzalloc_objs(*tbls->cpu_to_cid, nr_cpu_ids);
108 	tbls->cid_to_shard = kvzalloc_objs(*tbls->cid_to_shard, npossible);
109 	tbls->shard_node = kvzalloc_objs(*tbls->shard_node, npossible);
110 	tbls->shard_ranges = kvzalloc_objs(*tbls->shard_ranges, npossible);
111 	tbls->topo = kvzalloc_objs(*tbls->topo, npossible);
112 
113 	if (!tbls->cid_to_cpu || !tbls->cpu_to_cid || !tbls->cid_to_shard ||
114 	    !tbls->shard_node || !tbls->shard_ranges || !tbls->topo) {
115 		scx_cid_tables_free(tbls);
116 		return NULL;
117 	}
118 
119 	return tbls;
120 }
121 
122 /**
123  * scx_cid_publish_tables - Publish the tables scx_cid_init() built
124  *
125  * Called after ops.init_cids() where the layout is final.
126  */
127 void scx_cid_publish_tables(void)
128 {
129 	struct scx_cid_tables *tbls = scx_cid_tables;
130 
131 	lockdep_assert_held(&scx_enable_mutex);
132 
133 	scx_nr_cid_shards = tbls->nr_shards;
134 	rcu_assign_pointer(scx_cid_to_cpu_tbl, tbls->cid_to_cpu);
135 	rcu_assign_pointer(scx_cpu_to_cid_tbl, tbls->cpu_to_cid);
136 	rcu_assign_pointer(scx_cid_to_shard, tbls->cid_to_shard);
137 	rcu_assign_pointer(scx_shard_node, tbls->shard_node);
138 	rcu_assign_pointer(scx_cid_shard_ranges, tbls->shard_ranges);
139 	rcu_assign_pointer(scx_cid_topo, tbls->topo);
140 }
141 
142 /**
143  * scx_cid_retire_tables - Unpublish and retire the cid tables
144  *
145  * Called by root disable after the readers which dereference without NULL
146  * checks are drained, inside cpus_read_lock() to exclude the hotplug path.
147  */
148 void scx_cid_retire_tables(void)
149 {
150 	struct scx_cid_tables *tbls = scx_cid_tables;
151 
152 	lockdep_assert_held(&scx_enable_mutex);
153 	lockdep_assert_cpus_held();
154 
155 	if (!tbls)
156 		return;
157 
158 	scx_cid_tables = NULL;
159 	RCU_INIT_POINTER(scx_cid_to_cpu_tbl, NULL);
160 	RCU_INIT_POINTER(scx_cpu_to_cid_tbl, NULL);
161 	RCU_INIT_POINTER(scx_cid_to_shard, NULL);
162 	RCU_INIT_POINTER(scx_shard_node, NULL);
163 	RCU_INIT_POINTER(scx_cid_shard_ranges, NULL);
164 	RCU_INIT_POINTER(scx_cid_topo, NULL);
165 	call_rcu(&tbls->rcu, scx_cid_tables_free_rcufn);
166 }
167 
168 /**
169  * scx_cid_init - build the cid mapping
170  * @sch: the scx_sched being initialized; used as the scx_error() target
171  *
172  * Build a fresh table set. It becomes visible through scx_cid_publish_tables()
173  * and is retired by scx_cid_retire_tables() at disable.
174  *
175  * See "Topological CPU IDs" in cid.h for the model. Walk online cpus by
176  * intersection at each level (parent_scratch & this_level_mask), which keeps
177  * containment correct by construction and naturally splits a physical LLC
178  * straddling two NUMA nodes into two LLC units. The caller must hold
179  * cpus_read_lock.
180  */
181 s32 scx_cid_init(struct scx_sched *sch)
182 {
183 	cpumask_var_t to_walk __free(free_cpumask_var) = CPUMASK_VAR_NULL;
184 	cpumask_var_t node_scratch __free(free_cpumask_var) = CPUMASK_VAR_NULL;
185 	cpumask_var_t llc_scratch __free(free_cpumask_var) = CPUMASK_VAR_NULL;
186 	cpumask_var_t cluster_scratch __free(free_cpumask_var) = CPUMASK_VAR_NULL;
187 	cpumask_var_t core_scratch __free(free_cpumask_var) = CPUMASK_VAR_NULL;
188 	cpumask_var_t llc_fallback __free(free_cpumask_var) = CPUMASK_VAR_NULL;
189 	cpumask_var_t online_no_topo __free(free_cpumask_var) = CPUMASK_VAR_NULL;
190 	struct scx_cid_tables *tbls;
191 	u32 next_cid = 0;
192 	s32 next_node_idx = 0, next_llc_idx = 0, next_core_idx = 0;
193 	s32 next_cluster_idx = 0;
194 	s32 next_shard_idx = 0;
195 	u32 shard_size, max_cids;
196 	u32 notopo_in_shard;
197 	s32 notopo_shard_cid, notopo_shard_idx;
198 	s32 cpu, cid, si;
199 
200 	/* CMASK_MAX_WORDS in cid.bpf.h covers NR_CPUS up to 8192 */
201 	BUILD_BUG_ON(NR_CPUS > 8192);
202 
203 	lockdep_assert_cpus_held();
204 	lockdep_assert_held(&scx_enable_mutex);
205 
206 	shard_size = sch->ops.cid_shard_size ?: SCX_CID_SHARD_SIZE_DFL;
207 	max_cids = min_t(u32, shard_size, SCX_CID_SHARD_MAX_CPUS);
208 
209 	tbls = scx_cid_alloc_tables();
210 	if (!tbls)
211 		return -ENOMEM;
212 
213 	scx_cid_tables = tbls;
214 
215 	for (si = 0; si < num_possible_cpus(); si++)
216 		tbls->shard_node[si] = NUMA_NO_NODE;
217 
218 	if (!zalloc_cpumask_var(&to_walk, GFP_KERNEL) ||
219 	    !zalloc_cpumask_var(&node_scratch, GFP_KERNEL) ||
220 	    !zalloc_cpumask_var(&llc_scratch, GFP_KERNEL) ||
221 	    !zalloc_cpumask_var(&cluster_scratch, GFP_KERNEL) ||
222 	    !zalloc_cpumask_var(&core_scratch, GFP_KERNEL) ||
223 	    !zalloc_cpumask_var(&llc_fallback, GFP_KERNEL) ||
224 	    !zalloc_cpumask_var(&online_no_topo, GFP_KERNEL))
225 		return -ENOMEM;
226 
227 	/* -1 sentinels for sparse-possible cpu id holes (0 is a valid cid) */
228 	for (cpu = 0; cpu < nr_cpu_ids; cpu++)
229 		tbls->cpu_to_cid[cpu] = -1;
230 
231 	cpumask_copy(to_walk, cpu_online_mask);
232 
233 	while (!cpumask_empty(to_walk)) {
234 		s32 next_cpu = cpumask_first(to_walk);
235 		s32 nid = cpu_to_node(next_cpu);
236 		s32 node_cid = next_cid;
237 		s32 node_idx;
238 
239 		/*
240 		 * No NUMA info: skip and let the tail loop assign a no-topo
241 		 * cid. cpumask_of_node(-1) is undefined.
242 		 */
243 		if (nid < 0) {
244 			cpumask_clear_cpu(next_cpu, to_walk);
245 			continue;
246 		}
247 
248 		node_idx = next_node_idx++;
249 
250 		/* node_scratch = to_walk & this node */
251 		cpumask_and(node_scratch, to_walk, cpumask_of_node(nid));
252 		if (WARN_ON_ONCE(!cpumask_test_cpu(next_cpu, node_scratch)))
253 			return -EINVAL;
254 
255 		while (!cpumask_empty(node_scratch)) {
256 			s32 ncpu = cpumask_first(node_scratch);
257 			const struct cpumask *llc_mask = cpu_llc_mask(ncpu, llc_fallback);
258 			s32 llc_cid = next_cid;
259 			s32 llc_idx = next_llc_idx++;
260 			u32 cores_per_shard, nr_large;
261 			u32 shard_local = 0, cores_in_shard = 0, cids_in_shard = 0;
262 			s32 shard_cid, shard_idx;
263 			s32 cluster_cid = -1, cluster_idx = -1;
264 
265 			/* llc_scratch = node_scratch & this llc */
266 			cpumask_and(llc_scratch, node_scratch, llc_mask);
267 			if (WARN_ON_ONCE(!cpumask_test_cpu(ncpu, llc_scratch)))
268 				return -EINVAL;
269 
270 			calc_shard_layout(llc_scratch, shard_size, &cores_per_shard, &nr_large);
271 			shard_cid = next_cid;
272 			shard_idx = next_shard_idx++;
273 			tbls->shard_node[shard_idx] = nid;
274 
275 			while (!cpumask_empty(llc_scratch)) {
276 				const struct cpumask *sib;
277 				s32 core_cid, core_idx, lcpu, ccpu;
278 				u32 max_cores, cids_in_core;
279 
280 				/*
281 				 * Take the cores of one cluster before moving
282 				 * on, so that a cluster is a contiguous cid
283 				 * range like the core, LLC and node levels.
284 				 */
285 				if (cpumask_empty(cluster_scratch)) {
286 					s32 xcpu = cpumask_first(llc_scratch);
287 
288 					cpumask_or(cluster_scratch, topology_cluster_cpumask(xcpu),
289 						   topology_sibling_cpumask(xcpu));
290 					cpumask_and(cluster_scratch, cluster_scratch, llc_scratch);
291 					cluster_cid = next_cid;
292 					cluster_idx = next_cluster_idx++;
293 				}
294 
295 				lcpu = cpumask_first(cluster_scratch);
296 				sib = topology_sibling_cpumask(lcpu);
297 				core_cid = next_cid;
298 				core_idx = next_core_idx++;
299 
300 				/* core_scratch = cluster_scratch & this core */
301 				cpumask_and(core_scratch, cluster_scratch, sib);
302 				if (WARN_ON_ONCE(!cpumask_test_cpu(lcpu, core_scratch)))
303 					return -EINVAL;
304 
305 				/*
306 				 * Advance to a new shard when either core or
307 				 * cid count reaches max. The latter bounds
308 				 * shard sizes under uneven SMT. Never start an
309 				 * empty shard.
310 				 */
311 				cids_in_core = cpumask_weight(core_scratch);
312 				max_cores = cores_per_shard + (shard_local < nr_large ? 1 : 0);
313 				if (cores_in_shard &&
314 				    (cores_in_shard >= max_cores ||
315 				     cids_in_shard + cids_in_core > max_cids)) {
316 					shard_local++;
317 					cores_in_shard = 0;
318 					cids_in_shard = 0;
319 					shard_cid = next_cid;
320 					shard_idx = next_shard_idx++;
321 					tbls->shard_node[shard_idx] = nid;
322 				}
323 				cores_in_shard++;
324 				cids_in_shard += cids_in_core;
325 
326 				for_each_cpu(ccpu, core_scratch) {
327 					s32 cid = next_cid++;
328 
329 					tbls->cid_to_cpu[cid] = ccpu;
330 					tbls->cpu_to_cid[ccpu] = cid;
331 					tbls->cid_to_shard[cid] = shard_idx;
332 					tbls->topo[cid] = (struct scx_cid_topo){
333 						.core_cid = core_cid,
334 						.core_idx = core_idx,
335 						.llc_cid = llc_cid,
336 						.llc_idx = llc_idx,
337 						.node_cid = node_cid,
338 						.node_idx = node_idx,
339 						.shard_cid = shard_cid,
340 						.shard_idx = shard_idx,
341 						.cluster_cid = cluster_cid,
342 						.cluster_idx = cluster_idx,
343 					};
344 
345 					cpumask_clear_cpu(ccpu, cluster_scratch);
346 					cpumask_clear_cpu(ccpu, llc_scratch);
347 					cpumask_clear_cpu(ccpu, node_scratch);
348 					cpumask_clear_cpu(ccpu, to_walk);
349 				}
350 			}
351 		}
352 	}
353 
354 	/*
355 	 * No-topo section: any possible cpu without a cid - normally just the
356 	 * not-online ones. Pack into shards of up to min(@shard_size,
357 	 * SCX_CID_SHARD_MAX_CPUS) cids so that every cid has a valid shard
358 	 * assignment and the hard cap holds even with a large @shard_size.
359 	 * Collect any currently-online cpus that land here in @online_no_topo
360 	 * so we can warn about them at the end.
361 	 */
362 	notopo_in_shard = min_t(u32, shard_size, SCX_CID_SHARD_MAX_CPUS);
363 	notopo_shard_cid = -1;
364 	notopo_shard_idx = -1;
365 
366 	for_each_cpu(cpu, cpu_possible_mask) {
367 		if (tbls->cpu_to_cid[cpu] != -1)
368 			continue;
369 		if (cpu_online(cpu))
370 			cpumask_set_cpu(cpu, online_no_topo);
371 
372 		cid = next_cid++;
373 		tbls->cid_to_cpu[cid] = cpu;
374 		tbls->cpu_to_cid[cpu] = cid;
375 
376 		if (notopo_in_shard >= min_t(u32, shard_size, SCX_CID_SHARD_MAX_CPUS)) {
377 			notopo_shard_cid = cid;
378 			notopo_shard_idx = next_shard_idx++;
379 			notopo_in_shard = 0;
380 		}
381 		notopo_in_shard++;
382 
383 		tbls->cid_to_shard[cid] = notopo_shard_idx;
384 		tbls->topo[cid] = SCX_CID_TOPO_NEG;
385 		tbls->topo[cid].shard_cid = notopo_shard_cid;
386 		tbls->topo[cid].shard_idx = notopo_shard_idx;
387 	}
388 
389 	if (!cpumask_empty(llc_fallback))
390 		pr_warn("scx_cid: cpus without cacheinfo, using node mask as llc: %*pbl\n",
391 			cpumask_pr_args(llc_fallback));
392 	if (!cpumask_empty(online_no_topo))
393 		pr_warn("scx_cid: online cpus with no usable topology: %*pbl\n",
394 			cpumask_pr_args(online_no_topo));
395 
396 	/*
397 	 * Fill cid_shard_ranges[] from cid_to_shard[]. Shards are contiguous
398 	 * cid ranges by construction: base_cid is the first cid landing in a
399 	 * shard, nr_cids is the count.
400 	 */
401 	for (cid = 0; cid < next_cid; cid++) {
402 		s32 sidx = tbls->cid_to_shard[cid];
403 
404 		if (tbls->shard_ranges[sidx].nr_cids == 0)
405 			tbls->shard_ranges[sidx].base_cid = cid;
406 		tbls->shard_ranges[sidx].nr_cids++;
407 	}
408 
409 	tbls->nr_shards = next_shard_idx;
410 	return 0;
411 }
412 
413 /**
414  * scx_cmask_clear - Zero every bit in @m's active range
415  * @m: cmask to clear
416  *
417  * Storage past the active range is left as is.
418  */
419 void scx_cmask_clear(struct scx_cmask *m)
420 {
421 	u32 nr_words;
422 
423 	if (!m->nr_cids)
424 		return;
425 	nr_words = (m->base + m->nr_cids - 1) / 64 - m->base / 64 + 1;
426 	memset(m->bits, 0, nr_words * sizeof(u64));
427 }
428 
429 /**
430  * scx_cmask_fill - Set every bit in @m's active range
431  * @m: cmask to fill
432  *
433  * Counterpart to scx_cmask_clear(). Storage past the active range is left as is.
434  */
435 void scx_cmask_fill(struct scx_cmask *m)
436 {
437 	u32 nr_words, head_bits, tail_bits;
438 
439 	if (!m->nr_cids)
440 		return;
441 	nr_words = (m->base + m->nr_cids - 1) / 64 - m->base / 64 + 1;
442 	memset(m->bits, 0xff, nr_words * sizeof(u64));
443 
444 	/* clear word-0 bits below base */
445 	head_bits = m->base & 63;
446 	if (head_bits)
447 		m->bits[0] &= ~((1ULL << head_bits) - 1);
448 
449 	/* clear last-word bits at or past base + nr_cids */
450 	tail_bits = (m->base + m->nr_cids) & 63;
451 	if (tail_bits)
452 		m->bits[nr_words - 1] &= (1ULL << tail_bits) - 1;
453 }
454 
455 /*
456  * Return the index of the largest entry in @counts, or NUMA_NO_NODE if all
457  * entries are zero. Ties resolve to the lowest index.
458  */
459 static s32 pick_max_node(const u32 *counts, u32 n)
460 {
461 	s32 best = NUMA_NO_NODE;
462 	u32 best_count = 0, i;
463 
464 	for (i = 0; i < n; i++) {
465 		if (counts[i] > best_count) {
466 			best_count = counts[i];
467 			best = i;
468 		}
469 	}
470 	return best;
471 }
472 
473 __bpf_kfunc_start_defs();
474 
475 /**
476  * scx_bpf_cid_override - Install an explicit cpu->cid mapping with shard info
477  * @cpu_to_cid__arena: array of nr_cpu_ids s32 entries (cid for each cpu)
478  * @cpu_to_cid_cnt: number of entries, must be nr_cpu_ids
479  * @shard_start__arena: array of first-cid-of-each-shard, one entry per shard
480  * @shard_start_cnt: number of shards
481  * @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
482  *
483  * May only be called from ops.init_cids() of the root scheduler. Replace the
484  * topology-probed cid mapping and shard layout with caller-provided ones. Each
485  * possible cpu must map to a unique cid in [0, num_possible_cpus()). The shard
486  * starts must be strictly increasing with the first entry 0 and all values <
487  * num_possible_cpus(). The last shard extends to num_possible_cpus() and no
488  * shard may span more than SCX_CID_SHARD_MAX_CPUS cids. Topo info
489  * (core/cluster/LLC/node) is cleared and the shard layout is set from the
490  * input. On invalid input, abort the scheduler.
491  */
492 __bpf_kfunc void scx_bpf_cid_override(const s32 *cpu_to_cid__arena, u32 cpu_to_cid_cnt,
493 				      const s32 *shard_start__arena, u32 shard_start_cnt,
494 				      const struct bpf_prog_aux *aux)
495 {
496 	cpumask_var_t seen __free(free_cpumask_var) = CPUMASK_VAR_NULL;
497 	u32 *node_counts __free(kfree) = NULL;
498 	s32 *cpu_to_cid __free(kfree) = NULL;
499 	s32 *shard_start __free(kfree) = NULL;
500 	u32 npossible = num_possible_cpus();
501 	struct scx_cid_tables *tbls;
502 	struct scx_sched *sch;
503 	u32 nr_shards = shard_start_cnt;
504 	bool alloced;
505 	s32 cpu, cid, si;
506 
507 	/*
508 	 * GFP_KERNEL allocs must happen before the rcu read section. Snapshot
509 	 * the BPF-supplied arrays so a concurrent arena write can't change
510 	 * them between validation and use.
511 	 *
512 	 * The BPF-supplied counts size the snapshots and thus the arena reads.
513 	 * Gate the copies on the count bounds, reported below once @sch is
514 	 * available. The bounded reads, at most 32KB, stay within the guard
515 	 * region that arena fault recovery covers.
516 	 */
517 	alloced = zalloc_cpumask_var(&seen, GFP_KERNEL);
518 	node_counts = kzalloc_objs(*node_counts, nr_node_ids);
519 	if (cpu_to_cid_cnt == nr_cpu_ids)
520 		cpu_to_cid = kmemdup(cpu_to_cid__arena, cpu_to_cid_cnt * sizeof(s32),
521 				     GFP_KERNEL);
522 	if (nr_shards && nr_shards <= npossible)
523 		shard_start = kmemdup(shard_start__arena, nr_shards * sizeof(s32),
524 				      GFP_KERNEL);
525 
526 	guard(rcu)();
527 
528 	sch = scx_prog_sched(aux);
529 	if (unlikely(!sch))
530 		return;
531 
532 	/* called from ops.init_cids(), so the tables exist and are unpublished */
533 	lockdep_assert_held(&scx_enable_mutex);
534 	tbls = scx_cid_tables;
535 
536 	if (cpu_to_cid_cnt != nr_cpu_ids) {
537 		scx_error(sch, "scx_bpf_cid_override: cpu_to_cid expected %u entries, got %u",
538 			  nr_cpu_ids, cpu_to_cid_cnt);
539 		return;
540 	}
541 
542 	if (!nr_shards || nr_shards > npossible) {
543 		scx_error(sch, "scx_bpf_cid_override: invalid shard_start count %u",
544 			  nr_shards);
545 		return;
546 	}
547 
548 	if (!alloced || !node_counts || !cpu_to_cid || !shard_start) {
549 		scx_error(sch, "scx_bpf_cid_override: allocation failed");
550 		return;
551 	}
552 
553 	/* validate shard_start[]: starts at 0, strictly increasing, in range */
554 	if (shard_start[0] != 0) {
555 		scx_error(sch, "scx_bpf_cid_override: shard_start[0] must be 0, got %d",
556 			  shard_start[0]);
557 		return;
558 	}
559 	for (si = 1; si < nr_shards; si++) {
560 		if (shard_start[si] <= shard_start[si - 1]) {
561 			scx_error(sch, "scx_bpf_cid_override: shard_start not increasing at [%d]",
562 				  si);
563 			return;
564 		}
565 		if (shard_start[si] >= npossible) {
566 			scx_error(sch, "scx_bpf_cid_override: shard_start[%d]=%d >= %u",
567 				  si, shard_start[si], npossible);
568 			return;
569 		}
570 		if (shard_start[si] - shard_start[si - 1] > SCX_CID_SHARD_MAX_CPUS) {
571 			scx_error(sch, "scx_bpf_cid_override: shard[%d] span %d exceeds max %d",
572 				  si - 1, shard_start[si] - shard_start[si - 1],
573 				  SCX_CID_SHARD_MAX_CPUS);
574 			return;
575 		}
576 	}
577 	if (npossible - shard_start[nr_shards - 1] > SCX_CID_SHARD_MAX_CPUS) {
578 		scx_error(sch, "scx_bpf_cid_override: shard[%d] span %d exceeds max %d",
579 			  nr_shards - 1, npossible - shard_start[nr_shards - 1],
580 			  SCX_CID_SHARD_MAX_CPUS);
581 		return;
582 	}
583 
584 	/* validate first so that invalid input leaves the tables untouched */
585 	for_each_possible_cpu(cpu) {
586 		s32 c = cpu_to_cid[cpu];
587 
588 		if (!cid_valid(sch, c))
589 			return;
590 		if (cpumask_test_and_set_cpu(c, seen)) {
591 			scx_error(sch, "cid %d assigned to multiple cpus", c);
592 			return;
593 		}
594 	}
595 
596 	for_each_possible_cpu(cpu) {
597 		s32 c = cpu_to_cid[cpu];
598 
599 		tbls->cpu_to_cid[cpu] = c;
600 		tbls->cid_to_cpu[c] = cpu;
601 	}
602 
603 	/*
604 	 * Derive shard_node[] by majority count: an overridden shard may
605 	 * span NUMA nodes, so assign each to the node that owns the most cpus.
606 	 */
607 	for (si = 0; si < nr_shards; si++) {
608 		u32 end = (si + 1 < nr_shards) ? shard_start[si + 1] : npossible;
609 
610 		memset(node_counts, 0, nr_node_ids * sizeof(*node_counts));
611 		for (cid = shard_start[si]; cid < end; cid++) {
612 			s32 node = cpu_to_node(tbls->cid_to_cpu[cid]);
613 
614 			if (numa_valid_node(node))
615 				node_counts[node]++;
616 		}
617 		tbls->shard_node[si] = pick_max_node(node_counts, nr_node_ids);
618 	}
619 
620 	/*
621 	 * Invalidate stale topo info and install shard layout from
622 	 * @shard_start. Walk shards to derive shard_cid/shard_idx for each cid.
623 	 */
624 	si = 0;
625 	for (cid = 0; cid < npossible; cid++) {
626 		if (si + 1 < nr_shards && cid >= shard_start[si + 1])
627 			si++;
628 		tbls->cid_to_shard[cid] = si;
629 		tbls->topo[cid] = SCX_CID_TOPO_NEG;
630 		tbls->topo[cid].shard_cid = shard_start[si];
631 		tbls->topo[cid].shard_idx = si;
632 	}
633 
634 	/* Rebuild shard_ranges[] for the new layout. */
635 	memset(tbls->shard_ranges, 0, npossible * sizeof(*tbls->shard_ranges));
636 	for (si = 0; si < nr_shards; si++) {
637 		u32 end = (si + 1 < nr_shards) ? shard_start[si + 1] : npossible;
638 
639 		tbls->shard_ranges[si].base_cid = shard_start[si];
640 		tbls->shard_ranges[si].nr_cids = end - shard_start[si];
641 	}
642 
643 	tbls->nr_shards = nr_shards;
644 }
645 
646 /**
647  * scx_bpf_cid_to_cpu - Return the raw CPU id for @cid
648  * @cid: cid to look up
649  * @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
650  *
651  * Return the raw CPU id for @cid. Trigger scx_error() and return -EINVAL if
652  * @cid is invalid. The cid<->cpu mapping is static for the lifetime of the
653  * loaded scheduler, so the BPF side can cache the result to avoid repeated
654  * kfunc invocations.
655  */
656 __bpf_kfunc s32 scx_bpf_cid_to_cpu(s32 cid, const struct bpf_prog_aux *aux)
657 {
658 	struct scx_sched *sch;
659 
660 	guard(rcu)();
661 
662 	sch = scx_prog_sched(aux);
663 	if (unlikely(!sch))
664 		return -EINVAL;
665 	return scx_cid_to_cpu(sch, cid);
666 }
667 
668 /**
669  * scx_bpf_cid_node - Return the NUMA node the given @cid belongs to
670  * @cid: cid to look up
671  * @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
672  *
673  * Return NUMA_NO_NODE if @cid is invalid or there is no scheduler. Unlike
674  * scx_bpf_cpu_node(), this kfunc is available only with CID support, so BPF
675  * schedulers can use its presence to detect NUMA-aware CID lookups.
676  */
677 __bpf_kfunc s32 scx_bpf_cid_node(s32 cid, const struct bpf_prog_aux *aux)
678 {
679 	struct scx_sched *sch;
680 	s32 cpu;
681 
682 	guard(rcu)();
683 
684 	sch = scx_prog_sched(aux);
685 	if (unlikely(!sch))
686 		return NUMA_NO_NODE;
687 	cpu = scx_cid_to_cpu(sch, cid);
688 	return cpu < 0 ? NUMA_NO_NODE : cpu_to_node(cpu);
689 }
690 
691 /**
692  * scx_bpf_cpu_to_cid - Return the cid for @cpu
693  * @cpu: cpu to look up
694  * @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
695  *
696  * Return the cid for @cpu. Trigger scx_error() and return -EINVAL if @cpu is
697  * invalid. The cid<->cpu mapping is static for the lifetime of the loaded
698  * scheduler, so the BPF side can cache the result to avoid repeated kfunc
699  * invocations.
700  */
701 __bpf_kfunc s32 scx_bpf_cpu_to_cid(s32 cpu, const struct bpf_prog_aux *aux)
702 {
703 	struct scx_sched *sch;
704 
705 	guard(rcu)();
706 
707 	sch = scx_prog_sched(aux);
708 	if (unlikely(!sch))
709 		return -EINVAL;
710 	return scx_cpu_to_cid(sch, cpu);
711 }
712 
713 /*
714  * Set ops on cmasks. cmask_walk_op2() shares one walk across mutating
715  * (and/or/copy/andnot) and predicate (subset/intersects) two-cmask forms;
716  * cmask_walk_op1() does the same shape over a single cmask range. Every public
717  * entry passes a compile-time-constant @op; cmask_walk_op{1,2}() and
718  * cmask_word_op{1,2}() are __always_inline so the inner switch collapses to the
719  * selected op and cmask_op2_is_pred() folds the predicate early-exit out of
720  * mutating ops.
721  *
722  * Two-cmask ops only touch @dst bits inside the intersection of the two ranges;
723  * bits outside stay untouched. In particular, scx_cmask_copy() does NOT zero
724  * @dst bits that lie outside @src's range.
725  *
726  * Word accesses use READ_ONCE/WRITE_ONCE so a caller may read @src
727  * locklessly. Memory ordering against concurrent writers is the caller's
728  * responsibility.
729  */
730 enum cmask_op2 {
731 	/* mutating */
732 	CMASK_OP2_AND,
733 	CMASK_OP2_OR,
734 	CMASK_OP2_COPY,
735 	CMASK_OP2_ANDNOT,
736 	/* predicates - short-circuit when the per-word result is true */
737 	CMASK_OP2_SUBSET,
738 	CMASK_OP2_INTERSECTS,
739 	/*
740 	 * @a is a BPF-arena cmask. Words on @a use READ_ONCE/WRITE_ONCE since
741 	 * BPF may read/write concurrently. See scx_cmask_ref_or() / _copy().
742 	 */
743 	CMASK_OP2_REF_OR,
744 	CMASK_OP2_REF_COPY,
745 };
746 
747 static __always_inline bool cmask_op2_is_pred(const enum cmask_op2 op)
748 {
749 	return op == CMASK_OP2_SUBSET || op == CMASK_OP2_INTERSECTS;
750 }
751 
752 static __always_inline bool cmask_word_op2(u64 *av, const u64 *bp, u64 mask,
753 					   const enum cmask_op2 op)
754 {
755 	switch (op) {
756 	case CMASK_OP2_AND:
757 		WRITE_ONCE(*av, *av & (~mask | READ_ONCE(*bp)));
758 		return false;
759 	case CMASK_OP2_OR:
760 		WRITE_ONCE(*av, *av | (READ_ONCE(*bp) & mask));
761 		return false;
762 	case CMASK_OP2_COPY:
763 		WRITE_ONCE(*av, (*av & ~mask) | (READ_ONCE(*bp) & mask));
764 		return false;
765 	case CMASK_OP2_ANDNOT:
766 		WRITE_ONCE(*av, *av & ~(READ_ONCE(*bp) & mask));
767 		return false;
768 	case CMASK_OP2_SUBSET:
769 		/* stop on the first bit in @sub not set in @super */
770 		return (READ_ONCE(*bp) & ~READ_ONCE(*av)) & mask;
771 	case CMASK_OP2_INTERSECTS:
772 		return (READ_ONCE(*av) & READ_ONCE(*bp)) & mask;
773 	case CMASK_OP2_REF_OR:
774 		WRITE_ONCE(*av, READ_ONCE(*av) | (READ_ONCE(*bp) & mask));
775 		return false;
776 	case CMASK_OP2_REF_COPY:
777 		WRITE_ONCE(*av, (READ_ONCE(*av) & ~mask) | (READ_ONCE(*bp) & mask));
778 		return false;
779 	}
780 	unreachable();
781 }
782 
783 /*
784  * Walk the intersection of [@a_base, @a_base + @a_nr_cids) with [@b_base,
785  * @b_base + @b_nr_cids) word by word, applying @op. Mutating ops walk all words
786  * and return false; predicates return true on the first word whose per-word
787  * test is true. Empty intersection returns false (matches "no bits to consider"
788  * for both mutate and predicate).
789  *
790  * Base/nr_cids are taken as parameters so callers with snapshotted bounds can
791  * drive the walk with values independent of the cmask's header.
792  */
793 static __always_inline bool cmask_walk_op2(u64 *a_bits, u32 a_base, u32 a_nr_cids,
794 					   const u64 *b_bits, u32 b_base, u32 b_nr_cids,
795 					   const enum cmask_op2 op)
796 {
797 	u32 lo = max(a_base, b_base);
798 	u32 hi = min(a_base + a_nr_cids, b_base + b_nr_cids);
799 	u32 a_word_off = a_base / 64;
800 	u32 b_word_off = b_base / 64;
801 	u32 lo_word = lo / 64;
802 	u32 hi_word = (hi - 1) / 64;
803 	u64 head_mask = GENMASK_U64(63, lo & 63);
804 	u64 tail_mask = GENMASK_U64((hi - 1) & 63, 0);
805 	u32 w;
806 
807 	if (lo >= hi)
808 		return false;
809 
810 	if (lo_word == hi_word)
811 		return cmask_word_op2(&a_bits[lo_word - a_word_off],
812 				      &b_bits[lo_word - b_word_off],
813 				      head_mask & tail_mask, op);
814 
815 	if (cmask_word_op2(&a_bits[lo_word - a_word_off],
816 			   &b_bits[lo_word - b_word_off], head_mask, op) &&
817 	    cmask_op2_is_pred(op))
818 		return true;
819 
820 	for (w = lo_word + 1; w < hi_word; w++)
821 		if (cmask_word_op2(&a_bits[w - a_word_off],
822 				   &b_bits[w - b_word_off], ~0ULL, op) &&
823 		    cmask_op2_is_pred(op))
824 			return true;
825 
826 	return cmask_word_op2(&a_bits[hi_word - a_word_off],
827 			      &b_bits[hi_word - b_word_off], tail_mask, op);
828 }
829 
830 enum cmask_op1 {
831 	CMASK_OP1_ANY_SET,
832 };
833 
834 static __always_inline bool cmask_word_op1(const u64 *ap, u64 mask,
835 					   const enum cmask_op1 op)
836 {
837 	switch (op) {
838 	case CMASK_OP1_ANY_SET:
839 		return READ_ONCE(*ap) & mask;
840 	}
841 	unreachable();
842 }
843 
844 /*
845  * Walk [@a_base, @a_base + @a_nr_cids) of @a_bits word by word, applying @op.
846  * Returns true on the first word whose per-word test is true; returns false if
847  * no word matches or the range is empty. All current op1s short-circuit on
848  * per-word true; if a non-predicate op1 lands here, add a cmask_op1_is_pred()
849  * guard analogous to cmask_op2_is_pred().
850  */
851 static __always_inline bool cmask_walk_op1(const u64 *a_bits, u32 a_base,
852 					   u32 a_nr_cids,
853 					   const enum cmask_op1 op)
854 {
855 	u32 lo = a_base;
856 	u32 hi = a_base + a_nr_cids;
857 	u32 a_word_off = a_base / 64;
858 	u32 lo_word = lo / 64;
859 	u32 hi_word = (hi - 1) / 64;
860 	u64 head_mask = GENMASK_U64(63, lo & 63);
861 	u64 tail_mask = GENMASK_U64((hi - 1) & 63, 0);
862 	u32 w;
863 
864 	if (lo >= hi)
865 		return false;
866 
867 	if (lo_word == hi_word)
868 		return cmask_word_op1(&a_bits[lo_word - a_word_off],
869 				      head_mask & tail_mask, op);
870 
871 	if (cmask_word_op1(&a_bits[lo_word - a_word_off], head_mask, op))
872 		return true;
873 	for (w = lo_word + 1; w < hi_word; w++)
874 		if (cmask_word_op1(&a_bits[w - a_word_off], ~0ULL, op))
875 			return true;
876 	return cmask_word_op1(&a_bits[hi_word - a_word_off], tail_mask, op);
877 }
878 
879 void scx_cmask_and(struct scx_cmask *dst, const struct scx_cmask *src)
880 {
881 	cmask_walk_op2(dst->bits, dst->base, dst->nr_cids,
882 		       src->bits, src->base, src->nr_cids, CMASK_OP2_AND);
883 }
884 
885 void scx_cmask_or(struct scx_cmask *dst, const struct scx_cmask *src)
886 {
887 	cmask_walk_op2(dst->bits, dst->base, dst->nr_cids,
888 		       src->bits, src->base, src->nr_cids, CMASK_OP2_OR);
889 }
890 
891 void scx_cmask_copy(struct scx_cmask *dst, const struct scx_cmask *src)
892 {
893 	cmask_walk_op2(dst->bits, dst->base, dst->nr_cids,
894 		       src->bits, src->base, src->nr_cids, CMASK_OP2_COPY);
895 }
896 
897 void scx_cmask_andnot(struct scx_cmask *dst, const struct scx_cmask *src)
898 {
899 	cmask_walk_op2(dst->bits, dst->base, dst->nr_cids,
900 		       src->bits, src->base, src->nr_cids, CMASK_OP2_ANDNOT);
901 }
902 
903 /*
904  * Return true if @cm has any bit set in [@lo, @hi). Caller must ensure
905  * [@lo, @hi) is contained in @cm's range.
906  */
907 static bool cmask_any_set_in_range(const struct scx_cmask *cm, u32 lo, u32 hi)
908 {
909 	if (lo >= hi)
910 		return false;
911 	return cmask_walk_op1(&cm->bits[lo / 64 - cm->base / 64], lo, hi - lo,
912 			      CMASK_OP1_ANY_SET);
913 }
914 
915 /**
916  * scx_cmask_subset - test whether @sub is a subset of @super
917  * @sub: cmask to test
918  * @super: cmask to test against
919  *
920  * Return true iff every set bit of @sub is also set in @super.
921  */
922 bool scx_cmask_subset(const struct scx_cmask *sub, const struct scx_cmask *super)
923 {
924 	u32 super_end = super->base + super->nr_cids;
925 	u32 sub_end = sub->base + sub->nr_cids;
926 
927 	/*
928 	 * Set bits in @sub outside @super's range can't be in @super, so any
929 	 * such bit means not a subset. The walk below only visits words
930 	 * common to both ranges, so these need a separate scan.
931 	 */
932 	if (sub->base < super->base &&
933 	    cmask_any_set_in_range(sub, sub->base, min(super->base, sub_end)))
934 		return false;
935 	if (sub_end > super_end &&
936 	    cmask_any_set_in_range(sub, max(sub->base, super_end), sub_end))
937 		return false;
938 
939 	return !cmask_walk_op2((u64 *)super->bits, super->base, super->nr_cids,
940 			       sub->bits, sub->base, sub->nr_cids, CMASK_OP2_SUBSET);
941 }
942 
943 bool scx_cmask_intersects(const struct scx_cmask *a, const struct scx_cmask *b)
944 {
945 	return cmask_walk_op2((u64 *)a->bits, a->base, a->nr_cids,
946 			      b->bits, b->base, b->nr_cids, CMASK_OP2_INTERSECTS);
947 }
948 
949 /**
950  * scx_cmask_empty - Test whether @m has no bits set
951  * @m: cmask to test
952  *
953  * Return true iff @m's active range has no bits set.
954  */
955 bool scx_cmask_empty(const struct scx_cmask *m)
956 {
957 	return !cmask_any_set_in_range(m, m->base, m->base + m->nr_cids);
958 }
959 
960 /**
961  * scx_bpf_cid_topo - Copy out per-cid topology info
962  * @cid: cid to look up
963  * @out: where to copy the topology info
964  * @out__sz: size of @out, the program's sizeof(struct scx_cid_topo)
965  * @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
966  *
967  * Fill @out with the topology info for @cid. Trigger scx_error() if @cid is out
968  * of range. If @cid is valid but in the no-topo section, all fields are set to
969  * -1. All fields are also set to -1 when no cid tables have been published yet,
970  * which a program may observe while racing the root enable.
971  *
972  * The program's struct may be older or newer than the kernel's. The smaller of
973  * @out__sz and the kernel's size is copied and the rest of @out is set to -1.
974  */
975 __bpf_kfunc void scx_bpf_cid_topo(s32 cid, struct scx_cid_topo *out, size_t out__sz,
976 				  const struct bpf_prog_aux *aux)
977 {
978 	size_t len = min(out__sz, sizeof(*out));
979 	struct scx_cid_topo *topo;
980 	struct scx_sched *sch;
981 
982 	/* the error cases and fields the kernel lacks read as -1 */
983 	memset(out, 0xff, out__sz);
984 
985 	guard(rcu)();
986 
987 	sch = scx_prog_sched(aux);
988 	topo = rcu_dereference(scx_cid_topo);
989 	if (unlikely(!sch) || !cid_valid(sch, cid) || unlikely(!topo))
990 		return;
991 
992 	memcpy(out, &topo[cid], len);
993 }
994 
995 __bpf_kfunc_end_defs();
996 
997 BTF_KFUNCS_START(scx_kfunc_ids_init_cids)
998 BTF_ID_FLAGS(func, scx_bpf_cid_override, KF_IMPLICIT_ARGS | KF_SLEEPABLE)
999 BTF_KFUNCS_END(scx_kfunc_ids_init_cids)
1000 
1001 static const struct btf_kfunc_id_set scx_kfunc_set_init_cids = {
1002 	.owner	= THIS_MODULE,
1003 	.set	= &scx_kfunc_ids_init_cids,
1004 	.filter	= scx_kfunc_context_filter,
1005 };
1006 
1007 BTF_KFUNCS_START(scx_kfunc_ids_cid)
1008 BTF_ID_FLAGS(func, scx_bpf_cid_to_cpu, KF_IMPLICIT_ARGS)
1009 BTF_ID_FLAGS(func, scx_bpf_cid_node, KF_IMPLICIT_ARGS)
1010 BTF_ID_FLAGS(func, scx_bpf_cpu_to_cid, KF_IMPLICIT_ARGS)
1011 BTF_ID_FLAGS(func, scx_bpf_cid_topo, KF_IMPLICIT_ARGS)
1012 BTF_KFUNCS_END(scx_kfunc_ids_cid)
1013 
1014 static const struct btf_kfunc_id_set scx_kfunc_set_cid = {
1015 	.owner	= THIS_MODULE,
1016 	.set	= &scx_kfunc_ids_cid,
1017 };
1018 
1019 /**
1020  * scx_cmask_ref_init - Bind a scx_cmask_ref to a BPF-arena cmask
1021  * @sch: scheduler whose arena hosts @src
1022  * @src: BPF-supplied cmask, rebased to its kernel address
1023  * @ref: output ref
1024  *
1025  * Snapshot @src's @base, @nr_cids and @alloc_words. The snapshot is necessary
1026  * because BPF may mutate the live header asynchronously.
1027  *
1028  * Return 0 on success, -EINVAL if the range is out of bounds or @alloc_words
1029  * doesn't cover it.
1030  */
1031 int scx_cmask_ref_init(struct scx_sched *sch, const struct scx_cmask *src,
1032 		       struct scx_cmask_ref *ref)
1033 {
1034 	u32 base, nr_cids, alloc_words, npossible = num_possible_cpus();
1035 	s32 *cid_to_shard;
1036 
1037 	base = READ_ONCE(src->base);
1038 	nr_cids = READ_ONCE(src->nr_cids);
1039 	alloc_words = READ_ONCE(src->alloc_words);
1040 
1041 	if (unlikely(base >= npossible || nr_cids > npossible - base ||
1042 		     SCX_CMASK_NR_WORDS(nr_cids) > alloc_words))
1043 		return -EINVAL;
1044 
1045 	ref->sch = sch;
1046 	ref->src = (struct scx_cmask *)src;
1047 	ref->base = base;
1048 	ref->nr_cids = nr_cids;
1049 
1050 	cid_to_shard = rcu_dereference_all(scx_cid_to_shard);
1051 	ref->shard_first = cid_to_shard[base];
1052 	if (likely(nr_cids))
1053 		ref->shard_end = cid_to_shard[base + nr_cids - 1] + 1;
1054 	else
1055 		ref->shard_end = ref->shard_first;
1056 
1057 	return 0;
1058 }
1059 
1060 /**
1061  * scx_cmask_ref_init_kern - Bind a scx_cmask_ref to a kernel-owned cmask
1062  * @sch: scheduler the cmask belongs to
1063  * @m: kernel address of the target cmask, storage sized for @nr_cids at @base
1064  * @base: first cid of the active range
1065  * @nr_cids: active range length
1066  * @ref: output ref
1067  *
1068  * Like scx_cmask_ref_init() but the geometry is supplied by the caller, not
1069  * read from @m's header, so a concurrent BPF write to the header can't steer
1070  * later sizing or offsets. Rewrite the header from the trusted geometry and
1071  * bind @ref to it.
1072  */
1073 void scx_cmask_ref_init_kern(struct scx_sched *sch, struct scx_cmask *m,
1074 			     u32 base, u32 nr_cids, struct scx_cmask_ref *ref)
1075 {
1076 	s32 *cid_to_shard;
1077 
1078 	WRITE_ONCE(m->base, base);
1079 	WRITE_ONCE(m->nr_cids, nr_cids);
1080 	WRITE_ONCE(m->alloc_words, SCX_CMASK_NR_WORDS(nr_cids));
1081 
1082 	ref->sch = sch;
1083 	ref->src = m;
1084 	ref->base = base;
1085 	ref->nr_cids = nr_cids;
1086 
1087 	cid_to_shard = rcu_dereference_all(scx_cid_to_shard);
1088 	ref->shard_first = cid_to_shard[base];
1089 	if (likely(nr_cids))
1090 		ref->shard_end = cid_to_shard[base + nr_cids - 1] + 1;
1091 	else
1092 		ref->shard_end = ref->shard_first;
1093 }
1094 
1095 /**
1096  * scx_cmask_ref_shard - Read one shard from @ref into @out
1097  * @ref: validated ref
1098  * @shard_idx: target shard, in [@ref->shard_first, @ref->shard_end)
1099  * @out: output cmask whose @out->alloc_words must hold the shard
1100  *
1101  * Set @out to the intersection of @ref's range with @shard_idx's cid range,
1102  * with bits[] read from @ref->src via READ_ONCE. Empty intersection sets
1103  * @out->nr_cids to 0. scx_error()s on @ref's sched if @out can't hold the
1104  * shard.
1105  */
1106 void scx_cmask_ref_shard(const struct scx_cmask_ref *ref, s32 shard_idx,
1107 			 struct scx_cmask *out)
1108 {
1109 	const struct scx_cid_shard *shard =
1110 		&rcu_dereference_all(scx_cid_shard_ranges)[shard_idx];
1111 	u32 shard_base = shard->base_cid;
1112 	u32 shard_end = shard_base + shard->nr_cids;
1113 	u32 isect_base, isect_end, nr_words, src_off, wi;
1114 	u64 head_mask, tail_mask;
1115 
1116 	isect_base = max(ref->base, shard_base);
1117 	isect_end = min(ref->base + ref->nr_cids, shard_end);
1118 
1119 	if (isect_base >= isect_end) {
1120 		out->base = shard_base;
1121 		out->nr_cids = 0;
1122 		return;
1123 	}
1124 
1125 	nr_words = ((isect_end - 1) / 64) - (isect_base / 64) + 1;
1126 	if (nr_words > out->alloc_words) {
1127 		scx_error(ref->sch, "scx_cmask_ref_shard: out alloc_words=%u < %u for shard %d",
1128 			  out->alloc_words, nr_words, shard_idx);
1129 		out->base = shard_base;
1130 		out->nr_cids = 0;
1131 		return;
1132 	}
1133 
1134 	out->base = isect_base;
1135 	out->nr_cids = isect_end - isect_base;
1136 	src_off = (isect_base / 64) - (ref->base / 64);
1137 
1138 	for (wi = 0; wi < nr_words; wi++)
1139 		out->bits[wi] = READ_ONCE(ref->src->bits[src_off + wi]);
1140 
1141 	head_mask = GENMASK_U64(63, isect_base & 63);
1142 	out->bits[0] &= head_mask;
1143 	tail_mask = GENMASK_U64((isect_end - 1) & 63, 0);
1144 	out->bits[nr_words - 1] &= tail_mask;
1145 }
1146 
1147 /**
1148  * scx_cmask_ref_or - OR @src into the arena cmask referenced by @ref
1149  * @ref: validated ref
1150  * @src: stable kernel cmask
1151  *
1152  * Bits inside the intersection of @ref's snapshotted range with @src's range
1153  * are OR'd into @ref->src and bits outside are left unchanged. Stores on
1154  * @ref->src use WRITE_ONCE since BPF may read/write concurrently.
1155  */
1156 void scx_cmask_ref_or(const struct scx_cmask_ref *ref, const struct scx_cmask *src)
1157 {
1158 	cmask_walk_op2(ref->src->bits, ref->base, ref->nr_cids,
1159 		       src->bits, src->base, src->nr_cids, CMASK_OP2_REF_OR);
1160 }
1161 
1162 /**
1163  * scx_cmask_ref_copy - Copy @src into the arena cmask referenced by @ref
1164  * @ref: validated ref
1165  * @src: stable kernel cmask
1166  *
1167  * Bits inside the intersection of @ref's snapshotted range with @src's range
1168  * take @src's values and bits outside are left unchanged. Stores on @ref->src
1169  * use WRITE_ONCE since BPF may read/write concurrently.
1170  */
1171 void scx_cmask_ref_copy(const struct scx_cmask_ref *ref, const struct scx_cmask *src)
1172 {
1173 	cmask_walk_op2(ref->src->bits, ref->base, ref->nr_cids,
1174 		       src->bits, src->base, src->nr_cids, CMASK_OP2_REF_COPY);
1175 }
1176 
1177 /**
1178  * scx_cmask_ref_from_cpumask - Populate @ref's arena cmask from a cpumask
1179  * @ref: kern-bound ref, see scx_cmask_ref_init_kern()
1180  * @cpumask: cpus to translate into cids
1181  *
1182  * Write @ref's active range one word at a time, setting each cid's bit when
1183  * its cpu is in @cpumask. Offsets and length come from @ref's trusted geometry
1184  * and stores use WRITE_ONCE since BPF may read concurrently, so the arena
1185  * header is never read.
1186  */
1187 void scx_cmask_ref_from_cpumask(const struct scx_cmask_ref *ref,
1188 				const struct cpumask *cpumask)
1189 {
1190 	struct scx_cmask *m = ref->src;
1191 	u32 base = ref->base, nr_cids = ref->nr_cids;
1192 	u32 wi, nr_words;
1193 
1194 	if (!nr_cids)
1195 		return;
1196 
1197 	nr_words = (base + nr_cids - 1) / 64 - base / 64 + 1;
1198 	for (wi = 0; wi < nr_words; wi++) {
1199 		u32 word_first_cid = (base / 64 + wi) * 64;
1200 		u64 word = 0;
1201 		u32 bit;
1202 
1203 		for (bit = 0; bit < 64; bit++) {
1204 			u32 cid = word_first_cid + bit;
1205 
1206 			if (cid < base || cid >= base + nr_cids)
1207 				continue;
1208 			if (cpumask_test_cpu(__scx_cid_to_cpu(cid), cpumask))
1209 				word |= BIT_U64(bit);
1210 		}
1211 		WRITE_ONCE(m->bits[wi], word);
1212 	}
1213 }
1214 
1215 int scx_cid_kfunc_init(void)
1216 {
1217 	return register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS, &scx_kfunc_set_init_cids) ?:
1218 		register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS, &scx_kfunc_set_cid) ?:
1219 		register_btf_kfunc_id_set(BPF_PROG_TYPE_TRACING, &scx_kfunc_set_cid) ?:
1220 		register_btf_kfunc_id_set(BPF_PROG_TYPE_SYSCALL, &scx_kfunc_set_cid);
1221 }
1222