vyre_libs/compiler/
stack_layout.rs

1use crate::compiler::atomic_collect::atomic_collect_u32;
2use vyre::ir::{Expr, Program};
3
4const OP_ID: &str = "vyre-libs::parsing::opt_stack_layout_generation";
5
6/// GPU SIMT Stack Layout Generator (Prologue/Epilogue Spiller)
7///
8/// When variables cannot fit into the physical 16 registers of x86_64, they "spill".
9/// This module generates precise Main Memory stack offsets relative to `%rbp` (base pointer)
10/// and emits `push`/`pop` sequences using SIMT inclusive scan prefix offsets.
11#[must_use]
12pub fn opt_stack_layout_generation(
13    physical_registers: &str,
14    out_spill_offsets: &str,
15    num_regs: Expr,
16) -> Program {
17    atomic_collect_u32(
18        OP_ID,
19        physical_registers,
20        out_spill_offsets,
21        "tmp_stack_frame_size",
22        num_regs,
23        8,
24        None,
25        |reg_bound, _t| Expr::ge(reg_bound, Expr::u32(16)),
26        |t, _stack_offset| t,
27        |_t, stack_offset| stack_offset,
28    )
29}
30
31inventory::submit! {
32    crate::harness::OpEntry {
33        id: OP_ID,
34        build: || opt_stack_layout_generation("regs", "spills", Expr::u32(4)),
35        // 4 virtual registers: two sit in the 0..=15 physical window,
36        // two spill (reg_bound = 20, 30). Each spill claims an 8-byte
37        // stack slot via atomic_add on a workgroup-scoped counter.
38        // The reference interpreter runs one lane at a time in
39        // monotonic invocation order, so spill_offsets[1] = 0 and
40        // spill_offsets[3] = 8.
41        test_inputs: Some(|| {
42            let regs: [u32; 4] = [3, 20, 7, 30];
43            let bytes = vyre_primitives::wire::pack_u32_slice(&regs);
44            // regs (ReadOnly), out_spill_offsets (ReadWrite),
45            // tmp_stack_frame_size (ReadWrite, 1 slot).
46            vec![vec![bytes, vec![0u8; 4 * 4], vec![0u8; 4]]]
47        }),
48        expected_output: Some(|| {
49            let spills: [u32; 4] = [0, 0, 0, 8];
50            let frame_size: [u32; 1] = [16];
51            let to_bytes = vyre_primitives::wire::pack_u32_slice;
52            vec![vec![to_bytes(&spills), to_bytes(&frame_size)]]
53        }),
54        category: Some("compiler"),
55    }
56}
vyre_libs/compiler/stack_layout.rs

vyre_libs/compiler/
stack_layout.rs
