open-standards/wai/wai-rs/tests/ops_gpu.rs

//! The integer decoders with their graphs on a GPU (`int_ops::gpu`, `wai_ops_gpu`): every
//! committed vector and refusal fixture of every capability whose decode runs a
//! `wai-ops/1` graph decodes to the same bytes, or is refused with the same code, as the
//! decode whose graphs run on the CPU, on every adapter this process finds.
//!
//! - `int-payload-conformance/vectors.json`: its payload vectors and refusal fixtures
//!   (`wai.neural.int_hyper`, `wai.video.int_hyper`, `wai.audio.int_codec`,
//!   `wai.splat.int_codec`; `wai.video.int_motion` runs no graph, and is listed with no
//!   graph runs);
//! - `neural-conformance/`: the `wai.neural.int_synth` vectors;
//! - `wai-rs/tests/vectors/int_mlicpp`: the `wai.neural.int_mlicpp` vector;
//! - `wai-rs/tests/vectors/int_world{,2}`: the integer renderer of `wai.world.model`
//!   (informative, integer-operators §13.5);
//! - `wai.image.jpegai`'s exact decoder on the `jpegai/v1` inputs, when
//!   `WAI_JPEGAI_CODESTREAM`, `WAI_JPEGAI_WEIGHTS_DIR` and `WAI_JPEGAI_NORM` name them
//!   (`tools/wai_jpegai_fetch_e2e_inputs.sh`). Its scale graphs and hyper decoders run on
//!   the device; its synthesis graphs use operators outside `wai-ops/1-i32` and run on the
//!   CPU, which the line for it states.
//!
//! Each graph the device runs is also run on the interpreter, and its outputs compared
//! (`gpu::with_gpu_compared`): so the equality is shown graph by graph, and not only
//! through each decode's final bytes.
//!
//! **A skip is not a pass.** With no adapter the test prints `ops-gpu capabilities:
//! SKIPPED` and fails when `WAI_OPS_GPU_REQUIRE` is set; without the `jpegai/v1` inputs it
//! prints that `wai.image.jpegai` was not run, and fails when `WAI_REQUIRE_SOTA_VECTORS`
//! is set. A capability whose decodes ran no graph on the device fails. Each adapter's
//! result is one line per capability, `ops-gpu capability: key=…`, appended to
//! `WAI_OPS_GPU_REPORT` when it is set.
#![cfg(all(feature = "ops_gpu", feature = "video_int", feature = "audio_int", feature = "splat_int", feature = "world_int"))]

use std::collections::BTreeMap;
use std::io::Write;
use std::path::PathBuf;
use std::sync::Arc;

use serde_json::Value;
use wai::int_hyper_synth::IntModel;
use wai::int_ops::gpu::{self, Stats, wai_ops_gpu};
use wai::int_payload::{IntDecoded, IntParams, decode_int};
use wai::pin::HeldBytes;

fn wai() -> PathBuf {
    PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("..")
}
fn read(rel: &str) -> Vec<u8> {
    std::fs::read(wai().join(rel)).unwrap_or_else(|e| panic!("{rel}: {e}"))
}
fn json(rel: &str) -> Value {
    serde_json::from_slice(&read(rel)).unwrap_or_else(|e| panic!("{rel}: {e}"))
}

fn report(line: &str) {
    println!("{line}");
    if let Ok(p) = std::env::var("WAI_OPS_GPU_REPORT") {
        let mut f = std::fs::OpenOptions::new().create(true).append(true).open(&p).unwrap();
        writeln!(f, "{line}").unwrap();
    }
}

/// A decode's result as bytes: every output byte the decode gives, or its refusal code.
type Out = Result<Vec<u8>, String>;

fn bytes_of(d: &IntDecoded) -> Vec<u8> {
    let mut b = Vec::new();
    match d {
        IntDecoded::Image { rgb, width, height } => {
            b.extend_from_slice(&width.to_le_bytes());
            b.extend_from_slice(&height.to_le_bytes());
            b.extend_from_slice(rgb);
        }
        IntDecoded::Video { frames, width, height } => {
            b.extend_from_slice(&width.to_le_bytes());
            b.extend_from_slice(&height.to_le_bytes());
            for f in frames {
                b.extend_from_slice(f);
            }
        }
        IntDecoded::Pcm { samples, sample_rate, channels } => {
            b.extend_from_slice(&sample_rate.to_le_bytes());
            b.extend_from_slice(&channels.to_le_bytes());
            for s in samples {
                b.extend_from_slice(&s.to_le_bytes());
            }
        }
        IntDecoded::Attr { bytes, n_attr, s, n_gaussians, ranges } => {
            for v in [n_attr, s, n_gaussians] {
                b.extend_from_slice(&v.to_le_bytes());
            }
            for (lo, hi) in ranges {
                b.extend_from_slice(&lo.to_bits().to_le_bytes());
                b.extend_from_slice(&hi.to_bits().to_le_bytes());
            }
            b.extend_from_slice(bytes);
        }
        _ => panic!("an output form this test does not compare"),
    }
    b
}

/// One decode of the corpus: its capability, its name, and the decode.
struct Case {
    cap: &'static str,
    name: String,
    refusal: bool,
    run: Box<dyn Fn() -> Out>,
}

fn cap_static(c: &str) -> &'static str {
    for k in [
        "wai.video.int_motion", "wai.neural.int_hyper", "wai.video.int_hyper", "wai.audio.int_codec", "wai.splat.int_codec", "wai.neural.int_synth",
        "wai.neural.int_mlicpp", "wai.world.model", "wai.image.jpegai",
    ] {
        if k == c {
            return k;
        }
    }
    panic!("an unknown capability {c}")
}

fn payload_cases() -> Vec<Case> {
    let idx = json("int-payload-conformance/vectors.json");
    let mut out = Vec::new();
    for (list, refusal) in [("vectors", false), ("refusals", true)] {
        for e in idx[list].as_array().unwrap() {
            let cap = cap_static(e["capability"].as_str().unwrap());
            let payload = read(e["payload"].as_str().unwrap());
            let set: BTreeMap<String, Vec<u8>> = e["set"].as_object().into_iter().flatten().map(|(k, v)| (k.clone(), read(v.as_str().unwrap()))).collect();
            let run = move || -> Out {
                let view: BTreeMap<String, &[u8]> = set.iter().map(|(k, v)| (k.clone(), v.as_slice())).collect();
                let params = wai::int_payload::resolve(cap, None, &[HeldBytes::Set { listing: b"", files: &view }], None).unwrap_or_else(|_| IntParams::none());
                decode_int(cap, &payload, &params).map(|d| bytes_of(&d)).map_err(|e| e.code().to_string())
            };
            out.push(Case { cap, name: e["name"].as_str().unwrap().to_string(), refusal, run: Box::new(run) });
        }
    }
    out
}

fn synth_cases() -> Vec<Case> {
    let hashes = json("neural-conformance/expected_hashes.json");
    hashes
        .as_object()
        .unwrap()
        .keys()
        .map(|name| {
            let dir = format!("neural-conformance/vectors/{name}");
            let (payload, set) = (read(&format!("{dir}/latent.bin")), read(&format!("{dir}/int_synth.json")));
            let run = move || -> Out { wai::int_synth::decode_from_spec(&payload, &set).map(|i| [i.width.to_le_bytes().to_vec(), i.height.to_le_bytes().to_vec(), i.rgb].concat()).map_err(|e| e.to_string()) };
            Case { cap: "wai.neural.int_synth", name: name.clone(), refusal: false, run: Box::new(run) }
        })
        .collect()
}

fn mlicpp_cases() -> Vec<Case> {
    const DIR: &str = "wai-rs/tests/vectors/int_mlicpp";
    let m = wai::int_mlicpp::IntMlicModel::from_bundle(&read(&format!("{DIR}/model.wmb"))).unwrap();
    let p = wai::container::Wai::from_bytes(&read(&format!("{DIR}/stream.wai"))).unwrap().payload;
    let run = move || -> Out { wai::int_mlicpp::decode(&m, &p).map(|i| [(i.width as u64).to_le_bytes().to_vec(), (i.height as u64).to_le_bytes().to_vec(), i.rgb].concat()).map_err(|e| format!("{e:?}")) };
    vec![Case { cap: "wai.neural.int_mlicpp", name: "int_mlicpp".into(), refusal: false, run: Box::new(run) }]
}

fn world_cases() -> Vec<Case> {
    let mut out = Vec::new();
    for (name, dir, js, t, g) in [("world", "int_world", "world.json", "world_T.bin", "world_gs.bin"), ("world2", "int_world2", "world2.json", "world2_T.bin", "world2_gs.bin")] {
        let j = json(&format!("wai-rs/tests/vectors/{dir}/{js}"));
        let tm = IntModel::from_bytes(&read(&format!("wai-rs/tests/vectors/{dir}/{t}"))).unwrap();
        let gm = IntModel::from_bytes(&read(&format!("wai-rs/tests/vectors/{dir}/{g}"))).unwrap();
        let c = j["c"].as_u64().unwrap() as usize;
        let lat = j["lat"].as_u64().unwrap() as usize;
        let z0: Vec<i64> = j["z0"].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect();
        let afrac = j["afrac"].as_u64().unwrap() as u32;
        let q = |d: i64| wai_ops::stage::dequantize(d, afrac, 0).unwrap();
        let mv = |a: &Value| wai::int_world::ACTIONS[a.as_u64().unwrap() as usize];
        let rows: Vec<Vec<i64>> = match j.get("pairs") {
            Some(p) => p
                .as_array()
                .unwrap()
                .iter()
                .map(|p| {
                    let ((xa, ya), (xb, yb)) = (mv(&p[0]), mv(&p[1]));
                    vec![q(xa), q(ya), q(xb), q(yb)]
                })
                .collect(),
            None => j["actions"]
                .as_array()
                .unwrap()
                .iter()
                .map(|a| {
                    let (dx, dy) = mv(a);
                    vec![q(dx), q(dy), 0]
                })
                .collect(),
        };
        let run = move || -> Out { wai::int_world::decode_planes(&tm, &gm, &z0, &rows, c, lat).map(|w| w.frames.concat()).map_err(|e| format!("{e:?}")) };
        out.push(Case { cap: "wai.world.model", name: name.into(), refusal: false, run: Box::new(run) });
    }
    out
}

mod jpegai {
    use wai::jpegai_decode::{DecoderWeights, NormTables, PlaneWeights};
    use wai::jpegai_icci::{Conv, IcciNet, IcciPath, ResBlock};
    use wai::jpegai_synthesis::{IntNet, Weights};

    fn jload(p: &std::path::Path) -> serde_json::Value {
        serde_json::from_slice(&std::fs::read(p).unwrap()).unwrap()
    }
    fn w_of(dir: &str, name: &str) -> Weights {
        Weights::from_json(&jload(&std::path::Path::new(dir).join(name))).unwrap()
    }
    fn int_of(dir: &str, name: &str) -> IntNet {
        IntNet::from_json(&jload(&std::path::Path::new(dir).join(name))).unwrap()
    }
    fn i64v(v: &serde_json::Value, k: &str) -> Vec<i64> {
        v[k].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect()
    }
    fn i32v(v: &serde_json::Value, k: &str) -> Vec<i32> {
        v[k].as_array().unwrap().iter().map(|x| x.as_i64().unwrap() as i32).collect()
    }
    fn f3(v: &serde_json::Value, k: &str) -> [f32; 3] {
        let a = v[k].as_array().unwrap();
        [a[0].as_f64().unwrap() as f32, a[1].as_f64().unwrap() as f32, a[2].as_f64().unwrap() as f32]
    }
    fn i3(v: &serde_json::Value, k: &str) -> [i64; 3] {
        let a = v[k].as_array().unwrap();
        [a[0].as_i64().unwrap(), a[1].as_i64().unwrap(), a[2].as_i64().unwrap()]
    }
    fn conv_of(c: &serde_json::Value) -> Conv {
        let g = |k: &str| c[k].as_u64().unwrap() as usize;
        Conv {
            w: c["w"].as_array().unwrap().iter().map(|x| x.as_f64().unwrap() as f32).collect(),
            b: c["b"].as_array().unwrap().iter().map(|x| x.as_f64().unwrap() as f32).collect(),
            cin: g("cin"),
            cout: g("cout"),
            kh: g("kh"),
            kw: g("kw"),
            ph: g("ph"),
            pw: g("pw"),
        }
    }
    fn path_of(v: &serde_json::Value) -> IcciPath {
        IcciPath {
            conv_first: conv_of(&v["conv_first"]),
            blocks: v["blocks"]
                .as_array()
                .unwrap()
                .iter()
                .map(|b| ResBlock { conv1: conv_of(&b["c1"]), conv2: conv_of(&b["c2"]), scale: b["scale"].as_f64().unwrap() as f32 })
                .collect(),
            conv_last: conv_of(&v["conv_last"]),
            scale: v["scale"].as_f64().unwrap() as f32,
        }
    }

    /// The `jpegai/v1` inputs, as `tests/jpegai_e2e_full.rs` loads them.
    pub fn load(cs: &str, dir: &str, nf: &str) -> (Vec<u8>, DecoderWeights, NormTables) {
        let data = std::fs::read(cs).unwrap_or_else(|e| panic!("read {cs}: {e}"));
        let nv = jload(std::path::Path::new(nf));
        let iw = jload(&std::path::Path::new(dir).join("icci.json"));
        let icci = IcciNet { luma: path_of(&iw["luma"]), chroma: path_of(&iw["chroma"]), conv_last_v: conv_of(&iw["conv_last_v"]), scale_v: iw["scale_v"].as_f64().unwrap() as f32 };
        let weights = DecoderWeights {
            luma: PlaneWeights { hyper_scale: int_of(dir, "hs_y.json"), hyper_decoder: w_of(dir, "hyp_y.json"), synthesis: w_of(dir, "synth_y.json") },
            chroma: PlaneWeights { hyper_scale: int_of(dir, "hs_uv.json"), hyper_decoder: w_of(dir, "hyp_uv.json"), synthesis: w_of(dir, "synth_uv.json") },
            mcm_y: [w_of(dir, "mcm0_y.json"), w_of(dir, "mcm1_y.json"), w_of(dir, "mcm2_y.json"), w_of(dir, "mcm3_y.json")],
            icci,
        };
        let norm = NormTables {
            pmf: nv["pmf"].as_array().unwrap().iter().map(|r| r.as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect()).collect(),
            bounds: i64v(&nv, "bounds"),
            per_channel: nv["per_channel"].as_u64().unwrap() as usize,
            cdfs_y: i64v(&nv, "cdfs_y"),
            cdfs_uv: i64v(&nv, "cdfs_uv"),
            gain_y: i64v(&nv, "gain_y"),
            gain_uv: i64v(&nv, "gain_uv"),
            lsbs_t0: i32v(&nv, "lsbs_t0"),
            lsbs_t1: i32v(&nv, "lsbs_t1"),
            lef_thr: i3(&nv, "lef_thr0"),
            lef_mag: f3(&nv, "lef_mag0"),
        };
        (data, weights, norm)
    }
}

fn jpegai_cases() -> (Vec<Case>, Option<String>) {
    let env = |k: &str| std::env::var(k).ok().filter(|v| !v.is_empty());
    let (Some(cs), Some(dir), Some(nf)) = (env("WAI_JPEGAI_CODESTREAM"), env("WAI_JPEGAI_WEIGHTS_DIR"), env("WAI_JPEGAI_NORM")) else {
        return (vec![], Some("WAI_JPEGAI_CODESTREAM, WAI_JPEGAI_WEIGHTS_DIR and WAI_JPEGAI_NORM are not all set (tools/wai_jpegai_fetch_e2e_inputs.sh)".into()));
    };
    let (data, weights, norm) = jpegai::load(&cs, &dir, &nf);
    let run = move || -> Out { wai::jpegai_decode::decode(&data, &weights, &norm).map(|i| [(i.h as u64).to_le_bytes().to_vec(), (i.w as u64).to_le_bytes().to_vec(), i.data.clone()].concat()).map_err(|r| format!("{:?}", r.code())) };
    (vec![Case { cap: "wai.image.jpegai", name: "jpegai/v1 test256".into(), refusal: false, run: Box::new(run) }], None)
}

/// What one capability's cases did on one adapter.
#[derive(Default)]
struct Tally {
    cases: u32,
    refusals: u32,
    stats: Stats,
    differ: Vec<String>,
}

#[test]
fn every_capability_decodes_alike_with_its_graphs_on_every_adapter() {
    let d = wai_ops_gpu::blocking::discover();
    for (a, why) in &d.refused {
        report(&format!("ops-gpu adapter refused: key={} {} reason={why:?}", a.key(), a.describe()));
    }
    if d.found.is_empty() {
        report("ops-gpu capabilities: SKIPPED: no adapter");
        assert!(std::env::var_os("WAI_OPS_GPU_REQUIRE").is_none(), "WAI_OPS_GPU_REQUIRE is set and no adapter was found");
        return;
    }
    let mut cases = payload_cases();
    cases.extend(synth_cases());
    cases.extend(mlicpp_cases());
    cases.extend(world_cases());
    let (jp, missing) = jpegai_cases();
    cases.extend(jp);
    if let Some(why) = &missing {
        report(&format!("ops-gpu capability: wai.image.jpegai NOT RUN: {why}"));
        assert!(std::env::var_os("WAI_REQUIRE_SOTA_VECTORS").is_none(), "WAI_REQUIRE_SOTA_VECTORS is set and the jpegai/v1 inputs are absent");
    }
    // The CPU's result for every case, once.
    let want: Vec<Out> = cases.iter().map(|c| (c.run)()).collect();
    for (c, w) in cases.iter().zip(&want) {
        assert_eq!(c.refusal, w.is_err(), "{} {}: a vector decodes and a refusal fixture is refused, on the CPU: {:?}", c.cap, c.name, w.as_ref().err());
    }
    for g in d.found.into_iter().map(Arc::new) {
        let a = g.adapter().clone();
        let mut by_cap: BTreeMap<&'static str, Tally> = BTreeMap::new();
        for (c, w) in cases.iter().zip(&want) {
            let (got, stats) = gpu::with_gpu_compared(&g, || (c.run)());
            let t = by_cap.entry(c.cap).or_default();
            t.cases += 1;
            t.refusals += c.refusal as u32;
            if &got != w {
                t.differ.push(format!("{}: {}", c.name, match (w, &got) {
                    (Ok(x), Ok(y)) => format!("{} bytes vs {} bytes, first difference at {:?}", x.len(), y.len(), x.iter().zip(y).position(|(p, q)| p != q)),
                    _ => format!("{:?} vs {:?}", w.as_ref().map(|v| v.len()), got.as_ref().map(|v| v.len())),
                }));
            }
            t.stats.acc32 += stats.acc32;
            t.stats.acc64 += stats.acc64;
            for (k, n) in stats.cpu {
                *t.stats.cpu.entry(k).or_default() += n;
            }
            t.stats.errors.extend(stats.errors);
            t.stats.compared += stats.compared;
            t.stats.differed += stats.differed;
        }
        for (cap, t) in &by_cap {
            let cpu: Vec<String> = t.stats.cpu.iter().map(|(k, n)| format!("{n}x {k}")).collect();
            report(&format!(
                "ops-gpu capability: key={} {} cap={cap} cases={} refusals={} graphs_gpu_acc32={} graphs_gpu_acc64={} graphs_cpu={} [{}] graphs_compared={} graphs_differ={} differ={} errors={}",
                a.key(),
                a.describe(),
                t.cases,
                t.refusals,
                t.stats.acc32,
                t.stats.acc64,
                t.stats.on_cpu(),
                cpu.join("; "),
                t.stats.compared,
                t.stats.differed,
                t.differ.len(),
                t.stats.errors.len()
            ));
        }
        for (cap, t) in &by_cap {
            assert!(t.differ.is_empty(), "{} {cap}: {:?}", a.key(), t.differ);
            assert!(t.stats.errors.is_empty(), "{} {cap}: device errors {:?}", a.key(), t.stats.errors);
            assert_eq!(t.stats.differed, 0, "{} {cap}: a graph's outputs on the device differ from the interpreter's", a.key());
            assert_eq!(t.stats.compared, t.stats.on_gpu(), "{} {cap}: every graph the device ran is compared", a.key());
            if *cap != "wai.video.int_motion" {
                assert!(t.stats.on_gpu() > 0, "{} {cap}: no graph ran on the device", a.key());
            } else {
                assert_eq!(t.stats.on_gpu() + t.stats.on_cpu(), 0, "{} {cap}: int_motion runs no graph", a.key());
            }
            if *cap != "wai.image.jpegai" {
                assert_eq!(t.stats.on_cpu(), 0, "{} {cap}: every graph of this capability is in wai-ops/1-i32", a.key());
            }
        }
    }
}