Transaction-Science/open-standards
Watch
1
0
Fork
You've already forked open-standards
 0
Code Issues Pull requests Projects Releases 15 Packages Wiki Activity Actions 1

wai: keep loaded parameter sets between decodes, keyed by the SHA-256 of their bytes (O23) #268

Merged
dcharlot merged 20 commits from wai/o23-loaded-set-cache into main 2026-10-15 20:41:28 -04:00
Conversation 0 Commits 20 Files changed 43 +2978 -253
43 changed files with 2978 additions and 253 deletions
Download patch file Download diff file Expand all files Collapse all files

10
.github/workflows/wai-byte-equality.yml vendored
View file

@ -74,8 +74,9 @@ on:
# The fast kernels int_synth, int_transform and int_hyper_synth dispatch to
# by default: the goldens here decode through them.
- 'wai/wai-rs/src/int_kernels.rs'
# The integer operator set and the bridge to it: the goldens decode int_synth
# through its lowering too.
# The integer operator set and the bridge to it: the goldens decode int_synth,
# the audio and splat codecs and the conditional codec (WIH1, the WIV1
# keyframe) through their lowerings too.
- 'wai/wai-rs/src/int_ops.rs'
- 'wai/wai-ops/**'
- 'wai/wai-rs/tests/vectors/int_video/**'
@ -111,8 +112,9 @@ on:
# The fast kernels int_synth, int_transform and int_hyper_synth dispatch to
# by default: the goldens here decode through them.
- 'wai/wai-rs/src/int_kernels.rs'
# The integer operator set and the bridge to it: the goldens decode int_synth
# through its lowering too.
# The integer operator set and the bridge to it: the goldens decode int_synth,
# the audio and splat codecs and the conditional codec (WIH1, the WIV1
# keyframe) through their lowerings too.
- 'wai/wai-rs/src/int_ops.rs'
- 'wai/wai-ops/**'
- 'wai/wai-rs/tests/vectors/int_video/**'

18
.github/workflows/wai-conformance.yml vendored
View file

@ -253,7 +253,9 @@ jobs:
# 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).
# wai.audio.int_codec and its lowering wam1/1 (§13.3), and tests/int_hyper_ops.rs
# for wai.neural.int_hyper and the keyframe of wai.video.int_hyper (whm1/1: h_s,
# then g_s with the RGB8 step).
# 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
@ -271,8 +273,14 @@ jobs:
# which nothing may panic; tests/int_payload_preflight.rs (its own binary, a counting
# allocator) checks that a header past the caps is refused before anything sized from it
# is allocated; tests/int_payload_cost.rs (its own binary, a live-bytes allocator) checks
# that int_payload::declared_cost bounds what each decode holds, and, in the release
# step only, decodes a 2^30-byte clip a 3 KB payload codes; tests/int_video_enc_memory.rs
# that int_payload::declared_cost bounds what each decode holds, with the store of loaded
# sets off, then on (a decode that loads its set, and one that finds it held), that the
# store counts at least the bytes it holds, and, in the release step only, decodes a
# 2^30-byte clip a 3 KB payload codes; tests/int_set_cache.rs checks that store under
# the decodes: a second decode through a set is a hit with the same bytes on every
# committed vector, sets one byte apart never share an entry, a refused set stays
# refused, refusals keep their codes, and the least recently used set goes first;
# tests/int_video_enc_memory.rs
# (its own binary, the same allocator) checks that int_video_enc::working_bytes bounds
# what the encoder holds. tests/staged_resources.rs (its
# own binary, a counting allocator) checks that the staged sink's work and memory per
@ -311,10 +319,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_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
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_hyper_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_set_cache --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_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
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_hyper_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_set_cache --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

3
.github/workflows/wai-ops-differential.yml vendored
View file

@ -7,7 +7,8 @@ name: wai-ops differential (nightly)
# 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,
# wai.splat.int_codec and wai.audio.int_codec.
# wai.splat.int_codec, wai.audio.int_codec and wai.neural.int_hyper (whose decode is
# also the keyframe of wai.video.int_hyper).
#
# Each night starts from its own seed, the date, which the log prints; a failing case
# names its seed, and `WAI_DIFF_SEED=<seed>` re-runs it. The same tests run for a fixed

3
.github/workflows/wai-sota-vectors.yml vendored
View file

@ -117,6 +117,7 @@ on:
branches: [main]
paths:
- 'wai/wai-rs/src/int_hyper_synth.rs'
- 'wai/wai-rs/src/int_ops.rs'
- 'wai/wai-rs/src/int_kernels.rs'
- 'wai/wai-rs/src/int_entropy.rs'
- 'wai/wai-rs/src/int_transform.rs'
@ -148,6 +149,7 @@ on:
pull_request:
paths:
- 'wai/wai-rs/src/int_hyper_synth.rs'
- 'wai/wai-rs/src/int_ops.rs'
- 'wai/wai-rs/src/int_kernels.rs'
- 'wai/wai-rs/src/int_entropy.rs'
- 'wai/wai-rs/src/int_transform.rs'
@ -329,6 +331,7 @@ jobs:
for t in rust_integer_synthesis_reproduces_the_export_hash \
rust_integer_hyper_decoder_reproduces_the_export_hash \
rust_decodes_the_full_conditional_bitstream_to_the_export_image \
the_pretrained_set_lowers_to_its_certificates_graphs \
integer_decode_keeps_the_float_codec_quality \
decodes_the_committed_tiny_conformance_vector \
sota_int_hyper_keyframe_clip_decodes_byte_for_byte \

92
wai/CHANGELOG.md
View file

@ -6,6 +6,98 @@ wire format may change and such changes are recorded here.
## v0.4.0 — current
### Parameter sets kept loaded between decodes
- **The store.** A decode through the integer operator set lowers its parameter set and
loads each graph, the bounds pass included; for the pretrained conditional-codec set
that costs more than the decode. `int_set_cache::LoadedSets` now keeps a set once it
has loaded, and the next decode through the same bytes takes it from there
([integer-operators](extensions/integer-operators.md) §8.6, informative note).
`int_payload::decode_int` and `int_payload::declared_cost` use it for `WIH1`, a
`WIV1`'s keyframe, `WIA1`, `WIS1` and `wai.neural.int_synth`, and so the neural sink
and `wai-web`'s `decodeWiv1`, `decodeWih1`, `decodeWia1`, `decodeWis1` and
`decodeIntSynth` do too; `decodeIntSynth` now decodes through `decode_int`, and a
refusal's message starts with its code. `int_synth::decode_from_spec`, which the
byte-equality crate shares, is now three steps (`check_payload`, `load_spec`,
`decode_loaded`) and loads its set every time. The entry points that read the tools'
JSON forms do not use the store.
- **The key** is the capability, the lowering, the profile the graphs were loaded under,
and the SHA-256 of the listing of the members the decode reads, each by its own
SHA-256 (`pin::listing_from`): the digest a pin over exactly those members names.
Never a name: two sets whose members share their names and differ in one byte have
two entries. Parameters `int_payload::resolve` built carry the listing digest it
already formed for check 5, so such a decode does not hash its set a second time;
the digest is used only for a capability that reads exactly the members it was formed
over, and any other key hashes the bytes.
- **What does not change.** An entry is made only from what a load returned, after it
returned, so a refused set is never kept: it is lowered, loaded and refused again on
its next decode, with the same refusal. Every check before the load runs as before,
in integer-payloads §7's order. Every output, refusal and declared cost is unchanged,
and no golden moved.
- **Memory.** A kept set is the set's, outside every decode's declared working bytes
([integer-payloads](extensions/integer-payloads.md) §7, informative note). The store
keeps its entries within a byte budget, 256 MiB to start with, evicts the least
recently used first, and keeps no set larger than the whole budget; a budget of 0
turns it off, and the set is then not hashed. `LoadedSets::set_budget` sets it, and
`wai-web` exports `setLoadedSetBudget` and `loadedSetStats` for a page. An entry is
counted as the shared allocation that holds the set, its graphs' heap bytes
(`wai_ops::Graph::heap_bytes`, new) and a fixed share of the store's indexes, never
less than it holds (`tests/int_payload_cost.rs`).
- **Tests.** `tests/int_set_cache.rs`: a second decode through each committed set is a
hit, by the store's count, with the bytes of a decode with the store off; sets one
byte apart never share an entry; a set the load refuses (an `int_synth` set the bounds
pass refuses; a conditional-codec set whose lowering admits no input range) stays
refused while a set one byte away is held; every refusal fixture keeps its code and
detail with its set held; the least recently used set goes first; threads decoding
one set keep one entry. The browser crate's check (`demo/verify_int_payload_wasm.mjs`) decodes `tiny`
twice and counts the hit.
### `wai.neural.int_hyper` and `wai.video.int_hyper` through the integer operator set
- **The decode.** `int_hyper_synth::decode_conditional_checked` (so the `WIH1` dispatch,
the keyframe of a `WIV1` with keyframe 1, the staged sink and the browser crate) plans
both networks as [integer-payloads](extensions/integer-payloads.md) §7 orders, checks
the payload's shapes against them, lowers `h_s` by `whm1/1` and `g_s` by `whm1/1` with
the RGB8 step ([integer-operators](extensions/integer-operators.md) §13.2, §13.0),
loads each under `acc64` with its bounds certificate, decodes the hyper-latent, runs
`h_s` through the operator set's interpreter, splits its output into `σ` and `μ`,
decodes the latent against the tables `σ` selects, and runs `g_s`. Every
convolution, transposed convolution and inverse GDN runs on the fast kernels behind
the interpreter. The graphs' `check` ranges are the checks the direct decode made
before each layer, so the same streams are refused (`field_range`). Every committed
output is unchanged, byte for byte (the conformance vector, the neural-keyframe clip,
and the pretrained set's image and clip regenerated in `wai-sota-vectors.yml`), 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`, in either network: new refusals of sets no committed corpus
contains, each `model_mismatch` after the payload's shapes are checked and before
either stream is decoded. A scan of every committed set of the two capabilities, and
of every committed `WIH1` payload wherever it is carried, finds none affected.
- **The inverse GDN on the fast kernel.** The interpreter now lends an inverse GDN to the
fast IGDN kernel, as it lends convolutions, wherever the values present keep every
term and norm within `i128` and the root and output within `i64`, which every
certified graph does; the splat and audio decodes use it too. Their outputs are
unchanged.
- **The unchecked decode is removed.** `int_hyper_synth::decode_conditional`, which ran
both networks with no exactness check, is gone; its callers (the examples, the
benchmark, the tests and the byte-exact neural-keyframe golden) decode through
`decode_conditional_checked`, so through the lowering with its `check` ranges, with
every output equal. The layer-by-layer guard `run_net_checked` stays while
`wai.neural.int_mlicpp` still runs its networks through it.
- **The tools' JSON models.** The browser crate's debug decode of the tools' JSON
form now refuses (`model_mismatch`) a JSON model holding a weight past `i16` or an
inverse-GDN `γ` past `i32`, the widths `WHM1` gives them and its lowering writes;
the committed models hold none.
- **Their `ops` member in the registry** ([integer-operators](extensions/integer-operators.md)
§11). `registry/capabilities.json` gives `wai.neural.int_hyper` and
`wai.video.int_hyper` the member `{"set": "wai-ops/1", "profile": "acc64", "subset":
"full", "lowering": "whm1/1", "stage": {"rans64": {"table": "t4097", "escape":
"lenient", "end": "canonical"}, "bucket": "ceil"}}`. SPEC §5's rows cite the payloads
([integer-payloads](extensions/integer-payloads.md) §4, and §3 for the clip) and the
lowering (integer-operators §13.2), and SPEC §7 names both capabilities among those
that decode through the operator set.
### wai-quantum 0.5.1: neural-processor, matrix-extension, vector-unit, signal-processor and dataflow members
Additive. The default features, `full`, and every pinned digest are unchanged.

11
wai/SPEC.md
View file

@ -693,7 +693,7 @@ registration and makes no conforming envelope non-conforming (§8).
| `wai.neural.bmshj2018` | bmshj2018-factorized + rANS (byte-exact CompressAI port) | about 0.26 bpp at 29 dB PSNR (quality 3) | CompressAI |
| `wai.neural.mbt2018_mean` | mean-scale hyperprior (Gaussian-conditional entropy) | the conditional-entropy family (arXiv 1809.02736) | `neural/hyperprior.rs` |
| `wai.neural.int_synth` | integer fixed-point synthesis (all-i64 transposed-conv, no float/ONNX) | bit-identical on every machine — soundly backs `intent = replicate` | payload and set: [integer-payloads](extensions/integer-payloads.md) §13; decoded through its lowering, [integer-operators](extensions/integer-operators.md) §13.1 (`int_synth/1`); `int_synth.rs` (integer weights sink-supplied; the reference tools make them by training in float, then quantising after training with per-layer shifts calibrated on measured activations, `tools/wai_int_synth_train.py`. Of the committed vectors only `photo` comes from a trained model, `tools/wai_int_synth_real.py`; `gradient`, `bloom`, `facade` and the cross-architecture golden use seeded integer weights) |
| `wai.neural.int_hyper` | integer fixed-point **conditional** codec — whole bitstream→image path integer (integer hyper-decoder σ → integer-σ entropy → integer synthesis) | the integer-entropy contract fully assembled: a PTQ of mbt2018_mean, bit-identical on every machine, whose integer decode CI holds within 1 dB of the same model's float decode on a pinned photograph | payload `WIH1`; `int_hyper_synth.rs` (tables sink-supplied) |
| `wai.neural.int_hyper` | integer fixed-point **conditional** codec — whole bitstream→image path integer (integer hyper-decoder σ → integer-σ entropy → integer synthesis) | the integer-entropy contract fully assembled: a PTQ of mbt2018_mean, bit-identical on every machine, whose integer decode CI holds within 1 dB of the same model's float decode on a pinned photograph | payload `WIH1`; `int_hyper_synth.rs` (tables sink-supplied). Payload: [integer-payloads](extensions/integer-payloads.md) §4; its networks decoded through their lowering, [integer-operators](extensions/integer-operators.md) §13.2 (`whm1/1`) |
| `wai.neural.int_mlicpp` | integer fixed-point **MLIC++ entropy topology** — `K` channel slices of the latent, each decoded as a checkerboard anchor pass then a non-anchor pass, with an entropy network over concat[latent decoded so far, hyper] re-run before every pass to give σ‖μ; hyper-decoder, entropy network and synthesis all integer | the channel-autoregressive × checkerboard entropy model at **decode-equivalence**: the encoder and every decoder compute the same σ, so the same distribution for every element, at any thread count; bit-identical on every machine. The committed conformance vector has synthetic weights, so it demonstrates decode-equivalence and states no quality figure | `int_mlicpp.rs` (one-file model bundle sink-supplied; `tools/wai_int_mlicpp_export.py`) |
| `wai.image.jpegai` | JPEG AI (ISO/IEC 6048-1:2025) codestream → RGB — integer entropy front end (mANS, with an integer scale-index derivation) + 16-bit-PTQ integer hyper-decoder, context-model and synthesis networks + dequantisation, latent refinement, post-filters (the eICCI CNN among them) and colour conversion in IEEE-754 float with a fixed evaluation order. Scope: one configuration — 256×256 4:4:4, `model_id` 0 (tools_0), `beta_displacement_log` 0, the high-operation-point synthesis, the eICCI model pair the weight bundle carries, and the other header conditions in [determinism-tiers](extensions/determinism-tiers.md) §5; both decode paths refuse any other codestream. Within it, a sink without the integer weight bundle, or in a non-default floating-point environment, takes a float ONNX decode instead. Only checkpoint constants are sink-supplied | **entropy-consistency** (registry class `NeuralEntropyExact`): the entropy decode, including the scale-index derivation, is integer on both decode paths, so the decoded symbols reproduce on every machine. The exact decoder also pins one reconstruction to RGB, reproducible across machines under the build and floating-point conditions in [determinism-tiers](extensions/determinism-tiers.md) §5–§6 (it checks the floating-point environment before decoding) — a property of that decoder, not of the capability, so `wai.image.jpegai` does not back `replicate` (§7). On one reference codestream the exact decoder stays within max\|Δ\| ≤ 2 of the reference decoder (recorded max\|Δ\| = 1 on ~4.4% of samples); evidence in [determinism-tiers](extensions/determinism-tiers.md) §6. **Downgraded in v0.4.0:** this row previously claimed decode-equivalence, and the registry classed the capability `StandardDefined` | `jpegai_decode.rs` (+ `jpegai_{bitstream,mans,pipeline,synthesis,icci}.rs`) |
| `wai.generate.spectral` | deterministic integer GENERATIVE synthesis (`intent=create`) — prompt+seed → integer-LUT spectral composition, no float/ONNX | byte-identical on every machine, pinned to (model, prompt, seed, steps); the create-side of "conformance is a hash" (generation-equivalence). Payload: the instruction object of form v0, which names the model by its SHA-256 and lists every size a sink makes, each output checked by its SHA-256 (*Payload of `wai.generate.spectral`*, below) | `generate.rs` (small pinned model bundle) |
@ -713,7 +713,7 @@ registration and makes no conforming envelope non-conforming (§8).
| `wai.neural.dcvc_rt` | DCVC-RT (2025), inter-frame with reference-feature propagation | inter-frame neural video; its paper reports real-time 1080p and 4K decode on a GPU | research code, `neural/video_dcvc.rs` |
| `wai.neural.dcvc_fm` | DCVC-FM (feature modulation), wide single-model quality range (RGB + YUV) | a wide single-model quality range; shares the inter-frame route | research code, `neural/video_dcvc.rs` |
| `wai.video.int_motion` | DETERMINISTIC INTEGER inter-frame — I-frame + integer block motion-compensated P-frames + integer-quantized residual over rANS, no float/ONNX; motion is whole-pel or **half-pel** (`mv_precision`), the latter predicting at the half-grid by integer bilinear interpolation (`(a+b+1)>>1`, `(a+b+c+d+2)>>2`) so real footage — which pans a non-integer number of pixels per frame — costs fewer residual bits while staying byte-exact | byte-identical frame-by-frame on every machine (*frame-state-equivalence*, a per-frame portable hash); the temporal extension of the integer contract, where a float inter-frame decode desyncs down the GOP | payload `WIV1`, keyframe 0; `int_video.rs` (the encode tools write a JSON form a converter reads) |
| `wai.video.int_hyper` | NEURAL-KEYFRAME inter-frame — the I-frame is the LEARNED integer codec `int_hyper` decoded byte-exact, the P-frames are integer block-MC residuals on top | a learned keyframe you can put in a receipt + integer inter-frames; whole decode integer, byte-identical everywhere | payload `WIV1`, keyframe 1 (a nested `WIH1`); `int_video::decode_pframes_typed` ∘ `int_hyper_synth::decode_conditional_checked` |
| `wai.video.int_hyper` | NEURAL-KEYFRAME inter-frame — the I-frame is the LEARNED integer codec `int_hyper` decoded byte-exact, the P-frames are integer block-MC residuals on top | a learned keyframe you can put in a receipt + integer inter-frames; whole decode integer, byte-identical everywhere | payload `WIV1`, keyframe 1 (a nested `WIH1`); `int_video::decode_pframes_typed` ∘ `int_hyper_synth::decode_conditional_checked`. Payload: [integer-payloads](extensions/integer-payloads.md) §3, its keyframe §4; the keyframe's networks decoded through their lowering, [integer-operators](extensions/integer-operators.md) §13.2 (`whm1/1`) |
| `wai.video.manifest` | ATTESTED PER-SESSION MANIFEST — the OTT delivery layer, not a codec: a personalised HLS manifest (ads stitched at `EXT-X-DISCONTINUITY`, blackouts + geo enforced) rendered as a **pure function of the session *class*** `(content, entitlement, geo, blackout set, ad pod, window)` rather than the viewer, and sealed in an Ed25519 receipt that binds the JCP grant (capability, joule ceiling, `funds_ceiling` = ad budget), the composition, the stitched ads and the policy applied. A receipt over a WAI playlist MAY carry `render` and `extras`, the rendering parameters and per-segment lines that playlist adds, both bound through `manifest_hash` and through a render digest the signature covers ([http-delivery](extensions/http-delivery.md) §8.6) | **manifest-equivalence**: `manifest_hash == hash(render(segments))` and `session_class_hash` are recomputable by anyone holding the bytes, so the ads, blackouts and compose-work stop being dashboard claims; class-determinism means one generation serves every viewer of the class, and `uj_per_viewer` is the amortised energy that a 1:1-unique generator cannot report below its per-viewer cost. Only the measurement — `joules_micro`, its acquisition class and the class's evidence — is measured-attested: signed, and not re-measurable | `wai_video_manifest.rs` |
| `wai.video.ssai` | ATTESTED SERVER-SIDE AD INSERTION — a **deterministic** ad decision (greedy by price within the avail duration and the funds remaining) plus a per-impression Ed25519 receipt binding the candidate pod, the manifest it belongs to, the served `creative_hash`es and the VAST quartile each reached | **spend-equivalence**: `budget_spent_micros` is recomputed from the *billable* impressions and rejected if it exceeds the grant's `funds_ceiling` — an overspent receipt fails verification; every creative and quartile is checkable, so a buyer and seller reconcile impressions from the same bytes instead of two logs that only a dashboard joins. Only the measurement — `joules_micro`, its acquisition class and the class's evidence — is measured-attested: signed, and not re-measurable | `wai_video_ssai.rs` |
| `wai.video.keys` | ATTESTED CONTENT-KEY RELEASE — the DRM licensing decision as a record: every key release, and every refusal, is an Ed25519 receipt binding the KID, the DRM system, the entitlement demanded and the window enforced, with a BLAKE3 `key_commit` proving *which* key was released without carrying it. It is one of Lane A's records (§6.5): it covers the release decision, never a frame the key opens | **decision-equivalence**: `verify()` re-runs the release policy over the receipt's own `(grant capability, entitlement, issued, expiry)` and rejects any misreported outcome, and enforces granted-iff-committed — so "keys went only to entitled, unexpired sessions" is checkable against the signed log, not asserted. Only the measurement — `joules_micro`, its acquisition class and the class's evidence — is measured-attested: signed, and not re-measurable | `wai_video_keys.rs` |
@ -1512,8 +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`, `wai.splat.int_codec` and
`wai.audio.int_codec`.
refused. In this revision: `wai.neural.int_synth`, `wai.splat.int_codec`,
`wai.audio.int_codec`, `wai.neural.int_hyper` and `wai.video.int_hyper`.
- **`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 =
@ -2408,6 +2408,9 @@ envelope that uses any string a later revision registers.
- `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.
- `wai.neural.int_hyper` and the keyframe of `wai.video.int_hyper`: the same, for
either of its networks, `h_s` and `g_s`: such a set is refused after the payload's
shapes are checked and before either stream is decoded.
---

42
wai/byte-exact-conformance/src/int_payload_golden.rs
View file

@ -116,3 +116,45 @@ fn wih1_payload_is_byte_identical() {
assert_eq!(hex(&Sha256::digest(&rgb)), "dce7282d46f88f03d5dc0fd42f325756d310fe1fc982b6e1792a1fbfbeb482cf");
assert_eq!(hex(&Sha256::digest(&rgb)), jload("int_hyper/sha256.json")["tiny"]["sha256"].as_str().unwrap());
}
/// The tiny conditional codec's keyframe through its lowering (`int_hyper_synth::ops`:
/// `h_s` and `g_s` by `whm1/1`, loaded under `acc64` and run by the integer operator set),
/// from its committed set.
fn keyframe_through_the_lowering(k: &Wih1<'_>) -> (Vec<u8>, usize, usize) {
let g_s = IntModel::from_bytes(&read("int_hyper/tiny_g_s.bin")).expect("g_s");
let h_s = IntModel::from_bytes(&read("int_hyper/tiny_h_s.bin")).expect("h_s");
let ent = CondEntropy::try_from_value(&jload("int_hyper/tiny_entropy.json")).expect("entropy tables");
let s = CondStream {
m: k.m as usize, sh: k.sh as usize, sw: k.sw as usize, n_z: k.n_z as usize, zh: k.zh as usize, zw: k.zw as usize,
z_rans: k.z_rans, rans: k.rans,
};
crate::int_hyper_synth::ops::decode(&g_s, &h_s, &ent, &s).expect("the keyframe decodes through its lowering")
}
/// `WIH1` through its lowering, under the reference and the fast kernels: the same image
/// SHA-256 as the decode the sink runs.
#[test]
fn wih1_through_its_lowering_is_golden() {
use crate::int_kernels::{with_policy, Policy};
let p = read("int_hyper/tiny_bitstream.wih1");
assert_eq!(hex(&Sha256::digest(&p)), TINY_WIH1_SHA256);
let k = parse_wih1(&p).expect("the payload parses");
for pol in [Policy::reference(), Policy::fast(1), Policy::fast(0)] {
let (rgb, h, w) = with_policy(pol, || keyframe_through_the_lowering(&k));
assert_eq!((h, w), (16, 16));
assert_eq!(hex(&Sha256::digest(&rgb)), "dce7282d46f88f03d5dc0fd42f325756d310fe1fc982b6e1792a1fbfbeb482cf", "{pol:?}");
}
}
/// `WIV1` with keyframe 1, its keyframe through its lowering: the same clip.
#[test]
fn wiv1_keyframe_1_through_its_lowering_is_golden() {
let p = read("int_video/video_nk_bitstream.wiv1");
assert_eq!(hex(&Sha256::digest(&p)), VIDEO_NK_WIV1_SHA256);
let v = parse_wiv1(&p, Some(1)).expect("the payload parses");
let k = parse_wih1(v.wih1.expect("keyframe 1 nests a WIH1")).expect("the keyframe parses");
let (frame0, h, w) = keyframe_through_the_lowering(&k);
assert_eq!((h, w), (v.h as usize, v.w as usize), "the keyframe is the clip's size");
let clip = decode_pframes_typed(&frame0, &stream(&v)).expect("integer P-frames");
assert_eq!(assert_clip(&clip.frames, "video_nk"), "4a81f0c94ccd938a7ca80df203303d5735d9dfcdfa150da155d9f2e4ca283e8f");
}

10
wai/byte-exact-conformance/src/int_video_golden.rs
View file

@ -131,15 +131,13 @@ fn int_video_real_footage_exercises_the_half_pel_predictor() {
}
// NOT covered here: `video_nk` (`wai.video.int_hyper`, the NEURAL-keyframe
// variant). Its I-frame is produced by `int_hyper_synth::decode_conditional`
// variant). Its I-frame is produced by `int_hyper_synth::decode_conditional_checked`
// and only its P-frames go through `int_video::decode_pframes`, so it cannot be
// driven by `decode()` — an earlier draft of this file tried and got a slice
// overrun out of the rANS reader, which is what feeding one codec another
// codec'"'"'s bitstream looks like. Covering it means `#[path]`-including
// `int_hyper_synth.rs` as well; that is a real follow-on, not a line to add
// here, and until it lands the neural-keyframe leg of the temporal contract
// remains proven only where someone runs `--features neural_int,video_int`
// locally.
// codec'"'"'s bitstream looks like. It is covered next door, in
// `neural_keyframe_golden.rs` (the JSON forms) and `int_payload_golden.rs` (the
// `WIV1` payload), which include `int_hyper_synth.rs`.
//
// The synthetic/real pair above already separates the two failure modes this
// file needs to distinguish: the synthetic clip is whole-pel only, so a failure

10
wai/byte-exact-conformance/src/neural_keyframe_golden.rs
View file

@ -14,6 +14,10 @@
//! keyframe, it propagates through the whole clip. A learned keyframe that
//! reproduces bit-for-bit on four architectures is the claim that makes
//! "a learned codec you can put in a receipt" mean anything.
//!
//! The keyframe decodes as a sink decodes it: through the integer operator set
//! (`int_hyper_synth::decode_conditional_checked`, the networks lowered by `whm1/1` and
//! run with their `check` ranges), from the same JSON models and bitstream.
#![cfg(test)]
@ -21,7 +25,7 @@ use std::path::PathBuf;
use serde_json::Value;
use crate::int_hyper_synth::{decode_conditional, CondBitstream, CondEntropy, IntModel};
use crate::int_hyper_synth::{decode_conditional_checked, CondBitstream, CondEntropy, IntModel};
use sha2::Digest;
use crate::int_video::{decode_pframes, fnv1a};
@ -55,7 +59,7 @@ fn the_neural_keyframe_clip_is_byte_identical() {
let h_s = IntModel::from_json(&jload("int_hyper/tiny_h_s.json"));
let ent = CondEntropy::from_value(&jload("int_hyper/tiny_entropy.json"));
let bs = CondBitstream::from_value(&jload("int_hyper/tiny_bitstream.json"));
let (frame0, h, w) = decode_conditional(&g_s, &h_s, &ent, &bs).expect("learned I-frame");
let (frame0, h, w) = decode_conditional_checked(&g_s, &h_s, &ent, &bs.stream()).expect("learned I-frame");
let conf = jload("int_video/video_nk.json");
let vbs = jload("int_video/video_nk_bitstream.json");
@ -109,7 +113,7 @@ fn the_learned_keyframe_alone_is_byte_identical() {
let h_s = IntModel::from_json(&jload("int_hyper/tiny_h_s.json"));
let ent = CondEntropy::from_value(&jload("int_hyper/tiny_entropy.json"));
let bs = CondBitstream::from_value(&jload("int_hyper/tiny_bitstream.json"));
let (frame0, h, w) = decode_conditional(&g_s, &h_s, &ent, &bs).expect("learned I-frame");
let (frame0, h, w) = decode_conditional_checked(&g_s, &h_s, &ent, &bs.stream()).expect("learned I-frame");
let conf = jload("int_video/video_nk.json");
let want0 = conf["frame_hashes"].as_array().expect("frame_hashes")[0]

10
wai/extensions/integer-operators.md
View file

@ -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`, `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.
> 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`, `wai.audio.int_codec`, `wai.neural.int_hyper` and `wai.video.int_hyper`. 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.
@ -902,6 +902,12 @@ Before a sink decodes through a parameter set, it MUST establish the set's verdi
- A set whose verdict does not meet its capability's declared profile (§11) is refused (`model_mismatch`).
- So is a set whose pin claims a profile that its verdict does not meet.
*Informative: what a sink may keep between decodes.* A sink may keep more than the verdict: the set's graphs as loaded, so that a later decode through the same set skips the lowering and the load, whose bounds pass can cost more than a decode for a large set.
- **The key** is what makes this sound. It is the SHA-256 of the exact bytes the set was lowered from (for a set of several members, of their listing, which is the digest a pin over exactly those members names, [SPEC](../SPEC.md) §3.1), together with the capability, the lowering and the profile the graphs were loaded under. A name, a path or a label does not fix the bytes: keyed by one, a set could decode with the graphs of another set of the same name.
- **Only what loaded.** An entry is made after every load step of §10.5 has passed. A set that is refused is not kept, so it is lowered and loaded again on its next decode and refused again, with the same code.
- **Nothing else moves.** A decode makes every check before the load as before, in the order its payload extension gives ([integer-payloads](integer-payloads.md) §7), and the graphs run as they do after a fresh load, so the output is the same bytes.
- **Memory.** What such entries hold is the set's, the same for every payload, and outside a decode's declared working bytes ([integer-payloads](integer-payloads.md) §7). The sink bounds it with its own budget. The reference implementation keeps them in `int_set_cache`.
## 9. Subsets
### 9.1 `wai-ops/1` and `wai-ops/1-i32`
@ -1528,7 +1534,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`. 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.
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`, `wai.audio.int_codec`, `wai.neural.int_hyper` and `wai.video.int_hyper` 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. For the conditional codec either network can carry the change, and the set is refused after its plan and its shape checks and before either stream (integer-payloads §7), where the direct decoder refused with the code of whichever stream or network failed first.
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.

37
wai/extensions/integer-payloads.md
View file

@ -204,13 +204,18 @@ one fixed-point scale `2^afrac`; `entropy.json`, the tables):
symbol `s` from the hyper-latent stream against channel `k`'s table and offset;
the element is `s·2^afrac + median_k`.
2. `h_s` on the hyper-latent `[n_z, zh, zw]` gives `[2·m, sh, sw]`: its first `m`
channels are the scales `σ`, the rest the means `μ`.
channels are the scales `σ`, the rest the means `μ`. `h_s` decodes through its
lowering `whm1/1` ([integer-operators](integer-operators.md) §13.2); the split of its
output into `σ` and `μ` is a step of this decode, after the shape check, not part of
the lowered graph (integer-operators §13.5).
3. For each latent element `p` in order (channel, then row, then column), its table
is `b`, where, with the scale edges `e_0 … e_{N−1}` and `σ' = max(σ_p, e_0)`,
`b = N − 1 − #{ i < N − 1 : σ' ≤ e_i }`. Decode a symbol `r` from the latent
stream against table `b` and its offset; the element is `r·2^afrac + μ_p`.
4. `g_s` on the latent `[m, sh, sw]` gives three channels; each value `v` is clamped
to `[0, 2^afrac]` and becomes `(v·255 + 2^(afrac−1)) >> afrac`, as RGB8.
to `[0, 2^afrac]` and becomes `(v·255 + 2^(afrac−1)) >> afrac`, as RGB8. `g_s`
decodes through its lowering `whm1/1` (integer-operators §13.2), and this is the RGB8
step that ends its lowered graph (§13.0).
Every value is formed in `i64` and must fit it (§7).
@ -360,6 +365,17 @@ registry resolved; the C ABI's `wai_int_payload_declared`, given the set's path;
PCM or splat attributes, which it has no output kind for. This extension fixes the
caps; a sink's policy may be stricter.
*Informative.* A sink may keep a parameter set lowered and loaded between decodes,
keyed by the SHA-256 of its bytes ([integer-operators](integer-operators.md) §8.6),
so that only the first decode through it pays for the lowering and the bounds pass.
What it keeps is the set's, the same for every payload, and is not part of any
decode's declared working bytes; the sink's own budget for kept sets bounds it. The
reference implementation's store (`int_set_cache::LoadedSets`) starts with a budget of
256 MiB, evicts the least recently used set first, keeps no set larger than the whole
budget, and keeps nothing with a budget of 0 (`LoadedSets::set_budget`; `wai-web`'s
`setLoadedSetBudget`). A decode that finds its set kept makes every other check of this
section, in the order below, and gives the same output and the same refusals.
**The parameter set must fit the payload** (`model_mismatch`):
- every network parses (the `WHM1` and `WAM1` forms, §9), and the networks of one set
@ -373,8 +389,8 @@ caps; a sink's policy may be stricter.
geometry leaves an output (*Planning the networks*);
- 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).
`acc64` ([integer-operators](integer-operators.md) §13.2 and §13.3: for `WIH1` and a
`WIV1`'s keyframe, `h_s` and `g_s`; 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
@ -393,9 +409,11 @@ order and refuses with the first that fails:
the splat `g_s`), each layer by the three rules above (`model_mismatch`,
`too_large`): the caps come before any shape is compared;
4. what the networks give against the payload: for `WIH1`, `h_s`'s shape
(`geometry_mismatch`), then `g_s`'s three channels (`model_mismatch`); for `WIV1`
with keyframe 1, its keyframe's as for `WIH1`, then the keyframe's size against
`h × w` (`geometry_mismatch`); for `WIA1`,
(`geometry_mismatch`), then `g_s`'s three channels (`model_mismatch`), then `h_s`
and `g_s` lowered and loaded, in that order (`model_mismatch`); for `WIV1` with
keyframe 1, its keyframe's shape and channels as for `WIH1`, then the keyframe's size
against `h × w` (`geometry_mismatch`), then the keyframe's networks lowered and loaded
(`model_mismatch`); for `WIA1`,
one channel (`model_mismatch`), then at least `t_out` samples, then a latent no
longer than that needs (both `geometry_mismatch`), then the synthesis lowered and
loaded (`model_mismatch`); for `WIS1`, `g_s`'s shape
@ -418,8 +436,9 @@ 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. 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
([integer-operators](integer-operators.md) §13; for `WIH1`, a `WIV1`'s keyframe and
`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.

32
wai/registry/capabilities.json
View file

@ -724,7 +724,21 @@
"sets": [],
"spec_section": "§5 Neural capabilities",
"edition": null,
"registered": "2026-06-24"
"registered": "2026-06-24",
"ops": {
"set": "wai-ops/1",
"profile": "acc64",
"subset": "full",
"lowering": "whm1/1",
"stage": {
"rans64": {
"table": "t4097",
"escape": "lenient",
"end": "canonical"
},
"bucket": "ceil"
}
}
},
{
"capability": "wai.neural.int_mlicpp",
@ -1283,7 +1297,21 @@
"sets": [],
"spec_section": "§5 Neural capabilities",
"edition": null,
"registered": "2026-06-27"
"registered": "2026-06-27",
"ops": {
"set": "wai-ops/1",
"profile": "acc64",
"subset": "full",
"lowering": "whm1/1",
"stage": {
"rans64": {
"table": "t4097",
"escape": "lenient",
"end": "canonical"
},
"bucket": "ceil"
}
}
},
{
"capability": "wai.video.int_motion",

2
wai/tools/wai_int_video_keyframe_encode.py
View file

@ -5,7 +5,7 @@ on top — the whole decode integer, byte-exact on every machine.
The I-frame here is the committed tiny `int_hyper` vector (so the conformance is
self-contained and never-skips); the P-frames translate it and code the residual.
The Rust decoder composes `int_hyper::decode_conditional` (frame 0) with
The Rust decoder composes `int_hyper_synth::decode_conditional_checked` (frame 0) with
`int_video::decode_pframes` (the rest) — this script emits only the P-frame
bitstream + the expected per-frame / clip hashes (frame 0's hash is int_hyper's).

2
wai/tools/wai_int_video_sota_keyframe_encode.py
View file

@ -12,7 +12,7 @@ tools/wai_sota_vectors.sh runs the whole chain.
Emits <dir>/video_sota_bitstream.json + <dir>/video_sota.json (the P-frame stream +
expected per-frame / clip hashes; frame 0's hash is int_hyper's). The Rust test
`int_video_sota_keyframe` composes int_hyper::decode_conditional with
`int_video_sota_keyframe` composes int_hyper_synth::decode_conditional_checked with
int_video::decode_pframes and checks these (skipping loudly when they are absent,
failing instead under WAI_REQUIRE_SOTA_VECTORS=1).
"""

53
wai/wai-ops/src/graph.rs
View file

@ -863,6 +863,40 @@ impl Graph {
}
o
}
/// The bytes this graph holds on the heap: every list of its records, operators and
/// parameters, at its capacity times its element's size. With the size of the value
/// itself, this is what a sink that keeps a loaded graph between decodes (§8.6)
/// holds for it.
pub fn heap_bytes(&self) -> u64 {
fn vec<T>(v: &Vec<T>) -> u64 {
(v.capacity() as u64).saturating_mul(core::mem::size_of::<T>() as u64)
}
fn rq(r: &Rq) -> u64 {
vec(&r.mult) + vec(&r.shift) + vec(&r.bias)
}
// Every variant is named, so a parameter list added to one is counted here too.
let params = |p: &Params| match p {
Params::Conv { rq: r, w, .. } | Params::ConvT { rq: r, w, .. } => rq(r) + vec(w),
Params::Rescale(r) | Params::Mul(r) | Params::Gram { rq: r, .. } | Params::Apply { rq: r, .. } => rq(r),
Params::Chscale { m, .. } => vec(m),
Params::Lut16 { t, .. } => vec(t),
Params::Igdn { gamma, beta, .. } | Params::LayerNorm { gamma, beta, .. } => vec(gamma) + vec(beta),
Params::None
| Params::Split { .. }
| Params::Shuffle { .. }
| Params::Clamp { .. }
| Params::Leaky { .. }
| Params::PadRep { .. }
| Params::Crop { .. }
| Params::Subpos { .. }
| Params::Fill { .. }
| Params::L2norm { .. }
| Params::Elu { .. } => 0,
};
let ops: u64 = self.ops.iter().map(|o| vec(&o.ins) + vec(&o.outs) + params(&o.p)).sum();
vec(&self.scalars) + vec(&self.inputs) + vec(&self.outputs) + vec(&self.tensors) + vec(&self.ops) + ops
}
}
fn put_ints(o: &mut Vec<u8>, v: &[i64], w: u8) {
@ -1075,6 +1109,25 @@ mod tests {
assert_eq!(Graph::parse(&b).unwrap().to_bytes(), b);
}
/// The heap bytes are every list's capacity times its element's size: the records,
/// the operators, their operand lists, and a parameter's weights and requantisation
/// lists at the widths they are held in (`i64`, `u8`).
#[test]
fn heap_bytes_count_every_list() {
use core::mem::size_of;
let g = tiny();
let records = size_of::<InputRec>() + 4 + 2 * size_of::<TensorRec>() + size_of::<Op>() + 4 + 4;
assert_eq!(g.heap_bytes(), records as u64);
let mut c = tiny();
let rq = Rq { round: Round::Floor, bias_kind: BiasKind::Pre, m_kind: 2, s_kind: 1, m_bytes: 8, b_bytes: 8, mult: vec![1, 1], shift: vec![0], bias: vec![0, 0] };
c.ops[0].p = Params::Conv { cout: 2, kh: 1, kw: 1, sh: 1, sw: 1, ph: 0, pw: 0, groups: 1, w_bytes: 2, rq, w: vec![1, 2, 3] };
assert_eq!(c.heap_bytes() - g.heap_bytes(), 3 * 8 + 2 * 8 + 1 + 2 * 8);
// A list's spare capacity is held too.
let mut spare = tiny();
spare.outputs.reserve_exact(7);
assert_eq!(spare.heap_bytes() - g.heap_bytes(), 4 * (spare.outputs.capacity() as u64 - 1));
}
/// Load step 8: the graph loads under both profiles; with an input range past
/// i32 it loads under `acc64` only; with a `check` range that its producer cannot
/// meet, under neither.

11
wai/wai-rs/Cargo.toml
View file

@ -766,6 +766,11 @@ name = "int_splat_ops"
path = "tests/int_splat_ops.rs"
required-features = ["splat_int", "legacy_int"]
[[test]]
name = "int_hyper_ops"
path = "tests/int_hyper_ops.rs"
required-features = ["video_int", "legacy_int"]
[[test]]
name = "int_mlicpp"
path = "tests/int_mlicpp.rs"
@ -806,6 +811,12 @@ name = "int_payload_cost"
path = "tests/int_payload_cost.rs"
required-features = ["video_int", "audio_int", "splat_int"]
# The store of loaded parameter sets under the decodes that use it (src/int_set_cache.rs).
[[test]]
name = "int_set_cache"
path = "tests/int_set_cache.rs"
required-features = ["video_int", "audio_int", "splat_int"]
[[test]]
name = "int_video_enc_memory"
path = "tests/int_video_enc_memory.rs"

4
wai/wai-rs/examples/int_hyper_decode_frame.rs
View file

@ -8,7 +8,7 @@
//! `int_synth_sota_model.json`, `int_hs_model.json`, `cond_entropy.json`,
//! `cond_bitstream.json`.
use wai::int_hyper_synth::{decode_conditional, CondBitstream, CondEntropy, IntModel};
use wai::int_hyper_synth::{decode_conditional_checked, CondBitstream, CondEntropy, IntModel};
fn main() {
let a: Vec<String> = std::env::args().collect();
@ -24,7 +24,7 @@ fn main() {
let h_s = IntModel::from_json(&jl("int_hs_model.json"));
let ent = CondEntropy::from_value(&jl("cond_entropy.json"));
let bs = CondBitstream::from_value(&jl("cond_bitstream.json"));
let (rgb, h, w) = decode_conditional(&g_s, &h_s, &ent, &bs).expect("decode_conditional");
let (rgb, h, w) = decode_conditional_checked(&g_s, &h_s, &ent, &bs.stream()).expect("the conditional decode");
std::fs::write(&a[2], &rgb).unwrap();
eprintln!("{h}x{w} -> {} ({} bytes)", a[2], rgb.len());
}

14
wai/wai-rs/include/wai.h
View file

@ -153,6 +153,20 @@
*/
#define WAI_ERR_NO_RENDITION 7
/**
* The budget a process starts with: 256 MiB. A set holds its weights as `i64` once
* lowered, four times their bytes in a `WHM1` file, so this keeps a set of about
* 60 MB of `WHM1` networks with room beside it for small ones.
*/
#define LoadedSets_DEFAULT_BUDGET (256 << 20)
/**
* What an entry is counted for beyond its set: its share of the store's two indexes
* (by key, and by last use). It is more than either index's node for one entry, so
* the count is never under what the store holds.
*/
#define LoadedSets_INDEX_BYTES_PER_ENTRY 4096
/**
* Probability output scale of [`softmax`]: results are at `2^PROB_FRAC`
* (they sum to ~`2^PROB_FRAC`).

4
wai/wai-rs/src/bin/wai_decode_bench.rs
View file

@ -321,7 +321,7 @@ fn read_json(p: &str) -> Value {
}
fn vector_cases(cases: &mut Vec<Case>) {
use wai::int_hyper_synth::{decode_conditional, CondBitstream, CondEntropy};
use wai::int_hyper_synth::{decode_conditional_checked, CondBitstream, CondEntropy};
// int_hyper tiny conditional codec (entropy decode + h_s + g_s). Its weights are
// seeded random integers (tools/wai_int_hyper_tiny_vector.py), not a trained model.
{
@ -336,7 +336,7 @@ fn vector_cases(cases: &mut Vec<Case>) {
name: format!("vector · int_hyper tiny conditional codec, random weights ({px} px, committed)"),
per_run: (1.0, "images"),
macs: None,
run: Box::new(move || Out::B(decode_conditional(&g_s, &h_s, &ent, &bs).unwrap().0)),
run: Box::new(move || Out::B(decode_conditional_checked(&g_s, &h_s, &ent, &bs.stream()).unwrap().0)),
});
}
// int_splat: trained Gaussian-splat attribute decode.

4
wai/wai-rs/src/codecs/mod.rs
View file

@ -1359,6 +1359,10 @@ pub fn capability_ops(cap: &str) -> Option<OpsMember> {
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),
// The conditional codec buckets each latent's σ over the scale edges with E3's
// `ceil` rule (integer-payloads §4 step 3); the keyframe of
// `wai.video.int_hyper` is that codec, and its residual stream is §2's too.
CAP_NEURAL_INT_HYPER | CAP_VIDEO_INT_HYPER => ops("whm1/1", PAYLOAD_STREAM, Some("ceil")),
_ => None,
}
}

11
wai/wai-rs/src/codecs/registry.rs
View file

@ -1420,14 +1420,15 @@ 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`, `wai.splat.int_codec` and
/// `wai.audio.int_codec` decode through the operator set.
/// In this revision `wai.neural.int_synth`, `wai.splat.int_codec`,
/// `wai.audio.int_codec`, `wai.neural.int_hyper` and `wai.video.int_hyper` decode
/// through the operator set.
#[test]
fn ops_members_name_the_operator_set_a_capability_decodes_through() {
let regs = registrations();
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]);
assert_eq!(with, [CAP_AUDIO_INT_CODEC, CAP_NEURAL_INT_HYPER, CAP_NEURAL_INT_SYNTH, CAP_SPLAT_INT_CODEC, CAP_VIDEO_INT_HYPER]);
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);
@ -1446,6 +1447,10 @@ mod tests {
"set": "wai-ops/1", "profile": "acc64", "subset": "full", "lowering": "wam1/1",
"stage": { "rans64": { "table": "t4097", "escape": "lenient", "end": "canonical" } }
}),
CAP_NEURAL_INT_HYPER | CAP_VIDEO_INT_HYPER => serde_json::json!({
"set": "wai-ops/1", "profile": "acc64", "subset": "full", "lowering": "whm1/1",
"stage": { "rans64": { "table": "t4097", "escape": "lenient", "end": "canonical" }, "bucket": "ceil" }
}),
_ => serde_json::Value::Null,
};
assert_eq!(c["ops"], want, "{}", c["capability"]);

9
wai/wai-rs/src/int_audio.rs
View file

@ -527,6 +527,15 @@ pub mod ops {
/// 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> {
plan_with(m, s, load)
}
/// [`plan`], with the synthesis lowered and loaded by `load`, after every check
/// before it: a sink that keeps loaded sets (`crate::int_set_cache`) answers that one
/// step from its store.
pub(crate) fn plan_with<L>(
m: &AudioModel, s: &AudioStream<'_>, load: impl FnOnce(&AudioModel) -> Result<L, Refused>,
) -> Result<(L, u64, u128), Refused> {
let (_, working, macs) = audio_plan(m, s)?;
Ok((load(m)?, working, macs))
}

297
wai/wai-rs/src/int_hyper_synth.rs
View file

@ -484,9 +484,11 @@ impl std::fmt::Display for NetRefusal {
/// 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.
/// splat, audio and conditional decodes do (`int_splat`, `int_audio`,
/// [`decode_conditional_checked`]; their direct decodes, which call this or
/// [`igdn_check`], are test-only). Its one production caller left is
/// `int_mlicpp`, whose networks move to their lowering (`wmb1/1`) next, and it goes
/// when they have.
pub fn run_net_checked(
model: &IntModel, input: &[i64], mut c: usize, mut h: usize, mut w: usize,
) -> Result<(Vec<i64>, usize, usize, usize), NetRefusal> {
@ -820,52 +822,6 @@ impl CondBitstream {
}
}
/// Decode a full conditional-codec **bitstream → image**, end to end and fully
/// integer — the assembled `NeuralIntegerExact` codec (BOTH streams entropy-coded):
///
/// 1. factorized rANS-decode `z`: per channel, `z_q = s·2^afrac + median_q` (no
/// float) → integer hyper-decoder `h_s` → σ‖μ (all i64);
/// 2. each `y` element: σ → bucket ([`det_build_index`], no float) → rANS-decode
/// the residual `r` → `y_hat = r·2^afrac + μ` (μ from the integer `h_s`);
/// 3. integer synthesis `g_s` ([`decode`]) → RGB (HWC `u8`).
///
/// Encoder and decoder agree on every bucket because σ comes from the **integer**
/// `h_s` — the integer-entropy contract. Byte-identical on every machine.
pub fn decode_conditional(
g_s: &IntModel, h_s: &IntModel, ent: &CondEntropy, bs: &CondBitstream,
) -> Result<(Vec<u8>, usize, usize), &'static str> {
let afrac = g_s.afrac;
let one = 1i64 << afrac;
// 1. factorized-decode z (per channel), rebuild z_q = s·2^afrac + median_q.
let mut z_int = vec![0i64; bs.n_z * bs.zh * bs.zw];
{
let mut dz = RansDecoder::new(&bs.z_rans)?;
let plane = bs.zh * bs.zw;
for ch in 0..bs.n_z {
for p in 0..plane {
let s = dz.decode(&ent.z_cdfs[ch], ent.z_cdf_lengths[ch], ent.z_offsets[ch]);
z_int[ch * plane + p] = s as i64 * one + ent.z_median_q[ch];
}
}
}
let (sm, oc, sh, sw) = run_net(h_s, &z_int, bs.n_z, bs.zh, bs.zw);
if oc != 2 * bs.m || sh != bs.sh || sw != bs.sw {
return Err("h_s output shape disagrees with the bitstream header");
}
let n = bs.m * sh * sw;
let (sigma, mu) = sm.split_at(n); // σ = first M channels, μ = next M
// 2. entropy-decode y, conditioned on the integer σ; y_hat = r·2^afrac + μ.
let mut dec = RansDecoder::new(&bs.rans)?;
let mut y_hat = vec![0i64; n];
for p in 0..n {
let b = det_build_index(sigma[p], &ent.edges);
let r = dec.decode(&ent.cdfs[b], ent.cdf_lengths[b], ent.offsets[b]);
y_hat[p] = r as i64 * one + mu[p];
}
// 3. integer synthesis g_s → pixels.
Ok(decode(g_s, &y_hat, bs.m, sh, sw))
}
/// The kind of rule a checked decode refused on. The integer-payload decoders
/// (`int_payload`, and `int_audio` and `int_splat`'s typed decoders) report these as
/// the refusal codes of `extensions/integer-payloads.md` §8.
@ -1114,44 +1070,50 @@ fn plan_layers(net: &IntModel, c: usize, h: usize, w: usize) -> Result<NetPlan,
Ok(NetPlan { out: (c as usize, h as usize, w as usize), largest, live, macs })
}
/// [`run_net_checked`] with its refusal as a [`Refused`]: values that could overflow a
/// layer are the stream's (`FieldRange`); a layer that does not fit is the model's.
fn run_checked(net: &IntModel, x: &[i64], c: usize, h: usize, w: usize) -> Result<Vec<i64>, Refused> {
run_net_checked(net, x, c, h, w).map(|(o, ..)| o).map_err(|r| {
if r.out_of_range {
refused(RefusalKind::FieldRange, "a decoded latent could take a layer outside the exact integer range")
} else {
refused(RefusalKind::ModelMismatch, "a network's parameters do not fit the running shapes")
}
})
}
/// [`decode_conditional`], refusing instead of panicking or overflowing on any input:
/// the decode the `WIH1` payload (`wai.neural.int_hyper`, and the keyframe of
/// `wai.video.int_hyper`) runs. On a stream [`decode_conditional`] decodes, both give
/// the same image.
/// Decode a full conditional-codec **bitstream → image**, end to end and fully integer,
/// refusing instead of panicking or overflowing on any input: the decode the `WIH1`
/// payload (`wai.neural.int_hyper`, and the keyframe of `wai.video.int_hyper`) runs.
/// The hyper-latent `z` is decoded per channel against its factorized table and rebuilt
/// as `s·2^afrac + median_q`; the integer hyper-decoder `h_s` gives `σ‖μ`; each latent
/// element's table is `σ`'s bucket ([`det_build_index`]), and the element is
/// `r·2^afrac + μ`; the integer synthesis `g_s` gives RGB. Encoder and decoder agree on
/// every bucket because `σ` comes from the integer `h_s`: the integer-entropy contract.
/// Byte-identical on every machine.
///
/// Before anything is allocated, in this order: the networks share one fixed-point
/// scale in `1..=`[`MAX_AFRAC`]; the entropy model is consistent and has a `z` table
/// for each of the payload's channels; the latents are within [`MAX_TENSOR`]; `h_s` on
/// the hyper-latent's shape, then `g_s` on `[m, sh, sw]`, are planned ([`plan_net`]: each
/// layer fits, forms nothing over the cap, and leaves an output); then `h_s` must give
/// `[2·m, sh, sw]` (`GeometryMismatch`) and `g_s` three channels. Then the hyper-latent
/// is decoded symbol by symbol (a symbol that runs the stream out is refused before
/// its value is rebuilt with [`dequantize`]), its stream must end canonically, `h_s`
/// runs through [`run_net_checked`], the latent is decoded the same way, and `g_s`
/// runs. Returns RGB HWC with its height and width.
/// Before anything is allocated, in this order (`extensions/integer-payloads.md` §7): the
/// networks share one fixed-point scale in `1..=`[`MAX_AFRAC`]; the entropy model is
/// consistent and has a `z` table for each of the payload's channels; the latents are
/// within [`MAX_TENSOR`]; `h_s` on the hyper-latent's shape, then `g_s` on `[m, sh, sw]`,
/// are planned ([`plan_net`]: each layer fits, forms nothing over the cap, and leaves an
/// output); then `h_s` must give `[2·m, sh, sw]` (`GeometryMismatch`) and `g_s` three
/// channels; then `h_s` is lowered by `whm1/1` and `g_s` by `whm1/1` with the RGB8 step,
/// and each is loaded under `acc64` with its bounds certificate
/// (`extensions/integer-operators.md` §13.2, §13.5). Then the hyper-latent is decoded
/// symbol by symbol (a symbol that runs the stream out is refused before its value is
/// rebuilt with [`dequantize`](crate::int_entropy::dequantize)), its stream must end
/// canonically, `h_s` runs through the integer operator set ([`crate::int_ops`]), whose
/// `check` ranges refuse a value that could take a layer outside the exact integer range
/// (`FieldRange`), its output is split into `σ` and `μ`, the latent is decoded the same
/// way, and `g_s` runs, ending with the RGB8 step. Returns RGB HWC with its height and
/// width.
///
/// This is [`ops::decode`]. The direct decode it replaced is kept, test-only (feature
/// `legacy_int`), as `legacy::decode_conditional_checked`.
pub fn decode_conditional_checked(
g_s: &IntModel, h_s: &IntModel, ent: &CondEntropy, bs: &CondStream<'_>,
) -> Result<(Vec<u8>, usize, usize), Refused> {
ops::decode(g_s, h_s, ent, bs)
}
/// Step 1 of the decode, after its plan (`extensions/integer-payloads.md` §4): per
/// hyper-latent channel, the factorized rANS decode, each symbol rebuilt as
/// `s·2^afrac + median_q`. The hyper-latent grows a symbol at a time, to its declared
/// length and no further, and the first symbol that runs its stream out refuses the
/// payload before its value is used; the stream must then end canonically.
fn hyper_latent(ent: &CondEntropy, bs: &CondStream<'_>, afrac: u32) -> Result<Vec<i64>, Refused> {
use crate::int_entropy::dequantize;
use RefusalKind::*;
let p = cond_plan(g_s, h_s, ent, bs)?;
let (afrac, n_z, zh, zw, m, sh, sw) = (g_s.afrac, bs.n_z, bs.zh, bs.zw, bs.m, bs.sh, bs.sw);
let (zplane, n, (oc, oh, ow)) = (zh * zw, m * sh * sw, p.gs.out);
// 1. factorized-decode z, rebuild z_q = s·2^afrac + median_q. The latent grows a
// symbol at a time, to its declared length and no further, and the first symbol
// that runs its stream out refuses the payload before its value is used.
let (n_z, zplane) = (bs.n_z, bs.zh * bs.zw);
let mut z = Vec::new();
let mut dz = RansDecoder::new(bs.z_rans).map_err(|_| refused(FieldRange, "the z stream is not a rANS stream"))?;
for ch in 0..n_z {
@ -1166,10 +1128,18 @@ pub fn decode_conditional_checked(
if !dz.is_canonical_end() {
return Err(refused(StreamNoncanonical, "the z stream does not end canonically after its last symbol"));
}
let sm = run_checked(h_s, &z, n_z, zh, zw)?;
drop(z);
let (sigma, mu) = sm.split_at(n);
// 2. y, conditioned on the integer σ: y_hat = r·2^afrac + μ.
Ok(z)
}
/// Step 3 of the decode (integer-payloads §4), after `h_s` has run: each latent element,
/// in order, conditioned on the integer `σ` (its table is `σ`'s `ceil` bucket over the
/// scale edges, [`det_build_index`]) and rebuilt as `r·2^afrac + μ`, with the same
/// growth and stream checks as [`hyper_latent`]. `sigma` and `mu` are `h_s`'s channels
/// `[0, m)` and `[m, 2m)`.
fn latent(ent: &CondEntropy, bs: &CondStream<'_>, afrac: u32, sigma: &[i64], mu: &[i64]) -> Result<Vec<i64>, Refused> {
use crate::int_entropy::dequantize;
use RefusalKind::*;
let n = bs.m * bs.sh * bs.sw;
let mut dec = RansDecoder::new(bs.rans).map_err(|_| refused(FieldRange, "the y stream is not a rANS stream"))?;
let mut y_hat = Vec::new();
for p in 0..n {
@ -1183,11 +1153,50 @@ pub fn decode_conditional_checked(
if !dec.is_canonical_end() {
return Err(refused(StreamNoncanonical, "the y stream does not end canonically after its last symbol"));
}
drop(sm);
// 3. integer synthesis g_s → pixels.
let x = run_checked(g_s, &y_hat, m, sh, sw)?;
drop(y_hat);
Ok((to_rgb8(&x, oc, oh, ow, afrac), oh, ow))
Ok(y_hat)
}
/// The direct decode, as the conditional codec ran before it decoded through the integer
/// operator set: `h_s` and `g_s` run layer by layer by [`run_net_checked`], which checks
/// each layer's arithmetic on the values it is given before it runs, then the RGB8 bytes
/// formed here ([`to_rgb8`]). 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::*;
/// [`run_net_checked`] with its refusal as a [`Refused`]: values that could overflow a
/// layer are the stream's (`FieldRange`); a layer that does not fit is the model's.
fn run_checked(net: &IntModel, x: &[i64], c: usize, h: usize, w: usize) -> Result<Vec<i64>, Refused> {
run_net_checked(net, x, c, h, w).map(|(o, ..)| o).map_err(|r| {
if r.out_of_range {
refused(RefusalKind::FieldRange, "a decoded latent could take a layer outside the exact integer range")
} else {
refused(RefusalKind::ModelMismatch, "a network's parameters do not fit the running shapes")
}
})
}
/// The direct decode of [`super::decode_conditional_checked`]'s arguments.
pub fn decode_conditional_checked(
g_s: &IntModel, h_s: &IntModel, ent: &CondEntropy, bs: &CondStream<'_>,
) -> Result<(Vec<u8>, usize, usize), Refused> {
let p = cond_plan(g_s, h_s, ent, bs)?;
let (afrac, n, (oc, oh, ow)) = (g_s.afrac, bs.m * bs.sh * bs.sw, p.gs.out);
// 1. the hyper-latent, then h_s → σ‖μ.
let z = hyper_latent(ent, bs, afrac)?;
let sm = run_checked(h_s, &z, bs.n_z, bs.zh, bs.zw)?;
drop(z);
let (sigma, mu) = sm.split_at(n);
// 2. y, conditioned on the integer σ: y_hat = r·2^afrac + μ.
let y_hat = latent(ent, bs, afrac, sigma, mu)?;
drop(sm);
// 3. integer synthesis g_s → pixels.
let x = run_checked(g_s, &y_hat, bs.m, bs.sh, bs.sw)?;
drop(y_hat);
Ok((to_rgb8(&x, oc, oh, ow, afrac), oh, ow))
}
}
/// Append `x` to `v`, a buffer that will hold `total` elements, growing it by doubling
@ -1271,6 +1280,114 @@ impl CondPlan {
}
}
/// The conditional codec's networks through the integer operator set
/// (`extensions/integer-operators.md`): `h_s` lowered by `whm1/1`, and `g_s` lowered by
/// `whm1/1` with the RGB8 step (§13.2, §13.0; §13.0's roles, in that order), each loaded
/// under `acc64` with its bounds certificate ([`crate::int_ops`]).
pub mod ops {
use super::*;
use wai_ops::{Code, Graph, Loaded, Refusal};
/// A refusal of the operator set as the conditional codec 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)
}
const H_S: &str = "h_s does not lower to an integer-operator graph that loads under acc64";
const G_S: &str = "g_s does not lower to an integer-operator graph that loads under acc64";
/// One network lowered by `whm1/1`, with the RGB8 step where `rgb8` is set, from the
/// model this crate parsed ([`crate::int_ops::whm1_net`]).
fn lowered(m: &IntModel, rgb8: bool, what: &'static str) -> Result<Graph, Refused> {
let net = crate::int_ops::whm1_net(m).map_err(|r| kind(r, what))?;
wai_ops::lower::whm1_net(&net, rgb8).map_err(|r| kind(r, what))
}
/// The set's two graphs, `h_s` then `g_s` (with the RGB8 step), as `whm1/1` writes them
/// from its members ([`wai_ops::lower::int_hyper`]).
pub fn lower(g_s: &IntModel, h_s: &IntModel) -> Result<(Graph, Graph), Refused> {
Ok((lowered(h_s, false, H_S)?, lowered(g_s, true, G_S)?))
}
/// The set's two networks, lowered and loaded: the set as the decode holds it, beside
/// the parsed models.
pub struct Nets {
pub h_s: Loaded,
pub g_s: Loaded,
}
/// `h_s`, then `g_s`, each lowered and loaded under `acc64` in turn, so that no more
/// than one network's lowering is held at once.
pub fn load(g_s: &IntModel, h_s: &IntModel) -> Result<Nets, Refused> {
let h = crate::int_ops::load(&lowered(h_s, false, H_S)?).map_err(|r| kind(r, H_S))?;
let g = crate::int_ops::load(&lowered(g_s, true, G_S)?).map_err(|r| kind(r, G_S))?;
Ok(Nets { h_s: h, g_s: g })
}
/// Everything the decode checks before it decodes a symbol: [`cond_plan`], then the
/// two networks lowered and loaded ([`load`]). Returns the plan and the loaded
/// networks.
pub(crate) fn plan(g_s: &IntModel, h_s: &IntModel, ent: &CondEntropy, bs: &CondStream<'_>) -> Result<(CondPlan, Nets), Refused> {
let p = cond_plan(g_s, h_s, ent, bs)?;
Ok((p, load(g_s, h_s)?))
}
/// [`super::decode_conditional_checked`] through the integer operator set: the same
/// image, and the same refusals, except where integer-operators §15.3 lists a change.
pub fn decode(g_s: &IntModel, h_s: &IntModel, ent: &CondEntropy, bs: &CondStream<'_>) -> Result<(Vec<u8>, usize, usize), Refused> {
let (p, n) = plan(g_s, h_s, ent, bs)?;
let (rgb, h, w) = decode_loaded(&n, g_s.afrac, ent, bs)?;
debug_assert_eq!((3, h, w), p.gs.out, "the graphs give the image the plan planned");
Ok((rgb, h, w))
}
/// The decode after its plan, from the set's loaded networks ([`load`]): the
/// hyper-latent at `afrac`, then `h_s`'s graph, whose output the driver splits into
/// `σ` (channels `[0, m)`) and `μ` (`[m, 2m)`) (§13.5); then the latent, conditioned on
/// `σ`; then `g_s`'s graph, which ends with the RGB8 step; then the bytes, row-major
/// HWC. `bs` must be a payload 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(n: &Nets, afrac: u32, ent: &CondEntropy, bs: &CondStream<'_>) -> Result<(Vec<u8>, usize, usize), Refused> {
use wai_ops::Tensor;
let (m, sh, sw) = (bs.m, bs.sh, bs.sw);
// 1. The hyper-latent, then h_s. The plan admitted every size, so each fits u32.
let z = hyper_latent(ent, bs, afrac)?;
let x = Tensor::new(bs.n_z as u32, bs.zh as u32, bs.zw as u32, z);
let sm = crate::int_ops::run1(&n.h_s, x, &[]).map_err(|r| kind(r, "a decoded hyper-latent is outside a check range of h_s"))?;
if (sm.c as usize, sm.h as usize, sm.w as usize) != (2 * m, sh, sw) {
return Err(refused(RefusalKind::GeometryMismatch, "h_s does not give the latent shape the payload declares"));
}
// 2. The driver's split, then the latent conditioned on σ.
let (sigma, mu) = sm.data.split_at(m * sh * sw);
let y = latent(ent, bs, afrac, sigma, mu)?;
drop(sm);
// 3. g_s, ending with the RGB8 step, so every value is a byte; CHW → HWC.
let out = crate::int_ops::run1(&n.g_s, Tensor::new(m as u32, sh as u32, sw as u32, y), &[])
.map_err(|r| kind(r, "a decoded latent is outside a check range of g_s"))?;
if out.c != 3 {
return Err(refused(RefusalKind::ModelMismatch, "g_s does not give three channels"));
}
let (h, w) = (out.h as usize, out.w as usize);
let plane = h * w;
let mut rgb = vec![0u8; 3 * plane];
for (k, &v) in out.data.iter().enumerate() {
rgb[(k % plane) * 3 + k / plane] = v as u8;
}
Ok((rgb, h, w))
}
}
/// The fast kernels against the scalar reference, element by element. Operand
/// ranges cover the `i32` fast path, the `i64` path, and sparse (zero-heavy)
/// inputs; geometries cover strides, paddings, output paddings and kernels smaller

13
wai/wai-rs/src/int_kernels.rs
View file

@ -1113,6 +1113,9 @@ thread_local! {
/// a test tell which path a caller took, which the output alone cannot show
/// (both paths give the same integers). See [`testkit::fast_conv_runs`].
static FAST_CONV_RUNS: Cell<u64> = const { Cell::new(0) };
/// Test builds only: how many inverse GDNs the fast kernel ([`igdn`]) has computed on
/// this thread. See [`testkit::fast_igdn_runs`].
static FAST_IGDN_RUNS: Cell<u64> = const { Cell::new(0) };
}
/// Record that the fast convolution kernel runs (test builds only; a no-op otherwise).
@ -1225,6 +1228,8 @@ pub(crate) fn igdn(x: &[i64], c: usize, hw: usize, frac: u32, gamma: &[i64], bet
let narrow_out = xm <= i32::MAX as u128;
let narrow_mul = narrow_sum && gm <= i32::MAX as u128 && x2m <= i32::MAX as u128;
let threads = threads_for((c as u64) * (c as u64) * (hw as u64));
#[cfg(test)]
FAST_IGDN_RUNS.with(|k| k.set(k.get() + 1));
let mut out = vec![0i64; n];
let tiles = position_tiles(&mut out, hw, IGDN_TILE);
@ -1465,6 +1470,14 @@ pub(crate) mod testkit {
let r = f();
(r, super::FAST_CONV_RUNS.with(std::cell::Cell::get) - before)
}
/// `f`'s result, and how many inverse GDNs the fast kernel computed on this thread
/// while `f` ran.
pub fn fast_igdn_runs<R>(f: impl FnOnce() -> R) -> (R, u64) {
let before = super::FAST_IGDN_RUNS.with(std::cell::Cell::get);
let r = f();
(r, super::FAST_IGDN_RUNS.with(std::cell::Cell::get) - before)
}
}
#[cfg(test)]

234
wai/wai-rs/src/int_ops.rs
View file

@ -15,9 +15,10 @@
//! - under [`KernelPath::Fast`] a convolution and a transposed convolution run on the
//! fast kernels, which form the same exact sums (`int_kernels`' equivalence rule) and
//! finish each output with the operator set's own requantisation
//! ([`wai_ops::interp::rq`]). Where a fast kernel declines (operands past its narrowed
//! types, geometry it does not cover, a value that would leave `i64`), the interpreter
//! computes the operator, and refuses it where it refuses.
//! ([`wai_ops::interp::rq`]), and an inverse GDN runs on the fast IGDN kernel, whose
//! narrowed paths compute the same integers. Where a fast kernel declines (operands past
//! its narrowed types, geometry it does not cover, a value that would leave `i64`), the
//! interpreter computes the operator, and refuses it where it refuses.
//!
//! Either way the integers are the interpreter's: the tests below check it operator by
//! operator on random graphs, and `tests/int_synth_ops.rs` decoder by decoder on every
@ -34,7 +35,14 @@ use crate::int_kernels::{policy, ConvGeom, KernelPath};
/// The profile of every capability that decodes through the operator set
/// (`extensions/integer-operators.md` §11): `acc64`, which every lowering of §13
/// declares. A set may meet `acc32` as well; its certificate shows that.
const PROFILE: Profile = Profile::Acc64;
pub(crate) const PROFILE: Profile = Profile::Acc64;
/// The lowerings of §13 a decoder here lowers its set by, as the registry's `ops`
/// member names them (§11): the keys of loaded sets name them
/// ([`crate::int_set_cache::SetKey`]).
pub const WHM1: &str = "whm1/1";
pub const WAM1: &str = "wam1/1";
pub const INT_SYNTH: &str = "int_synth/1";
/// Load a lowered graph as a sink does: from its canonical bytes, under [`PROFILE`].
/// A refusal is `model_mismatch`, or `too_large` for a size (§10.5).
@ -49,24 +57,31 @@ pub fn load(g: &Graph) -> Result<Loaded, Refusal> {
/// 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`).
/// model parsed from bytes holds only 32-bit fields, weights that fit `i16` and `γ` that
/// fits `i32`, the widths `WHM1` gives them (`extensions/integer-payloads.md` §9) and its
/// lowering writes them at (§13.0). A model from another source (the tools' JSON form)
/// with a field past 32 bits, or a weight or `γ` past its width, is not a `WHM1` network,
/// and is refused here (`model_mismatch`), where its graph would otherwise hold other
/// values than the model.
pub fn whm1_net(m: &IntModel) -> Result<Net, Refusal> {
let f = |v: usize| u32::try_from(v).map_err(|_| Refusal::new(wai_ops::Code::ModelMismatch, wai_ops::Site::Lowering));
let refuse = || Refusal::new(wai_ops::Code::ModelMismatch, wai_ops::Site::Lowering);
let f = |v: usize| u32::try_from(v).map_err(|_| refuse());
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 i16s = |w: &[i64]| if w.iter().all(|&v| i16::try_from(v).is_ok()) { Ok(w.to_vec()) } else { Err(refuse()) };
let i32s = |g: &[i64]| if g.iter().all(|&v| i32::try_from(v).is_ok()) { Ok(g.to_vec()) } else { Err(refuse()) };
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(),
w: i16s(weight)?, 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(),
w: i16s(weight)?, mult: mult.clone(), bias: bias.clone(),
},
IntLayer::Igdn { gamma, beta } => NetLayer::Igdn { c: f(beta.len())?, gamma: gamma.clone(), beta: beta.clone() },
IntLayer::Igdn { gamma, beta } => NetLayer::Igdn { c: f(beta.len())?, gamma: i32s(gamma)?, beta: beta.clone() },
IntLayer::Relu => NetLayer::Relu,
IntLayer::Leaky { slope_q, slope_frac } => NetLayer::Leaky { slope_q: *slope_q, slope_frac: *slope_frac },
})
@ -94,6 +109,7 @@ impl Kernels for Host {
let (cin, h, w) = (x.c as usize, x.h as usize, x.w as usize);
let u = |v: &u32| *v as usize;
let g = match &op.p {
Params::Igdn { gf, frac, gamma, beta, .. } => return Some(vec![igdn(x, *gf, *frac, gamma, beta, (oc, oh, ow))?]),
Params::Conv { cout, kh, kw, sh, sw, ph, pw, groups, w: wt, .. } => {
ConvGeom::conv(cin, h, w, u(cout), u(groups), u(kh), u(kw), u(sh), u(sw), u(ph), u(pw), x.data.len(), wt.len())?
}
@ -110,6 +126,34 @@ impl Kernels for Host {
}
}
/// An inverse GDN on the fast kernel ([`crate::int_kernels::igdn`]), whose integers are the
/// interpreter's (§5.7: each term `γ·x²` floored by `2^gf` before it is summed, the
/// norm's integer square root, then `x·r` floored by `2^frac`) wherever no value leaves the
/// widths the interpreter holds it in. `None`, for the interpreter to compute the operator
/// (and refuse it), where one could, read from the values present as §13.2's `P_ℓ` reads
/// them: a term or the norm past `i128`, the root or the output past `i64`. No certified
/// graph forms one: the input's `check` range is that predicate's largest admissible
/// magnitude.
fn igdn(x: &Tensor, gf: u32, frac: u32, gamma: &[i64], beta: &[i64], shape: (u32, u32, u32)) -> Option<Tensor> {
use crate::int_kernels::max_abs;
const I64: u128 = i64::MAX as u128;
const I128: u128 = i128::MAX as u128;
let (c, hw) = (x.c as usize, x.h as usize * x.w as usize);
if shape != (x.c, x.h, x.w) || c.checked_mul(c) != Some(gamma.len()) || beta.len() != c || gf > 127 || frac > 127 {
return None;
}
let xm = max_abs(&x.data) as u128;
// x² ≤ 2^126 and x·root ≤ 2^126 fit u128.
let term = (max_abs(gamma) as u128).checked_mul(xm * xm).filter(|&t| t <= I128)?;
let norm = (max_abs(beta) as u128).checked_add((c as u128).checked_mul((term >> gf) + 1)?).filter(|&n| n <= I128)?;
let root = norm.isqrt();
if root > I64 || ((xm * root) >> frac) + 1 > I64 {
return None;
}
let d = crate::int_kernels::igdn(&x.data, c, hw, frac, gamma, beta, gf)?;
Some(Tensor::new(x.c, x.h, x.w, d))
}
/// A (transposed) convolution on the fast kernel, each output finished with the
/// operator set's requantisation. `None` where the kernel declines, or where a
/// requantised value would leave `i64`, which no certified graph forms: the interpreter
@ -127,6 +171,41 @@ fn conv(g: &ConvGeom, x: &[i64], wt: &[i64], rq: &wai_ops::graph::Rq) -> Option<
(!wide.load(Ordering::Relaxed)).then_some(out)
}
/// [`whm1_net`] at the edges of the widths `WHM1` gives a layer's weights (`i16`) and an
/// inverse GDN's `γ` (`i32`): a model at either end converts, and its lowered graph, as a
/// sink loads it from its bytes, holds the model's value; one past either end, which only
/// a model from another source than `WHM1` bytes can hold, is refused (`model_mismatch`).
#[cfg(test)]
mod conversion {
use super::*;
use wai_ops::graph::Params;
fn conv(w: i64) -> IntModel {
IntModel { afrac: 12, gf: 16, shift: 8, layers: vec![IntLayer::ConvT { out_ch: 1, k: 1, stride: 1, pad: 0, opad: 0, weight: vec![w], mult: vec![1], bias: vec![0] }] }
}
fn igdn(g: i64) -> IntModel {
IntModel { afrac: 12, gf: 16, shift: 8, layers: vec![IntLayer::Igdn { gamma: vec![g], beta: vec![1 << 24] }] }
}
/// The first operator's parameters, in the graph a sink parses from the lowering's bytes.
fn written(m: &IntModel) -> Params {
let g = wai_ops::lower::whm1_net(&whm1_net(m).unwrap(), false).unwrap();
Graph::parse(&g.to_bytes()).unwrap().ops[0].p.clone()
}
#[test]
fn a_weight_or_gamma_past_whm1s_widths_is_refused() {
for w in [i16::MIN as i64, i16::MAX as i64] {
assert!(matches!(written(&conv(w)), Params::ConvT { w: ref v, .. } if *v == [w]), "weight {w}");
}
for g in [i32::MIN as i64, i32::MAX as i64] {
assert!(matches!(written(&igdn(g)), Params::Igdn { gamma: ref v, .. } if *v == [g]), "gamma {g}");
}
for m in [conv(i16::MIN as i64 - 1), conv(i16::MAX as i64 + 1), igdn(i32::MIN as i64 - 1), igdn(i32::MAX as i64 + 1)] {
assert_eq!(whm1_net(&m).unwrap_err().code, wai_ops::Code::ModelMismatch);
}
}
}
/// Gate 5 of each decoder's migration, the per-operator half: the fast kernels behind
/// the interpreter against the interpreter alone, on random graphs of one convolution or
/// transposed convolution. Every output element must be the same integer under every
@ -252,12 +331,12 @@ mod equivalence {
}
/// 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.
/// every fast policy each of its two transposed convolutions and its inverse GDN run on
/// the fast kernels, and every attribute byte is the interpreter's.
#[cfg(feature = "splat_int")]
#[test]
fn the_splat_synthesis_runs_its_convolutions_on_the_fast_kernel() {
use crate::int_kernels::testkit::fast_igdn_runs;
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();
@ -267,22 +346,23 @@ mod equivalence {
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");
let ((want, none), no_igdn) = fast_igdn_runs(|| fast_conv_runs(|| with_policy(Policy::reference(), || ops::decode_typed(&gs, &st).unwrap().attr)));
assert_eq!((none, no_igdn), (0, 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));
let ((got, runs), igdns) = fast_igdn_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");
assert_eq!((runs, igdns), (2, 1), "{p:?}: both transposed convolutions and the inverse GDN run on the fast kernels");
}
}
/// 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.
/// every fast policy each of its four one-dimensional transposed convolutions and each of
/// its inverse GDNs run on the fast kernels, and every sample is the interpreter's. Its
/// PCM step runs 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_kernels::testkit::fast_igdn_runs;
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();
@ -292,15 +372,121 @@ mod equivalence {
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");
let igdns = ops::load(&m).unwrap().ops.iter().filter(|o| o.code == wai_ops::graph::opcode::IGDN).count() as u64;
assert!(igdns > 0, "the committed audio synthesis has an inverse GDN");
let ((want, none), no_igdn) = fast_igdn_runs(|| fast_conv_runs(|| with_policy(Policy::reference(), || ops::decode_typed(&m, &s).unwrap())));
assert_eq!((none, no_igdn), (0, 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()));
let ((got, runs), ig) = fast_igdn_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");
assert_eq!((runs, ig), (4, igdns), "{p:?}: every transposed convolution and inverse GDN runs on the fast kernels");
}
}
/// The committed conditional codec (`wai.neural.int_hyper`'s `tiny`, also the keyframe of
/// `wai.video.int_hyper`'s vector) through its lowering `whm1/1`: under every fast policy
/// its four convolutions (`h_s`'s transposed convolution and convolution, `g_s`'s two
/// transposed convolutions) and `g_s`'s inverse GDN run on the fast kernels, and every
/// byte is the interpreter's. `h_s`'s leaky ReLU and the RGB8 step run on the
/// interpreter under every policy.
#[test]
fn the_conditional_codec_runs_its_networks_on_the_fast_kernels() {
use crate::int_hyper_synth::{ops, CondEntropy, CondStream};
use crate::int_kernels::testkit::fast_igdn_runs;
let dir = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../wai-rs/tests/vectors/int_hyper");
let read = |f: &str| std::fs::read(dir.join(f)).unwrap();
let (g, h) = (IntModel::from_bytes(&read("tiny_g_s.bin")).unwrap(), IntModel::from_bytes(&read("tiny_h_s.bin")).unwrap());
let ent = CondEntropy::try_from_value(&serde_json::from_slice(&read("tiny_entropy.json")).unwrap()).unwrap();
let p = read("tiny_bitstream.wih1");
let k = crate::int_wire::parse_wih1(&p).unwrap();
let st = CondStream {
m: k.m as usize, sh: k.sh as usize, sw: k.sw as usize, n_z: k.n_z as usize, zh: k.zh as usize, zw: k.zw as usize,
z_rans: k.z_rans, rans: k.rans,
};
let decode = |pol: Policy| fast_igdn_runs(|| fast_conv_runs(|| with_policy(pol, || ops::decode(&g, &h, &ent, &st).unwrap())));
let ((want, none), no_igdn) = decode(Policy::reference());
assert_eq!((none, no_igdn), (0, 0), "the reference policy runs no fast kernel");
for pol in FAST {
let ((got, convs), igdns) = decode(pol);
assert!(got == want, "{pol:?}: the image differs from the interpreter alone");
assert_eq!((convs, igdns), (4, 1), "{pol:?}: every convolution and the inverse GDN run on the fast kernels");
}
}
/// The fast IGDN kernel behind the interpreter, against the interpreter alone, on random
/// inverse GDNs: 1 to 6 channels, `gf` and `frac` across `0..=127` (the fast kernel takes
/// those up to 62, the interpreter the rest), `γ` and `β` of either sign, and inputs whose
/// squares fit 32 bits, 64 bits, and neither: the kernel's narrow, 64-bit and 128-bit
/// paths. Every output element is the same integer under every fast policy. This is
/// `int_transform`'s equivalence of its fast IGDN with its scalar one, held to the
/// interpreter instead.
#[test]
fn inverse_gdns_equal_the_interpreter() {
use crate::int_kernels::testkit::fast_igdn_runs;
let mut r = Rng(0x0D2C_0004);
let (mut loaded, mut fast) = (0, 0);
for case in 0..400 {
let c = 1 + r.below(6) as u32;
let (h, w) = (1 + r.below(6) as u32, 1 + r.below(9) as u32);
let gf = [0u32, 8, 16, 24, 40, 62, 63, 100][r.below(8)];
let frac = [1u32, 12, 16, 30, 62, 63, 90][r.below(7)];
let (xmax, gmax, bmax) = [(1i64 << 12, 1i64 << 12, 1i64 << 24), (1 << 20, 1 << 16, 1 << 40), (1 << 31, 1 << 20, 1 << 50), (1 << 40, 1 << 30, 1 << 60)][case % 4];
let p = Params::Igdn { gf, frac, g_bytes: 4, gamma: r.vec((c * c) as usize, gmax, 10), beta: r.vec(c as usize, bmax, 0) };
let g = graph(TensorRec::check(c, -xmax, xmax), opcode::IGDN, c, p);
let x = Tensor::new(c, h, w, r.vec((c * h * w) as usize, xmax, 20));
let Ok(l) = load(&g) else { continue };
loaded += 1;
let want = with_policy(Policy::reference(), || run1(&l, x.clone(), &[]));
for pol in FAST {
let (got, runs) = fast_igdn_runs(|| with_policy(pol, || run1(&l, x.clone(), &[])));
assert_eq!(got, want, "{pol:?}: c={c} gf={gf} frac={frac} |x|<={xmax}");
fast += runs;
}
}
assert!(loaded > 200, "only {loaded} of 400 random inverse GDNs loaded");
assert!(fast > 400, "the fast IGDN kernel ran only {fast} times");
}
/// `int_hyper_synth`'s equivalence of its fast convolution kernels with its scalar ones,
/// held to the interpreter instead: random convolutions and transposed convolutions of
/// `whm1/1`'s shape (square kernels to 5, strides to 16 one case in three and to 3
/// otherwise, padding below the kernel, output padding below the stride) with `WHM1`'s
/// requantisation (a multiplier per channel, one shift, a `post` bias, `floor`), on
/// operands that fit `i32`, that do not, and mixed, some of them mostly zero.
#[test]
fn whm1_convolutions_equal_the_interpreter() {
const RANGES: [(i64, i64); 3] = [(1 << 14, 2047), (1 << 40, 1 << 10), (1 << 20, 1 << 20)];
let mut r = Rng(0x0D2C_0005);
let (mut loaded, mut fast) = (0, 0);
for case in 0..600 {
let k = 1 + r.below(5) as u32;
let s = 1 + if case % 3 == 0 { r.below(16) } else { r.below(3) } as u32;
let pad = r.below(k as usize) as u32;
let (h, w) = (1 + r.below(7) as u32, 1 + r.below(9) as u32);
let (cin, cout) = (1 + r.below(6) as u32, 1 + r.below(9) as u32);
let (xmax, wmax) = RANGES[case % 3];
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 << 20).abs()).collect(), shift: vec![16 + r.below(8) as u8], bias: r.vec(n, 1 << 20, 0),
};
let wt = r.vec((cin * cout * k * k) as usize, wmax, 10);
let (code, p) = if case % 2 == 1 {
let op = r.below(s as usize) as u32;
(opcode::CONV_TRANSPOSE, Params::ConvT { cout, kh: k, kw: k, sh: s, sw: s, ph: pad, pw: pad, oph: op, opw: op, w_bytes: 8, rq, w: wt })
} else {
(opcode::CONV, Params::Conv { cout, kh: k, kw: k, sh: s, sw: s, ph: pad, pw: pad, groups: 1, w_bytes: 8, rq, w: wt })
};
let g = graph(TensorRec::check(cin, -xmax, xmax), code, cout, p);
let x = Tensor::new(cin, h, w, r.vec((cin * h * w) as usize, xmax, [0, 50, 90][case % 3]));
let (l, f) = agree(&g, &x);
loaded += l as u32;
fast += f;
}
assert!(loaded > 400, "only {loaded} of 600 random graphs loaded");
assert!(fast > 800, "the fast kernel ran only {fast} times");
}
/// A value outside an input's `check` range is refused by the interpreter before any
/// kernel runs, under every policy alike.
#[test]

173
wai/wai-rs/src/int_payload/mod.rs
View file

@ -40,6 +40,7 @@ use crate::codecs::{
};
use crate::container::{PinRefusal, Prior};
use crate::pin::{HeldBytes, PinAuthority, Verified, VerifiedRef};
use std::sync::Arc;
/// Why [`decode_int`] or a converter refused: a code of the extension's closed list,
/// and what failed.
@ -127,7 +128,14 @@ pub fn set_members(cap: &str) -> &'static [&'static str] {
/// that no pin verified while reporting it pinned, nor an unpinned set other than the
/// one check 5 identified. [`IntParams::none`] holds none.
#[derive(Debug, Clone)]
pub struct IntParams<'a>(Params<'a>);
pub struct IntParams<'a>(Params<'a>, Option<ReadDigest>);
/// The SHA-256 of the listing of the members a decode reads
/// (`int_set_cache::SetKey::listing_digest`), with those members' names (`set_members`
/// of the capability it was formed for): kept where [`resolve`] formed it for check 5,
/// so that the key of the loaded set ([`loaded`]) does not hash the same bytes again.
/// The bytes are borrowed for as long as the parameters live, so they cannot change.
type ReadDigest = (&'static [&'static str], [u8; 32]);
#[derive(Debug, Clone)]
enum Params<'a> {
@ -145,14 +153,14 @@ impl<'a> IntParams<'a> {
/// No parameter set. A decode through a capability that reads one refuses with
/// `set_incomplete`.
pub fn none() -> IntParams<'static> {
IntParams(Params::None)
IntParams(Params::None, None)
}
/// The set `pin::verify` returned for a pin over `cap`: the neural sink's pinned
/// decode, whose registry ran SPEC §3.1's checks.
#[cfg(feature = "neural")]
pub(crate) fn verified(v: &'a Verified) -> Self {
IntParams(Params::Verified(v))
IntParams(Params::Verified(v), None)
}
/// The parameters a pin verified with `pin::verify_bytes` over `held`.
@ -162,7 +170,18 @@ impl<'a> IntParams<'a> {
VerifiedRef::Set { index, paths } => (*index, Some(paths.clone())),
};
let set = *held.get(index).ok_or_else(|| PinRefusal::NotHeld { capability: cap.to_string() })?;
Ok(IntParams(Params::Held { set, paths, pinned: true }))
Ok(IntParams(Params::Held { set, paths, pinned: true }, None))
}
/// The SHA-256 of the listing of the members `cap`'s decode reads from these
/// parameters: the one [`resolve`] formed where it formed one over those same members,
/// or else formed here from the bytes. The digest of the key of the loaded set
/// ([`loaded`]).
pub(crate) fn read_digest(&self, cap: &str) -> Option<[u8; 32]> {
match self.1 {
Some((members, digest)) if members == set_members(cap) => Some(digest),
_ => crate::int_set_cache::SetKey::listing_digest(&self.read(cap)),
}
}
/// Whether a pin verified these parameters.
@ -215,6 +234,12 @@ impl<'a> IntParams<'a> {
/// A set whose listing is malformed or does not match its files has no one identity,
/// so it is refused before check 5 runs.
pub(crate) fn unpinned_read(cap: &str, set: &HeldBytes<'_>) -> Result<(Vec<String>, Option<Vec<String>>), PinRefusal> {
unpinned_read_digest(cap, set).map(|(digests, paths, _)| (digests, paths))
}
/// [`unpinned_read`], and the SHA-256 of the listing of the members read, which is one of
/// the digests (`IntParams::read_digest`).
fn unpinned_read_digest(cap: &str, set: &HeldBytes<'_>) -> Result<(Vec<String>, Option<Vec<String>>, Option<[u8; 32]>), PinRefusal> {
use crate::pin::{parse_listing, sha256, sha256_hex};
let one_file = crate::container::pin_shape(cap) == crate::codecs::ParameterSetShape::OneFile;
let members = set_members(cap);
@ -250,10 +275,11 @@ pub(crate) fn unpinned_read(cap: &str, set: &HeldBytes<'_>) -> Result<(Vec<Strin
Some(read.keys().cloned().collect())
}
};
digests.extend(read_identities(cap, &read));
let (ids, listing) = identities_and_listing(cap, &read);
digests.extend(ids);
digests.sort_unstable();
digests.dedup();
Ok((digests, paths))
Ok((digests, paths, listing))
}
/// Every digest a pin over the bytes a decode of `cap` reads, `read` (member name →
@ -263,17 +289,23 @@ pub(crate) fn unpinned_read(cap: &str, set: &HeldBytes<'_>) -> Result<(Vec<Strin
/// pin names). Check 5 runs over these, pinned or unpinned, so a set is refused under
/// any identity of the bytes the decode reads (SPEC §3.1 check 5).
pub fn read_identities(cap: &str, read: &std::collections::BTreeMap<String, &[u8]>) -> Vec<String> {
use crate::pin::{listing_from, sha256_hex};
identities_and_listing(cap, read).0
}
/// [`read_identities`], and the SHA-256 of the listing of `read` among them.
fn identities_and_listing(cap: &str, read: &std::collections::BTreeMap<String, &[u8]>) -> (Vec<String>, Option<[u8; 32]>) {
use crate::pin::{listing_from, sha256, sha256_hex};
let mut out = Vec::new();
if crate::container::pin_shape(cap) == crate::codecs::ParameterSetShape::OneFile {
if let Some(b) = set_members(cap).first().and_then(|m| read.get(*m)) {
out.push(sha256_hex(b));
}
}
if let Ok(l) = listing_from(read) {
let listing = listing_from(read).ok().map(|l| {
out.push(sha256_hex(&l));
}
out
sha256(&l)
});
(out, listing)
}
/// The parameters `cap`'s decode reads, at a sink that holds the sets `held` in memory
@ -299,30 +331,31 @@ pub fn resolve<'a>(
authority: Option<&dyn PinAuthority>,
) -> Result<IntParams<'a>, PinRefusal> {
if crate::codecs::is_known_capability(cap) && !crate::codecs::decodes_against_parameter_set(cap) {
return Ok(IntParams(Params::None));
return Ok(IntParams(Params::None, None));
}
match prior {
Some(p) if crate::container::pin_applies_to(cap) => {
let r = crate::pin::verify_bytes_with(cap, p, held, crate::pin::READABLE_FORMS, authority)?;
let params = IntParams::from_verified_ref(cap, held, &r)?;
let mut params = IntParams::from_verified_ref(cap, held, &r)?;
// Check 5 ran over the pin's digest; it runs again over every identity of the
// bytes the decode will read, so a set revoked (or not granted) under its
// file's digest or under the listing of its members is not reached through
// a pin that names it another way.
if let Some(a) = authority {
let mut ids = read_identities(cap, &params.read(cap));
let (mut ids, listing) = identities_and_listing(cap, &params.read(cap));
ids.push(p.sha256.clone());
a.check_unpinned(cap, &ids)?;
params.1 = listing.map(|d| (set_members(cap), d));
}
Ok(params)
}
_ => {
let set = *held.first().ok_or_else(|| PinRefusal::NotHeld { capability: cap.to_string() })?;
let (digests, paths) = unpinned_read(cap, &set)?;
let (digests, paths, listing) = unpinned_read_digest(cap, &set)?;
if let Some(a) = authority {
a.check_unpinned(cap, &digests)?;
}
Ok(IntParams(Params::Held { set, paths, pinned: false }))
Ok(IntParams(Params::Held { set, paths, pinned: false }, listing.map(|d| (set_members(cap), d))))
}
}
}
@ -365,11 +398,17 @@ pub fn decode_int(cap: &str, payload: &[u8], params: &IntParams<'_>) -> Result<I
Some(wire::MAGIC_WIS1) => decode_wis1(cap, payload, params),
_ => match cap {
CAP_NEURAL_INT_SYNTH => {
use crate::int_synth::{check_payload, decode_loaded, load_spec, IntSynthError};
let spec = params.member(cap, "int_synth.json")?;
let img = crate::int_synth::decode_from_spec(payload, spec).map_err(|e| {
let crate::int_synth::IntSynthError::InvalidPayload(why) = e;
let synth = |e: IntSynthError| {
let IntSynthError::InvalidPayload(why) = e;
IntPayloadError::new(Refusal::PayloadMalformed, why)
})?;
};
// `int_synth::decode_from_spec`'s three steps, its set from the store when
// it holds it.
check_payload(payload).map_err(synth)?;
let g = loaded(cap, crate::int_ops::INT_SYNTH, params, || load_spec(spec)).map_err(synth)?;
let img = decode_loaded(payload, &g).map_err(synth)?;
Ok(IntDecoded::Image { rgb: img.rgb, width: img.width, height: img.height })
}
CAP_NEURAL_INT_MLICPP => {
@ -478,7 +517,7 @@ pub fn declared_cost(cap: &str, payload: &[u8], params: &IntParams<'_>) -> Resul
Some(wire::MAGIC_WIV1) => cost_wiv1(cap, payload, params),
Some(wire::MAGIC_WIH1) => {
let k = wire::parse_wih1(payload)?;
let p = cond_cost(cap, &k, params)?;
let (_, p, _) = cond_checked(cap, &k, params, None)?;
Ok(DeclaredCost { output_bytes: p.output_bytes(), symbols: p.z + p.y, working_bytes: p.working_bytes(), macs: p.macs() })
}
Some(wire::MAGIC_WIA1) => cost_wia1(cap, payload, params),
@ -491,15 +530,58 @@ pub fn declared_cost(cap: &str, payload: &[u8], params: &IntParams<'_>) -> Resul
#[cfg(any(feature = "audio_int", feature = "splat_int"))]
const VEC: u64 = std::mem::size_of::<Vec<i32>>() as u64;
/// The conditional codec's plan for a parsed `WIH1` against the capability's set.
fn cond_cost(cap: &str, k: &wire::Wih1<'_>, params: &IntParams<'_>) -> Result<crate::int_hyper_synth::CondPlan, IntPayloadError> {
/// The conditional codec's set and plan for a parsed `WIH1`: everything its decode checks
/// before the first symbol, in integer-payloads §7's order. The set's members, parsed;
/// the plan (`int_hyper_synth::cond_plan`); for a `WIV1` keyframe, its image's size
/// against the clip's `h × w` (`size`, `geometry_mismatch`); then `h_s` and `g_s`
/// lowered and loaded (`int_hyper_synth::ops::load`), or found loaded ([`loaded`]). The
/// decode and the declared cost both run it, so a payload the decode refuses before its
/// first symbol is refused by both with the same code.
fn cond_checked(
cap: &str, k: &wire::Wih1<'_>, params: &IntParams<'_>, size: Option<(u32, u32)>,
) -> Result<(CondSet, crate::int_hyper_synth::CondPlan, Arc<crate::int_hyper_synth::ops::Nets>), IntPayloadError> {
let (g_s, h_s, ent) = int_hyper_models(cap, params)?;
let s = crate::int_hyper_synth::CondStream {
let p = crate::int_hyper_synth::cond_plan(&g_s, &h_s, &ent, &cond_stream(k))?;
if let Some((h, w)) = size
&& p.gs.out != (3, h as usize, w as usize)
{
return Err(IntPayloadError::new(Refusal::GeometryMismatch, "the keyframe's size is not the clip's h × w"));
}
let nets = loaded(cap, crate::int_ops::WHM1, params, || crate::int_hyper_synth::ops::load(&g_s, &h_s))?;
Ok(((g_s.afrac, ent), p, nets))
}
/// The set `cap`'s decode reads from `params`, lowered by `lowering` and loaded under
/// the capability's profile by `load`: from the process's store of loaded sets
/// (`int_set_cache`) when it holds these bytes, and kept there when it loads. The key is
/// the SHA-256 of the members read, with the capability, the lowering and the profile
/// (`int_set_cache::SetKey::of`). Only the load step is answered from the store: every
/// check before it has run, and a set the load refuses is not kept, so it is refused
/// again, with the same refusal, on its next decode.
fn loaded<T: crate::int_set_cache::Resident + Send + Sync + 'static, E>(
cap: &str, lowering: &'static str, params: &IntParams<'_>, load: impl FnOnce() -> Result<T, E>,
) -> Result<Arc<T>, E> {
use crate::int_set_cache::{LoadedSets, SetKey};
// The key names the capability by the constant this crate registers it under.
let key = || {
let cap = [CAP_NEURAL_INT_HYPER, CAP_VIDEO_INT_HYPER, CAP_AUDIO_INT_CODEC, CAP_SPLAT_INT_CODEC, CAP_NEURAL_INT_SYNTH]
.into_iter().find(|c| *c == cap)?;
Some(SetKey { capability: cap, lowering, profile: crate::int_ops::PROFILE, digest: params.read_digest(cap)? })
};
LoadedSets::global().get_or_load(key, load)
}
/// What the conditional decode keeps of its set once the networks are loaded: the
/// fixed-point scale and the entropy tables.
type CondSet = (u32, crate::int_hyper_synth::CondEntropy);
/// A parsed `WIH1`'s typed form.
fn cond_stream<'a>(k: &wire::Wih1<'a>) -> crate::int_hyper_synth::CondStream<'a> {
crate::int_hyper_synth::CondStream {
m: k.m as usize, sh: k.sh as usize, sw: k.sw as usize,
n_z: k.n_z as usize, zh: k.zh as usize, zw: k.zw as usize,
z_rans: k.z_rans, rans: k.rans,
};
Ok(crate::int_hyper_synth::cond_plan(&g_s, &h_s, &ent, &s)?)
}
}
#[cfg(feature = "video_int")]
@ -519,10 +601,7 @@ fn cost_wiv1(cap: &str, payload: &[u8], params: &IntParams<'_>) -> Result<Declar
match v.wih1 {
None => Ok(DeclaredCost { output_bytes, symbols: output_bytes, working_bytes: v.keyframe0_working_bytes(), macs: 0 }),
Some(nested) => {
let p = cond_cost(cap, &wire::parse_wih1(nested)?, params)?;
if p.gs.out != (3, v.h as usize, v.w as usize) {
return Err(IntPayloadError::new(Refusal::GeometryMismatch, "the keyframe's size is not the clip's h × w"));
}
let (_, p, _) = cond_checked(cap, &wire::parse_wih1(nested)?, params, Some((v.h, v.w)))?;
// The keyframe's decode, then the clip with the keyframe's image beside it.
let working = p.working_bytes().max(working_bytes(&s) + frame);
Ok(DeclaredCost { output_bytes, symbols: frame * (v.n_frames as u64 - 1) + p.z + p.y, working_bytes: tables + working, macs: p.macs() })
@ -546,7 +625,7 @@ fn cost_wia1(cap: &str, payload: &[u8], params: &IntParams<'_>) -> Result<Declar
};
// 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)?;
let (_, working, macs) = crate::int_audio::ops::plan_with(&m, &s, |m| loaded(cap, crate::int_ops::WAM1, params, || crate::int_audio::ops::load(m)))?;
// 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 })
@ -569,7 +648,7 @@ fn cost_wis1(cap: &str, payload: &[u8], params: &IntParams<'_>) -> Result<Declar
};
// 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)?;
let (_, working, macs) = crate::int_splat::ops::plan_with(&gs, &s, |m| loaded(cap, crate::int_ops::WHM1, params, || crate::int_splat::ops::load(m)))?;
// 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;
@ -602,23 +681,13 @@ fn int_hyper_models(
Ok((g_s, h_s, ent))
}
/// A parsed `WIH1` decoded to RGB with the capability's set. A `WIV1` keyframe passes
/// the clip's `h × w` as `size`: the keyframe's planned image is checked against it
/// (`geometry_mismatch`) before any of its symbols is decoded (§7).
/// A parsed `WIH1` decoded to RGB with the capability's set, through the integer operator
/// set (`int_hyper_synth::ops`). A `WIV1` keyframe passes the clip's `h × w` as `size`: the
/// keyframe's planned image is checked against it (`geometry_mismatch`) before its
/// networks are lowered and before any of its symbols is decoded (§7).
fn wih1_image(cap: &str, k: &wire::Wih1<'_>, params: &IntParams<'_>, size: Option<(u32, u32)>) -> Result<(Vec<u8>, u32, u32), IntPayloadError> {
let (g_s, h_s, ent) = int_hyper_models(cap, params)?;
let s = crate::int_hyper_synth::CondStream {
m: k.m as usize, sh: k.sh as usize, sw: k.sw as usize,
n_z: k.n_z as usize, zh: k.zh as usize, zw: k.zw as usize,
z_rans: k.z_rans, rans: k.rans,
};
if let Some((h, w)) = size {
let p = crate::int_hyper_synth::cond_plan(&g_s, &h_s, &ent, &s)?;
if p.gs.out != (3, h as usize, w as usize) {
return Err(IntPayloadError::new(Refusal::GeometryMismatch, "the keyframe's size is not the clip's h × w"));
}
}
let (rgb, h, w) = crate::int_hyper_synth::decode_conditional_checked(&g_s, &h_s, &ent, &s)?;
let ((afrac, ent), _, nets) = cond_checked(cap, k, params, size)?;
let (rgb, h, w) = crate::int_hyper_synth::ops::decode_loaded(&nets, afrac, &ent, &cond_stream(k))?;
Ok((rgb, h as u32, w as u32))
}
@ -663,14 +732,16 @@ fn decode_wiv1(cap: &str, _: &[u8], _: &IntParams<'_>) -> Result<IntDecoded, Int
#[cfg(feature = "audio_int")]
fn decode_wia1(cap: &str, payload: &[u8], params: &IntParams<'_>) -> Result<IntDecoded, IntPayloadError> {
use crate::int_audio::{decode_typed, AudioModel, AudioStream};
use crate::int_audio::{ops, 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 samples = decode_typed(&m, &s)?;
// `int_audio::decode_typed`, with its set from the store when it holds it.
let (g, ..) = ops::plan_with(&m, &s, |m| loaded(cap, crate::int_ops::WAM1, params, || ops::load(m)))?;
let samples = ops::decode_loaded(&g, m.afrac, &s)?;
Ok(IntDecoded::Pcm { samples, sample_rate: a.sr, channels: 1 })
}
@ -682,14 +753,16 @@ fn decode_wia1(cap: &str, _: &[u8], _: &IntParams<'_>) -> Result<IntDecoded, Int
#[cfg(feature = "splat_int")]
fn decode_wis1(cap: &str, payload: &[u8], params: &IntParams<'_>) -> Result<IntDecoded, IntPayloadError> {
use crate::int_hyper_synth::IntModel;
use crate::int_splat::{decode_typed, SplatStream};
use crate::int_splat::{ops, 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 d = decode_typed(&gs, &s)?;
// `int_splat::decode_typed`, with its set from the store when it holds it.
let (g, ..) = ops::plan_with(&gs, &s, |m| loaded(cap, crate::int_ops::WHM1, params, || ops::load(m)))?;
let d = ops::decode_loaded(&g, gs.afrac, &s)?;
Ok(IntDecoded::Attr { bytes: d.attr, n_attr: v.attr, s: v.s, n_gaussians: v.n_gaussians, ranges: v.ranges })
}

26
wai/wai-rs/src/int_payload/tests.rs
View file

@ -111,11 +111,11 @@ fn wiv1_json_equivalence() {
other => panic!("video_nk: {other:?}"),
});
let legacy = {
use crate::int_hyper_synth::{decode_conditional, CondBitstream, CondEntropy, IntModel};
use crate::int_hyper_synth::{decode_conditional_checked, CondBitstream, CondEntropy, IntModel};
let g_s = IntModel::from_json(&jload("int_hyper/tiny_g_s.json"));
let h_s = IntModel::from_json(&jload("int_hyper/tiny_h_s.json"));
let ent = CondEntropy::from_value(&jload("int_hyper/tiny_entropy.json"));
let (f0, h, w) = decode_conditional(&g_s, &h_s, &ent, &CondBitstream::from_value(&jload("int_hyper/tiny_bitstream.json"))).unwrap();
let (f0, h, w) = decode_conditional_checked(&g_s, &h_s, &ent, &CondBitstream::from_value(&jload("int_hyper/tiny_bitstream.json")).stream()).unwrap();
decode_pframes(&f0, h, w, &js).unwrap().frames
};
assert_eq!(frames, legacy, "video_nk: the binary and JSON forms decode alike");
@ -131,7 +131,7 @@ fn wiv1_json_equivalence() {
/// §6.1 test 17, `WIH1`.
#[test]
fn wih1_json_equivalence() {
use crate::int_hyper_synth::{decode_conditional, fnv1a, CondBitstream, CondEntropy, IntModel};
use crate::int_hyper_synth::{decode_conditional_checked, fnv1a, CondBitstream, CondEntropy, IntModel};
let js = jload("int_hyper/tiny_bitstream.json");
let payload = read("int_hyper/tiny_bitstream.wih1");
assert_eq!(wih1_from_json(&js).unwrap(), payload, "the converter writes the committed payload");
@ -149,7 +149,7 @@ fn wih1_json_equivalence() {
let g_s = IntModel::from_json(&jload("int_hyper/tiny_g_s.json"));
let h_s = IntModel::from_json(&jload("int_hyper/tiny_h_s.json"));
let ent = CondEntropy::from_value(&jload("int_hyper/tiny_entropy.json"));
assert_eq!(rgb, decode_conditional(&g_s, &h_s, &ent, &b).unwrap().0, "the payload and the unchecked JSON decode agree");
assert_eq!(rgb, decode_conditional_checked(&g_s, &h_s, &ent, &b.stream()).unwrap().0, "the payload and the JSON form decode alike");
assert_eq!(sha(&rgb), jload("int_hyper/sha256.json")["tiny"]["sha256"].as_str().unwrap());
assert_eq!(format!("{:016x}", fnv1a(&rgb)), jload("int_hyper/tiny.json")["rgb_hash"].as_str().unwrap());
// The checked entropy reader accepts the committed tables and gives what the
@ -826,3 +826,21 @@ fn encode_wiv1_zero_motion_round_trips() {
let clip = vec![Vec::new(); 17];
assert_eq!(code(encode_wiv1_zero_motion(&clip, 4096, 5461, 4)), "too_large", "17 frames of a frame within the tensor cap pass the clip cap");
}
/// The digest `resolve` formed for check 5 (the SHA-256 of the listing of the members
/// read) keys the loaded set only for the members it was formed over: a decode through a
/// capability that reads other members hashes those.
#[test]
fn a_resolved_digest_keys_only_the_members_it_was_formed_over() {
use crate::int_set_cache::SetKey;
let set = int_hyper_set();
let v = view(&set);
let p = unpinned(CAP_NEURAL_INT_HYPER, HeldBytes::Set { listing: b"", files: &v });
let listing = |names: &[&str]| SetKey::listing_digest(&names.iter().map(|n| (n.to_string(), v[*n])).collect()).unwrap();
let three = listing(&["entropy.json", "g_s.whm", "h_s.whm"]);
assert_eq!(p.1, Some((set_members(CAP_NEURAL_INT_HYPER), three)), "resolve keeps the digest it formed");
assert_eq!(p.read_digest(CAP_NEURAL_INT_HYPER), Some(three));
assert_eq!(p.read_digest(CAP_VIDEO_INT_HYPER), Some(three), "the same three members");
assert_eq!(p.read_digest(CAP_SPLAT_INT_CODEC), Some(listing(&["g_s.whm"])), "another capability's one member is hashed");
assert_eq!(IntParams::none().read_digest(CAP_NEURAL_INT_HYPER), SetKey::listing_digest(&BTreeMap::new()));
}

391
wai/wai-rs/src/int_set_cache.rs Normal file
View file

@ -0,0 +1,391 @@
//! Loaded parameter sets, kept between decodes.
//!
//! A decode through the integer operator set lowers its parameter set to `WHG1` graphs
//! and loads each one (`extensions/integer-operators.md` §10.5), the bounds pass
//! included (load step 8, §8). For a large set that costs more than the decode itself,
//! and it is the same work for every payload the set decodes. A sink may keep the
//! verdict of a set it has loaded, by the set's digest (§8.6). This module keeps the
//! loaded graphs themselves, for every integer decode this crate runs from a set's
//! bytes through `int_payload::decode_int` and `int_payload::declared_cost`, so also for
//! the neural sink and the browser crate, which decode through them.
//! (`int_synth::decode_from_spec`, which the byte-equality crate shares, loads its set
//! every time.)
//!
//! - **Keyed by the set's bytes, never by a name.** An entry is found by its [`SetKey`]:
//! the capability, the lowering, the profile the graphs were loaded under, and the
//! SHA-256 of the listing of the members the decode reads, each listed by its own
//! SHA-256 (`pin::listing_from`). That is the digest a pin in the form
//! `wai.prior.set` over exactly those members names (SPEC §3.1). Two sets whose
//! members have the same names and whose bytes differ in one byte have two keys.
//! - **Only what loaded.** An entry is made from what a load returned, after it
//! returned. A [`wai_ops::Loaded`] graph is made only by `wai_ops::graph::load`, after
//! every load step, so nothing reaches the store that did not pass the bounds pass. A
//! refusal is not kept: a set that is refused is lowered and loaded again on its next
//! decode, and refused again with the same refusal.
//! - **Nothing else moves.** A decode makes every check before the load as it did, in
//! integer-payloads §7's order. A hit stands in for the load step alone, whose outcome
//! for those bytes it is. The graphs run as they would after a fresh load, so the
//! output is the same bytes.
//! - **Bounded.** The entries ([`Resident`], the shared allocation that holds each, and
//! [`LoadedSets::INDEX_BYTES_PER_ENTRY`]) are kept within a budget:
//! [`LoadedSets::DEFAULT_BUDGET`] unless the sink sets one
//! ([`LoadedSets::set_budget`]). The least recently used entry goes first, and a set
//! larger than the whole budget is not kept. A budget of 0 turns the
//! store off; the set is then not even hashed.
//! - **Memory and work.** A decode's declared cost (integer-payloads §7) counts what the
//! decode holds beyond the payload and its parameter set. A set lowered and loaded is
//! the set's, the same for every payload, so an entry here is outside every decode's
//! declared working bytes, and the budget is what bounds it.
//!
//! There is one store per process ([`LoadedSets::global`]), behind a lock. On `wasm32`
//! there is one thread, and the lock is never contended. A load runs outside the lock,
//! so two threads that miss the same set at once both load it, and one entry is kept.
use std::any::Any;
use std::collections::BTreeMap;
use std::sync::{Arc, Mutex, MutexGuard, PoisonError};
use wai_ops::{Loaded, Profile};
/// What a loaded set holds on the heap, beyond the value itself.
pub trait Resident {
fn heap_bytes(&self) -> u64;
}
impl Resident for Loaded {
fn heap_bytes(&self) -> u64 {
self.graph().heap_bytes()
}
}
/// The conditional codec's two networks (`int_hyper_synth::ops::Nets`). The impl is here,
/// not beside the type, because the byte-equality crate includes `int_hyper_synth.rs`
/// and not this module.
impl Resident for crate::int_hyper_synth::ops::Nets {
fn heap_bytes(&self) -> u64 {
self.h_s.graph().heap_bytes() + self.g_s.graph().heap_bytes()
}
}
/// The key of a loaded set: the capability whose driver loaded it, the lowering
/// (integer-operators §13), the profile its graphs were loaded under (§11), and the
/// SHA-256 of the members it was lowered from, as [`SetKey::of`] forms it.
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
pub struct SetKey {
pub capability: &'static str,
pub lowering: &'static str,
pub profile: Profile,
pub digest: [u8; 32],
}
impl SetKey {
/// The key of the set whose members a decode reads, `members` (name → bytes): the
/// capability, the lowering and the profile, and [`SetKey::listing_digest`] of the
/// members. `None` where the names cannot make a listing; no set this crate decodes
/// has such a name, and such a set is then loaded and not kept.
pub fn of(capability: &'static str, lowering: &'static str, profile: Profile, members: &BTreeMap<String, &[u8]>) -> Option<SetKey> {
Some(SetKey { capability, lowering, profile, digest: SetKey::listing_digest(members)? })
}
/// The SHA-256 of the listing of `members` (`pin::listing_from`, each member by its
/// own SHA-256): the digest a pin in the form `wai.prior.set` over exactly those
/// members names. It is the digest of a [`SetKey`].
pub fn listing_digest(members: &BTreeMap<String, &[u8]>) -> Option<[u8; 32]> {
crate::pin::listing_from(members).ok().map(|l| crate::pin::sha256(&l))
}
}
/// What a store has done, and what it holds.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
#[non_exhaustive]
pub struct Stats {
/// Lookups that found their set loaded.
pub hits: u64,
/// Lookups that did not, and loaded it.
pub misses: u64,
/// Sets kept after they loaded.
pub kept: u64,
/// Entries removed to make room, or to fit a smaller budget.
pub evicted: u64,
/// Sets that loaded and were not kept, being larger than the whole budget.
pub too_large: u64,
/// The entries held now.
pub entries: u64,
/// Their bytes, as counted against the budget.
pub resident_bytes: u64,
/// The budget.
pub budget: u64,
}
/// A store of loaded parameter sets (the module docs): an entry per [`SetKey`], within a
/// byte budget, the least recently used going first.
pub struct LoadedSets(Mutex<Inner>);
struct Inner {
budget: u64,
resident: u64,
clock: u64,
sets: BTreeMap<SetKey, Entry>,
by_use: BTreeMap<u64, SetKey>,
stats: Stats,
}
struct Entry {
set: Arc<dyn Any + Send + Sync>,
bytes: u64,
used: u64,
}
static GLOBAL: LoadedSets = LoadedSets::new(LoadedSets::DEFAULT_BUDGET);
impl LoadedSets {
/// The budget a process starts with: 256 MiB. A set holds its weights as `i64` once
/// lowered, four times their bytes in a `WHM1` file, so this keeps a set of about
/// 60 MB of `WHM1` networks with room beside it for small ones.
pub const DEFAULT_BUDGET: u64 = 256 << 20;
/// What an entry is counted for beyond its set: its share of the store's two indexes
/// (by key, and by last use). It is more than either index's node for one entry, so
/// the count is never under what the store holds.
pub const INDEX_BYTES_PER_ENTRY: u64 = 4096;
/// An empty store with the budget `budget` (bytes; 0 is off).
pub const fn new(budget: u64) -> LoadedSets {
let stats = Stats { hits: 0, misses: 0, kept: 0, evicted: 0, too_large: 0, entries: 0, resident_bytes: 0, budget: 0 };
LoadedSets(Mutex::new(Inner { budget, resident: 0, clock: 0, sets: BTreeMap::new(), by_use: BTreeMap::new(), stats }))
}
/// The process's store, which every decode in this crate uses. It starts with
/// [`LoadedSets::DEFAULT_BUDGET`].
pub fn global() -> &'static LoadedSets {
&GLOBAL
}
fn lock(&self) -> MutexGuard<'_, Inner> {
// Nothing panics while the lock is held; a poisoned lock still holds a whole store.
self.0.lock().unwrap_or_else(PoisonError::into_inner)
}
/// The budget, in bytes.
pub fn budget(&self) -> u64 {
self.lock().budget
}
/// Set the budget, in bytes, evicting the least recently used entries until the rest
/// fit. 0 empties the store and turns it off: a decode then loads its set every time,
/// as it would with no store.
pub fn set_budget(&self, bytes: u64) {
let mut i = self.lock();
i.budget = bytes;
i.evict_to(bytes);
}
/// Remove every entry, keeping the budget and the counts.
pub fn clear(&self) {
let mut i = self.lock();
i.sets.clear();
i.by_use.clear();
i.resident = 0;
}
/// Set the counts of [`Stats`] to 0.
pub fn reset_stats(&self) {
self.lock().stats = Stats::default();
}
/// The counts, and what the store holds now.
pub fn stats(&self) -> Stats {
let i = self.lock();
Stats { entries: i.sets.len() as u64, resident_bytes: i.resident, budget: i.budget, ..i.stats }
}
/// Whether the set `key` is held.
pub fn contains(&self, key: &SetKey) -> bool {
self.lock().sets.contains_key(key)
}
/// The set `key()` names, from the store when it is held (a hit), or else `load()`
/// (a miss), kept when it loads and fits the budget. A refusal is returned as it is
/// and keeps nothing. With a budget of 0, `load()` alone: `key` is not called.
///
/// `load` must lower and load the set from exactly the bytes `key` was formed from,
/// by the lowering and under the profile it names.
pub fn get_or_load<T, E>(&self, key: impl FnOnce() -> Option<SetKey>, load: impl FnOnce() -> Result<T, E>) -> Result<Arc<T>, E>
where
T: Resident + Send + Sync + 'static,
{
if self.budget() == 0 {
return load().map(Arc::new);
}
let Some(key) = key() else { return load().map(Arc::new) };
{
let mut i = self.lock();
if let Some(set) = i.touch(&key).and_then(|s| s.downcast::<T>().ok()) {
i.stats.hits += 1;
return Ok(set);
}
i.stats.misses += 1;
}
let set = Arc::new(load()?);
let bytes = entry_bytes(&*set);
self.lock().keep(key, set.clone(), bytes);
Ok(set)
}
}
/// The bytes an entry is counted for: the shared allocation that holds the set (two
/// counts, then the value), what the value holds on the heap, and its share of the
/// indexes.
fn entry_bytes<T: Resident>(t: &T) -> u64 {
use std::mem::{align_of, size_of};
let align = align_of::<T>().max(align_of::<usize>());
let shared = ((2 * size_of::<usize>()).next_multiple_of(align_of::<T>()) + size_of::<T>()).next_multiple_of(align);
shared as u64 + t.heap_bytes() + LoadedSets::INDEX_BYTES_PER_ENTRY
}
impl Inner {
fn tick(&mut self) -> u64 {
self.clock += 1;
self.clock
}
/// The entry `key`, marked as used now.
fn touch(&mut self, key: &SetKey) -> Option<Arc<dyn Any + Send + Sync>> {
let now = self.tick();
let Inner { sets, by_use, .. } = self;
let e = sets.get_mut(key)?;
by_use.remove(&e.used);
by_use.insert(now, *key);
e.used = now;
Some(e.set.clone())
}
/// Keep `set`, of `bytes`, under `key`, evicting the least recently used entries to
/// make room; a set larger than the budget is not kept.
fn keep(&mut self, key: SetKey, set: Arc<dyn Any + Send + Sync>, bytes: u64) {
if self.budget == 0 {
return;
}
if bytes > self.budget {
self.stats.too_large += 1;
return;
}
if let Some(old) = self.sets.remove(&key) {
self.by_use.remove(&old.used);
self.resident -= old.bytes;
}
self.evict_to(self.budget - bytes);
let used = self.tick();
self.by_use.insert(used, key);
self.sets.insert(key, Entry { set, bytes, used });
self.resident += bytes;
self.stats.kept += 1;
}
/// Evict the least recently used entries until the rest hold at most `limit` bytes.
fn evict_to(&mut self, limit: u64) {
while self.resident > limit {
let Some((_, key)) = self.by_use.pop_first() else { break };
if let Some(e) = self.sets.remove(&key) {
self.resident -= e.bytes;
self.stats.evicted += 1;
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::cell::Cell;
/// A set of a given size, and a set of another type.
struct Blob(u64);
impl Resident for Blob {
fn heap_bytes(&self) -> u64 {
self.0
}
}
struct Other;
impl Resident for Other {
fn heap_bytes(&self) -> u64 {
0
}
}
fn key(d: u8) -> SetKey {
SetKey { capability: "c", lowering: "l", profile: Profile::Acc64, digest: [d; 32] }
}
fn get(s: &LoadedSets, d: u8, size: u64) -> Arc<Blob> {
s.get_or_load(|| Some(key(d)), || Ok::<_, ()>(Blob(size))).unwrap()
}
/// The bytes one entry of `size` heap bytes is counted for.
fn one(size: u64) -> u64 {
entry_bytes(&Blob(size))
}
#[test]
fn a_second_lookup_is_a_hit_and_returns_the_kept_set() {
let s = LoadedSets::new(LoadedSets::DEFAULT_BUDGET);
let a = get(&s, 1, 100);
let loads = Cell::new(0);
let b = s.get_or_load(|| Some(key(1)), || { loads.set(loads.get() + 1); Ok::<_, ()>(Blob(100)) }).unwrap();
assert!(Arc::ptr_eq(&a, &b) && loads.get() == 0, "the second lookup is answered from the store");
let st = s.stats();
assert_eq!((st.hits, st.misses, st.kept, st.entries, st.resident_bytes), (1, 1, 1, 1, one(100)));
}
#[test]
fn a_refusal_keeps_nothing() {
let s = LoadedSets::new(LoadedSets::DEFAULT_BUDGET);
for _ in 0..2 {
assert_eq!(s.get_or_load(|| Some(key(1)), || Err::<Blob, _>("refused")).err(), Some("refused"));
}
let st = s.stats();
assert_eq!((st.hits, st.misses, st.kept, st.entries), (0, 2, 0, 0), "a refused set is loaded, and refused, each time");
}
#[test]
fn a_key_differs_by_its_digest_and_a_type_never_crosses() {
let s = LoadedSets::new(LoadedSets::DEFAULT_BUDGET);
get(&s, 1, 10);
get(&s, 2, 10);
assert_eq!(s.stats().entries, 2);
// A set of another type under a held key is loaded, never cast.
let other = s.get_or_load(|| Some(key(1)), || Ok::<_, ()>(Other)).map(|_| ());
assert!(other.is_ok());
assert_eq!(s.stats().misses, 3);
}
#[test]
fn the_least_recently_used_entry_goes_first() {
let s = LoadedSets::new(2 * one(10));
get(&s, 1, 10);
get(&s, 2, 10);
get(&s, 1, 10); // 1 is now the most recently used
get(&s, 3, 10);
assert!(s.contains(&key(1)) && s.contains(&key(3)) && !s.contains(&key(2)));
let st = s.stats();
assert_eq!((st.evicted, st.entries, st.resident_bytes), (1, 2, 2 * one(10)));
// A smaller budget evicts down to it; a set past the whole budget is not kept.
s.set_budget(one(10));
assert!(s.contains(&key(3)) && !s.contains(&key(1)));
get(&s, 4, 11);
assert!(!s.contains(&key(4)) && s.contains(&key(3)));
assert_eq!(s.stats().too_large, 1);
}
#[test]
fn a_budget_of_zero_neither_keeps_nor_hashes() {
let s = LoadedSets::new(LoadedSets::DEFAULT_BUDGET);
get(&s, 1, 10);
s.set_budget(0);
assert_eq!(s.stats().entries, 0, "a budget of 0 empties the store");
let keyed = Cell::new(false);
let r = s.get_or_load(|| { keyed.set(true); Some(key(1)) }, || Ok::<_, ()>(Blob(10)));
assert!(r.is_ok() && !keyed.get(), "no key is formed");
assert_eq!((s.stats().misses, s.stats().entries), (1, 0));
}
}

9
wai/wai-rs/src/int_splat.rs
View file

@ -233,6 +233,15 @@ pub mod ops {
/// 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> {
plan_with(gs, st, load)
}
/// [`plan`], with `g_s` lowered and loaded by `load`, after every check before it: a
/// sink that keeps loaded sets (`crate::int_set_cache`) answers that one step from
/// its store.
pub(crate) fn plan_with<L>(
gs: &IntModel, st: &SplatStream<'_>, load: impl FnOnce(&IntModel) -> Result<L, Refused>,
) -> Result<(L, u64, u128), Refused> {
let (working, macs) = splat_plan(gs, st)?;
Ok((load(gs)?, working, macs))
}

50
wai/wai-rs/src/int_synth.rs
View file

@ -221,22 +221,46 @@ pub fn decode_bytes(payload: &[u8], spec_json: &[u8]) -> Result<IntImage, IntSyn
/// integer-exact decode a browser / WASM sink runs directly, host-supplied
/// weights in, RGB out. Byte-identical on every machine.
///
/// It decodes through the integer operator set. The payload is checked first
/// (`extensions/integer-payloads.md` §13): its header, a reserved byte of 0, `C ≥ 1`,
/// `C·S·S ≤ 2^26` and its exact length. Then the set is lowered (integer-operators
/// §13.1, whose checks include at least one layer, fields below `2^32`, `out_ch`, `k`
/// and `stride` at least 1 and a shift of at most 127) and loaded under `acc64`: load
/// steps 1–8, the last being the bounds certificate ([`crate::int_ops`]). Then the graph
/// runs on the latents, through the operator set's run-time plan with its caps, and its
/// output must be `[3, H, W]`. Every refusal is [`IntSynthError::InvalidPayload`]
/// (`payload_malformed`).
/// It decodes through the integer operator set, in three steps a sink may also take
/// one by one: [`check_payload`], [`load_spec`] and [`decode_loaded`]. The payload is
/// checked first (`extensions/integer-payloads.md` §13): its header, a reserved byte of
/// 0, `C ≥ 1`, `C·S·S ≤ 2^26` and its exact length. Then the set is lowered
/// (integer-operators §13.1, whose checks include at least one layer, fields below
/// `2^32`, `out_ch`, `k` and `stride` at least 1 and a shift of at most 127) and loaded
/// under `acc64`: load steps 1–8, the last being the bounds certificate
/// ([`crate::int_ops`]). Then the graph runs on the latents, through the operator set's
/// run-time plan with its caps, and its output must be `[3, H, W]`. Every refusal is
/// [`IntSynthError::InvalidPayload`] (`payload_malformed`). A sink that keeps loaded
/// sets between decodes answers the middle step from its store (`int_payload`'s decode
/// of `wai.neural.int_synth`).
pub fn decode_from_spec(payload: &[u8], spec_json: &[u8]) -> Result<IntImage, IntSynthError> {
let (h, w, c, s) = header(payload)?;
let refused = |what: &str, r: wai_ops::Refusal| IntSynthError::InvalidPayload(format!("int_synth {what}: {} at {:?}", r.code.as_str(), r.at));
check_payload(payload)?;
let g = load_spec(spec_json)?;
decode_loaded(payload, &g)
}
/// The payload's checks, the first step of [`decode_from_spec`].
pub fn check_payload(payload: &[u8]) -> Result<(), IntSynthError> {
header(payload).map(|_| ())
}
fn refused(what: &str, r: wai_ops::Refusal) -> IntSynthError {
IntSynthError::InvalidPayload(format!("int_synth {what}: {} at {:?}", r.code.as_str(), r.at))
}
/// The set `int_synth.json` lowered by `int_synth/1` and loaded under `acc64`, the
/// second step of [`decode_from_spec`].
pub fn load_spec(spec_json: &[u8]) -> Result<wai_ops::Loaded, IntSynthError> {
let set = wai_ops::lower::int_synth(spec_json).map_err(|r| refused("int_synth.json does not lower", r))?;
let g = crate::int_ops::load(set.get("synth").expect("int_synth/1 lowers to the graph synth")).map_err(|r| refused("set does not load", r))?;
crate::int_ops::load(set.get("synth").expect("int_synth/1 lowers to the graph synth")).map_err(|r| refused("set does not load", r))
}
/// The decode with its set loaded ([`load_spec`]), the last step of
/// [`decode_from_spec`]: the payload's checks again, then the graph on the latents.
pub fn decode_loaded(payload: &[u8], g: &wai_ops::Loaded) -> Result<IntImage, IntSynthError> {
let (h, w, c, s) = header(payload)?;
let latents: Vec<i64> = payload[HEADER..].iter().map(|&b| b as i8 as i64).collect();
let out = crate::int_ops::run1(&g, wai_ops::Tensor::new(c, s, s, latents), &[]).map_err(|r| refused("decode", r))?;
let out = crate::int_ops::run1(g, wai_ops::Tensor::new(c, s, s, latents), &[]).map_err(|r| refused("decode", r))?;
if (out.c, out.h as usize, out.w as usize) != (3, h, w) {
return Err(IntSynthError::InvalidPayload(format!(
"int_synth output [{},{},{}] != [3,{h},{w}]", out.c, out.h, out.w)));

2
wai/wai-rs/src/int_video.rs
View file

@ -393,7 +393,7 @@ pub fn decode(bs: &Value) -> Result<DecodedClip, &'static str> {
/// Decode the **P-frames** of a JSON bitstream on top of an already-decoded I-frame
/// (`i_frame`, RGB HWC `u8`) at `h × w` — the integer block-MC + quantized residual
/// loop, seeded with frame 0 from any source. The NEURAL-KEYFRAME video composes this
/// with the learned `int_hyper` I-frame decode (`int_hyper_synth::decode_conditional`
/// with the learned `int_hyper` I-frame decode (`int_hyper_synth::decode_conditional_checked`
/// → frame 0): a learned, byte-exact keyframe + integer inter-frames, every frame
/// reproducible by hash. Pure i64. `bs`'s `n_frames` counts the I-frame; `mvs[0]`
/// is empty; the rANS stream holds only frames `1..n_frames`'s residuals. A debugging

4
wai/wai-rs/src/lib.rs
View file

@ -65,6 +65,10 @@ pub mod int_kernels;
// loaded under its capability's profile, the bounds certificate included.
#[cfg(feature = "neural_int")]
pub mod int_ops;
// Parameter sets lowered and loaded once, kept between decodes by the SHA-256 of their
// bytes, within a byte budget (extensions/integer-operators.md §8.6).
#[cfg(feature = "neural_int")]
pub mod int_set_cache;
// Integer Gaussian-conditional entropy — the reusable scale→bucket primitive
// that makes a conditional codec's entropy decode hardware-deterministic (the

602
wai/wai-rs/tests/int_hyper_ops.rs Normal file
View file

@ -0,0 +1,602 @@
//! `wai.neural.int_hyper` and the keyframe of `wai.video.int_hyper` through the integer
//! operator set: their lowering `whm1/1` (`extensions/integer-operators.md` §13.2; `h_s`,
//! then `g_s` with the RGB8 step of §13.0) and the certificate of every committed set
//! (§8.5), as this crate loads them (`int_ops`, `int_hyper_synth::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 models it parses (`int_ops::whm1_net`), must write those same bytes, and the load a
//! sink makes must give the verdict the certificate records.
use std::path::PathBuf;
use serde_json::Value;
use wai::int_hyper_synth::{self as hyper, IntModel};
use wai_ops::Profile;
const CAPS: [&str; 2] = ["wai.neural.int_hyper", "wai.video.int_hyper"];
fn wai() -> PathBuf {
PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("..")
}
fn read(rel: &str) -> Vec<u8> {
std::fs::read(wai().join(rel)).unwrap_or_else(|e| panic!("{rel}: {e}"))
}
fn json(rel: &str) -> Value {
serde_json::from_slice(&read(rel)).unwrap_or_else(|e| panic!("{rel}: {e}"))
}
fn sha256(b: &[u8]) -> String {
wai_ops::sha256::sha256(b).iter().map(|x| format!("{x:02x}")).collect()
}
/// A committed set of one of the two capabilities, as `ops-conformance/certificates/index.json`
/// lists it: its name, capability, members (name, path below `wai/`), digest and lowered
/// graphs (`h_s`, then `g_s`).
struct Indexed {
name: String,
cap: String,
members: Vec<(String, String)>,
digest: String,
roles: [String; 2],
}
fn indexed() -> Vec<Indexed> {
json("ops-conformance/certificates/index.json")["sets"]
.as_array()
.unwrap()
.iter()
.filter(|s| CAPS.contains(&s["capability"].as_str().unwrap()))
.map(|s| {
assert_eq!(s["lowering"], "whm1/1");
let roles = s["roles"].as_array().unwrap();
assert_eq!(roles.len(), 2, "whm1/1 lowers a conditional codec's set to two graphs, h_s and g_s");
let members = s["members"].as_object().unwrap().iter().map(|(k, v)| (k.clone(), v.as_str().unwrap().to_string())).collect();
Indexed {
name: s["name"].as_str().unwrap().to_string(),
cap: s["capability"].as_str().unwrap().to_string(),
members,
digest: s["set_sha256"].as_str().unwrap().to_string(),
roles: [roles[0].as_str().unwrap().to_string(), roles[1].as_str().unwrap().to_string()],
}
})
.collect()
}
impl Indexed {
fn member(&self, name: &str) -> Vec<u8> {
read(&self.members.iter().find(|(k, _)| k == name).unwrap_or_else(|| panic!("{}: no {name}", self.name)).1)
}
}
/// The two graphs this crate lowers from the models it parses (`int_hyper_synth::ops::lower`),
/// whose conversion (`int_ops::whm1_net`) must give the networks `wai-ops` parses from the
/// bytes.
fn lowered_from_the_parsed_models(g_s: &[u8], h_s: &[u8]) -> (wai_ops::Graph, wai_ops::Graph) {
let (g, h) = (IntModel::from_bytes(g_s).unwrap(), IntModel::from_bytes(h_s).unwrap());
for (m, b) in [(&g, g_s), (&h, h_s)] {
let net = wai::int_ops::whm1_net(m).expect("a model parsed from bytes converts");
assert!(net == wai_ops::lower::parse_whm1(b).unwrap(), "the converted model is the one wai-ops parses from the bytes");
}
hyper::ops::lower(&g, &h).expect("the set lowers")
}
/// Gate 1 of the migration: every committed set of `wai.neural.int_hyper` and
/// `wai.video.int_hyper` lowers to its committed graphs, from its members' bytes
/// (`lower::int_hyper`) and from the models this crate parses; the certificate beside them
/// is those graphs'; and loading each graph under the capability's profile gives the
/// certificate's verdict. The sets carry `check` ranges (§13.2), so they meet `acc64` by
/// construction and not `acc32`; the first `acc32` failure is `h_s`'s input range, at 53
/// bits for the conformance set `tiny` (Appendix C's prediction) and 52 for the synthetic
/// one. The tiny set's JSON models, the form the tools write beside the binary one,
/// convert to the same networks and lower to the same graphs.
#[test]
fn every_committed_set_lowers_to_its_committed_graphs_and_verdict() {
let sets = indexed();
let names: Vec<(&str, &str)> = sets.iter().map(|s| (s.cap.as_str(), s.name.as_str())).collect();
assert_eq!(names, [(CAPS[0], "tiny"), (CAPS[1], "tiny"), (CAPS[0], "hyper")], "the committed int_hyper sets");
for s in &sets {
let (name, cap) = (&s.name, &s.cap);
let (g_s, h_s, ent) = (s.member("g_s.whm"), s.member("h_s.whm"), s.member("entropy.json"));
let members = vec![("entropy.json", ent.clone()), ("g_s.whm", g_s.clone()), ("h_s.whm", h_s.clone())];
assert_eq!(wai_ops::generate::sets::set_digest(&members), s.digest, "{name}: the set's digest");
// The entropy tables are the stage operators' data, not lowered; the set's must load.
hyper::CondEntropy::try_from_value(&serde_json::from_slice(&ent).unwrap()).unwrap_or_else(|e| panic!("{name}: entropy.json: {e}"));
let lowered = wai_ops::lower::int_hyper(&h_s, &g_s).unwrap_or_else(|r| panic!("{name}: does not lower: {r:?}"));
let roles: Vec<&str> = lowered.roles.iter().map(|(r, _)| *r).collect();
assert_eq!(roles, ["h_s", "g_s"], "{name}");
let (h, g) = lowered_from_the_parsed_models(&g_s, &h_s);
let cert = json(&format!("ops-conformance/certificates/{cap}/{}.json", s.digest));
assert_eq!(cert["set_sha256"].as_str(), Some(s.digest.as_str()), "{name}");
assert_eq!(cert["capability"].as_str(), Some(cap.as_str()), "{name}");
assert_eq!(cert["lowering"], "whm1/1", "{name}");
for (k, (role, ours)) in [("h_s", &h), ("g_s", &g)].into_iter().enumerate() {
let from_bytes = lowered.get(role).unwrap();
// The verdict first: the load a sink makes.
let loaded = wai::int_ops::load(from_bytes).unwrap_or_else(|r| panic!("{name} {role}: refused under acc64: {r:?}"));
assert!(!loaded.meets(Profile::Acc32), "{name} {role}: a graph with check ranges inside it does not meet acc32");
// Then the committed bytes: the graph, both ways, and the certificate's record of it.
let graph = from_bytes.to_bytes();
assert!(graph == read(&s.roles[k]), "{name}: the lowering of the bytes differs from the committed {}", s.roles[k]);
assert!(ours.to_bytes() == graph, "{name}: the lowering of the parsed models differs from the committed {role}");
let c = &cert["graphs"][k];
assert_eq!(c["name"].as_str(), Some(role), "{name}");
assert_eq!(c["graph_sha256"].as_str(), Some(sha256(&graph).as_str()), "{name} {role}: the certificate is the graph's");
// The input's check range is the lowering's X_0 for the network's first layer (§13.2).
let x0 = &c["inputs"][0];
let t = &from_bytes.tensors[0];
assert_eq!((x0["policy"].as_str(), x0["lo"].as_str(), x0["hi"].as_str()), (Some("check"), Some(t.lo.to_string().as_str()), Some(t.hi.to_string().as_str())), "{name} {role}");
assert!(t.hi > 0 && t.lo == -t.hi, "{name} {role}: a symmetric check range");
}
let v = &cert["verdict"];
assert_eq!((v["acc64"].as_bool(), v["acc32"].as_bool()), (Some(true), Some(false)), "{name}: the certificate's verdict");
let bits = if name == "tiny" { "53" } else { "52" };
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("h_s"), Some("input"), Some("0"), Some("range"), Some(bits)),
"{name}: the first acc32 failure"
);
// The decode's own load: both networks, h_s first.
let (gm, hm) = (IntModel::from_bytes(&g_s).unwrap(), IntModel::from_bytes(&h_s).unwrap());
let nets = hyper::ops::load(&gm, &hm).unwrap_or_else(|r| panic!("{name}: {r}"));
assert!(nets.h_s.graph() == lowered.get("h_s").unwrap() && nets.g_s.graph() == lowered.get("g_s").unwrap(), "{name}: the decode loads the committed graphs");
}
// The tiny set's JSON models, which the byte-exact keyframe golden and the browser's
// debug decode read, convert to the networks of its binary models.
let dir = "wai-rs/tests/vectors/int_hyper";
let from_json = |f: &str| IntModel::try_from_json(&json(&format!("{dir}/{f}"))).unwrap();
let (g, h) = hyper::ops::lower(&from_json("tiny_g_s.json"), &from_json("tiny_h_s.json")).unwrap();
let (gb, hb) = lowered_from_the_parsed_models(&read(&format!("{dir}/tiny_g_s.bin")), &read(&format!("{dir}/tiny_h_s.bin")));
assert!(g == gb && h == hb, "the tiny set's JSON and binary models lower alike");
}
/// The set of the pretrained model (`mbt2018_mean`, exported to integers) is too large to
/// commit, so its certificate is committed by the set's digest and checked against the
/// regenerated set in `wai-sota-vectors.yml`, where `tests/int_hyper_synth.rs` also holds
/// this crate's lowering of it to the graphs the certificate records. Its verdict is
/// Appendix C's: `acc64`, first failing `acc32` at `h_s`'s input range, 41 bits.
#[test]
fn the_pretrained_sets_certificate_is_the_one_appendix_c_predicts() {
let ext = json("ops-conformance/certificates/external.json");
let s = ext["sets"].as_array().unwrap().iter().find(|s| s["capability"] == CAPS[0]).expect("the pretrained set is listed");
assert_eq!((s["lowering"].as_str(), s["profile"].as_str()), (Some("whm1/1"), Some("acc64")));
let digest = s["set_sha256"].as_str().unwrap();
let cert = json(&format!("ops-conformance/certificates/{}/{digest}.json", CAPS[0]));
let names: Vec<&str> = cert["graphs"].as_array().unwrap().iter().map(|g| g["name"].as_str().unwrap()).collect();
assert_eq!(names, ["h_s", "g_s"]);
let v = &cert["verdict"];
assert_eq!((v["acc64"].as_bool(), v["acc32"].as_bool()), (Some(true), Some(false)));
let first = &v["acc32_first_failure"];
assert_eq!(
(first["graph"].as_str(), first["where"].as_str(), first["key"].as_str(), first["bits"].as_str()),
(Some("h_s"), Some("input"), Some("range"), Some("41"))
);
}
/// The typed form of a parsed `WIH1` payload.
fn stream<'a>(k: &'a wai::int_wire::Wih1<'a>) -> hyper::CondStream<'a> {
hyper::CondStream {
m: k.m as usize, sh: k.sh as usize, sw: k.sw as usize, n_z: k.n_z as usize, zh: k.zh as usize, zw: k.zw as usize,
z_rans: k.z_rans, rans: k.rans,
}
}
/// A committed set's three members, parsed: `g_s`, `h_s` and the entropy tables.
fn set_of(v: &Value) -> (IntModel, IntModel, hyper::CondEntropy) {
let m = |k: &str| read(v["set"][k].as_str().unwrap());
let ent = hyper::CondEntropy::try_from_value(&serde_json::from_slice(&m("entropy.json")).unwrap()).unwrap();
(IntModel::from_bytes(&m("g_s.whm")).unwrap(), IntModel::from_bytes(&m("h_s.whm")).unwrap(), ent)
}
fn hex(b: &[u8]) -> String {
b.iter().map(|x| format!("{x:02x}")).collect()
}
/// Every frame's SHA-256 and the clip's, against the sidecar entry `key` of
/// `int_video/sha256.json`; every frame's FNV-1a and the clip's, against `conf`.
fn assert_clip(frames: &[Vec<u8>], key: &str, conf: &Value) {
use sha2::Digest;
let side = json("wai-rs/tests/vectors/int_video/sha256.json");
let (want, fnv) = (side[key]["frames_sha256"].as_array().unwrap(), conf["frame_hashes"].as_array().unwrap());
assert_eq!((want.len(), fnv.len()), (frames.len(), frames.len()), "{key}: frame count");
let (mut all, mut all_fnv) = (Vec::new(), Vec::new());
for (i, f) in frames.iter().enumerate() {
let d = sha2::Sha256::digest(f);
assert_eq!(hex(&d), want[i].as_str().unwrap(), "{key}: frame {i}");
let h = format!("{:016x}", wai::int_video::fnv1a(f));
assert_eq!(h, fnv[i].as_str().unwrap(), "{key}: frame {i}'s FNV-1a");
all.extend_from_slice(&d);
all_fnv.extend((0..8).map(|j| u8::from_str_radix(&h[2 * j..2 * j + 2], 16).unwrap()));
}
assert_eq!(hex(&sha2::Sha256::digest(&all)), side[key]["clip_sha256"].as_str().unwrap(), "{key}: the clip's SHA-256");
assert_eq!(format!("{:016x}", wai::int_video::fnv1a(&all_fnv)), conf["clip_hash"].as_str().unwrap(), "{key}: the clip's FNV-1a");
}
/// Gate 2 of the migration: golden equality. Every committed vector of
/// `wai.neural.int_hyper` and `wai.video.int_hyper` (`int-payload-conformance/vectors.json`)
/// decodes through the lowering (`int_hyper_synth::ops`) to the image the direct decode
/// gives (`int_hyper_synth::legacy`), and to every digest the corpus commits for it:
/// - `tiny` (`WIH1`): the SHA-256 sidecar and the export's FNV-1a, from the payload, from
/// the JSON form it is converted from, and with the set's JSON models, the form the
/// byte-exact keyframe golden reads;
/// - `video_nk` (`WIV1`, keyframe 1): its nested `WIH1` keyframe, then the clip's integer
/// P-frames on it, every frame's and the clip's SHA-256 and FNV-1a.
#[test]
fn every_committed_vector_decodes_through_the_lowering_to_its_golden() {
let index = json("int-payload-conformance/vectors.json");
let vs: Vec<&Value> = index["vectors"].as_array().unwrap().iter().filter(|v| CAPS.contains(&v["capability"].as_str().unwrap())).collect();
let names: Vec<&str> = vs.iter().map(|v| v["name"].as_str().unwrap()).collect();
assert_eq!(names, ["tiny", "video_nk"], "the committed int_hyper vectors");
let dir = "wai-rs/tests/vectors";
// tiny: WIH1, its JSON form, and the JSON models.
let v = vs[0];
let (g_s, h_s, ent) = set_of(v);
let p = read(v["payload"].as_str().unwrap());
assert_eq!(sha256(&p), v["payload_sha256"].as_str().unwrap());
let k = wai::int_wire::parse_wih1(&p).unwrap();
let bs = hyper::CondBitstream::try_from_value(&json(v["converted_from"]["json"].as_str().unwrap())).unwrap();
let from_json = |f: &str| IntModel::try_from_json(&json(&format!("{dir}/int_hyper/{f}"))).unwrap();
let (gj, hj) = (from_json("tiny_g_s.json"), from_json("tiny_h_s.json"));
let want_sha = json(v["sidecar"].as_str().unwrap())["tiny"]["sha256"].as_str().unwrap().to_string();
let conf = json(&format!("{dir}/int_hyper/tiny.json"));
for (form, g, h, st) in [("WIH1", &g_s, &h_s, stream(&k)), ("JSON", &g_s, &h_s, bs.stream()), ("JSON models", &gj, &hj, stream(&k))] {
let ops = hyper::ops::decode(g, h, &ent, &st).unwrap_or_else(|r| panic!("tiny {form}: {r}"));
let direct = hyper::legacy::decode_conditional_checked(g, h, &ent, &st).unwrap_or_else(|r| panic!("tiny {form}: {r}"));
assert!(ops == direct, "tiny {form}: the decode through the lowering differs from the direct decode");
assert_eq!((ops.1, ops.2), (16, 16), "tiny {form}");
assert_eq!(sha256(&ops.0), want_sha, "tiny {form}: the SHA-256 sidecar");
assert_eq!(format!("{:016x}", hyper::fnv1a(&ops.0)), conf["rgb_hash"].as_str().unwrap(), "tiny {form}: the export's FNV-1a");
}
// video_nk: WIV1 with keyframe 1.
let v = vs[1];
let (g_s, h_s, ent) = set_of(v);
let p = read(v["payload"].as_str().unwrap());
assert_eq!(sha256(&p), v["payload_sha256"].as_str().unwrap());
let w = wai::int_wire::parse_wiv1(&p, Some(1)).unwrap();
let k = wai::int_wire::parse_wih1(w.wih1.expect("keyframe 1 nests a WIH1")).unwrap();
let (frame0, h, wd) = hyper::ops::decode(&g_s, &h_s, &ent, &stream(&k)).unwrap();
assert!(hyper::legacy::decode_conditional_checked(&g_s, &h_s, &ent, &stream(&k)).unwrap() == (frame0.clone(), h, wd), "video_nk: the keyframe differs from the direct decode's");
assert_eq!((h, wd), (w.h as usize, w.w as usize), "video_nk: the keyframe is the clip's size");
let s = wai::int_video::VideoStream {
n_frames: w.n_frames as usize, h: w.h as usize, w: w.w as usize, block: w.block as usize,
q: w.q as i64, offset: w.offset, mv_precision: w.mv_precision, cdf: &w.cdf, mvs: &w.mvs, rans: w.rans,
};
let clip = wai::int_video::decode_pframes_typed(&frame0, &s).unwrap();
assert_clip(&clip.frames, "video_nk", &json(&format!("{dir}/int_video/video_nk.json")));
}
/// The code integer-payloads §8 gives a refusal of the conditional decode.
fn code_of(r: hyper::Refused) -> &'static str {
wai::int_payload::IntPayloadError::from(r).code()
}
/// The committed conformance set and payload (`int-payload-conformance`'s `tiny`).
const DIR: &str = "wai-rs/tests/vectors/int_hyper";
/// A typed `WIH1` payload that owns its streams, so a fixture can change any field.
#[derive(Clone)]
struct Owned {
m: usize,
sh: usize,
sw: usize,
n_z: usize,
zh: usize,
zw: usize,
z_rans: Vec<u8>,
rans: Vec<u8>,
}
impl Owned {
fn committed() -> Owned {
let p = read(&format!("{DIR}/tiny_bitstream.wih1"));
let k = wai::int_wire::parse_wih1(&p).unwrap();
Owned {
m: k.m as usize, sh: k.sh as usize, sw: k.sw as usize, n_z: k.n_z as usize, zh: k.zh as usize, zw: k.zw as usize,
z_rans: k.z_rans.to_vec(), rans: k.rans.to_vec(),
}
}
fn st(&self) -> hyper::CondStream<'_> {
hyper::CondStream { m: self.m, sh: self.sh, sw: self.sw, n_z: self.n_z, zh: self.zh, zw: self.zw, z_rans: &self.z_rans, rans: &self.rans }
}
}
/// The committed set, parsed, changed by `f`: `(g_s, h_s, entropy tables)`.
fn tiny(f: impl FnOnce(&mut IntModel, &mut IntModel, &mut hyper::CondEntropy)) -> (IntModel, IntModel, hyper::CondEntropy) {
let mut g = IntModel::from_bytes(&read(&format!("{DIR}/tiny_g_s.bin"))).unwrap();
let mut h = IntModel::from_bytes(&read(&format!("{DIR}/tiny_h_s.bin"))).unwrap();
let mut e = hyper::CondEntropy::try_from_value(&json(&format!("{DIR}/tiny_entropy.json"))).unwrap();
f(&mut g, &mut h, &mut e);
(g, h, e)
}
/// The committed payload, changed by `f`.
fn payload(f: impl FnOnce(&mut Owned)) -> Owned {
let mut o = Owned::committed();
f(&mut o);
o
}
/// Both decodes of a typed payload against a set: through the lowering, then directly.
fn both(set: &(IntModel, IntModel, hyper::CondEntropy), o: &Owned) -> (Result<Vec<u8>, &'static str>, Result<Vec<u8>, &'static str>) {
let (g, h, e) = set;
(
hyper::ops::decode(g, h, e, &o.st()).map(|d| d.0).map_err(code_of),
hyper::legacy::decode_conditional_checked(g, h, e, &o.st()).map(|d| d.0).map_err(code_of),
)
}
/// Gate 3 of the migration: refusal parity, on the committed fixtures. Each fixture of
/// `int-payload-conformance` for `wai.neural.int_hyper` and `wai.video.int_hyper` is
/// refused by a sink with its code. For one that parses (`decode`), its `WIH1` (the
/// payload, or a `WIV1`'s keyframe) decodes through the lowering and directly to the same
/// image, or is refused by both with the same code: the fixture's, where the keyframe is
/// what it refuses.
#[test]
fn every_committed_refusal_fixture_keeps_its_code() {
use std::collections::BTreeMap;
let (mut n, mut keyframes) = (0, 0);
for f in json("int-payload-conformance/vectors.json")["refusals"].as_array().unwrap().iter().filter(|f| CAPS.contains(&f["capability"].as_str().unwrap())) {
let (name, cap, want) = (f["name"].as_str().unwrap(), f["capability"].as_str().unwrap(), f["code"].as_str().unwrap());
let p = read(f["payload"].as_str().unwrap());
let set: BTreeMap<String, Vec<u8>> = f["set"].as_object().unwrap().iter().map(|(k, v)| (k.clone(), read(v.as_str().unwrap()))).collect();
let view: BTreeMap<String, &[u8]> = set.iter().map(|(k, v)| (k.clone(), v.as_slice())).collect();
let params = wai::int_payload::resolve(cap, None, &[wai::pin::HeldBytes::Set { listing: b"", files: &view }], None).expect("an unpinned set resolves");
let sink = wai::int_payload::decode_int(cap, &p, &params).map(|_| ()).map_err(|e| e.code());
assert_eq!(sink, Err(want), "{name}: the sink's code");
if f["layer"] == "decode" {
let nested = if cap == CAPS[0] { p.clone() } else { wai::int_wire::parse_wiv1(&p, Some(1)).unwrap().wih1.unwrap().to_vec() };
let k = wai::int_wire::parse_wih1(&nested).unwrap_or_else(|e| panic!("{name}: a decode fixture's keyframe parses: {e:?}"));
let g = IntModel::from_bytes(&set["g_s.whm"]).unwrap();
let h = IntModel::from_bytes(&set["h_s.whm"]).unwrap();
let e = hyper::CondEntropy::try_from_value(&serde_json::from_slice(&set["entropy.json"]).unwrap()).unwrap();
let ops = hyper::ops::decode(&g, &h, &e, &stream(&k)).map(|d| d.0).map_err(code_of);
let direct = hyper::legacy::decode_conditional_checked(&g, &h, &e, &stream(&k)).map(|d| d.0).map_err(code_of);
assert!(ops == direct, "{name}: the two decodes of the keyframe differ: {:?} / {:?}", ops.as_ref().err(), direct.as_ref().err());
if cap == CAPS[0] {
assert_eq!(ops, Err(want), "{name}: the WIH1 fixture's code, both ways");
}
keyframes += 1;
}
n += 1;
}
assert!(n >= 33 && keyframes >= 15, "the committed int_hyper fixtures: {n}, {keyframes} decoded");
}
/// Gate 3, beside the committed fixtures: a fixture for each refusal the direct decode
/// makes (integer-payloads §7, in its order: the set against the payload, the networks
/// planned, the shapes, then each stream and each network as it runs), each made from the
/// committed set and payload by one change. Each is refused by both decodes with the same
/// code, the one integer-payloads §8 gives that rule. The committed set and payload decode.
#[test]
fn every_refusal_of_the_direct_decode_is_refused_alike_through_the_lowering() {
use hyper::IntLayer;
let base = tiny(|_, _, _| {});
let (ops, direct) = both(&base, &Owned::committed());
assert!(ops.is_ok() && ops == direct, "the fixtures' base decodes alike");
let same = || tiny(|_, _, _| {});
let as_is = || payload(|_| {});
// A hyper-latent stream that runs out: two symbols of frequency 1 in 65536 each cost
// 16 bits, so the 8 symbols need more than the stream's 64-bit start, cut short.
let costly_z = |e: &mut hyper::CondEntropy| {
for k in 0..2 {
e.z_cdfs[k] = vec![0, 1, 2, 1 << 16];
e.z_cdf_lengths[k] = 4;
}
};
let costly_z_stream = || {
let mut enc = wai::rans::RansEncoder::new();
for _ in 0..8 {
assert!(enc.encode(-8, &[0, 1, 2, 1 << 16], 4, -8));
}
let mut s = enc.flush();
s.truncate(s.len() - 4);
s
};
let h_bias = |f: fn(&mut Vec<i64>, usize)| tiny(move |_, h, _| {
if let IntLayer::Conv { bias, .. } = &mut h.layers[2] { f(bias, 4) }
});
let fixtures: Vec<(&str, &str, (IntModel, IntModel, hyper::CondEntropy), Owned)> = vec![
("h_s's afrac not g_s's", "model_mismatch", tiny(|_, h, _| h.afrac = 11), as_is()),
("the set's afrac of 0", "model_mismatch", tiny(|g, h, _| (g.afrac, h.afrac) = (0, 0)), as_is()),
("the set's afrac of 31", "model_mismatch", tiny(|g, h, _| (g.afrac, h.afrac) = (31, 31)), as_is()),
("no scale edge", "model_mismatch", tiny(|_, _, e| { e.edges.clear(); e.cdfs.clear(); e.cdf_lengths.clear(); e.offsets.clear(); }), as_is()),
("a latent table short", "model_mismatch", tiny(|_, _, e| { e.cdfs.pop(); }), as_is()),
("a hyper-latent median short", "model_mismatch", tiny(|_, _, e| { e.z_median_q.pop(); }), as_is()),
("a table that does not end at 2^16", "model_mismatch", tiny(|_, _, e| { let n = e.cdf_lengths[1]; e.cdfs[1][n - 1] = 65535; }), as_is()),
("a hyper-latent table's length past its row", "model_mismatch", tiny(|_, _, e| e.z_cdf_lengths[0] = e.z_cdfs[0].len() + 1), as_is()),
("a latent dimension of 0", "field_range", same(), payload(|o| o.m = 0)),
("more hyper-latent channels than tables", "model_mismatch", same(), payload(|o| o.n_z = 3)),
("a hyper-latent past 2^26", "too_large", same(), payload(|o| (o.zh, o.zw) = (1 << 13, 1 << 13))),
("a latent past 2^26", "too_large", same(), payload(|o| (o.sh, o.sw) = (1 << 12, 1 << 12))),
("an uncropped plane of h_s past 2^26", "too_large", same(), payload(|o| (o.zh, o.zw) = (1 << 11, 1 << 11))),
("an h_s layer's bias short", "model_mismatch", h_bias(|b, _| { b.pop(); }), as_is()),
("padding that leaves g_s no output", "model_mismatch", tiny(|g, _, _| if let IntLayer::ConvT { pad, .. } = &mut g.layers[2] { *pad = 1000 }), as_is()),
("an h_s output that is not the declared latent's shape", "geometry_mismatch", same(), payload(|o| o.sh = 5)),
("a g_s that gives four channels", "model_mismatch", tiny(|g, _, _| if let IntLayer::ConvT { out_ch, weight, mult, bias, .. } = &mut g.layers[2] {
*out_ch = 4;
*weight = vec![1; 4 * 4 * 9];
mult.push(mult[0]);
bias.push(0);
}), as_is()),
("a hyper-latent stream too short to be one", "field_range", same(), payload(|o| o.z_rans.truncate(4))),
("a hyper-latent stream that runs out", "stream_exhausted", tiny(|_, _, e| costly_z(e)), payload(|o| o.z_rans = costly_z_stream())),
("a word left unread in the hyper-latent stream", "stream_noncanonical", same(), payload(|o| o.z_rans.extend([0, 0, 0, 0]))),
("a hyper-latent median at i64's end", "field_range", tiny(|_, _, e| e.z_median_q[0] = i64::MAX), as_is()),
("hyper-latents past h_s's range", "field_range", tiny(|_, _, e| e.z_median_q = vec![1 << 60; 2]), as_is()),
("a latent stream too short to be one", "field_range", same(), payload(|o| o.rans.truncate(4))),
("a latent stream that runs out", "stream_exhausted", same(), payload(|o| { let n = o.rans.len() / 2; o.rans.truncate(n - n % 4); })),
("a word left unread in the latent stream", "stream_noncanonical", same(), payload(|o| o.rans.extend([0, 0, 0, 0]))),
("means past g_s's range", "field_range", h_bias(|b, m| b[m..].iter_mut().for_each(|v| *v = 1 << 55)), as_is()),
("means near i64's end", "field_range", h_bias(|b, m| b[m..].iter_mut().for_each(|v| *v = i64::MAX - (1 << 20))), as_is()),
];
for (what, want, set, o) in &fixtures {
let (ops, direct) = both(set, 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: for each layer of the committed `h_s` and
/// `g_s`, a network of that layer alone (with the set's header; a leaky ReLU followed by a
/// zero 1×1 transposed convolution, since a network with no layer of tag 0, 1 or 2 does
/// not lower) 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 hyper::{run_net_checked, IntLayer};
use wai_ops::{Code, Site, Tensor};
let mut seen = 0;
for (net, rgb8) in [("tiny_h_s.bin", false), ("tiny_g_s.bin", true)] {
let set = read(&format!("{DIR}/{net}"));
let full = wai_ops::lower::whm1(&set, rgb8).unwrap();
let n = wai_ops::lower::parse_whm1(&set).unwrap().layers.len();
for l in 0..n {
let mut one = IntModel::from_bytes(&set).unwrap();
let layer = one.layers.swap_remove(l);
let leaky = matches!(layer, IntLayer::Leaky { .. });
one.layers = vec![layer];
let c = full.tensors[l].channels;
if leaky {
one.layers.push(IntLayer::ConvT { out_ch: 1, k: 1, stride: 1, pad: 0, opad: 0, weight: vec![0; c as usize], mult: vec![1], bias: vec![0] });
}
let g = wai_ops::lower::whm1_net(&wai::int_ops::whm1_net(&one).unwrap(), false).unwrap();
let (lo, hi) = (g.tensors[0].lo, g.tensors[0].hi);
assert_eq!((g.tensors[0].channels, lo, hi), (c, full.tensors[l].lo, full.tensors[l].hi), "{net} layer {l}: the range of tensor {l}");
assert!(hi > 0 && lo == -hi, "{net} 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!("{net} layer {l}: {v} runs: {r:?}"));
let (direct, ..) = run_net_checked(&one, &x(v), c as usize, 2, 2).unwrap_or_else(|r| panic!("{net} layer {l}: {v} runs directly: {r}"));
assert!(ops.data == direct, "{net} 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)), "{net} 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, "{net} layer {l}: {v}: {d}");
}
seen += 1;
}
}
assert_eq!(seen, 6, "every layer of the committed h_s and g_s");
}
/// 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 convolutions, the transposed convolutions and the inverse GDN on this
/// crate's fast kernels, serial and threaded) gives the same image, and the same
/// refusals, on the committed vector (as `WIH1`, as its JSON form, and with the set's JSON
/// models), and on sets whose hyper-latent channel 0 lies about the edge of `h_s`'s input
/// check range (its median at `±X_0` and one past it). The per-operator half is `int_ops`' equivalence
/// tests, one of which counts the fast kernels' 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 x0 = wai_ops::lower::whm1(&read(&format!("{DIR}/tiny_h_s.bin")), false).unwrap().tensors[0].hi;
let base = Owned::committed();
let bs = hyper::CondBitstream::try_from_value(&json(&format!("{DIR}/tiny_bitstream.json"))).unwrap();
let from_json = |f: &str| IntModel::try_from_json(&json(&format!("{DIR}/{f}"))).unwrap();
let json_models = (from_json("tiny_g_s.json"), from_json("tiny_h_s.json"), tiny(|_, _, _| {}).2);
let json_form = Owned { m: bs.m, sh: bs.sh, sw: bs.sw, n_z: bs.n_z, zh: bs.zh, zw: bs.zw, z_rans: bs.z_rans.clone(), rans: bs.rans.clone() };
let mut cases: Vec<(String, (IntModel, IntModel, hyper::CondEntropy), Owned)> =
vec![("WIH1".into(), tiny(|_, _, _| {}), base.clone()), ("JSON".into(), tiny(|_, _, _| {}), json_form), ("JSON models".into(), json_models, base.clone())];
// Channel 0's median moved to ±X_0 or one past it: its elements, the median plus each
// symbol at 2^12, lie about the edge of h_s's input check range.
for (what, t) in [("channel 0 at X_0", x0), ("at X_0 + 1", x0 + 1), ("at -X_0", -x0), ("at -X_0 - 1", -x0 - 1)] {
cases.push((what.into(), tiny(move |_, _, e| e.z_median_q[0] = t), base.clone()));
}
for (k, (what, set, o)) in cases.iter().enumerate() {
let (g, h, e) = set;
let want = with_policy(POLICIES[0], || hyper::ops::decode(g, h, e, &o.st())).map(|d| d.0).map_err(code_of);
assert!(k >= 3 || want.is_ok(), "{what}: the committed form decodes");
for pol in &POLICIES[1..] {
let got = with_policy(*pol, || hyper::ops::decode(g, h, e, &o.st())).map(|d| d.0).map_err(code_of);
assert!(got == want, "{what}: {pol:?} differs from the interpreter alone");
}
}
}
/// Every committed `WIH1` payload, alone or inside another file (a `WIV1`'s keyframe, an
/// envelope, a staged base, a counting fixture): each occurrence of the magic in the
/// corpora that hold integer payloads whose bytes from there parse as a `WIH1` of exactly
/// the length its header and stream lengths give. Returns `(where, payload)` pairs.
fn committed_wih1_payloads() -> Vec<(String, Vec<u8>)> {
fn files(d: &std::path::Path, out: &mut Vec<PathBuf>) {
let mut es: Vec<_> = std::fs::read_dir(d).unwrap().filter_map(|e| e.ok()).map(|e| e.path()).collect();
es.sort();
for p in es {
if p.is_dir() { files(&p, out) } else { out.push(p) }
}
}
let mut all = Vec::new();
for d in ["int-payload-conformance", "staged-conformance", "staged-measure-conformance", "http-conformance/objects", "wai-rs/tests/vectors"] {
files(&wai().join(d), &mut all);
}
let u32_at = |b: &[u8], i: usize| b.get(i..i + 4).map(|s| u32::from_le_bytes(s.try_into().unwrap()) as usize);
let mut v = Vec::new();
for f in all {
let b = std::fs::read(&f).unwrap();
for i in (0..b.len().saturating_sub(3)).filter(|&i| &b[i..i + 4] == b"WIH1") {
let Some(lz) = u32_at(&b, i + 32) else { continue };
let Some(ly) = i.checked_add(36 + lz).and_then(|j| u32_at(&b, j)) else { continue };
let end = i + 40 + lz + ly;
if let Some(p) = b.get(i..end).filter(|p| wai::int_wire::parse_wih1(p).is_ok()) {
v.push((format!("{}@{i}", f.strip_prefix(wai()).unwrap().display()), p.to_vec()));
}
}
}
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 of either capability has a
/// layer of tag 0, 1 or 2 in each network, header shifts under 128, no bias of magnitude
/// `2^63 − 1` or more, and its two networks lower and load under `acc64`; every committed
/// `WIH1` payload, wherever it is carried, decodes through the lowering to the bytes the
/// direct decode gives, or is refused by both with the same code, against the set every
/// committed one is decoded with (`tiny`). (The other files that hold the magic are code:
/// the browser crate's recorded build, its demo page and its scripts.)
#[test]
fn the_new_load_rules_affect_no_committed_set_or_payload() {
use wai_ops::lower::NetLayer as N;
for s in indexed() {
for (role, rgb8) in [("h_s.whm", false), ("g_s.whm", true)] {
let b = s.member(role);
let net = wai_ops::lower::parse_whm1(&b).unwrap();
assert!(net.layers.iter().any(|l| matches!(l, N::ConvT { .. } | N::Conv { .. } | N::Igdn { .. })), "{} {role}: no layer of tag 0, 1 or 2", s.name);
assert!(net.shift < 128 && net.gf < 128, "{} {role}: a header shift of 128 or more", s.name);
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), "{} {role}: a bias of magnitude 2^63 - 1 or more", s.name);
}
}
let g = wai_ops::lower::whm1(&b, rgb8).unwrap_or_else(|r| panic!("{} {role}: does not lower: {r:?}", s.name));
assert!(wai::int_ops::load(&g).is_ok(), "{} {role}: does not load under acc64", s.name);
}
}
let set = tiny(|_, _, _| {});
let (mut decoded, mut refused) = (0, 0);
let payloads = committed_wih1_payloads();
for (at, p) in &payloads {
let k = wai::int_wire::parse_wih1(p).unwrap();
let o = Owned { m: k.m as usize, sh: k.sh as usize, sw: k.sw as usize, n_z: k.n_z as usize, zh: k.zh as usize, zw: k.zw as usize, z_rans: k.z_rans.to_vec(), rans: k.rans.to_vec() };
let (ops, direct) = both(&set, &o);
assert!(ops == direct, "{at}: the decode through the lowering differs from the direct decode");
if ops.is_ok() { decoded += 1 } else { refused += 1 }
}
let staged = payloads.iter().filter(|(at, _)| at.starts_with("staged-conformance")).count();
assert!(decoded >= 10 && refused >= 5 && staged >= 5, "the committed WIH1 payloads, at least those of this revision: {decoded} decoded, {refused} refused, {staged} staged");
}

87
wai/wai-rs/tests/int_hyper_synth.rs
View file

@ -70,7 +70,8 @@ fn rust_integer_synthesis_reproduces_the_export_hash() {
let latent: Vec<i64> =
input["data"].as_array().unwrap().iter().map(|n| n.as_i64().unwrap()).collect();
let (rgb, oh, ow) = decode(&IntModel::from_json(&model), &latent, c, h, w);
let model = IntModel::from_json(&model);
let (rgb, oh, ow) = decode(&model, &latent, c, h, w);
assert_eq!((oh, ow), (conf["out_h"].as_u64().unwrap() as usize, conf["out_w"].as_u64().unwrap() as usize));
assert_eq!(rgb.len(), oh * ow * conf["out_c"].as_u64().unwrap() as usize);
@ -78,6 +79,24 @@ fn rust_integer_synthesis_reproduces_the_export_hash() {
let want = conf["rgb_hash"].as_str().unwrap();
println!("int_hyper integer synthesis: {oh}x{ow}, fnv1a {got} (psnr {:.2} dB)", conf["psnr_int"].as_f64().unwrap_or(0.0));
assert_eq!(got, want, "Rust integer synthesis must reproduce the export's reconstruction byte-for-byte");
// The same synthesis through its lowering (whm1/1 with the RGB8 step), as the
// decode a sink runs computes it: the same bytes, row-major HWC.
let t = through_the_lowering(&model, true, latent, c, h, w);
let plane = t.h as usize * t.w as usize;
let mut hwc = vec![0u8; t.data.len()];
for (k, &v) in t.data.iter().enumerate() {
hwc[(k % plane) * t.c as usize + k / plane] = v as u8;
}
assert_eq!(format!("{:016x}", fnv1a(&hwc)), want, "the synthesis through the integer operator set must reproduce the export's reconstruction");
}
/// One network of the pretrained set through its lowering `whm1/1` (with the RGB8 step
/// where `rgb8` is set), loaded under `acc64` and run by the integer operator set
/// (`int_ops`, its fast kernels behind it) on `x`, `[c, h, w]`.
fn through_the_lowering(m: &IntModel, rgb8: bool, x: Vec<i64>, c: usize, h: usize, w: usize) -> wai_ops::Tensor {
let g = wai_ops::lower::whm1_net(&wai::int_ops::whm1_net(m).expect("the model converts"), rgb8).expect("the network lowers");
let l = wai::int_ops::load(&g).expect("the network loads under acc64");
wai::int_ops::run1(&l, wai_ops::Tensor::new(c as u32, h as u32, w as u32, x), &[]).expect("the input is within its check range")
}
#[test]
@ -96,7 +115,8 @@ fn rust_integer_hyper_decoder_reproduces_the_export_hash() {
input["w"].as_u64().unwrap() as usize,
);
let z: Vec<i64> = input["data"].as_array().unwrap().iter().map(|n| n.as_i64().unwrap()).collect();
let (out, oc, oh, ow) = wai::int_hyper_synth::run_net(&IntModel::from_json(&model), &z, c, h, w);
let model = IntModel::from_json(&model);
let (out, oc, oh, ow) = wai::int_hyper_synth::run_net(&model, &z, c, h, w);
assert_eq!((oc, oh, ow), (
conf["out_c"].as_u64().unwrap() as usize,
conf["out_h"].as_u64().unwrap() as usize,
@ -105,6 +125,11 @@ fn rust_integer_hyper_decoder_reproduces_the_export_hash() {
let got = format!("{:016x}", fnv1a(&bytes));
println!("integer h_s: z_hat {c}x{h}x{w} -> sigma||mu {oc}x{oh}x{ow}, fnv1a {got}");
assert_eq!(got, conf["params_hash"].as_str().unwrap(), "integer h_s must reproduce the export byte-for-byte");
// The same hyper-decoder through its lowering (whm1/1), as the decode a sink runs.
let t = through_the_lowering(&model, false, z, c, h, w);
assert_eq!((t.c as usize, t.h as usize, t.w as usize), (oc, oh, ow));
let bytes: Vec<u8> = t.data.iter().flat_map(|v| v.to_le_bytes()).collect();
assert_eq!(format!("{:016x}", fnv1a(&bytes)), conf["params_hash"].as_str().unwrap(), "h_s through the integer operator set must reproduce the export");
}
#[test]
@ -114,7 +139,7 @@ fn decodes_the_committed_tiny_conformance_vector() {
// → integer g_s) with NO external artifacts — so CI verifies the int_hyper
// capability without the 23 MB pretrained model. Both the JSON and the compact
// binary model must reproduce the committed FNV-1a, byte-for-byte. Never skips.
use wai::int_hyper_synth::{decode_conditional, CondBitstream, CondEntropy};
use wai::int_hyper_synth::{decode_conditional_checked, CondBitstream, CondEntropy};
let d = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/vectors/int_hyper");
let conf = load_json(&d.join("tiny.json")).expect("committed tiny.json");
let ent = CondEntropy::from_value(&load_json(&d.join("tiny_entropy.json")).unwrap());
@ -124,13 +149,13 @@ fn decodes_the_committed_tiny_conformance_vector() {
let g_j = IntModel::from_json(&load_json(&d.join("tiny_g_s.json")).unwrap());
let h_j = IntModel::from_json(&load_json(&d.join("tiny_h_s.json")).unwrap());
let (rgb, h, w) = decode_conditional(&g_j, &h_j, &ent, &bs).expect("tiny json decode");
let (rgb, h, w) = decode_conditional_checked(&g_j, &h_j, &ent, &bs.stream()).expect("tiny json decode");
assert_eq!((h, w), (oh, ow));
assert_eq!(format!("{:016x}", fnv1a(&rgb)), want, "tiny int_hyper (json model) must reproduce the committed hash");
let g_b = IntModel::from_bytes(&std::fs::read(d.join("tiny_g_s.bin")).unwrap()).expect("tiny g_s bin");
let h_b = IntModel::from_bytes(&std::fs::read(d.join("tiny_h_s.bin")).unwrap()).expect("tiny h_s bin");
let (rgb2, _, _) = decode_conditional(&g_b, &h_b, &ent, &bs).expect("tiny binary decode");
let (rgb2, _, _) = decode_conditional_checked(&g_b, &h_b, &ent, &bs.stream()).expect("tiny binary decode");
assert_eq!(format!("{:016x}", fnv1a(&rgb2)), want, "tiny int_hyper (binary model) must reproduce the committed hash");
println!("tiny int_hyper conformance: {oh}x{ow}, {} z + {} y rANS bytes, fnv1a {want}", bs.z_rans.len(), bs.rans.len());
}
@ -142,7 +167,7 @@ fn rust_decodes_the_full_conditional_bitstream_to_the_export_image() {
// σ → integer bucket → rANS-decode r → y_hat = r·2^afrac + μ → integer g_s —
// to the SAME image the export tool produced, byte-for-byte (fully integer, so
// identical on every machine). This is the integer-entropy contract assembled.
use wai::int_hyper_synth::{decode_conditional, CondBitstream, CondEntropy};
use wai::int_hyper_synth::{decode_conditional_checked, CondBitstream, CondEntropy};
let d = dir();
let test = "int_hyper_synth::rust_decodes_the_full_conditional_bitstream_to_the_export_image";
let Some([conf, g_s, h_s, ent, bs]) = inputs(
@ -155,7 +180,7 @@ fn rust_decodes_the_full_conditional_bitstream_to_the_export_image() {
let (g_s, h_s) = (IntModel::from_json(&g_s), IntModel::from_json(&h_s));
let ent = CondEntropy::from_value(&ent);
let bs = CondBitstream::from_value(&bs);
let (rgb, oh, ow) = decode_conditional(&g_s, &h_s, &ent, &bs).expect("conditional decode");
let (rgb, oh, ow) = decode_conditional_checked(&g_s, &h_s, &ent, &bs.stream()).expect("conditional decode");
assert_eq!((oh, ow), (conf["out_h"].as_u64().unwrap() as usize, conf["out_w"].as_u64().unwrap() as usize));
let got = format!("{:016x}", fnv1a(&rgb));
println!(
@ -174,13 +199,55 @@ fn rust_decodes_the_full_conditional_bitstream_to_the_export_image() {
let (g_bin, h_bin) = (std::fs::read(&gb).unwrap(), std::fs::read(&hb).unwrap());
let g_s2 = IntModel::from_bytes(&g_bin).expect("g_s binary");
let h_s2 = IntModel::from_bytes(&h_bin).expect("h_s binary");
let (rgb2, _, _) = decode_conditional(&g_s2, &h_s2, &ent, &bs).expect("binary conditional decode");
let (rgb2, _, _) = decode_conditional_checked(&g_s2, &h_s2, &ent, &bs.stream()).expect("binary conditional decode");
assert_eq!(format!("{:016x}", fnv1a(&rgb2)), conf["rgb_hash"].as_str().unwrap(),
"the compact binary model must decode byte-identically to the JSON model");
println!("compact binary model ({} + {} bytes) decodes byte-identically", g_bin.len(), h_bin.len());
// Through the integer operator set (`int_hyper_synth::ops`: both networks lowered by
// whm1/1, loaded under acc64 and run with their check ranges): the same image.
let (rgb3, ..) = wai::int_hyper_synth::ops::decode(&g_s2, &h_s2, &ent, &bs.stream()).expect("the decode through the lowering");
assert_eq!(format!("{:016x}", fnv1a(&rgb3)), conf["rgb_hash"].as_str().unwrap(),
"the decode through the integer operator set must give the export image");
println!("the decode through the integer operator set gives the same image");
}
}
/// Gate 1 of the migration to the integer operator set (`extensions/integer-operators.md`
/// §13.2), for the pretrained set, which is too large to commit: the regenerated set is
/// the one whose certificate is committed by its digest
/// (`ops-conformance/certificates/external.json`), and this crate's lowering of the
/// models it parses (`int_hyper_synth::ops`) writes the two graphs that certificate
/// records, `h_s` then `g_s`, which load under `acc64` and not `acc32`.
/// (`wai-sota-vectors.yml` checks the certificate itself against the set, with both
/// certifiers.)
#[test]
fn the_pretrained_set_lowers_to_its_certificates_graphs() {
let d = dir();
let test = "int_hyper_synth::the_pretrained_set_lowers_to_its_certificates_graphs";
let names = ["cond_entropy.json", "int_synth_sota_model.bin", "int_hs_model.bin"];
let missing: Vec<&str> = names.iter().copied().filter(|n| !d.join(n).exists()).collect();
if !missing.is_empty() {
d.absent(test, &format!("{} not in {}", missing.join(", "), d.display()), PRODUCE);
return;
}
let [e, g, h] = names.map(|n| std::fs::read(d.join(n)).unwrap());
let digest = wai_ops::generate::sets::set_digest(&[("entropy.json", e), ("g_s.whm", g.clone()), ("h_s.whm", h.clone())]);
let certs = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../ops-conformance/certificates");
let ext = load_json(&certs.join("external.json")).expect("the external certificates' listing");
let listed = ext["sets"].as_array().unwrap().iter().find(|s| s["capability"] == "wai.neural.int_hyper").expect("the pretrained set is listed");
assert_eq!(listed["set_sha256"].as_str(), Some(digest.as_str()), "the regenerated set is the one whose certificate is committed");
let cert = load_json(&certs.join(format!("wai.neural.int_hyper/{digest}.json"))).expect("its certificate");
let (g_s, h_s) = (IntModel::from_bytes(&g).unwrap(), IntModel::from_bytes(&h).unwrap());
let (hg, gg) = wai::int_hyper_synth::ops::lower(&g_s, &h_s).expect("the pretrained set lowers");
for (k, (role, graph)) in [("h_s", &hg), ("g_s", &gg)].into_iter().enumerate() {
let sha: String = wai_ops::sha256::sha256(&graph.to_bytes()).iter().map(|x| format!("{x:02x}")).collect();
assert_eq!((cert["graphs"][k]["name"].as_str(), cert["graphs"][k]["graph_sha256"].as_str()), (Some(role), Some(sha.as_str())), "{role}");
}
let nets = wai::int_hyper_synth::ops::load(&g_s, &h_s).expect("the pretrained set loads under acc64");
assert!(!nets.h_s.meets(wai_ops::Profile::Acc32) && !nets.g_s.meets(wai_ops::Profile::Acc32), "the pretrained set meets acc64, not acc32");
println!("pretrained set {digest}: lowers to its certificate's graphs and loads under acc64");
}
/// The largest quality cost the integer decode may have against the float codec on
/// the same image. Measured costs: 0.19 dB for the full conditional decode on the
/// pinned photograph (arm64 macOS machine); 0.62 dB for the integer synthesis alone
@ -197,7 +264,7 @@ const MIN_FLOAT_PSNR_DB: f64 = 30.0;
fn integer_decode_keeps_the_float_codec_quality() {
// Measured by this decoder on its own output, against the exact source pixels
// the export tool coded — not the PSNR the tool reports about itself.
use wai::int_hyper_synth::{decode_conditional, CondBitstream, CondEntropy};
use wai::int_hyper_synth::{decode_conditional_checked, CondBitstream, CondEntropy};
let d = dir();
let test = "int_hyper_synth::integer_decode_keeps_the_float_codec_quality";
let Some([conf, g_s, h_s, ent, bs]) = inputs(
@ -219,7 +286,7 @@ fn integer_decode_keeps_the_float_codec_quality() {
let source = std::fs::read(&src_path).unwrap();
let (g_s, h_s) = (IntModel::from_json(&g_s), IntModel::from_json(&h_s));
let (rgb, h, w) = decode_conditional(&g_s, &h_s, &CondEntropy::from_value(&ent), &CondBitstream::from_value(&bs))
let (rgb, h, w) = decode_conditional_checked(&g_s, &h_s, &CondEntropy::from_value(&ent), &CondBitstream::from_value(&bs).stream())
.expect("conditional decode");
assert_eq!(source.len(), rgb.len(), "source.rgb is not the {h}x{w} RGB image the bitstream decodes to");

8
wai/wai-rs/tests/int_kernels_vectors.rs
View file

@ -45,14 +45,14 @@ fn same_under_every_policy<T: PartialEq + std::fmt::Debug>(what: &str, f: impl F
#[test]
fn int_hyper_tiny_conditional_codec() {
use wai::int_hyper_synth::{decode_conditional, fnv1a, CondBitstream, CondEntropy, IntModel};
use wai::int_hyper_synth::{decode_conditional_checked, fnv1a, CondBitstream, CondEntropy, IntModel};
let d = vroot().join("int_hyper");
let want = jload(&d.join("tiny.json"))["rgb_hash"].as_str().unwrap().to_string();
let ent = CondEntropy::from_value(&jload(&d.join("tiny_entropy.json")));
let bs = CondBitstream::from_value(&jload(&d.join("tiny_bitstream.json")));
let g_s = IntModel::from_bytes(&std::fs::read(d.join("tiny_g_s.bin")).unwrap()).unwrap();
let h_s = IntModel::from_bytes(&std::fs::read(d.join("tiny_h_s.bin")).unwrap()).unwrap();
let (rgb, _, _) = same_under_every_policy("int_hyper tiny", || decode_conditional(&g_s, &h_s, &ent, &bs).unwrap());
let (rgb, _, _) = same_under_every_policy("int_hyper tiny", || decode_conditional_checked(&g_s, &h_s, &ent, &bs.stream()).unwrap());
assert_eq!(format!("{:016x}", fnv1a(&rgb)), want);
sha256_sidecar::one(&d, "tiny", RGB8, &rgb);
}
@ -138,7 +138,7 @@ fn int_world_clips() {
#[test]
fn int_video_neural_keyframe_and_real_footage() {
use wai::int_hyper_synth::{decode_conditional, CondBitstream, CondEntropy, IntModel};
use wai::int_hyper_synth::{decode_conditional_checked, CondBitstream, CondEntropy, IntModel};
use wai::int_video::{decode, decode_pframes, fnv1a};
let (ih, iv) = (vroot().join("int_hyper"), vroot().join("int_video"));
let g_s = IntModel::from_json(&jload(&ih.join("tiny_g_s.json")));
@ -147,7 +147,7 @@ fn int_video_neural_keyframe_and_real_footage() {
let bs = CondBitstream::from_value(&jload(&ih.join("tiny_bitstream.json")));
let vbs = jload(&iv.join("video_nk_bitstream.json"));
let frames = same_under_every_policy("neural-keyframe video", || {
let (f0, h, w) = decode_conditional(&g_s, &h_s, &ent, &bs).unwrap();
let (f0, h, w) = decode_conditional_checked(&g_s, &h_s, &ent, &bs.stream()).unwrap();
decode_pframes(&f0, h, w, &vbs).unwrap().frames
});
let want = jload(&iv.join("video_nk.json"));

383
wai/wai-rs/tests/int_ops_differential.rs
View file

@ -7,7 +7,8 @@
//! 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`).
//! (`whm1/1`), `wai.audio.int_codec` (`wam1/1`) and `wai.neural.int_hyper` (`whm1/1`, its
//! two networks; the keyframe of `wai.video.int_hyper` is the same decode).
//!
//! 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
@ -982,3 +983,383 @@ fn int_audio_planned_breaks_are_new_refusals() {
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));
}
// ------------------------------------------------------------------ int_hyper (M4)
/// A conditional codec's entropy tables (`entropy.json`) and a `WIH1` payload's typed
/// fields, owned.
#[derive(Clone, Debug)]
struct Cond {
edges: Vec<i64>,
cdfs: Vec<Vec<i32>>,
cdf_lengths: Vec<usize>,
offsets: Vec<i32>,
z_cdfs: Vec<Vec<i32>>,
z_cdf_lengths: Vec<usize>,
z_offsets: Vec<i32>,
z_median_q: Vec<i64>,
m: usize,
sh: usize,
sw: usize,
n_z: usize,
zh: usize,
zw: usize,
z_rans: Vec<u8>,
rans: Vec<u8>,
}
impl Cond {
fn ent(&self) -> wai::int_hyper_synth::CondEntropy {
wai::int_hyper_synth::CondEntropy {
edges: self.edges.clone(), cdfs: self.cdfs.clone(), cdf_lengths: self.cdf_lengths.clone(), offsets: self.offsets.clone(),
z_cdfs: self.z_cdfs.clone(), z_cdf_lengths: self.z_cdf_lengths.clone(), z_offsets: self.z_offsets.clone(), z_median_q: self.z_median_q.clone(),
}
}
fn st(&self) -> wai::int_hyper_synth::CondStream<'_> {
wai::int_hyper_synth::CondStream { m: self.m, sh: self.sh, sw: self.sw, n_z: self.n_z, zh: self.zh, zw: self.zw, z_rans: &self.z_rans, rans: &self.rans }
}
}
/// The planned break of integer-operators §15.3 that a conditional codec's set falls
/// under, if any, read from its two networks' bytes alone: either network is refused at
/// load by its lowering (`h_s` by `whm1/1`, `g_s` by `whm1/1` with the RGB8 step, each
/// loaded under `acc64`) for a rule the direct decode does not make (see
/// [`int_splat_planned_break`]). The entropy tables are not lowered, and no payload field
/// changes meaning.
fn int_hyper_planned_break(g_s: &[u8], h_s: &[u8]) -> Option<&'static str> {
let (h, g) = (wai_ops::lower::parse_whm1(h_s).ok()?, wai_ops::lower::parse_whm1(g_s).ok()?);
load_break(&h, wai_ops::lower::whm1(h_s, false)).or_else(|| load_break(&g, wai_ops::lower::whm1(g_s, true)))
}
/// A random CDF table of 3 to 6 entries.
fn cdf_table(r: &mut Rng) -> Vec<i32> {
let len = 3 + r.below(4) as usize;
let mut cuts: Vec<i32> = (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);
t
}
/// A last layer of a random network: a convolution (stride 1) or a transposed convolution
/// from `c` channels to `out`, with `splat_layers`' parameter draws; one in 32 a ReLU.
fn last_layer(r: &mut Rng, c: u32, out: u32) -> Whm {
if r.below(32) == 0 {
return Whm::Relu;
}
let k = 1 + r.below(3) as u32;
let p = if r.below(16) == 0 { r.below(k as u64 + 2) } else { (k / 2) as u64 } as u32;
let wmax = if r.below(8) == 0 { 32767 } else { [16, 1024, 4096][r.below(3) as usize] };
let w: Vec<i16> = (0..c * out * k * k).map(|_| if r.below(4) == 0 { 0 } else { r.signed(wmax) as i16 }).collect();
let m: Vec<i64> = (0..out).map(|_| if r.below(32) == 0 { r.signed(1 << 40) } else { r.signed(1 << 12) }).collect();
let b: Vec<i64> = (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();
if r.below(2) == 0 {
Whm::Conv { inc: c, out, k, s: 1, p, w, m, b }
} else {
let s = 1 + r.below(2) as u32;
Whm::ConvT { inc: c, out, k, s, p: p.min(k.saturating_sub(1)), op: r.below(s as u64) as u32, w, m, b }
}
}
/// `n` values coded against one table each, with `RansEncoder`: `tables[i]` and its length
/// and offset for value `i`. `None` where a value is not codable against its table.
fn rans_of(values: &[(i32, &[i32], usize, i32)]) -> Option<Vec<u8>> {
let mut enc = wai::rans::RansEncoder::new();
for &(v, t, len, off) in values {
if !enc.encode(v, t, len, off) {
return None;
}
}
Some(enc.flush())
}
/// A stream damaged one case in sixteen (cut short), one in 32 (a word appended) or one in
/// 64 (a byte changed).
fn damage(r: &mut Rng, s: &mut Vec<u8>) {
match r.below(64) {
0..=3 => {
let cut = (4 * (1 + r.below(3)) as usize).min(s.len());
s.truncate(s.len() - cut);
}
4 | 5 => s.extend([0, 0, 0, 0]),
6 if !s.is_empty() => {
let i = r.below(s.len() as u64) as usize;
s[i] ^= 1 << r.below(8);
}
_ => {}
}
}
/// A random `wai.neural.int_hyper` set and `WIH1` payload. The set: `afrac` 12 mostly (one
/// in sixteen anywhere in 1..=30, one in sixteen at 0 or 31, and `h_s`'s not `g_s`'s one
/// in 32); each network's header `gf` or `shift` of 128 or more one set in 32 each; `h_s`
/// from `n_z` (1 to 3) channels, up to one of `splat_layers`' layers then a last layer to
/// `2·m` channels (one in sixteen to another count); `g_s` from `m` (1 to 3) channels, up
/// to two layers then a last layer to three channels (one in sixteen not); the tables: 1
/// to 4 ascending scale edges (one set in 32 none), a latent table, length and offset per
/// edge, and a hyper-latent table, offset and median per hyper-latent channel (one set in
/// 32 a channel short), one length in 32 past its row; medians mostly small, one in eight
/// at or past the 2^47 bound or near i64's ends. The payload: a hyper-latent of side 2 or
/// 3 (one in sixteen not square), its symbols mostly inside their table and some escaped,
/// one case in four with a whole channel at `h_s`'s input `check` range or one past it;
/// the latent's shape `h_s`'s output (one case in eight off by one), each symbol coded
/// against the table the direct decode's `σ` buckets it to (where `h_s` runs directly);
/// each stream damaged as [`damage`] says.
fn int_hyper_case(r: &mut Rng) -> (Vec<u8>, Vec<u8>, Cond) {
let afrac = match r.below(16) {
0 => 1 + r.below(30) as u32,
1 => [0, 31][r.below(2) as usize],
_ => 12,
};
let h_afrac = if r.below(32) == 0 { afrac ^ 1 } else { afrac };
let header = |r: &mut Rng| {
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 };
(gf, shift)
};
let (n_z, m) = (1 + r.below(3) as u32, 1 + r.below(3) as u32);
let (n, relu_only) = (r.below(2) as usize, r.below(32) == 0);
let (mut hl, c) = splat_layers(r, n_z, n, relu_only);
let out = if r.below(16) == 0 { 1 + r.below(6) as u32 } else { 2 * m };
hl.push(last_layer(r, c, out));
let (gf, shift) = header(r);
let h_s = whm1_bytes(h_afrac, gf, shift, &hl);
let (n, relu_only) = (r.below(3) as usize, r.below(32) == 0);
let (mut gl, c) = splat_layers(r, m, n, relu_only);
let out = if r.below(16) == 0 { 1 + r.below(4) as u32 } else { 3 };
gl.push(last_layer(r, c, out));
let (gf, shift) = header(r);
let g_s = whm1_bytes(afrac, gf, shift, &gl);
let a = afrac.min(30);
let n_edges = if r.below(32) == 0 { 0 } else { 1 + r.below(4) as usize };
let mut edges = Vec::new();
let mut e = 1 + r.below(1 << (a + 2)) as i64;
for _ in 0..n_edges {
edges.push(e);
e += 1 + r.below(1 << (a + 2)) as i64;
}
let cdfs: Vec<Vec<i32>> = (0..n_edges).map(|_| cdf_table(r)).collect();
let mut cdf_lengths: Vec<usize> = cdfs.iter().map(Vec::len).collect();
let offsets: Vec<i32> = (0..n_edges).map(|_| r.signed(3) as i32).collect();
let k = if r.below(32) == 0 { n_z as usize - 1 } else { n_z as usize };
let z_cdfs: Vec<Vec<i32>> = (0..k).map(|_| cdf_table(r)).collect();
let mut z_cdf_lengths: Vec<usize> = z_cdfs.iter().map(Vec::len).collect();
if r.below(32) == 0 {
match (n_edges, k) {
(1.., _) if r.below(2) == 0 => cdf_lengths[0] += 1,
(_, 1..) => z_cdf_lengths[0] += 1,
_ => {}
}
}
let z_offsets: Vec<i32> = (0..k).map(|_| r.signed(3) as i32).collect();
let mut z_median_q: Vec<i64> = (0..k)
.map(|_| match r.below(8) {
0 => [1i64 << 47, -(1 << 47), 1 << 60, -(1 << 60), i64::MAX, i64::MIN][r.below(6) as usize],
_ => r.signed(1 << (a + 4)),
})
.collect();
let zh = 2 + r.below(2) as usize;
let zw = if r.below(16) == 0 { 1 + r.below(3) as usize } else { zh };
let (nz, plane) = (n_z as usize, zh * zw);
// The hyper-latent's symbol values, per channel against its own table (channel 0's
// where a channel has none).
let zt = |ch: usize| ch.min(k.saturating_sub(1));
let mut zv: Vec<i32> = Vec::with_capacity(nz * plane);
for ch in 0..nz {
let (off, span) = if k == 0 { (0, 1) } else { (z_offsets[zt(ch)] as i64, z_cdfs[zt(ch)].len() as i64 - 2) };
for _ in 0..plane {
zv.push(match r.below(32) {
0 => r.signed(1 << 20),
1..=4 => off + span + r.below(40) as i64 - 20,
_ => off + r.below(span.max(1) as u64) as i64,
} as i32);
}
}
// One case in four: a whole hyper-latent channel at h_s's input check range, or one
// past it, where h_s lowers.
if r.below(4) == 0 && k > 0 {
if let Ok(g) = wai_ops::lower::whm1(&h_s, false) {
let x0 = g.tensors[0].hi as i128;
let ch = r.below(k as u64) as usize;
let v = z_offsets[ch];
let target = [x0, x0 + 1, -x0, -x0 - 1][r.below(4) as usize];
let med = target - ((v as i128) << a);
if let Ok(med) = i64::try_from(med) {
z_median_q[ch] = med;
zv[ch * plane..(ch + 1) * plane].iter_mut().for_each(|x| *x = v);
}
}
}
let z_rans = if k == 0 {
vec![0, 0, 0, 128, 0, 0, 0, 0]
} else {
let vals: Vec<(i32, &[i32], usize, i32)> =
(0..nz * plane).map(|i| { let t = zt(i / plane); (zv[i], z_cdfs[t].as_slice(), z_cdf_lengths[t], z_offsets[t]) }).collect();
rans_of(&vals).unwrap_or_else(|| vec![0, 0, 0, 128, 0, 0, 0, 0])
};
// The latent: h_s's output shape, and each symbol's table from the direct decode's σ.
let hm = IntModel::from_bytes(&h_s).expect("the generator writes WHM1 that parses");
let side = |n: usize| splat_side(&hl, n as i64).unwrap_or(1) as usize;
let (mut sh, sw) = (side(zh), side(zw));
if r.below(8) == 0 {
sh += 1;
}
let mu = m as usize;
let z_hat: Option<Vec<i64>> = (0..nz * plane)
.map(|i| {
let med = *z_median_q.get(zt(i / plane))?;
(zv[i] as i64).checked_mul(1 << a)?.checked_add(med)
})
.collect();
let sigma: Option<Vec<i64>> = z_hat.and_then(|z| {
let (o, oc, oh, ow) = wai::int_hyper_synth::run_net_checked(&hm, &z, nz, zh, zw).ok()?;
((oc, oh, ow) == (2 * mu, sh, sw)).then(|| o[..mu * sh * sw].to_vec())
});
let ny = mu * sh * sw;
let rans = match (&sigma, n_edges) {
(Some(s), 1..) => {
let mut vals = Vec::with_capacity(ny);
for p in 0..ny {
let b = wai::int_entropy::det_build_index(s[p], &edges);
let span = cdfs[b].len() as i64 - 2;
let v = match r.below(32) {
0 => r.signed(1 << 20),
1..=3 => offsets[b] as i64 + span + r.below(40) as i64 - 20,
_ => offsets[b] as i64 + r.below(span.max(1) as u64) as i64,
} as i32;
vals.push((v, cdfs[b].as_slice(), cdf_lengths[b], offsets[b]));
}
rans_of(&vals).unwrap_or_else(|| (0..8 + 4 * r.below(8)).map(|_| r.below(256) as u8).collect())
}
_ => (0..8 + 4 * r.below(8)).map(|_| r.below(256) as u8).collect(),
};
let mut cond = Cond {
edges, cdfs, cdf_lengths, offsets, z_cdfs, z_cdf_lengths, z_offsets, z_median_q,
m: mu, sh, sw, n_z: nz, zh, zw, z_rans, rans,
};
damage(r, &mut cond.z_rans);
damage(r, &mut cond.rans);
(g_s, h_s, cond)
}
/// The direct conditional decode, a panic (an overflow-checked build's arithmetic
/// overflow) caught.
fn hyper_direct(g: &IntModel, h: &IntModel, c: &Cond) -> Option<Result<Vec<u8>, RefusalKind>> {
let e = c.ent();
catch_unwind(AssertUnwindSafe(|| wai::int_hyper_synth::legacy::decode_conditional_checked(g, h, &e, &c.st()).map(|d| d.0).map_err(|r| r.kind))).ok()
}
/// Gate 4 of the migration, for `wai.neural.int_hyper` (and so the keyframe of
/// `wai.video.int_hyper`): the decode through the operator set against the direct decode
/// on random sets and payloads. Outside the planned breaks both decode to the same image
/// 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_hyper_random_sets_and_payloads_decode_alike() {
let (mut alike, mut refused, mut breaks) = (0u64, std::collections::BTreeMap::<String, u64>::new(), std::collections::BTreeMap::<&str, u64>::new());
let n = cases("int_hyper", 3000, |r, seed| {
let (g_s, h_s, c) = int_hyper_case(r);
let why = || format!("seed {seed}\npayload {c:?}");
let (g, h) = (IntModel::from_bytes(&g_s).unwrap(), IntModel::from_bytes(&h_s).unwrap());
let new = wai::int_hyper_synth::ops::decode(&g, &h, &c.ent(), &c.st()).map(|d| d.0).map_err(|r| r.kind);
match int_hyper_planned_break(&g_s, &h_s) {
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, hyper_direct(&g, &h, &c)) {
(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_hyper: {alike} decoded alike, refused by both {refused:?}, planned breaks {breaks:?}");
// Two networks and two streams compound the generator's refusal rates, so fewer cases
// decode than for one network; an eighth is still hundreds per push.
assert!(alike * 8 > 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 || ["StreamExhausted", "StreamNoncanonical", "GeometryMismatch", "ModelMismatch"].iter().all(|k| refused.contains_key(*k)), "a refusal kind was not reached: {refused:?}");
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.neural.int_hyper`, each
/// on a set and payload 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 either stream).
#[test]
fn int_hyper_planned_breaks_are_new_refusals() {
// A payload of two hyper-latent channels of side 2, every symbol the first of the
// table `t` (offset 2), so every hyper-latent is 2·2^12; a latent of one channel of
// side 2, every symbol the first of the one scale edge's table (offset 2).
let payload = |t: Vec<i32>| {
let z: Vec<(i32, &[i32], usize, i32)> = (0..8).map(|_| (2, t.as_slice(), t.len(), 2)).collect();
let y = [0, 1 << 15, 1 << 16];
let yv: Vec<(i32, &[i32], usize, i32)> = (0..4).map(|_| (2, y.as_slice(), 3, 2)).collect();
Cond {
edges: vec![1 << 12], cdfs: vec![y.to_vec()], cdf_lengths: vec![3], offsets: vec![2],
z_cdfs: vec![t.clone(); 2], z_cdf_lengths: vec![t.len(); 2], z_offsets: vec![2; 2], z_median_q: vec![0; 2],
m: 1, sh: 2, sw: 2, n_z: 2, zh: 2, zw: 2, z_rans: rans_of(&z).unwrap(), rans: rans_of(&yv).unwrap(),
}
};
let cl = payload(vec![0, 1 << 15, 1 << 16]);
// h_s: σ and μ are the two hyper-latent channels; g_s: the latent to each of three channels.
let ident = Whm::Conv { inc: 2, out: 2, k: 1, s: 1, p: 0, w: vec![1, 0, 0, 1], m: vec![1; 2], b: vec![0; 2] };
let to3 = |w: i16, b: i64| Whm::ConvT { inc: 1, out: 3, k: 1, s: 1, p: 0, op: 0, w: vec![w; 3], m: vec![1; 3], b: vec![b; 3] };
let (h0, g0) = (whm1_bytes(12, 16, 0, &[ident.clone()]), whm1_bytes(12, 16, 0, &[to3(1, 0)]));
let decode = |g_s: &[u8], h_s: &[u8], c: &Cond| {
let (g, h) = (IntModel::from_bytes(g_s).unwrap(), IntModel::from_bytes(h_s).unwrap());
(wai::int_hyper_synth::ops::decode(&g, &h, &c.ent(), &c.st()).map(|d| d.0).map_err(|r| r.kind), hyper_direct(&g, &h, c))
};
// The base decodes alike: every latent 2·2^12 + 2·2^12, clamped to 2^12: 255.
assert_eq!(int_hyper_planned_break(&g0, &h0), None);
let (new, old) = decode(&g0, &h0, &cl);
assert_eq!((new, old), (Ok(vec![255; 12]), Some(Ok(vec![255; 12]))));
let expect = |b: &str, g_s: &[u8], h_s: &[u8], c: &Cond, direct: Option<RefusalKind>| {
assert_eq!(int_hyper_planned_break(g_s, h_s), Some(b), "{b}: the classifier");
let (new, old) = decode(g_s, h_s, c);
match direct {
None => assert!(matches!(old, Some(Ok(_))), "{b}: the direct decode decodes it: {old:?}"),
Some(k) => assert_eq!(old, Some(Err(k)), "{b}: the direct decode's code"),
}
assert_eq!(new, Err(RefusalKind::ModelMismatch), "{b}: the decode through the operator set refuses it");
};
// No layer of tag 0, 1 or 2 in h_s: σ and μ are the hyper-latent's channels directly.
expect("no layer of tag 0, 1 or 2", &g0, &whm1_bytes(12, 16, 0, &[Whm::Relu]), &cl, None);
// A header shift no layer of h_s uses (only an inverse GDN that leaves the values as
// they are), and a gf no layer of g_s 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", &g0, &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, &[to3(1, 0)]), &h0, &cl, None);
// A bias at the end of i64 in g_s: refused directly as g_s runs (field_range), after
// both streams; at load through the lowering. In h_s, with a hyper-latent stream that
// runs out: refused directly for the stream, through the lowering for the set.
expect("a bias of magnitude 2^63 - 1 or more", &whm1_bytes(12, 16, 0, &[to3(1, i64::MAX)]), &h0, &cl, Some(RefusalKind::FieldRange));
let mut short = payload(vec![0, 1, 1 << 16]);
short.z_rans.truncate(8);
let h_min = Whm::Conv { inc: 2, out: 2, k: 1, s: 1, p: 0, w: vec![1, 0, 0, 1], m: vec![1; 2], b: vec![i64::MIN; 2] };
expect("a bias of magnitude 2^63 - 1 or more", &g0, &whm1_bytes(12, 16, 0, &[h_min]), &short, Some(RefusalKind::StreamExhausted));
// A check range its producer cannot meet, in g_s: its first layer gives every value
// 2^30 (weight 0, bias 2^30), into a layer of weight 2^11 under the multiplier 2^40
// and shift 0, whose range is ±4095. Refused directly at that layer (field_range).
let first = Whm::ConvT { inc: 1, out: 1, k: 1, s: 1, p: 0, op: 0, w: vec![0], m: vec![1], b: vec![1 << 30] };
let tight = Whm::ConvT { inc: 1, out: 3, k: 1, s: 1, p: 0, op: 0, w: vec![1 << 11; 3], m: vec![1 << 40; 3], b: vec![0; 3] };
expect("a check range its producer cannot meet", &whm1_bytes(12, 16, 0, &[first, tight]), &h0, &cl, Some(RefusalKind::FieldRange));
}

64
wai/wai-rs/tests/int_payload_cost.rs
View file

@ -10,6 +10,14 @@
//! 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.
//!
//! These checks run with the store of loaded sets (`int_set_cache`) off, so each decode
//! lowers and loads its set, as it did before there was a store. Then with the store on:
//! a decode that loads its set holds within the same bound, and leaves the set in the
//! store, which counts at least the bytes it holds for it (so its budget bounds what it
//! holds); a decode that finds its set held holds within the same bound, less than the
//! decode that loaded it, and leaves nothing more behind. A set kept in the store is the
//! set's, the same for every payload, and outside every decode's declared working bytes.
//!
//! The clip a few kilobytes can code: sixteen frames of 4729 × 4729 whose stream is
//! one symbol of frequency 65535/65536 over and over, about 3 KB, decoding to a clip at
//! the 2^30-byte cap. It is decoded in release builds only (CI's release step); a debug
@ -26,6 +34,7 @@ use std::sync::atomic::{AtomicUsize, Ordering::Relaxed};
use wai::int_kernels::{with_policy, Policy};
use wai::int_payload::{declared_cost, decode_int, resolve, wire, IntParams};
use wai::int_set_cache::LoadedSets;
use wai::pin::HeldBytes;
/// The system allocator, tracking the bytes live and the most live at once.
@ -110,6 +119,8 @@ fn constant_clip(n_frames: u32, h: u32, w: u32) -> Vec<u8> {
#[test]
fn the_declared_cost_bounds_what_the_decode_holds() {
use wai::codecs as c;
// The store off: every decode lowers and loads its set.
LoadedSets::global().set_budget(0);
let hyp: BTreeMap<String, Vec<u8>> = [("entropy.json", "int_hyper/tiny_entropy.json"), ("g_s.whm", "int_hyper/tiny_g_s.bin"), ("h_s.whm", "int_hyper/tiny_h_s.bin")]
.into_iter().map(|(k, v)| (k.to_string(), read(v))).collect();
let hyp_view: BTreeMap<String, &[u8]> = hyp.iter().map(|(k, v)| (k.clone(), v.as_slice())).collect();
@ -121,12 +132,16 @@ fn the_declared_cost_bounds_what_the_decode_holds() {
// The parameter set as a decode parses it, which the declared cost leaves out.
let set_bytes = |cap: &str| -> u64 {
let (r, peak) = peak_during(|| match cap {
// The conditional codec's set: its tables, and its two networks parsed, then
// lowered and loaded (`int_hyper_synth::ops::load`).
c::CAP_NEURAL_INT_HYPER | c::CAP_VIDEO_INT_HYPER => {
let e: serde_json::Value = serde_json::from_slice(&hyp["entropy.json"]).unwrap();
let t = wai::int_hyper_synth::CondEntropy::try_from_value(&e).unwrap();
drop(e);
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<dyn std::any::Any>
let (g, h) = (wai::int_hyper_synth::IntModel::from_bytes(&hyp["g_s.whm"]).unwrap(),
wai::int_hyper_synth::IntModel::from_bytes(&hyp["h_s.whm"]).unwrap());
let n = wai::int_hyper_synth::ops::load(&g, &h).unwrap();
Box::new((g, h, t, n)) as Box<dyn std::any::Any>
}
c::CAP_AUDIO_INT_CODEC => {
let m = wai::int_audio::AudioModel::from_bytes(&model).unwrap();
@ -161,8 +176,10 @@ fn the_declared_cost_bounds_what_the_decode_holds() {
assert!(big.len() < 4096, "the clip at the cap is coded in {} bytes", big.len());
cases.push(("16 frames of 4729 × 4729 in about 3 KB", c::CAP_VIDEO_INT_MOTION, big, &none));
}
let mut costs = BTreeMap::new();
for (what, cap, payload, params) in cases {
let cost = declared_cost(cap, &payload, params).unwrap_or_else(|e| panic!("{what}: {e}"));
costs.insert(what, cost);
// The header-only reading agrees where the header fixes the output.
let size = wai::container::declared_size(cap, &payload).unwrap().unwrap();
assert_eq!(size.symbols, cost.symbols, "{what}: symbols");
@ -188,6 +205,33 @@ fn the_declared_cost_bounds_what_the_decode_holds() {
cost.working_bytes, cost.output_bytes, cost.symbols, cost.macs);
}
// The store on, and empty. For each payload that decodes against a set: the decode
// that loads it, then the decode that finds it held.
let store = LoadedSets::global();
store.set_budget(LoadedSets::DEFAULT_BUDGET);
for (what, cap, rel, params) in [
("video_nk", c::CAP_VIDEO_INT_HYPER, "int_video/video_nk_bitstream.wiv1", &hyp_p),
("tiny", c::CAP_NEURAL_INT_HYPER, "int_hyper/tiny_bitstream.wih1", &hyp_p),
("audio", c::CAP_AUDIO_INT_CODEC, "int_audio/audio_bitstream.wia1", &aud),
("splat", c::CAP_SPLAT_INT_CODEC, "int_splat/splat.wis1", &spl),
] {
let (payload, cost, set) = (read(rel), costs[what], set_bytes(cap));
let counted0 = store.stats().resident_bytes;
let base = CUR.load(Relaxed);
let (r, miss) = with_policy(Policy::reference(), || peak_during(|| decode_int(cap, &payload, params)));
drop(r.unwrap_or_else(|e| panic!("{what}: {e}")));
let held = (CUR.load(Relaxed) - base) as u64;
let counted = store.stats().resident_bytes - counted0;
assert!(miss <= cost.working_bytes + set, "{what}, its set loaded and kept: held {miss}, over the declared {} and the set's {set}", cost.working_bytes);
assert!(held <= counted && counted - held <= LoadedSets::INDEX_BYTES_PER_ENTRY, "{what}: the store holds {held} bytes for the set and counts {counted}");
let (r, hit) = with_policy(Policy::reference(), || peak_during(|| decode_int(cap, &payload, params)));
drop(r.unwrap_or_else(|e| panic!("{what}: {e}")));
assert_eq!((CUR.load(Relaxed) - base) as u64, held, "{what}: a decode that finds its set held keeps nothing more");
assert!(hit <= cost.working_bytes + set && hit < miss, "{what}, its set held: held {hit}; loading it held {miss}");
println!("{what}, store on: the decode that loads the set held {miss}, the one that finds it held {hit}; the store holds {held} and counts {counted}");
}
store.set_budget(0);
// 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
@ -221,4 +265,20 @@ fn the_declared_cost_bounds_what_the_decode_holds() {
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);
// And for the conditional codec, its two networks loaded beforehand.
let p = read("int_hyper/tiny_bitstream.wih1");
let cost = declared_cost(c::CAP_NEURAL_INT_HYPER, &p, &hyp_p).unwrap();
let k = wire::parse_wih1(&p).unwrap();
let st = wai::int_hyper_synth::CondStream {
m: k.m as usize, sh: k.sh as usize, sw: k.sw as usize, n_z: k.n_z as usize, zh: k.zh as usize, zw: k.zw as usize,
z_rans: k.z_rans, rans: k.rans,
};
let (g, h) = (wai::int_hyper_synth::IntModel::from_bytes(&hyp["g_s.whm"]).unwrap(), wai::int_hyper_synth::IntModel::from_bytes(&hyp["h_s.whm"]).unwrap());
let ent = wai::int_hyper_synth::CondEntropy::try_from_value(&serde_json::from_slice(&hyp["entropy.json"]).unwrap()).unwrap();
let nets = wai::int_hyper_synth::ops::load(&g, &h).unwrap();
let (r, peak) = with_policy(Policy::reference(), || peak_during(|| wai::int_hyper_synth::ops::decode_loaded(&nets, g.afrac, &ent, &st)));
r.unwrap();
assert!(peak >= cost.output_bytes && peak <= cost.working_bytes, "int_hyper, set loaded: held {peak}, declared {}", cost.working_bytes);
println!("int_hyper, its set loaded beforehand: declared {} working bytes; held {peak}", cost.working_bytes);
}

391
wai/wai-rs/tests/int_set_cache.rs Normal file
View file

@ -0,0 +1,391 @@
//! The store of loaded parameter sets (`int_set_cache`) under the decodes that use it,
//! `int_payload::decode_int` and `declared_cost`, for every capability that decodes
//! against a set through the integer operator set.
//!
//! - A second decode through the same set is a hit, by the store's count, and gives the
//! same bytes as a decode with the store off, for every committed vector that decodes
//! against a set: `tiny`, `video_nk`, `audio` and `splat` of the int-payload corpus,
//! and the int_synth corpus.
//! - Two sets whose members have the same names and whose bytes differ in one byte never
//! share an entry.
//! - A set the load refuses is refused again, with the same refusal, while a set one
//! byte away from it is held, and it is never kept: for int_synth, a set the bounds
//! pass refuses (load step 8); for the conditional codec, a set whose lowering leaves
//! a layer no admissible input range (integer-operators §13.2).
//! - Every refusal fixture of the int-payload corpus keeps its code and its detail with
//! its set held, and a payload refused while it is parsed never reaches the store.
//! - Under a small budget the least recently used set goes first, a set past the budget
//! is not kept, and a budget of 0 neither keeps nor looks up.
//! - Threads decoding one set at once get the same bytes, and one entry is kept.
//!
//! The store is one per process, so each test takes a lock and starts from an empty
//! store with its counts at 0.
#![cfg(all(feature = "video_int", feature = "audio_int", feature = "splat_int"))]
use std::collections::BTreeMap;
use std::path::PathBuf;
use std::sync::{Mutex, MutexGuard};
use serde_json::Value;
use wai::codecs as c;
use wai::int_payload::{declared_cost, decode_int, resolve, IntDecoded};
use wai::int_set_cache::{LoadedSets, SetKey, Stats};
const DEFAULT_BUDGET: u64 = LoadedSets::DEFAULT_BUDGET;
use wai::pin::HeldBytes;
static SERIAL: Mutex<()> = Mutex::new(());
/// The store, to this test alone: empty, its counts at 0, its budget `budget`.
fn store(budget: u64) -> MutexGuard<'static, ()> {
let guard = SERIAL.lock().unwrap_or_else(|e| e.into_inner());
let s = LoadedSets::global();
s.set_budget(0);
s.set_budget(budget);
s.reset_stats();
guard
}
fn stats() -> Stats {
LoadedSets::global().stats()
}
fn wai_dir() -> PathBuf {
PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("..")
}
fn read(rel: &str) -> Vec<u8> {
let p = wai_dir().join(rel);
std::fs::read(&p).unwrap_or_else(|e| panic!("read {}: {e}", p.display()))
}
fn index() -> Value {
serde_json::from_slice(&read("int-payload-conformance/vectors.json")).expect("the index parses")
}
/// A set: member name → bytes.
type Set = BTreeMap<String, Vec<u8>>;
fn set_of(entry: &Value) -> Set {
entry["set"].as_object().into_iter().flatten().map(|(k, v)| (k.clone(), read(v.as_str().unwrap()))).collect()
}
fn view(set: &Set) -> BTreeMap<String, &[u8]> {
set.iter().map(|(k, v)| (k.clone(), v.as_slice())).collect()
}
/// A decode's output as bytes, or its refusal's code and detail.
type Out = Result<Vec<u8>, (String, String)>;
fn bytes(d: IntDecoded) -> Vec<u8> {
match d {
IntDecoded::Image { rgb, .. } => rgb,
IntDecoded::Video { frames, .. } => frames.concat(),
IntDecoded::Pcm { samples, .. } => samples.iter().flat_map(|s| s.to_le_bytes()).collect(),
IntDecoded::Attr { bytes, .. } => bytes,
other => panic!("{other:?}"),
}
}
fn decode(cap: &str, payload: &[u8], set: &Set) -> Out {
let v = view(set);
let params = resolve(cap, None, &[HeldBytes::Set { listing: b"", files: &v }], None).expect("an unpinned set resolves");
decode_int(cap, payload, &params).map(bytes).map_err(|e| (e.code().to_string(), e.detail().to_string()))
}
/// The same decode with the store off, then on again with the budget it had. A budget of
/// 0 empties the store, so a test takes what it needs with the store off before it fills it.
fn decode_off(cap: &str, payload: &[u8], set: &Set) -> Out {
let s = LoadedSets::global();
let budget = s.budget();
s.set_budget(0);
let out = decode(cap, payload, set);
s.set_budget(budget);
out
}
/// The capability constant `cap` names, and the lowering its driver lowers by.
fn registered(cap: &str) -> (&'static str, &'static str) {
match cap {
c::CAP_NEURAL_INT_HYPER => (c::CAP_NEURAL_INT_HYPER, wai::int_ops::WHM1),
c::CAP_VIDEO_INT_HYPER => (c::CAP_VIDEO_INT_HYPER, wai::int_ops::WHM1),
c::CAP_AUDIO_INT_CODEC => (c::CAP_AUDIO_INT_CODEC, wai::int_ops::WAM1),
c::CAP_SPLAT_INT_CODEC => (c::CAP_SPLAT_INT_CODEC, wai::int_ops::WHM1),
c::CAP_NEURAL_INT_SYNTH => (c::CAP_NEURAL_INT_SYNTH, wai::int_ops::INT_SYNTH),
other => panic!("{other} decodes against no set"),
}
}
/// The key the decode of `cap` keeps `set` under.
fn key(cap: &str, set: &Set) -> SetKey {
let (cap, lowering) = registered(cap);
// Every lowering of integer-operators §13 declares acc64, which each graph loads under.
SetKey::of(cap, lowering, wai_ops::Profile::Acc64, &view(set)).expect("a set's names make a listing")
}
/// An `int_synth` payload (integer-payloads §13): `h`, `w`, `C` channels of side
/// `2^s_log2`, a reserved 0, then the latents.
fn synth_payload(h: u32, w: u32, ch: u16, s_log2: u8, latents: &[i8]) -> Vec<u8> {
let mut p = Vec::new();
p.extend_from_slice(&h.to_le_bytes());
p.extend_from_slice(&w.to_le_bytes());
p.extend_from_slice(&ch.to_le_bytes());
p.extend_from_slice(&[s_log2, 0]);
p.extend(latents.iter().map(|&v| v as u8));
p
}
fn synth(payload: &[u8], spec: &[u8]) -> Out {
decode(c::CAP_NEURAL_INT_SYNTH, payload, &synth_set(spec))
}
fn synth_off(payload: &[u8], spec: &[u8]) -> Out {
decode_off(c::CAP_NEURAL_INT_SYNTH, payload, &synth_set(spec))
}
fn synth_set(spec: &[u8]) -> Set {
BTreeMap::from([("int_synth.json".to_string(), spec.to_vec())])
}
/// The byte offsets of each convolution's multipliers and biases in a `WHM1` network
/// (integer-payloads §9), with its output channels: `(out, mult, bias)`.
fn whm1_rq(b: &[u8]) -> Vec<(usize, usize, usize)> {
let u = |p: usize| u32::from_le_bytes(b[p..p + 4].try_into().unwrap()) as usize;
let (mut p, mut out) = (20, Vec::new());
for _ in 0..u(16) {
let tag = b[p];
p += 1;
match tag {
0 | 1 => {
let (inc, o, k) = (u(p), u(p + 4), u(p + 8));
p += if tag == 0 { 24 } else { 20 } + 2 * inc * o * k * k;
out.push((o, p, p + 8 * o));
p += 16 * o;
}
2 => p += 4 + 4 * u(p) * u(p) + 8 * u(p),
3 => {}
4 => p += 12,
t => panic!("layer tag {t}"),
}
}
assert_eq!(p, b.len(), "the walk ends at the network's last byte");
out
}
/// The committed vectors that decode against a set, with their capability, payload and set.
fn vectors() -> Vec<(String, String, Vec<u8>, Set)> {
let mut v: Vec<_> = index()["vectors"].as_array().unwrap().iter().filter(|v| !set_of(v).is_empty()).map(|v| {
(v["name"].as_str().unwrap().to_string(), v["capability"].as_str().unwrap().to_string(), read(v["payload"].as_str().unwrap()), set_of(v))
}).collect();
let names: Vec<String> = v.iter().map(|x| x.0.clone()).collect();
assert_eq!(names, ["tiny", "video_nk", "audio", "splat"], "the corpus's vectors with a set");
let hashes: Value = serde_json::from_slice(&read("neural-conformance/expected_hashes.json")).unwrap();
for name in hashes.as_object().unwrap().keys() {
let dir = format!("neural-conformance/vectors/{name}");
v.push((format!("int_synth {name}"), c::CAP_NEURAL_INT_SYNTH.to_string(), read(&format!("{dir}/latent.bin")), synth_set(&read(&format!("{dir}/int_synth.json")))));
}
v
}
#[test]
fn a_second_decode_through_a_set_is_a_hit_and_gives_the_same_bytes() {
let _s = store(DEFAULT_BUDGET);
let vectors = vectors();
let offs: Vec<Vec<u8>> = vectors.iter().map(|(name, cap, payload, set)| decode_off(cap, payload, set).unwrap_or_else(|e| panic!("{name}: {e:?}"))).collect();
let mut n = 0;
for ((name, cap, payload, set), off) in vectors.into_iter().zip(offs) {
let s0 = stats();
let first = decode(&cap, &payload, &set);
let s1 = stats();
let second = decode(&cap, &payload, &set);
let s2 = stats();
assert_eq!(first.as_ref(), Ok(&off), "{name}: the decode that loads its set gives the bytes of the decode with the store off");
assert_eq!(second.as_ref(), Ok(&off), "{name}: the decode from the store gives the same bytes");
assert_eq!((s1.misses - s0.misses, s1.hits - s0.hits, s1.kept - s0.kept), (1, 0, 1), "{name}: the first decode loads its set and keeps it");
assert_eq!((s2.misses - s1.misses, s2.hits - s1.hits, s2.kept - s1.kept), (0, 1, 0), "{name}: the second is a hit");
assert!(LoadedSets::global().contains(&key(&cap, &set)), "{name}: kept under the SHA-256 of its members");
// The declared cost reads the same entry (int_synth has none).
if cap != c::CAP_NEURAL_INT_SYNTH {
let v = view(&set);
let params = resolve(&cap, None, &[HeldBytes::Set { listing: b"", files: &v }], None).unwrap();
declared_cost(&cap, &payload, &params).unwrap();
assert_eq!((stats().hits - s2.hits, stats().misses - s2.misses), (1, 0), "{name}: the declared cost finds the set loaded");
}
println!("{name}: {} bytes alike with the store off, on (miss) and on (hit)", off.len());
n += 1;
}
assert_eq!(n, 8);
// tiny and video_nk read the same three members under two capabilities: two entries.
assert_eq!(stats().entries, 8);
}
#[test]
fn sets_one_byte_apart_never_share_an_entry() {
let _s = store(DEFAULT_BUDGET);
// The conditional codec: g_s's last bias, channel 0, its fourth byte flipped. The
// members keep their names.
let tiny = index()["vectors"].as_array().unwrap().iter().find(|v| v["name"] == "tiny").cloned().unwrap();
let (a, payload) = (set_of(&tiny), read(tiny["payload"].as_str().unwrap()));
let mut b = a.clone();
let g = b.get_mut("g_s.whm").unwrap();
let &(_, _, bias) = whm1_rq(g).last().unwrap();
g[bias + 3] ^= 0x01;
let cap = c::CAP_NEURAL_INT_HYPER;
// int_synth: the first digit of the last layer's first bias.
let base = read("neural-conformance/vectors/facade/int_synth.json");
let latent = read("neural-conformance/vectors/facade/latent.bin");
let at = base.windows(8).rposition(|w| w == b"\"bias\":[").unwrap() + 8;
let mut other = base.clone();
let digit = (at..other.len()).find(|&i| other[i].is_ascii_digit()).unwrap();
other[digit] = if other[digit] == b'9' { b'1' } else { b'9' };
// Each with the store off first.
let (a_off, b_off) = (decode_off(cap, &payload, &a).unwrap(), decode_off(cap, &payload, &b).unwrap());
let (x_off, y_off) = (synth_off(&latent, &base).unwrap(), synth_off(&latent, &other).unwrap());
assert_ne!(a_off, b_off, "the byte changes the image");
assert_ne!(x_off, y_off, "the digit changes the image");
assert_eq!(decode(cap, &payload, &a).unwrap(), a_off);
assert_eq!(decode(cap, &payload, &b).unwrap(), b_off, "the set one byte away is decoded with its own networks");
assert_eq!((stats().misses, stats().hits, stats().entries), (2, 0, 2), "each set is loaded and kept under its own key");
assert_eq!(decode(cap, &payload, &b).unwrap(), b_off);
assert_eq!(decode(cap, &payload, &a).unwrap(), a_off);
assert_eq!((stats().misses, stats().hits), (2, 2));
assert!(key(cap, &a) != key(cap, &b) && LoadedSets::global().contains(&key(cap, &a)) && LoadedSets::global().contains(&key(cap, &b)));
assert_eq!(synth(&latent, &base).unwrap(), x_off);
assert_eq!(synth(&latent, &other).unwrap(), y_off, "the set one byte away is decoded with its own graph");
assert_eq!(synth(&latent, &other).unwrap(), y_off);
assert_eq!((stats().misses, stats().hits, stats().entries), (4, 3, 4));
let (kx, ky) = (key(c::CAP_NEURAL_INT_SYNTH, &synth_set(&base)), key(c::CAP_NEURAL_INT_SYNTH, &synth_set(&other)));
assert!(kx != ky && LoadedSets::global().contains(&kx) && LoadedSets::global().contains(&ky));
}
#[test]
fn a_set_the_load_refuses_stays_refused_beside_a_held_neighbour() {
let _s = store(DEFAULT_BUDGET);
// int_synth: the falsification set that passes i64 by one unit (integer-operators
// §14.3), refused by the bounds pass, and its neighbour with a bias of 0 in place of
// 1, which loads.
let overflow = read("ops-conformance/sets/overflow_acc64/int_synth.json");
const BIAS: &[u8] = b"\"bias\":[1,0,0]";
let at = overflow.windows(BIAS.len()).position(|w| w == BIAS).expect("the overflow set's bias") + 8;
let mut fits = overflow.clone();
fits[at] = b'0';
assert_eq!(overflow.iter().zip(&fits).filter(|(x, y)| x != y).count(), 1, "one byte apart");
let p = synth_payload(1, 1, 1, 0, &[1]);
// The conditional codec: tiny with layer 0's first bias of g_s at 2^63 − 1, so no
// input range admits that layer (P_0(0) fails, integer-operators §13.2) and g_s is
// refused as it is lowered, after the plan; tiny itself is held.
let tiny = index()["vectors"].as_array().unwrap().iter().find(|v| v["name"] == "tiny").cloned().unwrap();
let (good, payload) = (set_of(&tiny), read(tiny["payload"].as_str().unwrap()));
let mut bad = good.clone();
let g = bad.get_mut("g_s.whm").unwrap();
let (_, _, bias) = whm1_rq(g)[0];
g[bias..bias + 8].copy_from_slice(&i64::MAX.to_le_bytes());
let cap = c::CAP_NEURAL_INT_HYPER;
// Each refusal with the store off first.
let synth_refused = synth_off(&p, &overflow).expect_err("the overflow set is refused");
assert!(synth_refused.1.contains("set does not load: model_mismatch at Load(8)"), "{synth_refused:?}");
let hyper_refused = decode_off(cap, &payload, &bad).expect_err("tiny with that bias is refused");
assert_eq!(hyper_refused, ("model_mismatch".to_string(), "g_s does not lower to an integer-operator graph that loads under acc64".to_string()));
assert!(synth(&p, &fits).is_ok(), "the neighbour loads and decodes");
let held = stats();
for _ in 0..2 {
assert_eq!(synth(&p, &overflow), Err(synth_refused.clone()), "refused again, with the same refusal");
}
let after = stats();
assert_eq!((after.misses - held.misses, after.hits - held.hits, after.kept - held.kept), (2, 0, 0), "loaded each time, never kept");
assert!(!LoadedSets::global().contains(&key(c::CAP_NEURAL_INT_SYNTH, &synth_set(&overflow))));
decode(cap, &payload, &good).expect("tiny decodes, and is kept");
let held = stats();
for _ in 0..2 {
assert_eq!(decode(cap, &payload, &bad), Err(hyper_refused.clone()), "refused again, with the same refusal");
}
let after = stats();
assert_eq!((after.misses - held.misses, after.hits - held.hits, after.kept - held.kept), (2, 0, 0), "loaded each time, never kept");
assert!(!LoadedSets::global().contains(&key(cap, &bad)) && LoadedSets::global().contains(&key(cap, &good)));
}
#[test]
fn every_refusal_fixture_keeps_its_code_and_detail_with_its_set_held() {
let _s = store(DEFAULT_BUDGET);
let idx = index();
let vectors = idx["vectors"].as_array().unwrap();
let mut n = 0;
for f in idx["refusals"].as_array().unwrap() {
let set = set_of(f);
if set.is_empty() {
continue;
}
let (name, cap) = (f["name"].as_str().unwrap(), f["capability"].as_str().unwrap());
let payload = read(f["payload"].as_str().unwrap());
let off = decode_off(cap, &payload, &set);
assert_eq!(off.as_ref().map_err(|e| e.0.as_str()).map(|_| ()), Err(f["code"].as_str().unwrap()), "{name}: the store off");
// Twice from an empty store: the first may load and keep the set (a refusal after
// the load step), the second then finds it held.
LoadedSets::global().clear();
let s0 = stats();
assert_eq!(decode(cap, &payload, &set), off, "{name}: from an empty store");
assert_eq!(decode(cap, &payload, &set), off, "{name}: again");
if f["layer"] == "parse" {
assert_eq!(stats(), s0, "{name}: a payload refused while it is parsed never reaches the store");
}
// With the set held by a committed vector's decode, where one decodes against it.
if let Some(v) = vectors.iter().find(|v| v["capability"] == cap && set_of(v) == set) {
decode(cap, &read(v["payload"].as_str().unwrap()), &set).expect("the committed vector decodes");
assert!(LoadedSets::global().contains(&key(cap, &set)));
assert_eq!(decode(cap, &payload, &set), off, "{name}: with its set held");
}
n += 1;
}
assert!(n >= 40, "{n} fixtures with a set");
}
#[test]
fn the_least_recently_used_set_goes_first_under_a_small_budget() {
let _s = store(DEFAULT_BUDGET);
// Three sets of one shape, so of one size.
let spec = |w: u8| format!(r#"{{"layers":[{{"in_ch":1,"out_ch":3,"k":1,"stride":1,"pad":0,"weight":[1,2,{w}],"bias":[0,0,0]}}]}}"#).into_bytes();
let (s1, s2, s3) = (spec(3), spec(4), spec(5));
let k = |s: &[u8]| key(c::CAP_NEURAL_INT_SYNTH, &synth_set(s));
let p = synth_payload(1, 1, 1, 0, &[5]);
let has = |s: &[u8]| LoadedSets::global().contains(&k(s));
assert_eq!(synth(&p, &s1).unwrap(), [5, 10, 15]);
assert_eq!(synth(&p, &s2).unwrap(), [5, 10, 20]);
let two = stats().resident_bytes;
assert_eq!((stats().entries, two % 2), (2, 0));
let one = two / 2;
// Room for two: s1 is used again, so s2 is the least recently used when s3 comes.
LoadedSets::global().set_budget(two);
assert_eq!(synth(&p, &s1).unwrap(), [5, 10, 15]);
assert_eq!(synth(&p, &s3).unwrap(), [5, 10, 25]);
assert!(has(&s1) && has(&s3) && !has(&s2), "s2 went first");
assert_eq!((stats().evicted, stats().entries, stats().resident_bytes), (1, 2, two));
assert_eq!(synth(&p, &s2).unwrap(), [5, 10, 20], "an evicted set is loaded again");
assert!(has(&s2) && has(&s3) && !has(&s1));
// A budget under one entry: everything goes, and a set past it is decoded, not kept.
LoadedSets::global().set_budget(one - 1);
assert_eq!(stats().entries, 0);
let before = stats();
assert_eq!(synth(&p, &s1).unwrap(), [5, 10, 15]);
assert_eq!((stats().too_large - before.too_large, stats().entries), (1, 0));
// A budget of 0: no lookup and nothing kept.
LoadedSets::global().set_budget(0);
let before = stats();
assert_eq!(synth(&p, &s1).unwrap(), [5, 10, 15]);
assert_eq!((stats().hits, stats().misses, stats().entries), (before.hits, before.misses, 0));
}
#[test]
fn threads_decoding_one_set_at_once_keep_one_entry() {
let _s = store(DEFAULT_BUDGET);
let tiny = index()["vectors"].as_array().unwrap().iter().find(|v| v["name"] == "tiny").cloned().unwrap();
let (set, payload) = (set_of(&tiny), read(tiny["payload"].as_str().unwrap()));
let cap = c::CAP_NEURAL_INT_HYPER;
let off = decode_off(cap, &payload, &set).unwrap();
let outs: Vec<Out> = std::thread::scope(|sc| {
let hs: Vec<_> = (0..8).map(|_| sc.spawn(|| decode(cap, &payload, &set))).collect();
hs.into_iter().map(|h| h.join().unwrap()).collect()
});
assert!(outs.iter().all(|o| o.as_ref() == Ok(&off)), "every thread decodes the same bytes");
let s = stats();
assert_eq!((s.hits + s.misses, s.entries), (8, 1));
}

4
wai/wai-rs/tests/int_video_keyframe.rs
View file

@ -7,7 +7,7 @@
use std::path::PathBuf;
use wai::int_hyper_synth::{decode_conditional, CondBitstream, CondEntropy, IntModel};
use wai::int_hyper_synth::{decode_conditional_checked, CondBitstream, CondEntropy, IntModel};
use wai::int_video::{decode_pframes, fnv1a};
fn vroot() -> PathBuf {
@ -34,7 +34,7 @@ fn neural_keyframe_clip_decodes_byte_for_byte() {
let h_s = IntModel::from_json(&jload(&ih.join("tiny_h_s.json")));
let ent = CondEntropy::from_value(&jload(&ih.join("tiny_entropy.json")));
let bs = CondBitstream::from_value(&jload(&ih.join("tiny_bitstream.json")));
let (frame0, h, w) = decode_conditional(&g_s, &h_s, &ent, &bs).expect("int_hyper I-frame");
let (frame0, h, w) = decode_conditional_checked(&g_s, &h_s, &ent, &bs.stream()).expect("int_hyper I-frame");
// frames 1..N: integer block-MC residuals on top.
let conf = jload(&iv.join("video_nk.json"));

4
wai/wai-rs/tests/int_video_sota_keyframe.rs
View file

@ -15,7 +15,7 @@ mod require_vectors;
use require_vectors::{env_dir, load_all};
use wai::int_hyper_synth::{decode_conditional, CondBitstream, CondEntropy, IntModel};
use wai::int_hyper_synth::{decode_conditional_checked, CondBitstream, CondEntropy, IntModel};
use wai::int_video::{decode_pframes, fnv1a};
fn hex16(s: &str) -> [u8; 8] {
@ -48,7 +48,7 @@ fn sota_int_hyper_keyframe_clip_decodes_byte_for_byte() {
let h_s = IntModel::from_json(&hh);
let ent = CondEntropy::from_value(&e);
let bs = CondBitstream::from_value(&b);
let (frame0, h, w) = decode_conditional(&g_s, &h_s, &ent, &bs).expect("int_hyper I-frame");
let (frame0, h, w) = decode_conditional_checked(&g_s, &h_s, &ent, &bs.stream()).expect("int_hyper I-frame");
// frames 1..N: integer motion-compensated P-frames panning the keyframe.
let clip = decode_pframes(&frame0, h, w, &vb).expect("p-frames");

25
wai/wai-web/demo/verify_int_payload_wasm.mjs
View file

@ -2,7 +2,9 @@
// the native decode: each committed WIV1, WIH1, WIA1 and WIS1 payload, decoded by
// decodeWiv1 / decodeWih1 / decodeWia1 / decodeWis1 in the browser crate built for
// wasm32, must give the SHA-256 its vector's sha256.json sidecar states, and a
// tool's JSON form must be refused with bad_magic.
// tool's JSON form must be refused with bad_magic. A second decode through the same
// set is answered from the page's store of loaded sets (loadedSetStats: one more hit,
// no more misses) with the same bytes, and a budget of 0 keeps and looks up nothing.
//
// node verify_int_payload_wasm.mjs [path/to/wai_web.js] [path/to/wai-rs/tests/vectors]
//
@ -42,6 +44,27 @@ const checks = [
side('int_video/sha256.json', 'video_nk').clip_sha256],
['tiny (WIH1)', () => hex(sha(wai.decodeWih1(rd('int_hyper/tiny_bitstream.wih1'), g, h, e).rgb)),
side('int_hyper/sha256.json', 'tiny').sha256],
['tiny again, its set held', () => {
const before = wai.loadedSetStats();
const got = hex(sha(wai.decodeWih1(rd('int_hyper/tiny_bitstream.wih1'), g, h, e).rgb));
const after = wai.loadedSetStats();
if (after.hits !== before.hits + 1 || after.misses !== before.misses) {
return `hits ${before.hits} -> ${after.hits}, misses ${before.misses} -> ${after.misses}`;
}
return got;
}, side('int_hyper/sha256.json', 'tiny').sha256],
['tiny with a budget of 0', () => {
const budget = wai.loadedSetStats().budget;
wai.setLoadedSetBudget(0);
const before = wai.loadedSetStats();
const got = hex(sha(wai.decodeWih1(rd('int_hyper/tiny_bitstream.wih1'), g, h, e).rgb));
const after = wai.loadedSetStats();
wai.setLoadedSetBudget(budget);
if (before.entries !== 0 || after.entries !== 0 || after.hits !== before.hits || after.misses !== before.misses) {
return `entries ${before.entries} -> ${after.entries}, a lookup was made`;
}
return got;
}, side('int_hyper/sha256.json', 'tiny').sha256],
['audio (WIA1)', () => {
const a = wai.decodeWia1(rd('int_audio/audio_bitstream.wia1'), rd('int_audio/audio_model.bin'));
if (a.sampleRate !== 16000) return `sample rate ${a.sampleRate}`;

60
wai/wai-web/src/int_payload.rs
View file

@ -91,6 +91,66 @@ pub fn int_payload_cost(
Ok(IntPayloadCost { output_bytes: c.output_bytes as f64, symbols: c.symbols as f64, working_bytes: c.working_bytes as f64, macs: c.macs })
}
/// The page's store of loaded parameter sets (`wai::int_set_cache`): a set that
/// `decodeWiv1`, `decodeWih1`, `decodeWia1`, `decodeWis1` or `decodeIntSynth` lowered
/// and loaded is kept, by the SHA-256 of its bytes, and a later decode through the same
/// bytes skips its lowering and bounds pass. The output is the same bytes either way,
/// and a refused set is never kept. The entry points that read the tools' JSON forms,
/// for debugging and conversion, load their models each time.
#[wasm_bindgen]
pub struct LoadedSetStats {
s: wai::int_set_cache::Stats,
}
#[wasm_bindgen]
impl LoadedSetStats {
/// Decodes that found their set loaded.
#[wasm_bindgen(getter)]
pub fn hits(&self) -> f64 {
self.s.hits as f64
}
/// Decodes that lowered and loaded their set.
#[wasm_bindgen(getter)]
pub fn misses(&self) -> f64 {
self.s.misses as f64
}
/// The sets held now.
#[wasm_bindgen(getter)]
pub fn entries(&self) -> f64 {
self.s.entries as f64
}
/// The bytes they are counted for against the budget.
#[wasm_bindgen(getter, js_name = residentBytes)]
pub fn resident_bytes(&self) -> f64 {
self.s.resident_bytes as f64
}
/// Entries removed to make room.
#[wasm_bindgen(getter)]
pub fn evicted(&self) -> f64 {
self.s.evicted as f64
}
/// The budget, in bytes (0: no set is kept).
#[wasm_bindgen(getter)]
pub fn budget(&self) -> f64 {
self.s.budget as f64
}
}
/// What the page's store of loaded sets has done and holds.
#[wasm_bindgen(js_name = loadedSetStats)]
pub fn loaded_set_stats() -> LoadedSetStats {
LoadedSetStats { s: wai::int_set_cache::LoadedSets::global().stats() }
}
/// Set the store's budget in bytes (256 MiB to start with), evicting the least recently
/// used sets until the rest fit; 0 empties it and keeps no set from then on. A negative
/// or non-finite budget is 0.
#[wasm_bindgen(js_name = setLoadedSetBudget)]
pub fn set_loaded_set_budget(bytes: f64) {
let b = if bytes.is_finite() && bytes > 0.0 { bytes as u64 } else { 0 };
wai::int_set_cache::LoadedSets::global().set_budget(b);
}
/// The conditional codec's three set members, from the host's bytes.
pub(crate) fn int_hyper_set<'a>(g_s: &'a [u8], h_s: &'a [u8], entropy: &'a [u8]) -> [(&'static str, &'a [u8]); 3] {
[("entropy.json", entropy), ("g_s.whm", g_s), ("h_s.whm", h_s)]

12
wai/wai-web/src/neural_int.rs
View file

@ -46,12 +46,16 @@ impl DecodedImage {
/// Decode a `wai.neural.int_synth` payload to RGB in the browser, integer-exact.
/// `payload` is the envelope's payload bytes; `spec_json` is the model spec (the
/// `int_synth.json` weights) the host has installed. No float, no ONNX — the
/// output is byte-identical to the native reference decode.
/// output is byte-identical to the native reference decode. It runs
/// `wai::int_payload::decode_int`, as the other payload entry points do, so a set it
/// has loaded is kept for the next decode (`loadedSetStats`), and a refusal's message
/// starts with its code (`payload_malformed`).
#[wasm_bindgen(js_name = decodeIntSynth)]
pub fn decode_int_synth(payload: &[u8], spec_json: &[u8]) -> Result<DecodedImage, JsError> {
let img = wai::int_synth::decode_from_spec(payload, spec_json)
.map_err(|e| JsError::new(&e.to_string()))?;
Ok(DecodedImage { rgb: img.rgb, width: img.width, height: img.height })
match crate::int_payload::decode("wai.neural.int_synth", payload, &[("int_synth.json", spec_json)]).map_err(|e| JsError::new(&e))? {
wai::int_payload::IntDecoded::Image { rgb, width, height } => Ok(DecodedImage { rgb, width, height }),
_ => Err(JsError::new("unsupported_capability: not an image")),
}
}
/// Decode a `WIH1` payload (`wai.neural.int_hyper`, `extensions/integer-payloads.md`