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