vyre_runtime/megakernel/builder/
jit.rs

1use super::persistent_lane_prologue;
2use super::{
3    claimed_slot_bindings, direct_slot_base_binding, process_io_requests, slot_tenant_id_load,
4    tenant_authorized_claim_body, wrap_persistent_megakernel_program,
5};
6use crate::megakernel::ir_util::atomic_load_relaxed;
7use crate::megakernel::protocol::{control, slot, STATUS_WORD};
8use vyre_foundation::ir::{Expr, Node, Program};
9
10/// Build the JIT Megakernel IR where payload processor logic is fused into the body stream.
11#[must_use]
12pub fn build_program_jit(workgroup_size_x: u32, payload_processor: &[Node]) -> Program {
13    build_program_jit_slots(workgroup_size_x, workgroup_size_x.max(1), payload_processor)
14}
15
16/// Build the JIT megakernel IR for an explicit number of ring slots.
17#[must_use]
18pub fn build_program_jit_slots(
19    workgroup_size_x: u32,
20    slot_count: u32,
21    payload_processor: &[Node],
22) -> Program {
23    wrap_persistent_megakernel_program(
24        workgroup_size_x,
25        slot_count,
26        persistent_body_jit(workgroup_size_x, payload_processor),
27    )
28}
29
30fn execute_slot_body_jit(payload_processor: &[Node]) -> Vec<Node> {
31    vec![
32        Node::let_bind(
33            "status_index",
34            Expr::add(Expr::var("slot_base"), Expr::u32(STATUS_WORD)),
35        ),
36        Node::let_bind(
37            "observed_status",
38            atomic_load_relaxed("ring_buffer", Expr::var("status_index")),
39        ),
40        Node::if_then(
41            Expr::eq(Expr::var("observed_status"), Expr::u32(slot::PUBLISHED)),
42            tenant_authorized_claim_body(
43                slot_tenant_id_load(),
44                claimed_slot_body_jit(payload_processor),
45            ),
46        ),
47    ]
48}
49
50// ---- JIT variant ----
51
52/// The JIT body that runs once per iteration per lane.
53#[must_use]
54pub fn persistent_body_jit(workgroup_size_x: u32, payload_processor: &[Node]) -> Vec<Node> {
55    match try_persistent_body_jit(workgroup_size_x, payload_processor) {
56        Ok(body) => body,
57        Err(error) => panic!("{error}"),
58    }
59}
60
61/// Fallible JIT body builder with explicit staging-allocation reporting.
62pub(super) fn try_persistent_body_jit(
63    workgroup_size_x: u32,
64    payload_processor: &[Node],
65) -> Result<Vec<Node>, String> {
66    let mut body = persistent_lane_prologue(workgroup_size_x);
67    let body_capacity = body.len().checked_add(3).ok_or_else(|| {
68        "megakernel JIT body node reservation overflowed usize. Fix: reduce fused payload/body staging before building the JIT megakernel."
69            .to_string()
70    })?;
71    vyre_foundation::allocation::try_reserve_vec_to_capacity(&mut body, body_capacity).map_err(|error| {
72        format!(
73            "megakernel JIT body node reservation failed: {error}. Fix: reduce fused payload/body staging before building the JIT megakernel."
74        )
75    })?;
76    body.push(direct_slot_base_binding());
77    body.push(Node::Block(execute_slot_body_jit(payload_processor)));
78    body.push(Node::Block(process_io_requests()));
79    Ok(body)
80}
81
82fn claimed_slot_body_jit(payload_processor: &[Node]) -> Vec<Node> {
83    let mut nodes = claimed_slot_bindings();
84
85    // Wire the statically JIT-compiled rule/payload evaluation graph.
86    nodes.extend(payload_processor.iter().cloned());
87
88    nodes.push(Node::let_bind(
89        "done_prev",
90        Expr::atomic_add("control", Expr::u32(control::DONE_COUNT), Expr::u32(1)),
91    ));
92    nodes.push(Node::store(
93        "ring_buffer",
94        Expr::var("status_index"),
95        Expr::u32(slot::DONE),
96    ));
97    nodes
98}
vyre_runtime/megakernel/builder/jit.rs

vyre_runtime/megakernel/builder/
jit.rs
