vyre_runtime/megakernel/scheduler.rs
1//! Work scheduler - priority-aware slot scanning for the persistent megakernel.
2//!
3//! Extends the base slot-claim logic with priority partitioning:
4//! each priority level occupies a contiguous partition of the ring buffer.
5//! Workers scan from highest priority (0=CRITICAL) to lowest (4=IDLE),
6//! claiming the first PUBLISHED slot found. This ensures latency-sensitive
7//! work is processed before background tasks without true preemption.
8//!
9//! ## Slot Layout Extension
10//!
11//! The priority is encoded in `ring_buffer[slot_base + PRIORITY_WORD]`.
12//! The host sets this when publishing; the scheduler reads it to
13//! sort work into the right scan order.
14//!
15//! ## Starvation Guard
16//!
17//! After `STARVATION_THRESHOLD` consecutive high-priority claims, the
18//! scheduler forcibly scans lower-priority partitions for one iteration.
19//! This prevents priority inversion where a flood of CRITICAL slots
20//! starves NORMAL/"background" work indefinitely.
21
22use super::ir_util::{atomic_load_relaxed, atomic_store_relaxed};
23use super::protocol::*;
24use vyre_foundation::ir::{Expr, Node};
25
26mod offsets;
27pub use offsets::{default_priority_offsets_array, write_default_priority_offsets};
28
29/// Number of priority levels the scheduler supports.
30pub const PRIORITY_LEVELS: u32 = 5;
31
32/// Priority discriminants.
33pub mod priority {
34 /// Highest priority - interactive/latency-critical work.
35 pub const CRITICAL: u32 = 0;
36 /// High priority - important but not latency-critical.
37 pub const HIGH: u32 = 1;
38 /// Normal priority - the default for all work.
39 pub const NORMAL: u32 = 2;
40 /// Low priority - background, non-urgent work.
41 pub const LOW: u32 = 3;
42 /// Idle priority - processed only when no other work exists.
43 pub const IDLE: u32 = 4;
44}
45
46/// After this many consecutive claims at the same (or higher) priority,
47/// the scheduler forcibly scans lower-priority partitions for one iteration.
48pub const STARVATION_THRESHOLD: u32 = 16;
49
50/// After this many claims by a single tenant in a single worker's "epoch",
51/// the tenant is considered "greedy" and may be throttled.
52pub const TENANT_FAIRNESS_THRESHOLD: u32 = 64;
53
54/// Control word storing the priority partition offsets.
55/// `control[PRIORITY_OFFSETS_BASE + pri]` = first slot index for priority `pri`.
56/// `control[PRIORITY_OFFSETS_BASE + PRIORITY_LEVELS]` = total slot count (sentinel).
57pub const PRIORITY_OFFSETS_BASE: u32 = control::PRIORITY_OFFSETS_BASE;
58
59/// Control word storing consecutive high-priority claims.
60pub const PRIORITY_STARVATION_COUNTER: u32 = control::PRIORITY_STARVATION_COUNTER;
61
62/// Policy helper: select the next slot to probe within a partition.
63///
64/// Offsetting the start by `lane_id` reduces CAS contention on the first
65/// few slots of a partition when many workers wake up simultaneously.
66#[must_use]
67pub fn policy_offset_start(partition_start: Expr, partition_end: Expr, lane_id: Expr) -> Expr {
68 let range = Expr::sub(partition_end.clone(), partition_start.clone());
69 let nonzero_range = Expr::max(range, Expr::u32(1));
70 Expr::add(partition_start, Expr::rem(lane_id, nonzero_range))
71}
72
73/// Number of strided probes each lane needs to cover a priority partition.
74///
75/// The scheduler has `worker_width` lanes scanning one partition in lockstep.
76/// Bounding this as a ceiling division keeps the generated scan work linear in
77/// slot count instead of priority_levels * total_slots.
78#[must_use]
79pub fn priority_partition_probe_count(partition_slots: u32, worker_width: u32) -> u32 {
80 if partition_slots == 0 {
81 return 0;
82 }
83 let width = worker_width.max(1);
84 partition_slots.div_ceil(width)
85}
86
87/// Number of lanes that should actively probe one priority partition.
88///
89/// Lanes outside `partition_slots` cannot discover additional work when the
90/// worker set is wider than the partition; masking them avoids duplicate slot
91/// probes across every priority band.
92#[must_use]
93pub fn priority_partition_active_lane_count(partition_slots: u32, worker_width: u32) -> u32 {
94 partition_slots.min(worker_width.max(1))
95}
96
97/// Upper bound on slot status probes for one priority partition.
98#[must_use]
99pub fn priority_partition_probe_budget(partition_slots: u32, worker_width: u32) -> u32 {
100 priority_partition_active_lane_count(partition_slots, worker_width).saturating_mul(
101 priority_partition_probe_count(partition_slots, worker_width),
102 )
103}
104
105/// Policy helper: check if a tenant has exceeded its fairness quota.
106#[must_use]
107pub fn check_tenant_fairness(tenant_id: Expr) -> Expr {
108 let tenant_counter = Expr::rem(tenant_id, Expr::u32(control::TENANT_FAIRNESS_SLOTS));
109 let count = atomic_load_relaxed(
110 "control",
111 Expr::add(Expr::u32(control::TENANT_FAIRNESS_BASE), tenant_counter),
112 );
113 Expr::lt(count, Expr::u32(TENANT_FAIRNESS_THRESHOLD))
114}
115
116/// Policy helper: check if a priority level has exceeded its fairness quota.
117#[must_use]
118pub fn check_priority_fairness(priority: Expr) -> Expr {
119 let count = atomic_load_relaxed(
120 "control",
121 Expr::add(Expr::u32(control::PRIORITY_FAIRNESS_BASE), priority),
122 );
123 Expr::lt(count, Expr::u32(STARVATION_THRESHOLD))
124}
125
126/// Build the priority-aware scan loop as `Vec<Node>` for composition.
127///
128/// The scan checks priorities from `start_priority` to `PRIORITY_LEVELS - 1`.
129/// For each priority level, it scans the corresponding ring partition
130/// for a PUBLISHED slot. If found, claims it via CAS and yields
131/// the slot base to the caller.
132///
133/// Variables set on success:
134/// - `claimed_slot_base`: the slot_base of the claimed slot (u32::MAX if none found)
135/// - `claimed_priority`: the priority level of the claimed slot
136/// - `claimed_tenant`: the tenant id of the claimed slot
137///
138/// Requires `lane_id` and `workgroup_size_x` in scope.
139#[must_use]
140pub fn priority_scan_body(total_slots: u32) -> Vec<Node> {
141 priority_scan_body_with_stride(total_slots, total_slots.max(1))
142}
143
144/// Build the priority-aware scan loop with an explicit global worker stride.
145///
146/// Each lane probes its own congruence class inside each priority partition.
147/// Across all launched workers this changes the scan from every worker probing
148/// every slot to the worker set covering the partition once per priority pass.
149#[must_use]
150pub fn priority_scan_body_with_stride(total_slots: u32, worker_stride: u32) -> Vec<Node> {
151 let worker_stride = worker_stride.max(1);
152 vec![
153 // Initialize output: no slot claimed
154 Node::let_bind("claimed_slot_base", Expr::u32(u32::MAX)),
155 Node::let_bind("claimed_priority", Expr::u32(u32::MAX)),
156 Node::let_bind("claimed_tenant", Expr::u32(u32::MAX)),
157 Node::let_bind(
158 "priority_starvation_count",
159 atomic_load_relaxed("control", Expr::u32(PRIORITY_STARVATION_COUNTER)),
160 ),
161 Node::let_bind(
162 "priority_force_lower",
163 Expr::ge(
164 Expr::var("priority_starvation_count"),
165 Expr::u32(STARVATION_THRESHOLD),
166 ),
167 ),
168 // Scan each priority level in order
169 Node::loop_for(
170 "scan_pri",
171 Expr::u32(0),
172 Expr::u32(PRIORITY_LEVELS),
173 vec![
174 // Skip if we already claimed a slot
175 Node::if_then(
176 Expr::and(
177 Expr::eq(Expr::var("claimed_slot_base"), Expr::u32(u32::MAX)),
178 Expr::or(
179 Expr::not(Expr::var("priority_force_lower")),
180 Expr::gt(Expr::var("scan_pri"), Expr::u32(priority::HIGH)),
181 ),
182 ),
183 vec![
184 // Load partition boundaries from control buffer
185 Node::let_bind(
186 "part_start",
187 atomic_load_relaxed(
188 "control",
189 Expr::add(Expr::u32(PRIORITY_OFFSETS_BASE), Expr::var("scan_pri")),
190 ),
191 ),
192 Node::let_bind(
193 "part_end",
194 atomic_load_relaxed(
195 "control",
196 Expr::add(
197 Expr::u32(PRIORITY_OFFSETS_BASE),
198 Expr::add(Expr::var("scan_pri"), Expr::u32(1)),
199 ),
200 ),
201 ),
202 Node::let_bind(
203 "part_len",
204 Expr::sub(Expr::var("part_end"), Expr::var("part_start")),
205 ),
206 Node::let_bind(
207 "probe_count",
208 Expr::div(
209 Expr::add(
210 Expr::var("part_len"),
211 Expr::u32(worker_stride.saturating_sub(1)),
212 ),
213 Expr::u32(worker_stride),
214 ),
215 ),
216 // Scan slots within this priority partition
217 Node::if_then(
218 Expr::gt(Expr::var("part_len"), Expr::u32(0)),
219 vec![
220 Node::let_bind(
221 "partition_lane",
222 Expr::rem(Expr::var("lane_id"), Expr::u32(worker_stride)),
223 ),
224 Node::if_then(
225 Expr::lt(Expr::var("partition_lane"), Expr::var("part_len")),
226 vec![Node::loop_for(
227 "scan_idx",
228 Expr::u32(0),
229 Expr::var("probe_count"),
230 vec![
231 Node::let_bind(
232 "scan_slot",
233 Expr::add(
234 Expr::var("part_start"),
235 Expr::rem(
236 Expr::add(
237 Expr::var("partition_lane"),
238 Expr::mul(
239 Expr::var("scan_idx"),
240 Expr::u32(worker_stride),
241 ),
242 ),
243 Expr::var("part_len"),
244 ),
245 ),
246 ),
247 Node::if_then(
248 Expr::and(
249 Expr::eq(
250 Expr::var("claimed_slot_base"),
251 Expr::u32(u32::MAX),
252 ),
253 Expr::lt(
254 Expr::var("scan_slot"),
255 Expr::u32(total_slots),
256 ),
257 ),
258 vec![
259 Node::let_bind(
260 "probe_base",
261 Expr::mul(
262 Expr::var("scan_slot"),
263 Expr::u32(SLOT_WORDS),
264 ),
265 ),
266 Node::let_bind(
267 "probe_status",
268 atomic_load_relaxed(
269 "ring_buffer",
270 Expr::var("probe_base"),
271 ),
272 ),
273 Node::let_bind(
274 "probe_schedulable",
275 Expr::or(
276 Expr::eq(
277 Expr::var("probe_status"),
278 Expr::u32(slot::PUBLISHED),
279 ),
280 Expr::or(
281 Expr::eq(
282 Expr::var("probe_status"),
283 Expr::u32(slot::YIELD),
284 ),
285 Expr::eq(
286 Expr::var("probe_status"),
287 Expr::u32(slot::REQUEUE),
288 ),
289 ),
290 ),
291 ),
292 Node::if_then(
293 Expr::var("probe_schedulable"),
294 vec![
295 Node::let_bind(
296 "probe_tenant",
297 Expr::load(
298 "ring_buffer",
299 Expr::add(
300 Expr::var("probe_base"),
301 Expr::u32(TENANT_WORD),
302 ),
303 ),
304 ),
305 Node::let_bind(
306 "probe_tenant_base",
307 atomic_load_relaxed(
308 "control",
309 Expr::u32(
310 control::TENANT_BASE,
311 ),
312 ),
313 ),
314 Node::let_bind(
315 "probe_tenant_mask",
316 atomic_load_relaxed(
317 "control",
318 Expr::add(
319 Expr::var(
320 "probe_tenant_base",
321 ),
322 Expr::var("probe_tenant"),
323 ),
324 ),
325 ),
326 Node::if_then(
327 Expr::ne(
328 Expr::var(
329 "probe_tenant_mask",
330 ),
331 Expr::u32(0),
332 ),
333 vec![
334 Node::let_bind(
335 "probe_expected",
336 Expr::var("probe_status"),
337 ),
338 Node::let_bind(
339 "probe_prev",
340 Expr::atomic_compare_exchange(
341 "ring_buffer",
342 Expr::var("probe_base"),
343 Expr::var("probe_expected"),
344 Expr::u32(slot::CLAIMED),
345 ),
346 ),
347 Node::if_then(
348 Expr::eq(
349 Expr::var("probe_prev"),
350 Expr::var("probe_expected"),
351 ),
352 vec![
353 Node::assign(
354 "claimed_slot_base",
355 Expr::var("probe_base"),
356 ),
357 Node::assign(
358 "claimed_priority",
359 Expr::var("scan_pri"),
360 ),
361 Node::assign(
362 "claimed_tenant",
363 Expr::var("probe_tenant"),
364 ),
365 ],
366 ),
367 ],
368 ),
369 ],
370 ),
371 ],
372 ),
373 ],
374 )],
375 ),
376 ],
377 ),
378 ],
379 ),
380 ],
381 ),
382 // Post-claim: Update fairness accounting
383 Node::if_then(
384 Expr::ne(Expr::var("claimed_priority"), Expr::u32(u32::MAX)),
385 vec![
386 // Update priority starvation counter atomically
387 Node::if_then_else(
388 Expr::le(Expr::var("claimed_priority"), Expr::u32(priority::HIGH)),
389 vec![Node::let_bind(
390 "priority_starvation_prev",
391 Expr::atomic_add(
392 "control",
393 Expr::u32(PRIORITY_STARVATION_COUNTER),
394 Expr::u32(1),
395 ),
396 )],
397 vec![atomic_store_relaxed(
398 "priority_starvation_prev",
399 "control",
400 Expr::u32(PRIORITY_STARVATION_COUNTER),
401 Expr::u32(0),
402 )],
403 ),
404 // Update per-tenant fairness counter
405 Node::let_bind(
406 "tenant_fairness_prev",
407 Expr::atomic_add(
408 "control",
409 Expr::add(
410 Expr::u32(control::TENANT_FAIRNESS_BASE),
411 Expr::rem(
412 Expr::var("claimed_tenant"),
413 Expr::u32(control::TENANT_FAIRNESS_SLOTS),
414 ),
415 ),
416 Expr::u32(1),
417 ),
418 ),
419 // Update per-priority fairness counter (telemetry)
420 Node::let_bind(
421 "priority_fairness_prev",
422 Expr::atomic_add(
423 "control",
424 Expr::add(
425 Expr::u32(control::PRIORITY_FAIRNESS_BASE),
426 Expr::var("claimed_priority"),
427 ),
428 Expr::u32(1),
429 ),
430 ),
431 ],
432 ),
433 ]
434}
435
436#[cfg(test)]
437mod tests;