weirflow 0.1.0

//! Parity test: GPU live-variables closure matches reference oracle.

#![cfg(test)]

use vyre::DispatchConfig;
use vyre::VyreBackend;
use vyre_driver_cuda::{CudaBackend, CudaBackendRegistration};
use vyre_primitives::predicate::edge_kind;
use weir::fixed_point_resident::{
    FixedPointResidentFrontierScratch, FixedPointResidentGraph, FixedPointResidentPlan,
};
use weir::live::{live_closure_borrowed_via, prepare_live_plan};
use weir::oracle::graph::live_closure as live_closure_cpu;

fn run_dispatch(
    backend: &CudaBackend,
) -> impl Fn(&vyre::ir::Program, &[&[u8]], Option<[u32; 3]>) -> Result<Vec<Vec<u8>>, String> + '_ {
    move |program, inputs, grid_override| {
        let mut config = DispatchConfig::default();
        config.grid_override = grid_override;
        backend
            .dispatch_borrowed(program, inputs, &config)
            .map_err(|err| err.to_string())
    }
}

fn live_dispatcher() -> CudaBackend {
    CudaBackend::acquire()
        .expect("CudaBackend::acquire failed on a host that must have an NVIDIA GPU.")
}

#[test]
fn cuda_live_chain_propagates_back() {
    let backend = live_dispatcher();
    let dispatch = run_dispatch(&backend);
    // Forward CFG: 0 -> 1 -> 2 -> 3, all CONTROL edges.
    let edge_offsets = vec![0u32, 1, 2, 3, 3];
    let edge_targets = vec![1u32, 2, 3];
    let edge_kind_mask = vec![edge_kind::CONTROL; 3];
    // Live at exit (node 3). Backward should make all nodes live.
    let seed_bits = vec![0b1000u32];
    let cpu = live_closure_cpu(4, &edge_offsets, &edge_targets, &edge_kind_mask, &seed_bits);
    let gpu = live_closure_borrowed_via(
        &dispatch,
        4,
        &edge_offsets,
        &edge_targets,
        &edge_kind_mask,
        &seed_bits,
        16,
    )
    .expect("dispatch");
    assert_eq!(gpu, cpu, "chain backward closure divergence");
    assert_eq!(gpu, vec![0b1111u32]);
}

#[test]
fn cuda_live_diamond() {
    let backend = live_dispatcher();
    let dispatch = run_dispatch(&backend);
    // Diamond 0 -> {1, 2} -> 3, CONTROL edges. Live at 3 → all nodes live.
    let edge_offsets = vec![0u32, 2, 3, 4, 4];
    let edge_targets = vec![1u32, 2, 3, 3];
    let edge_kind_mask = vec![edge_kind::CONTROL; 4];
    let seed_bits = vec![0b1000u32];
    let cpu = live_closure_cpu(4, &edge_offsets, &edge_targets, &edge_kind_mask, &seed_bits);
    let gpu = live_closure_borrowed_via(
        &dispatch,
        4,
        &edge_offsets,
        &edge_targets,
        &edge_kind_mask,
        &seed_bits,
        16,
    )
    .expect("dispatch");
    assert_eq!(gpu, cpu);
    assert_eq!(gpu, vec![0b1111u32]);
}

#[test]
fn cuda_live_disconnected_pre_succ() {
    let backend = live_dispatcher();
    let dispatch = run_dispatch(&backend);
    // 0 -> 1, 2 -> 3 (two disjoint chains).
    let edge_offsets = vec![0u32, 1, 1, 2, 2];
    let edge_targets = vec![1u32, 3];
    let edge_kind_mask = vec![edge_kind::CONTROL; 2];
    // Live at 1 only. Should backprop: 0 also live (was a pred of 1), 2 and 3 stay zero.
    let seed_bits = vec![0b0010u32];
    let cpu = live_closure_cpu(4, &edge_offsets, &edge_targets, &edge_kind_mask, &seed_bits);
    let gpu = live_closure_borrowed_via(
        &dispatch,
        4,
        &edge_offsets,
        &edge_targets,
        &edge_kind_mask,
        &seed_bits,
        16,
    )
    .expect("dispatch");
    assert_eq!(gpu, cpu);
    assert_eq!(gpu, vec![0b0011u32]);
}

#[test]
fn cuda_live_self_loop_terminates() {
    let backend = live_dispatcher();
    let dispatch = run_dispatch(&backend);
    let edge_offsets = vec![0u32, 1, 1];
    let edge_targets = vec![0u32];
    let edge_kind_mask = vec![edge_kind::CONTROL];
    let seed_bits = vec![0b01u32];
    let cpu = live_closure_cpu(2, &edge_offsets, &edge_targets, &edge_kind_mask, &seed_bits);
    let gpu = live_closure_borrowed_via(
        &dispatch,
        2,
        &edge_offsets,
        &edge_targets,
        &edge_kind_mask,
        &seed_bits,
        50,
    )
    .expect("dispatch");
    assert_eq!(gpu, cpu);
}

#[test]
fn cuda_live_non_control_edge_blocked() {
    let backend = live_dispatcher();
    let dispatch = run_dispatch(&backend);
    // Non-CONTROL edge 0 -> 1. Even though dst=1 is live, src=0 stays
    // not-live because the edge is filtered out.
    let edge_offsets = vec![0u32, 1, 1];
    let edge_targets = vec![1u32];
    let edge_kind_mask = vec![edge_kind::ASSIGNMENT];
    let seed_bits = vec![0b10u32];
    let cpu = live_closure_cpu(2, &edge_offsets, &edge_targets, &edge_kind_mask, &seed_bits);
    let gpu = live_closure_borrowed_via(
        &dispatch,
        2,
        &edge_offsets,
        &edge_targets,
        &edge_kind_mask,
        &seed_bits,
        16,
    )
    .expect("dispatch");
    assert_eq!(gpu, cpu);
    assert_eq!(gpu, vec![0b10u32]); // src=0 stays not-live
}

#[test]
fn cuda_live_resident_graph_parity() {
    let backend = CudaBackendRegistration::new(live_dispatcher());
    let edge_offsets = vec![0u32, 2, 3, 4, 4];
    let edge_targets = vec![1u32, 2, 3, 3];
    let edge_kind_mask = vec![edge_kind::CONTROL; 4];
    let seed_bits = vec![0b1000u32];
    let cpu = live_closure_cpu(4, &edge_offsets, &edge_targets, &edge_kind_mask, &seed_bits);
    let plan = prepare_live_plan(4, &edge_offsets, &edge_targets, &edge_kind_mask)
        .expect("live plan must prepare for resident CUDA parity");
    let config = DispatchConfig::default();
    let pipeline = backend
        .compile_native(plan.program(), &config)
        .expect("CUDA native compile must succeed for resident live parity")
        .expect("CUDA backend registration must return a native compiled pipeline for resident live parity");
    let resident = FixedPointResidentGraph::upload(&backend, plan.graph())
        .expect("CUDA backend must upload live graph as resident resources");
    let mut scratch = weir::fixed_point_scratch::FixedPointScratch::default();

    let gpu = resident
        .live_resident_frontier(&backend, &*pipeline, &seed_bits, 16, &config, &mut scratch)
        .expect("resident CUDA live dispatch must converge");

    assert_eq!(gpu, cpu);
    assert_eq!(gpu, vec![0b1111u32]);
    resident
        .free(&backend)
        .expect("resident live graph resources must free cleanly");
}

#[test]
fn cuda_live_resident_frontier_reuse_parity() {
    let backend = CudaBackendRegistration::new(live_dispatcher());
    let edge_offsets = vec![0u32, 2, 3, 4, 4];
    let edge_targets = vec![1u32, 2, 3, 3];
    let edge_kind_mask = vec![edge_kind::CONTROL; 4];
    let first_seed = vec![0b1000u32];
    let second_seed = vec![0b0010u32];
    let first_cpu = live_closure_cpu(
        4,
        &edge_offsets,
        &edge_targets,
        &edge_kind_mask,
        &first_seed,
    );
    let second_cpu = live_closure_cpu(
        4,
        &edge_offsets,
        &edge_targets,
        &edge_kind_mask,
        &second_seed,
    );
    let plan = prepare_live_plan(4, &edge_offsets, &edge_targets, &edge_kind_mask)
        .expect("live plan must prepare for resident CUDA frontier reuse parity");
    let config = DispatchConfig::default();
    let pipeline = backend
        .compile_native(plan.program(), &config)
        .expect("CUDA native compile must succeed for resident live frontier reuse parity")
        .expect("CUDA backend registration must return a native compiled pipeline for resident live frontier reuse parity");
    let resident_plan = FixedPointResidentPlan::new(&backend, plan.graph(), pipeline)
        .expect("CUDA backend must upload live resident plan resources");
    let mut scratch = weir::fixed_point_scratch::FixedPointScratch::default();
    let mut resident_frontier = FixedPointResidentFrontierScratch::default();

    let first_gpu = resident_plan
        .live_reusing_frontier(
            &backend,
            &first_seed,
            16,
            &config,
            &mut scratch,
            &mut resident_frontier,
        )
        .expect("first resident CUDA live frontier reuse run must converge");
    let second_gpu = resident_plan
        .live_reusing_frontier(
            &backend,
            &second_seed,
            16,
            &config,
            &mut scratch,
            &mut resident_frontier,
        )
        .expect("second resident CUDA live frontier reuse run must converge");

    assert!(resident_frontier.is_allocated());
    assert_eq!(resident_frontier.byte_len(), 4);
    assert_eq!(first_gpu, first_cpu);
    assert_eq!(second_gpu, second_cpu);
    assert_eq!(first_gpu, vec![0b1111u32]);
    assert_eq!(second_gpu, vec![0b0011u32]);
    resident_plan
        .free_with_frontier(&backend, &mut resident_frontier)
        .expect("resident live plan and frontier resources must free cleanly");
}

#[test]
fn cuda_live_resident_sequence_window_uses_final_convergence_readback_only() {
    let backend = CudaBackendRegistration::new(live_dispatcher());
    let edge_offsets = vec![0u32, 1, 2, 3, 4, 4];
    let edge_targets = vec![1u32, 2, 3, 4];
    let edge_kind_mask = vec![edge_kind::CONTROL; 4];
    let seed_bits = vec![0b10000u32];
    let cpu = live_closure_cpu(5, &edge_offsets, &edge_targets, &edge_kind_mask, &seed_bits);
    let plan = prepare_live_plan(5, &edge_offsets, &edge_targets, &edge_kind_mask)
        .expect("live plan must prepare for sequence-window resident CUDA parity");
    let config = DispatchConfig::default();
    let pipeline = backend
        .compile_native(plan.program(), &config)
        .expect("CUDA native compile must succeed for resident live sequence-window parity")
        .expect("CUDA backend registration must return a native compiled pipeline for resident live sequence-window parity");
    let resident_plan = FixedPointResidentPlan::new(&backend, plan.graph(), pipeline)
        .expect("CUDA backend must upload live sequence-window resident plan resources");
    let mut scratch = weir::fixed_point_scratch::FixedPointScratch::default();
    let mut resident_frontier = FixedPointResidentFrontierScratch::default();

    backend.reset_telemetry();
    let gpu = resident_plan
        .live_reusing_frontier_sequence_window(
            &backend,
            &plan,
            &seed_bits,
            16,
            &mut scratch,
            &mut resident_frontier,
        )
        .expect("sequence-window resident CUDA live dispatch must converge");
    let telemetry = backend.telemetry_snapshot();

    assert_eq!(gpu, cpu);
    assert_eq!(gpu, vec![0b11111u32]);
    assert_eq!(
        telemetry.readback_bytes, 8,
        "Fix: sequence-window live resident solve must read only final and previous one-word frontiers, not poll every iteration."
    );
    assert!(
        telemetry.sync_points < 16,
        "Fix: sequence-window live resident solve must not synchronize once per fixed-point iteration; observed {} sync points.",
        telemetry.sync_points
    );
    assert!(
        telemetry.param_upload_bytes < telemetry.kernel_launches.saturating_mul(4),
        "Fix: sequence-window live resident solve must hoist repeated fixed-point launch parameter uploads; observed {} parameter bytes across {} launches.",
        telemetry.param_upload_bytes,
        telemetry.kernel_launches
    );
    resident_plan
        .free_with_frontier(&backend, &mut resident_frontier)
        .expect("sequence-window resident live plan and frontier resources must free cleanly");
}