diff --git a/.github/workflows/wai-byte-equality.yml b/.github/workflows/wai-byte-equality.yml index feaad6d1e..5f7600932 100644 --- a/.github/workflows/wai-byte-equality.yml +++ b/.github/workflows/wai-byte-equality.yml @@ -14,7 +14,10 @@ name: wai cross-architecture byte-equality # (`wai.video.int_hyper`) whose I-frame comes out of a learned conditional codec, # and the MLIC++-topology codec (`wai.neural.int_mlicpp`), whose entropy network runs # before every checkerboard pass and must give every platform the encoder's scale -# bucket for every element. The binary payloads SPEC §5 pins (WIV1, WIH1) are parsed +# bucket for every element. The learned integer audio and Gaussian-splat attribute codecs +# (`wai.audio.int_codec`, `wai.splat.int_codec`) decode their committed WIA1 and WIS1 +# payloads to their PCM and attribute digests. The binary payloads SPEC §5 pins (WIV1, +# WIH1, WIA1, WIS1) are parsed # with the wire parser a sink runs and decoded to the same digests, and the staged-delivery # residual layers (WIR1) of the staged corpus are applied stage by stage to pinned digests. # Video is the strictest leg — inter-frame decode is a @@ -61,6 +64,8 @@ on: - 'wai/wai-rs/src/int_transform.rs' - 'wai/wai-rs/src/int_hyper_synth.rs' - 'wai/wai-rs/src/int_mlicpp.rs' + - 'wai/wai-rs/src/int_audio.rs' + - 'wai/wai-rs/src/int_splat.rs' - 'wai/wai-rs/src/int_payload/wire.rs' - 'wai/wai-rs/src/staged/predict.rs' - 'wai/wai-rs/src/staged/wir1.rs' @@ -76,6 +81,8 @@ on: - 'wai/wai-rs/tests/vectors/int_video/**' - 'wai/wai-rs/tests/vectors/int_hyper/**' - 'wai/wai-rs/tests/vectors/int_mlicpp/**' + - 'wai/wai-rs/tests/vectors/int_audio/**' + - 'wai/wai-rs/tests/vectors/int_splat/**' - '.github/workflows/wai-byte-equality.yml' push: branches: @@ -94,6 +101,8 @@ on: - 'wai/wai-rs/src/int_transform.rs' - 'wai/wai-rs/src/int_hyper_synth.rs' - 'wai/wai-rs/src/int_mlicpp.rs' + - 'wai/wai-rs/src/int_audio.rs' + - 'wai/wai-rs/src/int_splat.rs' - 'wai/wai-rs/src/int_payload/wire.rs' - 'wai/wai-rs/src/staged/predict.rs' - 'wai/wai-rs/src/staged/wir1.rs' @@ -109,6 +118,8 @@ on: - 'wai/wai-rs/tests/vectors/int_video/**' - 'wai/wai-rs/tests/vectors/int_hyper/**' - 'wai/wai-rs/tests/vectors/int_mlicpp/**' + - 'wai/wai-rs/tests/vectors/int_audio/**' + - 'wai/wai-rs/tests/vectors/int_splat/**' - '.github/workflows/wai-byte-equality.yml' jobs: diff --git a/.github/workflows/wai-conformance.yml b/.github/workflows/wai-conformance.yml index 080668132..b4660a500 100644 --- a/.github/workflows/wai-conformance.yml +++ b/.github/workflows/wai-conformance.yml @@ -249,7 +249,11 @@ jobs: # test-only) decodes random sets and payloads both ways: the same output or both # refused, except on the planned breaks of integer-operators §15.3, where the decode # through the operator set refuses. wai-ops-differential.yml runs it for an hour in - # each profile, nightly. + # each profile, nightly. tests/int_splat_ops.rs holds wai.splat.int_codec to its + # lowering whm1/1 the same way (integer-operators §13.2): its committed set lowers, + # from its bytes and from the model this crate parses, to its committed graph, and + # loads with its certificate's verdict. tests/int_audio_ops.rs does the same for + # wai.audio.int_codec and its lowering wam1/1 (§13.3). # tests/int_mlicpp.rs decodes the committed wai.neural.int_mlicpp vector against # its pinned hashes (image, every scale bucket, every symbol) under every kernel # policy and thread count, and checks that damaged payloads are refused rather than @@ -307,10 +311,10 @@ jobs: # codestream. No step named any of them until scripts/check-tests-in-ci.py asked. - name: Fast integer kernels ≡ reference kernels (overflow checks on) working-directory: wai/wai-rs - run: cargo test --no-default-features --features neural_int,video_int,audio_int,world_int,splat_int,generate,score,staged,stage_measure,legacy_int --lib --test int_kernels_vectors --test int_kernels_geometry --test int_synth_ops --test int_ops_differential --test int_mlicpp --test int_mlicpp_preflight --test int_audio --test int_splat --test generate --test cli_generate_score --test int_payload --test int_payload_preflight --test int_payload_cost --test int_video_enc_memory --test cli_stage --test staged_resources --test stage_measure --test stage_measure_bounds --test int_video --test int_video_halfpel --test int_video_keyframe --test int_world --test int_world2 --test jpegai_bitstream_calibration + run: cargo test --no-default-features --features neural_int,video_int,audio_int,world_int,splat_int,generate,score,staged,stage_measure,legacy_int --lib --test int_kernels_vectors --test int_kernels_geometry --test int_synth_ops --test int_ops_differential --test int_splat_ops --test int_audio_ops --test int_mlicpp --test int_mlicpp_preflight --test int_audio --test int_splat --test generate --test cli_generate_score --test int_payload --test int_payload_preflight --test int_payload_cost --test int_video_enc_memory --test cli_stage --test staged_resources --test stage_measure --test stage_measure_bounds --test int_video --test int_video_halfpel --test int_video_keyframe --test int_world --test int_world2 --test jpegai_bitstream_calibration - name: Fast integer kernels ≡ reference kernels (release, overflow checks off) working-directory: wai/wai-rs - run: cargo test --release --no-default-features --features neural_int,video_int,audio_int,world_int,splat_int,generate,score,staged,stage_measure,legacy_int --lib --test int_kernels_vectors --test int_kernels_geometry --test int_synth_ops --test int_ops_differential --test int_mlicpp --test int_mlicpp_preflight --test int_audio --test int_splat --test generate --test cli_generate_score --test int_payload --test int_payload_preflight --test int_payload_cost --test int_video_enc_memory --test cli_stage --test staged_resources --test stage_measure --test stage_measure_bounds --test manifest_parse_memory --test manifest_reader_equivalence --test int_video --test int_video_halfpel --test int_video_keyframe --test int_world --test int_world2 --test jpegai_bitstream_calibration + run: cargo test --release --no-default-features --features neural_int,video_int,audio_int,world_int,splat_int,generate,score,staged,stage_measure,legacy_int --lib --test int_kernels_vectors --test int_kernels_geometry --test int_synth_ops --test int_ops_differential --test int_splat_ops --test int_audio_ops --test int_mlicpp --test int_mlicpp_preflight --test int_audio --test int_splat --test generate --test cli_generate_score --test int_payload --test int_payload_preflight --test int_payload_cost --test int_video_enc_memory --test cli_stage --test staged_resources --test stage_measure --test stage_measure_bounds --test manifest_parse_memory --test manifest_reader_equivalence --test int_video --test int_video_halfpel --test int_video_keyframe --test int_world --test int_world2 --test jpegai_bitstream_calibration # The integer decode path builds under each of its features alone, so a sink can take # one medium without the others, and for wasm32 (a browser sink); `int_payload` arms a # medium only under its feature, and none needs `neural`. So does the staged-measurement diff --git a/.github/workflows/wai-ops-differential.yml b/.github/workflows/wai-ops-differential.yml index 2ab257513..847660324 100644 --- a/.github/workflows/wai-ops-differential.yml +++ b/.github/workflows/wai-ops-differential.yml @@ -6,7 +6,8 @@ name: wai-ops differential (nightly) # build profile: with overflow checks on and with them off. Outside the planned breaks # of integer-operators §15.3 the two decode to the same bytes or both refuse, and the # direct decode never panics; under a planned break the decode through the operator set -# refuses (wai/wai-rs/tests/int_ops_differential.rs). Today: wai.neural.int_synth. +# refuses (wai/wai-rs/tests/int_ops_differential.rs). Today: wai.neural.int_synth, +# wai.splat.int_codec and wai.audio.int_codec. # # Each night starts from its own seed, the date, which the log prints; a failing case # names its seed, and `WAI_DIFF_SEED=` re-runs it. The same tests run for a fixed @@ -60,4 +61,4 @@ jobs: if [ "${{ matrix.profile }}" = release ]; then P="--release"; fi export WAI_DIFF_SEED="$(date -u +%Y%m%d)" echo "WAI_DIFF_SECS=$WAI_DIFF_SECS WAI_DIFF_SEED=$WAI_DIFF_SEED" - cargo test $P --locked --no-default-features --features legacy_int --test int_ops_differential -- --nocapture + cargo test $P --locked --no-default-features --features legacy_int,splat_int,audio_int --test int_ops_differential -- --nocapture diff --git a/wai/CHANGELOG.md b/wai/CHANGELOG.md index b088b8124..865be38e0 100644 --- a/wai/CHANGELOG.md +++ b/wai/CHANGELOG.md @@ -6,6 +6,73 @@ wire format may change and such changes are recorded here. ## v0.4.0 — current +### `wai.audio.int_codec` through the integer operator set + +- **The decode.** `int_audio::decode_typed` (so the `WIA1` dispatch, the staged sink and + the browser crate) plans the synthesis as + [integer-payloads](extensions/integer-payloads.md) §7 orders, lowers `model.wam` by + `wam1/1` ([integer-operators](extensions/integer-operators.md) §13.3, with its PCM + step), loads the graph under `acc64` with its bounds certificate, decodes the stream, + and runs the graph through the operator set's interpreter with `t_out` as its scalar + and the fast transposed-convolution kernel behind it. Its `check` ranges are the + checks the direct decode made before each layer, and the PCM step's range is the one + in which a sample's scaling fits `i64`, so the same clips are refused + (`field_range`). Every committed output is unchanged, byte for byte, and every + committed refusal fixture keeps its code. The direct decode is kept, test-only + (feature `legacy_int`), in the differential tests and their nightly run. +- **Planned breaks** (SPEC §8, integer-operators §15.3 item 3), as for + `wai.splat.int_codec`: new refusals of sets no committed corpus contains, each + `model_mismatch` before any stream is decoded, the PCM step's range among the `check` + ranges whose producer must meet them. A scan of every committed audio set and payload + finds none affected. +- **Its `ops` member in the registry** ([integer-operators](extensions/integer-operators.md) + §11). `registry/capabilities.json` gives `wai.audio.int_codec` the member + `{"set": "wai-ops/1", "profile": "acc64", "subset": "full", "lowering": "wam1/1", + "stage": {"rans64": {"table": "t4097", "escape": "lenient", "end": "canonical"}}}`. + SPEC §5's row cites the payload ([integer-payloads](extensions/integer-payloads.md) §5) + and the lowering (integer-operators §13.3), and SPEC §7 names the capability among + those that decode through the operator set. + +### `wai.splat.int_codec` through the integer operator set + +- **The decode.** `int_splat::decode_typed` (so the `WIS1` dispatch, the staged sink and + the browser crate) plans the networks as + [integer-payloads](extensions/integer-payloads.md) §7 orders, lowers `g_s.whm` by + `whm1/1` ([integer-operators](extensions/integer-operators.md) §13.2, with the RGB8 + step of §13.0), loads the graph under `acc64` with its bounds certificate, decodes the + stream, and runs the graph through the operator set's interpreter with the fast + transposed-convolution kernel behind it. Its `check` ranges are the checks the direct + decode made before each layer, so the same latents are refused (`field_range`). + Every committed output is unchanged, byte for byte, and every committed refusal + fixture keeps its code. The direct decode is kept, test-only (feature + `legacy_int`): `tests/int_ops_differential.rs` decodes random sets and clouds both + ways, every push and for an hour per build profile each night. +- **Planned breaks** (SPEC §8, integer-operators §15.3 item 3): new refusals of sets no + committed corpus contains, each `model_mismatch` before any stream is decoded. A layer + with a bias of magnitude `2^63 − 1` or more, or a `check` range its producer cannot + meet, refused every decode with `field_range` when the network ran; a network with no + layer of tag 0, 1 or 2, or with a header `shift` or `gf` of 128 or more that no layer + uses, decoded where the shapes allowed. A scan of every committed splat set and + payload finds none affected. +- **The interpreter's memory.** A run releases each tensor once its last reader has + run, takes its inputs by value, and holds a transposed convolution's uncropped plane + in `i64` (each cell is within the certified `convT.acc`). With its set loaded, a splat + decode through the operator set holds no more than the payload's declared working + bytes (integer-payloads §7, *Memory and work*); the declared cost of a `WIS1` payload + is unchanged. `int_payload::declared_cost` now refuses, with the decode's code, a set + whose lowering does not load. +- **The byte-equality matrix** (`wai-byte-equality.yml`, `byte-exact-conformance`) now + decodes the committed `WIA1` and `WIS1` payloads of `wai.audio.int_codec` and + `wai.splat.int_codec` on every leg, under the reference and the fast kernels, to their + pinned PCM and attribute digests; it compiled neither decoder before. +- **Its `ops` member in the registry** ([integer-operators](extensions/integer-operators.md) + §11). `registry/capabilities.json` gives `wai.splat.int_codec` the member + `{"set": "wai-ops/1", "profile": "acc64", "subset": "full", "lowering": "whm1/1", + "stage": {"rans64": {"table": "t4097", "escape": "lenient", "end": "canonical"}}}`. + SPEC §5's row cites the payload ([integer-payloads](extensions/integer-payloads.md) §6) + and the lowering (integer-operators §13.2), and SPEC §7 names the capability among + those that decode through the operator set. + ### wai-quantum 0.5.0: provenance, sampling members, device programs, and three more members Breaking: `Certificate` gains a public field, `provenance`, and `FabricClass` is now diff --git a/wai/SPEC.md b/wai/SPEC.md index 2a3762c10..6de3d6628 100644 --- a/wai/SPEC.md +++ b/wai/SPEC.md @@ -811,12 +811,12 @@ standard. See the linked extension for each. | `wai.splat.spz` | SPZ (an open-source Gaussian-splat format) | asset, content-addressed; presentation only | | `wai.audio.scene` | WSCN scene (objects + trajectories) | mixdown-equivalence · [spatial-audio](extensions/spatial-audio.md) | | `wai.audio.score` | WSCR container (note events) | music as instructions; sink synthesizes, mixdown-equivalence · [score](extensions/score.md) | -| `wai.audio.int_codec` | `WIA1` — a LEARNED 1-D transform codec: integer transposed-conv synthesis + integer IGDN over a factorized-rANS latent (`tools/wai_int_audio_export.py` writes the model and a JSON form a converter reads) | **sample-equivalence**: int16 PCM byte-identical on every machine (a portable PCM hash); the audio analog of `wai.neural.int_hyper`, where a float learned-audio decode desyncs on ~1e-5 round-off. Pure i64, no float/ONNX (`int_audio.rs`) | +| `wai.audio.int_codec` | `WIA1` — a LEARNED 1-D transform codec: integer transposed-conv synthesis + integer IGDN over a factorized-rANS latent (`tools/wai_int_audio_export.py` writes the model and a JSON form a converter reads) | **sample-equivalence**: int16 PCM byte-identical on every machine (a portable PCM hash); the audio analog of `wai.neural.int_hyper`, where a float learned-audio decode desyncs on ~1e-5 round-off. Pure i64, no float/ONNX (`int_audio.rs`). Payload: [integer-payloads](extensions/integer-payloads.md) §5; decoded through its lowering, [integer-operators](extensions/integer-operators.md) §13.3 (`wam1/1`) | | `wai.haptic.signal` | WHAP envelope (intensity/sharpness keyframes) | sample-equivalence · [haptics](extensions/haptics.md) | | `wai.feed.state` | WWLD, sim `wai.world.sim.feed0` | replay-equivalence · [state-feeds](extensions/state-feeds.md) | | `wai.avatar.pose` | WWLD, sim `wai.world.sim.skeletal0` | pose-equivalence · [avatars](extensions/avatars.md) | | `wai.splat.4d` | W4DS container (per-Gaussian keyframed tracks) | frame-state-equivalence · [volumetric](extensions/volumetric.md) | -| `wai.splat.int_codec` | `WIS1` — a LEARNED Gaussian-splat ATTRIBUTE codec: a learned transform over the `[14,S,S]` attribute image (pos/scale/rot/color/opacity), integer synthesis (int_hyper machinery, 14 channels) over a factorized-rANS latent (`int_splat.rs`; `tools/wai_int_splat_export.py` writes the model and a JSON form a converter reads) | **splat-equivalence**: the decoded attribute tensor is byte-identical on every machine (a portable cloud hash) — the int_hyper of 3D, where a float splat-attribute decode desyncs. The float rasterizer stays sink-supplied presentation; the *attributes* are the contract. Pure i64, no float/ONNX | +| `wai.splat.int_codec` | `WIS1` — a LEARNED Gaussian-splat ATTRIBUTE codec: a learned transform over the `[14,S,S]` attribute image (pos/scale/rot/color/opacity), integer synthesis (int_hyper machinery, 14 channels) over a factorized-rANS latent (`int_splat.rs`; `tools/wai_int_splat_export.py` writes the model and a JSON form a converter reads) | **splat-equivalence**: the decoded attribute tensor is byte-identical on every machine (a portable cloud hash) — the int_hyper of 3D, where a float splat-attribute decode desyncs. The float rasterizer stays sink-supplied presentation; the *attributes* are the contract. Pure i64, no float/ONNX. Payload: [integer-payloads](extensions/integer-payloads.md) §6; decoded through its lowering, [integer-operators](extensions/integer-operators.md) §13.2 (`whm1/1`) | | `wai.film.linear` | WFLM container (sealed world + camera track + frame schedule) | film-equivalence · [linear-film](extensions/linear-film.md) | | `wai.film.reel` | WREL container (sealed world + a sequence of shots = camera track + schedule each) | reel-equivalence (film-equivalence with cuts) · [linear-film](extensions/linear-film.md) | | `wai.ternary.tt` | WQTT container (balanced-ternary MPS cores — packed trits + integer per-column scales) | a **quantum-inspired** (tensor-network) transform ported from the ternary/MPS fabric with ONE change: the per-column scale is an **integer**, so the whole decode is `i64`/`i128` and the reconstruction is byte-identical on every machine — **reconstruction-equivalence** (a portable recon hash), receiptable. Scope: a **provenance** transform, measured at 13.3 dB SNR for 3.8× with six residual layers (`Lossy` tier; the extension's §5). Encoder-side float SVD only · [ternary-tt](extensions/ternary-tt.md) | @@ -1512,7 +1512,8 @@ registry as `codecs::capability_replicate_class(cap) -> ReplicateClass`: ([integer-operators](extensions/integer-operators.md) §11): the lowering its sets are lowered by, the profile every conforming set's certificate meets, and the modes of its stage operators. A set whose certificate does not meet that profile is - refused. In this revision: `wai.neural.int_synth`. + refused. In this revision: `wai.neural.int_synth`, `wai.splat.int_codec` and + `wai.audio.int_codec`. - **`NeuralEntropyExact`** — the entropy decode is integer by construction, so the decoded symbols are byte-identical on every machine; later stages run float. Registered at `entropy-consistency`. An encoder MUST NOT emit `intent = @@ -2399,6 +2400,14 @@ envelope that uses any string a later revision registers. convolution's uncropped output, past `2^26` elements); and a set whose certificate does not meet `acc64`. A shift from 64 to 127, which a build used to mask or panic on, now has its exact meaning. + - `wai.splat.int_codec`: a set refused at load by its lowering, before any stream is + decoded (`model_mismatch`): a layer with a bias of magnitude `2^63 − 1` or more, or + a `check` range its producer cannot meet, each of which refused every decode + (`field_range`) when the network ran; a network with no layer of tag 0, 1 or 2; and + a header `shift` or `gf` of 128 or more that no layer uses. + - `wai.audio.int_codec`: the same, for its synthesis, whose `check` ranges include + its PCM step's: a set whose last layer cannot give a sample inside that range, + which refused every decode at that step, is refused at load. --- diff --git a/wai/byte-exact-conformance/Cargo.toml b/wai/byte-exact-conformance/Cargo.toml index 99dae5607..8fb8813f7 100644 --- a/wai/byte-exact-conformance/Cargo.toml +++ b/wai/byte-exact-conformance/Cargo.toml @@ -32,3 +32,9 @@ wai-ops = { path = "../wai-ops" } # as a sink does, through the integer operator set; the crate's own tests compile the # direct decode's unit tests under `cfg(test)`. legacy_int = [] +# wai-rs's features for the splat and audio decoders (`int_splat`, `int_audio`), which +# this crate always includes: on by default, so the included sources' tests that need +# them run here too. +splat_int = [] +audio_int = [] +default = ["splat_int", "audio_int"] diff --git a/wai/byte-exact-conformance/src/int_codec_golden.rs b/wai/byte-exact-conformance/src/int_codec_golden.rs new file mode 100644 index 000000000..d385226af --- /dev/null +++ b/wai/byte-exact-conformance/src/int_codec_golden.rs @@ -0,0 +1,143 @@ +//! Cross-architecture byte-equality for the learned integer audio and Gaussian-splat +//! attribute codecs (`wai.audio.int_codec`, `wai.splat.int_codec`; +//! `extensions/integer-payloads.md` §5 and §6). +//! +//! The committed `WIA1` and `WIS1` payloads are parsed with the wire parser a sink runs +//! (`int_payload/wire.rs`) and decoded through the checked typed decoders, under the +//! scalar reference kernels and under the fast kernels a sink runs by default; the JSON +//! forms they are converted from decode to the same bytes. Each output must be the +//! digest pinned below, which is also the one its vector's `sha256.json` sidecar states +//! (and, for the PCM, the FNV-1a the export's Python integer reference computed), and the +//! payloads themselves are pinned by their SHA-256. One differing sample or attribute on +//! any leg fails the matrix. + +#![cfg(test)] + +use std::path::PathBuf; + +use serde_json::Value; +use sha2::{Digest, Sha256}; + +use crate::int_audio::{decode_typed as decode_audio, pcm_fnv1a, AudioJson, AudioModel, AudioStream}; +use crate::int_hyper_synth::IntModel; +use crate::int_kernels::{with_policy, Policy}; +use crate::int_payload_wire::{parse_wia1, parse_wis1}; +use crate::int_splat::{decode_typed as decode_splat, SplatJson, SplatStream}; + +const AUDIO_WIA1_SHA256: &str = "e4343c9468f6c7b652f911e24e273223ef14feaa7014df014f2d99e9f28e1b3b"; +const SPLAT_WIS1_SHA256: &str = "473ea0223840758b26ff7d540d861e3b2c1565dbc2bec68b0235b3f50f19d422"; +/// The PCM, as signed 16-bit little-endian samples. +const AUDIO_PCM_SHA256: &str = "4b0344c87c67fc30b7d977ae2077fc310822199598e5d2f9f8fc06314b3a2875"; +const AUDIO_PCM_FNV1A: u64 = 0xe5ff_66ad_1b86_117c; +/// The attribute tensor's bytes, plane after plane. +const SPLAT_ATTR_SHA256: &str = "f07679de87ff7c3d4fc4890b80919ef857fb85bd4d9e4cfa16105211d580427f"; + +fn vectors() -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../wai-rs/tests/vectors") +} +fn read(rel: &str) -> Vec { + let p = vectors().join(rel); + std::fs::read(&p).unwrap_or_else(|e| panic!("read {}: {e}", p.display())) +} +fn jload(rel: &str) -> Value { + serde_json::from_slice(&read(rel)).unwrap_or_else(|e| panic!("parse {rel}: {e}")) +} +fn hex(b: &[u8]) -> String { + b.iter().map(|x| format!("{x:02x}")).collect() +} +fn sha256(b: &[u8]) -> String { + hex(&Sha256::digest(b)) +} + +/// `f`, under the scalar reference kernels and under the default (fast) kernels: the two +/// must agree, and the reference's result is returned. +fn both_kernels(what: &str, f: impl Fn() -> T) -> T { + let reference = with_policy(Policy::reference(), &f); + let fast = f(); + assert_eq!(fast, reference, "{what}: the fast kernels and the reference kernels differ"); + reference +} + +/// `WIA1`: the committed payload, parsed and decoded, gives the pinned PCM; so does the +/// JSON bitstream it is converted from. +#[test] +fn audio_wia1_decodes_to_its_pcm_golden() { + let payload = read("int_audio/audio_bitstream.wia1"); + assert_eq!(sha256(&payload), AUDIO_WIA1_SHA256, "the committed WIA1 payload"); + let model = AudioModel::from_bytes(&read("int_audio/audio_model.bin")).expect("the WAM1 synthesis parses"); + let a = parse_wia1(&payload).expect("the WIA1 payload parses"); + let s = AudioStream { + afrac: a.afrac as u32, n_ch: a.n_ch as usize, lat_len: a.lat_len as usize, t_out: a.t_out as usize, + offset: a.offset, median_q: &a.median_q, cdfs: &a.cdfs, rans: a.rans, + }; + let pcm = both_kernels("WIA1", || decode_audio(&model, &s).expect("the WIA1 payload decodes")); + let s16le: Vec = pcm.iter().flat_map(|v| v.to_le_bytes()).collect(); + assert_eq!(sha256(&s16le), AUDIO_PCM_SHA256, "the PCM diverged"); + assert_eq!(pcm_fnv1a(&pcm), AUDIO_PCM_FNV1A, "the PCM's FNV-1a diverged"); + assert_eq!(jload("int_audio/sha256.json")["audio"]["sha256"].as_str(), Some(AUDIO_PCM_SHA256), "the sidecar"); + assert_eq!(jload("int_audio/audio.json")["pcm_fnv1a"].as_str(), Some(format!("{AUDIO_PCM_FNV1A:016x}").as_str())); + + let j = AudioJson::from_value(&jload("int_audio/audio_bitstream.json")).expect("the JSON form reads"); + let from_json = both_kernels("the audio JSON form", || decode_audio(&model, &j.stream()).expect("the JSON form decodes")); + assert_eq!(from_json, pcm, "the JSON form and the payload decode alike"); +} + +/// `WIS1`: the committed payload, parsed and decoded, gives the pinned attribute tensor; +/// so does the JSON bitstream it is converted from. +#[test] +fn splat_wis1_decodes_to_its_attribute_golden() { + let payload = read("int_splat/splat.wis1"); + assert_eq!(sha256(&payload), SPLAT_WIS1_SHA256, "the committed WIS1 payload"); + let gs = IntModel::from_bytes(&read("int_splat/splat_gs.bin")).expect("the WHM1 synthesis parses"); + let v = parse_wis1(&payload).expect("the WIS1 payload parses"); + let s = SplatStream { + afrac: v.afrac as u32, c: v.c as usize, sh: v.sh as usize, sw: v.sw as usize, s: v.s as usize, + n_attr: v.attr as usize, cdfs: &v.cdfs, offsets: &v.offsets, median_q: &v.median_q, rans: v.rans, + }; + let attr = both_kernels("WIS1", || decode_splat(&gs, &s).expect("the WIS1 payload decodes").attr); + assert_eq!(attr.len(), v.attr as usize * v.s as usize * v.s as usize, "the attribute tensor's size"); + assert_eq!(sha256(&attr), SPLAT_ATTR_SHA256, "the attribute tensor diverged"); + assert_eq!(jload("int_splat/sha256.json")["splat"]["sha256"].as_str(), Some(SPLAT_ATTR_SHA256), "the sidecar"); + + let bs = jload("int_splat/splat.json"); + let j = SplatJson::from_value(&bs).expect("the JSON form reads"); + let from_json = both_kernels("the splat JSON form", || decode_splat(&gs, &j.stream()).expect("the JSON form decodes").attr); + assert_eq!(from_json, attr, "the JSON form and the payload decode alike"); + let fnv = crate::int_hyper_synth::fnv1a(&attr); + assert_eq!(Some(format!("{fnv:016x}").as_str()), bs["attr_hash"].as_str(), "the export's attribute hash"); +} + +/// `WIA1` through the synthesis's lowering `wam1/1` (`extensions/integer-operators.md` +/// §13.3, with the PCM step): loaded under `acc64` and run by the integer operator set's +/// interpreter, with the fast kernels behind it, it gives the same pinned PCM on every +/// leg. +#[test] +fn audio_through_its_lowering_is_golden() { + let payload = read("int_audio/audio_bitstream.wia1"); + let model = AudioModel::from_bytes(&read("int_audio/audio_model.bin")).expect("the WAM1 synthesis parses"); + let a = parse_wia1(&payload).expect("the WIA1 payload parses"); + let s = AudioStream { + afrac: a.afrac as u32, n_ch: a.n_ch as usize, lat_len: a.lat_len as usize, t_out: a.t_out as usize, + offset: a.offset, median_q: &a.median_q, cdfs: &a.cdfs, rans: a.rans, + }; + let pcm = both_kernels("WIA1 through the lowering", || crate::int_audio::ops::decode_typed(&model, &s).expect("decodes through the lowering")); + let s16le: Vec = pcm.iter().flat_map(|v| v.to_le_bytes()).collect(); + assert_eq!(sha256(&s16le), AUDIO_PCM_SHA256, "the PCM through the lowering diverged"); + assert_eq!(pcm_fnv1a(&pcm), AUDIO_PCM_FNV1A); +} + +/// `WIS1` through `g_s`'s lowering `whm1/1` (`extensions/integer-operators.md` §13.2): +/// loaded under `acc64` and run by the integer operator set's interpreter, with the fast +/// kernels behind it, it gives the same pinned attribute tensor on every leg. +#[test] +fn splat_through_its_lowering_is_golden() { + let payload = read("int_splat/splat.wis1"); + let gs = IntModel::from_bytes(&read("int_splat/splat_gs.bin")).expect("the WHM1 synthesis parses"); + let v = parse_wis1(&payload).expect("the WIS1 payload parses"); + let s = SplatStream { + afrac: v.afrac as u32, c: v.c as usize, sh: v.sh as usize, sw: v.sw as usize, s: v.s as usize, + n_attr: v.attr as usize, cdfs: &v.cdfs, offsets: &v.offsets, median_q: &v.median_q, rans: v.rans, + }; + let attr = both_kernels("WIS1 through the lowering", || crate::int_splat::ops::decode_typed(&gs, &s).expect("decodes through the lowering").attr); + assert_eq!(sha256(&attr), SPLAT_ATTR_SHA256, "the attribute tensor through the lowering diverged"); +} diff --git a/wai/byte-exact-conformance/src/lib.rs b/wai/byte-exact-conformance/src/lib.rs index f66085f04..91735d4be 100644 --- a/wai/byte-exact-conformance/src/lib.rs +++ b/wai/byte-exact-conformance/src/lib.rs @@ -75,11 +75,22 @@ pub mod int_hyper_synth; #[path = "../../wai-rs/src/int_mlicpp.rs"] pub mod int_mlicpp; +// The learned integer audio and Gaussian-splat attribute codecs (`wai.audio.int_codec`, +// `wai.splat.int_codec`): one-dimensional and two-dimensional integer transposed +// convolutions and inverse GDN, over the networks, kernels and rANS decoder already +// included here, so the matrix checks their PCM and attribute bytes on every platform. +#[path = "../../wai-rs/src/int_audio.rs"] +pub mod int_audio; +#[path = "../../wai-rs/src/int_splat.rs"] +pub mod int_splat; + // The wire parser of the integer payloads SPEC §5 pins (WIV1, WIH1, WIA1, WIS1). // Self-contained — std's Vec only, nothing from the crate — so the matrix parses the // committed binary payloads with the parser a sink runs. #[path = "../../wai-rs/src/int_payload/wire.rs"] pub mod int_payload_wire; +// The name wai-rs gives the same module (`int_wire`), which int_audio's own tests use. +pub use int_payload_wire as int_wire; // The staged-delivery residual (`WIR1`, extensions/staged-delivery.md Appendix A): its // parser, predictors and arithmetic step. Self-contained beside the rANS coder and the @@ -93,4 +104,5 @@ mod int_video_golden; mod neural_keyframe_golden; mod jpegai_golden; mod int_mlicpp_golden; +mod int_codec_golden; mod staged_golden; diff --git a/wai/extensions/integer-operators.md b/wai/extensions/integer-operators.md index 5bbf742de..7a8738bbc 100644 --- a/wai/extensions/integer-operators.md +++ b/wai/extensions/integer-operators.md @@ -14,7 +14,7 @@ > > **What a reader can build from this text.** With the payload extensions it cites, an implementation of this text decodes `wai.video.int_motion`, `wai.audio.int_codec`, `wai.splat.int_codec`, `wai.neural.int_synth`, `wai.neural.int_hyper`, `wai.video.int_hyper`, `wai.neural.int_mlicpp` and `wai.{image,video,audio,splat}.int_refine`, and reproduces the symbols and scale indices that `wai.image.jpegai` entropy-decodes (§13.6.1). The payload text four of these depend on, `entropy.json`, the `int_synth` payload and the `WIM1` payload with its `WMB1` bundle, is in [integer-payloads](integer-payloads.md), and §15.2 says where. The integer renderer of `wai.world.model` and the other integer networks of `wai.image.jpegai` are informative (§13.5, §13.6.2). > -> Reference implementation: the `wai-ops` crate (`wai/wai-ops`): parse, lowering, interpreter, stage operators and bounds. Tool: `wai-verify bounds`. Conformance: `wai/ops-conformance/`. A capability's reference decoder decodes through this operator set from the revision that registers the capability's `ops` member (§11); in this revision, `wai.neural.int_synth`. Until then, a capability's reference implementation is the decoder its payload extension names, and §15.3 lists where this text and that decoder differ. +> Reference implementation: the `wai-ops` crate (`wai/wai-ops`): parse, lowering, interpreter, stage operators and bounds. Tool: `wai-verify bounds`. Conformance: `wai/ops-conformance/`. A capability's reference decoder decodes through this operator set from the revision that registers the capability's `ops` member (§11); in this revision, `wai.neural.int_synth`, `wai.splat.int_codec` and `wai.audio.int_codec`. Until then, a capability's reference implementation is the decoder its payload extension names, and §15.3 lists where this text and that decoder differ. > > The keywords **MUST**, **MUST NOT**, **SHOULD** and **MAY** are as defined in RFC 2119 and RFC 8174. @@ -1528,7 +1528,7 @@ No committed output changes, and no committed refusal fixture changes its code. - A shift from 64 to 127 now has its exact meaning. - The caps apply. 2. **`wai.world.model`'s integer renderer** refuses values outside its lowered `check` ranges, and a set whose `afrac` is outside `1..=30` or differs between its two networks. -3. **`WHM1` and `WAM1` layers.** A layer whose predicate fails at 0 (a bias of magnitude `2^63 − 1` or more) is refused at load, not at decode. So is a set in which a `check` range cannot hold any value its producer can give (§8.1), which today refuses every decode with `field_range`. +3. **`WHM1` and `WAM1` layers.** A layer whose predicate fails at 0 (a bias of magnitude `2^63 − 1` or more) is refused at load, not at decode. So is a set in which a `check` range cannot hold any value its producer can give (§8.1), which today refuses every decode with `field_range`. Each is refused with `model_mismatch`, before any stream is decoded, where the direct decoder refused with `field_range`, or with a stream's code when the stream failed first. A network with no layer of tag 0, 1 or 2, and one whose header `shift` or `gf` is 128 or more where no layer uses it, are refused too (§13.2); the direct decoder decoded each where the shapes allowed. These take effect for `wai.splat.int_codec` and `wai.audio.int_codec` in this revision; for `wai.audio.int_codec` the `check` ranges include the PCM step's (§13.3), so a set whose last layer cannot give a sample inside it is refused at load, where every decode was refused at that step. 4. **`wai.image.jpegai`.** - Its exact decoder refuses a malformed latent or hyper-latent stream (`stream_exhausted`) where it panics today. - It refuses a bundle that fails §13.6.1's load checks (`model_mismatch`), including a bound above 255 or a shift outside `0..=127`, which today's loader wraps. diff --git a/wai/extensions/integer-payloads.md b/wai/extensions/integer-payloads.md index ff6c367c2..c476705f3 100644 --- a/wai/extensions/integer-payloads.md +++ b/wai/extensions/integer-payloads.md @@ -83,6 +83,9 @@ escape of 4-bit groups): hashes. A stream that codes no symbol (the residual stream of a one-frame `WIV1` with keyframe 1) is the state `2^31` alone: the words `2^31` and `0`. +The decode these rules define is [integer-operators](integer-operators.md) E1, +`rans64_decode`, with the modes `t4097`, `lenient` and `canonical`. + **One clip, several payloads.** These rules fix how a payload is read, not the only payload of a decoded output: tables that differ in entries the stream never selects (a symbol the clip does not use, say), or a different escape layout, code @@ -239,9 +242,13 @@ The header's field rules are checked in the order `afrac`, then `sr`, `n_ch`, 1. For each channel `k` and each of its `lat_len` elements, decode a symbol `s` against channel `k`'s table at `offset`; the element is `s·2^afrac + median_k`. 2. The synthesis (one-dimensional integer transposed convolutions and integer - inverse GDN) on `[n_ch, lat_len]` gives one channel of `t ≥ t_out` samples. + inverse GDN) on `[n_ch, lat_len]` gives one channel of `t ≥ t_out` samples. It + decodes through its lowering `wam1/1` ([integer-operators](integer-operators.md) + §13.3), with `t_out` as its scalar. 3. The first `t_out` samples `v` become `(v·32767 + 2^(afrac−1)) >> afrac`, clamped - to `[−32768, 32767]`. + to `[−32768, 32767]`: the PCM step that ends the lowered graph (integer-operators + §13.3). A sample for which `v·32767 + 2^(afrac−1)` leaves `i64` is refused + (`field_range`). The output is `t_out` samples of `i16` PCM, one channel, at `sr` hertz: the buffer of the capability's sample-equivalence digest (signed 16-bit little-endian samples). @@ -275,8 +282,11 @@ The header's field rules are checked in the order `afrac`, the dimensions, 1. For each channel `k` and each of its `sh·sw` elements, decode a symbol `s` against channel `k`'s table and offset; the element is `s·2^afrac + median_k`. 2. `g_s` on `[c, sh, sw]` gives `[attr, s, s]` (`geometry_mismatch` otherwise). + `g_s` decodes through its lowering `whm1/1` + ([integer-operators](integer-operators.md) §13.2). 3. Each value `v` is clamped to `[0, 2^afrac]` and becomes - `(v·255 + 2^(afrac−1)) >> afrac`, a byte. + `(v·255 + 2^(afrac−1)) >> afrac`, a byte: the RGB8 step that ends the lowered + graph (integer-operators §13.0). The output, plane after plane, is the buffer of the capability's splat-equivalence digest. The first `n_gaussians` cells of the planes, in raster order, hold one @@ -361,7 +371,10 @@ caps; a sink's policy may be stricter. hyper-latent channels than the tables; - each layer's parameters have the lengths the running channel count needs, and its geometry leaves an output (*Planning the networks*); -- the audio synthesis gives one channel; the conditional codec's `g_s` gives three. +- the audio synthesis gives one channel; the conditional codec's `g_s` gives three; +- a network that decodes through its lowering lowers, and its graph loads under + `acc64` ([integer-operators](integer-operators.md) §13.2 and §13.3: for `WIS1`, + `g_s`; for `WIA1`, the synthesis). **The shapes the payload declares must be the ones the decode gives** (`geometry_mismatch`): `h_s`'s output and `[2·m, sh, sw]`; a `WIV1` keyframe and @@ -384,8 +397,9 @@ order and refuses with the first that fails: with keyframe 1, its keyframe's as for `WIH1`, then the keyframe's size against `h × w` (`geometry_mismatch`); for `WIA1`, one channel (`model_mismatch`), then at least `t_out` samples, then a latent no - longer than that needs (both `geometry_mismatch`); for `WIS1`, `g_s`'s shape - (`geometry_mismatch`); + longer than that needs (both `geometry_mismatch`), then the synthesis lowered and + loaded (`model_mismatch`); for `WIS1`, `g_s`'s shape + (`geometry_mismatch`), then `g_s` lowered and loaded (`model_mismatch`); 5. each stream in decode order (`WIH1`: the hyper-latent's, then, after `h_s` runs, the latent's; `WIV1` with keyframe 1: the keyframe's streams, then the residual stream), symbol by symbol: a @@ -403,7 +417,11 @@ network layer runs, the decode checks on the values present that no product, par sum, requantised value or inverse-GDN term can leave the type its arithmetic is formed in. A stream whose values fail either check is refused (`field_range`). So a decode with overflow checks and one without compute the same integers, and neither -panics. +panics. A network that decodes through its lowering +([integer-operators](integer-operators.md) §13; for `WIS1`, §13.2; for `WIA1`, §13.3) +makes the second check as the `check` ranges of its graph: a value outside one is +refused (`field_range`), and the refusals are the same. So does `WIA1`'s PCM step +(§5 step 3), whose range is its graph's last `check` range. ## 8. Refusal codes @@ -581,8 +599,8 @@ was made. `wai_int_payload_vectors gen` and checked by `wai_int_payload_vectors verify`; the same corpus checked by `int_payload_verify.py`, which shares no code with wai-rs and decodes the `WIV1` keyframe-0 payloads in full; and the cross-architecture goldens of - `byte-exact-conformance`, which parse and decode the `WIV1` and `WIH1` payloads on - every platform of its matrix. + `byte-exact-conformance`, which parse and decode the `WIV1`, `WIH1`, `WIA1` and `WIS1` + payloads on every platform of its matrix. ## 13. The `wai.neural.int_synth` payload and set diff --git a/wai/int-payload-conformance/README.md b/wai/int-payload-conformance/README.md index 70879ef72..cc265de1a 100644 --- a/wai/int-payload-conformance/README.md +++ b/wai/int-payload-conformance/README.md @@ -75,8 +75,7 @@ compared: **Not decoded independently.** The Python verifier does not decode `WIH1`, `WIV1` with keyframe 1 (its keyframe is a `WIH1`), `WIA1` or `WIS1`: their decodes run integer networks it does not implement. Their outputs are checked by wai-rs against -the sidecars above, and the `WIV1` and `WIH1` payloads also by the -cross-architecture goldens in +the sidecars above, and also by the cross-architecture goldens in [`../byte-exact-conformance/`](../byte-exact-conformance/). For their `decode` fixtures the Python verifier checks that they parse; their codes are checked by wai-rs. diff --git a/wai/registry/capabilities.json b/wai/registry/capabilities.json index 8b2b6cb73..db15e3767 100644 --- a/wai/registry/capabilities.json +++ b/wai/registry/capabilities.json @@ -99,7 +99,20 @@ "sets": [], "spec_section": "§5 Interactive, temporal & computed content", "edition": null, - "registered": "2026-06-27" + "registered": "2026-06-27", + "ops": { + "set": "wai-ops/1", + "profile": "acc64", + "subset": "full", + "lowering": "wam1/1", + "stage": { + "rans64": { + "table": "t4097", + "escape": "lenient", + "end": "canonical" + } + } + } }, { "capability": "wai.audio.int_refine", @@ -1011,7 +1024,20 @@ "sets": [], "spec_section": "§5 Interactive, temporal & computed content", "edition": null, - "registered": "2026-06-28" + "registered": "2026-06-28", + "ops": { + "set": "wai-ops/1", + "profile": "acc64", + "subset": "full", + "lowering": "whm1/1", + "stage": { + "rans64": { + "table": "t4097", + "escape": "lenient", + "end": "canonical" + } + } + } }, { "capability": "wai.splat.int_refine", diff --git a/wai/tools/wai_int_audio_export.py b/wai/tools/wai_int_audio_export.py index ef7085cbd..2b1592bd3 100644 --- a/wai/tools/wai_int_audio_export.py +++ b/wai/tools/wai_int_audio_export.py @@ -297,7 +297,8 @@ def build_codec(C): # ---- integer reference decode (numpy; mirrors the Rust i64 arithmetic) ------- def conv_t_1d_int(xi, wi, mult, bias, stride, pad, opad): """Integer 1-D transposed conv + per-out-channel (mult>>SHIFT)+bias — byte-for- - byte what Rust int_audio::conv_t_1d does. xi [cin,T]; wi [cin,cout,k].""" + byte what the Rust decoder's transposed convolution does (integer-operators §13.3, + wam1/1). xi [cin,T]; wi [cin,cout,k].""" import numpy as np cin, cout, k = wi.shape; Tin = xi.shape[1] To = (Tin - 1) * stride + k + opad - 2 * pad diff --git a/wai/wai-ops/src/interp.rs b/wai/wai-ops/src/interp.rs index a8d5c9f91..67424778c 100644 --- a/wai/wai-ops/src/interp.rs +++ b/wai/wai-ops/src/interp.rs @@ -1,7 +1,9 @@ //! The reference interpreter of the graph operators (§5). //! //! Tensors hold `i64`. Every intermediate is formed in `i128`, which holds every -//! value a graph certified for `acc64` forms (§7.1), and every step is still checked: +//! value a graph certified for `acc64` forms (§7.1), except `conv_transpose`'s uncropped +//! plane, whose cells are `acc`-class partial sums and so are held in `i64` (see +//! `k_conv_t`). Every step is still checked: //! a graph whose arithmetic would leave `i128`, or whose producer result would leave //! `i64`, is one that load step 8 refuses, so [`run`], which takes only a loaded graph, //! never meets one. Were it to, the decode would be refused (`model_mismatch`) at that @@ -189,40 +191,46 @@ pub fn rq(rq: &Rq, o: usize, v: i128) -> Option { /// from `seq` (§10.3). `scalars` holds every scalar's value. Returns the graph's /// outputs, in output-table order. A refused decode leaves `seq` as it was. pub fn run(g: &Loaded, inputs: &[Tensor], scalars: &[i64], seq: &mut Sequence) -> Res> { - run_inner(g, inputs, scalars, seq, &mut Quiet, &Reference) + run_inner(g, inputs.to_vec(), scalars, seq, &mut Quiet, &Reference) } -/// [`run`], with the host's kernels `k` asked first for each operator ([`Kernels`]). -/// The outputs and refusals are [`run`]'s, provided `k` keeps its contract. -pub fn run_with(g: &Loaded, inputs: &[Tensor], scalars: &[i64], seq: &mut Sequence, k: &K) -> Res> { +/// [`run`], with the host's kernels `k` asked first for each operator ([`Kernels`]), +/// taking the data inputs by value, so that a decode holds no second copy of them. The +/// outputs and refusals are [`run`]'s, provided `k` keeps its contract. +pub fn run_with(g: &Loaded, inputs: Vec, scalars: &[i64], seq: &mut Sequence, k: &K) -> Res> { run_inner(g, inputs, scalars, seq, &mut Quiet, k) } /// [`run`], reporting every named intermediate and producer result to `obs` /// ([`Observe`]). The outputs and refusals are [`run`]'s. pub fn run_observed(g: &Loaded, inputs: &[Tensor], scalars: &[i64], seq: &mut Sequence, obs: &mut O) -> Res> { - run_inner(g, inputs, scalars, seq, obs, &Reference) + run_inner(g, inputs.to_vec(), scalars, seq, obs, &Reference) } -fn run_inner(g: &Loaded, inputs: &[Tensor], scalars: &[i64], seq: &mut Sequence, obs: &mut O, k: &K) -> Res> { +/// One decode. A tensor is held from the operator that writes it (or the start, for an +/// input) until the last operator that reads it has run, and an output until the end: +/// the integers are the same whatever is held, and a decode then holds, beside its +/// outputs, only what the operator running needs. +fn run_inner(g: &Loaded, inputs: Vec, scalars: &[i64], seq: &mut Sequence, obs: &mut O, k: &K) -> Res> { let g: &Graph = g.graph(); // The inputs, data and state, in input-table order. let n_data = g.inputs.iter().filter(|i| i.kind == 0).count(); if inputs.len() != n_data { return Err(Refusal::new(Code::ModelMismatch, Site::Input(inputs.len().min(n_data) as u32))); } + let first = inputs.first().map(|t| (t.h, t.w)); + let mut data = inputs.into_iter(); let mut all: Vec = Vec::with_capacity(g.inputs.len()); - let (mut d, mut s) = (0usize, 0usize); + let mut s = 0usize; for rec in &g.inputs { if rec.kind == 0 { - all.push(inputs[d].clone()); - d += 1; + all.push(data.next().expect("one data input per data record")); } else { let t = match &seq.state { Some(st) => st[s].clone(), None => { - let i0 = &inputs[0]; - Tensor::filled(g.tensors[rec.tensor as usize].channels, i0.h, i0.w, rec.init) + let (h0, w0) = first.expect("a state input's first decode takes the first data input's sizes"); + Tensor::filled(g.tensors[rec.tensor as usize].channels, h0, w0, rec.init) } }; all.push(t); @@ -238,6 +246,18 @@ fn run_inner(g: &Loaded, inputs: &[Tensor], sca let hw = run_plan(g, &shapes, scalars)?; let mut vals: Vec> = vec![None; g.tensors.len()]; + // The last operator that reads each tensor, and the graph's outputs, which are held + // to the end. + let mut last: Vec> = vec![None; g.tensors.len()]; + for (i, op) in g.ops.iter().enumerate() { + for &t in &op.ins { + last[t as usize] = Some(i); + } + } + let mut output = vec![false; g.tensors.len()]; + for &t in &g.outputs { + output[t as usize] = true; + } // Step 4 begins with the inputs: each input's policy applies to it. for (k, (rec, mut t)) in g.inputs.iter().zip(all).enumerate() { let tr = &g.tensors[rec.tensor as usize]; @@ -269,8 +289,20 @@ fn run_inner(g: &Loaded, inputs: &[Tensor], sca } vals[t as usize] = Some(ten); } + // Release what no later operator reads: this operator's inputs whose last reader + // it is, and any output of it that nothing reads. + for &t in op.ins.iter().chain(&op.outs) { + let t = t as usize; + if !output[t] && last[t].is_none_or(|l| l <= i) { + vals[t] = None; + } + } + } + let mut outputs: Vec = Vec::with_capacity(g.outputs.len()); + for (k, &t) in g.outputs.iter().enumerate() { + let v = if g.outputs[k + 1..].contains(&t) { vals[t as usize].clone() } else { vals[t as usize].take() }; + outputs.push(v.expect("every tensor is defined")); } - let outputs: Vec = g.outputs.iter().map(|&t| vals[t as usize].clone().expect("every tensor is defined")).collect(); let next: Vec = g.inputs.iter().filter(|i| i.kind == 1).map(|i| outputs[i.src as usize].clone()).collect(); seq.state = Some(next); Ok(outputs) @@ -421,6 +453,12 @@ pub(crate) fn k_conv( /// `conv_transpose` (§5.3), in the scatter form: the same products as §5.3's gather, /// summed exactly. +/// +/// The uncropped plane holds `i64`. Every cell, cropped or kept, is at every step a sum +/// of some of the taps of one residue (`i mod sh`, `j mod sw`), so it lies in that +/// residue's partial-sum interval, which §8.3's `convT.acc` covers; and a loaded graph +/// meets `acc64`, under which `convT.acc` and `convT.term` fit `i64` (§7.1). A term or a +/// cell that would leave `i64`, which no loaded graph forms, refuses (`model_mismatch`). #[allow(clippy::too_many_arguments)] pub(crate) fn k_conv_t( x: &Tensor, cout: usize, kh: usize, kw: usize, sh: usize, sw: usize, ph: usize, pw: usize, oph: usize, opw: usize, rq: &Rq, wt: &[i64], @@ -429,7 +467,7 @@ pub(crate) fn k_conv_t( let (c, h, w) = (x.c as usize, x.h as usize, x.w as usize); let uh = (h - 1) * sh + kh + oph; let uw = (w - 1) * sw + kw + opw; - let mut plane = vec![0i128; cout * uh * uw]; + let mut plane = vec![0i64; cout * uh * uw]; for ci in 0..c { for u in 0..h { for v in 0..w { @@ -444,8 +482,8 @@ pub(crate) fn k_conv_t( let wv = wt[((ci * cout + o) * kh + i) * kw + j] as i128; obs.see(at, "convT.term", wv * xv); let cell = &mut plane[row + j]; - *cell = cell.checked_add(wv * xv).ok_or(wide(at))?; - obs.see(at, "convT.acc", *cell); + *cell = i64::try_from(wv * xv).ok().and_then(|t| cell.checked_add(t)).ok_or(wide(at))?; + obs.see(at, "convT.acc", *cell as i128); } } } @@ -456,7 +494,7 @@ pub(crate) fn k_conv_t( for o in 0..cout { for y in 0..oh { for xx in 0..ow { - let acc = plane[(o * uh + y + ph) * uw + xx + pw]; + let acc = plane[(o * uh + y + ph) * uw + xx + pw] as i128; d[(o * oh + y) * ow + xx] = to64(rq_apply(rq, o, acc, at, obs).ok_or(wide(at))?, at)?; } } diff --git a/wai/wai-ops/src/lower/mod.rs b/wai/wai-ops/src/lower/mod.rs index 702abbf7f..271c6d8c6 100644 --- a/wai/wai-ops/src/lower/mod.rs +++ b/wai/wai-ops/src/lower/mod.rs @@ -16,7 +16,7 @@ mod whm; pub use jpegai::jpegai; pub use synth::int_synth; -pub use whm::{Net, NetLayer, int_hyper, parse_wam1, parse_whm1, splat, wam1, whm1, wmb1, world}; +pub use whm::{Net, NetLayer, int_hyper, parse_wam1, parse_whm1, splat, wam1, wam1_net, whm1, whm1_net, wmb1, world}; #[cfg(kani)] pub(crate) use whm::{p_conv, p_igdn, p_leaky}; diff --git a/wai/wai-ops/src/lower/whm.rs b/wai/wai-ops/src/lower/whm.rs index d3833782a..e1db5a231 100644 --- a/wai/wai-ops/src/lower/whm.rs +++ b/wai/wai-ops/src/lower/whm.rs @@ -312,7 +312,17 @@ fn net_graph(net: &Net, output_step: bool) -> Res { /// `whm1/1`: one `WHM1` network, with the RGB8 step when `output_step` is set. pub fn whm1(bytes: &[u8], output_step: bool) -> Res { - net_graph(&parse_whm1(bytes)?, output_step) + whm1_net(&parse_whm1(bytes)?, output_step) +} + +/// [`whm1`] of a network already parsed: a decoder that holds its network parsed lowers +/// it from that form, which is [`parse_whm1`] of its bytes. The graph is the one +/// [`whm1`] writes from those bytes. +pub fn whm1_net(net: &Net, output_step: bool) -> Res { + if net.audio { + return refuse(); + } + net_graph(net, output_step) } /// `whm1/1` for `wai.neural.int_hyper` and the keyframe of `wai.video.int_hyper`: @@ -328,14 +338,23 @@ pub fn splat(g_s: &[u8]) -> Res { /// `wam1/1` (§13.3): the `WAM1` synthesis with the PCM step. pub fn wam1(bytes: &[u8]) -> Res { - let net = parse_wam1(bytes)?; - header_ok(&net)?; - let c0 = in_channels(&net)?; + wam1_net(&parse_wam1(bytes)?) +} + +/// [`wam1`] of a network already parsed: a decoder that holds its network parsed lowers +/// it from that form, which is [`parse_wam1`] of its bytes. The graph is the one [`wam1`] +/// writes from those bytes. +pub fn wam1_net(net: &Net) -> Res { + if !net.audio { + return refuse(); + } + header_ok(net)?; + let c0 = in_channels(net)?; let mut b = Builder::new(); let first = net.layers.first().ok_or(()).or_else(|_| refuse())?; - let x0 = b.input(x_record(&net, first, c0, true)?, 0, 0); + let x0 = b.input(x_record(net, first, c0, true)?, 0, 0); let t_out = b.scalar(1, 1 << 26); - let x = layers(&mut b, &net, x0)?; + let x = layers(&mut b, net, x0)?; let c = b.channels(x); // The range in which v·32767 + 2^(afrac−1) fits i64. let lo = -((1i128 << 63) / 32767) as i64; // ⌈−2^63 / 32767⌉ diff --git a/wai/wai-rs/Cargo.toml b/wai/wai-rs/Cargo.toml index 1e71b366c..af472d6f0 100644 --- a/wai/wai-rs/Cargo.toml +++ b/wai/wai-rs/Cargo.toml @@ -749,13 +749,23 @@ required-features = ["neural_int"] [[test]] name = "int_ops_differential" path = "tests/int_ops_differential.rs" -required-features = ["legacy_int"] +required-features = ["legacy_int", "splat_int", "audio_int"] [[test]] name = "int_synth_ops" path = "tests/int_synth_ops.rs" required-features = ["legacy_int"] +[[test]] +name = "int_audio_ops" +path = "tests/int_audio_ops.rs" +required-features = ["audio_int", "legacy_int"] + +[[test]] +name = "int_splat_ops" +path = "tests/int_splat_ops.rs" +required-features = ["splat_int", "legacy_int"] + [[test]] name = "int_mlicpp" path = "tests/int_mlicpp.rs" diff --git a/wai/wai-rs/src/codecs/mod.rs b/wai/wai-rs/src/codecs/mod.rs index 9b941602c..52a4f9e31 100644 --- a/wai/wai-rs/src/codecs/mod.rs +++ b/wai/wai-rs/src/codecs/mod.rs @@ -1353,8 +1353,12 @@ pub struct OpsMember { /// networks through their lowering (`extensions/integer-operators.md` §15.1). pub fn capability_ops(cap: &str) -> Option { let ops = |lowering, rans64, bucket| Some(OpsMember { profile: "acc64", subset: "full", lowering, rans64, bucket }); + // integer-payloads §2's stream: `rans64_decode` with `t4097`, `lenient` and `canonical`. + const PAYLOAD_STREAM: Option<(&str, &str, &str)> = Some(("t4097", "lenient", "canonical")); match cap { CAP_NEURAL_INT_SYNTH => ops("int_synth/1", None, None), + CAP_SPLAT_INT_CODEC => ops("whm1/1", PAYLOAD_STREAM, None), + CAP_AUDIO_INT_CODEC => ops("wam1/1", PAYLOAD_STREAM, None), _ => None, } } diff --git a/wai/wai-rs/src/codecs/registry.rs b/wai/wai-rs/src/codecs/registry.rs index f22f8389e..ddb12dd2a 100644 --- a/wai/wai-rs/src/codecs/registry.rs +++ b/wai/wai-rs/src/codecs/registry.rs @@ -1420,13 +1420,14 @@ mod tests { /// The `ops` members (`extensions/integer-operators.md` §11). Only a /// `NeuralIntegerExact` capability carries one, every one is at `acc64` with the /// full subset and names a lowering of §13, and the export writes exactly these. - /// In this revision `wai.neural.int_synth` is the one capability that decodes - /// through the operator set. + /// In this revision `wai.neural.int_synth`, `wai.splat.int_codec` and + /// `wai.audio.int_codec` decode through the operator set. #[test] fn ops_members_name_the_operator_set_a_capability_decodes_through() { let regs = registrations(); - let with: Vec<&str> = regs.iter().filter(|r| r.ops.is_some()).map(|r| r.capability).collect(); - assert_eq!(with, [CAP_NEURAL_INT_SYNTH]); + let mut with: Vec<&str> = regs.iter().filter(|r| r.ops.is_some()).map(|r| r.capability).collect(); + with.sort_unstable(); + assert_eq!(with, [CAP_AUDIO_INT_CODEC, CAP_NEURAL_INT_SYNTH, CAP_SPLAT_INT_CODEC]); for r in regs.iter().filter(|r| r.ops.is_some()) { let o = r.ops.unwrap(); assert_eq!(r.replicate, Some(ReplicateClass::NeuralIntegerExact), "{}", r.capability); @@ -1437,6 +1438,14 @@ mod tests { for c in export["capabilities"].as_array().unwrap() { let want = match c["capability"].as_str().unwrap() { CAP_NEURAL_INT_SYNTH => serde_json::json!({ "set": "wai-ops/1", "profile": "acc64", "subset": "full", "lowering": "int_synth/1" }), + CAP_SPLAT_INT_CODEC => serde_json::json!({ + "set": "wai-ops/1", "profile": "acc64", "subset": "full", "lowering": "whm1/1", + "stage": { "rans64": { "table": "t4097", "escape": "lenient", "end": "canonical" } } + }), + CAP_AUDIO_INT_CODEC => serde_json::json!({ + "set": "wai-ops/1", "profile": "acc64", "subset": "full", "lowering": "wam1/1", + "stage": { "rans64": { "table": "t4097", "escape": "lenient", "end": "canonical" } } + }), _ => serde_json::Value::Null, }; assert_eq!(c["ops"], want, "{}", c["capability"]); diff --git a/wai/wai-rs/src/int_audio.rs b/wai/wai-rs/src/int_audio.rs index 3b4a118f6..5aca02912 100644 --- a/wai/wai-rs/src/int_audio.rs +++ b/wai/wai-rs/src/int_audio.rs @@ -16,7 +16,6 @@ use serde_json::Value; -use crate::int_transform::igdn; use crate::rans::RansDecoder; /// One layer of the integer synthesis `g_s`. `ConvT1d` weight is `[cin·cout·k]` @@ -112,27 +111,11 @@ impl AudioModel { } } -/// One integer 1-D transposed conv with per-output-channel requant — see -/// [`reference::conv_t_1d`] for the definition. Runs the fast kernel -/// ([`crate::int_kernels`], as a 1-row 2-D transposed conv) unless the current -/// [`crate::int_kernels::Policy`] selects the reference; both give the same integer -/// for every output element. -#[allow(clippy::too_many_arguments)] -fn conv_t_1d( - x: &[i64], cin: usize, t: usize, weight: &[i64], cout: usize, k: usize, - stride: usize, pad: usize, opad: usize, mult: &[i64], bias: &[i64], shift: u32, -) -> (Vec, usize) { - use crate::int_kernels::{conv_acc_i64, policy, ConvGeom, KernelPath}; - if policy().path == KernelPath::Fast && mult.len() >= cout && bias.len() >= cout - && let Some(g) = ConvGeom::transposed(cin, 1, t, cout, 1, k, 1, stride, 0, pad, 0, opad, x.len(), weight.len()) - && let Some(out) = conv_acc_i64(&g, x, weight, |oc, raw| ((raw as i128 * mult[oc] as i128) >> shift) as i64 + bias[oc]) - { - return (out, g.wo); - } - reference::conv_t_1d(x, cin, t, weight, cout, k, stride, pad, opad, mult, bias, shift) -} - -/// The scalar reference kernel the fast path is checked against, kept verbatim. +/// The direct decode's one-dimensional transposed convolution, kept verbatim: the scalar +/// kernel its fast dispatch was checked against. That dispatch is gone: the decode runs +/// through the integer operator set, whose interpreter lends the convolution to the fast +/// kernel behind it (`int_ops`), held to the interpreter there. +#[cfg(any(test, feature = "legacy_int"))] #[doc(hidden)] pub mod reference { /// One integer 1-D transposed conv with per-output-channel requant @@ -308,11 +291,13 @@ fn plan(model: &AudioModel, c: usize, t: usize) -> Result<(usize, usize, u64, u1 Ok((c as usize, t as usize, live, macs)) } -/// What [`decode_typed`] checks before it decodes a symbol, and the buffers it then -/// forms. Returns the synthesis's length, the most bytes the decode holds at once -/// on the reference kernels beyond the payload and the parsed model (the latent and -/// each layer's buffers, `i64` each, then the synthesis's output, `i64`, with the -/// PCM, `i16`), and the synthesis's multiply-accumulates on the latent. +/// What [`decode_typed`] checks before it decodes a symbol, short of lowering the +/// synthesis ([`ops::plan`] adds that), and the buffers it then forms. Returns the +/// synthesis's length, the most bytes the decode holds at once on the reference kernels +/// beyond the payload and the set as it holds it (the parsed model, and its loaded +/// lowering): the latent and each layer's buffers (its input, its output and a +/// transposed convolution's uncropped plane), `i64` each, then the synthesis's output, +/// `i64`, with the PCM, `i16`; and the synthesis's multiply-accumulates on the latent. /// /// The latent's length is the one `t_out` needs and no longer: the synthesis of /// `lat_len` latent samples gives at least `t_out` samples (`GeometryMismatch` @@ -359,34 +344,6 @@ pub(crate) fn audio_plan(model: &AudioModel, s: &AudioStream<'_>) -> Result<(usi Ok((t, working, macs)) } -/// One layer's arithmetic on the values `x` holds, checked before it runs, as -/// [`crate::int_hyper_synth::run_net_checked`] checks a 2-D layer: every product and -/// partial sum and the requantised value of a convolution fit `i64`; IGDN's terms, -/// norm and root fit their types. -fn layer_in_range(model: &AudioModel, layer: &ALayer, x: &[i64], c: usize) -> Result<(), Refused> { - use crate::int_kernels::max_abs; - const I64: u128 = i64::MAX as u128; - match layer { - ALayer::ConvT1d { k, weight, mult, bias, .. } => { - if model.shift > 127 { - return Err(refused(RefusalKind::ModelMismatch, "requantisation shift of 128 or more")); - } - let terms = (c as u128).saturating_mul(*k as u128); - let acc = terms.saturating_mul(max_abs(weight) as u128).saturating_mul(max_abs(x) as u128); - if acc > I64 { - return Err(refused(RefusalKind::FieldRange, "a latent could take an accumulation outside i64")); - } - let q = ((acc * max_abs(mult) as u128) >> model.shift) + 1; - if q.saturating_add(max_abs(bias) as u128) > I64 { - return Err(refused(RefusalKind::FieldRange, "a latent could take a requantised value outside i64")); - } - Ok(()) - } - ALayer::Igdn { gamma, beta, .. } => crate::int_hyper_synth::igdn_check(x, c, gamma, beta, model.gf, model.afrac) - .map_err(|(why, values)| refused(if values { RefusalKind::FieldRange } else { RefusalKind::ModelMismatch }, why)), - } -} - /// Decode a typed clip → int16 PCM, mono, byte-exact on every machine, refusing /// rather than panicking or overflowing on any input: /// @@ -397,17 +354,28 @@ fn layer_in_range(model: &AudioModel, layer: &ALayer, x: &[i64], c: usize) -> Re /// 2. per channel, rANS-decode the symbol `s`; the latent is `s·2^afrac + median_q` /// ([`crate::int_entropy::dequantize`]), and a stream that runs out before its last /// symbol is refused; -/// 3. integer synthesis `g_s`, each layer's arithmetic checked before it runs; -/// 4. the first `t_out` samples, `(x·32767 + 2^(afrac−1)) >> afrac` clamped to `i16`. +/// 3. integer synthesis `g_s` through the integer operator set ([`crate::int_ops`]), +/// lowered by `wam1/1` and loaded under `acc64` with its bounds certificate before +/// the stream is decoded (`extensions/integer-operators.md` §13.3); its `check` ranges +/// refuse a latent that could take a layer outside the exact integer range +/// (`field_range`); +/// 4. the graph's PCM step: the first `t_out` samples, `(x·32767 + 2^(afrac−1)) >> afrac` +/// clamped to `i16`, a sample whose scaling would leave `i64` refused (`field_range`). +/// +/// This is [`ops::decode_typed`]. The direct decode it replaced is kept, test-only +/// (feature `legacy_int`), as `legacy::decode_typed`. pub fn decode_typed(model: &AudioModel, s: &AudioStream<'_>) -> Result, Refused> { + ops::decode_typed(model, s) +} + +/// Step 2 of the decode, after its plan: the factorized rANS decode of the per-channel +/// latent (channel-major), each symbol rebuilt as `s·2^afrac + median_q`. The latent +/// grows a symbol at a time, and the first symbol that runs the stream out refuses the +/// payload before its value is used; the stream must then end canonically. +fn latent(s: &AudioStream<'_>, afrac: u32) -> Result, Refused> { use crate::int_entropy::dequantize; use RefusalKind::*; - audio_plan(model, s)?; - let (afrac, c, l, len) = (model.afrac, s.n_ch, s.lat_len, s.cdfs[0].len()); - - // 1. factorized rANS-decode the per-channel latent (channel-major), rebuild y_q. - // The latent grows a symbol at a time, and the first symbol that runs the stream - // out refuses the payload before its value is used. + let (c, l, len) = (s.n_ch, s.lat_len, s.cdfs[0].len()); let mut y = Vec::new(); let mut dec = RansDecoder::new(s.rans).map_err(|_| refused(FieldRange, "the latent stream is not a rANS stream"))?; for ch in 0..c { @@ -422,31 +390,168 @@ pub fn decode_typed(model: &AudioModel, s: &AudioStream<'_>) -> Result, if !dec.is_canonical_end() { return Err(refused(StreamNoncanonical, "the latent stream does not end canonically after its last symbol")); } - // 2. integer synthesis → waveform, each layer checked on the values it is given. - let (mut x, mut cc, mut tt) = (y, c, l); - for layer in &model.layers { - layer_in_range(model, layer, &x, cc)?; - (x, cc, tt) = run_layer(model, layer, x, cc, tt); - } - // 3. waveform → int16 PCM (round-half-up, clamp): the first t_out samples. - let half = 1i64 << (afrac - 1); - let mut pcm = Vec::with_capacity(s.t_out); - for &v in &x[..s.t_out] { - let num = v.checked_mul(32767).and_then(|n| n.checked_add(half)) - .ok_or(refused(FieldRange, "a synthesised sample leaves i64 when scaled"))?; - pcm.push((num >> afrac).clamp(-32768, 32767) as i16); - } - Ok(pcm) + Ok(y) } -/// One layer of the synthesis on `x` (`[c, t]`): the new activations and shape. -fn run_layer(model: &AudioModel, layer: &ALayer, x: Vec, c: usize, t: usize) -> (Vec, usize, usize) { - match layer { - ALayer::ConvT1d { out_ch, k, stride, pad, opad, weight, mult, bias } => { - let (nx, to) = conv_t_1d(&x, c, t, weight, *out_ch, *k, *stride, *pad, *opad, mult, bias, model.shift); - (nx, *out_ch, to) +/// The direct decode, as the decoder ran before it decoded through the integer operator +/// set: each layer's arithmetic checked on the values it is given before it runs +/// (`layer_in_range`, the guard the graph's `check` ranges now carry), each transposed +/// convolution on the scalar kernel ([`reference::conv_t_1d`]), then the PCM formed here. +/// Test-only (`legacy_int`): the differential tests decode both ways, and the two differ +/// only on the planned breaks of `extensions/integer-operators.md` §15.3. +#[cfg(any(test, feature = "legacy_int"))] +pub mod legacy { + use super::*; + use crate::int_transform::igdn; + + /// One layer's arithmetic on the values `x` holds, checked before it runs, as + /// [`crate::int_hyper_synth::run_net_checked`] checks a 2-D layer: every product and + /// partial sum and the requantised value of a convolution fit `i64`; IGDN's terms, + /// norm and root fit their types. + fn layer_in_range(model: &AudioModel, layer: &ALayer, x: &[i64], c: usize) -> Result<(), Refused> { + use crate::int_kernels::max_abs; + const I64: u128 = i64::MAX as u128; + match layer { + ALayer::ConvT1d { k, weight, mult, bias, .. } => { + if model.shift > 127 { + return Err(refused(RefusalKind::ModelMismatch, "requantisation shift of 128 or more")); + } + let terms = (c as u128).saturating_mul(*k as u128); + let acc = terms.saturating_mul(max_abs(weight) as u128).saturating_mul(max_abs(x) as u128); + if acc > I64 { + return Err(refused(RefusalKind::FieldRange, "a latent could take an accumulation outside i64")); + } + let q = ((acc * max_abs(mult) as u128) >> model.shift) + 1; + if q.saturating_add(max_abs(bias) as u128) > I64 { + return Err(refused(RefusalKind::FieldRange, "a latent could take a requantised value outside i64")); + } + Ok(()) + } + ALayer::Igdn { gamma, beta, .. } => crate::int_hyper_synth::igdn_check(x, c, gamma, beta, model.gf, model.afrac) + .map_err(|(why, values)| refused(if values { RefusalKind::FieldRange } else { RefusalKind::ModelMismatch }, why)), } - ALayer::Igdn { gamma, beta, .. } => (igdn(&x, c, t, model.afrac, gamma, beta, model.gf), c, t), + } + + /// One layer of the synthesis on `x` (`[c, t]`): the new activations and shape. + fn run_layer(model: &AudioModel, layer: &ALayer, x: Vec, c: usize, t: usize) -> (Vec, usize, usize) { + match layer { + ALayer::ConvT1d { out_ch, k, stride, pad, opad, weight, mult, bias } => { + let (nx, to) = reference::conv_t_1d(&x, c, t, weight, *out_ch, *k, *stride, *pad, *opad, mult, bias, model.shift); + (nx, *out_ch, to) + } + ALayer::Igdn { gamma, beta, .. } => (igdn(&x, c, t, model.afrac, gamma, beta, model.gf), c, t), + } + } + + /// The direct decode of [`super::decode_typed`]'s arguments. + pub fn decode_typed(model: &AudioModel, s: &AudioStream<'_>) -> Result, Refused> { + use RefusalKind::*; + audio_plan(model, s)?; + let afrac = model.afrac; + let y = latent(s, afrac)?; + // Integer synthesis → waveform, each layer checked on the values it is given. + let (mut x, mut cc, mut tt) = (y, s.n_ch, s.lat_len); + for layer in &model.layers { + layer_in_range(model, layer, &x, cc)?; + (x, cc, tt) = run_layer(model, layer, x, cc, tt); + } + // Waveform → int16 PCM (round-half-up, clamp): the first t_out samples. + let half = 1i64 << (afrac - 1); + let mut pcm = Vec::with_capacity(s.t_out); + for &v in &x[..s.t_out] { + let num = v.checked_mul(32767).and_then(|n| n.checked_add(half)) + .ok_or(refused(FieldRange, "a synthesised sample leaves i64 when scaled"))?; + pcm.push((num >> afrac).clamp(-32768, 32767) as i16); + } + Ok(pcm) + } +} + +/// The synthesis through the integer operator set (`extensions/integer-operators.md`): +/// `WAM1`'s lowering `wam1/1` (§13.3), with the PCM step. +pub mod ops { + use super::*; + use wai_ops::lower::{Net, NetLayer}; + use wai_ops::{Code, Refusal, Site}; + + /// The model this crate parses ([`AudioModel`]) in the form `wai-ops` lowers it from + /// ([`Net`]): the one [`wai_ops::lower::parse_wam1`] gives of the bytes it was parsed + /// from, so [`wai_ops::lower::wam1_net`] of it writes the graph + /// [`wai_ops::lower::wam1`] writes from those bytes. The model does not keep a + /// transposed convolution's input channels: they are its weights' count over its + /// output channels and kernel, as the bytes give them (0 where those hold a 0, which + /// the lowering refuses). A model parsed from bytes holds only 32-bit fields; one with + /// a field past 32 bits is refused here, as the lowering refuses one + /// (`model_mismatch`). + pub fn wam1_net(m: &AudioModel) -> Result { + let f = |v: usize| u32::try_from(v).map_err(|_| Refusal::new(Code::ModelMismatch, Site::Lowering)); + let layers = m.layers.iter().map(|l| { + Ok(match l { + ALayer::ConvT1d { out_ch, k, stride, pad, opad, weight, mult, bias } => NetLayer::ConvT { + inc: f(out_ch.checked_mul(*k).and_then(|d| weight.len().checked_div(d)).unwrap_or(0))?, + out: f(*out_ch)?, k: f(*k)?, stride: f(*stride)?, pad: f(*pad)?, opad: f(*opad)?, + w: weight.clone(), mult: mult.clone(), bias: bias.clone(), + }, + ALayer::Igdn { c, gamma, beta } => NetLayer::Igdn { c: f(*c)?, gamma: gamma.clone(), beta: beta.clone() }, + }) + }); + Ok(Net { afrac: m.afrac, gf: m.gf, shift: m.shift, layers: layers.collect::>()?, audio: true }) + } + + /// A refusal of the operator set as the audio decode reports it + /// (`extensions/integer-payloads.md` §8): a value outside a `check` range is + /// `field_range`; the other codes are the payload extension's of the same name. + fn kind(r: Refusal, what: &'static str) -> Refused { + let k = match r.code { + Code::FieldRange => RefusalKind::FieldRange, + Code::TooLarge => RefusalKind::TooLarge, + Code::GeometryMismatch => RefusalKind::GeometryMismatch, + Code::StreamExhausted => RefusalKind::StreamExhausted, + Code::StreamNoncanonical => RefusalKind::StreamNoncanonical, + Code::ModelMismatch => RefusalKind::ModelMismatch, + }; + refused(k, what) + } + + /// The synthesis lowered by `wam1/1` with the PCM step and loaded under `acc64`: the + /// set as the decode holds it, beside the parsed model. + pub fn load(m: &AudioModel) -> Result { + const WHAT: &str = "the synthesis does not lower to an integer-operator graph that loads under acc64"; + let net = wam1_net(m).map_err(|r| kind(r, WHAT))?; + let g = wai_ops::lower::wam1_net(&net).map_err(|r| kind(r, WHAT))?; + drop(net); + crate::int_ops::load(&g).map_err(|r| kind(r, WHAT)) + } + + /// Everything the decode checks before it decodes a symbol: [`audio_plan`], then the + /// synthesis lowered and loaded. Returns the loaded graph, and [`audio_plan`]'s + /// working bytes and multiply-accumulates. + pub(crate) fn plan(m: &AudioModel, s: &AudioStream<'_>) -> Result<(wai_ops::Loaded, u64, u128), Refused> { + let (_, working, macs) = audio_plan(m, s)?; + Ok((load(m)?, working, macs)) + } + + /// [`super::decode_typed`] through the integer operator set: the same PCM, and the + /// same refusals, except where integer-operators §15.3 lists a change. + pub fn decode_typed(m: &AudioModel, s: &AudioStream<'_>) -> Result, Refused> { + let (g, ..) = plan(m, s)?; + decode_loaded(&g, m.afrac, s) + } + + /// The decode after its plan: the stream at `afrac`, then `g`, the set's loaded + /// lowering ([`load`]), with `t_out` as its scalar 0, then the samples. `s` must be a + /// clip the plan admitted against that set. What it holds beyond the set is what the + /// payload declares (integer-payloads §7, *Memory and work*). + pub fn decode_loaded(g: &wai_ops::Loaded, afrac: u32, s: &AudioStream<'_>) -> Result, Refused> { + let y = latent(s, afrac)?; + // The plan admitted every size, so each fits u32. + let x = wai_ops::Tensor::new(s.n_ch as u32, 1, s.lat_len as u32, y); + let out = crate::int_ops::run1(g, x, &[s.t_out as i64]).map_err(|r| kind(r, "a decoded latent or a synthesised sample is outside a check range of the synthesis"))?; + if (out.c, out.h, out.w as usize) != (1, 1, s.t_out) { + return Err(refused(RefusalKind::GeometryMismatch, "the synthesis does not give t_out samples")); + } + // The PCM step ends the graph with clamp(−32768, 32767), so every value is a sample. + Ok(out.data.iter().map(|&v| v as i16).collect()) } } @@ -500,7 +605,7 @@ fn b64_decode(s: &str) -> Result, &'static str> { Ok(out) } -/// The fast 1-D transposed conv against the scalar reference, element by element. +/// The plan's multiply-accumulates against the layers' own terms. #[cfg(test)] mod plan_tests { use super::*; @@ -511,7 +616,9 @@ mod plan_tests { /// every element times `c`. #[test] fn the_planned_macs_are_every_layers_terms() { - let dir = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/vectors/int_audio"); + // Through `../wai-rs`, so the path holds in wai-rs and in the byte-equality crate, + // which includes this file. + let dir = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../wai-rs/tests/vectors/int_audio"); let model = AudioModel::from_bytes(&std::fs::read(dir.join("audio_model.bin")).unwrap()).unwrap(); let wia1 = std::fs::read(dir.join("audio_bitstream.wia1")).unwrap(); let a = crate::int_wire::parse_wia1(&wia1).unwrap(); @@ -533,134 +640,101 @@ mod plan_tests { } } +/// Edge geometry through the decode's graph (`wam1/1`, run by the interpreter with the +/// fast kernels behind it), under every kernel policy: the cases the direct decode's +/// fast one-dimensional dispatch was tested on beside its kernel, now that the dispatch +/// is gone. Each graph is the lowering of a one-layer model; its output is the PCM step's. #[cfg(test)] -mod kernel_equivalence { +mod geometry { use super::*; - use crate::int_kernels::testkit::{Rng, FAST}; - use crate::int_kernels::with_policy; + use crate::int_kernels::testkit::{fast_conv_runs, FAST}; + use crate::int_kernels::{with_policy, Policy}; + use wai_ops::{Code, Tensor}; - #[test] - fn conv_t_1d_equals_reference() { - let mut g = Rng(0xA0D_0001); - let mut ran = 0; - for case in 0..500 { - let k = 1 + g.below(9); - let stride = 1 + g.below(5); - let pad = g.below(k); - let opad = if stride > 1 { g.below(stride) } else { 0 }; - let t = 1 + g.below(40); - let (cin, cout) = (1 + g.below(6), 1 + g.below(6)); - if (t - 1) * stride + k + opad <= 2 * pad { - continue; - } - let (xr, wr) = [(1 << 15, 2047), (1 << 40, 1 << 10)][case % 2]; - let x = g.vec(cin * t, xr, 30); - let wt = g.vec(cin * cout * k, wr, 5); - let mult: Vec = (0..cout).map(|_| 1 + g.signed(1 << 16).abs()).collect(); - let bias = g.vec(cout, 1 << 16, 0); - let want = reference::conv_t_1d(&x, cin, t, &wt, cout, k, stride, pad, opad, &mult, &bias, 24); - for p in FAST { - let got = with_policy(p, || conv_t_1d(&x, cin, t, &wt, cout, k, stride, pad, opad, &mult, &bias, 24)); - assert_eq!(got, want, "conv_t_1d {p:?} k={k} s={stride} p={pad} op={opad} t={t}"); - } - ran += 1; - } - assert!(ran > 250, "too few valid cases ({ran})"); + /// A one-layer model: a transposed convolution of `k` taps `w` from `cin` channels to + /// one, its multiplier `m`, bias `b` and the header's shift `shift`, at `afrac` 12. + #[allow(clippy::too_many_arguments)] + fn one(cin: usize, k: usize, stride: usize, pad: usize, opad: usize, w: Vec, m: i64, b: i64, shift: u32) -> wai_ops::Loaded { + assert_eq!(w.len(), cin * k); + let model = AudioModel { afrac: 12, gf: 16, shift, layers: vec![ALayer::ConvT1d { out_ch: 1, k, stride, pad, opad, weight: w, mult: vec![m], bias: vec![b] }] }; + ops::load(&model).expect("the one-layer model lowers and loads") } - /// Overflow only at a position the reference accumulates into and then crops: - /// a length-1 input, kernel 3, pad 1 keeps only the middle of the reference's - /// 3-sample raw buffer (`2·2^40·1 = 2^41`); both outer samples are - /// `2·2^40·2^22 = 2^63` and overflow `i64`. The fast path must panic where the - /// reference does and otherwise give the same 2^41. - #[test] - fn overflow_only_at_a_cropped_border_behaves_like_reference() { - use crate::int_kernels::testkit::overflow_checks; - use crate::int_kernels::Policy; - let (cin, cout, k) = (2usize, 1usize, 3usize); - let x = vec![1i64 << 40; cin]; - let wt: Vec = (0..cin * cout).flat_map(|_| [1i64 << 22, 1, 1 << 22]).collect(); - let (mult, bias) = (vec![1i64; cout], vec![0i64; cout]); - let run = |p: Policy| { - std::panic::catch_unwind(|| with_policy(p, || conv_t_1d(&x, cin, 1, &wt, cout, k, 1, 1, 0, &mult, &bias, 0))).ok() - }; - let want = run(Policy::reference()); - if overflow_checks() { - assert!(want.is_none(), "the reference overflows at a cropped sample in a checked build"); - } else { - assert_eq!(want, Some((vec![1i64 << 41], 1))); - } + /// `x` (`cin` channels of one sample) through `g`, presenting `t_out` samples, under + /// the reference and each fast policy: the outcomes, which must agree, and the fast + /// kernel's runs under each fast policy. + fn run(g: &wai_ops::Loaded, x: &[i64], t_out: i64) -> (Result, Code>, Vec) { + let go = || crate::int_ops::run1(g, Tensor::new(x.len() as u32, 1, 1, x.to_vec()), &[t_out]).map(|t| t.data).map_err(|r| r.code); + let want = with_policy(Policy::reference(), go); + let mut runs = Vec::new(); for p in FAST { - assert_eq!(run(p), want, "fast must match the reference, panic included ({p:?})"); + let (got, n) = fast_conv_runs(|| with_policy(p, go)); + assert_eq!(got, want, "{p:?} differs from the interpreter alone"); + runs.push(n); } + (want, runs) } - /// A one-sample input upsampled by any stride is one output sample (k = 3, - /// pad 1): a handful of operations for the reference. The fast path must return - /// the same sample without overflowing, or allocating or looping per stride step. + /// The PCM step's sample for a synthesised value `v` at `afrac` 12. + fn pcm(v: i64) -> i64 { + ((v * 32767 + 2048) >> 12).clamp(-32768, 32767) + } + + /// A one-sample input upsampled by any stride a `WAM1` field holds is one output + /// sample (k = 3, pad 1): the middle tap, 5·2 = 10, under every policy, with no + /// allocation or loop in proportion to the stride. #[test] fn a_one_sample_input_upsampled_by_any_stride() { - use crate::int_kernels::Policy; - let (x, wt) = ([5i64], [1i64, 2, 3]); - // usize::MAX first: a kernel allocating a table per phase fails there at once - // ("capacity overflow") instead of committing memory for 2^29 tables later. - for s in [usize::MAX, 1, 2, 1 << 13, 1 << 16, 1 << 29] { - let run = |p: Policy| with_policy(p, || conv_t_1d(&x, 1, 1, &wt, 1, 3, s, 1, 0, &[1], &[0], 0)); - let want = run(Policy::reference()); - assert_eq!(want, (vec![10], 1), "stride {s}: the middle of 5·[1, 2, 3]"); - for p in FAST { - assert_eq!(run(p), want, "stride {s} {p:?}"); - } + for s in [u32::MAX as usize, 1, 2, 1 << 13, 1 << 16, 1 << 29] { + let g = one(1, 3, s, 1, 0, vec![1, 2, 3], 1, 0, 0); + assert_eq!(run(&g, &[5], 1).0, Ok(vec![pcm(10)]), "stride {s}"); } } - /// Several input samples upsampled by a stride and cropped by padding so large - /// that the reference's uncropped buffer `cout · (to + 2·pad)` cannot be sized: - /// three samples, kernel 1, stride = pad = 2^62, one output sample. With two - /// output channels the size overflows `usize` (the reference panics with - /// overflow checks on, and with them off indexes out of the wrapped buffer); - /// with one it is 2^63 + 1 samples, more than any allocation can hold ("capacity - /// overflow"). The reference panics in every build, so every policy must, and - /// the fast kernel must not be entered. - #[cfg(target_pointer_width = "64")] + /// Several samples upsampled by a stride and cropped by padding so large that the + /// direct decode's uncropped buffer could not be sized: here the stride and the + /// padding are 2^31 (a `WAM1` field is 32 bits). The run-time plan caps the uncropped + /// length at 2^26 before anything is formed, so every policy refuses (`too_large`), + /// and the fast kernel is not entered. #[test] - fn several_samples_whose_uncropped_buffer_cannot_be_sized() { - use crate::int_kernels::testkit::fast_conv_runs; - use crate::int_kernels::Policy; - let s = 1usize << 62; - for cout in [2usize, 1] { - let (x, wt) = ([1i64, 2, 3], vec![7i64; cout]); - let (mult, bias) = (vec![1i64; cout], vec![0i64; cout]); - let run = |p: Policy| { - fast_conv_runs(|| std::panic::catch_unwind(|| with_policy(p, || conv_t_1d(&x, 1, 3, &wt, cout, 1, s, s, 0, &mult, &bias, 0))).ok()) - }; - assert_eq!(run(Policy::reference()), (None, 0), "{cout} channels: the reference panics"); - for p in FAST { - assert_eq!(run(p), (None, 0), "{cout} channels {p:?}"); - } + fn several_samples_whose_uncropped_buffer_is_past_the_cap() { + let s = 1usize << 31; + let model = AudioModel { afrac: 12, gf: 16, shift: 0, layers: vec![ALayer::ConvT1d { out_ch: 1, k: 1, stride: s, pad: s, opad: 0, weight: vec![7], mult: vec![1], bias: vec![0] }] }; + let g = ops::load(&model).expect("the model lowers and loads"); + let go = |p: Policy| fast_conv_runs(|| with_policy(p, || crate::int_ops::run1(&g, Tensor::new(1, 1, 3, vec![1, 2, 3]), &[1]).map_err(|r| r.code))); + assert!(matches!(go(Policy::reference()), (Err(Code::TooLarge), 0))); + for p in FAST { + assert!(matches!(go(p), (Err(Code::TooLarge), 0)), "{p:?}"); } } - /// Geometry whose padded input would be more than `PAD_FACTOR` (16) times what - /// the reference reads and writes is refused by the fast path, and the reference - /// runs: the fast kernel is not entered and the samples are the reference's. At - /// exactly the factor the fast kernel runs. A one-sample input, kernel 1, stride - /// 1, cropped by `pad` with output padding `2·pad` restoring one output sample - /// that no tap reaches: the padded input is `cin · (1 + pad)` samples against - /// `cin + cin + 1` the reference touches (input, weights, output). + /// Geometry whose zero-padded input would be more than the fast kernels' pad factor + /// (16) times what the reference reads and writes is declined by the host, and the + /// interpreter computes it; at exactly the factor the fast kernel runs. A one-sample + /// input, kernel 1, stride 1, cropped by `pad` with output padding `2·pad` restoring + /// one output sample that no tap reaches: its value is the bias, 7. #[test] - fn padding_beyond_the_pad_factor_runs_the_reference() { - use crate::int_kernels::testkit::fast_conv_runs; - // (channels, pad, expected fast-kernel runs): 1 · 48 = 16 · 3 and - // 2 · 40 = 16 · 5 at the factor; 1 · 49 and 2 · 41 beyond it. + fn padding_beyond_the_pad_factor_runs_on_the_interpreter() { + // (channels, pad, fast-kernel runs): 1 · 48 = 16 · 3 and 2 · 40 = 16 · 5 at the + // factor; 1 · 49 and 2 · 41 beyond it. for (cin, pad, runs) in [(1usize, 47usize, 1u64), (1, 48, 0), (2, 39, 1), (2, 40, 0)] { - let (x, wt) = (vec![5i64; cin], vec![3i64; cin]); - let want = reference::conv_t_1d(&x, cin, 1, &wt, 1, 1, 1, pad, 2 * pad, &[3], &[7], 1); - assert_eq!(want, (vec![7], 1), "reference, {cin} channels, pad {pad}"); - for p in FAST { - let got = fast_conv_runs(|| with_policy(p, || conv_t_1d(&x, cin, 1, &wt, 1, 1, 1, pad, 2 * pad, &[3], &[7], 1))); - assert_eq!(got, (want.clone(), runs), "{cin} channels, pad {pad} {p:?}"); - } + let g = one(cin, 1, 1, pad, 2 * pad, vec![3; cin], 3, 7, 1); + let (out, n) = run(&g, &vec![5; cin], 1); + assert_eq!(out, Ok(vec![pcm(7)]), "{cin} channels, pad {pad}"); + assert!(n.iter().all(|&r| r == runs), "{cin} channels, pad {pad}: fast runs {n:?}"); } } + + /// An input on which the scalar kernel's sums would overflow `i64` only at the + /// positions it accumulates into and then crops (two channels of 2^50, taps 2^14, 1, + /// 2^14 — `WAM1`'s weights are 16-bit — pad 1: the kept middle sums to 2^51, each + /// cropped border to 2^65): through the lowering, the layer's `check` range refuses it + /// before any arithmetic, under every policy (`field_range`), as the direct decode's + /// guard did; one within the range runs. + #[test] + fn an_input_that_could_overflow_is_refused_before_the_kernel() { + let g = one(2, 3, 1, 1, 0, [1 << 14, 1, 1 << 14, 1 << 14, 1, 1 << 14].to_vec(), 1, 0, 0); + assert_eq!(run(&g, &[1 << 50, 1 << 50], 1).0, Err(Code::FieldRange)); + assert_eq!(run(&g, &[1, 1], 1).0, Ok(vec![pcm(2)])); + } } diff --git a/wai/wai-rs/src/int_hyper_synth.rs b/wai/wai-rs/src/int_hyper_synth.rs index c1acd61c7..321ef3db1 100644 --- a/wai/wai-rs/src/int_hyper_synth.rs +++ b/wai/wai-rs/src/int_hyper_synth.rs @@ -480,6 +480,13 @@ impl std::fmt::Display for NetRefusal { /// a network whose worst case over all inputs would. A decoder that feeds values it /// read from an untrusted stream into a network uses this, and refuses the stream on /// `Err`. +/// +/// A network that decodes through its lowering in the integer operator set +/// (`extensions/integer-operators.md` §13.2) makes these checks as the `check` ranges of +/// its graph instead, which refuse exactly what these refuse, at the same layer. The +/// splat and audio decodes do (`int_splat`, `int_audio`; their direct decodes, which +/// call this or [`igdn_check`], are test-only); the conditional codec's checked decode +/// ([`decode_conditional_checked`]) still calls this, and goes with it when it moves. pub fn run_net_checked( model: &IntModel, input: &[i64], mut c: usize, mut h: usize, mut w: usize, ) -> Result<(Vec, usize, usize, usize), NetRefusal> { @@ -647,7 +654,7 @@ fn check_layer(model: &IntModel, layer: &IntLayer, x: &[i64], c: usize, h: usize /// The IGDN checks [`run_net_checked`] documents, for activations `x` of `c` channels /// at `2^afrac` against `gamma` (`c × c`, at `2^gf`) and `beta` (`c`). Shared with the -/// 1-D synthesis of `int_audio`, whose IGDN is the same [`igdn`]. +/// test-only direct decode of `int_audio`, whose IGDN is the same [`igdn`]. pub(crate) fn igdn_check(x: &[i64], c: usize, gamma: &[i64], beta: &[i64], gf: u32, afrac: u32) -> Result<(), Failed> { use crate::int_kernels::max_abs; const I64: u128 = i64::MAX as u128; diff --git a/wai/wai-rs/src/int_ops.rs b/wai/wai-rs/src/int_ops.rs index 18744dd26..ccc78cdaf 100644 --- a/wai/wai-rs/src/int_ops.rs +++ b/wai/wai-rs/src/int_ops.rs @@ -25,8 +25,10 @@ use wai_ops::graph::{Graph, Loaded, Op, Params}; use wai_ops::interp::{Kernels, Sequence, Tensor}; +use wai_ops::lower::{Net, NetLayer}; use wai_ops::{Profile, Refusal}; +use crate::int_hyper_synth::{IntLayer, IntModel}; use crate::int_kernels::{policy, ConvGeom, KernelPath}; /// The profile of every capability that decodes through the operator set @@ -40,10 +42,42 @@ pub fn load(g: &Graph) -> Result { wai_ops::graph::load(&g.to_bytes(), PROFILE) } +/// A `WHM1` network parsed by this crate ([`IntModel`]) in the form `wai-ops` lowers it +/// from ([`Net`]): the one [`wai_ops::lower::parse_whm1`] gives of the bytes the model was +/// parsed from, so [`wai_ops::lower::whm1_net`] of it writes the graph +/// [`wai_ops::lower::whm1`] writes from those bytes. The model does not keep a layer's +/// input channels: a convolution's are its weights' count over its output channels and +/// kernel, as the bytes give them, and an inverse GDN's are its `β`'s count. A layer +/// with no output channels or a kernel of 0 is given 0, and the lowering refuses it. A +/// model parsed from bytes holds only 32-bit fields; one with a field past 32 bits +/// (another source) is refused here, as the lowering refuses one (`model_mismatch`). +pub fn whm1_net(m: &IntModel) -> Result { + let f = |v: usize| u32::try_from(v).map_err(|_| Refusal::new(wai_ops::Code::ModelMismatch, wai_ops::Site::Lowering)); + let inc = |w: &[i64], out: usize, k: usize| { + f(out.checked_mul(k).and_then(|v| v.checked_mul(k)).and_then(|d| w.len().checked_div(d)).unwrap_or(0)) + }; + let layers = m.layers.iter().map(|l| { + Ok(match l { + IntLayer::ConvT { out_ch, k, stride, pad, opad, weight, mult, bias } => NetLayer::ConvT { + inc: inc(weight, *out_ch, *k)?, out: f(*out_ch)?, k: f(*k)?, stride: f(*stride)?, pad: f(*pad)?, opad: f(*opad)?, + w: weight.clone(), mult: mult.clone(), bias: bias.clone(), + }, + IntLayer::Conv { out_ch, k, stride, pad, weight, mult, bias } => NetLayer::Conv { + inc: inc(weight, *out_ch, *k)?, out: f(*out_ch)?, k: f(*k)?, stride: f(*stride)?, pad: f(*pad)?, + w: weight.clone(), mult: mult.clone(), bias: bias.clone(), + }, + IntLayer::Igdn { gamma, beta } => NetLayer::Igdn { c: f(beta.len())?, gamma: gamma.clone(), beta: beta.clone() }, + IntLayer::Relu => NetLayer::Relu, + IntLayer::Leaky { slope_q, slope_frac } => NetLayer::Leaky { slope_q: *slope_q, slope_frac: *slope_frac }, + }) + }); + Ok(Net { afrac: m.afrac, gf: m.gf, shift: m.shift, layers: layers.collect::>()?, audio: false }) +} + /// One decode of a loaded graph with one data input and no state, through the /// interpreter with this crate's kernels behind it ([`Host`]): its first output. pub fn run1(g: &Loaded, x: Tensor, scalars: &[i64]) -> Result { - let mut out = wai_ops::run_with(g, &[x], scalars, &mut Sequence::new(), &Host)?; + let mut out = wai_ops::run_with(g, vec![x], scalars, &mut Sequence::new(), &Host)?; Ok(out.swap_remove(0)) } @@ -178,6 +212,95 @@ mod equivalence { assert!(fast > 600, "the fast kernel ran only {fast} times"); } + /// The one-dimensional transposed convolutions of `wam1/1` (§13.3: `kh = 1`, `sh = 1`, + /// `ph = 0`, `oph = 0` over `[cin, 1, T]`), with `WAM1`'s requantisation (a multiplier + /// per channel, one shift, a `post` bias, `floor`): random kernels up to 9, strides up + /// to 5, lengths up to 40, and operands that fit `i32` and that do not. The fast + /// kernel behind the interpreter gives the interpreter's integers. This is + /// `int_audio`'s equivalence of its fast one-dimensional kernel with its scalar one, + /// held to the interpreter instead. + #[test] + fn one_dimensional_transposed_convolutions_equal_the_interpreter() { + let mut r = Rng(0x0D2C_0003); + let (mut loaded, mut fast) = (0, 0); + for case in 0..400 { + let k = 1 + r.below(9) as u32; + let s = 1 + r.below(5) as u32; + let pad = r.below(k as usize) as u32; + let op = if s > 1 { r.below(s as usize) as u32 } else { 0 }; + let t = 1 + r.below(40) as u32; + let (cin, cout) = (1 + r.below(6) as u32, 1 + r.below(6) as u32); + if (t - 1) * s + k + op <= 2 * pad { + continue; + } + let (xmax, wmax) = [(1i64 << 15, 2047i64), (1 << 40, 1 << 10)][case % 2]; + let n = cout as usize; + let rq = Rq { + round: Round::Floor, bias_kind: BiasKind::Post, m_kind: 2, s_kind: 1, m_bytes: 8, b_bytes: 8, + mult: (0..n).map(|_| 1 + r.signed(1 << 16).abs()).collect(), shift: vec![24], bias: r.vec(n, 1 << 16, 0), + }; + let wt = r.vec((cin * cout * k) as usize, wmax, 5); + let p = Params::ConvT { cout, kh: 1, kw: k, sh: 1, sw: s, ph: 0, pw: pad, oph: 0, opw: op, w_bytes: 8, rq, w: wt }; + let g = graph(TensorRec::check(cin, -xmax, xmax), opcode::CONV_TRANSPOSE, cout, p); + let x = Tensor::new(cin, 1, t, r.vec((cin * t) as usize, xmax, 30)); + let (l, f) = agree(&g, &x); + loaded += l as u32; + fast += f; + } + assert!(loaded > 250, "only {loaded} of the random graphs loaded"); + assert!(fast > 500, "the fast kernel ran only {fast} times"); + } + + /// The committed `wai.splat.int_codec` synthesis through its lowering `whm1/1`: under + /// every fast policy each of its two transposed convolutions runs on the fast kernel, + /// and every attribute byte is the interpreter's. Its inverse GDN runs on the + /// interpreter under every policy: this host lends no IGDN kernel. + #[cfg(feature = "splat_int")] + #[test] + fn the_splat_synthesis_runs_its_convolutions_on_the_fast_kernel() { + use crate::int_splat::{ops, SplatStream}; + let dir = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../wai-rs/tests/vectors/int_splat"); + let gs = IntModel::from_bytes(&std::fs::read(dir.join("splat_gs.bin")).unwrap()).unwrap(); + let p = std::fs::read(dir.join("splat.wis1")).unwrap(); + let v = crate::int_wire::parse_wis1(&p).unwrap(); + let st = SplatStream { + afrac: v.afrac as u32, c: v.c as usize, sh: v.sh as usize, sw: v.sw as usize, s: v.s as usize, + n_attr: v.attr as usize, cdfs: &v.cdfs, offsets: &v.offsets, median_q: &v.median_q, rans: v.rans, + }; + let (want, none) = fast_conv_runs(|| with_policy(Policy::reference(), || ops::decode_typed(&gs, &st).unwrap().attr)); + assert_eq!(none, 0, "the reference policy runs no fast kernel"); + for p in FAST { + let (got, runs) = fast_conv_runs(|| with_policy(p, || ops::decode_typed(&gs, &st).unwrap().attr)); + assert!(got == want, "{p:?}: the attribute bytes differ from the interpreter alone"); + assert_eq!(runs, 2, "{p:?}: both transposed convolutions run on the fast kernel"); + } + } + + /// The committed `wai.audio.int_codec` synthesis through its lowering `wam1/1`: under + /// every fast policy each of its four one-dimensional transposed convolutions runs on + /// the fast kernel, and every sample is the interpreter's. Its inverse GDNs and its + /// PCM step run on the interpreter under every policy. + #[cfg(feature = "audio_int")] + #[test] + fn the_audio_synthesis_runs_its_convolutions_on_the_fast_kernel() { + use crate::int_audio::{ops, AudioModel, AudioStream}; + let dir = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../wai-rs/tests/vectors/int_audio"); + let m = AudioModel::from_bytes(&std::fs::read(dir.join("audio_model.bin")).unwrap()).unwrap(); + let p = std::fs::read(dir.join("audio_bitstream.wia1")).unwrap(); + let a = crate::int_wire::parse_wia1(&p).unwrap(); + let s = AudioStream { + afrac: a.afrac as u32, n_ch: a.n_ch as usize, lat_len: a.lat_len as usize, t_out: a.t_out as usize, + offset: a.offset, median_q: &a.median_q, cdfs: &a.cdfs, rans: a.rans, + }; + let (want, none) = fast_conv_runs(|| with_policy(Policy::reference(), || ops::decode_typed(&m, &s).unwrap())); + assert_eq!(none, 0, "the reference policy runs no fast kernel"); + for p in FAST { + let (got, runs) = fast_conv_runs(|| with_policy(p, || ops::decode_typed(&m, &s).unwrap())); + assert!(got == want, "{p:?}: the samples differ from the interpreter alone"); + assert_eq!(runs, 4, "{p:?}: all four transposed convolutions run on the fast kernel"); + } + } + /// A value outside an input's `check` range is refused by the interpreter before any /// kernel runs, under every policy alike. #[test] diff --git a/wai/wai-rs/src/int_payload/mod.rs b/wai/wai-rs/src/int_payload/mod.rs index 62cd667a4..b3fcf1e30 100644 --- a/wai/wai-rs/src/int_payload/mod.rs +++ b/wai/wai-rs/src/int_payload/mod.rs @@ -537,14 +537,16 @@ fn cost_wiv1(cap: &str, _: &[u8], _: &IntParams<'_>) -> Result) -> Result { - use crate::int_audio::{audio_plan, AudioModel, AudioStream}; + use crate::int_audio::{AudioModel, AudioStream}; let a = wire::parse_wia1(payload)?; let m = AudioModel::from_bytes(params.member(cap, "model.wam")?).map_err(|e| model(format!("model.wam: {e}")))?; let s = AudioStream { afrac: a.afrac as u32, n_ch: a.n_ch as usize, lat_len: a.lat_len as usize, t_out: a.t_out as usize, offset: a.offset, median_q: &a.median_q, cdfs: &a.cdfs, rans: a.rans, }; - let (_, working, macs) = audio_plan(&m, &s)?; + // Everything the decode checks before its first symbol, its set lowered and loaded + // included, so a payload the decode refuses there is refused here with its code. + let (_, working, macs) = crate::int_audio::ops::plan(&m, &s)?; // The parsed medians and tables, held for the whole decode. let tables = a.n_ch as u64 * (8 + VEC + 4 * a.cdfs[0].len() as u64); Ok(DeclaredCost { output_bytes: 2 * a.t_out as u64, symbols: a.n_ch as u64 * a.lat_len as u64, working_bytes: tables + working, macs }) @@ -558,14 +560,16 @@ fn cost_wia1(cap: &str, _: &[u8], _: &IntParams<'_>) -> Result) -> Result { use crate::int_hyper_synth::IntModel; - use crate::int_splat::{splat_plan, SplatStream}; + use crate::int_splat::SplatStream; let v = wire::parse_wis1(payload)?; let gs = IntModel::from_bytes(params.member(cap, "g_s.whm")?).map_err(|e| model(format!("g_s.whm: {e}")))?; let s = SplatStream { afrac: v.afrac as u32, c: v.c as usize, sh: v.sh as usize, sw: v.sw as usize, s: v.s as usize, n_attr: v.attr as usize, cdfs: &v.cdfs, offsets: &v.offsets, median_q: &v.median_q, rans: v.rans, }; - let (working, macs) = splat_plan(&gs, &s)?; + // Everything the decode checks before its first symbol, its set lowered and loaded + // included, so a payload the decode refuses there is refused here with its code. + let (_, working, macs) = crate::int_splat::ops::plan(&gs, &s)?; // The parsed ranges, offsets, medians and tables, held for the whole decode. let cdfs: u64 = v.cdfs.iter().map(|t| 4 * t.len() as u64).sum(); let tables = 8 * v.attr as u64 + v.c as u64 * (4 + 8 + VEC) + cdfs; diff --git a/wai/wai-rs/src/int_splat.rs b/wai/wai-rs/src/int_splat.rs index 9eb63a0cb..3f065f0e8 100644 --- a/wai/wai-rs/src/int_splat.rs +++ b/wai/wai-rs/src/int_splat.rs @@ -3,15 +3,16 @@ //! color[3], opacity), arranged as a `[14, S, S]` attribute image. A decode that runs //! float arithmetic can round differently on different machines, and then the //! reconstructed cloud — and its hash — differ between them. Here the synthesis is -//! integer (the [`crate::int_hyper_synth`] machinery, 14 channels instead of RGB's 3) -//! over a per-channel factorized-rANS latent, all i64 — so the decoded attribute -//! tensor is byte-identical on every machine: **splat-equivalence**, a portable cloud -//! hash. The float rasterizer is sink-supplied presentation; the *attributes* are the -//! conformance contract. +//! integer (a `WHM1` network, 14 channels instead of RGB's 3) over a per-channel +//! factorized-rANS latent, all exact integers — so the decoded attribute tensor is +//! byte-identical on every machine: **splat-equivalence**, a portable cloud hash. The +//! synthesis decodes through the integer operator set, lowered by `whm1/1` +//! (`extensions/integer-operators.md` §13.2; [`crate::int_ops`]). The float rasterizer is +//! sink-supplied presentation; the *attributes* are the conformance contract. use serde_json::Value; -use crate::int_hyper_synth::{fnv1a, refused, run_net_checked, IntModel, Refused, RefusalKind, MAX_AFRAC, MAX_TENSOR}; +use crate::int_hyper_synth::{fnv1a, refused, IntModel, Refused, RefusalKind, MAX_AFRAC, MAX_TENSOR}; use crate::rans::RansDecoder; /// A decoded splat: the `[attr · s · s]` `u8` attribute image (channel-major), to be @@ -117,23 +118,32 @@ impl SplatJson { /// 1. the stream's `afrac` is the model's, in `1..=30`; the medians, offsets and /// tables are one per channel and every table is well formed; `g_s` maps the /// latent's shape to `[n_attr, s, s]` with no buffer over -/// [`crate::int_hyper_synth::MAX_TENSOR`] elements — all before anything is -/// allocated; +/// [`crate::int_hyper_synth::MAX_TENSOR`] elements; then `g_s` is lowered by +/// `whm1/1` and loaded under `acc64` with its bounds certificate +/// (`extensions/integer-operators.md` §13.2) — all before anything is allocated +/// for the payload; /// 2. per channel, rANS-decode the symbol and rebuild `s·2^afrac + median_q` /// ([`crate::int_entropy::dequantize`]); a stream that runs out before its last /// symbol is refused; -/// 3. integer synthesis `g_s` through [`crate::int_hyper_synth::run_net_checked`]; -/// 4. round-half-up → `u8`, channel-major `[attr, s, s]`. +/// 3. integer synthesis `g_s` through the integer operator set ([`crate::int_ops`]), +/// whose `check` ranges refuse a latent that could take a layer outside the exact +/// integer range (`field_range`), ending with the RGB8 step: each value clamped to +/// `[0, 2^afrac]` and rounded half up to `u8`, channel-major `[attr, s, s]`. +/// +/// This is [`ops::decode_typed`]. The direct decode it replaced is kept, test-only +/// (feature `legacy_int`), as `legacy::decode_typed`. pub fn decode_typed(gs: &IntModel, st: &SplatStream<'_>) -> Result { + ops::decode_typed(gs, st) +} + +/// Step 2 of the decode, after its plan: per channel, the factorized rANS decode of the +/// latent, each symbol rebuilt as `s·2^afrac + median_q`. The latent grows a symbol at a +/// time, and the first symbol that runs the stream out refuses the payload before its +/// value is used; the stream must then end canonically. +fn latent(st: &SplatStream<'_>, afrac: u32) -> Result, Refused> { use crate::int_entropy::dequantize; use RefusalKind::*; - splat_plan(gs, st)?; - let (afrac, c, sh, sw, s, n_attr) = (gs.afrac, st.c, st.sh, st.sw, st.s, st.n_attr); - - // 1. per-channel factorized rANS-decode the latent, rebuild y_q. The latent grows - // a symbol at a time, and the first symbol that runs the stream out refuses the - // payload before its value is used. - let plane = sh * sw; + let (c, plane) = (st.c, st.sh * st.sw); let mut y = Vec::new(); let mut dec = RansDecoder::new(st.rans).map_err(|_| refused(FieldRange, "the latent stream is not a rANS stream"))?; for ch in 0..c { @@ -148,26 +158,118 @@ pub fn decode_typed(gs: &IntModel, st: &SplatStream<'_>) -> Result> afrac).clamp(0, 255) as u8).collect(); - Ok(DecodedSplat { attr, n_attr, s }) + Ok(y) } -/// What [`decode_typed`] checks before it decodes a symbol, and the most bytes the -/// decode then holds at once on the reference kernels beyond the payload and the -/// parsed model: the latent (`i64`) with `g_s`'s buffers while it runs, then `g_s`'s -/// output (`i64`) with the attribute tensor (`u8`); and `g_s`'s multiply-accumulates -/// ([`crate::int_hyper_synth::NetPlan`]). +/// The direct decode, as the decoder ran before it decoded through the integer operator +/// set: the synthesis `g_s` run layer by layer by +/// [`crate::int_hyper_synth::run_net_checked`], which checks each layer's arithmetic on +/// the values it is given before it runs, then the attribute bytes formed here. +/// Test-only (`legacy_int`): the differential tests decode both ways, and the two differ +/// only on the planned breaks of `extensions/integer-operators.md` §15.3. +#[cfg(any(test, feature = "legacy_int"))] +pub mod legacy { + use super::*; + use crate::int_hyper_synth::run_net_checked; + + /// The direct decode of [`super::decode_typed`]'s arguments. + pub fn decode_typed(gs: &IntModel, st: &SplatStream<'_>) -> Result { + use RefusalKind::*; + splat_plan(gs, st)?; + let (afrac, c, sh, sw, s, n_attr) = (gs.afrac, st.c, st.sh, st.sw, st.s, st.n_attr); + let y = latent(st, afrac)?; + // Integer synthesis → attribute tensor at 2^afrac. + let (x, ..) = run_net_checked(gs, &y, c, sh, sw).map_err(|r| { + if r.out_of_range { + refused(FieldRange, "a decoded latent could take a layer outside the exact integer range") + } else { + refused(ModelMismatch, "g_s's parameters do not fit the running shapes") + } + })?; + drop(y); + // Round-half-up → u8 (channel-major [attr, s, s]). + let (one, half) = (1i64 << afrac, 1i64 << (afrac - 1)); + let attr = x.iter().map(|&v| ((v.clamp(0, one) * 255 + half) >> afrac).clamp(0, 255) as u8).collect(); + Ok(DecodedSplat { attr, n_attr, s }) + } +} + +/// The decode through the integer operator set (`extensions/integer-operators.md`): +/// `g_s` lowered by `whm1/1` with the RGB8 step (§13.2, §13.0), loaded under `acc64` with +/// its bounds certificate, and run by the operator set's interpreter +/// ([`crate::int_ops`]), whose `check` ranges carry the exactness checks of +/// integer-payloads §7. The driver keeps integer-payloads §7's order: the plan of +/// [`splat_plan`], then the set lowered and loaded, then the stream, then the graph. +pub mod ops { + use super::*; + use wai_ops::{Code, Loaded, Refusal, Tensor}; + + /// A refusal of the operator set as the splat decode reports it + /// (`extensions/integer-payloads.md` §8): a value outside a `check` range is + /// `field_range`; the other codes are the payload extension's of the same name. + fn kind(r: Refusal, what: &'static str) -> Refused { + let k = match r.code { + Code::FieldRange => RefusalKind::FieldRange, + Code::TooLarge => RefusalKind::TooLarge, + Code::GeometryMismatch => RefusalKind::GeometryMismatch, + Code::StreamExhausted => RefusalKind::StreamExhausted, + Code::StreamNoncanonical => RefusalKind::StreamNoncanonical, + Code::ModelMismatch => RefusalKind::ModelMismatch, + }; + refused(k, what) + } + + /// `g_s` lowered by `whm1/1` with the RGB8 step and loaded under `acc64`: the set as + /// the decode holds it, beside the parsed model. + pub fn load(gs: &IntModel) -> Result { + const WHAT: &str = "g_s does not lower to an integer-operator graph that loads under acc64"; + let net = crate::int_ops::whm1_net(gs).map_err(|r| kind(r, WHAT))?; + let g = wai_ops::lower::whm1_net(&net, true).map_err(|r| kind(r, WHAT))?; + drop(net); + crate::int_ops::load(&g).map_err(|r| kind(r, WHAT)) + } + + /// Everything the decode checks before it decodes a symbol: [`splat_plan`], then `g_s` + /// lowered and loaded. Returns the loaded graph, and [`splat_plan`]'s working bytes and + /// multiply-accumulates. + pub(crate) fn plan(gs: &IntModel, st: &SplatStream<'_>) -> Result<(Loaded, u64, u128), Refused> { + let (working, macs) = splat_plan(gs, st)?; + Ok((load(gs)?, working, macs)) + } + + /// [`super::decode_typed`] through the integer operator set: the same attribute bytes, + /// and the same refusals, except where integer-operators §15.3 lists a change. + pub fn decode_typed(gs: &IntModel, st: &SplatStream<'_>) -> Result { + let (g, ..) = plan(gs, st)?; + decode_loaded(&g, gs.afrac, st) + } + + /// The decode after its plan: the stream at `afrac`, then `g`, the set's loaded + /// lowering ([`load`]), then the bytes. `st` must be a cloud the plan admitted + /// against that set. What it holds beyond the set is what the payload declares + /// (integer-payloads §7, *Memory and work*). + pub fn decode_loaded(g: &Loaded, afrac: u32, st: &SplatStream<'_>) -> Result { + let (c, sh, sw, s, n_attr) = (st.c, st.sh, st.sw, st.s, st.n_attr); + let y = latent(st, afrac)?; + // The plan admitted every size, so each fits u32. + let x = Tensor::new(c as u32, sh as u32, sw as u32, y); + let out = crate::int_ops::run1(g, x, &[]).map_err(|r| kind(r, "a decoded latent is outside a check range of g_s"))?; + if (out.c as usize, out.h as usize, out.w as usize) != (n_attr, s, s) { + return Err(refused(RefusalKind::GeometryMismatch, "g_s does not give the attribute shape the payload declares")); + } + // The RGB8 step ends the graph, so every value is a byte. + let attr = out.data.iter().map(|&v| v as u8).collect(); + Ok(DecodedSplat { attr, n_attr, s }) + } +} + +/// What [`decode_typed`] checks before it decodes a symbol, short of lowering `g_s` +/// ([`ops::plan`] adds that), and the most bytes the decode then holds at once on the +/// reference kernels beyond the payload and the set as it holds it (the parsed model, +/// and its loaded lowering): the latent (`i64`) with `g_s`'s buffers while it runs (a +/// layer's input, its output and a transposed convolution's uncropped plane), then +/// `g_s`'s output (`i64`) with the attribute tensor (`u8`); and `g_s`'s +/// multiply-accumulates ([`crate::int_hyper_synth::NetPlan`]). pub(crate) fn splat_plan(gs: &IntModel, st: &SplatStream<'_>) -> Result<(u64, u128), Refused> { use RefusalKind::*; let afrac = gs.afrac; diff --git a/wai/wai-rs/tests/int_audio_ops.rs b/wai/wai-rs/tests/int_audio_ops.rs new file mode 100644 index 000000000..257af6dc8 --- /dev/null +++ b/wai/wai-rs/tests/int_audio_ops.rs @@ -0,0 +1,499 @@ +//! `wai.audio.int_codec` through the integer operator set: its lowering `wam1/1` +//! (`extensions/integer-operators.md` §13.3, with the PCM step) and the certificate of +//! every committed set (§8.5), 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, from +//! the model it parses (`int_audio::ops::wam1_net`), must write those same bytes, and the +//! load a sink makes must give the verdict the certificate records. +//! +//! The decode through the lowering (`int_audio::ops`) must give, on every committed +//! vector, the samples of the direct decode (`int_audio::legacy`) and the vector's +//! committed digests. + +use std::path::PathBuf; + +use serde_json::Value; +use wai::int_audio::{self, AudioJson, AudioModel, AudioStream}; +use wai_ops::Profile; + +const CAP: &str = "wai.audio.int_codec"; + +fn wai() -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("..") +} + +fn read(rel: &str) -> Vec { + std::fs::read(wai().join(rel)).unwrap_or_else(|e| panic!("{rel}: {e}")) +} + +fn json(rel: &str) -> Value { + serde_json::from_slice(&read(rel)).unwrap_or_else(|e| panic!("{rel}: {e}")) +} + +fn sha256(b: &[u8]) -> String { + wai_ops::sha256::sha256(b).iter().map(|x| format!("{x:02x}")).collect() +} + +/// The `wai.audio.int_codec` sets `ops-conformance/certificates/index.json` lists: +/// `(name, synthetic, set file, set digest, lowered graph file)`, each relative to `wai/`. +fn indexed() -> Vec<(String, bool, String, String, String)> { + json("ops-conformance/certificates/index.json")["sets"] + .as_array() + .unwrap() + .iter() + .filter(|s| s["capability"] == CAP) + .map(|s| { + assert_eq!(s["lowering"], "wam1/1"); + let roles = s["roles"].as_array().unwrap(); + assert_eq!(roles.len(), 1, "wam1/1 lowers an audio set to one graph, g_s"); + ( + s["name"].as_str().unwrap().to_string(), + s["synthetic"].as_bool().unwrap(), + s["members"]["model.wam"].as_str().unwrap().to_string(), + s["set_sha256"].as_str().unwrap().to_string(), + roles[0].as_str().unwrap().to_string(), + ) + }) + .collect() +} + +/// The graph this crate lowers from the model it parses: `int_audio::ops::wam1_net`, +/// then `wam1/1` with the PCM step. +fn lowered_from_the_parsed_model(set: &[u8]) -> wai_ops::Graph { + let m = AudioModel::from_bytes(set).expect("the WAM1 set parses"); + let net = wai::int_audio::ops::wam1_net(&m).expect("a model parsed from bytes converts"); + assert!(net == wai_ops::lower::parse_wam1(set).unwrap(), "the converted model is the one wai-ops parses from the bytes"); + wai_ops::lower::wam1_net(&net).expect("the set lowers") +} + +/// Gate 1 of the migration: every committed `wai.audio.int_codec` set (the conformance +/// vector's and the synthetic one of `ops-conformance/sets/`) lowers to its committed +/// graph, from its bytes (`lower::wam1`) and from the model this crate parses; the +/// certificate beside it is that graph's; and loading the graph under the capability's +/// profile gives the certificate's verdict. Each set carries `check` ranges (§13.2, +/// §13.3), so it meets `acc64` by construction and not `acc32`, first failing at `g_s`'s +/// input range (the vector's at 47 bits, as Appendix C predicts); and the PCM step's +/// `crop`, whose length is a scalar, keeps it out of the accelerator subset. +#[test] +fn every_committed_set_lowers_to_its_committed_graph_and_verdict() { + let sets = indexed(); + let mut names: Vec<(&str, bool)> = sets.iter().map(|s| (s.0.as_str(), s.1)).collect(); + names.sort(); + assert_eq!(names, [("audio", false), ("audio", true)], "the committed audio sets"); + for (name, synthetic, set, digest, role) in &sets { + let name = if *synthetic { format!("{name} (synthetic)") } else { name.clone() }; + let bytes = read(set); + let g = wai_ops::lower::wam1(&bytes).unwrap_or_else(|r| panic!("{name}: does not lower: {r:?}")); + // The verdict first: the load a sink makes. + let loaded = wai::int_ops::load(&g).unwrap_or_else(|r| panic!("{name}: refused under acc64: {r:?}")); + assert!(!loaded.meets(Profile::Acc32), "{name}: a set with check ranges inside its graph does not meet acc32"); + // Then the committed bytes: the set, its lowered graph (both ways) 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 of the bytes differs from the committed {role}"); + assert!(lowered_from_the_parsed_model(&bytes).to_bytes() == graph, "{name}: the lowering of the parsed model differs from the committed {role}"); + let cert = json(&format!("ops-conformance/certificates/{CAP}/{digest}.json")); + assert_eq!(cert["set_sha256"].as_str(), Some(digest.as_str()), "{name}"); + assert_eq!(cert["lowering"], "wam1/1", "{name}"); + assert_eq!(cert["graphs"][0]["name"], "g_s", "{name}"); + assert_eq!(cert["graphs"][0]["graph_sha256"].as_str(), Some(sha256(&graph).as_str()), "{name}: the certificate is the graph's"); + let v = &cert["verdict"]; + assert_eq!((v["acc64"].as_bool(), v["acc32"].as_bool()), (Some(true), Some(false)), "{name}: the certificate's verdict"); + let first = &v["acc32_first_failure"]; + assert_eq!( + (first["graph"].as_str(), first["where"].as_str(), first["index"].as_str(), first["key"].as_str()), + (Some("g_s"), Some("input"), Some("0"), Some("range")), + "{name}: the first acc32 failure is g_s's input range" + ); + assert!(v["npu_reasons"].as_array().unwrap().iter().any(|r| r == "crop_scalar"), "{name}: the PCM step's crop"); + let x0 = &cert["graphs"][0]["inputs"][0]; + assert_eq!(x0["policy"].as_str(), Some("check"), "{name}"); + assert_eq!(x0["hi"].as_str(), Some(g.tensors[0].hi.to_string().as_str()), "{name}: the input's range is X_0"); + if !synthetic { + assert_eq!(first["bits"].as_str(), Some("47"), "{name}: Appendix C's 47 bits"); + assert_eq!(g.tensors[0].hi, 56322496561155, "{name}"); + } + // The scalar is t_out, `[1, 2^26]` (§13.0's driver contract). + assert_eq!((g.scalars.len(), g.scalars[0].lo, g.scalars[0].hi), (1, 1, 1 << 26), "{name}: scalar 0 is t_out"); + } +} + +/// The typed form of a parsed `WIA1` payload. +fn stream<'a>(a: &'a wai::int_wire::Wia1<'a>) -> AudioStream<'a> { + AudioStream { + afrac: a.afrac as u32, n_ch: a.n_ch as usize, lat_len: a.lat_len as usize, t_out: a.t_out as usize, + offset: a.offset, median_q: &a.median_q, cdfs: &a.cdfs, rans: a.rans, + } +} + +/// The `wai.audio.int_codec` vectors `int-payload-conformance/vectors.json` indexes: +/// `(name, payload, set member model.wam, JSON form, sidecar)`, each relative to `wai/`. +fn vectors() -> Vec<(String, String, String, String, String)> { + json("int-payload-conformance/vectors.json")["vectors"] + .as_array() + .unwrap() + .iter() + .filter(|v| v["capability"] == CAP) + .map(|v| { + let s = |k: &Value| k.as_str().unwrap().to_string(); + (s(&v["name"]), s(&v["payload"]), s(&v["set"]["model.wam"]), s(&v["converted_from"]["json"]), s(&v["sidecar"])) + }) + .collect() +} + +fn le_bytes(pcm: &[i16]) -> Vec { + pcm.iter().flat_map(|v| v.to_le_bytes()).collect() +} + +/// Gate 2 of the migration: golden equality. Every committed `wai.audio.int_codec` +/// vector decodes through the lowering (`int_audio::ops`) to the samples the direct +/// decode gives (`int_audio::legacy`), and to every digest the corpus commits for it: +/// the SHA-256 sidecar of the PCM and the export's FNV-1a. The payload (`WIA1`) and the +/// JSON form it is converted from both decode, and alike. The staged audio base carries +/// the same payload, so it is covered too. +#[test] +fn every_committed_vector_decodes_through_the_lowering_to_its_golden() { + let vs = vectors(); + assert_eq!(vs.len(), 1, "the committed audio vectors"); + for (name, payload, set, from, sidecar) in &vs { + let m = AudioModel::from_bytes(&read(set)).unwrap(); + let p = read(payload); + let a = wai::int_wire::parse_wia1(&p).unwrap_or_else(|e| panic!("{name}: {e:?}")); + let j = AudioJson::from_value(&json(from)).unwrap(); + let want = json(sidecar)[name]["sha256"].as_str().unwrap().to_string(); + let fnv = json("wai-rs/tests/vectors/int_audio/audio.json")["pcm_fnv1a"].as_str().unwrap().to_string(); + for (form, st) in [("WIA1", stream(&a)), ("JSON", j.stream())] { + let ops = int_audio::ops::decode_typed(&m, &st).unwrap_or_else(|r| panic!("{name} {form}: {r}")); + let direct = int_audio::legacy::decode_typed(&m, &st).unwrap_or_else(|r| panic!("{name} {form}: {r}")); + assert!(ops == direct, "{name} {form}: the decode through the lowering differs from the direct decode"); + assert_eq!(ops.len(), st.t_out, "{name} {form}: t_out samples"); + assert_eq!(sha256(&le_bytes(&ops)), want, "{name} {form}: the SHA-256 sidecar"); + assert_eq!(format!("{:016x}", int_audio::pcm_fnv1a(&ops)), fnv, "{name} {form}: the export's FNV-1a"); + } + let base = read("staged-conformance/cases/audio_int_codec_1l/base.wai"); + assert!(base.windows(p.len()).any(|w| w == p.as_slice()), "{name}: the staged audio base carries this payload"); + } +} + +/// The committed vector's set and payload (`int-payload-conformance`'s `audio`). +const SET: &str = "wai-rs/tests/vectors/int_audio/audio_model.bin"; +const PAYLOAD: &str = "wai-rs/tests/vectors/int_audio/audio_bitstream.wia1"; + +/// A typed clip that owns its tables, so a fixture can change any field. +#[derive(Clone)] +struct Owned { + afrac: u32, + n_ch: usize, + lat_len: usize, + t_out: usize, + offset: i32, + median_q: Vec, + cdfs: Vec>, + rans: Vec, +} + +impl Owned { + fn committed() -> Owned { + let p = read(PAYLOAD); + let a = wai::int_wire::parse_wia1(&p).unwrap(); + Owned { + afrac: a.afrac as u32, n_ch: a.n_ch as usize, lat_len: a.lat_len as usize, t_out: a.t_out as usize, + offset: a.offset, median_q: a.median_q.clone(), cdfs: a.cdfs.clone(), rans: a.rans.to_vec(), + } + } + /// `n_ch` channels of `lat_len` latents at `afrac` 12, every latent of channel `k` + /// equal to `values[k]` (one table symbol each, its median placing it), presenting + /// `t_out` samples. + fn constant(values: &[i64], lat_len: usize, t_out: usize) -> Owned { + let (t, offset) = (vec![0, 1 << 15, 1 << 16], 3); + let mut enc = wai::rans::RansEncoder::new(); + for _ in 0..values.len() * lat_len { + assert!(enc.encode(offset, &t, 3, offset)); + } + Owned { + afrac: 12, n_ch: values.len(), lat_len, t_out, offset, + median_q: values.iter().map(|v| v - ((offset as i64) << 12)).collect(), cdfs: vec![t; values.len()], rans: enc.flush(), + } + } + fn st(&self) -> AudioStream<'_> { + AudioStream { + afrac: self.afrac, n_ch: self.n_ch, lat_len: self.lat_len, t_out: self.t_out, offset: self.offset, + median_q: &self.median_q, cdfs: &self.cdfs, rans: &self.rans, + } + } +} + +/// The committed set, parsed, changed by `f`. +fn model(f: impl FnOnce(&mut AudioModel)) -> AudioModel { + let mut m = AudioModel::from_bytes(&read(SET)).unwrap(); + f(&mut m); + m +} + +/// The committed clip, changed by `f`. +fn clip(f: impl FnOnce(&mut Owned)) -> Owned { + let mut o = Owned::committed(); + f(&mut o); + o +} + +/// The code integer-payloads §8 gives a refusal of the audio decode. +fn code_of(r: wai::int_hyper_synth::Refused) -> &'static str { + wai::int_payload::IntPayloadError::from(r).code() +} + +/// Both decodes of a clip against a model: through the lowering, then directly. +fn both(m: &AudioModel, o: &Owned) -> (Result, &'static str>, Result, &'static str>) { + (int_audio::ops::decode_typed(m, &o.st()).map_err(code_of), int_audio::legacy::decode_typed(m, &o.st()).map_err(code_of)) +} + +/// Gate 3 of the migration: refusal parity, on the committed fixtures. Each +/// `wai.audio.int_codec` fixture of `int-payload-conformance` is refused by a sink with +/// its code; one that parses (`decode`) is refused by the decode through the lowering and +/// by the direct decode alike, with that code. +#[test] +fn every_committed_refusal_fixture_keeps_its_code() { + let mut n = 0; + for f in json("int-payload-conformance/vectors.json")["refusals"].as_array().unwrap().iter().filter(|f| f["capability"] == CAP) { + let (name, want) = (f["name"].as_str().unwrap(), f["code"].as_str().unwrap()); + let (p, set) = (read(f["payload"].as_str().unwrap()), read(f["set"]["model.wam"].as_str().unwrap())); + let params = wai::int_payload::resolve(CAP, None, &[wai::pin::HeldBytes::File(&set)], None).expect("an unpinned set resolves"); + let sink = wai::int_payload::decode_int(CAP, &p, ¶ms).map(|_| ()).map_err(|e| e.code()); + assert_eq!(sink, Err(want), "{name}: the sink's code"); + if f["layer"] == "decode" { + let m = AudioModel::from_bytes(&set).unwrap(); + let a = wai::int_wire::parse_wia1(&p).unwrap_or_else(|e| panic!("{name}: a decode fixture parses: {e:?}")); + let ops = int_audio::ops::decode_typed(&m, &stream(&a)).err().map(code_of); + let direct = int_audio::legacy::decode_typed(&m, &stream(&a)).err().map(code_of); + assert_eq!((ops, direct), (Some(want), Some(want)), "{name}: through the lowering, then directly"); + } + n += 1; + } + assert!(n >= 14, "the committed audio fixtures: {n}"); +} + +/// Gate 3, beside the committed fixtures: a fixture for each refusal the direct decode +/// makes (integer-payloads §5 and §7, in their order), each made from the committed set +/// and clip by one change. Each is refused by both decodes with the same code; the code +/// it states is integer-payloads §8's for that rule. The committed set and clip decode. +#[test] +fn every_refusal_of_the_direct_decode_is_refused_alike_through_the_lowering() { + use int_audio::ALayer; + let (m, base) = (model(|_| {}), Owned::committed()); + assert!(int_audio::ops::decode_typed(&m, &base.st()).is_ok(), "the fixtures' base decodes"); + let same = || model(|_| {}); + let last = |m: &mut AudioModel| m.layers.len() - 1; + let fixtures: Vec<(&str, &str, AudioModel, Owned)> = vec![ + ("the set's afrac of 0", "model_mismatch", model(|m| m.afrac = 0), clip(|_| {})), + ("the set's afrac of 31", "model_mismatch", model(|m| m.afrac = 31), clip(|_| {})), + ("a clip afrac not the set's", "model_mismatch", same(), clip(|o| o.afrac = 11)), + ("no latent channels", "field_range", same(), clip(|o| o.n_ch = 0)), + ("a latent length of 0", "field_range", same(), clip(|o| o.lat_len = 0)), + ("a t_out of 0", "field_range", same(), clip(|o| o.t_out = 0)), + ("a median short", "field_range", same(), clip(|o| { o.median_q.pop(); })), + ("a table short", "field_range", same(), clip(|o| { o.cdfs.pop(); })), + ("a median past 2^47", "field_range", same(), clip(|o| o.median_q[3] = -(1 << 47) - 1)), + ("a table that does not end at 2^16", "cdf_malformed", same(), clip(|o| *o.cdfs[2].last_mut().unwrap() = 65535)), + ("a table of another length", "cdf_malformed", same(), clip(|o| { o.cdfs[5].pop(); *o.cdfs[5].last_mut().unwrap() = 1 << 16; })), + ("a latent past 2^26", "too_large", same(), clip(|o| o.lat_len = 1 << 23)), + ("a t_out past 2^26", "too_large", same(), clip(|o| o.t_out = (1 << 26) + 1)), + ("an uncropped buffer past 2^26", "too_large", same(), clip(|o| (o.lat_len, o.t_out) = (1 << 22, 1 << 25))), + ("a layer's bias short", "model_mismatch", model(|m| { + if let ALayer::ConvT1d { bias, .. } = &mut m.layers[0] { bias.pop(); } + }), clip(|_| {})), + ("a stride of 0", "model_mismatch", model(|m| { + if let ALayer::ConvT1d { stride, .. } = &mut m.layers[2] { *stride = 0; } + }), clip(|_| {})), + ("an inverse GDN short of a beta", "model_mismatch", model(|m| { + if let ALayer::Igdn { beta, .. } = &mut m.layers[1] { beta.pop(); } + }), clip(|_| {})), + ("padding that leaves no output", "model_mismatch", model(|m| { + let l = last(m); + if let ALayer::ConvT1d { pad, .. } = &mut m.layers[l] { *pad = 1 << 20; } + }), clip(|_| {})), + ("a synthesis of 16 channels", "model_mismatch", model(|m| { m.layers.pop(); }), clip(|_| {})), + ("a t_out past the synthesis", "geometry_mismatch", same(), clip(|o| o.t_out += 1)), + ("a latent longer than t_out needs", "geometry_mismatch", same(), clip(|o| o.t_out -= 512)), + ("a stream too short to be one", "field_range", same(), clip(|o| o.rans.truncate(4))), + ("a stream that runs out", "stream_exhausted", same(), clip(|o| { let n = o.rans.len() / 2; o.rans.truncate(n - n % 4); })), + ("a word left unread", "stream_noncanonical", same(), clip(|o| o.rans.extend([0, 0, 0, 0]))), + ("every latent past layer 0's range", "field_range", same(), clip(|o| o.median_q.iter_mut().for_each(|m| *m = 1 << 47))), + ("every latent below it", "field_range", same(), clip(|o| o.median_q.iter_mut().for_each(|m| *m = -(1 << 47)))), + ]; + for (what, want, m, o) in &fixtures { + let (ops, direct) = both(m, o); + assert_eq!((ops.err(), direct.err()), (Some(*want), Some(*want)), "{what}: through the lowering, then directly"); + } +} + +/// Gate 3, at the edge of every `check` range of the committed set: for each of its +/// layers, a network of that layer alone (with the set's header), followed by a +/// transposed convolution to one channel whose weights and bias are 0 (whose own range +/// is all of `i64`), decodes a clip whose latents all lie in `[−X_ℓ, X_ℓ]`, with one +/// channel at `X_ℓ` (or `−X_ℓ`), to the same samples both ways, and refuses one at +/// `X_ℓ + 1` (or `−X_ℓ − 1`) both ways (`field_range`), where `X_ℓ` is tensor `ℓ`'s +/// range in the committed graph. And the PCM step's range: a network adding a bias to +/// its one channel decodes a sample at the top of the range in which +/// `v·32767 + 2^(afrac−1)` fits `i64`, and the bottom, and refuses one past each, both +/// ways. +#[test] +fn every_check_range_refuses_exactly_what_the_direct_decode_refuses() { + use int_audio::ALayer; + let full = wai_ops::lower::wam1(&read(SET)).unwrap(); + let n = AudioModel::from_bytes(&read(SET)).unwrap().layers.len(); + for l in 0..n { + let one = model(|m| { + let layer = m.layers.swap_remove(l); + let c = match &layer { + ALayer::ConvT1d { out_ch, .. } => *out_ch, + ALayer::Igdn { c, .. } => *c, + }; + m.layers = vec![layer, ALayer::ConvT1d { out_ch: 1, k: 1, stride: 1, pad: 0, opad: 0, weight: vec![0; c], mult: vec![1], bias: vec![0] }]; + }); + let g = wai_ops::lower::wam1_net(&int_audio::ops::wam1_net(&one).unwrap()).unwrap(); + let (c, x) = (g.tensors[0].channels as usize, g.tensors[0].hi); + assert_eq!((g.tensors[0].lo, x), (full.tensors[l].lo, full.tensors[l].hi), "layer {l}: the range of tensor {l}"); + // The clip's length: the layer's output length on 4 latents. + let t = match &one.layers[0] { + ALayer::ConvT1d { k, stride, pad, opad, .. } => 3 * stride + k + opad - 2 * pad, + ALayer::Igdn { .. } => 4, + }; + let at = |v: i64| { + let mut vals = vec![0i64; c]; + vals[c / 2] = v; + Owned::constant(&vals, 4, t) + }; + for v in [x, -x] { + let (ops, direct) = both(&one, &at(v)); + assert!(ops.is_ok() && ops == direct, "layer {l}: at {v}: {ops:?} through the lowering, {direct:?} directly"); + } + for v in [x + 1, -x - 1] { + assert_eq!(both(&one, &at(v)), (Err("field_range"), Err("field_range")), "layer {l}: one past the range, {v}"); + } + } + // The PCM step: one channel plus a bias. + let lo = -((1i128 << 63) / 32767) as i64; + let hi = ((i64::MAX as i128 - (1 << 11)) / 32767) as i64; + let add = |b: i64| AudioModel { + afrac: 12, gf: 16, shift: 0, + layers: vec![ALayer::ConvT1d { out_ch: 1, k: 1, stride: 1, pad: 0, opad: 0, weight: vec![1], mult: vec![1], bias: vec![b] }], + }; + let one = Owned::constant(&[5 << 12], 4, 4); + for (b, ok) in [(hi - (5 << 12), Ok(vec![32767; 4])), (lo - (5 << 12), Ok(vec![-32768; 4])), (hi - (5 << 12) + 1, Err("field_range")), (lo - (5 << 12) - 1, Err("field_range"))] { + let (ops, direct) = both(&add(b), &one); + assert_eq!((&ops, &direct), (&ok, &ok), "a sample at {}", 5 * 4096 + b); + } +} + +/// 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 samples, and the same refusals, on every committed vector, +/// as `WIA1` and as its JSON form, and on clips with a channel moved to the input's +/// check range and one past it, at each end. The per-operator half is `int_ops`' +/// equivalence tests, one of which counts the fast kernel's runs on this set. +#[test] +fn every_kernel_policy_decodes_through_the_lowering_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 m = model(|_| {}); + let x0 = wai_ops::lower::wam1(&read(SET)).unwrap().tensors[0].hi; + let at = |t: i64| clip(move |o| o.median_q[1] = t - ((o.offset as i64) << o.afrac)); + let j = AudioJson::from_value(&json("wai-rs/tests/vectors/int_audio/audio_bitstream.json")).unwrap(); + let mut clips: Vec<(String, Owned)> = vec![ + ("WIA1".into(), Owned::committed()), + ("JSON".into(), Owned { + afrac: j.afrac, n_ch: j.n_ch, lat_len: j.lat_len, t_out: j.t_out, offset: j.offset, + median_q: j.median_q.clone(), cdfs: j.cdfs.clone(), rans: j.rans.clone(), + }), + ]; + for (what, t) in [("channel 1's median at X_0", x0), ("at X_0 + 1", x0 + 1), ("at -X_0", -x0), ("at -X_0 - 1", -x0 - 1)] { + clips.push((what.into(), at(t))); + } + for (k, (what, o)) in clips.iter().enumerate() { + let want = with_policy(POLICIES[0], || int_audio::ops::decode_typed(&m, &o.st())).map_err(code_of); + assert!(k >= 2 || want.is_ok(), "{what}: the committed form decodes"); + for pol in &POLICIES[1..] { + let got = with_policy(*pol, || int_audio::ops::decode_typed(&m, &o.st())).map_err(code_of); + assert!(got == want, "{what}: {pol:?} differs from the interpreter alone"); + } + } +} + +/// The committed `wai.audio.int_codec` sets and payloads, relative to `wai/`: every set +/// and payload `int-payload-conformance/vectors.json` names for the capability (its +/// vectors and its refusal fixtures), and every base of a staged case decoded through +/// the capability, with the set its case names and the payload its envelope carries. +/// Returns `(set, (where, payload))` pairs. +fn committed_sets_and_payloads() -> Vec<(String, (String, Vec))> { + let mut v = Vec::new(); + let index = json("int-payload-conformance/vectors.json"); + for e in index["vectors"].as_array().unwrap().iter().chain(index["refusals"].as_array().unwrap()) { + if e["capability"] == CAP { + let (set, p) = (e["set"]["model.wam"].as_str().unwrap(), e["payload"].as_str().unwrap()); + v.push((set.to_string(), (p.to_string(), read(p)))); + } + } + let cases = wai().join("staged-conformance/cases"); + let mut dirs: Vec<_> = std::fs::read_dir(&cases).unwrap().filter_map(|e| e.ok()).map(|e| e.path()).collect(); + dirs.sort(); + for d in dirs { + let Ok(b) = std::fs::read(d.join("case.json")) else { continue }; + let case: Value = serde_json::from_slice(&b).unwrap(); + for base in case["bases"].as_array().into_iter().flatten().filter(|b| b["decoded_via"] == CAP) { + let f = d.join(base["file"].as_str().unwrap()); + let env = wai::container::Wai::from_bytes(&std::fs::read(&f).unwrap()).unwrap(); + assert_eq!(env.manifest.model_requirement.capability, CAP, "{}", f.display()); + v.push((base["decode"]["set"]["model.wam"].as_str().unwrap().to_string(), (f.display().to_string(), env.payload))); + } + } + v +} + +/// The scan for the load rules the switch brings (integer-operators §15.3 item 3): no +/// committed set or payload is affected. Every committed audio set (the one the +/// payloads decode against, and the synthetic one of `ops-conformance/sets/`) has +/// layers, header shifts under 128, no bias of magnitude `2^63 − 1` or more, and lowers +/// and loads under `acc64`, its PCM step's range included; every committed payload that +/// the direct decode decodes decodes through the lowering to the same samples, and every +/// one it refuses is refused with the same code. (These are the files that carry a `WIA1` +/// payload or name the capability's set; `WIA1` appears in no other committed payload or +/// envelope.) +#[test] +fn the_new_load_rules_affect_no_committed_set_or_payload() { + use wai_ops::lower::NetLayer as N; + let all = committed_sets_and_payloads(); + let mut sets: Vec = all.iter().map(|(s, _)| s.clone()).collect(); + sets.extend(indexed().into_iter().map(|s| s.2)); + sets.sort_unstable(); + sets.dedup(); + for set in &sets { + let bytes = read(set); + let net = wai_ops::lower::parse_wam1(&bytes).unwrap_or_else(|r| panic!("{set}: {r:?}")); + assert!(!net.layers.is_empty(), "{set}: no layers"); + assert!(net.shift < 128 && net.gf < 128, "{set}: a header shift of 128 or more"); + for l in &net.layers { + if let N::ConvT { bias, .. } = l { + assert!(bias.iter().all(|b| b.unsigned_abs() < i64::MAX as u64), "{set}: a bias of magnitude 2^63 - 1 or more"); + } + } + let g = wai_ops::lower::wam1(&bytes).unwrap_or_else(|r| panic!("{set}: does not lower: {r:?}")); + assert!(wai::int_ops::load(&g).is_ok(), "{set}: does not load under acc64"); + } + let (mut decoded, mut refused) = (0, 0); + for (set, (at, p)) in &all { + let m = AudioModel::from_bytes(&read(set)).unwrap(); + let Ok(a) = wai::int_wire::parse_wia1(p) else { + refused += 1; + continue; + }; + let ops = int_audio::ops::decode_typed(&m, &stream(&a)).map_err(code_of); + let direct = int_audio::legacy::decode_typed(&m, &stream(&a)).map_err(code_of); + assert!(ops == direct, "{at}: the decode through the lowering differs from the direct decode"); + if ops.is_ok() { decoded += 1 } else { refused += 1 } + } + assert_eq!(sets.len(), 2, "the committed audio sets: {sets:?}"); + assert!(decoded >= 2 && refused >= 14, "the committed audio payloads, at least those of this revision: {decoded} decoded, {refused} refused"); +} diff --git a/wai/wai-rs/tests/int_kernels_geometry.rs b/wai/wai-rs/tests/int_kernels_geometry.rs index 3aa4479e6..ccda016fc 100644 --- a/wai/wai-rs/tests/int_kernels_geometry.rs +++ b/wai/wai-rs/tests/int_kernels_geometry.rs @@ -18,9 +18,9 @@ //! and every policy must do the same. int_synth decodes through its lowering //! (integer-operators §13.1), whose fields are 32 bits and whose run-time plan caps //! every size it forms at 2^26: where a case is past either, every policy refuses it, -//! one of the planned breaks of integer-operators §15.3. The int_audio 1-D transposed -//! convolution has the -//! same stride check beside its kernel (`int_audio` unit tests). CI runs this file +//! one of the planned breaks of integer-operators §15.3. int_audio decodes through its +//! lowering `wam1/1` too; its 1-D transposed convolution's edge geometry is tested +//! through that graph under every policy (`int_audio` unit tests). CI runs this file //! beside `int_kernels_vectors` on every `wai-conformance` leg, with overflow checks //! on and off. #![cfg(feature = "neural_int")] @@ -224,7 +224,7 @@ fn synth_payload_2x2(h: u32, w: u32, latents: [i8; 4]) -> Vec { // below.) A fast policy must have // the same outcome: it may not return a value where the reference panics. A 1×1 // input cannot show this, since `(h − 1)·stride` is 0. The int_audio 1-D form is -// tested beside its kernel. +// tested through its lowering's graph (`int_audio` unit tests). /// A 3×3 input, and a 2×2 one with fields that fit a WHM1 model's `u32`. const X: [i64; 9] = [1, 2, 3, 4, 5, 6, 7, 8, 9]; diff --git a/wai/wai-rs/tests/int_ops_differential.rs b/wai/wai-rs/tests/int_ops_differential.rs index 9d8b7c80c..ff1c2b9b0 100644 --- a/wai/wai-rs/tests/int_ops_differential.rs +++ b/wai/wai-rs/tests/int_ops_differential.rs @@ -6,16 +6,22 @@ //! alone, and a case that falls under a planned break must be refused by the decode //! through the operator set, whatever the direct decode does there. //! +//! The decoders: `wai.neural.int_synth` (`int_synth/1`), `wai.splat.int_codec` +//! (`whm1/1`) and `wai.audio.int_codec` (`wam1/1`). +//! //! The random cases run for a fixed count by default. `WAI_DIFF_SECS` runs each random //! test for that many seconds instead (the nightly run, `wai-ops-differential.yml`), and //! `WAI_DIFF_SEED` sets the first seed; a failing case names its seed. -#![cfg(feature = "legacy_int")] +#![cfg(all(feature = "legacy_int", feature = "splat_int", feature = "audio_int"))] use std::panic::catch_unwind; use std::path::PathBuf; use std::time::{Duration, Instant}; +use std::panic::AssertUnwindSafe; + use serde_json::{json, Value}; +use wai::int_hyper_synth::{IntModel, RefusalKind}; use wai::int_synth::{decode_from_spec, legacy, IntImage, IntSynthError}; fn wai() -> PathBuf { @@ -296,3 +302,683 @@ fn int_synth_planned_breaks_are_new_refusals() { assert_eq!(int_synth_planned_break(&five, &s64), Some("shift 64 to 127")); assert_eq!(decode_from_spec(&five, &s64).unwrap().rgb, [0, 0, 0]); } + +// ------------------------------------------------------------------ int_splat (M2) + +/// One layer of a `WHM1` network, as a set's bytes hold it (integer-payloads §9). +#[derive(Clone)] +enum Whm { + ConvT { inc: u32, out: u32, k: u32, s: u32, p: u32, op: u32, w: Vec, m: Vec, b: Vec }, + Conv { inc: u32, out: u32, k: u32, s: u32, p: u32, w: Vec, m: Vec, b: Vec }, + Igdn { c: u32, g: Vec, b: Vec }, + Relu, + Leaky { q: i64, f: u32 }, +} + +/// A `WHM1` network's bytes. +fn whm1_bytes(afrac: u32, gf: u32, shift: u32, layers: &[Whm]) -> Vec { + let mut v = b"WHM1".to_vec(); + for x in [afrac, gf, shift, layers.len() as u32] { + v.extend(x.to_le_bytes()); + } + let words = |v: &mut Vec, xs: &[u32]| xs.iter().for_each(|x| v.extend(x.to_le_bytes())); + for l in layers { + match l { + Whm::ConvT { inc, out, k, s, p, op, w, m, b } => { + v.push(0); + words(&mut v, &[*inc, *out, *k, *s, *p, *op]); + w.iter().for_each(|x| v.extend(x.to_le_bytes())); + m.iter().chain(b).for_each(|x| v.extend(x.to_le_bytes())); + } + Whm::Conv { inc, out, k, s, p, w, m, b } => { + v.push(1); + words(&mut v, &[*inc, *out, *k, *s, *p]); + w.iter().for_each(|x| v.extend(x.to_le_bytes())); + m.iter().chain(b).for_each(|x| v.extend(x.to_le_bytes())); + } + Whm::Igdn { c, g, b } => { + v.push(2); + words(&mut v, &[*c]); + g.iter().for_each(|x| v.extend(x.to_le_bytes())); + b.iter().for_each(|x| v.extend(x.to_le_bytes())); + } + Whm::Relu => v.push(3), + Whm::Leaky { q, f } => { + v.push(4); + v.extend(q.to_le_bytes()); + v.extend(f.to_le_bytes()); + } + } + } + v +} + +/// A `WIS1` cloud's typed fields, owned. +#[derive(Clone, Debug)] +struct Cloud { + afrac: u32, + c: usize, + sh: usize, + sw: usize, + s: usize, + n_attr: usize, + cdfs: Vec>, + offsets: Vec, + median_q: Vec, + rans: Vec, +} + +impl Cloud { + fn st(&self) -> wai::int_splat::SplatStream<'_> { + wai::int_splat::SplatStream { + afrac: self.afrac, c: self.c, sh: self.sh, sw: self.sw, s: self.s, n_attr: self.n_attr, + cdfs: &self.cdfs, offsets: &self.offsets, median_q: &self.median_q, rans: &self.rans, + } + } +} + +/// The planned break of integer-operators §15.3 that a splat set falls under, if any, +/// read from the set's bytes alone: the set is refused at load by its lowering +/// (`whm1/1` with the RGB8 step, loaded under `acc64`), for a rule the direct decode +/// does not make. A set refused at load for a rule integer-payloads §7's plan also makes +/// (the chain, a zero dimension, a shift a layer uses) is not a break: the plan, which +/// both decodes run first, refuses it alike. No payload field changes meaning. +fn int_splat_planned_break(set: &[u8]) -> Option<&'static str> { + let net = wai_ops::lower::parse_whm1(set).ok()?; + load_break(&net, wai_ops::lower::whm1(set, true)) +} + +/// The planned break a parsed `WHM1` or `WAM1` network falls under, given its lowering +/// (see [`int_splat_planned_break`]). +fn load_break(net: &wai_ops::lower::Net, lowered: Result) -> Option<&'static str> { + use wai_ops::lower::NetLayer as N; + let loaded = lowered.as_ref().map_err(|r| *r).and_then(wai::int_ops::load); + if loaded.is_ok() { + return None; + } + let conv = net.layers.iter().any(|l| matches!(l, N::ConvT { .. } | N::Conv { .. })); + let igdn = net.layers.iter().any(|l| matches!(l, N::Igdn { .. })); + if !conv && !igdn { + return Some("no layer of tag 0, 1 or 2"); + } + if (net.shift > 127 && !conv) || (net.gf > 127 && !igdn) { + return Some("a shift or gf of 128 or more that no layer uses"); + } + if (1..=30).contains(&net.afrac) && net.shift <= 127 && net.gf <= 127 { + let huge = |b: &[i64]| b.iter().any(|v| v.unsigned_abs() >= i64::MAX as u64); + if net.layers.iter().any(|l| matches!(l, N::ConvT { bias, .. } | N::Conv { bias, .. } if huge(bias))) { + return Some("a bias of magnitude 2^63 - 1 or more"); + } + } + match loaded { + Err(r) if lowered.is_ok() && r.at == wai_ops::Site::Load(8) => Some("a check range its producer cannot meet"), + _ => None, + } +} + +/// A random layer stack for `int_splat_case`, from `c` input channels: transposed +/// convolutions, convolutions, inverse GDNs, ReLUs and leaky ReLUs; kernels 1 to 4, +/// strides 1 to 3, paddings mostly below the kernel. Weights of up to 2^4, 2^10 or 2^12 +/// (one set in eight up to i16's range), a quarter 0; multipliers mostly up to 2^12 (one +/// in 32 up to 2^40), biases mostly up to 2^20, with one in 64 at an end of i64; IGDN `γ` mostly +/// non-negative and up to 2^12, `β` up to 2^28. Returns the layers and the output +/// channels. +fn splat_layers(r: &mut Rng, mut c: u32, n: usize, only_relu: bool) -> (Vec, u32) { + let big = r.below(8) == 0; + let mut layers = Vec::new(); + for _ in 0..n { + let kind = if only_relu { 3 } else { [0, 0, 0, 1, 2, 2, 3, 4][r.below(8) as usize] }; + let l = match kind { + 0 | 1 => { + let out = 1 + r.below(4) as u32; + let k = 1 + r.below(4) as u32; + let s = 1 + r.below(3) as u32; + let p = if r.below(16) == 0 { r.below(k as u64 + 3) } else { r.below(k as u64) } as u32; + let wmax = if big { 32767 } else { [16, 1024, 4096][r.below(3) as usize] }; + let w: Vec = (0..c * out * k * k).map(|_| if r.below(4) == 0 { 0 } else { r.signed(wmax) as i16 }).collect(); + let m: Vec = (0..out).map(|_| if r.below(32) == 0 { r.signed(1 << 40) } else { r.signed(1 << 12) }).collect(); + let b: Vec = (0..out) + .map(|_| match r.below(64) { + 0 => [i64::MAX, i64::MIN, i64::MIN + 1, i64::MAX - 1][r.below(4) as usize], + 1 => r.signed(1 << 40), + _ => r.signed(1 << 20), + }) + .collect(); + let inc = c; + c = out; + if kind == 0 { + let op = if r.below(16) == 0 { r.below(s as u64 + 1) } else { r.below(s as u64) } as u32; + Whm::ConvT { inc, out, k, s, p, op, w, m, b } + } else { + Whm::Conv { inc, out, k, s, p, w, m, b } + } + } + 2 => { + let g: Vec = (0..c * c).map(|_| if r.below(16) == 0 { r.signed(1 << 12) as i32 } else { r.below(1 << 12) as i32 }).collect(); + let b: Vec = (0..c).map(|_| if r.below(16) == 0 { r.signed(1 << 28) } else { r.below(1 << 28) as i64 }).collect(); + Whm::Igdn { c, g, b } + } + 3 => Whm::Relu, + _ => Whm::Leaky { q: if r.below(32) == 0 { r.signed(1 << 40) } else { r.signed(1 << 10) }, f: if r.below(32) == 0 { 64 } else { r.below(16) as u32 } }, + }; + layers.push(l); + } + (layers, c) +} + +/// The output side of `layers` on a square input of side `n`, or `None` where the +/// geometry leaves no output. +fn splat_side(layers: &[Whm], mut n: i64) -> Option { + for l in layers { + n = match l { + Whm::ConvT { k, s, p, op, .. } => (n - 1) * *s as i64 + *k as i64 + *op as i64 - 2 * *p as i64, + Whm::Conv { k, s, p, .. } => { + let np = n + 2 * *p as i64 - *k as i64; + if np < 0 { return None } else { np / *s as i64 + 1 } + } + _ => n, + }; + if n < 1 { + return None; + } + } + Some(n) +} + +/// A random `wai.splat.int_codec` set and cloud: a network of 1 to 4 layers (one set in +/// 32 with no layers, one in 32 of ReLUs only) at `afrac` 12 mostly (one set in sixteen +/// anywhere in 1..=30, one in sixteen at 0 or 31), a shift of 8 to 23, with a header `gf` +/// or `shift` of 128 or more one set in 32 each; a latent of 1 to 3 channels of side 1 to 4 (one in sixteen +/// not square), each channel with its own table of 3 to 6 entries, offset and median; the +/// symbols mostly inside their table, one in eight escaped past it and one in 32 far +/// past it; medians mostly small, one in eight at or past the 2^47 bound, and one in four +/// placing a whole channel's latents at the input's `check` range or one past it. The +/// declared attribute shape is the network's output (one case in eight off by one); the +/// cloud's `afrac` is the set's (one in 32 not); one stream in sixteen is cut short, one +/// in 32 has a word appended, one in 64 a byte changed. +fn int_splat_case(r: &mut Rng) -> (Vec, Cloud) { + let afrac = match r.below(16) { + 0 => 1 + r.below(30) as u32, + 1 => [0, 31][r.below(2) as usize], + _ => 12, + }; + let gf = if r.below(32) == 0 { 128 + r.below(64) as u32 } else { 16 }; + let shift = if r.below(32) == 0 { 128 + r.below(64) as u32 } else { 8 + r.below(16) as u32 }; + let c0 = 1 + r.below(3) as u32; + let n_layers = if r.below(32) == 0 { 0 } else { 1 + r.below(4) as usize }; + let only_relu = r.below(32) == 0; + let (layers, c_out) = splat_layers(r, c0, n_layers, only_relu); + let set = whm1_bytes(afrac, gf, shift, &layers); + + let sh = 1 + r.below(4) as usize; + let sw = if r.below(16) == 0 { 1 + r.below(4) as usize } else { sh }; + let mut s = splat_side(&layers, sh as i64).unwrap_or(1) as usize; + let mut n_attr = c_out as usize; + if r.below(8) == 0 { + if r.below(2) == 0 { s += 1 } else { n_attr += 1 } + } + let c = c0 as usize; + let a = afrac.min(30); + let mut cdfs = Vec::new(); + let mut offsets = Vec::new(); + let mut median_q = Vec::new(); + for _ in 0..c { + let len = 3 + r.below(4) as usize; + let mut cuts: Vec = (0..len - 2).map(|_| 1 + r.below(65_535) as i32).collect(); + cuts.sort_unstable(); + cuts.dedup(); + let mut t = vec![0]; + t.extend(cuts); + t.push(1 << 16); + cdfs.push(t); + offsets.push(r.signed(3) as i32); + median_q.push(match r.below(8) { + 0 => [1i64 << 47, -(1 << 47), (1 << 47) + 1, -(1 << 47) - 1][r.below(4) as usize], + _ => r.signed(1 << (a + 4)), + }); + } + let plane = sh * sw; + let mut values: Vec = Vec::with_capacity(c * plane); + for ch in 0..c { + let span = cdfs[ch].len() as i64 - 2; + for _ in 0..plane { + let v = match r.below(32) { + 0 => r.signed(1 << 20), + 1..=4 => offsets[ch] as i64 + span + r.below(40) as i64 - 20, + _ => offsets[ch] as i64 + r.below(span as u64) as i64, + }; + values.push(v as i32); + } + } + // One case in four: a whole channel's latents at the input's check range, or one + // past it, where the set lowers. + if r.below(4) == 0 { + if let Ok(g) = wai_ops::lower::whm1(&set, true) { + let x0 = g.tensors[0].hi as i128; + let ch = r.below(c as u64) as usize; + let v = offsets[ch]; + let target = [x0, x0 + 1, -x0, -x0 - 1][r.below(4) as usize]; + let med = target - ((v as i128) << a); + if med.unsigned_abs() <= 1 << 47 { + median_q[ch] = med as i64; + values[ch * plane..(ch + 1) * plane].iter_mut().for_each(|x| *x = v); + } + } + } + let mut enc = wai::rans::RansEncoder::new(); + for ch in 0..c { + for &v in &values[ch * plane..(ch + 1) * plane] { + assert!(enc.encode(v, &cdfs[ch], cdfs[ch].len(), offsets[ch]), "a value the generator writes is codable"); + } + } + let mut rans = enc.flush(); + match r.below(64) { + 0..=3 => { + let cut = (4 * (1 + r.below(3)) as usize).min(rans.len()); + rans.truncate(rans.len() - cut); + } + 4 | 5 => rans.extend([0, 0, 0, 0]), + 6 => { + let i = r.below(rans.len() as u64) as usize; + rans[i] ^= 1 << r.below(8); + } + _ => {} + } + let cloud_afrac = if r.below(32) == 0 { afrac ^ 1 } else { afrac }; + (set, Cloud { afrac: cloud_afrac, c, sh, sw, s, n_attr, cdfs, offsets, median_q, rans }) +} + +/// The direct splat decode, a panic (an overflow-checked build's arithmetic overflow) +/// caught. +fn splat_direct(gs: &IntModel, cl: &Cloud) -> Option, RefusalKind>> { + catch_unwind(AssertUnwindSafe(|| wai::int_splat::legacy::decode_typed(gs, &cl.st()).map(|d| d.attr).map_err(|r| r.kind))).ok() +} + +/// Gate 4 of the migration, for `wai.splat.int_codec`: the decode through the operator +/// set against the direct decode on random sets and clouds. Outside the planned breaks +/// both decode to the same attribute bytes or both refuse with the same code, and the +/// direct decode never panics; under a planned break the decode through the operator set +/// refuses. +#[test] +fn int_splat_random_sets_and_clouds_decode_alike() { + let (mut alike, mut refused, mut breaks) = (0u64, std::collections::BTreeMap::::new(), std::collections::BTreeMap::<&str, u64>::new()); + let n = cases("int_splat", 3000, |r, seed| { + let (set, cl) = int_splat_case(r); + let why = || format!("seed {seed}\ncloud {cl:?}"); + let gs = IntModel::from_bytes(&set).expect("the generator writes WHM1 that parses"); + let new = wai::int_splat::ops::decode_typed(&gs, &cl.st()).map(|d| d.attr).map_err(|r| r.kind); + match int_splat_planned_break(&set) { + Some(b) => { + assert!(new.is_err(), "{b}: the decode through the operator set must refuse it\n{}", why()); + *breaks.entry(b).or_default() += 1; + } + None => match (&new, splat_direct(&gs, &cl)) { + (Ok(a), Some(Ok(b))) => { + assert!(*a == b, "the two decodes differ\n{}", why()); + alike += 1; + } + (Err(a), Some(Err(b))) => { + assert_eq!(*a, b, "the two decodes refuse with different codes\n{}", why()); + *refused.entry(format!("{a:?}")).or_default() += 1; + } + (_, old) => panic!("new {new:?}, direct {old:?}\n{}", why()), + }, + } + }); + eprintln!("int_splat: {alike} decoded alike, refused by both {refused:?}, planned breaks {breaks:?}"); + assert!(alike * 5 > n, "the generator must produce decodable cases: {alike} of {n}"); + assert!(n < 1000 || refused.contains_key("FieldRange"), "no case reached a check range"); + assert!(n < 1000 || breaks.len() >= 3, "too few kinds of planned break reached: {breaks:?}"); +} + +/// The planned breaks of integer-operators §15.3 item 3 for `wai.splat.int_codec`, each +/// on a set and cloud the direct decode decodes, or refuses with another code: the +/// classifier names it, and the decode through the operator set refuses it +/// (`model_mismatch`, before the stream). +#[test] +fn int_splat_planned_breaks_are_new_refusals() { + // A cloud of 2 channels of side 2, every latent 2·2^12, each symbol the first of the + // table `t` (offset 2). + let cloud_with = |t: Vec| { + let (c, n, s, n_attr) = (2usize, 2usize, 2usize, 2usize); + let mut enc = wai::rans::RansEncoder::new(); + for _ in 0..c * n * n { + assert!(enc.encode(2, &t, 3, 2)); + } + Cloud { afrac: 12, c, sh: n, sw: n, s, n_attr, cdfs: vec![t; c], offsets: vec![2; c], median_q: vec![0; c], rans: enc.flush() } + }; + let conv1 = |inc: u32, out: u32, w: i16, b: i64| Whm::ConvT { inc, out, k: 1, s: 1, p: 0, op: 0, w: vec![w; (inc * out) as usize], m: vec![1; out as usize], b: vec![b; out as usize] }; + let expect = |b: &str, set: &[u8], cl: &Cloud, direct: Option| { + assert_eq!(int_splat_planned_break(set), Some(b), "{b}: the classifier"); + let gs = IntModel::from_bytes(set).unwrap(); + match direct { + None => assert!(matches!(splat_direct(&gs, cl), Some(Ok(_))), "{b}: the direct decode decodes it"), + Some(k) => assert_eq!(splat_direct(&gs, cl), Some(Err(k)), "{b}: the direct decode's code"), + } + let new = wai::int_splat::ops::decode_typed(&gs, &cl.st()).map(|_| ()).map_err(|r| r.kind); + assert_eq!(new, Err(RefusalKind::ModelMismatch), "{b}: the decode through the operator set refuses it"); + }; + // The base decodes alike: one 1×1 transposed convolution, weight 1, shift 0. + let base = whm1_bytes(12, 16, 0, &[conv1(2, 2, 1, 0)]); + assert_eq!(int_splat_planned_break(&base), None); + let gs = IntModel::from_bytes(&base).unwrap(); + let cl = cloud_with(vec![0, 1 << 15, 1 << 16]); + assert_eq!(splat_direct(&gs, &cl), Some(Ok(vec![255; 8])), "2·2^12 clamped to 2^12: 255"); + assert_eq!(wai::int_splat::ops::decode_typed(&gs, &cl.st()).unwrap().attr, vec![255; 8]); + // No layer of tag 0, 1 or 2: the direct decode presents the latents themselves. + expect("no layer of tag 0, 1 or 2", &whm1_bytes(12, 16, 0, &[Whm::Relu]), &cl, None); + expect("no layer of tag 0, 1 or 2", &whm1_bytes(12, 16, 0, &[]), &cl, None); + // A header shift no layer uses (only an inverse GDN), or a gf no layer uses. + let igdn = Whm::Igdn { c: 2, g: vec![0; 4], b: vec![1 << 24; 2] }; + expect("a shift or gf of 128 or more that no layer uses", &whm1_bytes(12, 16, 200, &[igdn]), &cl, None); + expect("a shift or gf of 128 or more that no layer uses", &whm1_bytes(12, 200, 0, &[conv1(2, 2, 1, 0)]), &cl, None); + // A bias at the end of i64: the direct decode refuses every input at that layer + // (field_range), after the stream; the lowering refuses the set before it, so a + // stream that runs out is refused for the set, not the stream. + expect("a bias of magnitude 2^63 - 1 or more", &whm1_bytes(12, 16, 0, &[conv1(2, 2, 1, i64::MAX)]), &cl, Some(RefusalKind::FieldRange)); + let mut short = cloud_with(vec![0, 1, 1 << 16]); + short.rans.truncate(8); + expect("a bias of magnitude 2^63 - 1 or more", &whm1_bytes(12, 16, 0, &[conv1(2, 2, 1, i64::MIN)]), &short, Some(RefusalKind::StreamExhausted)); + // A check range its producer cannot meet: layer 0 gives every value 2^30 (weights 0, + // bias 2^30); layer 1, two input channels of weight 2^11 under the multiplier 2^40 and + // shift 0, has X_1 = 2047. The direct decode refuses every input at layer 1 + // (field_range); the load refuses the set (an empty check intersection, §8.1). + let tight = Whm::ConvT { inc: 2, out: 2, k: 1, s: 1, p: 0, op: 0, w: vec![1 << 11; 4], m: vec![1 << 40; 2], b: vec![0; 2] }; + let set = whm1_bytes(12, 16, 0, &[conv1(2, 2, 0, 1 << 30), tight]); + expect("a check range its producer cannot meet", &set, &cl, Some(RefusalKind::FieldRange)); +} + +// ------------------------------------------------------------------ int_audio (M3) + +/// One layer of a `WAM1` network, as a set's bytes hold it (integer-payloads §9). +#[derive(Clone)] +enum Wam { + ConvT { inc: u32, out: u32, k: u32, s: u32, p: u32, op: u32, w: Vec, m: Vec, b: Vec }, + Igdn { c: u32, g: Vec, b: Vec }, +} + +/// A `WAM1` network's bytes. +fn wam1_bytes(afrac: u32, gf: u32, shift: u32, layers: &[Wam]) -> Vec { + let mut v = b"WAM1".to_vec(); + for x in [afrac, gf, shift, layers.len() as u32] { + v.extend(x.to_le_bytes()); + } + for l in layers { + match l { + Wam::ConvT { inc, out, k, s, p, op, w, m, b } => { + v.push(0); + [*inc, *out, *k, *s, *p, *op].iter().for_each(|x| v.extend(x.to_le_bytes())); + w.iter().for_each(|x| v.extend(x.to_le_bytes())); + m.iter().chain(b).for_each(|x| v.extend(x.to_le_bytes())); + } + Wam::Igdn { c, g, b } => { + v.push(2); + v.extend(c.to_le_bytes()); + g.iter().for_each(|x| v.extend(x.to_le_bytes())); + b.iter().for_each(|x| v.extend(x.to_le_bytes())); + } + } + } + v +} + +/// A `WIA1` clip's typed fields, owned. +#[derive(Clone, Debug)] +struct Clip { + afrac: u32, + n_ch: usize, + lat_len: usize, + t_out: usize, + offset: i32, + median_q: Vec, + cdfs: Vec>, + rans: Vec, +} + +impl Clip { + fn st(&self) -> wai::int_audio::AudioStream<'_> { + wai::int_audio::AudioStream { + afrac: self.afrac, n_ch: self.n_ch, lat_len: self.lat_len, t_out: self.t_out, offset: self.offset, + median_q: &self.median_q, cdfs: &self.cdfs, rans: &self.rans, + } + } +} + +/// The planned break of integer-operators §15.3 that an audio set falls under, if any, +/// read from the set's bytes alone: the set is refused at load by its lowering (`wam1/1` +/// with the PCM step, loaded under `acc64`), for a rule the direct decode does not make +/// (as for splat, [`int_splat_planned_break`]; a `WAM1` network has layers of tags 0 and +/// 2 only, so one with no such layer has no layers). No payload field changes meaning. +fn int_audio_planned_break(set: &[u8]) -> Option<&'static str> { + let net = wai_ops::lower::parse_wam1(set).ok()?; + load_break(&net, wai_ops::lower::wam1(set)) +} + +/// The length a stack of layers gives on `n` latent samples, or `None` where one leaves +/// no output. +fn audio_len(layers: &[Wam], mut n: i64) -> Option { + for l in layers { + if let Wam::ConvT { k, s, p, op, .. } = l { + n = (n - 1) * *s as i64 + *k as i64 + *op as i64 - 2 * *p as i64; + if n < 1 { + return None; + } + } + } + Some(n) +} + +/// A random `wai.audio.int_codec` set and clip: a network of 1 to 4 layers (one set in +/// 32 with none), one-dimensional transposed convolutions and inverse GDNs, the last a +/// transposed convolution to one channel (one set in sixteen not), at `afrac` 12 mostly +/// (one in sixteen anywhere in 1..=30, one in sixteen at 0 or 31), a shift of 8 to 23, +/// and a header `gf` or `shift` of 128 or more one set in 32 each; kernels 1 to 5, +/// strides 1 to 3, paddings mostly below the kernel; weights up to 2^4, 2^10 or 2^12 (one +/// set in eight up to i16's range), multipliers mostly up to 2^12, biases mostly up to +/// 2^20 with one in 64 at an end of i64; IGDN `γ` up to 2^12, `β` up to 2^28. A latent +/// of 1 to 3 channels of 1 to 6 samples, one table length and offset for every channel; +/// symbols mostly inside the table, some escaped and far past it; medians mostly small, +/// one in eight at or past 2^47, and one case in four with a whole channel at the input's +/// `check` range or one past it. `t_out` is one the latent is the shortest for (one case +/// in eight not); the clip's `afrac` is the set's (one in 32 not); one stream in sixteen +/// is cut short, one in 32 has a word appended, one in 64 a byte changed. +fn int_audio_case(r: &mut Rng) -> (Vec, Clip) { + let afrac = match r.below(16) { + 0 => 1 + r.below(30) as u32, + 1 => [0, 31][r.below(2) as usize], + _ => 12, + }; + let gf = if r.below(32) == 0 { 128 + r.below(64) as u32 } else { 16 }; + let shift = if r.below(32) == 0 { 128 + r.below(64) as u32 } else { 8 + r.below(16) as u32 }; + let c0 = 1 + r.below(3) as u32; + let n_layers = if r.below(32) == 0 { 0 } else { 1 + r.below(4) as usize }; + let big = r.below(8) == 0; + let (mut c, mut layers) = (c0, Vec::new()); + for i in 0..n_layers { + let last = i + 1 == n_layers; + if !last && r.below(3) == 0 { + let g: Vec = (0..c * c).map(|_| if r.below(16) == 0 { r.signed(1 << 12) as i32 } else { r.below(1 << 12) as i32 }).collect(); + let b: Vec = (0..c).map(|_| if r.below(16) == 0 { r.signed(1 << 28) } else { r.below(1 << 28) as i64 }).collect(); + layers.push(Wam::Igdn { c, g, b }); + continue; + } + let out = if last && r.below(16) != 0 { 1 } else { 1 + r.below(3) as u32 }; + let k = 1 + r.below(5) as u32; + let s = 1 + r.below(3) as u32; + let p = if r.below(16) == 0 { r.below(k as u64 + 3) } else { r.below(k as u64) } as u32; + let op = if r.below(16) == 0 { r.below(s as u64 + 1) } else { r.below(s as u64) } as u32; + let wmax = if big { 32767 } else { [16, 1024, 4096][r.below(3) as usize] }; + let w: Vec = (0..c * out * k).map(|_| if r.below(4) == 0 { 0 } else { r.signed(wmax) as i16 }).collect(); + let m: Vec = (0..out).map(|_| if r.below(32) == 0 { r.signed(1 << 40) } else { r.signed(1 << 12) }).collect(); + let b: Vec = (0..out) + .map(|_| match r.below(64) { + 0 => [i64::MAX, i64::MIN, i64::MIN + 1, i64::MAX - 1][r.below(4) as usize], + 1 => r.signed(1 << 40), + _ => r.signed(1 << 20), + }) + .collect(); + layers.push(Wam::ConvT { inc: c, out, k, s, p, op, w, m, b }); + c = out; + } + let set = wam1_bytes(afrac, gf, shift, &layers); + + let lat_len = 1 + r.below(6) as usize; + let full = audio_len(&layers, lat_len as i64).unwrap_or(1).max(1) as usize; + let shorter = if lat_len > 1 { audio_len(&layers, lat_len as i64 - 1).unwrap_or(0).max(0) as usize } else { 0 }; + let mut t_out = if full > shorter { shorter + 1 + r.below((full - shorter) as u64) as usize } else { full }; + if r.below(8) == 0 { + t_out = if r.below(2) == 0 { full + 1 } else { shorter.max(1) }; + } + let n_ch = c0 as usize; + let a = afrac.min(30); + let len = 3 + r.below(4) as usize; + let mut cuts: Vec = (0..len - 2).map(|_| 1 + r.below(65_535) as i32).collect(); + cuts.sort_unstable(); + cuts.dedup(); + let mut t = vec![0]; + t.extend(cuts); + t.push(1 << 16); + let offset = r.signed(3) as i32; + let span = t.len() as i64 - 2; + let mut median_q: Vec = (0..n_ch) + .map(|_| match r.below(8) { + 0 => [1i64 << 47, -(1 << 47), (1 << 47) + 1, -(1 << 47) - 1][r.below(4) as usize], + _ => r.signed(1 << (a + 4)), + }) + .collect(); + let mut values: Vec = (0..n_ch * lat_len) + .map(|_| match r.below(32) { + 0 => r.signed(1 << 20), + 1..=4 => offset as i64 + span + r.below(40) as i64 - 20, + _ => offset as i64 + r.below(span as u64) as i64, + } as i32) + .collect(); + if r.below(4) == 0 { + if let Ok(g) = wai_ops::lower::wam1(&set) { + let x0 = g.tensors[0].hi as i128; + let ch = r.below(n_ch as u64) as usize; + let target = [x0, x0 + 1, -x0, -x0 - 1][r.below(4) as usize]; + let med = target - ((offset as i128) << a); + if med.unsigned_abs() <= 1 << 47 { + median_q[ch] = med as i64; + values[ch * lat_len..(ch + 1) * lat_len].iter_mut().for_each(|x| *x = offset); + } + } + } + let mut enc = wai::rans::RansEncoder::new(); + for &v in &values { + assert!(enc.encode(v, &t, t.len(), offset), "a value the generator writes is codable"); + } + let mut rans = enc.flush(); + match r.below(64) { + 0..=3 => { + let cut = (4 * (1 + r.below(3)) as usize).min(rans.len()); + rans.truncate(rans.len() - cut); + } + 4 | 5 => rans.extend([0, 0, 0, 0]), + 6 => { + let i = r.below(rans.len() as u64) as usize; + rans[i] ^= 1 << r.below(8); + } + _ => {} + } + let clip_afrac = if r.below(32) == 0 { afrac ^ 1 } else { afrac }; + (set, Clip { afrac: clip_afrac, n_ch, lat_len, t_out, offset, median_q, cdfs: vec![t; n_ch], rans }) +} + +/// The direct audio decode, a panic (an overflow-checked build's arithmetic overflow) +/// caught. +fn audio_direct(m: &wai::int_audio::AudioModel, cl: &Clip) -> Option, RefusalKind>> { + catch_unwind(AssertUnwindSafe(|| wai::int_audio::legacy::decode_typed(m, &cl.st()).map_err(|r| r.kind))).ok() +} + +/// Gate 4 of the migration, for `wai.audio.int_codec`: the decode through the operator +/// set against the direct decode on random sets and clips. Outside the planned breaks +/// both decode to the same samples or both refuse with the same code, and the direct +/// decode never panics; under a planned break the decode through the operator set +/// refuses. +#[test] +fn int_audio_random_sets_and_clips_decode_alike() { + let (mut alike, mut refused, mut breaks) = (0u64, std::collections::BTreeMap::::new(), std::collections::BTreeMap::<&str, u64>::new()); + let n = cases("int_audio", 3000, |r, seed| { + let (set, cl) = int_audio_case(r); + let why = || format!("seed {seed}\nclip {cl:?}"); + let m = wai::int_audio::AudioModel::from_bytes(&set).expect("the generator writes WAM1 that parses"); + let new = wai::int_audio::ops::decode_typed(&m, &cl.st()).map_err(|r| r.kind); + match int_audio_planned_break(&set) { + Some(b) => { + assert!(new.is_err(), "{b}: the decode through the operator set must refuse it\n{}", why()); + *breaks.entry(b).or_default() += 1; + } + None => match (&new, audio_direct(&m, &cl)) { + (Ok(a), Some(Ok(b))) => { + assert!(*a == b, "the two decodes differ\n{}", why()); + alike += 1; + } + (Err(a), Some(Err(b))) => { + assert_eq!(*a, b, "the two decodes refuse with different codes\n{}", why()); + *refused.entry(format!("{a:?}")).or_default() += 1; + } + (_, old) => panic!("new {new:?}, direct {old:?}\n{}", why()), + }, + } + }); + eprintln!("int_audio: {alike} decoded alike, refused by both {refused:?}, planned breaks {breaks:?}"); + assert!(alike * 5 > n, "the generator must produce decodable cases: {alike} of {n}"); + assert!(n < 1000 || refused.contains_key("FieldRange"), "no case reached a check range"); + assert!(n < 1000 || breaks.len() >= 3, "too few kinds of planned break reached: {breaks:?}"); +} + +/// The planned breaks of integer-operators §15.3 item 3 for `wai.audio.int_codec`, each +/// on a set and clip the direct decode decodes, or refuses with another code: the +/// classifier names it, and the decode through the operator set refuses it +/// (`model_mismatch`, before the stream). +#[test] +fn int_audio_planned_breaks_are_new_refusals() { + use wai::int_audio::AudioModel; + // A clip of one channel of 4 latents, each 2·2^12, presenting 4 samples; with + // `freq1`, each symbol costs 16 bits, so a stream cut to its seed runs out. + let clip_with = |t: Vec| { + let mut enc = wai::rans::RansEncoder::new(); + for _ in 0..4 { + assert!(enc.encode(2, &t, 3, 2)); + } + Clip { afrac: 12, n_ch: 1, lat_len: 4, t_out: 4, offset: 2, median_q: vec![0], cdfs: vec![t], rans: enc.flush() } + }; + let conv1 = |w: i16, m: i64, b: i64| Wam::ConvT { inc: 1, out: 1, k: 1, s: 1, p: 0, op: 0, w: vec![w], m: vec![m], b: vec![b] }; + let expect = |b: &str, set: &[u8], cl: &Clip, direct: Option| { + assert_eq!(int_audio_planned_break(set), Some(b), "{b}: the classifier"); + let m = AudioModel::from_bytes(set).unwrap(); + match direct { + None => assert!(matches!(audio_direct(&m, cl), Some(Ok(_))), "{b}: the direct decode decodes it"), + Some(k) => assert_eq!(audio_direct(&m, cl), Some(Err(k)), "{b}: the direct decode's code"), + } + let new = wai::int_audio::ops::decode_typed(&m, &cl.st()).map(|_| ()).map_err(|r| r.kind); + assert_eq!(new, Err(RefusalKind::ModelMismatch), "{b}: the decode through the operator set refuses it"); + }; + // The base decodes alike: one 1×1 transposed convolution, weight 1, shift 0, so each + // sample is (2·2^12·32767 + 2^11) >> 12 = 65534, clamped to 32767. + let cl = clip_with(vec![0, 1 << 15, 1 << 16]); + let base = wam1_bytes(12, 16, 0, &[conv1(1, 1, 0)]); + assert_eq!(int_audio_planned_break(&base), None); + let m = AudioModel::from_bytes(&base).unwrap(); + assert_eq!(audio_direct(&m, &cl), Some(Ok(vec![32767; 4]))); + assert_eq!(wai::int_audio::ops::decode_typed(&m, &cl.st()).unwrap(), vec![32767; 4]); + // No layers: the direct decode presents the latents themselves. + expect("no layer of tag 0, 1 or 2", &wam1_bytes(12, 16, 0, &[]), &cl, None); + // A header shift no layer uses (only an inverse GDN), or a gf no layer uses. + let igdn = Wam::Igdn { c: 1, g: vec![0], b: vec![1 << 24] }; + expect("a shift or gf of 128 or more that no layer uses", &wam1_bytes(12, 16, 200, &[igdn]), &cl, None); + expect("a shift or gf of 128 or more that no layer uses", &wam1_bytes(12, 200, 0, &[conv1(1, 1, 0)]), &cl, None); + // A bias at the end of i64: refused at that layer, after the stream, directly; at + // load through the lowering, so a stream that runs out is refused for the set. + expect("a bias of magnitude 2^63 - 1 or more", &wam1_bytes(12, 16, 0, &[conv1(1, 1, i64::MAX)]), &cl, Some(RefusalKind::FieldRange)); + let mut short = clip_with(vec![0, 1, 1 << 16]); + short.rans.truncate(8); + expect("a bias of magnitude 2^63 - 1 or more", &wam1_bytes(12, 16, 0, &[conv1(1, 1, i64::MIN)]), &short, Some(RefusalKind::StreamExhausted)); + // A check range its producer cannot meet: every value 2^30 into a layer whose range is + // ±4095 (weight 2^11, multiplier 2^40, shift 0); and every sample 2^62, past the PCM + // step's range. The direct decode refuses each at that step (field_range). + expect("a check range its producer cannot meet", &wam1_bytes(12, 16, 0, &[conv1(0, 1, 1 << 30), conv1(1 << 11, 1 << 40, 0)]), &cl, Some(RefusalKind::FieldRange)); + expect("a check range its producer cannot meet", &wam1_bytes(12, 16, 0, &[conv1(0, 1, 1 << 62)]), &cl, Some(RefusalKind::FieldRange)); +} diff --git a/wai/wai-rs/tests/int_payload_cost.rs b/wai/wai-rs/tests/int_payload_cost.rs index 27494e371..f71e75c7d 100644 --- a/wai/wai-rs/tests/int_payload_cost.rs +++ b/wai/wai-rs/tests/int_payload_cost.rs @@ -5,7 +5,9 @@ //! A sink reads the declared cost before it decodes and applies its own size policy, //! so the cost must be one it can rely on: the decode, on the reference kernels, holds //! at least the output and at most the declared working bytes beside the payload and -//! the parameter set as parsed. For a `WIV1` with keyframe 0, which runs no network, +//! the parameter set as parsed: for a capability that decodes through the integer +//! operator set, parsed, lowered and loaded, which is the same for every payload the +//! set decodes. For a `WIV1` with keyframe 0, which runs no network, //! the declared working bytes are what the decode holds, to within a few bytes. //! //! The clip a few kilobytes can code: sixteen frames of 4729 × 4729 whose stream is @@ -126,8 +128,19 @@ fn the_declared_cost_bounds_what_the_decode_holds() { Box::new((wai::int_hyper_synth::IntModel::from_bytes(&hyp["g_s.whm"]).unwrap(), wai::int_hyper_synth::IntModel::from_bytes(&hyp["h_s.whm"]).unwrap(), t)) as Box } - c::CAP_AUDIO_INT_CODEC => Box::new(wai::int_audio::AudioModel::from_bytes(&model).unwrap()) as Box, - c::CAP_SPLAT_INT_CODEC => Box::new(wai::int_hyper_synth::IntModel::from_bytes(&gs).unwrap()) as Box, + c::CAP_AUDIO_INT_CODEC => { + let m = wai::int_audio::AudioModel::from_bytes(&model).unwrap(); + let g = wai::int_audio::ops::load(&m).unwrap(); + Box::new((m, g)) as Box + } + // A decode through the integer operator set holds its set parsed and its + // lowered graph loaded (`int_splat::ops::load`), and the lowering forms each + // in turn. + c::CAP_SPLAT_INT_CODEC => { + let m = wai::int_hyper_synth::IntModel::from_bytes(&gs).unwrap(); + let g = wai::int_splat::ops::load(&m).unwrap(); + Box::new((m, g)) as Box + } _ => Box::new(()) as Box, }); drop(r); @@ -174,4 +187,38 @@ fn the_declared_cost_bounds_what_the_decode_holds() { println!("{what}: declared {} working bytes ({} output, {} symbols, {} multiply-accumulates); held {peak}", cost.working_bytes, cost.output_bytes, cost.symbols, cost.macs); } + + // The decode through the integer operator set, with its set already parsed and + // loaded: what it then holds for the payload (the latent, each operator's input, + // output and uncropped plane, the attribute bytes) is within the declared working + // bytes alone. The set's own lowering, above, holds more for a moment; that is the + // set's, the same for every payload. + let p = read("int_splat/splat.wis1"); + let cost = declared_cost(c::CAP_SPLAT_INT_CODEC, &p, &spl).unwrap(); + let v = wire::parse_wis1(&p).unwrap(); + let st = wai::int_splat::SplatStream { + afrac: v.afrac as u32, c: v.c as usize, sh: v.sh as usize, sw: v.sw as usize, s: v.s as usize, + n_attr: v.attr as usize, cdfs: &v.cdfs, offsets: &v.offsets, median_q: &v.median_q, rans: v.rans, + }; + let m = wai::int_hyper_synth::IntModel::from_bytes(&gs).unwrap(); + let g = wai::int_splat::ops::load(&m).unwrap(); + let (r, peak) = with_policy(Policy::reference(), || peak_during(|| wai::int_splat::ops::decode_loaded(&g, m.afrac, &st))); + r.unwrap(); + assert!(peak >= cost.output_bytes && peak <= cost.working_bytes, "splat, set loaded: held {peak}, declared {}", cost.working_bytes); + println!("splat, its set loaded beforehand: declared {} working bytes; held {peak}", cost.working_bytes); + + // The same for the audio synthesis. + let p = read("int_audio/audio_bitstream.wia1"); + let cost = declared_cost(c::CAP_AUDIO_INT_CODEC, &p, &aud).unwrap(); + let a = wire::parse_wia1(&p).unwrap(); + let st = wai::int_audio::AudioStream { + afrac: a.afrac as u32, n_ch: a.n_ch as usize, lat_len: a.lat_len as usize, t_out: a.t_out as usize, + offset: a.offset, median_q: &a.median_q, cdfs: &a.cdfs, rans: a.rans, + }; + let m = wai::int_audio::AudioModel::from_bytes(&model).unwrap(); + let g = wai::int_audio::ops::load(&m).unwrap(); + let (r, peak) = with_policy(Policy::reference(), || peak_during(|| wai::int_audio::ops::decode_loaded(&g, m.afrac, &st))); + r.unwrap(); + assert!(peak >= cost.output_bytes && peak <= cost.working_bytes, "audio, set loaded: held {peak}, declared {}", cost.working_bytes); + println!("audio, its set loaded beforehand: declared {} working bytes; held {peak}", cost.working_bytes); } diff --git a/wai/wai-rs/tests/int_splat_ops.rs b/wai/wai-rs/tests/int_splat_ops.rs new file mode 100644 index 000000000..0501fd1c3 --- /dev/null +++ b/wai/wai-rs/tests/int_splat_ops.rs @@ -0,0 +1,440 @@ +//! `wai.splat.int_codec` through the integer operator set: its lowering `whm1/1` +//! (`extensions/integer-operators.md` §13.2, with the RGB8 step of §13.0) and the +//! certificate of every committed set (§8.5), 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, from +//! the model it parses (`int_ops::whm1_net`), must write those same bytes, and the load a +//! sink makes must give the verdict the certificate records. +//! +//! The decode through the lowering (`int_splat::ops`) must give, on every committed +//! vector, the bytes of the direct decode (`int_splat::legacy`) and the vector's +//! committed digests. + +use std::path::PathBuf; + +use serde_json::Value; +use wai::int_hyper_synth::IntModel; +use wai::int_splat::{self, SplatJson, SplatStream}; +use wai_ops::Profile; + +const CAP: &str = "wai.splat.int_codec"; + +fn wai() -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("..") +} + +fn read(rel: &str) -> Vec { + std::fs::read(wai().join(rel)).unwrap_or_else(|e| panic!("{rel}: {e}")) +} + +fn json(rel: &str) -> Value { + serde_json::from_slice(&read(rel)).unwrap_or_else(|e| panic!("{rel}: {e}")) +} + +fn sha256(b: &[u8]) -> String { + wai_ops::sha256::sha256(b).iter().map(|x| format!("{x:02x}")).collect() +} + +/// The `wai.splat.int_codec` sets `ops-conformance/certificates/index.json` lists: +/// `(name, set file, set digest, lowered graph file)`, each relative to `wai/`. +fn indexed() -> Vec<(String, String, String, String)> { + json("ops-conformance/certificates/index.json")["sets"] + .as_array() + .unwrap() + .iter() + .filter(|s| s["capability"] == CAP) + .map(|s| { + assert_eq!(s["lowering"], "whm1/1"); + let roles = s["roles"].as_array().unwrap(); + assert_eq!(roles.len(), 1, "whm1/1 lowers a splat set to one graph, g_s"); + ( + s["name"].as_str().unwrap().to_string(), + s["members"]["g_s.whm"].as_str().unwrap().to_string(), + s["set_sha256"].as_str().unwrap().to_string(), + roles[0].as_str().unwrap().to_string(), + ) + }) + .collect() +} + +/// The graph this crate lowers from the model it parses: `int_ops::whm1_net`, then +/// `whm1/1` with the RGB8 step. +fn lowered_from_the_parsed_model(set: &[u8]) -> wai_ops::Graph { + let m = IntModel::from_bytes(set).expect("the WHM1 set parses"); + let net = wai::int_ops::whm1_net(&m).expect("a model parsed from bytes converts"); + assert!(net == wai_ops::lower::parse_whm1(set).unwrap(), "the converted model is the one wai-ops parses from the bytes"); + wai_ops::lower::whm1_net(&net, true).expect("the set lowers") +} + +/// Gate 1 of the migration: every committed `wai.splat.int_codec` set lowers to its +/// committed graph, from its bytes (`lower::splat`) and from the model this crate parses; +/// the certificate beside it is that graph's; and loading the graph under the +/// capability's profile gives the certificate's verdict. The set carries `check` ranges +/// (§13.2), so it meets `acc64` by construction and not `acc32`; its first `acc32` failure +/// is `g_s`'s input range, at 44 bits. +#[test] +fn every_committed_set_lowers_to_its_committed_graph_and_verdict() { + let sets = indexed(); + let names: Vec<&str> = sets.iter().map(|s| s.0.as_str()).collect(); + assert_eq!(names, ["splat"], "the committed splat sets"); + for (name, set, digest, role) in &sets { + let bytes = read(set); + let lowered = wai_ops::lower::splat(&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, ["g_s"], "{name}"); + let g = lowered.get("g_s").unwrap(); + // The verdict first: the load a sink makes. + let loaded = wai::int_ops::load(g).unwrap_or_else(|r| panic!("{name}: refused under acc64: {r:?}")); + assert!(!loaded.meets(Profile::Acc32), "{name}: a set with check ranges inside its graph does not meet acc32"); + // Then the committed bytes: the set, its lowered graph (both ways) 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 of the bytes differs from the committed {role}"); + assert!(lowered_from_the_parsed_model(&bytes).to_bytes() == graph, "{name}: the lowering of the parsed model differs from the committed {role}"); + let cert = json(&format!("ops-conformance/certificates/{CAP}/{digest}.json")); + assert_eq!(cert["set_sha256"].as_str(), Some(digest.as_str()), "{name}"); + assert_eq!(cert["lowering"], "whm1/1", "{name}"); + assert_eq!(cert["graphs"][0]["name"], "g_s", "{name}"); + assert_eq!(cert["graphs"][0]["graph_sha256"].as_str(), Some(sha256(&graph).as_str()), "{name}: the certificate is the graph's"); + let v = &cert["verdict"]; + assert_eq!((v["acc64"].as_bool(), v["acc32"].as_bool()), (Some(true), Some(false)), "{name}: the certificate's verdict"); + let first = &v["acc32_first_failure"]; + assert_eq!( + (first["graph"].as_str(), first["where"].as_str(), first["index"].as_str(), first["key"].as_str(), first["bits"].as_str()), + (Some("g_s"), Some("input"), Some("0"), Some("range"), Some("44")), + "{name}: the first acc32 failure Appendix C predicts" + ); + // The input's check range is the lowering's X_0 for the first layer (§13.2). + let x0 = &cert["graphs"][0]["inputs"][0]; + assert_eq!((x0["policy"].as_str(), x0["lo"].as_str(), x0["hi"].as_str()), (Some("check"), Some("-5632249656115"), Some("5632249656115")), "{name}"); + assert_eq!(g.tensors[0].lo, -5632249656115, "{name}"); + assert_eq!(g.tensors[0].hi, 5632249656115, "{name}"); + } +} + +/// The typed form of a parsed `WIS1` payload. +fn stream<'a>(v: &'a wai::int_wire::Wis1<'a>) -> SplatStream<'a> { + SplatStream { + afrac: v.afrac as u32, c: v.c as usize, sh: v.sh as usize, sw: v.sw as usize, s: v.s as usize, + n_attr: v.attr as usize, cdfs: &v.cdfs, offsets: &v.offsets, median_q: &v.median_q, rans: v.rans, + } +} + +/// The `wai.splat.int_codec` vectors `int-payload-conformance/vectors.json` indexes: +/// `(name, payload, set member g_s.whm, JSON form, sidecar)`, each relative to `wai/`. +fn vectors() -> Vec<(String, String, String, String, String)> { + json("int-payload-conformance/vectors.json")["vectors"] + .as_array() + .unwrap() + .iter() + .filter(|v| v["capability"] == CAP) + .map(|v| { + let s = |k: &Value| k.as_str().unwrap().to_string(); + (s(&v["name"]), s(&v["payload"]), s(&v["set"]["g_s.whm"]), s(&v["converted_from"]["json"]), s(&v["sidecar"])) + }) + .collect() +} + +/// Gate 2 of the migration: golden equality. Every committed `wai.splat.int_codec` +/// vector decodes through the lowering (`int_splat::ops`) to the bytes the direct decode +/// gives (`int_splat::legacy`), and to every digest the corpus commits for it: the +/// SHA-256 sidecar and the export's FNV-1a. The payload (`WIS1`) and the JSON form it is +/// converted from both decode, and alike. The staged corpus's splat base carries the +/// same payload, so it is covered too. +#[test] +fn every_committed_vector_decodes_through_the_lowering_to_its_golden() { + let vs = vectors(); + assert_eq!(vs.len(), 1, "the committed splat vectors"); + for (name, payload, set, from, sidecar) in &vs { + let gs = IntModel::from_bytes(&read(set)).unwrap(); + let p = read(payload); + let v = wai::int_wire::parse_wis1(&p).unwrap_or_else(|e| panic!("{name}: {e:?}")); + let bs = json(from); + let j = SplatJson::from_value(&bs).unwrap(); + let want = json(sidecar)[name]["sha256"].as_str().unwrap().to_string(); + for (form, st) in [("WIS1", stream(&v)), ("JSON", j.stream())] { + let ops = int_splat::ops::decode_typed(&gs, &st).unwrap_or_else(|r| panic!("{name} {form}: {r}")); + let direct = int_splat::legacy::decode_typed(&gs, &st).unwrap_or_else(|r| panic!("{name} {form}: {r}")); + assert!(ops.attr == direct.attr, "{name} {form}: the decode through the lowering differs from the direct decode"); + assert_eq!((ops.n_attr, ops.s), (direct.n_attr, direct.s), "{name} {form}"); + assert_eq!(sha256(&ops.attr), want, "{name} {form}: the SHA-256 sidecar"); + assert_eq!(Some(format!("{:016x}", ops.attr_hash()).as_str()), bs["attr_hash"].as_str(), "{name} {form}: the export's FNV-1a"); + } + let base = read("staged-conformance/cases/splat_attr_1l/base.wai"); + assert!(base.windows(p.len()).any(|w| w == p.as_slice()), "{name}: the staged splat base carries this payload"); + } +} + +/// The committed vector's set and payload (`int-payload-conformance`'s `splat`). +const SET: &str = "wai-rs/tests/vectors/int_splat/splat_gs.bin"; +const PAYLOAD: &str = "wai-rs/tests/vectors/int_splat/splat.wis1"; + +/// A typed cloud that owns its tables, so a fixture can change any field. +#[derive(Clone)] +struct Owned { + afrac: u32, + c: usize, + sh: usize, + sw: usize, + s: usize, + n_attr: usize, + cdfs: Vec>, + offsets: Vec, + median_q: Vec, + rans: Vec, +} + +impl Owned { + fn committed() -> Owned { + let p = read(PAYLOAD); + let v = wai::int_wire::parse_wis1(&p).unwrap(); + Owned { + afrac: v.afrac as u32, c: v.c as usize, sh: v.sh as usize, sw: v.sw as usize, s: v.s as usize, n_attr: v.attr as usize, + cdfs: v.cdfs.clone(), offsets: v.offsets.clone(), median_q: v.median_q.clone(), rans: v.rans.to_vec(), + } + } + fn st(&self) -> SplatStream<'_> { + SplatStream { + afrac: self.afrac, c: self.c, sh: self.sh, sw: self.sw, s: self.s, n_attr: self.n_attr, + cdfs: &self.cdfs, offsets: &self.offsets, median_q: &self.median_q, rans: &self.rans, + } + } +} + +/// The committed set, parsed, changed by `f`. +fn model(f: impl FnOnce(&mut IntModel)) -> IntModel { + let mut m = IntModel::from_bytes(&read(SET)).unwrap(); + f(&mut m); + m +} + +/// The committed cloud, changed by `f`. +fn cloud(f: impl FnOnce(&mut Owned)) -> Owned { + let mut o = Owned::committed(); + f(&mut o); + o +} + +/// The code integer-payloads §8 gives a refusal of the splat decode. +fn code_of(r: wai::int_hyper_synth::Refused) -> &'static str { + wai::int_payload::IntPayloadError::from(r).code() +} + +/// Gate 3 of the migration: refusal parity, on the committed fixtures. Each +/// `wai.splat.int_codec` fixture of `int-payload-conformance` is refused by a sink with +/// its code; one that parses (`decode`) is refused by the decode through the lowering and +/// by the direct decode alike, with that code. +#[test] +fn every_committed_refusal_fixture_keeps_its_code() { + let mut n = 0; + for f in json("int-payload-conformance/vectors.json")["refusals"].as_array().unwrap().iter().filter(|f| f["capability"] == CAP) { + let (name, want) = (f["name"].as_str().unwrap(), f["code"].as_str().unwrap()); + let (p, set) = (read(f["payload"].as_str().unwrap()), read(f["set"]["g_s.whm"].as_str().unwrap())); + let params = wai::int_payload::resolve(CAP, None, &[wai::pin::HeldBytes::File(&set)], None).expect("an unpinned set resolves"); + let sink = wai::int_payload::decode_int(CAP, &p, ¶ms).map(|_| ()).map_err(|e| e.code()); + assert_eq!(sink, Err(want), "{name}: the sink's code"); + if f["layer"] == "decode" { + let gs = IntModel::from_bytes(&set).unwrap(); + let v = wai::int_wire::parse_wis1(&p).unwrap_or_else(|e| panic!("{name}: a decode fixture parses: {e:?}")); + let ops = int_splat::ops::decode_typed(&gs, &stream(&v)).err().map(code_of); + let direct = int_splat::legacy::decode_typed(&gs, &stream(&v)).err().map(code_of); + assert_eq!((ops, direct), (Some(want), Some(want)), "{name}: through the lowering, then directly"); + } + n += 1; + } + assert!(n >= 11, "the committed splat fixtures: {n}"); +} + +/// Gate 3, beside the committed fixtures: a fixture for each refusal the direct decode +/// makes (integer-payloads §7, in its order), each made from the committed set and cloud +/// by one change. Each is refused by both decodes with the same code; the code it states +/// is integer-payloads §8's for that rule. The committed set and cloud decode. +#[test] +fn every_refusal_of_the_direct_decode_is_refused_alike_through_the_lowering() { + let (gs, base) = (model(|_| {}), Owned::committed()); + assert!(int_splat::ops::decode_typed(&gs, &base.st()).is_ok(), "the fixtures' base decodes"); + let same = || model(|_| {}); + let fixtures: Vec<(&str, &str, IntModel, Owned)> = vec![ + ("the set's afrac of 0", "model_mismatch", model(|m| m.afrac = 0), cloud(|_| {})), + ("the set's afrac of 31", "model_mismatch", model(|m| m.afrac = 31), cloud(|_| {})), + ("a cloud afrac not the set's", "model_mismatch", same(), cloud(|o| o.afrac = 11)), + ("no latent channels", "field_range", same(), cloud(|o| o.c = 0)), + ("a median short", "field_range", same(), cloud(|o| { o.median_q.pop(); })), + ("an offset short", "field_range", same(), cloud(|o| { o.offsets.pop(); })), + ("a table short", "field_range", same(), cloud(|o| { o.cdfs.pop(); })), + ("a median past 2^47", "field_range", same(), cloud(|o| o.median_q[3] = (1 << 47) + 1)), + ("a table that does not end at 2^16", "cdf_malformed", same(), cloud(|o| *o.cdfs[2].last_mut().unwrap() = 65535)), + ("a latent past 2^26", "too_large", same(), cloud(|o| (o.sh, o.sw) = (1 << 13, 1 << 13))), + ("an attribute tensor past 2^26", "too_large", same(), cloud(|o| o.s = 1 << 14)), + ("an uncropped plane past 2^26", "too_large", same(), cloud(|o| (o.sh, o.sw) = (1100, 1100))), + ("a layer's bias short", "model_mismatch", model(|m| { + if let wai::int_hyper_synth::IntLayer::ConvT { bias, .. } = &mut m.layers[0] { bias.pop(); } + }), cloud(|_| {})), + ("padding that leaves no output", "model_mismatch", model(|m| { + if let wai::int_hyper_synth::IntLayer::ConvT { pad, .. } = &mut m.layers[2] { *pad = 1000; } + }), cloud(|_| {})), + ("an attribute side not g_s's", "geometry_mismatch", same(), cloud(|o| o.s = 31)), + ("a stream too short to be one", "field_range", same(), cloud(|o| o.rans.truncate(4))), + ("a stream that runs out", "stream_exhausted", same(), cloud(|o| { let n = o.rans.len() / 2; o.rans.truncate(n - n % 4); })), + ("a word left unread", "stream_noncanonical", same(), cloud(|o| o.rans.extend([0, 0, 0, 0]))), + ("every latent past layer 0's range", "field_range", same(), cloud(|o| o.median_q.iter_mut().for_each(|m| *m = 1 << 47))), + ("every latent below it", "field_range", same(), cloud(|o| o.median_q.iter_mut().for_each(|m| *m = -(1 << 47)))), + ]; + for (what, want, gs, o) in &fixtures { + let ops = int_splat::ops::decode_typed(gs, &o.st()).err().map(code_of); + let direct = int_splat::legacy::decode_typed(gs, &o.st()).err().map(code_of); + assert_eq!((ops, direct), (Some(*want), Some(*want)), "{what}: through the lowering, then directly"); + } +} + +/// Gate 3, at the edge of every `check` range: for each layer of the committed set, a +/// network of that layer alone (with the set's header) is lowered, and its input's +/// `check` range is `[−X_ℓ, X_ℓ]`, the range the full graph gives tensor `ℓ` (§13.2). An +/// input whose largest magnitude is `X_ℓ` runs, through the lowering and directly, to the +/// same output; one at `X_ℓ + 1` is refused by both: by the graph (`field_range`, at +/// input 0) and by the direct decode's check of that layer (`run_net_checked`, a value +/// out of range at layer 0). So the ranges refuse exactly what integer-payloads §7 +/// *Exactness* refuses, at the same layer. +#[test] +fn every_check_range_refuses_exactly_what_the_direct_decode_refuses() { + use wai::int_hyper_synth::run_net_checked; + use wai_ops::{Code, Site, Tensor}; + let set = read(SET); + let full = wai_ops::lower::whm1(&set, true).unwrap(); + let net = wai_ops::lower::parse_whm1(&set).unwrap(); + for l in 0..net.layers.len() { + let one = model(|m| { + let layer = m.layers.swap_remove(l); + m.layers = vec![layer]; + }); + let g = wai_ops::lower::whm1_net(&wai::int_ops::whm1_net(&one).unwrap(), false).unwrap(); + let (c, lo, hi) = (g.tensors[0].channels, g.tensors[0].lo, g.tensors[0].hi); + assert_eq!((c, lo, hi), (full.tensors[l].channels, full.tensors[l].lo, full.tensors[l].hi), "layer {l}: the range of tensor {l}"); + assert!(hi > 0 && lo == -hi, "layer {l}: a symmetric check range [{lo}, {hi}]"); + let loaded = wai::int_ops::load(&g).unwrap(); + let x = |v: i64| { + let mut d = vec![0i64; c as usize * 4]; + d[c as usize * 2 + 1] = v; + d + }; + for v in [hi, -hi] { + let ops = wai::int_ops::run1(&loaded, Tensor::new(c, 2, 2, x(v)), &[]).unwrap_or_else(|r| panic!("layer {l}: {v} runs: {r:?}")); + let (direct, ..) = run_net_checked(&one, &x(v), c as usize, 2, 2).unwrap_or_else(|r| panic!("layer {l}: {v} runs directly: {r}")); + assert!(ops.data == direct, "layer {l}: at {v} the two differ"); + } + for v in [hi + 1, -hi - 1] { + let r = wai::int_ops::run1(&loaded, Tensor::new(c, 2, 2, x(v)), &[]).expect_err("one past the range"); + assert_eq!((r.code, r.at), (Code::FieldRange, Site::Input(0)), "layer {l}: {v}"); + let d = run_net_checked(&one, &x(v), c as usize, 2, 2).expect_err("one past the range, directly"); + assert!(d.out_of_range && d.layer == 0, "layer {l}: {v}: {d}"); + } + } +} + +/// 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 attribute bytes, and the same refusals, on every committed +/// vector, as `WIS1` and as its JSON form, and on clouds at the edge of the input's check +/// range. The per-operator half is `int_ops`' equivalence tests, one of which counts the +/// fast kernel's runs on this set. +#[test] +fn every_kernel_policy_decodes_through_the_lowering_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 gs = model(|_| {}); + let x0 = wai_ops::lower::whm1(&read(SET), true).unwrap().tensors[0].hi; + let at = |t: i64| cloud(move |o| o.median_q[1] = t - ((o.offsets[1] as i64) << o.afrac)); + let mut clouds: Vec<(String, Owned)> = vec![("WIS1".into(), Owned::committed())]; + let j = SplatJson::from_value(&json("wai-rs/tests/vectors/int_splat/splat.json")).unwrap(); + clouds.push(("JSON".into(), Owned { + afrac: j.afrac, c: j.c, sh: j.sh, sw: j.sw, s: j.s, n_attr: j.n_attr, + cdfs: j.cdfs.clone(), offsets: j.offsets.clone(), median_q: j.median_q.clone(), rans: j.rans.clone(), + })); + for (what, t) in [("channel 1's median at X_0", x0), ("at X_0 + 1", x0 + 1), ("at -X_0", -x0), ("at -X_0 - 1", -x0 - 1)] { + clouds.push((what.into(), at(t))); + } + for (k, (what, o)) in clouds.iter().enumerate() { + let want = with_policy(POLICIES[0], || int_splat::ops::decode_typed(&gs, &o.st())).map(|d| d.attr).map_err(code_of); + assert!(k >= 2 || want.is_ok(), "{what}: the committed form decodes"); + for pol in &POLICIES[1..] { + let got = with_policy(*pol, || int_splat::ops::decode_typed(&gs, &o.st())).map(|d| d.attr).map_err(code_of); + assert!(got == want, "{what}: {pol:?} differs from the interpreter alone"); + } + } +} + +/// The committed `wai.splat.int_codec` sets and payloads, relative to `wai/`: every set +/// and payload `int-payload-conformance/vectors.json` names for the capability (its +/// vectors and its refusal fixtures), and every base of a staged case decoded through +/// the capability, with the set its case names and the payload its envelope carries. +/// Returns `(set, (where, payload))` pairs. +fn committed_sets_and_payloads() -> Vec<(String, (String, Vec))> { + let mut v = Vec::new(); + let index = json("int-payload-conformance/vectors.json"); + for e in index["vectors"].as_array().unwrap().iter().chain(index["refusals"].as_array().unwrap()) { + if e["capability"] == CAP { + let (set, p) = (e["set"]["g_s.whm"].as_str().unwrap(), e["payload"].as_str().unwrap()); + v.push((set.to_string(), (p.to_string(), read(p)))); + } + } + let cases = wai().join("staged-conformance/cases"); + let mut dirs: Vec<_> = std::fs::read_dir(&cases).unwrap().filter_map(|e| e.ok()).map(|e| e.path()).collect(); + dirs.sort(); + for d in dirs { + let Ok(b) = std::fs::read(d.join("case.json")) else { continue }; + let case: Value = serde_json::from_slice(&b).unwrap(); + for base in case["bases"].as_array().into_iter().flatten().filter(|b| b["decoded_via"] == CAP) { + let f = d.join(base["file"].as_str().unwrap()); + let env = wai::container::Wai::from_bytes(&std::fs::read(&f).unwrap()).unwrap(); + assert_eq!(env.manifest.model_requirement.capability, CAP, "{}", f.display()); + v.push((base["decode"]["set"]["g_s.whm"].as_str().unwrap().to_string(), (f.display().to_string(), env.payload))); + } + } + v +} + +/// The scan for the load rules the switch brings (integer-operators §15.3 item 3): no +/// committed set or payload is affected. Every committed set has a layer of tag 0, 1 or +/// 2, header shifts under 128, no bias of magnitude `2^63 − 1` or more, and lowers and +/// loads under `acc64`; every committed payload that the direct decode decodes decodes +/// through the lowering to the same bytes, and every one it refuses is refused with the +/// same code. (These are the files that carry a `WIS1` payload or name the capability's +/// set; `WIS1` appears in no other committed payload or envelope.) +#[test] +fn the_new_load_rules_affect_no_committed_set_or_payload() { + use wai_ops::lower::NetLayer as N; + let all = committed_sets_and_payloads(); + let mut sets: Vec<&str> = all.iter().map(|(s, _)| s.as_str()).collect(); + sets.sort_unstable(); + sets.dedup(); + for set in &sets { + let bytes = read(set); + let net = wai_ops::lower::parse_whm1(&bytes).unwrap_or_else(|r| panic!("{set}: {r:?}")); + assert!(net.layers.iter().any(|l| matches!(l, N::ConvT { .. } | N::Conv { .. } | N::Igdn { .. })), "{set}: no layer of tag 0, 1 or 2"); + assert!(net.shift < 128 && net.gf < 128, "{set}: a header shift of 128 or more"); + for l in &net.layers { + if let N::ConvT { bias, .. } | N::Conv { bias, .. } = l { + assert!(bias.iter().all(|b| b.unsigned_abs() < i64::MAX as u64), "{set}: a bias of magnitude 2^63 - 1 or more"); + } + } + let g = wai_ops::lower::whm1(&bytes, true).unwrap_or_else(|r| panic!("{set}: does not lower: {r:?}")); + assert!(wai::int_ops::load(&g).is_ok(), "{set}: does not load under acc64"); + } + let (mut decoded, mut refused) = (0, 0); + for (set, (at, p)) in &all { + let gs = IntModel::from_bytes(&read(set)).unwrap(); + let Ok(v) = wai::int_wire::parse_wis1(p) else { + refused += 1; + continue; + }; + let ops = int_splat::ops::decode_typed(&gs, &stream(&v)).map(|d| d.attr).map_err(code_of); + let direct = int_splat::legacy::decode_typed(&gs, &stream(&v)).map(|d| d.attr).map_err(code_of); + assert!(ops == direct, "{at}: the decode through the lowering differs from the direct decode"); + if ops.is_ok() { decoded += 1 } else { refused += 1 } + } + assert_eq!(sets.len(), 1, "the committed splat sets"); + assert!(decoded >= 2 && refused >= 11, "the committed splat payloads, at least those of this revision: {decoded} decoded, {refused} refused"); +}