// 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");
}