// wai-conformance: int-kernels int-kernels-amd //! `wai.image.jpegai`'s scale path and entropy decodes through the integer operator set: //! its lowering `jpegai-int/1` (`extensions/integer-operators.md` §13.6.1) and the //! certificate of each bundle (§8.5), as this crate loads them (`jpegai_ops`). //! //! The committed bundle (`ops-conformance/sets/jpegai`, synthetic) has its lowered graphs //! and certificate under `ops-conformance/certificates/`, where `wai-verify` and the //! independent checker (`tools/wai_ops_check.py`) reproduce them byte for byte. The //! `jpegai/v1` bundle of the models root is not committed; its certificate is, by digest, //! and the tests that read it take its files from `WAI_JPEGAI_WEIGHTS_DIR` and //! `WAI_JPEGAI_NORM` (`tools/wai_jpegai_fetch_e2e_inputs.sh`), with //! `WAI_REQUIRE_SOTA_VECTORS` deciding whether their absence is a skip or a failure. #![cfg(all(feature = "neural_int", feature = "legacy_int"))] #[path = "support/require_vectors.rs"] mod require_vectors; use std::collections::BTreeMap; use std::path::{Path, PathBuf}; use serde_json::Value; use wai::jpegai_decode::NormTables; use wai::jpegai_ops; use wai::jpegai_synthesis::IntNet; use wai_ops::Profile; const CAP: &str = "wai.image.jpegai"; const SYNTHETIC: &str = "1bc0daae512797a8817f2e77e57d009ace0d22acf4fe6945af8f4e0ee95244b9"; const V1: &str = "3e243dc11566922f5cb7f4c628afd440b18d56a08b99b07cc33d5c8522aeec37"; const PRODUCE: &str = "tools/wai_jpegai_fetch_e2e_inputs.sh (pinned by tools/jpegai-e2e-v1.sha256)"; fn wai() -> PathBuf { PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("..") } fn read(p: &Path) -> Vec { std::fs::read(p).unwrap_or_else(|e| panic!("{}: {e}", p.display())) } fn json(b: &[u8]) -> Value { serde_json::from_slice(b).unwrap() } fn sha256(b: &[u8]) -> String { wai_ops::sha256::sha256(b).iter().map(|x| format!("{x:02x}")).collect() } fn i64s(v: &Value) -> Vec { v.as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect() } /// The tables of `norm.json` as the decoders parse them. A synthetic set has only the /// members §13.6.1 reads; the others are left empty. fn norm_tables(nv: &Value) -> NormTables { let opt = |k: &str| nv.get(k).map(i64s).unwrap_or_default(); let three = |k: &str, d: f64| -> [f64; 3] { nv.get(k).map(|a| std::array::from_fn(|i| a[i].as_f64().unwrap())).unwrap_or([d; 3]) }; let t = three("lef_thr0", 0.0); let m = three("lef_mag0", 0.0); NormTables { pmf: nv["pmf"].as_array().unwrap().iter().map(i64s).collect(), bounds: i64s(&nv["bounds"]), per_channel: nv["per_channel"].as_u64().unwrap() as usize, cdfs_y: i64s(&nv["cdfs_y"]), cdfs_uv: i64s(&nv["cdfs_uv"]), gain_y: i64s(&nv["gain_y"]), gain_uv: i64s(&nv["gain_uv"]), lsbs_t0: opt("lsbs_t0").iter().map(|&x| x as i32).collect(), lsbs_t1: opt("lsbs_t1").iter().map(|&x| x as i32).collect(), lef_thr: t.map(|x| x as i64), lef_mag: m.map(|x| x as f32), } } /// A bundle's three scale-path members, as bytes: `norm.json`, `hs_y.json`, `hs_uv.json`. struct Bundle { norm: Vec, hs_y: Vec, hs_uv: Vec, } impl Bundle { fn synthetic() -> Bundle { let d = wai().join("ops-conformance/sets/jpegai"); Bundle { norm: read(&d.join("norm.json")), hs_y: read(&d.join("weights/hs_y.json")), hs_uv: read(&d.join("weights/hs_uv.json")) } } /// The `jpegai/v1` bundle, from the variables the whole-decode tests read. fn pinned(test: &str) -> Option<(Bundle, PathBuf)> { let [dir, norm] = require_vectors::env_paths(test, ["WAI_JPEGAI_WEIGHTS_DIR", "WAI_JPEGAI_NORM"], PRODUCE)?; let d = PathBuf::from(dir); Some((Bundle { norm: read(Path::new(&norm)), hs_y: read(&d.join("hs_y.json")), hs_uv: read(&d.join("hs_uv.json")) }, d)) } fn parsed(&self) -> (NormTables, IntNet, IntNet) { (norm_tables(&json(&self.norm)), IntNet::from_json(&json(&self.hs_y)).unwrap(), IntNet::from_json(&json(&self.hs_uv)).unwrap()) } } /// Gate 1, on one bundle: the lowering of its bytes; the bundle as this crate parses it, /// converted (`jpegai_ops::net`), is the bundle `wai-ops` parses from those bytes and /// lowers to the same graphs; the scale path this crate loads (`jpegai_ops::load`) holds /// those graphs and meets `acc32`; and the committed certificate `digest` is those graphs', /// each of whose nodes meets `acc32` with no accelerator flag (the scale path's /// `role_profiles`, §11). fn lowers_to_its_certificate(name: &str, b: &Bundle, digest: &str) { let lowered = wai_ops::lower::jpegai_scale(&b.norm, &b.hs_y, &b.hs_uv).unwrap_or_else(|r| panic!("{name}: does not lower: {r:?}")); let roles: Vec<&str> = lowered.roles.iter().map(|(r, _)| *r).collect(); assert_eq!(roles, ["hs_y", "hs_uv"], "{name}"); let (n, y, uv) = b.parsed(); let net = jpegai_ops::net(&n, &y, &uv).unwrap_or_else(|r| panic!("{name}: does not convert: {r:?}")); assert!(net == wai_ops::lower::parse_jpegai(&b.norm, &b.hs_y, &b.hs_uv).unwrap(), "{name}: the converted bundle is the one wai-ops parses from the bytes"); assert!(wai_ops::lower::jpegai_net(&net).unwrap() == lowered, "{name}: the lowering of the parsed bundle differs from the lowering of its bytes"); let sp = jpegai_ops::load(&n, &y, &uv).unwrap_or_else(|r| panic!("{name}: refused under acc32: {r:?}")); let cert = json(&read(&wai().join(format!("ops-conformance/certificates/{CAP}/{digest}.json")))); assert_eq!(cert["set_sha256"].as_str(), Some(digest), "{name}"); assert_eq!(cert["lowering"], jpegai_ops::LOWERING, "{name}"); for (i, (role, plane)) in [("hs_y", jpegai_ops::Plane::Luma), ("hs_uv", jpegai_ops::Plane::Chroma)].into_iter().enumerate() { let g = lowered.get(role).unwrap(); assert!(*sp.graph(plane).graph() == *g, "{name}: {role} as loaded"); assert!(sp.graph(plane).meets(Profile::Acc32), "{name}: {role} meets acc32"); assert_eq!(cert["graphs"][i]["name"], role, "{name}"); assert_eq!(cert["graphs"][i]["graph_sha256"].as_str(), Some(sha256(&g.to_bytes()).as_str()), "{name}: the certificate is {role}'s"); let x = &cert["graphs"][i]["inputs"][0]; assert_eq!((x["policy"].as_str(), x["lo"].as_str(), x["hi"].as_str()), (Some("check"), Some("-31"), Some("32")), "{name}: z's range"); } // The set's verdict covers its other networks too (§13.6.2), which meet acc64 only; // the scale path's two graphs meet acc32, in the accelerator subset, with no flags // (§11's `role_profiles`). assert_eq!(cert["verdict"]["acc64"].as_bool(), Some(true), "{name}: the certificate's verdict"); for i in 0..2 { for n in cert["graphs"][i]["nodes"].as_array().unwrap() { assert!(n["fits"].as_array().unwrap().iter().any(|p| p == "acc32"), "{name}: node {} of the scale path fits acc32", n["index"]); assert_eq!(n["npu_flags"].as_array().map(Vec::len), Some(0), "{name}: no accelerator flags"); } } } /// The committed synthetic bundle: its lowered scale graphs are the committed ones, and /// its set digest is the listing of its eleven members. #[test] fn the_committed_bundle_lowers_to_its_committed_graphs_and_verdict() { let b = Bundle::synthetic(); let dir = wai().join(format!("ops-conformance/certificates/{CAP}/{SYNTHETIC}")); let lowered = wai_ops::lower::jpegai_scale(&b.norm, &b.hs_y, &b.hs_uv).unwrap(); for role in ["hs_y", "hs_uv"] { assert!(lowered.get(role).unwrap().to_bytes() == read(&dir.join(format!("{role}.whg1"))), "the lowering of the bytes differs from the committed {role}.whg1"); } let nets: Vec<(String, Vec)> = wai_ops::lower::jpegai_nets::NET_MEMBERS.iter().map(|n| (n.to_string(), read(&wai().join("ops-conformance/sets/jpegai").join(n)))).collect(); let mut members: BTreeMap = [("norm.json", &b.norm), ("weights/hs_y.json", &b.hs_y), ("weights/hs_uv.json", &b.hs_uv)].into_iter().map(|(k, v)| (k.to_string(), v.as_slice())).collect(); members.extend(nets.iter().map(|(k, v)| (k.clone(), v.as_slice()))); let digest = wai::int_set_cache::SetKey::listing_digest(&members).unwrap(); assert_eq!(digest.iter().map(|x| format!("{x:02x}")).collect::(), SYNTHETIC, "the set's digest"); lowers_to_its_certificate("synthetic", &b, SYNTHETIC); } /// The `jpegai/v1` bundle: its set digest is the listing of its twelve members, the one /// its set-form pin and its committed certificate name. #[test] fn the_pinned_bundle_lowers_to_its_certificate_and_meets_acc32() { let Some((b, dir)) = Bundle::pinned("jpegai_ops::the_pinned_bundle_lowers_to_its_certificate_and_meets_acc32") else { return; }; let names = ["hs_y", "hs_uv", "hyp_y", "hyp_uv", "synth_y", "synth_uv", "mcm0_y", "mcm1_y", "mcm2_y", "mcm3_y", "icci"]; let files: Vec<(String, Vec)> = names.iter().map(|n| (format!("weights/{n}.json"), read(&dir.join(format!("{n}.json"))))).collect(); let mut members: BTreeMap = files.iter().map(|(k, v)| (k.clone(), v.as_slice())).collect(); members.insert("norm.json".into(), &b.norm); let digest = wai::int_set_cache::SetKey::listing_digest(&members).unwrap(); assert_eq!(digest.iter().map(|x| format!("{x:02x}")).collect::(), V1, "the set's digest"); lowers_to_its_certificate("jpegai/v1", &b, V1); eprintln!("jpegai/v1: hs_y and hs_uv lower to the committed certificate {V1} and meet acc32"); } /// The whole `jpegai/v1` bundle and the test codestream, as the decoders parse them. struct Whole { data: Vec, weights: wai::jpegai_decode::DecoderWeights, norm: NormTables, } impl Whole { fn pinned(test: &str) -> Option { use wai::jpegai_decode::{DecoderWeights, PlaneWeights}; use wai::jpegai_synthesis::Weights; let [cs, dir, nf] = require_vectors::env_paths(test, ["WAI_JPEGAI_CODESTREAM", "WAI_JPEGAI_WEIGHTS_DIR", "WAI_JPEGAI_NORM"], PRODUCE)?; let d = PathBuf::from(dir); let j = |n: &str| json(&read(&d.join(n))); let w = |n: &str| Weights::from_json(&j(n)).unwrap(); let net = |n: &str| IntNet::from_json(&j(n)).unwrap(); let weights = DecoderWeights { luma: PlaneWeights { hyper_scale: net("hs_y.json"), hyper_decoder: w("hyp_y.json"), synthesis: w("synth_y.json") }, chroma: PlaneWeights { hyper_scale: net("hs_uv.json"), hyper_decoder: w("hyp_uv.json"), synthesis: w("synth_uv.json") }, mcm_y: [w("mcm0_y.json"), w("mcm1_y.json"), w("mcm2_y.json"), w("mcm3_y.json")], icci: wai::jpegai_icci::IcciNet::from_json(&j("icci.json")).unwrap(), }; Some(Whole { data: read(Path::new(&cs)), weights, norm: norm_tables(&json(&read(Path::new(&nf)))) }) } } /// A codestream's segment payloads and a plane's GRFS flags, as the decoder reads them. fn segments(data: &[u8], plane: jpegai_ops::Plane) -> (Vec, Vec, Vec) { use wai::jpegai_bitstream::{find_segment, parse_container, parse_picture_header, Marker}; let segs = parse_container(data); let pay = |m| { let s = find_segment(&segs, m).unwrap(); data[s.start..s.start + s.len].to_vec() }; let lat = if plane == jpegai_ops::Plane::Luma { Marker::Sorp } else { Marker::Sors }; let ph = parse_picture_header(data, find_segment(&segs, Marker::Pih).unwrap().start_bit()); let grfs = ph.components[plane as usize].grfs_channel_flags.clone(); (pay(Marker::Soz), pay(lat), grfs) } /// Gate 2, the integers: on the test codestream, each plane's hyper-latent, scale, /// mask, `likely`, RVS addition, `σ_idx` and latent through the operator set are the /// direct computation's, every value. The chroma plane's hyper-latent comes after the 160 /// luma rows its decode reads and discards. #[test] fn the_planes_decode_through_the_lowering_as_directly() { let Some(wh) = Whole::pinned("jpegai_ops::the_planes_decode_through_the_lowering_as_directly") else { return; }; let sp = wai::jpegai_decode::load(&wh.weights, &wh.norm).unwrap(); for (plane, pw, gain) in [ (jpegai_ops::Plane::Luma, &wh.weights.luma, &wh.norm.gain_y), (jpegai_ops::Plane::Chroma, &wh.weights.chroma, &wh.norm.gain_uv), ] { let (soz, lat, grfs) = segments(&wh.data, plane); let ops = jpegai_ops::plane(&sp, plane, &soz, &lat, &grfs).unwrap(); let direct = jpegai_ops::legacy::plane(&wh.norm, &pw.hyper_scale, gain, plane, &soz, &lat, &grfs); for (what, a, b) in [("z", &ops.z, &direct.z), ("v", &ops.v, &direct.v), ("likely", &ops.likely, &direct.likely), ("rvs", &ops.rvs, &direct.rvs), ("σ_idx", &ops.sigma, &direct.sigma), ("y_q", &ops.y_q, &direct.y_q)] { assert!(a == b, "{plane:?}: {what} through the lowering differs from the direct computation"); } assert!(ops.mask == direct.mask, "{plane:?}: the mask"); let coded = ops.mask.iter().filter(|&&m| m).count(); eprintln!("{plane:?}: z, v, mask ({coded} coded), likely, rvs, σ_idx and y_q equal the direct computation's"); } } /// Gate 2, the decode: the whole decode through the operator set gives the direct /// decode's float values after every stage, bit for bit, and its RGB bytes, which hash to /// `WAI_JPEGAI_EXPECT_SHA256` where that is set. #[test] fn the_whole_decode_through_the_lowering_is_the_direct_decode() { let Some(wh) = Whole::pinned("jpegai_ops::the_whole_decode_through_the_lowering_is_the_direct_decode") else { return; }; let bits = |planes: &[&[f32]]| -> String { sha256(&planes.iter().flat_map(|p| p.iter().flat_map(|v| v.to_bits().to_le_bytes())).collect::>()) }; let mut direct = vec![]; let a = wai::jpegai_decode::legacy::decode_with_stages(&wh.data, &wh.weights, &wh.norm, &mut |s, p| direct.push((s.name(), bits(p)))).expect("the direct decode"); let mut ops = vec![]; let b = wai::jpegai_decode::decode_with_stages(&wh.data, &wh.weights, &wh.norm, &mut |s, p| ops.push((s.name(), bits(p)))).expect("the decode through the operator set"); assert_eq!(ops, direct, "the stages through the lowering"); assert!(a.data == b.data, "the RGB bytes through the lowering"); let rgb = sha256(&b.data); match std::env::var("WAI_JPEGAI_EXPECT_SHA256") { Ok(want) if !want.is_empty() => { assert_eq!(rgb, want, "the decode through the operator set is not the pinned decode"); eprintln!("the decode through the operator set matches the pinned WAI_JPEGAI_EXPECT_SHA256 ({} stages equal)", ops.len()); } _ => eprintln!("the decode through the operator set equals the direct decode ({} stages), rgb {rgb}", ops.len()), } } /// One member of the committed bundle changed by `edit`: the bundle's three members as /// bytes, and whether this crate's parser can read it at all (a non-integer value it /// cannot). fn edited(edit: impl Fn(&mut Value, &mut Value)) -> (Bundle, bool) { let b = Bundle::synthetic(); let (mut n, mut y) = (json(&b.norm), json(&b.hs_y)); edit(&mut n, &mut y); let ints = |v: &Value| v.as_array().is_none_or(|a| a.iter().all(|x| x.is_i64() || x.is_array() || x.is_object())); let parseable = y.as_object().unwrap().values().all(|m| ints(&m["v"])); (Bundle { norm: serde_json::to_vec(&n).unwrap(), hs_y: serde_json::to_vec(&y).unwrap(), hs_uv: b.hs_uv }, parseable) } /// Gate 3, the bundle: each load check of §13.6.1, broken one change at a time from the /// committed bundle, is refused `model_mismatch`, by the load from the members' bytes and /// by the load from the bundle as this crate parses it; and so is a bundle whose scale /// graphs do not meet `acc32`. The direct decode made none of these checks /// (integer-operators §15.3 item 4): it read a table value into a byte, shifted by /// whatever a shift said, and accumulated in `i64`. #[test] fn every_load_check_of_the_bundle_refuses_model_mismatch() { let arr = |v: &mut Value, k: &str| -> Vec { v[k].as_array().unwrap().clone() }; let mut cases: Vec<(&str, (Bundle, bool))> = vec![ ("pmf: 31 rows", edited(|n, _| { let mut p = arr(n, "pmf"); p.truncate(31); n["pmf"] = Value::from(p); })), ("pmf: a row summing to 257", edited(|n, _| { let r = &mut n["pmf"][0][0]; *r = Value::from(r.as_i64().unwrap() + 1); })), ("pmf: a row of another width", edited(|n, _| { n["pmf"][1].as_array_mut().unwrap().push(Value::from(0)); })), ("pmf: a negative count", edited(|n, _| { let a = n["pmf"][0][0].as_i64().unwrap(); let b = n["pmf"][0][1].as_i64().unwrap(); n["pmf"][0][0] = Value::from(-1); n["pmf"][0][1] = Value::from(a + b + 1); })), ("bounds: 31", edited(|n, _| { let mut b = arr(n, "bounds"); b.truncate(31); n["bounds"] = Value::from(b); })), ("bounds: one of 256", edited(|n, _| { n["bounds"][0] = Value::from(256); })), ("bounds: one of -1", edited(|n, _| { n["bounds"][0] = Value::from(-1); })), ("per_channel: 15", edited(|n, _| { n["per_channel"] = Value::from(15); })), ("cdfs_y: one short", edited(|n, _| { let mut c = arr(n, "cdfs_y"); c.pop(); n["cdfs_y"] = Value::from(c); })), ("cdfs_y: a count of 256", edited(|n, _| { n["cdfs_y"][62] = Value::from(256); })), ("cdfs_y: a row that decreases", edited(|n, _| { let v = n["cdfs_y"][1].as_i64().unwrap(); n["cdfs_y"][0] = Value::from(v + 1); n["cdfs_y"][1] = Value::from(v); })), ("cdfs_uv: one long", edited(|n, _| { let mut c = arr(n, "cdfs_uv"); c.push(Value::from(255)); n["cdfs_uv"] = Value::from(c); })), ("gain_y: 159", edited(|n, _| { let mut g = arr(n, "gain_y"); g.pop(); n["gain_y"] = Value::from(g); })), ("gain_uv: 97", edited(|n, _| { let mut g = arr(n, "gain_uv"); g.push(Value::from(0)); n["gain_uv"] = Value::from(g); })), ("hs_y: a member missing", edited(|_, y| { y.as_object_mut().unwrap().remove("onnx::Pow_91"); })), ("hs_y: weights of other dims", edited(|_, y| { y["conv1.weight"]["dims"] = serde_json::json!([160, 2, 1, 1]); })), ("hs_y: a shift of 128", edited(|_, y| { y["onnx::Pow_85"]["v"][0] = Value::from(128); })), ("hs_y: a shift of -1", edited(|_, y| { y["onnx::Pow_85"]["v"][0] = Value::from(-1); })), ("hs_y: a bias that is not an integer", edited(|_, y| { y["conv1.bias"]["v"][0] = Value::from(1.5); })), ]; // Past acc32 and within acc64: block 3's sums exceed i32 with one weight of 2^30 over // inputs in [0, 127]. let big = edited(|_, y| { y["pointwise.weight"]["v"][0] = Value::from(1i64 << 30); }); let g = wai_ops::lower::jpegai_scale(&big.0.norm, &big.0.hs_y, &big.0.hs_uv).expect("the bundle lowers"); assert!(wai_ops::graph::load(&g.get("hs_y").unwrap().to_bytes(), Profile::Acc64).is_ok(), "it meets acc64"); assert!(wai_ops::graph::load(&g.get("hs_y").unwrap().to_bytes(), Profile::Acc32).is_err(), "and not acc32"); cases.push(("hs_y: a sum past i32 (acc64, not acc32)", big)); for (what, (b, parseable)) in &cases { let r = jpegai_ops::load_bytes(&b.norm, &b.hs_y, &b.hs_uv).err().unwrap_or_else(|| panic!("{what}: the load from bytes accepts it")); assert_eq!(r.code, wai_ops::Code::ModelMismatch, "{what}"); if *parseable { let (n, y, uv) = b.parsed(); let r = jpegai_ops::load(&n, &y, &uv).err().unwrap_or_else(|| panic!("{what}: the load from the parsed bundle accepts it")); assert_eq!(r.code, wai_ops::Code::ModelMismatch, "{what}"); } } assert!(jpegai_ops::load_bytes(&Bundle::synthetic().norm, &Bundle::synthetic().hs_y, &Bundle::synthetic().hs_uv).is_ok(), "the committed bundle loads"); eprintln!("{} load checks refuse model_mismatch from bytes and from the parsed bundle", cases.len()); } /// How the direct front half came out on one input: its integers, or a panic. fn direct(n: &NormTables, hs: &IntNet, gain: &[i64], p: jpegai_ops::Plane, soz: &[u8], lat: &[u8], grfs: &[bool]) -> Option { std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| jpegai_ops::legacy::plane(n, hs, gain, p, soz, lat, grfs))).ok() } /// The refusal parity of the entropy decodes on one bundle: each plane, with the SOZ /// payload cut at every length, then the plane's latent payload cut at every length, then /// each lengthened at its front by 16, 256 and 4096 bytes. /// Where the operator set decodes, the direct decode gives the same integers. Where the /// operator set refuses, the refusal is `stream_exhausted`, and the direct decode /// panicked: in a build with overflow checks it panics at the read or the refill that /// E8 refuses; in one without, it may instead compute on from a window that wrapped, and /// those cases are counted. Returns (decoded alike, refused and panicked, refused where /// the unchecked build computed on). fn cut_streams(sp: &jpegai_ops::ScalePath, n: &NormTables, hs: [&IntNet; 2], data: &[u8]) -> (usize, usize, usize) { let mut tally = (0, 0, 0); for (k, p) in [jpegai_ops::Plane::Luma, jpegai_ops::Plane::Chroma].into_iter().enumerate() { let gain = if k == 0 { &n.gain_y } else { &n.gain_uv }; let (soz, lat, grfs) = segments(data, p); // Bytes put before a payload are the last its backward window reaches: a payload // so lengthened decodes further before it runs out. let filler = |n: usize, seed: u64| -> Vec { (0..n as u64).map(|i| ((i.wrapping_add(seed)).wrapping_mul(0x9E37_79B9_7F4A_7C15) >> 56) as u8).collect() }; let longer = |pre: Vec, p: &[u8]| -> Vec { let mut v = pre; v.extend_from_slice(p); v }; let cuts = (0..=soz.len()).map(|c| (soz[..c].to_vec(), lat.clone())) .chain((0..lat.len()).map(|c| (soz.clone(), lat[..c].to_vec()))) .chain([16usize, 256, 4096].into_iter().flat_map(|n| [ (longer(filler(n, 1), &soz), lat.clone()), (soz.clone(), longer(filler(n, 2), &lat)), (longer(filler(n, 3), &soz), longer(filler(n, 4), &lat)), ])) .collect::>(); for (s, l) in cuts { let ops = jpegai_ops::plane(sp, p, &s, &l, &grfs); match (ops, direct(n, hs[k], gain, p, &s, &l, &grfs)) { (Ok(a), Some(b)) => { assert!(a == b, "{p:?} ({}, {}): the decodes differ", s.len(), l.len()); tally.0 += 1 } (Err(r), None) => { assert_eq!(r.code, wai_ops::Code::StreamExhausted, "{p:?}"); tally.1 += 1 } (Err(r), Some(_)) => { assert_eq!(r.code, wai_ops::Code::StreamExhausted, "{p:?}"); if cfg!(debug_assertions) { panic!("{p:?} ({}, {}): refused, and the checked build's direct decode did not panic", s.len(), l.len()); } tally.2 += 1 } (Ok(_), None) => panic!("{p:?} ({}, {}): the direct decode panicked where the operator set decodes", s.len(), l.len()), } } } tally } /// Gate 3, the streams, on the committed bundle and the committed test codestream. The /// committed bundle's tables are not the ones the codestream was coded with, so the /// decodes run out early unless a payload is lengthened; the integers they give differ /// from the pinned bundle's, and both decodes compute them alike. #[test] fn a_stream_cut_short_is_refused_where_the_direct_decode_panics() { let b = Bundle::synthetic(); let (n, y, uv) = b.parsed(); let sp = jpegai_ops::load(&n, &y, &uv).unwrap(); let data = read(&wai().join("wai-rs/tests/vectors/jpegai/test256.bin")); let t = cut_streams(&sp, &n, [&y, &uv], &data); eprintln!("committed bundle: {} decoded alike, {} refused where the direct decode panicked, {} refused where the unchecked build computed on", t.0, t.1, t.2); assert!(t.0 > 0 && t.1 > 0, "both outcomes are reached"); } /// Gate 3, the streams, on the pinned bundle. #[test] fn a_stream_cut_short_is_refused_where_the_direct_decode_panics_pinned() { let Some(wh) = Whole::pinned("jpegai_ops::a_stream_cut_short_is_refused_where_the_direct_decode_panics_pinned") else { return; }; let sp = wai::jpegai_decode::load(&wh.weights, &wh.norm).unwrap(); let t = cut_streams(&sp, &wh.norm, [&wh.weights.luma.hyper_scale, &wh.weights.chroma.hyper_scale], &wh.data); eprintln!("jpegai/v1: {} decoded alike, {} refused where the direct decode panicked, {} refused where the unchecked build computed on", t.0, t.1, t.2); assert!(t.0 > 0 && t.1 > 0, "both outcomes are reached"); } /// Gate 3, the codestream: a codestream outside the decoder's configuration is refused /// through the operator set with the reason `check_scope` gives, before the bundle is /// read, as the direct decode refused it. #[test] fn a_codestream_out_of_scope_is_refused_with_its_reason() { let Some(wh) = Whole::pinned("jpegai_ops::a_codestream_out_of_scope_is_refused_with_its_reason") else { return; }; use wai::jpegai_bitstream::{find_segment, parse_container, Marker}; let pih = find_segment(&parse_container(&wh.data), Marker::Pih).unwrap().start_bit(); let set = |d: &mut [u8], pos: usize, n: usize, value: u64| { for i in 0..n { let bit = (value >> (n - 1 - i)) & 1; let (byte, sh) = ((pos + i) / 8, 7 - (pos + i) % 8); d[byte] = (d[byte] & !(1 << sh)) | ((bit as u8) << sh); } }; // model_id (4 bits at 85), the colour transform (2 bits at 83, the reserved 3), β // (12 bits at 90, past the encoder's range), the first synthesis transform (4 bits at // 12, the base-operation-point one). for (at, n, v) in [(85usize, 4usize, 1u64), (83, 2, 3), (90, 12, 2048 + 703), (12, 4, 1)] { let mut d = wh.data.clone(); set(&mut d, pih + at, n, v); let why = wai::jpegai_decode::check_scope(&d).unwrap_err(); assert!(wai::jpegai_decode::legacy::decode(&d, &wh.weights, &wh.norm).is_none(), "the direct decode refuses it"); assert_eq!(wai::jpegai_decode::decode(&d, &wh.weights, &wh.norm).err(), Some(wai::jpegai_decode::Refusal::Scope(why)), "through the operator set"); } // Out of scope and with a bundle that fails a load check: the scope is checked first. let mut d = wh.data.clone(); set(&mut d, pih + 85, 4, 1); let mut bad = NormTables { bounds: vec![256; 32], ..norm_tables(&json(&read(Path::new(&std::env::var("WAI_JPEGAI_NORM").unwrap())))) }; assert!(matches!(wai::jpegai_decode::decode(&d, &wh.weights, &bad), Err(wai::jpegai_decode::Refusal::Scope(_)))); bad.bounds = wh.norm.bounds.clone(); bad.gain_y.pop(); assert_eq!(wai::jpegai_decode::decode(&wh.data, &wh.weights, &bad).err(), Some(wai::jpegai_decode::Refusal::Ops(wai_ops::Code::ModelMismatch)), "in scope, the bundle's load check"); } /// Gate 5: under every kernel policy, the reference kernels and the fast ones at each /// thread count, both planes' integers are the same, on the committed bundle and, where /// it is supplied, the pinned one, with the test codestream's payloads whole and cut. #[test] fn every_kernel_policy_decodes_the_planes_alike() { use wai::int_kernels::{with_policy, Policy}; static POLICIES: std::sync::LazyLock> = std::sync::LazyLock::new(wai::int_kernels::equivalence_policies); let mut sets: Vec<(&str, jpegai_ops::ScalePath, Vec)> = vec![]; let b = Bundle::synthetic(); let (n, y, uv) = b.parsed(); sets.push(("committed", jpegai_ops::load(&n, &y, &uv).unwrap(), read(&wai().join("wai-rs/tests/vectors/jpegai/test256.bin")))); if let Some(wh) = Whole::pinned("jpegai_ops::every_kernel_policy_decodes_the_planes_alike") { sets.push(("jpegai/v1", wai::jpegai_decode::load(&wh.weights, &wh.norm).unwrap(), wh.data)); } let mut compared = 0; for (name, sp, data) in &sets { for p in [jpegai_ops::Plane::Luma, jpegai_ops::Plane::Chroma] { let (soz, lat, grfs) = segments(data, p); for (s, l) in [(soz.clone(), lat.clone()), (soz[..soz.len() / 2].to_vec(), lat.clone()), (soz.clone(), lat[..lat.len() / 2].to_vec())] { let want = with_policy(POLICIES[0], || jpegai_ops::plane(sp, p, &s, &l, &grfs)); for pol in &POLICIES[1..] { let got = with_policy(*pol, || jpegai_ops::plane(sp, p, &s, &l, &grfs)); assert!(got == want, "{name} {p:?}: {pol:?} differs from the interpreter alone"); compared += 1; } } } } eprintln!("{compared} plane decodes under the fast policies equal the reference policy's"); }