inferencelayer 0.2.12

//! CUDA-tower gate + 3-way race: CPU oracle vs wgpu tower vs CUDA tower on the same page grid.
//!
//!   GLM_DIR=<checkpoint> cuda-tower-gate
//!
//! Prints cosine/max|Δ| of both GPU arms against the CPU tower (the transformer-oracle-gated
//! reference) and per-forward wall times. This binary IS the accept/reject evidence for the
//! CUDA arm — quality first, speed second.

use inferencelayer::vision::ImagePatches;

fn fill(n: usize) -> Vec<f32> {
    let mut s = 7u32;
    (0..n)
        .map(|_| {
            s = s.wrapping_mul(1664525).wrapping_add(1013904223);
            ((s >> 8) as f32 / 16_777_216.0 - 0.5) * 0.2
        })
        .collect()
}

fn cos(a: &[f32], b: &[f32]) -> (f64, f32) {
    let (mut dot, mut na, mut nb, mut mx) = (0f64, 0f64, 0f64, 0f32);
    for (x, y) in a.iter().zip(b) {
        dot += (*x as f64) * (*y as f64);
        na += (*x as f64) * (*x as f64);
        nb += (*y as f64) * (*y as f64);
        mx = mx.max((x - y).abs());
    }
    (dot / (na.sqrt() * nb.sqrt()), mx)
}

fn main() -> anyhow::Result<()> {
    let dir = std::env::var("GLM_DIR").expect("GLM_DIR");
    let cpu = inferencelayer::vision_glm::GlmVisionTower::load(std::path::Path::new(&dir))?;
    let pd = cpu.cfg.patch_dim();
    let grids: &[(&str, u32, u32)] = if std::env::var("GATE_BIG").is_ok() {
        &[("~760x980", 54, 70), ("~1240x1600", 88, 114)]
    } else {
        &[("~760x980", 54, 70)]
    };
    for &(label, gh, gw) in grids {
        let n = (gh * gw) as usize;
        let img = ImagePatches { patches: fill(n * pd), grid: [1, gh, gw] };
        let want = cpu.forward(std::slice::from_ref(&img))?;

        // CUDA arm
        let cuda = inferencelayer::cuda_tower::CudaGlmTower::new(&cpu)?;
        let got = cuda.forward(std::slice::from_ref(&img))?;
        let (c, d) = cos(&got, &want);
        let t = std::time::Instant::now();
        for _ in 0..3 {
            let _ = cuda.forward(std::slice::from_ref(&img))?;
        }
        let cuda_ms = t.elapsed().as_secs_f64() * 1000.0 / 3.0;
        println!("CUDA  {label}: cos={c:.9} max|d|={d:.2e}  {cuda_ms:.0} ms");
        drop(cuda);

        // wgpu arm (prod env mirror comes from the job env)
        if let Ok(ctx) = inferencelayer::GpuCtx::new() {
            let gpu = inferencelayer::vision_glm_gpu::GlmVisionGpu::new(&ctx, &cpu)?;
            let got = gpu.forward(&ctx, std::slice::from_ref(&img))?;
            let (c, d) = cos(&got, &want);
            let t = std::time::Instant::now();
            for _ in 0..3 {
                let _ = gpu.forward(&ctx, std::slice::from_ref(&img))?;
            }
            let wgpu_ms = t.elapsed().as_secs_f64() * 1000.0 / 3.0;
            println!("wgpu  {label}: cos={c:.9} max|d|={d:.2e}  {wgpu_ms:.0} ms");
        }
    }
    Ok(())
}