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

// wai-conformance: int-kernels int-kernels-amd
//! `wai.neural.int_synth` through the integer operator set: its lowering `int_synth/1`
//! (`extensions/integer-operators.md` §13.1) and the certificate of every committed
//! set (§8.5, §14.3), as this crate loads them (`int_ops`).
//!
//! The lowered graphs and certificates are committed under `ops-conformance/certificates/`,
//! where `wai-verify` and the independent checker (`tools/wai_ops_check.py`) each
//! reproduce them byte for byte (`wai-ops.yml`). Here the lowering this crate runs must
//! write those same bytes, and the load a sink makes must give the verdict the
//! certificate records.
//!
//! The decode through the lowering (`int_synth::decode_from_spec`, the one a sink runs) must give, on
//! every committed vector, the bytes of the direct decode it replaced
//! (`int_synth::legacy`, behind the test-only feature `legacy_int`, which this file
//! needs) and the vector's committed digests. Every refusal of the direct decode is a
//! refusal of the decode through the lowering, with the same code: integer-payloads §13
//! gives one, `payload_malformed`.

use std::path::PathBuf;

use serde_json::Value;
use wai_ops::interp::Tensor;
use wai_ops::{Code, Profile, Site};

const CAP: &str = "wai.neural.int_synth";

fn wai() -> PathBuf {
    PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("..")
}

fn read(rel: &str) -> Vec<u8> {
    std::fs::read(wai().join(rel)).unwrap_or_else(|e| panic!("{rel}: {e}"))
}

fn json(rel: &str) -> Value {
    serde_json::from_slice(&read(rel)).unwrap_or_else(|e| panic!("{rel}: {e}"))
}

fn sha256(b: &[u8]) -> String {
    wai_ops::sha256::sha256(b).iter().map(|x| format!("{x:02x}")).collect()
}

/// The `int_synth` sets an index of `ops-conformance/certificates/` lists:
/// `(name, set file, set digest, lowered graph file)`, each relative to `wai/`.
fn indexed(index: &str) -> Vec<(String, String, String, String)> {
    json(index)["sets"]
        .as_array()
        .unwrap()
        .iter()
        .filter(|s| s["capability"] == CAP)
        .map(|s| {
            assert_eq!(s["lowering"], "int_synth/1");
            let roles = s["roles"].as_array().unwrap();
            assert_eq!(roles.len(), 1, "int_synth/1 lowers to one graph, synth");
            (
                s["name"].as_str().unwrap().to_string(),
                s["members"]["int_synth.json"].as_str().unwrap().to_string(),
                s["set_sha256"].as_str().unwrap().to_string(),
                roles[0].as_str().unwrap().to_string(),
            )
        })
        .collect()
}

/// The certificate of the set with digest `d`.
fn certificate(d: &str) -> Value {
    json(&format!("ops-conformance/certificates/{CAP}/{d}.json"))
}

/// Gate 1 of the migration: every committed `int_synth` set lowers to its committed
/// graph, the certificate beside it is that graph's, and loading the graph under the
/// capability's profile gives the certificate's verdict. The neural-conformance corpus
/// meets `acc64`; facade, gradient and bloom meet `acc32` too, and photo does not.
#[test]
fn every_committed_set_lowers_to_its_committed_graph_and_verdict() {
    let sets = indexed("ops-conformance/certificates/index.json");
    let mut names: Vec<&str> = sets.iter().map(|s| s.0.as_str()).collect();
    names.sort();
    assert_eq!(names, ["bloom", "facade", "gradient", "photo", "synth"], "the committed int_synth sets");
    for (name, set, digest, role) in &sets {
        let bytes = read(set);
        let lowered = wai_ops::lower::int_synth(&bytes).unwrap_or_else(|r| panic!("{name}: does not lower: {r:?}"));
        let roles: Vec<&str> = lowered.roles.iter().map(|(r, _)| *r).collect();
        assert_eq!(roles, ["synth"], "{name}");
        let g = lowered.get("synth").unwrap();
        // The verdict first: the load a sink makes, then the profile each set meets.
        let loaded = wai::int_ops::load(g).unwrap_or_else(|r| panic!("{name}: refused under acc64: {r:?}"));
        let acc32 = matches!(name.as_str(), "facade" | "gradient" | "bloom" | "synth");
        assert_eq!(loaded.meets(Profile::Acc32), acc32, "{name}: whether it meets acc32");
        // Then the committed bytes: the set, its lowered graph and its certificate.
        assert_eq!(&sha256(&bytes), digest, "{name}: the set's digest");
        let graph = g.to_bytes();
        assert!(graph == read(role), "{name}: the lowering differs from the committed {role}");
        let cert = certificate(digest);
        assert_eq!(cert["set_sha256"].as_str(), Some(digest.as_str()), "{name}");
        assert_eq!(cert["graphs"][0]["name"], "synth", "{name}");
        assert_eq!(cert["graphs"][0]["graph_sha256"].as_str(), Some(sha256(&graph).as_str()), "{name}: the certificate is the graph's");
        assert_eq!(cert["verdict"]["acc64"], true, "{name}");
        assert_eq!(Value::Bool(loaded.meets(Profile::Acc32)), cert["verdict"]["acc32"], "{name}: the certificate's verdict on acc32");
    }
    // The browser demo serves byte copies of the corpus's sets.
    for name in ["bloom", "facade", "gradient", "photo"] {
        let copy = sha256(&read(&format!("wai-web/demo/int_synth/{name}/int_synth.json")));
        assert!(sets.iter().any(|s| s.2 == copy), "wai-web/demo/int_synth/{name}: not a committed set");
    }
}

/// The two falsification sets of §14.3. `overflow_acc64`'s accumulator passes i64 by
/// one unit, so the capability's load refuses it at load step 8; `claim_acc32`'s
/// passes i32 by one unit, so it loads under `acc64` and does not meet `acc32`.
#[test]
fn the_falsification_sets_fail_the_profile_they_claim() {
    let sets = indexed("ops-conformance/certificates/falsification.json");
    let mut names: Vec<&str> = sets.iter().map(|s| s.0.as_str()).collect();
    names.sort();
    assert_eq!(names, ["claim_acc32", "overflow_acc64"]);
    for (name, set, digest, role) in &sets {
        let bytes = read(set);
        assert_eq!(&sha256(&bytes), digest, "{name}");
        let lowered = wai_ops::lower::int_synth(&bytes).unwrap();
        let g = lowered.get("synth").unwrap();
        assert!(g.to_bytes() == read(role), "{name}: the lowering differs from the committed {role}");
        let cert = certificate(digest);
        match name.as_str() {
            "overflow_acc64" => {
                assert_eq!(cert["verdict"]["acc64"], false);
                let r = wai::int_ops::load(g).expect_err("the overflow set is refused");
                assert_eq!((r.code, r.at), (Code::ModelMismatch, Site::Load(8)));
            }
            _ => {
                assert_eq!((cert["verdict"]["acc64"].as_bool(), cert["verdict"]["acc32"].as_bool()), (Some(true), Some(false)));
                assert!(!wai::int_ops::load(g).expect("acc64").meets(Profile::Acc32));
            }
        }
    }
}

/// Gate 2 of the migration: golden equality. Every committed `int_synth` vector decodes
/// through the lowering to the bytes the direct decode gives, and to every digest the
/// corpus commits for it: the SHA-256 sidecar, the BLAKE3 `rgb_hash` with its sizes,
/// and photo's reference RGB from the Python integer forward. The browser demo serves
/// byte copies of the corpus's payloads and digests.
#[test]
fn every_committed_vector_decodes_through_the_lowering_to_its_golden() {
    let hashes = json("neural-conformance/expected_hashes.json");
    let sidecar = json("neural-conformance/sha256.json");
    assert!(read("wai-web/demo/int_synth/expected_hashes.json") == read("neural-conformance/expected_hashes.json"));
    let mut n = 0;
    for (name, want) in hashes.as_object().unwrap() {
        let dir = format!("neural-conformance/vectors/{name}");
        let (payload, set) = (read(&format!("{dir}/latent.bin")), read(&format!("{dir}/int_synth.json")));
        assert!(payload == read(&format!("wai-web/demo/int_synth/{name}/latent.bin")), "{name}: the demo's copy");
        let ops = wai::int_synth::decode_from_spec(&payload, &set).unwrap_or_else(|e| panic!("{name}: {e}"));
        let direct = wai::int_synth::legacy::decode_from_spec(&payload, &set).unwrap_or_else(|e| panic!("{name}: {e}"));
        assert!(ops == direct, "{name}: the decode through the lowering differs from the direct decode");
        assert_eq!((ops.width as u64, ops.height as u64), (want["width"].as_u64().unwrap(), want["height"].as_u64().unwrap()), "{name}");
        assert_eq!(sha256(&ops.rgb), sidecar[name]["sha256"].as_str().unwrap(), "{name}: the SHA-256 sidecar");
        let mut h = blake3::Hasher::new();
        h.update(b"wai:int-synth-rgb\x01");
        h.update(&ops.rgb);
        assert_eq!(h.finalize().to_hex().as_str(), want["rgb_hash"].as_str().unwrap(), "{name}: rgb_hash");
        if let Ok(r) = std::fs::read(wai().join(format!("{dir}/ref_rgb.bin"))) {
            assert!(ops.rgb == r, "{name}: the Python integer-forward reference");
        }
        n += 1;
    }
    assert_eq!(n, 4, "the four committed int_synth vectors");
}

/// A payload of `C` channels at side `2^s_log2` for an `H`×`W` output: the latents are
/// `fill`, and `extra` bytes are added (negative: removed) at the end.
fn payload(h: u32, w: u32, c: u16, s_log2: u8, fill: u8, extra: i64) -> Vec<u8> {
    let mut p = Vec::new();
    p.extend(h.to_le_bytes());
    p.extend(w.to_le_bytes());
    p.extend(c.to_le_bytes());
    p.extend([s_log2, 0]);
    let n = (c as usize) << (2 * s_log2 as usize);
    p.resize((12 + n as i64 + extra) as usize, fill);
    p
}

/// One layer of `int_synth.json`, every weight `w`.
fn layer(in_ch: u64, out_ch: u64, k: u64, stride: u64, pad: u64, w: i64) -> Value {
    serde_json::json!({
        "in_ch": in_ch, "out_ch": out_ch, "k": k, "stride": stride, "pad": pad,
        "weight": vec![w; (in_ch * out_ch * k * k) as usize], "bias": vec![1; out_ch as usize], "shift": 1, "relu": true,
    })
}

fn set(layers: Vec<Value>) -> Vec<u8> {
    serde_json::to_vec(&serde_json::json!({ "layers": layers })).unwrap()
}

/// The fixtures: one for each refusal the direct decode makes (integer-payloads §13),
/// each beside a payload and set it decodes. One layer, 1 → 3 channels, k 2, stride 2,
/// on a 1×1 latent: a 2×2 output.
fn refusal_fixtures() -> Vec<(&'static str, Vec<u8>, Vec<u8>)> {
    let good_set = set(vec![layer(1, 3, 2, 2, 0, 3)]);
    let good = payload(2, 2, 1, 0, 9, 0);
    let edit = |f: &dyn Fn(&mut Value)| {
        let mut v: Value = serde_json::from_slice(&good_set).unwrap();
        f(&mut v);
        serde_json::to_vec(&v).unwrap()
    };
    let p = |what, pay: Vec<u8>| (what, pay, good_set.clone());
    let q = |what, s: Vec<u8>| (what, good.clone(), s);
    vec![
        p("no payload", vec![]),
        p("a payload shorter than its header", good[..11].to_vec()),
        p("a latent short", payload(2, 2, 1, 0, 9, -1)),
        p("a byte past the latents", payload(2, 2, 1, 0, 9, 1)),
        p("a side of 2^64", { let mut v = good.clone(); v[10] = 64; v }),
        p("a side of 2^255", { let mut v = good.clone(); v[10] = 255; v }),
        p("C·S·S past the word", { let mut v = payload(2, 2, 1, 0, 9, 0); v[8..10].copy_from_slice(&u16::MAX.to_le_bytes()); v[10] = 31; v }),
        p("an output taller than H", payload(1, 2, 1, 0, 9, 0)),
        p("an output narrower than W", payload(2, 3, 1, 0, 9, 0)),
        p("H of 0", payload(0, 2, 1, 0, 9, 0)),
        p("first in_ch is not C", payload(2, 2, 2, 0, 9, 0)),
        q("a set that is not JSON", b"{".to_vec()),
        q("a set with no layers member", b"{}".to_vec()),
        q("layers that is not a list", br#"{"layers":{}}"#.to_vec()),
        q("a layer that is not an object", br#"{"layers":[1]}"#.to_vec()),
        q("a layer without k", edit(&|v| { v["layers"][0].as_object_mut().unwrap().remove("k"); })),
        q("a negative stride", edit(&|v| v["layers"][0]["stride"] = (-2).into())),
        q("a k with a fraction", edit(&|v| v["layers"][0]["k"] = 2.5.into())),
        q("a negative shift", edit(&|v| v["layers"][0]["shift"] = (-1).into())),
        q("a relu that is not a boolean", edit(&|v| v["layers"][0]["relu"] = 1.into())),
        q("a weight short", edit(&|v| { v["layers"][0]["weight"].as_array_mut().unwrap().pop(); })),
        q("a bias past out_ch", edit(&|v| v["layers"][0]["bias"].as_array_mut().unwrap().push(0.into()))),
        q("a weight past i32", edit(&|v| v["layers"][0]["weight"][0] = (1i64 << 31).into())),
        q("a bias below i32", edit(&|v| v["layers"][0]["bias"][0] = (-(1i64 << 31) - 1).into())),
        q("layers that do not chain", set(vec![layer(1, 2, 1, 1, 0, 1), layer(3, 3, 2, 2, 0, 1)])),
        q("four output channels", set(vec![layer(1, 4, 2, 2, 0, 3)])),
        q("an output side of 0", set(vec![layer(1, 3, 2, 2, 1, 3)])),
    ]
}

/// The code a sink reports for a decode of `wai.neural.int_synth` (`int_payload`, the
/// integer decode path the dispatch takes).
fn code(payload: &[u8], set: &[u8]) -> Result<(), &'static str> {
    let params = wai::int_payload::resolve(CAP, None, &[wai::pin::HeldBytes::File(set)], None).expect("an unpinned set resolves");
    wai::int_payload::decode_int(CAP, payload, &params).map(|_| ()).map_err(|e| e.code())
}

/// Gate 3 of the migration: refusal parity. Each fixture, refused by the direct decode,
/// is refused by the decode through the lowering, and a sink reports it as
/// `payload_malformed`, the one code integer-payloads §13 gives. Beside them, the
/// payload and set they are edited from decode both ways.
#[test]
fn every_refusal_of_the_direct_decode_is_refused_through_the_lowering() {
    let good = (payload(2, 2, 1, 0, 9, 0), set(vec![layer(1, 3, 2, 2, 0, 3)]));
    let a = wai::int_synth::decode_from_spec(&good.0, &good.1).expect("the fixtures' base decodes");
    assert_eq!(Ok(a.clone()), wai::int_synth::legacy::decode_from_spec(&good.0, &good.1));
    assert_eq!(a.rgb, [14u8; 12], "(9·3 + 1) >> 1 at every sample");
    assert_eq!(code(&good.0, &good.1), Ok(()));
    let fixtures = refusal_fixtures();
    for (what, pay, s) in &fixtures {
        let direct = wai::int_synth::legacy::decode_from_spec(pay, s);
        let ops = wai::int_synth::decode_from_spec(pay, s);
        assert!(direct.is_err(), "{what}: the direct decode refuses it: {direct:?}");
        assert!(ops.is_err(), "{what}: the decode through the lowering refuses it: {ops:?}");
        assert_eq!(code(pay, s), Err("payload_malformed"), "{what}");
    }
}

/// The lowered graph's one `check` range, its input's `[−128, 127]` (§13.0's driver
/// contract), at its ends and one past them: a payload's latents are signed bytes, so
/// only a graph-level decode reaches past them, and the graph refuses it (`field_range`,
/// at input 0).
#[test]
fn the_latent_range_refuses_one_past_each_end() {
    let lowered = wai_ops::lower::int_synth(&read("neural-conformance/vectors/gradient/int_synth.json")).unwrap();
    let g = wai::int_ops::load(lowered.get("synth").unwrap()).unwrap();
    let (c, s) = (4u32, 8u32);
    let x = |v: i64| {
        let mut d = vec![0i64; (c * s * s) as usize];
        d[37] = v;
        Tensor::new(c, s, s, d)
    };
    for v in [-128, 127] {
        assert!(wai::int_ops::run1(&g, x(v), &[]).is_ok(), "{v} is a latent");
    }
    for v in [-129, 128] {
        let r = wai::int_ops::run1(&g, x(v), &[]).expect_err("one past the latent range");
        assert_eq!((r.code, r.at), (Code::FieldRange, Site::Input(0)), "{v}");
    }
}

/// Gate 5 of the migration, the whole-decoder half: the decode through the lowering
/// under the reference kernels (the interpreter computes every operator) and under each
/// fast policy (the transposed convolutions on this crate's fast kernels, serial and
/// threaded) gives the same bytes on every committed vector, and on the fixtures' base
/// case. The per-operator half is `int_ops`' equivalence tests.
#[test]
fn every_kernel_policy_decodes_through_the_lowering_alike() {
    use wai::int_kernels::{with_policy, Policy};
    static POLICIES: std::sync::LazyLock<Vec<Policy>> = std::sync::LazyLock::new(wai::int_kernels::equivalence_policies);
    let mut cases: Vec<(String, Vec<u8>, Vec<u8>)> = json("neural-conformance/expected_hashes.json")
        .as_object()
        .unwrap()
        .keys()
        .map(|name| {
            let dir = format!("neural-conformance/vectors/{name}");
            (name.clone(), read(&format!("{dir}/latent.bin")), read(&format!("{dir}/int_synth.json")))
        })
        .collect();
    cases.push(("the fixtures' base".into(), payload(2, 2, 1, 0, 9, 0), set(vec![layer(1, 3, 2, 2, 0, 3)])));
    for (name, p, s) in &cases {
        let want = with_policy(POLICIES[0], || wai::int_synth::decode_from_spec(p, s)).unwrap_or_else(|e| panic!("{name}: {e}"));
        for pol in &POLICIES[1..] {
            let got = with_policy(*pol, || wai::int_synth::decode_from_spec(p, s));
            assert!(got.as_ref() == Ok(&want), "{name}: {pol:?} differs from the interpreter alone");
        }
    }
}

/// The committed `int_synth` sets and payloads: every directory under
/// `neural-conformance/vectors/`, the browser demo's `wai-web/demo/int_synth/` and
/// `ops-conformance/sets/` that holds an `int_synth.json`, with its `latent.bin` where
/// it has one. Relative to `wai/`.
fn committed_sets_and_payloads() -> Vec<(String, Option<String>)> {
    let mut v = Vec::new();
    for root in ["neural-conformance/vectors", "wai-web/demo/int_synth", "ops-conformance/sets"] {
        let mut dirs: Vec<String> = std::fs::read_dir(wai().join(root))
            .unwrap()
            .filter_map(|e| e.ok())
            .filter(|e| e.path().join("int_synth.json").exists())
            .map(|e| format!("{root}/{}", e.file_name().to_string_lossy()))
            .collect();
        dirs.sort();
        for d in dirs {
            let p = format!("{d}/latent.bin");
            let payload = wai().join(&p).exists().then_some(p);
            v.push((format!("{d}/int_synth.json"), payload));
        }
    }
    v
}

/// O12: the load rules the switch brings (integer-operators §15.3 item 1, and the two
/// the lowering makes that its text now states) affect no committed set or payload.
/// Every committed set has layers, each with `out_ch`, `k` and `stride` at least 1,
/// fields that fit 32 bits and a shift of at most 63 (so none from 64 to 127 changes
/// meaning either), and lowers and loads under `acc64`; the one exception is
/// `overflow_acc64`, the falsification set made to fail its certificate. Every
/// committed payload has a reserved byte of 0, `C ≥ 1` and `C·S·S ≤ 2^26`, is exactly
/// its length, and decodes through the lowering against its set.
#[test]
fn the_new_load_rules_affect_no_committed_set_or_payload() {
    let all = committed_sets_and_payloads();
    let (mut sets, mut payloads) = (0, 0);
    for (set_path, payload_path) in &all {
        let bytes = read(set_path);
        let v: Value = serde_json::from_slice(&bytes).unwrap();
        let layers = v["layers"].as_array().unwrap();
        assert!(!layers.is_empty(), "{set_path}: no layers");
        for (i, l) in layers.iter().enumerate() {
            let f = |k: &str, d: u64| l.get(k).map_or(Some(d), Value::as_u64).unwrap_or_else(|| panic!("{set_path} layer {i}: {k}"));
            for k in ["in_ch", "out_ch", "k", "stride", "pad"] {
                assert!(f(k, u64::MAX) <= u32::MAX as u64, "{set_path} layer {i}: {k} past 32 bits");
            }
            assert!(f("out_pad", 0) <= u32::MAX as u64, "{set_path} layer {i}: out_pad");
            assert!(f("out_ch", 0) >= 1 && f("k", 0) >= 1 && f("stride", 0) >= 1, "{set_path} layer {i}: out_ch, k or stride of 0");
            assert!(f("shift", 0) <= 63, "{set_path} layer {i}: a shift of 64 or more");
        }
        let lowered = wai_ops::lower::int_synth(&bytes).unwrap_or_else(|r| panic!("{set_path}: does not lower: {r:?}"));
        let loaded = wai::int_ops::load(lowered.get("synth").unwrap());
        if set_path.contains("/overflow_acc64/") {
            assert!(loaded.is_err(), "{set_path}: the falsification set fails acc64");
        } else {
            assert!(loaded.is_ok(), "{set_path}: does not load under acc64: {loaded:?}");
        }
        sets += 1;
        let Some(p) = payload_path else { continue };
        let b = read(p);
        let c = u16::from_le_bytes([b[8], b[9]]) as u64;
        let n = c << (2 * b[10] as u64);
        assert_eq!(b[11], 0, "{p}: the reserved byte");
        assert!(c >= 1 && n <= 1 << 26 && b.len() as u64 == 12 + n, "{p}: C, the cap and the length");
        assert!(wai::int_synth::decode_from_spec(&b, &bytes).is_ok(), "{p}: decodes through the lowering");
        payloads += 1;
    }
    assert!(sets >= 11 && payloads >= 8, "the committed int_synth sets and payloads, at least those of this revision: {all:?}");
}