1 /* SPDX-License-Identifier: GPL-2.0 */
2 /*
3 * BPF-side helpers for cids and cmasks. See kernel/sched/ext/cid.h for the
4 * authoritative layout and semantics. The BPF-side helpers use the cmask_*
5 * naming (no scx_ prefix); cmask is the SCX bitmap type so the prefix is
6 * redundant in BPF code. Atomics use __sync_val_compare_and_swap and every
7 * helper is inline (no .c counterpart).
8 *
9 * Included by scx/common.bpf.h; don't include directly.
10 *
11 * Copyright (c) 2026 Meta Platforms, Inc. and affiliates.
12 * Copyright (c) 2026 Tejun Heo <tj@kernel.org>
13 */
14 #ifndef __SCX_CID_BPF_H
15 #define __SCX_CID_BPF_H
16
17 #include "bpf_arena_common.bpf.h"
18
19 #ifndef BIT_U64
20 #define BIT_U64(nr) (1ULL << (nr))
21 #endif
22 #ifndef GENMASK_U64
23 #define GENMASK_U64(h, l) ((~0ULL << (l)) & (~0ULL >> (63 - (h))))
24 #endif
25
26 /*
27 * Storage cap for bounded loops over bits[]. Sized to cover NR_CPUS=8192 with
28 * one extra word for head-misalignment. Increase if deployment targets larger
29 * NR_CPUS.
30 */
31 #ifndef CMASK_MAX_WORDS
32 #define CMASK_MAX_WORDS 129
33 #endif
34
35 /*
36 * Mirrors SCX_CMASK_NR_WORDS in kernel/sched/ext/types.h. The u64 cast keeps
37 * the +63 from wrapping when @nr_cids is near U32_MAX, so cmask_reframe()
38 * bounds-checking the result against alloc_words catches the overflow instead
39 * of seeing a small value.
40 */
41 #define CMASK_NR_WORDS(nr_cids) ((u32)(((u64)(nr_cids) + 63) / 64 + 1))
42
__cmask_contains(u32 cid,const struct scx_cmask __arena * m)43 static __always_inline bool __cmask_contains(u32 cid, const struct scx_cmask __arena *m)
44 {
45 return cid >= m->base && cid < m->base + m->nr_cids;
46 }
47
__cmask_word(u32 cid,const struct scx_cmask __arena * m)48 static __always_inline u64 __arena *__cmask_word(u32 cid, const struct scx_cmask __arena *m)
49 {
50 return (u64 __arena *)&m->bits[cid / 64 - m->base / 64];
51 }
52
53 /**
54 * __cmask_init - Initialize @m with explicit storage capacity
55 * @m: cmask to initialize
56 * @base: first cid of the active range
57 * @nr_cids: number of cids in the active range
58 * @alloc_cids: storage capacity in cids, at least @nr_cids
59 *
60 * Use when storage is sized larger than the initial active range. All of
61 * bits[] is zeroed.
62 */
__cmask_init(struct scx_cmask __arena * m,u32 base,u32 nr_cids,u32 alloc_cids)63 static __always_inline void __cmask_init(struct scx_cmask __arena *m, u32 base,
64 u32 nr_cids, u32 alloc_cids)
65 {
66 u32 alloc_words, i;
67
68 if (unlikely(nr_cids > alloc_cids)) {
69 scx_bpf_error("__cmask_init: nr_cids=%u exceeds alloc_cids=%u",
70 nr_cids, alloc_cids);
71 return;
72 }
73 alloc_words = CMASK_NR_WORDS(alloc_cids);
74
75 m->base = base;
76 m->nr_cids = nr_cids;
77 m->alloc_words = alloc_words;
78
79 bpf_for(i, 0, CMASK_MAX_WORDS) {
80 if (i >= alloc_words)
81 break;
82 m->bits[i] = 0;
83 }
84 }
85
86 /**
87 * cmask_init - Initialize @m on tight storage
88 * @m: cmask to initialize
89 * @base: first cid of the active range
90 * @nr_cids: number of cids in the active range
91 *
92 * All of bits[] is zeroed.
93 */
cmask_init(struct scx_cmask __arena * m,u32 base,u32 nr_cids)94 static __always_inline void cmask_init(struct scx_cmask __arena *m, u32 base, u32 nr_cids)
95 {
96 __cmask_init(m, base, nr_cids, nr_cids);
97 }
98
99 /**
100 * cmask_reframe - Reshape @m's active range without resizing storage
101 * @m: cmask to reframe
102 * @base: new active range base
103 * @nr_cids: new active range length, must fit within @m->alloc_words
104 *
105 * Body bits within the new range become garbage - only the head and tail
106 * words are zeroed to keep the padding invariant.
107 */
cmask_reframe(struct scx_cmask __arena * m,u32 base,u32 nr_cids)108 static __always_inline void cmask_reframe(struct scx_cmask __arena *m, u32 base, u32 nr_cids)
109 {
110 if (CMASK_NR_WORDS(nr_cids) > m->alloc_words) {
111 scx_bpf_error("cmask_reframe: nr_cids=%u exceeds alloc_words=%u",
112 nr_cids, m->alloc_words);
113 return;
114 }
115 if (nr_cids) {
116 u32 last_word = ((base & 63) + nr_cids - 1) / 64;
117
118 m->bits[0] = 0;
119 m->bits[last_word] = 0;
120 }
121 m->base = base;
122 m->nr_cids = nr_cids;
123 }
124
cmask_test(u32 cid,const struct scx_cmask __arena * m)125 static __always_inline bool cmask_test(u32 cid, const struct scx_cmask __arena *m)
126 {
127 if (!__cmask_contains(cid, m))
128 return false;
129 return *__cmask_word(cid, m) & BIT_U64(cid & 63);
130 }
131
132 /*
133 * x86 BPF JIT rejects BPF_OR | BPF_FETCH and BPF_AND | BPF_FETCH on arena
134 * pointers (see bpf_jit_supports_insn() in arch/x86/net/bpf_jit_comp.c). Only
135 * BPF_CMPXCHG / BPF_XCHG / BPF_ADD with FETCH are allowed. Implement
136 * test_and_{set,clear} and the atomic set/clear via a cmpxchg loop.
137 *
138 * CMASK_CAS_TRIES is sized so exhausting it means seconds of real spinning
139 * on one word - past any plausible contention. Abort hard.
140 */
141 #define CMASK_CAS_TRIES (1U << 23)
142
cmask_set(u32 cid,struct scx_cmask __arena * m)143 static __always_inline void cmask_set(u32 cid, struct scx_cmask __arena *m)
144 {
145 u64 __arena *w;
146 u64 bit, old, new;
147 u32 i;
148
149 if (!__cmask_contains(cid, m))
150 return;
151 w = __cmask_word(cid, m);
152 bit = BIT_U64(cid & 63);
153 bpf_for(i, 0, CMASK_CAS_TRIES) {
154 old = *w;
155 if (old & bit)
156 return;
157 new = old | bit;
158 if (__sync_val_compare_and_swap(w, old, new) == old)
159 return;
160 }
161 scx_bpf_error("cmask_set CAS exhausted at cid %u", cid);
162 }
163
cmask_clear(u32 cid,struct scx_cmask __arena * m)164 static __always_inline void cmask_clear(u32 cid, struct scx_cmask __arena *m)
165 {
166 u64 __arena *w;
167 u64 bit, old, new;
168 u32 i;
169
170 if (!__cmask_contains(cid, m))
171 return;
172 w = __cmask_word(cid, m);
173 bit = BIT_U64(cid & 63);
174 bpf_for(i, 0, CMASK_CAS_TRIES) {
175 old = *w;
176 if (!(old & bit))
177 return;
178 new = old & ~bit;
179 if (__sync_val_compare_and_swap(w, old, new) == old)
180 return;
181 }
182 scx_bpf_error("cmask_clear CAS exhausted at cid %u", cid);
183 }
184
cmask_test_and_set(u32 cid,struct scx_cmask __arena * m)185 static __always_inline bool cmask_test_and_set(u32 cid, struct scx_cmask __arena *m)
186 {
187 u64 __arena *w;
188 u64 bit, old, new;
189 u32 i;
190
191 if (!__cmask_contains(cid, m))
192 return false;
193 w = __cmask_word(cid, m);
194 bit = BIT_U64(cid & 63);
195 bpf_for(i, 0, CMASK_CAS_TRIES) {
196 old = *w;
197 if (old & bit)
198 return true;
199 new = old | bit;
200 if (__sync_val_compare_and_swap(w, old, new) == old)
201 return false;
202 }
203 scx_bpf_error("cmask_test_and_set CAS exhausted at cid %u", cid);
204 return false;
205 }
206
cmask_test_and_clear(u32 cid,struct scx_cmask __arena * m)207 static __always_inline bool cmask_test_and_clear(u32 cid, struct scx_cmask __arena *m)
208 {
209 u64 __arena *w;
210 u64 bit, old, new;
211 u32 i;
212
213 if (!__cmask_contains(cid, m))
214 return false;
215 w = __cmask_word(cid, m);
216 bit = BIT_U64(cid & 63);
217 bpf_for(i, 0, CMASK_CAS_TRIES) {
218 old = *w;
219 if (!(old & bit))
220 return false;
221 new = old & ~bit;
222 if (__sync_val_compare_and_swap(w, old, new) == old)
223 return true;
224 }
225 scx_bpf_error("cmask_test_and_clear CAS exhausted at cid %u", cid);
226 return false;
227 }
228
__cmask_set(u32 cid,struct scx_cmask __arena * m)229 static __always_inline void __cmask_set(u32 cid, struct scx_cmask __arena *m)
230 {
231 if (!__cmask_contains(cid, m))
232 return;
233 *__cmask_word(cid, m) |= BIT_U64(cid & 63);
234 }
235
__cmask_clear(u32 cid,struct scx_cmask __arena * m)236 static __always_inline void __cmask_clear(u32 cid, struct scx_cmask __arena *m)
237 {
238 if (!__cmask_contains(cid, m))
239 return;
240 *__cmask_word(cid, m) &= ~BIT_U64(cid & 63);
241 }
242
__cmask_test_and_set(u32 cid,struct scx_cmask __arena * m)243 static __always_inline bool __cmask_test_and_set(u32 cid, struct scx_cmask __arena *m)
244 {
245 u64 bit = BIT_U64(cid & 63);
246 u64 __arena *w;
247 u64 prev;
248
249 if (!__cmask_contains(cid, m))
250 return false;
251 w = __cmask_word(cid, m);
252 prev = *w & bit;
253 *w |= bit;
254 return prev;
255 }
256
__cmask_test_and_clear(u32 cid,struct scx_cmask __arena * m)257 static __always_inline bool __cmask_test_and_clear(u32 cid, struct scx_cmask __arena *m)
258 {
259 u64 bit = BIT_U64(cid & 63);
260 u64 __arena *w;
261 u64 prev;
262
263 if (!__cmask_contains(cid, m))
264 return false;
265 w = __cmask_word(cid, m);
266 prev = *w & bit;
267 *w &= ~bit;
268 return prev;
269 }
270
cmask_zero(struct scx_cmask __arena * m)271 static __always_inline void cmask_zero(struct scx_cmask __arena *m)
272 {
273 u32 nr_words = CMASK_NR_WORDS(m->nr_cids), i;
274
275 bpf_for(i, 0, CMASK_MAX_WORDS) {
276 if (i >= nr_words)
277 break;
278 m->bits[i] = 0;
279 }
280 }
281
282 /*
283 * BPF_-prefixed to avoid colliding with the kernel's anonymous CMASK_OP_*
284 * enum in ext/cid.c, which is exported via BTF and reachable through
285 * vmlinux.h.
286 */
287 enum {
288 BPF_CMASK_OP_AND,
289 BPF_CMASK_OP_OR,
290 BPF_CMASK_OP_COPY,
291 BPF_CMASK_OP_ANDNOT,
292 };
293
cmask_op_word(struct scx_cmask __arena * dst,const struct scx_cmask __arena * src,u32 di,u32 si,u64 mask,int op)294 static __always_inline void cmask_op_word(struct scx_cmask __arena *dst,
295 const struct scx_cmask __arena *src,
296 u32 di, u32 si, u64 mask, int op)
297 {
298 u64 dv = dst->bits[di];
299 u64 sv = src->bits[si];
300 u64 rv;
301
302 if (op == BPF_CMASK_OP_AND)
303 rv = dv & sv;
304 else if (op == BPF_CMASK_OP_OR)
305 rv = dv | sv;
306 else if (op == BPF_CMASK_OP_ANDNOT)
307 rv = dv & ~sv;
308 else
309 rv = sv;
310
311 dst->bits[di] = (dv & ~mask) | (rv & mask);
312 }
313
cmask_op(struct scx_cmask __arena * dst,const struct scx_cmask __arena * src,int op)314 static __always_inline void cmask_op(struct scx_cmask __arena *dst,
315 const struct scx_cmask __arena *src, int op)
316 {
317 u32 d_end = dst->base + dst->nr_cids;
318 u32 s_end = src->base + src->nr_cids;
319 u32 lo = dst->base > src->base ? dst->base : src->base;
320 u32 hi = d_end < s_end ? d_end : s_end;
321 u32 d_base = dst->base / 64;
322 u32 s_base = src->base / 64;
323 u32 lo_word, hi_word, i;
324 u64 head_mask, tail_mask;
325
326 if (lo >= hi)
327 return;
328
329 lo_word = lo / 64;
330 hi_word = (hi - 1) / 64;
331 head_mask = GENMASK_U64(63, lo & 63);
332 tail_mask = GENMASK_U64((hi - 1) & 63, 0);
333
334 bpf_for(i, 0, CMASK_MAX_WORDS) {
335 u32 w = lo_word + i;
336 u64 m;
337
338 if (w > hi_word)
339 break;
340
341 m = GENMASK_U64(63, 0);
342 if (w == lo_word)
343 m &= head_mask;
344 if (w == hi_word)
345 m &= tail_mask;
346
347 cmask_op_word(dst, src, w - d_base, w - s_base, m, op);
348 }
349 }
350
351 /*
352 * cmask_and/or/copy only modify @dst bits that lie in the intersection of
353 * [@dst->base, @dst->base + @dst->nr_cids) and [@src->base,
354 * @src->base + @src->nr_cids). Bits in @dst outside that window
355 * keep their prior values - in particular, cmask_copy() does NOT zero @dst
356 * bits that lie outside @src's range.
357 */
cmask_and(struct scx_cmask __arena * dst,const struct scx_cmask __arena * src)358 static __always_inline void cmask_and(struct scx_cmask __arena *dst,
359 const struct scx_cmask __arena *src)
360 {
361 cmask_op(dst, src, BPF_CMASK_OP_AND);
362 }
363
cmask_or(struct scx_cmask __arena * dst,const struct scx_cmask __arena * src)364 static __always_inline void cmask_or(struct scx_cmask __arena *dst,
365 const struct scx_cmask __arena *src)
366 {
367 cmask_op(dst, src, BPF_CMASK_OP_OR);
368 }
369
cmask_copy(struct scx_cmask __arena * dst,const struct scx_cmask __arena * src)370 static __always_inline void cmask_copy(struct scx_cmask __arena *dst,
371 const struct scx_cmask __arena *src)
372 {
373 cmask_op(dst, src, BPF_CMASK_OP_COPY);
374 }
375
cmask_andnot(struct scx_cmask __arena * dst,const struct scx_cmask __arena * src)376 static __always_inline void cmask_andnot(struct scx_cmask __arena *dst,
377 const struct scx_cmask __arena *src)
378 {
379 cmask_op(dst, src, BPF_CMASK_OP_ANDNOT);
380 }
381
382 /*
383 * True iff @a and @b have identical bits over their (assumed equal) range.
384 * Callers are expected to pass same-shape cmasks; differing shapes always
385 * compare unequal.
386 */
cmask_equal(const struct scx_cmask __arena * a,const struct scx_cmask __arena * b)387 static __always_inline bool cmask_equal(const struct scx_cmask __arena *a,
388 const struct scx_cmask __arena *b)
389 {
390 u32 nr_words, i;
391
392 if (a->base != b->base || a->nr_cids != b->nr_cids)
393 return false;
394 if (a->nr_cids == 0)
395 return true;
396 nr_words = (a->base + a->nr_cids - 1) / 64 - a->base / 64 + 1;
397
398 bpf_for(i, 0, CMASK_MAX_WORDS) {
399 if (i >= nr_words)
400 break;
401 if (a->bits[i] != b->bits[i])
402 return false;
403 }
404 return true;
405 }
406
407 /**
408 * cmask_next_set - find the first set bit at or after @cid
409 * @m: cmask to search
410 * @cid: starting cid (clamped to @m->base if below)
411 *
412 * Returns the smallest set cid in [@cid, @m->base + @m->nr_cids), or
413 * @m->base + @m->nr_cids if none (the out-of-range sentinel matches the
414 * termination condition used by cmask_for_each()).
415 */
cmask_next_set(const struct scx_cmask __arena * m,u32 cid)416 static __always_inline u32 cmask_next_set(const struct scx_cmask __arena *m, u32 cid)
417 {
418 u32 end = m->base + m->nr_cids;
419 u32 base = m->base / 64;
420 u32 last_wi = (end - 1) / 64 - base;
421 u32 start_wi, start_bit, i;
422
423 if (cid < m->base)
424 cid = m->base;
425 if (cid >= end)
426 return end;
427
428 start_wi = cid / 64 - base;
429 start_bit = cid & 63;
430
431 bpf_for(i, 0, CMASK_MAX_WORDS) {
432 u32 wi = start_wi + i;
433 u64 word;
434 u32 found;
435
436 if (wi > last_wi)
437 break;
438
439 word = m->bits[wi];
440 if (i == 0)
441 word &= GENMASK_U64(63, start_bit);
442 if (!word)
443 continue;
444
445 found = (base + wi) * 64 + ctzll(word);
446 if (found >= end)
447 return end;
448 return found;
449 }
450 return end;
451 }
452
cmask_first_set(const struct scx_cmask __arena * m)453 static __always_inline u32 cmask_first_set(const struct scx_cmask __arena *m)
454 {
455 return cmask_next_set(m, m->base);
456 }
457
458 #define cmask_for_each(cid, m) \
459 for ((cid) = cmask_first_set(m); \
460 (cid) < (m)->base + (m)->nr_cids; \
461 (cid) = cmask_next_set((m), (cid) + 1))
462
463 /*
464 * True iff every bit set in @a is also set in @b. Matches the kernel-side
465 * scx_cmask_subset(): ranges don't need to nest, and set bits of @a outside
466 * @b's range fail the test.
467 */
cmask_subset(const struct scx_cmask __arena * a,const struct scx_cmask __arena * b)468 static __always_inline bool cmask_subset(const struct scx_cmask __arena *a,
469 const struct scx_cmask __arena *b)
470 {
471 u32 a_end = a->base + a->nr_cids;
472 u32 b_end = b->base + b->nr_cids;
473 u32 a_wbase = a->base / 64;
474 u32 b_wbase = b->base / 64;
475 u32 lo = a->base > b->base ? a->base : b->base;
476 u32 hi = a_end < b_end ? a_end : b_end;
477 u32 lo_word, hi_word, i;
478
479 /* set bits of @a outside @b's range can't be in @b */
480 if (a->base < b->base &&
481 cmask_next_set(a, a->base) < (b->base < a_end ? b->base : a_end))
482 return false;
483 if (a_end > b_end &&
484 cmask_next_set(a, a->base > b_end ? a->base : b_end) < a_end)
485 return false;
486
487 if (lo >= hi)
488 return true;
489
490 /*
491 * Walk the words the range intersection spans. Plain word tests
492 * suffice: the scans above guarantee @a has no set bit outside @b's
493 * range and padding bits are kept clear by all cmask helpers.
494 */
495 lo_word = lo / 64;
496 hi_word = (hi - 1) / 64;
497
498 bpf_for(i, 0, CMASK_MAX_WORDS) {
499 u32 w = lo_word + i;
500
501 if (w > hi_word)
502 break;
503 if (a->bits[w - a_wbase] & ~b->bits[w - b_wbase])
504 return false;
505 }
506 return true;
507 }
508
509 /*
510 * Population count over [base, base + nr_cids). Padding bits in the head/tail
511 * words are guaranteed zero by the mutating helpers, so a flat popcount over
512 * the words the range spans is correct.
513 */
cmask_weight(const struct scx_cmask __arena * m)514 static __always_inline u32 cmask_weight(const struct scx_cmask __arena *m)
515 {
516 u32 nr_words, i;
517 u32 count = 0;
518
519 if (!m->nr_cids)
520 return 0;
521 nr_words = (m->base + m->nr_cids - 1) / 64 - m->base / 64 + 1;
522
523 bpf_for(i, 0, CMASK_MAX_WORDS) {
524 if (i >= nr_words)
525 break;
526 count += __builtin_popcountll(m->bits[i]);
527 }
528 return count;
529 }
530
531 /*
532 * True if @a and @b share any set bit. Walk only the intersection of their
533 * ranges, matching the semantics of cmask_and().
534 */
cmask_intersects(const struct scx_cmask __arena * a,const struct scx_cmask __arena * b)535 static __always_inline bool cmask_intersects(const struct scx_cmask __arena *a,
536 const struct scx_cmask __arena *b)
537 {
538 u32 a_end = a->base + a->nr_cids;
539 u32 b_end = b->base + b->nr_cids;
540 u32 lo = a->base > b->base ? a->base : b->base;
541 u32 hi = a_end < b_end ? a_end : b_end;
542 u32 a_base = a->base / 64;
543 u32 b_base = b->base / 64;
544 u32 lo_word, hi_word, i;
545 u64 head_mask, tail_mask;
546
547 if (lo >= hi)
548 return false;
549
550 lo_word = lo / 64;
551 hi_word = (hi - 1) / 64;
552 head_mask = GENMASK_U64(63, lo & 63);
553 tail_mask = GENMASK_U64((hi - 1) & 63, 0);
554
555 bpf_for(i, 0, CMASK_MAX_WORDS) {
556 u32 w = lo_word + i;
557 u64 mask, av, bv;
558
559 if (w > hi_word)
560 break;
561
562 mask = GENMASK_U64(63, 0);
563 if (w == lo_word)
564 mask &= head_mask;
565 if (w == hi_word)
566 mask &= tail_mask;
567
568 av = a->bits[w - a_base] & mask;
569 bv = b->bits[w - b_base] & mask;
570 if (av & bv)
571 return true;
572 }
573 return false;
574 }
575
576 /*
577 * Find the next cid set in both @a and @b at or after @start, bounded by the
578 * intersection of the two ranges. Return a->base + a->nr_cids if none found.
579 *
580 * Building block for cmask_next_and_set_wrap(). Callers that want a bounded
581 * scan without wrap call this directly.
582 */
cmask_next_and_set(const struct scx_cmask __arena * a,const struct scx_cmask __arena * b,u32 start)583 static __always_inline u32 cmask_next_and_set(const struct scx_cmask __arena *a,
584 const struct scx_cmask __arena *b,
585 u32 start)
586 {
587 u32 a_end = a->base + a->nr_cids;
588 u32 b_end = b->base + b->nr_cids;
589 u32 a_wbase = a->base / 64;
590 u32 b_wbase = b->base / 64;
591 u32 lo = a->base > b->base ? a->base : b->base;
592 u32 hi = a_end < b_end ? a_end : b_end;
593 u32 last_wi, start_wi, start_bit, i;
594
595 if (lo >= hi)
596 return a_end;
597 if (start < lo)
598 start = lo;
599 if (start >= hi)
600 return a_end;
601
602 last_wi = (hi - 1) / 64;
603 start_wi = start / 64;
604 start_bit = start & 63;
605
606 bpf_for(i, 0, CMASK_MAX_WORDS) {
607 u32 abs_wi = start_wi + i;
608 u64 word;
609 u32 found;
610
611 if (abs_wi > last_wi)
612 break;
613
614 word = a->bits[abs_wi - a_wbase] & b->bits[abs_wi - b_wbase];
615 if (i == 0)
616 word &= GENMASK_U64(63, start_bit);
617 if (!word)
618 continue;
619
620 found = abs_wi * 64 + ctzll(word);
621 if (found >= hi)
622 return a_end;
623 return found;
624 }
625 return a_end;
626 }
627
628 /*
629 * Find the next set cid in @m at or after @start, wrapping to @m->base if no
630 * set bit is found in [start, m->base + m->nr_cids). Return m->base +
631 * m->nr_cids if @m is empty.
632 *
633 * Callers do round-robin distribution by passing (last_cid + 1) as @start.
634 */
cmask_next_set_wrap(const struct scx_cmask __arena * m,u32 start)635 static __always_inline u32 cmask_next_set_wrap(const struct scx_cmask __arena *m,
636 u32 start)
637 {
638 u32 end = m->base + m->nr_cids;
639 u32 found;
640
641 found = cmask_next_set(m, start);
642 if (found < end || start <= m->base)
643 return found;
644
645 found = cmask_next_set(m, m->base);
646 return found < start ? found : end;
647 }
648
649 /*
650 * Find the next cid set in both @a and @b at or after @start, wrapping to
651 * @a->base if none found in the forward half. Return a->base + a->nr_cids
652 * if the intersection is empty.
653 *
654 * Callers do round-robin distribution by passing (last_cid + 1) as @start.
655 */
cmask_next_and_set_wrap(const struct scx_cmask __arena * a,const struct scx_cmask __arena * b,u32 start)656 static __always_inline u32 cmask_next_and_set_wrap(const struct scx_cmask __arena *a,
657 const struct scx_cmask __arena *b,
658 u32 start)
659 {
660 u32 a_end = a->base + a->nr_cids;
661 u32 found;
662
663 found = cmask_next_and_set(a, b, start);
664 if (found < a_end || start <= a->base)
665 return found;
666
667 found = cmask_next_and_set(a, b, a->base);
668 return found < start ? found : a_end;
669 }
670
671 /*
672 * Like cmask_next_and_set() but over the intersection of THREE masks. Return
673 * a->base + a->nr_cids if no cid is set in all three at or after @start.
674 */
cmask_next_and2_set(const struct scx_cmask __arena * a,const struct scx_cmask __arena * b,const struct scx_cmask __arena * c,u32 start)675 static __always_inline u32 cmask_next_and2_set(const struct scx_cmask __arena *a,
676 const struct scx_cmask __arena *b,
677 const struct scx_cmask __arena *c,
678 u32 start)
679 {
680 u32 a_end = a->base + a->nr_cids;
681 u32 b_end = b->base + b->nr_cids;
682 u32 c_end = c->base + c->nr_cids;
683 u32 a_wbase = a->base / 64;
684 u32 b_wbase = b->base / 64;
685 u32 c_wbase = c->base / 64;
686 u32 lo = a->base > b->base ? a->base : b->base;
687 u32 hi = a_end < b_end ? a_end : b_end;
688 u32 last_wi, start_wi, start_bit, i;
689
690 lo = lo > c->base ? lo : c->base;
691 hi = hi < c_end ? hi : c_end;
692
693 if (lo >= hi)
694 return a_end;
695 if (start < lo)
696 start = lo;
697 if (start >= hi)
698 return a_end;
699
700 last_wi = (hi - 1) / 64;
701 start_wi = start / 64;
702 start_bit = start & 63;
703
704 bpf_for(i, 0, CMASK_MAX_WORDS) {
705 u32 abs_wi = start_wi + i;
706 u64 word;
707 u32 found;
708
709 if (abs_wi > last_wi)
710 break;
711
712 word = a->bits[abs_wi - a_wbase] & b->bits[abs_wi - b_wbase] &
713 c->bits[abs_wi - c_wbase];
714 if (i == 0)
715 word &= GENMASK_U64(63, start_bit);
716 if (!word)
717 continue;
718
719 found = abs_wi * 64 + ctzll(word);
720 if (found >= hi)
721 return a_end;
722 return found;
723 }
724 return a_end;
725 }
726
727 /*
728 * Round-robin variant of cmask_next_and2_set(): wrap to @a->base if the
729 * three-way intersection has no cid in the forward half. Return a->base +
730 * a->nr_cids if empty.
731 */
cmask_next_and2_set_wrap(const struct scx_cmask __arena * a,const struct scx_cmask __arena * b,const struct scx_cmask __arena * c,u32 start)732 static __always_inline u32 cmask_next_and2_set_wrap(const struct scx_cmask __arena *a,
733 const struct scx_cmask __arena *b,
734 const struct scx_cmask __arena *c,
735 u32 start)
736 {
737 u32 a_end = a->base + a->nr_cids;
738 u32 found;
739
740 found = cmask_next_and2_set(a, b, c, start);
741 if (found < a_end || start <= a->base)
742 return found;
743
744 found = cmask_next_and2_set(a, b, c, a->base);
745 return found < start ? found : a_end;
746 }
747
748 /**
749 * cmask_from_cpumask - translate a kernel cpumask to a cid-space cmask
750 * @m: cmask to fill. Zeroed first; only bits within [@m->base, @m->base +
751 * @m->nr_cids) are updated - cpus mapping to cids outside that range
752 * are ignored.
753 * @cpumask: kernel cpumask to translate
754 *
755 * For each cpu in @cpumask, set the cpu's cid in @m. Caller must ensure
756 * @cpumask stays stable across the call (e.g. RCU read lock for
757 * task->cpus_ptr).
758 */
cmask_from_cpumask(struct scx_cmask __arena * m,const struct cpumask * cpumask)759 static __always_inline void cmask_from_cpumask(struct scx_cmask __arena *m,
760 const struct cpumask *cpumask)
761 {
762 u32 nr_cpu_ids = scx_bpf_nr_cpu_ids();
763 s32 cpu;
764
765 cmask_zero(m);
766 bpf_for(cpu, 0, nr_cpu_ids) {
767 s32 cid;
768
769 if (!bpf_cpumask_test_cpu(cpu, cpumask))
770 continue;
771 cid = scx_bpf_cpu_to_cid(cpu);
772 if (cid >= 0)
773 __cmask_set(cid, m);
774 }
775 }
776
777 #endif /* __SCX_CID_BPF_H */
778