vyre_runtime/megakernel/builder/
priority.rs

1use super::{
2    claimed_slot_body, execute_already_claimed_slot_body, persistent_lane_prologue,
3    process_io_requests, wrap_persistent_megakernel_program,
4};
5use super::{Expr, Node, OpcodeHandler, Program};
6
7// ---- Priority-aware variant ----
8
9/// Build a priority-aware megakernel IR.
10///
11/// Unlike `build_program_sharded` where each lane owns exactly one slot,
12/// the priority variant has workers scan across priority-partitioned ring
13/// regions, claiming the highest-priority PUBLISHED slot available. This
14/// ensures latency-sensitive work (CRITICAL, HIGH) is processed before
15/// background tasks (LOW, IDLE).
16///
17/// The control buffer is extended with `PRIORITY_OFFSETS_BASE..+6` words
18/// that the host sets to define partition boundaries. The host can
19/// dynamically resize partitions by updating these offsets between batches.
20#[must_use]
21pub fn build_program_priority(workgroup_size_x: u32, opcodes: &[OpcodeHandler]) -> Program {
22    build_program_priority_slots(workgroup_size_x, workgroup_size_x.max(1), opcodes)
23}
24
25/// Build a priority-aware megakernel IR for an explicit ring slot count.
26#[must_use]
27pub fn build_program_priority_slots(
28    workgroup_size_x: u32,
29    slot_count: u32,
30    opcodes: &[OpcodeHandler],
31) -> Program {
32    wrap_persistent_megakernel_program(
33        workgroup_size_x,
34        slot_count.max(1),
35        persistent_body_priority_slots(workgroup_size_x, slot_count.max(1), opcodes),
36    )
37}
38
39/// Priority-aware loop body. Replaces the per-lane 1:1 slot mapping
40/// with the scheduler's priority scan.
41#[must_use]
42pub fn persistent_body_priority(workgroup_size_x: u32, opcodes: &[OpcodeHandler]) -> Vec<Node> {
43    persistent_body_priority_slots(workgroup_size_x, workgroup_size_x.max(1), opcodes)
44}
45
46/// Priority-aware loop body for an explicit ring slot count.
47#[must_use]
48pub fn persistent_body_priority_slots(
49    workgroup_size_x: u32,
50    slot_count: u32,
51    opcodes: &[OpcodeHandler],
52) -> Vec<Node> {
53    use crate::megakernel::scheduler;
54
55    let slot_count = slot_count.max(1);
56    let mut body = persistent_lane_prologue(workgroup_size_x);
57
58    // -- Priority scan: find and claim the best available slot. --------
59    body.extend(scheduler::priority_scan_body(slot_count));
60
61    // -- If claimed, execute the slot. ---------------------------------
62    body.push(Node::if_then(
63        Expr::ne(Expr::var("claimed_slot_base"), Expr::u32(u32::MAX)),
64        {
65            // Rebind `slot_base` to the claimed slot so downstream
66            // handler code works unchanged.
67            let mut exec = vec![Node::let_bind("slot_base", Expr::var("claimed_slot_base"))];
68            exec.extend(execute_already_claimed_slot_body(
69                Expr::var("claimed_tenant"),
70                claimed_slot_body(opcodes),
71            ));
72            exec
73        },
74    ));
75
76    // -- IO poll (same as base variant). --------------------------------
77    body.push(Node::Block(process_io_requests()));
78
79    body
80}
vyre_runtime/megakernel/builder/priority.rs

vyre_runtime/megakernel/builder/
priority.rs
