//! 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` and `int_mlicpp_natural`: the `wai.neural.int_mlicpp` vectors; //! - `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 { 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, String>; fn bytes_of(d: &IntDecoded) -> Vec { 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.0.to_le_bytes()); b.extend_from_slice(&hi.0.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 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 { 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> = 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 = 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, ¶ms).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 { 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 { // The synthetic vector, and `natural`, the set trained on openly licensed natural pictures. [("int_mlicpp", "wai-rs/tests/vectors/int_mlicpp"), ("int_mlicpp_natural", "wai-rs/tests/vectors/int_mlicpp_natural")] .into_iter() .map(|(name, dir)| { 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:?}")) }; Case { cap: "wai.neural.int_mlicpp", name: name.into(), refusal: false, run: Box::new(run) } }) .collect() } fn world_cases() -> Vec { 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 = 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> = 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 { v[k].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect() } fn i32v(v: &serde_json::Value, k: &str) -> Vec { 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, 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, Option) { 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, } #[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 = 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 = 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()); } } } } /// A sequence of pictures decoded as a sink decodes a stream (`int_payload::decode_int_seq`): /// each payload's synthesis on a thread of its own while the next payload's entropy stage runs. /// Its results are decoding each payload in turn, on the CPU and with the graphs on every /// adapter, for the committed `natural` vectors of `wai.neural.int_hyper` and /// `wai.neural.int_mlicpp` and for payloads of other sizes drawn from each set's own entropy /// model (`int_hyper_synth::ops::synth_payload`, `int_mlicpp::ops::synth_payload`), a refusal /// (a payload cut short) among them; and every graph the device ran gives the interpreter's /// outputs. #[test] fn a_sequence_decodes_pipelined_as_in_turn_on_every_adapter() { let d = wai_ops_gpu::blocking::discover(); if d.found.is_empty() { report("ops-gpu sequence: 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 dir = "wai-rs/tests/vectors/int_hyper"; let hyper: BTreeMap> = [("entropy.json", "natural_entropy.json"), ("g_s.whm", "natural_g_s.bin"), ("h_s.whm", "natural_h_s.bin")].into_iter().map(|(k, f)| (k.to_string(), read(&format!("{dir}/{f}")))).collect(); let mlic: BTreeMap> = [("model.wmb".to_string(), read("wai-rs/tests/vectors/int_mlicpp_natural/model.wmb"))].into_iter().collect(); let g_s = IntModel::from_bytes(&hyper["g_s.whm"]).unwrap(); let h_s = IntModel::from_bytes(&hyper["h_s.whm"]).unwrap(); let ent = wai::int_hyper_synth::CondEntropy::read(&hyper["entropy.json"]).unwrap(); let m = wai::int_mlicpp::IntMlicModel::from_bundle(&mlic["model.wmb"]).unwrap(); let mut hp = vec![read(&format!("{dir}/natural_bitstream.wih1"))]; let mut mp = vec![wai::container::Wai::from_bytes(&read("wai-rs/tests/vectors/int_mlicpp_natural/stream.wai")).unwrap().payload]; for (seed, (h, w)) in [(1u64, (128, 192)), (2, (192, 128)), (3, (64, 64)), (4, (128, 128))] { hp.push(wai::int_hyper_synth::ops::synth_payload(&g_s, &h_s, &ent, h, w, seed).expect("a drawn int_hyper payload")); mp.push(wai::int_mlicpp::ops::synth_payload(&m, h, w, 4, seed).expect("a drawn int_mlicpp payload")); } // A payload cut short in the middle of each sequence: refused, the rest decoded. for p in [&mut hp, &mut mp] { let cut = p[1][..p[1].len() - 9].to_vec(); p.insert(2, cut); } let seqs = [("wai.neural.int_hyper", hyper, hp), ("wai.neural.int_mlicpp", mlic, mp)]; let devs: Vec> = d.found.into_iter().map(Arc::new).collect(); for (cap, set, payloads) in &seqs { let view: BTreeMap = 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(); let ps: Vec<&[u8]> = payloads.iter().map(|p| p.as_slice()).collect(); let out = |r: Vec>| -> Vec { r.into_iter().map(|x| x.map(|d| bytes_of(&d)).map_err(|e| e.code().to_string())).collect() }; let want: Vec = out(ps.iter().map(|p| decode_int(cap, p, ¶ms)).collect()); assert!(want.iter().filter(|w| w.is_err()).count() == 1 && want.len() == 6, "{cap}: one refusal among six payloads on the CPU: {want:?}"); assert_eq!(out(wai::int_payload::decode_int_seq(cap, &ps, ¶ms)), want, "{cap}: the sequence on the CPU"); for g in &devs { let (got, s) = gpu::with_gpu_compared(g, || out(wai::int_payload::decode_int_seq(cap, &ps, ¶ms))); report(&format!( "ops-gpu sequence: key={} cap={cap} payloads={} refused={} graphs_gpu={} graphs_compared={} graphs_differ={} errors={} equal={}", g.adapter().key(), ps.len(), want.iter().filter(|w| w.is_err()).count(), s.on_gpu(), s.compared, s.differed, s.errors.len(), got == want )); assert_eq!(got, want, "{} {cap}: the pipelined sequence with its graphs on the device", g.adapter().key()); assert!(s.errors.is_empty() && s.differed == 0 && s.on_gpu() > 0, "{} {cap}: {:?}", g.adapter().key(), s); } } }