xref: /linux/kernel/sched/ext/cid.c (revision 729beb5aa80ce4bcf40f9ff74ed845269393a6ae)
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 				if (cids_in_core > SCX_CID_SHARD_MAX_CPUS) {
313 					scx_error(sch, "core with %u cpus exceeds max shard size %d",
314 						  cids_in_core, SCX_CID_SHARD_MAX_CPUS);
315 					return -E2BIG;
316 				}
317 
318 				max_cores = cores_per_shard + (shard_local < nr_large ? 1 : 0);
319 				if (cores_in_shard &&
320 				    (cores_in_shard >= max_cores ||
321 				     cids_in_shard + cids_in_core > max_cids)) {
322 					shard_local++;
323 					cores_in_shard = 0;
324 					cids_in_shard = 0;
325 					shard_cid = next_cid;
326 					shard_idx = next_shard_idx++;
327 					tbls->shard_node[shard_idx] = nid;
328 				}
329 				cores_in_shard++;
330 				cids_in_shard += cids_in_core;
331 
332 				for_each_cpu(ccpu, core_scratch) {
333 					s32 cid = next_cid++;
334 
335 					tbls->cid_to_cpu[cid] = ccpu;
336 					tbls->cpu_to_cid[ccpu] = cid;
337 					tbls->cid_to_shard[cid] = shard_idx;
338 					tbls->topo[cid] = (struct scx_cid_topo){
339 						.core_cid = core_cid,
340 						.core_idx = core_idx,
341 						.llc_cid = llc_cid,
342 						.llc_idx = llc_idx,
343 						.node_cid = node_cid,
344 						.node_idx = node_idx,
345 						.shard_cid = shard_cid,
346 						.shard_idx = shard_idx,
347 						.cluster_cid = cluster_cid,
348 						.cluster_idx = cluster_idx,
349 					};
350 
351 					cpumask_clear_cpu(ccpu, cluster_scratch);
352 					cpumask_clear_cpu(ccpu, llc_scratch);
353 					cpumask_clear_cpu(ccpu, node_scratch);
354 					cpumask_clear_cpu(ccpu, to_walk);
355 				}
356 			}
357 		}
358 	}
359 
360 	/*
361 	 * No-topo section: any possible cpu without a cid - normally just the
362 	 * not-online ones. Pack into shards of up to min(@shard_size,
363 	 * SCX_CID_SHARD_MAX_CPUS) cids so that every cid has a valid shard
364 	 * assignment and the hard cap holds even with a large @shard_size.
365 	 * Collect any currently-online cpus that land here in @online_no_topo
366 	 * so we can warn about them at the end.
367 	 */
368 	notopo_in_shard = min_t(u32, shard_size, SCX_CID_SHARD_MAX_CPUS);
369 	notopo_shard_cid = -1;
370 	notopo_shard_idx = -1;
371 
372 	for_each_cpu(cpu, cpu_possible_mask) {
373 		if (tbls->cpu_to_cid[cpu] != -1)
374 			continue;
375 		if (cpu_online(cpu))
376 			cpumask_set_cpu(cpu, online_no_topo);
377 
378 		cid = next_cid++;
379 		tbls->cid_to_cpu[cid] = cpu;
380 		tbls->cpu_to_cid[cpu] = cid;
381 
382 		if (notopo_in_shard >= min_t(u32, shard_size, SCX_CID_SHARD_MAX_CPUS)) {
383 			notopo_shard_cid = cid;
384 			notopo_shard_idx = next_shard_idx++;
385 			notopo_in_shard = 0;
386 		}
387 		notopo_in_shard++;
388 
389 		tbls->cid_to_shard[cid] = notopo_shard_idx;
390 		tbls->topo[cid] = SCX_CID_TOPO_NEG;
391 		tbls->topo[cid].shard_cid = notopo_shard_cid;
392 		tbls->topo[cid].shard_idx = notopo_shard_idx;
393 	}
394 
395 	if (!cpumask_empty(llc_fallback))
396 		pr_warn("scx_cid: cpus without cacheinfo, using node mask as llc: %*pbl\n",
397 			cpumask_pr_args(llc_fallback));
398 	if (!cpumask_empty(online_no_topo))
399 		pr_warn("scx_cid: online cpus with no usable topology: %*pbl\n",
400 			cpumask_pr_args(online_no_topo));
401 
402 	/*
403 	 * Fill cid_shard_ranges[] from cid_to_shard[]. Shards are contiguous
404 	 * cid ranges by construction: base_cid is the first cid landing in a
405 	 * shard, nr_cids is the count.
406 	 */
407 	for (cid = 0; cid < next_cid; cid++) {
408 		s32 sidx = tbls->cid_to_shard[cid];
409 
410 		if (tbls->shard_ranges[sidx].nr_cids == 0)
411 			tbls->shard_ranges[sidx].base_cid = cid;
412 		tbls->shard_ranges[sidx].nr_cids++;
413 	}
414 
415 	tbls->nr_shards = next_shard_idx;
416 	return 0;
417 }
418 
419 /**
420  * scx_cmask_clear - Zero every bit in @m's active range
421  * @m: cmask to clear
422  *
423  * Storage past the active range is left as is.
424  */
425 void scx_cmask_clear(struct scx_cmask *m)
426 {
427 	u32 nr_words;
428 
429 	if (!m->nr_cids)
430 		return;
431 	nr_words = (m->base + m->nr_cids - 1) / 64 - m->base / 64 + 1;
432 	memset(m->bits, 0, nr_words * sizeof(u64));
433 }
434 
435 /**
436  * scx_cmask_fill - Set every bit in @m's active range
437  * @m: cmask to fill
438  *
439  * Counterpart to scx_cmask_clear(). Storage past the active range is left as is.
440  */
441 void scx_cmask_fill(struct scx_cmask *m)
442 {
443 	u32 nr_words, head_bits, tail_bits;
444 
445 	if (!m->nr_cids)
446 		return;
447 	nr_words = (m->base + m->nr_cids - 1) / 64 - m->base / 64 + 1;
448 	memset(m->bits, 0xff, nr_words * sizeof(u64));
449 
450 	/* clear word-0 bits below base */
451 	head_bits = m->base & 63;
452 	if (head_bits)
453 		m->bits[0] &= ~((1ULL << head_bits) - 1);
454 
455 	/* clear last-word bits at or past base + nr_cids */
456 	tail_bits = (m->base + m->nr_cids) & 63;
457 	if (tail_bits)
458 		m->bits[nr_words - 1] &= (1ULL << tail_bits) - 1;
459 }
460 
461 /*
462  * Return the index of the largest entry in @counts, or NUMA_NO_NODE if all
463  * entries are zero. Ties resolve to the lowest index.
464  */
465 static s32 pick_max_node(const u32 *counts, u32 n)
466 {
467 	s32 best = NUMA_NO_NODE;
468 	u32 best_count = 0, i;
469 
470 	for (i = 0; i < n; i++) {
471 		if (counts[i] > best_count) {
472 			best_count = counts[i];
473 			best = i;
474 		}
475 	}
476 	return best;
477 }
478 
479 __bpf_kfunc_start_defs();
480 
481 /**
482  * scx_bpf_cid_override - Install an explicit cpu->cid mapping with shard info
483  * @cpu_to_cid__arena: array of nr_cpu_ids s32 entries (cid for each cpu)
484  * @cpu_to_cid_cnt: number of entries, must be nr_cpu_ids
485  * @shard_start__arena: array of first-cid-of-each-shard, one entry per shard
486  * @shard_start_cnt: number of shards
487  * @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
488  *
489  * May only be called from ops.init_cids() of the root scheduler. Replace the
490  * topology-probed cid mapping and shard layout with caller-provided ones. Each
491  * possible cpu must map to a unique cid in [0, num_possible_cpus()). The shard
492  * starts must be strictly increasing with the first entry 0 and all values <
493  * num_possible_cpus(). The last shard extends to num_possible_cpus() and no
494  * shard may span more than SCX_CID_SHARD_MAX_CPUS cids. Topo info
495  * (core/cluster/LLC/node) is cleared and the shard layout is set from the
496  * input. On invalid input, abort the scheduler.
497  */
498 __bpf_kfunc void scx_bpf_cid_override(const s32 *cpu_to_cid__arena, u32 cpu_to_cid_cnt,
499 				      const s32 *shard_start__arena, u32 shard_start_cnt,
500 				      const struct bpf_prog_aux *aux)
501 {
502 	cpumask_var_t seen __free(free_cpumask_var) = CPUMASK_VAR_NULL;
503 	u32 *node_counts __free(kfree) = NULL;
504 	s32 *cpu_to_cid __free(kfree) = NULL;
505 	s32 *shard_start __free(kfree) = NULL;
506 	u32 npossible = num_possible_cpus();
507 	struct scx_cid_tables *tbls;
508 	struct scx_sched *sch;
509 	u32 nr_shards = shard_start_cnt;
510 	bool alloced;
511 	s32 cpu, cid, si;
512 
513 	/*
514 	 * GFP_KERNEL allocs must happen before the rcu read section. Snapshot
515 	 * the BPF-supplied arrays so a concurrent arena write can't change
516 	 * them between validation and use.
517 	 *
518 	 * The BPF-supplied counts size the snapshots and thus the arena reads.
519 	 * Gate the copies on the count bounds, reported below once @sch is
520 	 * available. The bounded reads, at most 32KB, stay within the guard
521 	 * region that arena fault recovery covers.
522 	 */
523 	alloced = zalloc_cpumask_var(&seen, GFP_KERNEL);
524 	node_counts = kzalloc_objs(*node_counts, nr_node_ids);
525 	if (cpu_to_cid_cnt == nr_cpu_ids)
526 		cpu_to_cid = kmemdup(cpu_to_cid__arena, cpu_to_cid_cnt * sizeof(s32),
527 				     GFP_KERNEL);
528 	if (nr_shards && nr_shards <= npossible)
529 		shard_start = kmemdup(shard_start__arena, nr_shards * sizeof(s32),
530 				      GFP_KERNEL);
531 
532 	guard(rcu)();
533 
534 	sch = scx_prog_sched(aux);
535 	if (unlikely(!sch))
536 		return;
537 
538 	/* called from ops.init_cids(), so the tables exist and are unpublished */
539 	lockdep_assert_held(&scx_enable_mutex);
540 	tbls = scx_cid_tables;
541 
542 	if (cpu_to_cid_cnt != nr_cpu_ids) {
543 		scx_error(sch, "scx_bpf_cid_override: cpu_to_cid expected %u entries, got %u",
544 			  nr_cpu_ids, cpu_to_cid_cnt);
545 		return;
546 	}
547 
548 	if (!nr_shards || nr_shards > npossible) {
549 		scx_error(sch, "scx_bpf_cid_override: invalid shard_start count %u",
550 			  nr_shards);
551 		return;
552 	}
553 
554 	if (!alloced || !node_counts || !cpu_to_cid || !shard_start) {
555 		scx_error(sch, "scx_bpf_cid_override: allocation failed");
556 		return;
557 	}
558 
559 	/* validate shard_start[]: starts at 0, strictly increasing, in range */
560 	if (shard_start[0] != 0) {
561 		scx_error(sch, "scx_bpf_cid_override: shard_start[0] must be 0, got %d",
562 			  shard_start[0]);
563 		return;
564 	}
565 	for (si = 1; si < nr_shards; si++) {
566 		if (shard_start[si] <= shard_start[si - 1]) {
567 			scx_error(sch, "scx_bpf_cid_override: shard_start not increasing at [%d]",
568 				  si);
569 			return;
570 		}
571 		if (shard_start[si] >= npossible) {
572 			scx_error(sch, "scx_bpf_cid_override: shard_start[%d]=%d >= %u",
573 				  si, shard_start[si], npossible);
574 			return;
575 		}
576 		if (shard_start[si] - shard_start[si - 1] > SCX_CID_SHARD_MAX_CPUS) {
577 			scx_error(sch, "scx_bpf_cid_override: shard[%d] span %d exceeds max %d",
578 				  si - 1, shard_start[si] - shard_start[si - 1],
579 				  SCX_CID_SHARD_MAX_CPUS);
580 			return;
581 		}
582 	}
583 	if (npossible - shard_start[nr_shards - 1] > SCX_CID_SHARD_MAX_CPUS) {
584 		scx_error(sch, "scx_bpf_cid_override: shard[%d] span %d exceeds max %d",
585 			  nr_shards - 1, npossible - shard_start[nr_shards - 1],
586 			  SCX_CID_SHARD_MAX_CPUS);
587 		return;
588 	}
589 
590 	/* validate first so that invalid input leaves the tables untouched */
591 	for_each_possible_cpu(cpu) {
592 		s32 c = cpu_to_cid[cpu];
593 
594 		if (!cid_valid(sch, c))
595 			return;
596 		if (cpumask_test_and_set_cpu(c, seen)) {
597 			scx_error(sch, "cid %d assigned to multiple cpus", c);
598 			return;
599 		}
600 	}
601 
602 	for_each_possible_cpu(cpu) {
603 		s32 c = cpu_to_cid[cpu];
604 
605 		tbls->cpu_to_cid[cpu] = c;
606 		tbls->cid_to_cpu[c] = cpu;
607 	}
608 
609 	/*
610 	 * Derive shard_node[] by majority count: an overridden shard may
611 	 * span NUMA nodes, so assign each to the node that owns the most cpus.
612 	 */
613 	for (si = 0; si < nr_shards; si++) {
614 		u32 end = (si + 1 < nr_shards) ? shard_start[si + 1] : npossible;
615 
616 		memset(node_counts, 0, nr_node_ids * sizeof(*node_counts));
617 		for (cid = shard_start[si]; cid < end; cid++) {
618 			s32 node = cpu_to_node(tbls->cid_to_cpu[cid]);
619 
620 			if (numa_valid_node(node))
621 				node_counts[node]++;
622 		}
623 		tbls->shard_node[si] = pick_max_node(node_counts, nr_node_ids);
624 	}
625 
626 	/*
627 	 * Invalidate stale topo info and install shard layout from
628 	 * @shard_start. Walk shards to derive shard_cid/shard_idx for each cid.
629 	 */
630 	si = 0;
631 	for (cid = 0; cid < npossible; cid++) {
632 		if (si + 1 < nr_shards && cid >= shard_start[si + 1])
633 			si++;
634 		tbls->cid_to_shard[cid] = si;
635 		tbls->topo[cid] = SCX_CID_TOPO_NEG;
636 		tbls->topo[cid].shard_cid = shard_start[si];
637 		tbls->topo[cid].shard_idx = si;
638 	}
639 
640 	/* Rebuild shard_ranges[] for the new layout. */
641 	memset(tbls->shard_ranges, 0, npossible * sizeof(*tbls->shard_ranges));
642 	for (si = 0; si < nr_shards; si++) {
643 		u32 end = (si + 1 < nr_shards) ? shard_start[si + 1] : npossible;
644 
645 		tbls->shard_ranges[si].base_cid = shard_start[si];
646 		tbls->shard_ranges[si].nr_cids = end - shard_start[si];
647 	}
648 
649 	tbls->nr_shards = nr_shards;
650 }
651 
652 /**
653  * scx_bpf_cid_to_cpu - Return the raw CPU id for @cid
654  * @cid: cid to look up
655  * @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
656  *
657  * Return the raw CPU id for @cid. Trigger scx_error() and return -EINVAL if
658  * @cid is invalid. The cid<->cpu mapping is static for the lifetime of the
659  * loaded scheduler, so the BPF side can cache the result to avoid repeated
660  * kfunc invocations.
661  */
662 __bpf_kfunc s32 scx_bpf_cid_to_cpu(s32 cid, const struct bpf_prog_aux *aux)
663 {
664 	struct scx_sched *sch;
665 
666 	guard(rcu)();
667 
668 	sch = scx_prog_sched(aux);
669 	if (unlikely(!sch))
670 		return -EINVAL;
671 	return scx_cid_to_cpu(sch, cid);
672 }
673 
674 /**
675  * scx_bpf_cid_node - Return the NUMA node the given @cid belongs to
676  * @cid: cid to look up
677  * @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
678  *
679  * Return NUMA_NO_NODE if @cid is invalid or there is no scheduler. Unlike
680  * scx_bpf_cpu_node(), this kfunc is available only with CID support, so BPF
681  * schedulers can use its presence to detect NUMA-aware CID lookups.
682  */
683 __bpf_kfunc s32 scx_bpf_cid_node(s32 cid, const struct bpf_prog_aux *aux)
684 {
685 	struct scx_sched *sch;
686 	s32 cpu;
687 
688 	guard(rcu)();
689 
690 	sch = scx_prog_sched(aux);
691 	if (unlikely(!sch))
692 		return NUMA_NO_NODE;
693 	cpu = scx_cid_to_cpu(sch, cid);
694 	return cpu < 0 ? NUMA_NO_NODE : cpu_to_node(cpu);
695 }
696 
697 /**
698  * scx_bpf_cpu_to_cid - Return the cid for @cpu
699  * @cpu: cpu to look up
700  * @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
701  *
702  * Return the cid for @cpu. Trigger scx_error() and return -EINVAL if @cpu is
703  * invalid. The cid<->cpu mapping is static for the lifetime of the loaded
704  * scheduler, so the BPF side can cache the result to avoid repeated kfunc
705  * invocations.
706  */
707 __bpf_kfunc s32 scx_bpf_cpu_to_cid(s32 cpu, const struct bpf_prog_aux *aux)
708 {
709 	struct scx_sched *sch;
710 
711 	guard(rcu)();
712 
713 	sch = scx_prog_sched(aux);
714 	if (unlikely(!sch))
715 		return -EINVAL;
716 	return scx_cpu_to_cid(sch, cpu);
717 }
718 
719 /*
720  * Set ops on cmasks. cmask_walk_op2() shares one walk across mutating
721  * (and/or/copy/andnot) and predicate (subset/intersects) two-cmask forms;
722  * cmask_walk_op1() does the same shape over a single cmask range. Every public
723  * entry passes a compile-time-constant @op; cmask_walk_op{1,2}() and
724  * cmask_word_op{1,2}() are __always_inline so the inner switch collapses to the
725  * selected op and cmask_op2_is_pred() folds the predicate early-exit out of
726  * mutating ops.
727  *
728  * Two-cmask ops only touch @dst bits inside the intersection of the two ranges;
729  * bits outside stay untouched. In particular, scx_cmask_copy() does NOT zero
730  * @dst bits that lie outside @src's range.
731  *
732  * Word accesses use READ_ONCE/WRITE_ONCE so a caller may read @src
733  * locklessly. Memory ordering against concurrent writers is the caller's
734  * responsibility.
735  */
736 enum cmask_op2 {
737 	/* mutating */
738 	CMASK_OP2_AND,
739 	CMASK_OP2_OR,
740 	CMASK_OP2_COPY,
741 	CMASK_OP2_ANDNOT,
742 	/* predicates - short-circuit when the per-word result is true */
743 	CMASK_OP2_SUBSET,
744 	CMASK_OP2_INTERSECTS,
745 	/*
746 	 * @a is a BPF-arena cmask. Words on @a use READ_ONCE/WRITE_ONCE since
747 	 * BPF may read/write concurrently. See scx_cmask_ref_or() / _copy().
748 	 */
749 	CMASK_OP2_REF_OR,
750 	CMASK_OP2_REF_COPY,
751 };
752 
753 static __always_inline bool cmask_op2_is_pred(const enum cmask_op2 op)
754 {
755 	return op == CMASK_OP2_SUBSET || op == CMASK_OP2_INTERSECTS;
756 }
757 
758 static __always_inline bool cmask_word_op2(u64 *av, const u64 *bp, u64 mask,
759 					   const enum cmask_op2 op)
760 {
761 	switch (op) {
762 	case CMASK_OP2_AND:
763 		WRITE_ONCE(*av, *av & (~mask | READ_ONCE(*bp)));
764 		return false;
765 	case CMASK_OP2_OR:
766 		WRITE_ONCE(*av, *av | (READ_ONCE(*bp) & mask));
767 		return false;
768 	case CMASK_OP2_COPY:
769 		WRITE_ONCE(*av, (*av & ~mask) | (READ_ONCE(*bp) & mask));
770 		return false;
771 	case CMASK_OP2_ANDNOT:
772 		WRITE_ONCE(*av, *av & ~(READ_ONCE(*bp) & mask));
773 		return false;
774 	case CMASK_OP2_SUBSET:
775 		/* stop on the first bit in @sub not set in @super */
776 		return (READ_ONCE(*bp) & ~READ_ONCE(*av)) & mask;
777 	case CMASK_OP2_INTERSECTS:
778 		return (READ_ONCE(*av) & READ_ONCE(*bp)) & mask;
779 	case CMASK_OP2_REF_OR:
780 		WRITE_ONCE(*av, READ_ONCE(*av) | (READ_ONCE(*bp) & mask));
781 		return false;
782 	case CMASK_OP2_REF_COPY:
783 		WRITE_ONCE(*av, (READ_ONCE(*av) & ~mask) | (READ_ONCE(*bp) & mask));
784 		return false;
785 	}
786 	unreachable();
787 }
788 
789 /*
790  * Walk the intersection of [@a_base, @a_base + @a_nr_cids) with [@b_base,
791  * @b_base + @b_nr_cids) word by word, applying @op. Mutating ops walk all words
792  * and return false; predicates return true on the first word whose per-word
793  * test is true. Empty intersection returns false (matches "no bits to consider"
794  * for both mutate and predicate).
795  *
796  * Base/nr_cids are taken as parameters so callers with snapshotted bounds can
797  * drive the walk with values independent of the cmask's header.
798  */
799 static __always_inline bool cmask_walk_op2(u64 *a_bits, u32 a_base, u32 a_nr_cids,
800 					   const u64 *b_bits, u32 b_base, u32 b_nr_cids,
801 					   const enum cmask_op2 op)
802 {
803 	u32 lo = max(a_base, b_base);
804 	u32 hi = min(a_base + a_nr_cids, b_base + b_nr_cids);
805 	u32 a_word_off = a_base / 64;
806 	u32 b_word_off = b_base / 64;
807 	u32 lo_word = lo / 64;
808 	u32 hi_word = (hi - 1) / 64;
809 	u64 head_mask = GENMASK_U64(63, lo & 63);
810 	u64 tail_mask = GENMASK_U64((hi - 1) & 63, 0);
811 	u32 w;
812 
813 	if (lo >= hi)
814 		return false;
815 
816 	if (lo_word == hi_word)
817 		return cmask_word_op2(&a_bits[lo_word - a_word_off],
818 				      &b_bits[lo_word - b_word_off],
819 				      head_mask & tail_mask, op);
820 
821 	if (cmask_word_op2(&a_bits[lo_word - a_word_off],
822 			   &b_bits[lo_word - b_word_off], head_mask, op) &&
823 	    cmask_op2_is_pred(op))
824 		return true;
825 
826 	for (w = lo_word + 1; w < hi_word; w++)
827 		if (cmask_word_op2(&a_bits[w - a_word_off],
828 				   &b_bits[w - b_word_off], ~0ULL, op) &&
829 		    cmask_op2_is_pred(op))
830 			return true;
831 
832 	return cmask_word_op2(&a_bits[hi_word - a_word_off],
833 			      &b_bits[hi_word - b_word_off], tail_mask, op);
834 }
835 
836 enum cmask_op1 {
837 	CMASK_OP1_ANY_SET,
838 };
839 
840 static __always_inline bool cmask_word_op1(const u64 *ap, u64 mask,
841 					   const enum cmask_op1 op)
842 {
843 	switch (op) {
844 	case CMASK_OP1_ANY_SET:
845 		return READ_ONCE(*ap) & mask;
846 	}
847 	unreachable();
848 }
849 
850 /*
851  * Walk [@a_base, @a_base + @a_nr_cids) of @a_bits word by word, applying @op.
852  * Returns true on the first word whose per-word test is true; returns false if
853  * no word matches or the range is empty. All current op1s short-circuit on
854  * per-word true; if a non-predicate op1 lands here, add a cmask_op1_is_pred()
855  * guard analogous to cmask_op2_is_pred().
856  */
857 static __always_inline bool cmask_walk_op1(const u64 *a_bits, u32 a_base,
858 					   u32 a_nr_cids,
859 					   const enum cmask_op1 op)
860 {
861 	u32 lo = a_base;
862 	u32 hi = a_base + a_nr_cids;
863 	u32 a_word_off = a_base / 64;
864 	u32 lo_word = lo / 64;
865 	u32 hi_word = (hi - 1) / 64;
866 	u64 head_mask = GENMASK_U64(63, lo & 63);
867 	u64 tail_mask = GENMASK_U64((hi - 1) & 63, 0);
868 	u32 w;
869 
870 	if (lo >= hi)
871 		return false;
872 
873 	if (lo_word == hi_word)
874 		return cmask_word_op1(&a_bits[lo_word - a_word_off],
875 				      head_mask & tail_mask, op);
876 
877 	if (cmask_word_op1(&a_bits[lo_word - a_word_off], head_mask, op))
878 		return true;
879 	for (w = lo_word + 1; w < hi_word; w++)
880 		if (cmask_word_op1(&a_bits[w - a_word_off], ~0ULL, op))
881 			return true;
882 	return cmask_word_op1(&a_bits[hi_word - a_word_off], tail_mask, op);
883 }
884 
885 void scx_cmask_and(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_AND);
889 }
890 
891 void scx_cmask_or(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_OR);
895 }
896 
897 void scx_cmask_copy(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_COPY);
901 }
902 
903 void scx_cmask_andnot(struct scx_cmask *dst, const struct scx_cmask *src)
904 {
905 	cmask_walk_op2(dst->bits, dst->base, dst->nr_cids,
906 		       src->bits, src->base, src->nr_cids, CMASK_OP2_ANDNOT);
907 }
908 
909 /*
910  * Return true if @cm has any bit set in [@lo, @hi). Caller must ensure
911  * [@lo, @hi) is contained in @cm's range.
912  */
913 static bool cmask_any_set_in_range(const struct scx_cmask *cm, u32 lo, u32 hi)
914 {
915 	if (lo >= hi)
916 		return false;
917 	return cmask_walk_op1(&cm->bits[lo / 64 - cm->base / 64], lo, hi - lo,
918 			      CMASK_OP1_ANY_SET);
919 }
920 
921 /**
922  * scx_cmask_subset - test whether @sub is a subset of @super
923  * @sub: cmask to test
924  * @super: cmask to test against
925  *
926  * Return true iff every set bit of @sub is also set in @super.
927  */
928 bool scx_cmask_subset(const struct scx_cmask *sub, const struct scx_cmask *super)
929 {
930 	u32 super_end = super->base + super->nr_cids;
931 	u32 sub_end = sub->base + sub->nr_cids;
932 
933 	/*
934 	 * Set bits in @sub outside @super's range can't be in @super, so any
935 	 * such bit means not a subset. The walk below only visits words
936 	 * common to both ranges, so these need a separate scan.
937 	 */
938 	if (sub->base < super->base &&
939 	    cmask_any_set_in_range(sub, sub->base, min(super->base, sub_end)))
940 		return false;
941 	if (sub_end > super_end &&
942 	    cmask_any_set_in_range(sub, max(sub->base, super_end), sub_end))
943 		return false;
944 
945 	return !cmask_walk_op2((u64 *)super->bits, super->base, super->nr_cids,
946 			       sub->bits, sub->base, sub->nr_cids, CMASK_OP2_SUBSET);
947 }
948 
949 bool scx_cmask_intersects(const struct scx_cmask *a, const struct scx_cmask *b)
950 {
951 	return cmask_walk_op2((u64 *)a->bits, a->base, a->nr_cids,
952 			      b->bits, b->base, b->nr_cids, CMASK_OP2_INTERSECTS);
953 }
954 
955 /**
956  * scx_cmask_empty - Test whether @m has no bits set
957  * @m: cmask to test
958  *
959  * Return true iff @m's active range has no bits set.
960  */
961 bool scx_cmask_empty(const struct scx_cmask *m)
962 {
963 	return !cmask_any_set_in_range(m, m->base, m->base + m->nr_cids);
964 }
965 
966 /**
967  * scx_bpf_cid_topo - Copy out per-cid topology info
968  * @cid: cid to look up
969  * @out: where to copy the topology info
970  * @out__sz: size of @out, the program's sizeof(struct scx_cid_topo)
971  * @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
972  *
973  * Fill @out with the topology info for @cid. Trigger scx_error() if @cid is out
974  * of range. If @cid is valid but in the no-topo section, all fields are set to
975  * -1. All fields are also set to -1 when no cid tables have been published yet,
976  * which a program may observe while racing the root enable.
977  *
978  * The program's struct may be older or newer than the kernel's. The smaller of
979  * @out__sz and the kernel's size is copied and the rest of @out is set to -1.
980  */
981 __bpf_kfunc void scx_bpf_cid_topo(s32 cid, struct scx_cid_topo *out, size_t out__sz,
982 				  const struct bpf_prog_aux *aux)
983 {
984 	size_t len = min(out__sz, sizeof(*out));
985 	struct scx_cid_topo *topo;
986 	struct scx_sched *sch;
987 
988 	/* the error cases and fields the kernel lacks read as -1 */
989 	memset(out, 0xff, out__sz);
990 
991 	guard(rcu)();
992 
993 	sch = scx_prog_sched(aux);
994 	topo = rcu_dereference(scx_cid_topo);
995 	if (unlikely(!sch) || !cid_valid(sch, cid) || unlikely(!topo))
996 		return;
997 
998 	memcpy(out, &topo[cid], len);
999 }
1000 
1001 __bpf_kfunc_end_defs();
1002 
1003 BTF_KFUNCS_START(scx_kfunc_ids_init_cids)
1004 BTF_ID_FLAGS(func, scx_bpf_cid_override, KF_IMPLICIT_ARGS | KF_SLEEPABLE)
1005 BTF_KFUNCS_END(scx_kfunc_ids_init_cids)
1006 
1007 static const struct btf_kfunc_id_set scx_kfunc_set_init_cids = {
1008 	.owner	= THIS_MODULE,
1009 	.set	= &scx_kfunc_ids_init_cids,
1010 	.filter	= scx_kfunc_context_filter,
1011 };
1012 
1013 BTF_KFUNCS_START(scx_kfunc_ids_cid)
1014 BTF_ID_FLAGS(func, scx_bpf_cid_to_cpu, KF_IMPLICIT_ARGS)
1015 BTF_ID_FLAGS(func, scx_bpf_cid_node, KF_IMPLICIT_ARGS)
1016 BTF_ID_FLAGS(func, scx_bpf_cpu_to_cid, KF_IMPLICIT_ARGS)
1017 BTF_ID_FLAGS(func, scx_bpf_cid_topo, KF_IMPLICIT_ARGS)
1018 BTF_KFUNCS_END(scx_kfunc_ids_cid)
1019 
1020 static const struct btf_kfunc_id_set scx_kfunc_set_cid = {
1021 	.owner	= THIS_MODULE,
1022 	.set	= &scx_kfunc_ids_cid,
1023 };
1024 
1025 /**
1026  * scx_cmask_ref_init - Bind a scx_cmask_ref to a BPF-arena cmask
1027  * @sch: scheduler whose arena hosts @src
1028  * @src: BPF-supplied cmask, rebased to its kernel address
1029  * @ref: output ref
1030  *
1031  * Snapshot @src's @base, @nr_cids and @alloc_words. The snapshot is necessary
1032  * because BPF may mutate the live header asynchronously.
1033  *
1034  * Return 0 on success, -EINVAL if the range is out of bounds or @alloc_words
1035  * doesn't cover it.
1036  */
1037 int scx_cmask_ref_init(struct scx_sched *sch, const struct scx_cmask *src,
1038 		       struct scx_cmask_ref *ref)
1039 {
1040 	u32 base, nr_cids, alloc_words, npossible = num_possible_cpus();
1041 	s32 *cid_to_shard;
1042 
1043 	base = READ_ONCE(src->base);
1044 	nr_cids = READ_ONCE(src->nr_cids);
1045 	alloc_words = READ_ONCE(src->alloc_words);
1046 
1047 	if (unlikely(base >= npossible || nr_cids > npossible - base ||
1048 		     SCX_CMASK_NR_WORDS(nr_cids) > alloc_words))
1049 		return -EINVAL;
1050 
1051 	ref->sch = sch;
1052 	ref->src = (struct scx_cmask *)src;
1053 	ref->base = base;
1054 	ref->nr_cids = nr_cids;
1055 
1056 	cid_to_shard = rcu_dereference_all(scx_cid_to_shard);
1057 	ref->shard_first = cid_to_shard[base];
1058 	if (likely(nr_cids))
1059 		ref->shard_end = cid_to_shard[base + nr_cids - 1] + 1;
1060 	else
1061 		ref->shard_end = ref->shard_first;
1062 
1063 	return 0;
1064 }
1065 
1066 /**
1067  * scx_cmask_ref_init_kern - Bind a scx_cmask_ref to a kernel-owned cmask
1068  * @sch: scheduler the cmask belongs to
1069  * @m: kernel address of the target cmask, storage sized for @nr_cids at @base
1070  * @base: first cid of the active range
1071  * @nr_cids: active range length
1072  * @ref: output ref
1073  *
1074  * Like scx_cmask_ref_init() but the geometry is supplied by the caller, not
1075  * read from @m's header, so a concurrent BPF write to the header can't steer
1076  * later sizing or offsets. Rewrite the header from the trusted geometry and
1077  * bind @ref to it.
1078  */
1079 void scx_cmask_ref_init_kern(struct scx_sched *sch, struct scx_cmask *m,
1080 			     u32 base, u32 nr_cids, struct scx_cmask_ref *ref)
1081 {
1082 	s32 *cid_to_shard;
1083 
1084 	WRITE_ONCE(m->base, base);
1085 	WRITE_ONCE(m->nr_cids, nr_cids);
1086 	WRITE_ONCE(m->alloc_words, SCX_CMASK_NR_WORDS(nr_cids));
1087 
1088 	ref->sch = sch;
1089 	ref->src = m;
1090 	ref->base = base;
1091 	ref->nr_cids = nr_cids;
1092 
1093 	cid_to_shard = rcu_dereference_all(scx_cid_to_shard);
1094 	ref->shard_first = cid_to_shard[base];
1095 	if (likely(nr_cids))
1096 		ref->shard_end = cid_to_shard[base + nr_cids - 1] + 1;
1097 	else
1098 		ref->shard_end = ref->shard_first;
1099 }
1100 
1101 /**
1102  * scx_cmask_ref_shard - Read one shard from @ref into @out
1103  * @ref: validated ref
1104  * @shard_idx: target shard, in [@ref->shard_first, @ref->shard_end)
1105  * @out: output cmask whose @out->alloc_words must hold the shard
1106  *
1107  * Set @out to the intersection of @ref's range with @shard_idx's cid range,
1108  * with bits[] read from @ref->src via READ_ONCE. Empty intersection sets
1109  * @out->nr_cids to 0. scx_error()s on @ref's sched if @out can't hold the
1110  * shard.
1111  */
1112 void scx_cmask_ref_shard(const struct scx_cmask_ref *ref, s32 shard_idx,
1113 			 struct scx_cmask *out)
1114 {
1115 	const struct scx_cid_shard *shard =
1116 		&rcu_dereference_all(scx_cid_shard_ranges)[shard_idx];
1117 	u32 shard_base = shard->base_cid;
1118 	u32 shard_end = shard_base + shard->nr_cids;
1119 	u32 isect_base, isect_end, nr_words, src_off, wi;
1120 	u64 head_mask, tail_mask;
1121 
1122 	isect_base = max(ref->base, shard_base);
1123 	isect_end = min(ref->base + ref->nr_cids, shard_end);
1124 
1125 	if (isect_base >= isect_end) {
1126 		out->base = shard_base;
1127 		out->nr_cids = 0;
1128 		return;
1129 	}
1130 
1131 	nr_words = ((isect_end - 1) / 64) - (isect_base / 64) + 1;
1132 	if (nr_words > out->alloc_words) {
1133 		scx_error(ref->sch, "scx_cmask_ref_shard: out alloc_words=%u < %u for shard %d",
1134 			  out->alloc_words, nr_words, shard_idx);
1135 		out->base = shard_base;
1136 		out->nr_cids = 0;
1137 		return;
1138 	}
1139 
1140 	out->base = isect_base;
1141 	out->nr_cids = isect_end - isect_base;
1142 	src_off = (isect_base / 64) - (ref->base / 64);
1143 
1144 	for (wi = 0; wi < nr_words; wi++)
1145 		out->bits[wi] = READ_ONCE(ref->src->bits[src_off + wi]);
1146 
1147 	head_mask = GENMASK_U64(63, isect_base & 63);
1148 	out->bits[0] &= head_mask;
1149 	tail_mask = GENMASK_U64((isect_end - 1) & 63, 0);
1150 	out->bits[nr_words - 1] &= tail_mask;
1151 }
1152 
1153 /**
1154  * scx_cmask_ref_or - OR @src into the arena cmask referenced by @ref
1155  * @ref: validated ref
1156  * @src: stable kernel cmask
1157  *
1158  * Bits inside the intersection of @ref's snapshotted range with @src's range
1159  * are OR'd into @ref->src and bits outside are left unchanged. Stores on
1160  * @ref->src use WRITE_ONCE since BPF may read/write concurrently.
1161  */
1162 void scx_cmask_ref_or(const struct scx_cmask_ref *ref, const struct scx_cmask *src)
1163 {
1164 	cmask_walk_op2(ref->src->bits, ref->base, ref->nr_cids,
1165 		       src->bits, src->base, src->nr_cids, CMASK_OP2_REF_OR);
1166 }
1167 
1168 /**
1169  * scx_cmask_ref_copy - Copy @src into the arena cmask referenced by @ref
1170  * @ref: validated ref
1171  * @src: stable kernel cmask
1172  *
1173  * Bits inside the intersection of @ref's snapshotted range with @src's range
1174  * take @src's values and bits outside are left unchanged. Stores on @ref->src
1175  * use WRITE_ONCE since BPF may read/write concurrently.
1176  */
1177 void scx_cmask_ref_copy(const struct scx_cmask_ref *ref, const struct scx_cmask *src)
1178 {
1179 	cmask_walk_op2(ref->src->bits, ref->base, ref->nr_cids,
1180 		       src->bits, src->base, src->nr_cids, CMASK_OP2_REF_COPY);
1181 }
1182 
1183 /**
1184  * scx_cmask_ref_from_cpumask - Populate @ref's arena cmask from a cpumask
1185  * @ref: kern-bound ref, see scx_cmask_ref_init_kern()
1186  * @cpumask: cpus to translate into cids
1187  *
1188  * Write @ref's active range one word at a time, setting each cid's bit when
1189  * its cpu is in @cpumask. Offsets and length come from @ref's trusted geometry
1190  * and stores use WRITE_ONCE since BPF may read concurrently, so the arena
1191  * header is never read.
1192  */
1193 void scx_cmask_ref_from_cpumask(const struct scx_cmask_ref *ref,
1194 				const struct cpumask *cpumask)
1195 {
1196 	struct scx_cmask *m = ref->src;
1197 	u32 base = ref->base, nr_cids = ref->nr_cids;
1198 	u32 wi, nr_words;
1199 
1200 	if (!nr_cids)
1201 		return;
1202 
1203 	nr_words = (base + nr_cids - 1) / 64 - base / 64 + 1;
1204 	for (wi = 0; wi < nr_words; wi++) {
1205 		u32 word_first_cid = (base / 64 + wi) * 64;
1206 		u64 word = 0;
1207 		u32 bit;
1208 
1209 		for (bit = 0; bit < 64; bit++) {
1210 			u32 cid = word_first_cid + bit;
1211 
1212 			if (cid < base || cid >= base + nr_cids)
1213 				continue;
1214 			if (cpumask_test_cpu(__scx_cid_to_cpu(cid), cpumask))
1215 				word |= BIT_U64(bit);
1216 		}
1217 		WRITE_ONCE(m->bits[wi], word);
1218 	}
1219 }
1220 
1221 int scx_cid_kfunc_init(void)
1222 {
1223 	return register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS, &scx_kfunc_set_init_cids) ?:
1224 		register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS, &scx_kfunc_set_cid) ?:
1225 		register_btf_kfunc_id_set(BPF_PROG_TYPE_TRACING, &scx_kfunc_set_cid) ?:
1226 		register_btf_kfunc_id_set(BPF_PROG_TYPE_SYSCALL, &scx_kfunc_set_cid);
1227 }
1228