e6453a9566
Gate 7 made these seven tests require `legacy_int` in Cargo.toml, but their files were still gated on `neural_int` or `neural` alone, which `check-tests-in-ci.py` reports: a file's `#![cfg]` must name what its `[[test]]` requires. Each file is now gated on `all(..., feature = "legacy_int")`, as `jpegai_ops.rs` is. Three test comments named the design's appendix and a plan item. They now name the text's sections (§11's `role_profiles`, §13.6.2's table). Co-Authored-By: Claude Opus 5.5 <[email protected]>
465 lines
28 KiB
Rust
465 lines
28 KiB
Rust
//! `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 <dir> (pinned by tools/jpegai-e2e-v1.sha256)";
|
||
|
||
fn wai() -> PathBuf {
|
||
PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("..")
|
||
}
|
||
|
||
fn read(p: &Path) -> Vec<u8> {
|
||
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<i64> {
|
||
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<u8>,
|
||
hs_y: Vec<u8>,
|
||
hs_uv: Vec<u8>,
|
||
}
|
||
|
||
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<u8>)> = 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<String, &[u8]> =
|
||
[("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::<String>(), 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<u8>)> = names.iter().map(|n| (format!("weights/{n}.json"), read(&dir.join(format!("{n}.json"))))).collect();
|
||
let mut members: BTreeMap<String, &[u8]> = 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::<String>(), 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<u8>,
|
||
weights: wai::jpegai_decode::DecoderWeights,
|
||
norm: NormTables,
|
||
}
|
||
|
||
impl Whole {
|
||
fn pinned(test: &str) -> Option<Whole> {
|
||
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<u8>, Vec<u8>, Vec<bool>) {
|
||
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::<Vec<u8>>()) };
|
||
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<Value> { 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<jpegai_ops::PlaneInts> {
|
||
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<u8> { (0..n as u64).map(|i| ((i.wrapping_add(seed)).wrapping_mul(0x9E37_79B9_7F4A_7C15) >> 56) as u8).collect() };
|
||
let longer = |pre: Vec<u8>, p: &[u8]| -> Vec<u8> { 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::<Vec<_>>();
|
||
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:?}");
|
||
assert!(!cfg!(debug_assertions), "{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), β (12 bits at 90).
|
||
for (at, n, v) in [(85usize, 4usize, 1u64), (83, 2, 0), (90, 12, 2048 + 7)] {
|
||
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};
|
||
const POLICIES: [Policy; 5] = [Policy::reference(), Policy::fast(1), Policy::fast(2), Policy::fast(5), Policy::fast(0)];
|
||
let mut sets: Vec<(&str, jpegai_ops::ScalePath, Vec<u8>)> = 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");
|
||
}
|