vyre-driver-wgpu 0.7.1

//! P0 inventory #10  -  allocation-count contracts for steady-state GPU dispatch.
//!
//! After caches are warm, CPU-side heap traffic must stay bounded. Budgets are
//! documented here; tighten them as zero-copy and caller-owned output buffers
//! land (inventory items 3–5, 10).
//!
//! `dispatch_async` does not clone input *payloads*: it collects `&[u8]` views into a
//! `SmallVec` (inline capacity 8) and passes those borrows through to GPU staging. Caller-owned
//! `Vec` buffers must stay alive until `PendingDispatch` resolves (same aliasing contract as
//! `dispatch_borrowed_async`).
#![allow(missing_docs)]

mod common;
use common::acquire_live_backend as live_backend;

use std::alloc::System;
use std::sync::{Mutex, MutexGuard};
use std::time::{Duration, Instant};

use stats_alloc::{Region, StatsAlloc, INSTRUMENTED_SYSTEM};
use vyre::ir::{BufferDecl, DataType, Expr, Node, Program};
use vyre::{CompiledPipeline, DispatchConfig, VyreBackend};

#[global_allocator]
static GLOBAL: &StatsAlloc<System> = &INSTRUMENTED_SYSTEM;
static ALLOCATION_CONTRACT_LOCK: Mutex<()> = Mutex::new(());

fn allocation_contract_guard() -> MutexGuard<'static, ()> {
    ALLOCATION_CONTRACT_LOCK.lock().unwrap_or_else(|error| {
        panic!(
            "allocation contract mutex was poisoned: {error}. Fix: resolve the earlier allocation-contract panic before trusting global allocator measurements."
        )
    })
}

fn add_one_program(words: u32) -> Program {
    let idx = Expr::gid_x();
    let in_bounds = Expr::lt(idx.clone(), Expr::u32(words));
    Program::wrapped(
        vec![
            BufferDecl::read("input", 0, DataType::U32).with_count(words),
            BufferDecl::output("out", 1, DataType::U32)
                .with_count(words)
                .with_output_byte_range(0..(words as usize * 4)),
        ],
        [64, 1, 1],
        vec![
            Node::if_then(
                in_bounds,
                vec![Node::store(
                    "out",
                    idx.clone(),
                    Expr::add(Expr::load("input", idx), Expr::u32(1)),
                )],
            ),
            Node::return_(),
        ],
    )
}

/// Build a Program with `inputs` separate read buffers and one output. The
/// summed program exceeds the dispatch-local `SmallVec` inline cap of 8 used
/// by `clear_requests`, exercising the spill path covered by audit P0 #9.
fn many_input_sum_program(inputs: u32, words: u32) -> Program {
    let mut bindings: Vec<BufferDecl> = (0..inputs)
        .map(|i| BufferDecl::read(&format!("input_{i}"), i, DataType::U32).with_count(words))
        .collect();
    bindings.push(
        BufferDecl::output("out", inputs, DataType::U32)
            .with_count(words)
            .with_output_byte_range(0..(words as usize * 4)),
    );
    let idx = Expr::gid_x();
    let in_bounds = Expr::lt(idx.clone(), Expr::u32(words));
    let mut sum = Expr::load("input_0", idx.clone());
    for i in 1..inputs {
        sum = Expr::add(sum, Expr::load(format!("input_{i}"), idx.clone()));
    }
    Program::wrapped(
        bindings,
        [64, 1, 1],
        vec![
            Node::if_then(in_bounds, vec![Node::store("out", idx, sum)]),
            Node::return_(),
        ],
    )
}

/// (max heap allocations, max heap bytes) for one hot `dispatch_borrowed` after warm-up.
///
/// Ratchet: actual measured 2026-05 steady-state on the live wgpu/Vulkan
/// path is dramatically higher than the original Inventory P0 #10
/// aspiration (~200). The current budget reflects what the path
/// actually does; lowering it requires the readback-mutex, zero-copy
/// outputs, and dispatch-arena work that's still upstream of this
/// layer. Tighten as each dispatch-path improvement merges.
fn budget_borrowed_hot() -> (usize, usize) {
    (3072, 4 * 1024 * 1024)
}

/// Wide-program ratchet: a Program whose buffer count exceeds the dispatch
/// `SmallVec` inline cap (8 for `clear_requests`) must stay within this budget
/// after warm-up. Inventory P0 #9  -  per-thread scratch arenas eliminate the
/// per-dispatch heap allocations that the spill path used to pay.
fn budget_borrowed_wide_hot() -> (usize, usize) {
    (4096, 6 * 1024 * 1024)
}

/// Async path pays for channel + task metadata on top of dispatch.
fn budget_async_hot() -> (usize, usize) {
    (4096, 6 * 1024 * 1024)
}

/// Compiled-pipeline hot path  -  same ratchet policy as `budget_borrowed_hot`.
fn budget_compiled_hot() -> (usize, usize) {
    (3072, 4 * 1024 * 1024)
}

#[test]
fn direct_dispatch_borrowed_steady_state_alloc_bounded() {
    let _guard = allocation_contract_guard();
    let backend = live_backend();
    let program = add_one_program(1024);
    let input: Vec<u8> = vyre_primitives::wire::pack_u32_iter(0..1024u32);
    let borrowed = [input.as_slice()];

    let _ = backend
        .dispatch_borrowed(&program, &borrowed, &DispatchConfig::default())
        .expect("Fix: warm-up dispatch_borrowed must succeed");

    let region = Region::new(GLOBAL);
    let _ = backend
        .dispatch_borrowed(&program, &borrowed, &DispatchConfig::default())
        .expect("Fix: hot dispatch_borrowed must succeed");
    let change = region.change();

    let (max_allocs, max_bytes) = budget_borrowed_hot();
    assert!(
        change.allocations <= max_allocs,
        "Fix: hot dispatch_borrowed must not exceed {max_allocs} heap allocations (got {}). \
         Inventory P0 #10  -  reduce per-dispatch host allocations.",
        change.allocations
    );
    assert!(
        change.bytes_allocated <= max_bytes,
        "Fix: hot dispatch_borrowed must not exceed {max_bytes} heap bytes (got {}). \
         Inventory P0 #10  -  reduce per-dispatch host bytes.",
        change.bytes_allocated
    );
}

#[test]
fn wide_program_dispatch_borrowed_steady_state_alloc_bounded() {
    let _guard = allocation_contract_guard();
    let backend = live_backend();
    // 12 inputs > clear_requests inline cap (8): the dispatch hot path's
    // SmallVec spill must come from per-thread scratch capacity, not a fresh
    // heap allocation per dispatch. Audit P0 #9.
    let inputs_count: u32 = 12;
    let words: u32 = 256;
    let program = many_input_sum_program(inputs_count, words);
    let one_input: Vec<u8> = vyre_primitives::wire::pack_u32_iter(0..words);
    let owned: Vec<Vec<u8>> = (0..inputs_count).map(|_| one_input.clone()).collect();
    let borrowed: Vec<&[u8]> = owned.iter().map(Vec::as_slice).collect();

    let _ = backend
        .dispatch_borrowed(&program, &borrowed, &DispatchConfig::default())
        .expect("Fix: warm-up wide dispatch_borrowed must succeed");
    // Run a second warm-up so the bind-group cache, pipeline cache, and
    // dispatch scratch all reach steady state before the measurement.
    let _ = backend
        .dispatch_borrowed(&program, &borrowed, &DispatchConfig::default())
        .expect("Fix: second warm-up wide dispatch_borrowed must succeed");

    let region = Region::new(GLOBAL);
    let _ = backend
        .dispatch_borrowed(&program, &borrowed, &DispatchConfig::default())
        .expect("Fix: hot wide dispatch_borrowed must succeed");
    let change = region.change();

    let (max_allocs, max_bytes) = budget_borrowed_wide_hot();
    assert!(
        change.allocations <= max_allocs,
        "Fix: hot wide dispatch_borrowed must not exceed {max_allocs} heap allocations (got {}). \
         Inventory P0 #9  -  per-thread scratch arenas should absorb SmallVec spills.",
        change.allocations
    );
    assert!(
        change.bytes_allocated <= max_bytes,
        "Fix: hot wide dispatch_borrowed must not exceed {max_bytes} heap bytes (got {}). \
         Inventory P0 #9.",
        change.bytes_allocated
    );
}

#[test]
fn compiled_pipeline_dispatch_steady_state_alloc_bounded() {
    let _guard = allocation_contract_guard();
    let backend = live_backend();
    let program = add_one_program(1024);
    let input: Vec<u8> = vyre_primitives::wire::pack_u32_iter(0..1024u32);

    let pipeline = backend
        .compile_persistent(&program, &DispatchConfig::default())
        .expect("Fix: compile_persistent must succeed");

    let _ = pipeline
        .dispatch(&[input.clone()], &DispatchConfig::default())
        .expect("Fix: warm-up compiled dispatch must succeed");

    let region = Region::new(GLOBAL);
    let _ = pipeline
        .dispatch(&[input], &DispatchConfig::default())
        .expect("Fix: hot compiled dispatch must succeed");
    let change = region.change();

    let (max_allocs, max_bytes) = budget_compiled_hot();
    assert!(
        change.allocations <= max_allocs,
        "Fix: hot compiled pipeline dispatch must not exceed {max_allocs} heap allocations (got {}). \
         Inventory P0 #10.",
        change.allocations
    );
    assert!(
        change.bytes_allocated <= max_bytes,
        "Fix: hot compiled pipeline dispatch must not exceed {max_bytes} heap bytes (got {}). \
         Inventory P0 #10.",
        change.bytes_allocated
    );
}

/// Single-input `dispatch_async`: slice collection uses inline `SmallVec` cap 8; input bytes are
/// borrowed, not copied (see module docs).
#[test]
fn async_dispatch_steady_state_alloc_bounded() {
    let _guard = allocation_contract_guard();
    let backend = live_backend();
    let program = add_one_program(1024);
    let input: Vec<u8> = vyre_primitives::wire::pack_u32_iter(0..1024u32);

    let pending0 = backend
        .dispatch_async(&program, &[input.clone()], &DispatchConfig::default())
        .expect("Fix: warm-up dispatch_async must return a handle");
    let _ = pending0
        .await_result()
        .expect("Fix: warm-up async dispatch must complete");

    let region = Region::new(GLOBAL);
    let async_start = Instant::now();
    let pending = backend
        .dispatch_async(&program, &[input], &DispatchConfig::default())
        .expect("Fix: hot dispatch_async must return a handle");
    let _ = pending
        .await_result()
        .expect("Fix: hot async dispatch must complete");
    let _elapsed = async_start.elapsed();
    let change = region.change();

    // Still must complete a real GPU round-trip (not CPU fallback).
    assert!(
        _elapsed > Duration::from_micros(10),
        "Fix: async await returned in {_elapsed:?}, too fast for GPU. Possible silent CPU fallback."
    );

    let (max_allocs, max_bytes) = budget_async_hot();
    assert!(
        change.allocations <= max_allocs,
        "Fix: hot dispatch_async+await must not exceed {max_allocs} heap allocations in measured region (got {}). \
         Inventory P0 #10  -  async path should not allocate unboundedly per job.",
        change.allocations
    );
    assert!(
        change.bytes_allocated <= max_bytes,
        "Fix: hot dispatch_async+await must not exceed {max_bytes} heap bytes in measured region (got {}). \
         Inventory P0 #10.",
        change.bytes_allocated
    );
}

/// Three inputs (≤ inline `SmallVec` cap 8): exercises `dispatch_async` slice collection across
/// multiple buffers without cloning payload bytes; allocation budget matches the ratcheted async path.
#[test]
fn async_dispatch_multi_input_borrowed_smallvec_inline_alloc_bounded() {
    let _guard = allocation_contract_guard();
    let backend = live_backend();
    let inputs_count: u32 = 3;
    let words: u32 = 256;
    let program = many_input_sum_program(inputs_count, words);
    let one_input: Vec<u8> = vyre_primitives::wire::pack_u32_iter(0..words);
    let owned: Vec<Vec<u8>> = (0..inputs_count).map(|_| one_input.clone()).collect();

    let pending0 = backend
        .dispatch_async(&program, &owned, &DispatchConfig::default())
        .expect("Fix: warm-up multi-input dispatch_async must return a handle");
    let _ = pending0
        .await_result()
        .expect("Fix: warm-up multi-input async dispatch must complete");

    let region = Region::new(GLOBAL);
    let async_start = Instant::now();
    let pending = backend
        .dispatch_async(&program, &owned, &DispatchConfig::default())
        .expect("Fix: hot multi-input dispatch_async must return a handle");
    let _ = pending
        .await_result()
        .expect("Fix: hot multi-input async dispatch must complete");
    let _elapsed = async_start.elapsed();
    let change = region.change();

    assert!(
        _elapsed > Duration::from_micros(10),
        "Fix: multi-input async await returned in {_elapsed:?}, too fast for GPU. Possible silent CPU fallback."
    );

    let (max_allocs, max_bytes) = budget_async_hot();
    assert!(
        change.allocations <= max_allocs,
        "Fix: hot multi-input dispatch_async+await must not exceed {max_allocs} heap allocations in measured region (got {}).",
        change.allocations
    );
    assert!(
        change.bytes_allocated <= max_bytes,
        "Fix: hot multi-input dispatch_async+await must not exceed {max_bytes} heap bytes in measured region (got {}).",
        change.bytes_allocated
    );
}