wai-rs: fast integer decode kernels that give the same integer for every output #108
46 changed files with 6565 additions and 465 deletions
6
.github/workflows/wai-byte-equality.yml
vendored
6
.github/workflows/wai-byte-equality.yml
vendored
|
|
@ -34,6 +34,9 @@ on:
|
|||
- 'wai/wai-rs/src/int_video.rs'
|
||||
- 'wai/wai-rs/src/int_transform.rs'
|
||||
- 'wai/wai-rs/src/int_hyper_synth.rs'
|
||||
# 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'
|
||||
- 'wai/wai-rs/tests/vectors/int_video/**'
|
||||
- 'wai/wai-rs/tests/vectors/int_hyper/**'
|
||||
- '.github/workflows/wai-byte-equality.yml'
|
||||
|
|
@ -53,6 +56,9 @@ on:
|
|||
- 'wai/wai-rs/src/int_video.rs'
|
||||
- 'wai/wai-rs/src/int_transform.rs'
|
||||
- 'wai/wai-rs/src/int_hyper_synth.rs'
|
||||
# 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'
|
||||
- 'wai/wai-rs/tests/vectors/int_video/**'
|
||||
- 'wai/wai-rs/tests/vectors/int_hyper/**'
|
||||
- '.github/workflows/wai-byte-equality.yml'
|
||||
|
|
|
|||
20
.github/workflows/wai-conformance.yml
vendored
20
.github/workflows/wai-conformance.yml
vendored
|
|
@ -98,6 +98,26 @@ jobs:
|
|||
set -euo pipefail
|
||||
cargo clippy --lib --no-default-features --features neural_int -- -A clippy::all
|
||||
cargo test --no-default-features --features world,audio_scene,haptic,splat4d,score,provenance,neural_int,film --test jpegai_decode_symbols
|
||||
|
||||
# The integer decoders run fast kernels (src/int_kernels.rs) that must give
|
||||
# the same integer as the scalar reference kernels for every output element.
|
||||
# The feature list above leaves out audio_int, video_int, world_int and
|
||||
# splat_int, so without these two steps int_audio's kernel-equivalence tests
|
||||
# and tests/int_kernels_vectors.rs (every committed integer-decoder vector
|
||||
# decoded under the reference and under each fast policy, against its frozen
|
||||
# hash) would run nowhere. `--lib` also carries the other kernels' equivalence
|
||||
# modules; tests/int_kernels_geometry.rs runs edge geometry (strides up to
|
||||
# usize::MAX, padding far beyond the kernel) through the public decoders
|
||||
# under every policy. Both profiles, because the overflow-parity tests assert
|
||||
# different things in each: with overflow checks on (dev) that the fast path panics
|
||||
# where the reference panics, with them off (release) that it returns the
|
||||
# same integers. All pure Rust, so every leg runs them.
|
||||
- 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 --lib --test int_kernels_vectors --test int_kernels_geometry
|
||||
- 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 --lib --test int_kernels_vectors --test int_kernels_geometry
|
||||
- name: Carriage conformance — every registered capability round-trips byte-exact
|
||||
working-directory: wai/wai-rs
|
||||
run: cargo test --no-default-features --test carriage
|
||||
|
|
|
|||
|
|
@ -33,6 +33,10 @@ on:
|
|||
- 'wai/wai-rs/src/int_entropy.rs'
|
||||
- 'wai/wai-rs/src/int_transform.rs'
|
||||
- 'wai/wai-rs/src/wai_video_policy.rs'
|
||||
# Not verified by any harness, but #[path]-included because int_transform's
|
||||
# IGDN dispatches to it, so it must still compile under the model checker
|
||||
# (its cfg(kani) branches are compiled nowhere else).
|
||||
- 'wai/wai-rs/src/int_kernels.rs'
|
||||
- '.github/workflows/wai-formal-verification.yml'
|
||||
push:
|
||||
branches:
|
||||
|
|
@ -43,6 +47,7 @@ on:
|
|||
- 'wai/wai-rs/src/int_entropy.rs'
|
||||
- 'wai/wai-rs/src/int_transform.rs'
|
||||
- 'wai/wai-rs/src/wai_video_policy.rs'
|
||||
- 'wai/wai-rs/src/int_kernels.rs'
|
||||
- '.github/workflows/wai-formal-verification.yml'
|
||||
|
||||
jobs:
|
||||
|
|
|
|||
170
.github/workflows/wai-jpegai-e2e.yml
vendored
Normal file
170
.github/workflows/wai-jpegai-e2e.yml
vendored
Normal file
|
|
@ -0,0 +1,170 @@
|
|||
name: WAI JPEG-AI whole decode (pinned inputs, required)
|
||||
|
||||
# WAI states that the exact decoder, `wai::jpegai_decode::decode`, decodes an
|
||||
# ISO/IEC 6048 JPEG-AI codestream to RGB deterministically, so a hard binding
|
||||
# taken over one decode verifies a second one, and that on one reference
|
||||
# codestream its output stays within max|Δ| ≤ 2 of the reference decoder's
|
||||
# (recorded max|Δ| = 1; wai/SPEC.md §5, the wai.image.jpegai row).
|
||||
# Across machines the registry states less: it classes wai.image.jpegai as
|
||||
# NeuralEntropyExact, a conservative class that guarantees only the entropy decode
|
||||
# (the decoded symbols and their entropy parameters) to be byte-identical on every
|
||||
# machine. That the exact decoder's pinned output is the same on other machines is
|
||||
# a measurement, recorded in wai/extensions/determinism-tiers.md §6. The tests that
|
||||
# carry these statements (tests/jpegai_e2e_full.rs,
|
||||
# tests/jpegai_reconstruction_binding.rs) read their inputs from outside the tree
|
||||
# and used to report PASS without them.
|
||||
#
|
||||
# This job fetches those inputs from WAI's versioned models root (jpegai/v1/),
|
||||
# checks every file against the SHA-256 manifest committed in
|
||||
# wai/tools/jpegai-e2e-v1.sha256, and runs both tests with
|
||||
# WAI_REQUIRE_SOTA_VECTORS=1, so a missing input is a failure. What passes here:
|
||||
# - the whole chain, codestream → RGB, through the one public `decode()` call,
|
||||
# lands within max|Δ| ≤ 2 of the reference decoder's output, and at most 5%
|
||||
# of samples differ from it at all (the test's bar);
|
||||
# - two decodes in one process are byte-identical and the reconstruction binding
|
||||
# verifies the second and rejects a one-byte change;
|
||||
# - the decoded pixels hash to WAI_JPEGAI_EXPECT_SHA256 below: the value measured
|
||||
# from these inputs on an Apple Silicon Mac (aarch64 macOS, native) and
|
||||
# reproduced by an x86_64 Linux build (glibc, run under QEMU user-mode
|
||||
# emulation), and the value the four-platform comparison in determinism-tiers
|
||||
# §6 (macOS arm64, wasm32-wasip2, Linux arm64, Linux x86_64 under user-mode
|
||||
# emulation) found on every platform. This runner is x86_64 hardware, so the
|
||||
# job adds a decode on one more machine to that measurement. A match is a
|
||||
# measurement on this runner, not something the registry class promises; a
|
||||
# mismatch is a departure from the measured agreement.
|
||||
# The pinned decode is max|Δ| = 1 from the reference on 4.446% of samples
|
||||
# (on the Mac and the emulated x86_64 build), so a hash match also confirms
|
||||
# that recorded max|Δ| = 1 for this codestream; the test's own tolerance
|
||||
# above is looser than that.
|
||||
#
|
||||
# What this job does not cover:
|
||||
# - Other codestreams: it decodes one 256×256 test codestream, so the
|
||||
# max|Δ| ≤ 2 bound and the determinism statement are checked on that input
|
||||
# only.
|
||||
# - The extraction: the inputs are an extraction of the JPEG-AI
|
||||
# reference-software assets (tools/wai_jpegai_build_e2e_inputs.py), which are
|
||||
# not in the repository, so the extraction is not re-run here; the manifest
|
||||
# pins its published bytes.
|
||||
# - The ONNX-runtime decode (`neural::jpegai`, tests/jpegai_live.rs), which needs
|
||||
# those local assets and is not run in CI. It also needs the `neural` feature,
|
||||
# whose prebuilt ONNX Runtime does not link on the `linux` runner's image
|
||||
# (glibc 2.35; the library references glibc 2.38 symbols), which is where
|
||||
# wai-neural-e2e.yml's runs fail once they reach the Rust build. The
|
||||
# features built here (neural_int, provenance) link no prebuilt native
|
||||
# library.
|
||||
# The job depends on the models root being reachable; an outage fails the job.
|
||||
#
|
||||
# The path filters list the source files the two tests' decode path imports,
|
||||
# among them src/float_repro.rs (the exact decoder's `exp` and its roundings to
|
||||
# an integer) and src/int_kernels.rs (the fast convolution and parallel-map
|
||||
# kernels the integer networks run through). An entry for a file the tree does
|
||||
# not have matches nothing until the file exists. A module the decode path
|
||||
# starts importing belongs here too; otherwise a change to that module alone is
|
||||
# first checked by the weekly run.
|
||||
|
||||
on:
|
||||
push:
|
||||
branches: [main]
|
||||
paths:
|
||||
- 'wai/wai-rs/src/jpegai_*.rs'
|
||||
- 'wai/wai-rs/src/reconstruction_binding.rs'
|
||||
- 'wai/wai-rs/src/int_transformer.rs'
|
||||
- 'wai/wai-rs/src/int_hyper_synth.rs'
|
||||
- 'wai/wai-rs/src/int_kernels.rs'
|
||||
- 'wai/wai-rs/src/float_repro.rs'
|
||||
- 'wai/wai-rs/src/int_entropy.rs'
|
||||
- 'wai/wai-rs/src/int_transform.rs'
|
||||
- 'wai/wai-rs/src/rans.rs'
|
||||
- 'wai/wai-rs/src/lib.rs'
|
||||
- 'wai/wai-rs/Cargo.toml'
|
||||
- 'wai/wai-rs/Cargo.lock'
|
||||
- 'wai/wai-rs/tests/jpegai_e2e_full.rs'
|
||||
- 'wai/wai-rs/tests/jpegai_reconstruction_binding.rs'
|
||||
- 'wai/wai-rs/tests/jpegai_cross_platform.rs'
|
||||
- 'wai/wai-rs/tests/support/**'
|
||||
- 'wai/tools/wai_jpegai_fetch_e2e_inputs.sh'
|
||||
- 'wai/tools/jpegai-e2e-v1.sha256'
|
||||
- '.github/workflows/wai-jpegai-e2e.yml'
|
||||
pull_request:
|
||||
paths:
|
||||
- 'wai/wai-rs/src/jpegai_*.rs'
|
||||
- 'wai/wai-rs/src/reconstruction_binding.rs'
|
||||
- 'wai/wai-rs/src/int_transformer.rs'
|
||||
- 'wai/wai-rs/src/int_hyper_synth.rs'
|
||||
- 'wai/wai-rs/src/int_kernels.rs'
|
||||
- 'wai/wai-rs/src/float_repro.rs'
|
||||
- 'wai/wai-rs/src/int_entropy.rs'
|
||||
- 'wai/wai-rs/src/int_transform.rs'
|
||||
- 'wai/wai-rs/src/rans.rs'
|
||||
- 'wai/wai-rs/src/lib.rs'
|
||||
- 'wai/wai-rs/Cargo.toml'
|
||||
- 'wai/wai-rs/Cargo.lock'
|
||||
- 'wai/wai-rs/tests/jpegai_e2e_full.rs'
|
||||
- 'wai/wai-rs/tests/jpegai_reconstruction_binding.rs'
|
||||
- 'wai/wai-rs/tests/jpegai_cross_platform.rs'
|
||||
- 'wai/wai-rs/tests/support/**'
|
||||
- 'wai/tools/wai_jpegai_fetch_e2e_inputs.sh'
|
||||
- 'wai/tools/jpegai-e2e-v1.sha256'
|
||||
- '.github/workflows/wai-jpegai-e2e.yml'
|
||||
# Weekly, independent of changes: the inputs live outside the repository, and
|
||||
# their disappearing should surface here, not at the next unrelated PR.
|
||||
schedule:
|
||||
- cron: '41 6 * * 1'
|
||||
|
||||
jobs:
|
||||
jpegai-e2e:
|
||||
name: fetch pinned JPEG-AI inputs → whole decode + pixel hash, required
|
||||
runs-on: linux
|
||||
timeout-minutes: 30
|
||||
env:
|
||||
WAI_REQUIRE_SOTA_VECTORS: '1'
|
||||
WAI_JPEGAI_DIR: ${{ github.workspace }}/jpegai_e2e
|
||||
# sha256 of the decoded RGB8 pixels for the jpegai/v1 inputs (see header).
|
||||
WAI_JPEGAI_EXPECT_SHA256: dbc494a5a670914edf55c07c48afc0dd348cd05d3a82fff293bf384b2bb69ea3
|
||||
steps:
|
||||
- uses: actions/checkout@v5
|
||||
# Bootstrap rustup: the docker runner image ships no Rust.
|
||||
- name: Install Rust 1.98.0
|
||||
shell: bash
|
||||
run: |
|
||||
export PATH="$HOME/.cargo/bin:$PATH"
|
||||
command -v rustup >/dev/null 2>&1 || \
|
||||
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --profile minimal --default-toolchain 1.98.0
|
||||
rustup toolchain install 1.98.0 --profile minimal
|
||||
rustup default 1.98.0
|
||||
echo "$HOME/.cargo/bin" >> "$GITHUB_PATH"
|
||||
- name: Fetch the JPEG-AI inputs and check each against the pinned manifest
|
||||
shell: bash
|
||||
working-directory: wai
|
||||
run: tools/wai_jpegai_fetch_e2e_inputs.sh "$WAI_JPEGAI_DIR"
|
||||
- name: Whole decode vs the reference, determinism, pinned pixel hash (missing input = failure)
|
||||
shell: bash
|
||||
working-directory: wai/wai-rs
|
||||
run: |
|
||||
set -o pipefail
|
||||
log="${RUNNER_TEMP:-/tmp}/jpegai-tests.log"
|
||||
# Without the switch both tests would skip and report ok.
|
||||
[ "${WAI_REQUIRE_SOTA_VECTORS:-}" = 1 ] \
|
||||
|| { echo "::error::WAI_REQUIRE_SOTA_VECTORS must be 1 in this job"; exit 1; }
|
||||
export WAI_JPEGAI_CODESTREAM="$WAI_JPEGAI_DIR/test256.bin"
|
||||
export WAI_JPEGAI_WEIGHTS_DIR="$WAI_JPEGAI_DIR"
|
||||
export WAI_JPEGAI_NORM="$WAI_JPEGAI_DIR/norm.json"
|
||||
# --show-output rather than --nocapture: each test's output is printed
|
||||
# after the results, so it cannot split a `test X ... ok` line below.
|
||||
cargo test --release --no-default-features --features neural_int,provenance \
|
||||
--test jpegai_e2e_full --test jpegai_reconstruction_binding --test jpegai_cross_platform \
|
||||
-- --show-output 2>&1 | tee "$log"
|
||||
# A test that was renamed or compiled out would not fail — it would just
|
||||
# not run. Require each one by name, and the pinned-hash check by its line.
|
||||
for t in codestream_to_rgb_matches_reference_png real_decode_to_reconstruction_binding \
|
||||
decode_matches_pinned_hashes; do
|
||||
grep -q "^test $t \.\.\. ok$" "$log" || { echo "::error::$t did not run and pass"; exit 1; }
|
||||
done
|
||||
grep -q "matches the pinned WAI_JPEGAI_EXPECT_SHA256" "$log" \
|
||||
|| { echo "::error::the pinned pixel hash was not checked"; exit 1; }
|
||||
# A skip also prints `test X ... ok`; every skip writes a banner with
|
||||
# "SKIPPED" and "NOT VERIFIED" (tests/support/require_vectors.rs).
|
||||
if grep -nE 'SKIPPED|NOT VERIFIED' "$log"; then
|
||||
echo "::error::a JPEG-AI test skipped instead of checking its input"
|
||||
exit 1
|
||||
fi
|
||||
254
.github/workflows/wai-sota-vectors.yml
vendored
Normal file
254
.github/workflows/wai-sota-vectors.yml
vendored
Normal file
|
|
@ -0,0 +1,254 @@
|
|||
name: WAI mbt2018_mean integer decode (regenerated vectors)
|
||||
|
||||
# `wai.neural.int_hyper` is an integer export of one pretrained model: the
|
||||
# mbt2018_mean mean-scale hyperprior from the CompressAI model zoo (quality 6,
|
||||
# MSE), post-training-quantized to integers. WAI states that its whole
|
||||
# bitstream→image decode is integer, so the Rust decoder reproduces the export
|
||||
# tool's integer reference byte-for-byte, and that `wai.video.int_hyper` builds a
|
||||
# clip on that keyframe. The model is too large to commit, so the tests that carry
|
||||
# these statements read vectors from outside the tree; without them they used to
|
||||
# return early and report PASS.
|
||||
#
|
||||
# This job checks them on every run: it regenerates the vectors from pinned
|
||||
# inputs (a public-domain photograph shipped in the pinned scikit-image wheel and
|
||||
# pinned by the SHA-256 of its pixels; the published mbt2018_mean checkpoint,
|
||||
# pinned by its full SHA-256 and loaded as tensors only; a pinned Python stack),
|
||||
# then runs the int_hyper vector tests with WAI_REQUIRE_SOTA_VECTORS=1, under which
|
||||
# a missing vector is a failure. It also runs tests/require_vectors_rule.rs, which
|
||||
# checks in child processes what the switch does not cover here: that without it a
|
||||
# variable naming absent or incomplete inputs still fails, and a test skips only
|
||||
# when no variable names what is missing (an unset variable, or an absent or
|
||||
# incomplete default location). What passes here:
|
||||
# - the Rust integer decode reproduces the export tool's integer reference
|
||||
# byte-for-byte (synthesis, hyper-decoder, the whole conditional bitstream, the
|
||||
# compact binary model, the eight-frame clip);
|
||||
# - on the pinned photograph, the integer decode costs at most 1 dB against the
|
||||
# same model's float decode (the bar CI enforces), measured by the Rust decoder
|
||||
# on its own output. Measured on an Apple Silicon Mac: 35.985 dB for the integer
|
||||
# decode (8-bit PSNR against the source pixels) against 36.175 dB for the float
|
||||
# decode (the export tool's PSNR of the float reconstruction, before 8-bit
|
||||
# rounding), a cost of 0.19 dB, from 6328 bytes of rANS payload for the
|
||||
# 256×256 crop (about 0.77 bits per pixel).
|
||||
#
|
||||
# What this job does not cover:
|
||||
# - How mbt2018_mean compares with any other codec. Nothing here measures rate or
|
||||
# distortion against another model or standard; the job checks only that the
|
||||
# integer decode is exact and within 1 dB of the same model's float decode.
|
||||
# - Quality on other images. PSNR is image-dependent and the job measures one
|
||||
# crop. The 42.6 dB development figure was measured on a photograph that is not
|
||||
# in the repository.
|
||||
# - Cross-architecture byte-equality of these vectors. They are regenerated on
|
||||
# this runner, so the job checks Rust against the Python reference on one
|
||||
# machine (x86_64 Linux). The cross-architecture check is the committed tiny
|
||||
# int_hyper vector in wai-byte-equality.yml.
|
||||
# - Portable or repeatable vector hashes. No expected hash of the regenerated
|
||||
# vectors is pinned in the repository (the committed tiny vector's is, in
|
||||
# tests/vectors/int_hyper/tiny.json). Each is written by the same run's export
|
||||
# and checked in that run, and a hash taken from one run's vectors must not be
|
||||
# pinned for another run, on this machine or another: it can differ with both
|
||||
# decoders correct. What was measured, comparing all sixteen files by SHA-256:
|
||||
# - Between an Apple Silicon Mac and this runner's x86_64 image, four files
|
||||
# differed. int_synth_sota_input.json (the synthesis-only test's input
|
||||
# latent, which the float encoder and hyper-decoder compute) differed, and
|
||||
# with it the rgb_hash in int_synth_sota.json (fd9ca73a… on the Mac,
|
||||
# f0fb6813… here); each side wrote the same latent in all of its runs.
|
||||
# cond.json and int_synth_sota.json also differed. Beyond that rgb_hash,
|
||||
# their difference is not established as a machine difference: both files
|
||||
# also differ from run to run on this runner, most likely in psnr_float
|
||||
# (next point). source.png differed in its compressed bytes only: step 1
|
||||
# writes it with the PNG encoder from the pinned pixels before any model
|
||||
# runs, and its pixels (source.rgb) were identical, so that difference is
|
||||
# the encoder's output, not float computation.
|
||||
# - Two runs of the same export on this runner (same image, pins and tree,
|
||||
# minutes apart) wrote different cond.json and int_synth_sota.json; the
|
||||
# other fourteen files and both rgb_hash values matched. The field those two
|
||||
# files share that comes from the float model is psnr_float, the float
|
||||
# decode's PSNR, so it most likely varies from run to run here, although the
|
||||
# export sets torch to deterministic algorithms on one thread. The runner's
|
||||
# files were not kept, so the field is inferred, not diffed. Repeated runs on
|
||||
# the Mac wrote identical files.
|
||||
# Nothing the tests compare exactly reads psnr_float; only the quality test
|
||||
# does, against its thresholds (the 1 dB bar and a 30 dB floor). The files
|
||||
# that matched (both models, the conditional bitstream, frame0, the clip) are
|
||||
# derived from float computation too, so matching here does not make them
|
||||
# portable. The pins that hold across runs and machines are on the inputs: the
|
||||
# photograph's pixels and the checkpoint's bytes. The Python stack is pinned
|
||||
# by version, but each platform installs its own build of it.
|
||||
# The JPEG-AI whole-decoder tests run in wai-jpegai-e2e.yml (see
|
||||
# wai/wai-rs/tests/support/require_vectors.rs).
|
||||
#
|
||||
# The path filters list the source files the int_hyper decode imports, among
|
||||
# them src/int_kernels.rs (the fast convolution kernels int_hyper_synth.rs runs
|
||||
# through). An entry for a file the tree does not have matches nothing until the
|
||||
# file exists. A module the decode path starts importing belongs here too;
|
||||
# otherwise a change to that module alone is first checked by the weekly run.
|
||||
#
|
||||
# Licences of what is fetched: CompressAI and its model zoo, BSD-3-Clause-Clear;
|
||||
# the astronaut photograph, public domain (NASA), per the scikit-image docstring.
|
||||
|
||||
on:
|
||||
push:
|
||||
branches: [main]
|
||||
paths:
|
||||
- 'wai/wai-rs/src/int_hyper_synth.rs'
|
||||
- 'wai/wai-rs/src/int_kernels.rs'
|
||||
- 'wai/wai-rs/src/int_entropy.rs'
|
||||
- 'wai/wai-rs/src/int_transform.rs'
|
||||
- 'wai/wai-rs/src/int_video.rs'
|
||||
- 'wai/wai-rs/src/rans.rs'
|
||||
- 'wai/wai-rs/src/lib.rs'
|
||||
- 'wai/wai-rs/Cargo.toml'
|
||||
- 'wai/wai-rs/Cargo.lock'
|
||||
- 'wai/wai-rs/examples/int_hyper_decode_frame.rs'
|
||||
- 'wai/wai-rs/tests/int_hyper_synth.rs'
|
||||
- 'wai/wai-rs/tests/int_video_sota_keyframe.rs'
|
||||
- 'wai/wai-rs/tests/support/**'
|
||||
- 'wai/wai-rs/tests/require_vectors_rule.rs'
|
||||
- 'wai/tools/wai_sota_*'
|
||||
- 'wai/tools/wai_int_video_sota_keyframe_encode.py'
|
||||
- 'wai/tools/requirements-sota-int.txt'
|
||||
- '.github/workflows/wai-sota-vectors.yml'
|
||||
pull_request:
|
||||
paths:
|
||||
- 'wai/wai-rs/src/int_hyper_synth.rs'
|
||||
- 'wai/wai-rs/src/int_kernels.rs'
|
||||
- 'wai/wai-rs/src/int_entropy.rs'
|
||||
- 'wai/wai-rs/src/int_transform.rs'
|
||||
- 'wai/wai-rs/src/int_video.rs'
|
||||
- 'wai/wai-rs/src/rans.rs'
|
||||
- 'wai/wai-rs/src/lib.rs'
|
||||
- 'wai/wai-rs/Cargo.toml'
|
||||
- 'wai/wai-rs/Cargo.lock'
|
||||
- 'wai/wai-rs/examples/int_hyper_decode_frame.rs'
|
||||
- 'wai/wai-rs/tests/int_hyper_synth.rs'
|
||||
- 'wai/wai-rs/tests/int_video_sota_keyframe.rs'
|
||||
- 'wai/wai-rs/tests/support/**'
|
||||
- 'wai/wai-rs/tests/require_vectors_rule.rs'
|
||||
- 'wai/tools/wai_sota_*'
|
||||
- 'wai/tools/wai_int_video_sota_keyframe_encode.py'
|
||||
- 'wai/tools/requirements-sota-int.txt'
|
||||
- '.github/workflows/wai-sota-vectors.yml'
|
||||
# Weekly, independent of changes: the job depends on artifacts fetched from
|
||||
# outside the repository (the checkpoint, the pinned wheels), and their
|
||||
# disappearing should surface as a failure here, not at the next unrelated PR.
|
||||
schedule:
|
||||
- cron: '17 6 * * 1'
|
||||
|
||||
jobs:
|
||||
sota-vectors:
|
||||
name: regenerate mbt2018_mean integer vectors → Rust decode, required
|
||||
# The `linux` runner is a docker executor on catthehacker/ubuntu:act-22.04
|
||||
# (glibc 2.35, no Rust). The Python half of this setup — setup-python 3.11,
|
||||
# CPU torch, compressai — completes there: every logged wai-neural-e2e.yml
|
||||
# run got through it and generated its fixtures. Each of those runs that
|
||||
# reached the Rust build then failed linking `ort`'s prebuilt ONNX Runtime,
|
||||
# which references glibc 2.38 and newer-libstdc++ symbols (`__isoc23_strtol`,
|
||||
# `_M_replace_cold`). The features built here (neural_int, video_int) are
|
||||
# pure Rust and link no prebuilt native library; do not add `neural` here.
|
||||
runs-on: linux
|
||||
timeout-minutes: 45
|
||||
env:
|
||||
WAI_SOTA_INT_DIR: ${{ github.workspace }}/sota_int
|
||||
WAI_REQUIRE_SOTA_VECTORS: '1'
|
||||
steps:
|
||||
- uses: actions/checkout@v5
|
||||
- uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: '3.11'
|
||||
- name: Pinned ML stack (CPU torch first, then the rest)
|
||||
shell: bash
|
||||
working-directory: wai
|
||||
run: |
|
||||
python -m pip install --upgrade pip
|
||||
# CPU wheels: the default index resolves the multi-GB CUDA builds.
|
||||
python -m pip install torch==2.14.0 torchvision==0.29.0 --index-url https://download.pytorch.org/whl/cpu
|
||||
python -m pip install -r tools/requirements-sota-int.txt
|
||||
# Fail here, not mid-export, if the resolver installed anything but the
|
||||
# pins, or if torchvision's compiled ops do not match torch (the
|
||||
# `torchvision::nms does not exist` mismatch wai-neural-e2e.yml once hit).
|
||||
python - tools/requirements-sota-int.txt <<'PY'
|
||||
import sys, importlib.metadata as md
|
||||
bad = []
|
||||
for line in open(sys.argv[1]):
|
||||
line = line.split("#")[0].strip()
|
||||
if not line:
|
||||
continue
|
||||
name, want = line.split("==")
|
||||
got = md.version(name)
|
||||
print(f"{name}=={got}")
|
||||
if got.split("+")[0] != want:
|
||||
bad.append(f"{name} {got}, pinned {want}")
|
||||
if bad:
|
||||
sys.exit("not the pinned stack: " + "; ".join(bad))
|
||||
import torch, torchvision, compressai, skimage
|
||||
torchvision.ops.nms(torch.zeros(1, 4), torch.zeros(1), 0.5)
|
||||
PY
|
||||
# Bootstrap rustup: the docker runner image ships no Rust.
|
||||
- name: Install Rust 1.98.0
|
||||
shell: bash
|
||||
run: |
|
||||
export PATH="$HOME/.cargo/bin:$PATH"
|
||||
command -v rustup >/dev/null 2>&1 || \
|
||||
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --profile minimal --default-toolchain 1.98.0
|
||||
rustup toolchain install 1.98.0 --profile minimal
|
||||
rustup default 1.98.0
|
||||
echo "$HOME/.cargo/bin" >> "$GITHUB_PATH"
|
||||
- name: Regenerate the int_hyper vectors (source photo → integer export → keyframe → clip)
|
||||
shell: bash
|
||||
working-directory: wai
|
||||
run: |
|
||||
tools/wai_sota_vectors.sh "$WAI_SOTA_INT_DIR"
|
||||
# Recorded so a result can be tied to the exact vectors it checked. It
|
||||
# identifies this run's vectors, not this machine's: two runs on this
|
||||
# runner wrote different cond.json and int_synth_sota.json, and other
|
||||
# files differ between machines (see the header). It is not a list any
|
||||
# other run must reproduce, here or elsewhere.
|
||||
(cd "$WAI_SOTA_INT_DIR" && sha256sum *)
|
||||
- name: Rust decodes every int_hyper vector byte-for-byte (missing input = failure), and the missing-input rule holds
|
||||
shell: bash
|
||||
working-directory: wai/wai-rs
|
||||
run: |
|
||||
set -o pipefail
|
||||
log="${RUNNER_TEMP:-/tmp}/sota-tests.log"
|
||||
# Without the switch every vector test would skip and report ok.
|
||||
[ "${WAI_REQUIRE_SOTA_VECTORS:-}" = 1 ] \
|
||||
|| { echo "::error::WAI_REQUIRE_SOTA_VECTORS must be 1 in this job"; exit 1; }
|
||||
# --show-output rather than --nocapture: each test's output is printed
|
||||
# after the results, so it cannot split a `test X ... ok` line below.
|
||||
# require_vectors_rule runs its cases in child processes with the switch
|
||||
# and the input variables set per case, so the job's own env does not
|
||||
# reach them; it prints nothing on success, so it cannot trip the skip
|
||||
# backstop below.
|
||||
cargo test --release --no-default-features --features neural_int,video_int \
|
||||
--test int_hyper_synth --test int_video_sota_keyframe --test require_vectors_rule \
|
||||
-- --show-output 2>&1 | tee "$log"
|
||||
# A test that was renamed or compiled out would not fail — it would just
|
||||
# not run. Require each one by name.
|
||||
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 \
|
||||
integer_decode_keeps_the_float_codec_quality \
|
||||
decodes_the_committed_tiny_conformance_vector \
|
||||
sota_int_hyper_keyframe_clip_decodes_byte_for_byte \
|
||||
without_the_switch_a_named_directory_that_does_not_exist_fails \
|
||||
without_the_switch_a_named_path_that_is_a_file_fails_as_a_directory \
|
||||
without_the_switch_a_named_empty_directory_fails \
|
||||
without_the_switch_an_unset_or_empty_directory_variable_skips_on_an_absent_default \
|
||||
without_the_switch_an_incomplete_default_directory_skips \
|
||||
a_named_directory_holding_the_input_loads_it \
|
||||
without_the_switch_a_path_variable_naming_nothing_fails \
|
||||
without_the_switch_a_path_variable_naming_nothing_fails_when_another_is_unset \
|
||||
without_the_switch_an_unset_path_variable_skips \
|
||||
path_variables_naming_files_that_exist_load_them \
|
||||
switch_values_that_leave_it_off_behave_as_unset \
|
||||
with_the_switch_an_absent_default_fails; do
|
||||
grep -q "^test $t \.\.\. ok$" "$log" || { echo "::error::$t did not run and pass"; exit 1; }
|
||||
done
|
||||
# A skipped test ALSO prints `test X ... ok`, so the name check alone
|
||||
# cannot tell a skip from a pass. Every skip writes a banner containing
|
||||
# "SKIPPED" and "NOT VERIFIED" (tests/support/require_vectors.rs); either
|
||||
# word anywhere in the log fails the job, whatever the switch says.
|
||||
if grep -nE 'SKIPPED|NOT VERIFIED' "$log"; then
|
||||
echo "::error::a vector test skipped instead of checking its input"
|
||||
exit 1
|
||||
fi
|
||||
132
wai/CHANGELOG.md
132
wai/CHANGELOG.md
|
|
@ -6,6 +6,123 @@ wire format may change and such changes are recorded here.
|
|||
|
||||
## v0.4.0 — current
|
||||
|
||||
### Integer decoders — fast kernels that keep every output integer
|
||||
|
||||
- `wai-rs` adds `int_kernels`: register-blocked convolution, transposed
|
||||
convolution and IGDN kernels (NEON widening multiply-accumulate on aarch64,
|
||||
portable Rust elsewhere), scoped worker threads (none on wasm), and a run-time
|
||||
`Policy` that selects them or the retained scalar reference kernels. The
|
||||
int_synth, int_hyper, int_audio, int_splat, int_world and JPEG-AI synthesis
|
||||
decoders dispatch to them. No bitstream, conformance vector or hash changes.
|
||||
- The claim, precisely: for every input the reference decodes, every policy
|
||||
returns the same integer for every output element; with overflow checks on, a
|
||||
fast policy panics on the inputs where the reference panics, though the panic
|
||||
message can differ. Geometry the fast path does not handle, including a
|
||||
transposed convolution whose sizes the reference's own arithmetic cannot form,
|
||||
goes to the reference. The fast path's own allocations and loops stay in
|
||||
proportion to what the reference reads and writes, whatever the stride or
|
||||
padding. The two paths still allocate different buffers, so an allocation that
|
||||
fails is outside the claim. It is argued by construction and tested, not
|
||||
machine-proved. The rules, the tests, their limits and the measured throughput
|
||||
(one machine) are in
|
||||
[`research/decode-throughput.md`](research/decode-throughput.md).
|
||||
- CI: `wai-conformance` runs the kernel-equivalence tests, every committed
|
||||
vector that reaches the kernels and edge geometry through the public decoders
|
||||
(strides up to `usize::MAX`, padding far beyond the kernel, and strides and
|
||||
paddings at which the reference's size arithmetic overflows), under every
|
||||
policy, with overflow checks on and off, on each platform leg. Geometry on
|
||||
either side of the padded-input bound is pinned too: at the bound the fast
|
||||
kernel runs, beyond it the reference does, and a test-build count of
|
||||
fast-kernel entries shows which one ran.
|
||||
`wai-byte-equality` and `wai-formal-verification` fire on changes to
|
||||
`int_kernels.rs`.
|
||||
|
||||
### Headline decode tests fail on a missing input instead of passing
|
||||
|
||||
- The `int_hyper` tests over the pretrained mbt2018_mean (`int_hyper_synth`,
|
||||
`int_video_sota_keyframe`) and the JPEG-AI whole-decoder tests
|
||||
(`jpegai_e2e_full`, `jpegai_reconstruction_binding`, `jpegai_live`) read inputs
|
||||
from outside the tree, and used to `return` — reported as a pass — when those
|
||||
were absent. With
|
||||
`WAI_REQUIRE_SOTA_VECTORS=1` a missing input is now a failure that names the
|
||||
tool producing it; without it the skip is written to the process's real stderr,
|
||||
so it shows in an ordinary `cargo test` run. Either way, a corrupt input file
|
||||
fails, and so does a variable that names inputs which are not all there: a
|
||||
path that does not exist (`WAI_JPEGAI_CODESTREAM`/`_WEIGHTS_DIR`/`_NORM`, even
|
||||
when another of the three is unset), or a directory that does not exist or
|
||||
lacks a file the test reads (`WAI_SOTA_INT_DIR`, `WAI_JPEGAI_ASSETS`, the
|
||||
JPEG-AI weights directory). So without the switch a test skips only when no
|
||||
variable names what is missing: one of the three path variables is unset, or
|
||||
`WAI_SOTA_INT_DIR` or `WAI_JPEGAI_ASSETS` is unset or empty and its default
|
||||
directory (`/tmp/sota_int`, `../jpegai-vm-assets`) is absent or incomplete.
|
||||
The CI jobs always set the switch, so the new test
|
||||
`require_vectors_rule` checks this rule without it: it runs the helpers in
|
||||
child processes with the switch and the variables set per case, and
|
||||
`wai-sota-vectors.yml` runs it.
|
||||
- New CI job `wai-sota-vectors.yml` regenerates the mbt2018_mean integer vectors
|
||||
from pinned inputs (`tools/wai_sota_vectors.sh`, the Python stack in
|
||||
`tools/requirements-sota-int.txt`, a public-domain photograph pinned by SHA-256,
|
||||
the published checkpoint pinned by its full SHA-256) and runs those tests
|
||||
with the switch set. It is triggered by changes to the decode path and by a
|
||||
weekly schedule. Besides requiring each test by name, it fails if the test log
|
||||
contains a skip banner, since a skipped test also reports `ok`. The export tool
|
||||
checks the checkpoint's bytes before loading them (as tensors only), because the
|
||||
model zoo's own download path checks no hash. The job checks the Rust decode
|
||||
against the Python integer reference on one runner; cross-architecture
|
||||
byte-equality remains the committed tiny vector's job in
|
||||
`wai-byte-equality.yml`. The regenerated vectors are not portable, and part
|
||||
of them is not repeatable run to run on every machine. Between an Apple
|
||||
Silicon Mac and the x86_64 runner, the synthesis-only test's input latent,
|
||||
which the float model computes, differed, and with it its expected
|
||||
`rgb_hash`; each side wrote the same latent in all of its runs. Two runs on
|
||||
the one runner wrote different `cond.json` and `int_synth_sota.json`, most
|
||||
likely in the float decode's PSNR, which only the quality test reads, against
|
||||
thresholds. The PNG bytes of the source also differed between the Mac and the
|
||||
runner; the PNG encoder writes them from the pinned pixels before any model
|
||||
runs, and the pixels matched. So no expected hash of the regenerated vectors
|
||||
is pinned in the repository (the committed tiny vector's is, in
|
||||
`tests/vectors/int_hyper/tiny.json`); each run checks the hashes its own export
|
||||
wrote, and a hash from one run's vectors must not be pinned for another run,
|
||||
on the same machine or another. The pins that hold across machines are on the
|
||||
inputs: the photograph's pixels and the checkpoint's bytes.
|
||||
- New test `integer_decode_keeps_the_float_codec_quality`: the Rust decoder
|
||||
measures its own reconstruction against the source pixels and must stay within
|
||||
1 dB of the same model's float decode on the same image. On the pinned
|
||||
photograph, measured on an Apple Silicon Mac: integer decode 35.985 dB (8-bit
|
||||
PSNR of the Rust decode), float decode 36.175 dB (the export tool's PSNR of the
|
||||
float reconstruction), a cost of 0.19 dB, from 6328 bytes of rANS payload for
|
||||
256×256 pixels. This is a relative bar on one image: nothing here compares
|
||||
mbt2018_mean with another codec, and quality is image-dependent — the 42.6 dB
|
||||
figure quoted for this codec was measured on a different photograph, which is
|
||||
not in the repository.
|
||||
- New CI job `wai-jpegai-e2e.yml` runs `jpegai_e2e_full` and
|
||||
`jpegai_reconstruction_binding` with the switch set, with the same by-name and
|
||||
skip-banner checks. Their inputs — the decoder
|
||||
weights, normative tables and test codestream, extracted from the JPEG-AI
|
||||
reference-software assets (which are not in this repository) — are fetched from
|
||||
WAI's versioned models root (`jpegai/v1/`) by
|
||||
`tools/wai_jpegai_fetch_e2e_inputs.sh` and checked against the SHA-256 manifest
|
||||
`tools/jpegai-e2e-v1.sha256`. The job also pins the SHA-256 of the decoded
|
||||
pixels (`WAI_JPEGAI_EXPECT_SHA256`): measured natively on an Apple Silicon Mac
|
||||
and reproduced by an x86_64 Linux build under QEMU user-mode emulation, and
|
||||
required again on the x86_64 runner. That extends a measurement, not a
|
||||
registry guarantee: the registry classes `wai.image.jpegai` as
|
||||
`NeuralEntropyExact`, a conservative class that guarantees only the entropy
|
||||
decode to be byte-identical on every machine, while the exact decoder's pinned
|
||||
output for these inputs was measured identical on four platforms
|
||||
(determinism-tiers §6).
|
||||
- `jpegai_live` (the ONNX-runtime `wai.image.jpegai` decode) follows the same
|
||||
missing-input rule, but no CI job runs it: its model bundle exists only in the
|
||||
local JPEG-AI asset directory, and it needs the `neural` and `codecs` features.
|
||||
- determinism-tiers §5–§6 record the two jobs. The full-scale `int_hyper` row and
|
||||
the JPEG AI whole-decode row (split from the index-derivation row, whose test CI
|
||||
still does not run) now name their workflow, state that a missing input fails
|
||||
there instead of skipping, and state that neither workflow has run through the
|
||||
forge: each passed under `forgejo-runner exec` on the Linux runner's host in an
|
||||
earlier form, and its current steps have run on macOS arm64 only. §5 quotes the
|
||||
`int_hyper` cost measured on the pinned photograph in place of the 42.6 dB
|
||||
figure, and names the bar each job re-checks.
|
||||
|
||||
### Prior carriage and derived-track operations — registered with their limits
|
||||
|
||||
- New [`extensions/prior-carriage.md`]: a `Prior` MAY declare the wire form of
|
||||
|
|
@ -261,7 +378,7 @@ wire format may change and such changes are recorded here.
|
|||
this change's final source, and on all four also before the math-library calls
|
||||
were replaced. determinism-tiers §6 records the toolchain and the method for each.
|
||||
Not measured: macOS x86_64, Windows, x86_64 hardware. The test skips without its
|
||||
inputs, and no CI job supplies them yet; `float_repro`'s own pinned-output test
|
||||
inputs; `wai-jpegai-e2e.yml` supplies them and requires it to pass; `float_repro`'s own pinned-output test
|
||||
needs no inputs and runs in the library tests.
|
||||
- Two checks keep math-library calls off the decode path. Clippy's
|
||||
`disallowed_methods` (`wai-rs/clippy.toml`), denied in each module on the path,
|
||||
|
|
@ -560,9 +677,12 @@ wire format may change and such changes are recorded here.
|
|||
external standard's codestream deterministically.
|
||||
- Modules: `jpegai_decode` (the composed entry point) over `jpegai_bitstream`,
|
||||
`jpegai_mans`, `jpegai_pipeline`, `jpegai_synthesis`, `jpegai_icci`
|
||||
(feature `neural_int`). Inputs are regenerable from the shipped assets via
|
||||
`tools/wai_jpegai_build_e2e_inputs.py`. The whole chain — codestream → RGB —
|
||||
is exercised end to end against the reference PNG.
|
||||
(feature `neural_int`). Inputs are regenerable via
|
||||
`tools/wai_jpegai_build_e2e_inputs.py` from the JPEG-AI reference-software
|
||||
assets, which are kept locally and are not in this repository; the extraction is
|
||||
published under WAI's versioned models root (`jpegai/v1/`). The whole chain —
|
||||
codestream → RGB — is exercised end to end against the reference PNG by
|
||||
`wai-jpegai-e2e.yml`, which fetches those inputs pinned by SHA-256.
|
||||
|
||||
### JPEG AI — a live, self-sufficient `wai.image.jpegai` decode (2026-06-22)
|
||||
|
||||
|
|
@ -577,7 +697,9 @@ wire format may change and such changes are recorded here.
|
|||
exact, 0 violations / 65536 positions); the dequant is
|
||||
`resi × rvs_dq[grfs[ch]][lk] / 2^16 / scaler_vec[ch]`, with `grfs` determined
|
||||
the same way the σ rvs does so the per-channel `scaler_vec` is consistent.
|
||||
Matches the reference decoder's PNG to within rounding (max diff 1).
|
||||
Matches the reference decoder's PNG to within rounding (max diff 1), checked by
|
||||
`tests/jpegai_live.rs` on a machine holding the local JPEG-AI assets; no CI job
|
||||
runs it (its model bundle is not in the repository).
|
||||
|
||||
### Film reel — `wai.film.reel`, cuts over one world (2026-06-17)
|
||||
|
||||
|
|
|
|||
|
|
@ -1249,6 +1249,14 @@ can be called from any language. The crate wraps:
|
|||
Run `cargo test --lib` (with `RUSTFLAGS="-L /opt/homebrew/opt/flac/lib"`
|
||||
on macOS, or via the committed `.cargo/config.toml`).
|
||||
|
||||
Decode speed is an implementation property, not a conformance property: the
|
||||
integer decoders run faster kernels (SIMD, threads) built to give the same
|
||||
integer for every output element as the retained scalar reference kernels. That
|
||||
equivalence is argued by construction and tested, not machine-proved (the
|
||||
harnesses of §7 do not reach these kernels) — see
|
||||
[`research/decode-throughput.md`](research/decode-throughput.md) for the
|
||||
argument, its tests, its limits and measured throughput.
|
||||
|
||||
## Appendix B — Why a *capability menu*, not a custom codec
|
||||
|
||||
Earlier drafts of WAI defined custom transform + entropy stages for
|
||||
|
|
|
|||
|
|
@ -13,6 +13,11 @@ rust-version = "1.97"
|
|||
license = "Apache-2.0"
|
||||
publish = false
|
||||
|
||||
# The included wai-rs sources carry `cfg(kani)` items (set only by the model
|
||||
# checker in wai-verify). Declared so this crate's builds do not warn on it.
|
||||
[lints.rust]
|
||||
unexpected_cfgs = { level = "warn", check-cfg = ["cfg(kani)"] }
|
||||
|
||||
[dependencies]
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json = "1"
|
||||
|
|
|
|||
|
|
@ -17,6 +17,15 @@ pub mod neural;
|
|||
#[path = "../../wai-rs/src/int_entropy.rs"]
|
||||
pub mod int_entropy;
|
||||
|
||||
// The fast kernels the integer decoders dispatch to (and the policy that can
|
||||
// select the scalar reference instead). Included so the cross-arch matrix checks
|
||||
// the code path sinks actually run by default. Pure std, self-contained. Some of
|
||||
// its helpers serve modules this crate does not compile (the attention nets), so
|
||||
// their dead-code warnings are expected here and silenced at the include.
|
||||
#[allow(dead_code)]
|
||||
#[path = "../../wai-rs/src/int_kernels.rs"]
|
||||
pub mod int_kernels;
|
||||
|
||||
// The merged byte-exact JPEG-AI decoder's entropy stage (crate-root, self-contained —
|
||||
// no shims needed). Included so the cross-arch byte-equality matrix covers the decoder
|
||||
// that `wai.image.jpegai` actually reconstructs with, not only the older `neural/` port.
|
||||
|
|
|
|||
|
|
@ -433,14 +433,19 @@ same ones for its integer networks:
|
|||
|
||||
- **Fidelity to the float model.** The quantised networks do not reproduce the
|
||||
float model's output, so the pinned reconstruction differs from it. For
|
||||
`wai.neural.int_hyper`, `wai-rs/src/int_hyper_synth.rs` records 42.6 dB PSNR for
|
||||
one pretrained model (`mbt2018_mean`), within ~0.6 dB of the float model; the
|
||||
test image, quality level and machine are not recorded there. For the JPEG AI
|
||||
exact decoder, `wai-rs/tests/jpegai_e2e_full.rs` compares against the reference
|
||||
decoder's output on one 256×256 reference codestream and asserts max|Δ| ≤ 2 with
|
||||
at most 5% of samples differing; the recorded result is max|Δ| = 1 on 4.446% of
|
||||
samples (macOS arm64, release build, 2026-09-27). Neither figure is re-derived in
|
||||
CI (§6).
|
||||
`wai.neural.int_hyper`, over one pretrained model (`mbt2018_mean`, quality 6)
|
||||
and one pinned test photograph, the integer decode of the full conditional
|
||||
bitstream measured 35.985 dB PSNR against 36.175 dB for the same model's float
|
||||
decode, a cost of 0.19 dB (macOS arm64; `wai-rs/src/int_hyper_synth.rs`).
|
||||
Quality is image-dependent, and this compares the model only with its own float
|
||||
decode. For the JPEG AI exact decoder, `wai-rs/tests/jpegai_e2e_full.rs` compares
|
||||
against the reference decoder's output on one 256×256 reference codestream and
|
||||
asserts max|Δ| ≤ 2 with at most 5% of samples differing; the recorded result is
|
||||
max|Δ| = 1 on 4.446% of samples (macOS arm64, release build, 2026-09-27). CI
|
||||
jobs re-check both bars: `wai-sota-vectors.yml` requires the `int_hyper` cost to
|
||||
stay within 1 dB on that photograph, and `wai-jpegai-e2e.yml` runs the JPEG AI
|
||||
comparison and requires the decoded pixels to hash to a pinned value. Neither
|
||||
workflow has run through the forge yet (§6).
|
||||
- **Compute.** The reference decoders are single-threaded CPU code with no
|
||||
explicit SIMD, in integer arithmetic (i64, with i128 intermediates in places;
|
||||
plus fixed-order float for the JPEG AI exact decoder's float stages). They
|
||||
|
|
@ -478,13 +483,14 @@ day.
|
|||
|---|---|---|
|
||||
| Cross-architecture identity of the rANS coder, `int_synth`, and the JPEG AI symbol decode (a reference-encoder codestream, given its distribution indices) | golden-byte assertions passed on Linux x86_64, Linux arm64 and macOS arm64 on main at `afc7ece` (2026-09-27, `wai-conformance.yml`; `wai-byte-equality.yml` last ran on main at `3f50097`, 2026-08-29, all three legs passing), and on Windows x86_64 last on main at `3e7811f` (2026-08-27). Between `3e7811f` and `afc7ece` the suite only gained the video and keyframe goldens; these goldens, the sources they decode and the suite's lockfile are unchanged | `byte-exact-conformance/`; `.github/workflows/wai-conformance.yml`, `wai-byte-equality.yml` |
|
||||
| Cross-architecture identity of integer inter-frame video, and of `int_hyper` on a small learned keyframe | the same suite, passing on Linux x86_64, Linux arm64 and macOS arm64 on main at `afc7ece` (2026-09-27). Windows x86_64 never decoded the keyframe clip: it was added in `d4c922b` (2026-08-28), whose Windows legs were cancelled. Windows decoded the integer video goldens only on the pre-merge commits `c9be2ad`, `256e5dd`, `7dfca26` and `3d34860` (2026-08-25/26), whose video sources, vectors and golden file match main but whose `byte-exact-conformance/Cargo.lock` differs; the change reached main at `3f50097`, after the last Windows run on main | as above |
|
||||
| `wai.image.jpegai` exact decoder: the integer derivation of the distribution indices matches the reference decoder exactly; the whole decode stays within a stated bound of it (max\|Δ\| ≤ 2, ≤ 5% of samples) | tests on one machine that skip unless their fixtures (sink-supplied weights, normative tables, codestream) are provided; CI does not provide them. The whole-decode test last ran on macOS arm64 on 2026-09-27 with a local bundle | `wai-rs/tests/jpegai_hyperscale_calibration.rs`, `jpegai_e2e_full.rs` |
|
||||
| `wai.image.jpegai` exact decoder across platforms: the `jpegai/v1` codestream decodes to the same RGB8 bytes, and to the same bits in the float values the decoder reports after each stage (values inside a stage are not hashed), on each platform measured | measured 2026-09-28 with rustc 1.98.0, release builds, inputs checked against their SHA-256, all four on the source of the change that added this row: macOS arm64 (native, on the development machine); wasm32-wasip2 (wasmtime 49, on the same machine); Linux arm64 (glibc 2.36, a container in an arm64 Linux VM on that machine); Linux x86_64 (glibc 2.36, the x86_64 build run in that container under user-mode emulation, QEMU 7.0, not on x86_64 hardware). All nine pinned hashes matched on each platform. On all four they also matched on an earlier state of that change, before the decode's `exp` and rounding moved to `float_repro`. Before the move the decode called the platform `exp` on every target, and its roundings compiled to calls to the symbols `round`, `roundf` and `rint` on x86_64 and `round` and `roundf` on wasm32, which the link resolves: an x86_64 Linux test executable inspected after the move, whose other code still rounds this way, took them from Rust's compiler builtins, not the C library (for wasm32 this was not checked). Not measured: macOS x86_64 (no x86_64 translation layer is installed on the development machine), Windows, x86_64 hardware, 32-bit targets. The test skips unless the inputs are supplied, and no CI job supplies them to it | `wai-rs/tests/jpegai_cross_platform.rs` |
|
||||
| `wai.image.jpegai` exact decoder: the integer derivation of the distribution indices matches the reference decoder exactly | a test on one machine that skips unless its fixtures (sink-supplied weights, normative tables, codestream) are provided; CI does not provide them | `wai-rs/tests/jpegai_hyperscale_calibration.rs` |
|
||||
| `wai.image.jpegai` exact decoder, whole decode through the one public `decode()` call: it stays within a stated bound of the reference decoder's output (max\|Δ\| ≤ 2, ≤ 5% of samples); two decodes in one process are byte-identical, and a reconstruction binding taken over one verifies the other | `wai-jpegai-e2e.yml` fetches the pinned `jpegai/v1` inputs (an extraction of the reference-software assets, published under WAI's models root), checks each file against the SHA-256 manifest `tools/jpegai-e2e-v1.sha256`, and runs both tests with `WAI_REQUIRE_SOTA_VECTORS=1`, under which a missing input fails instead of skipping. The job also fails on any skip banner in the test log, and requires the decoded pixels to hash to a pinned SHA-256. One 256×256 codestream. The workflow has not run through the forge. Its job passed under `forgejo-runner exec` on the Linux runner's host (x86_64 hardware, 2026-09-27), with the exact decoder as it stood on main at `afc7ece`, which calls the platform `exp`; there the decoded pixels matched the pinned hash, the value measured on macOS arm64. The job's steps have otherwise run on macOS arm64 only, outside a runner | `wai-rs/tests/jpegai_e2e_full.rs`, `jpegai_reconstruction_binding.rs`; `.github/workflows/wai-jpegai-e2e.yml` |
|
||||
| `wai.image.jpegai` exact decoder across platforms: the `jpegai/v1` codestream decodes to the same RGB8 bytes, and to the same bits in the float values the decoder reports after each stage (values inside a stage are not hashed), on each platform measured | measured 2026-09-28 with rustc 1.98.0, release builds, inputs checked against their SHA-256, all four on the source of the change that added this row: macOS arm64 (native, on the development machine); wasm32-wasip2 (wasmtime 49, on the same machine); Linux arm64 (glibc 2.36, a container in an arm64 Linux VM on that machine); Linux x86_64 (glibc 2.36, the x86_64 build run in that container under user-mode emulation, QEMU 7.0, not on x86_64 hardware). All nine pinned hashes matched on each platform. On all four they also matched on an earlier state of that change, before the decode's `exp` and rounding moved to `float_repro`. Before the move the decode called the platform `exp` on every target, and its roundings compiled to calls to the symbols `round`, `roundf` and `rint` on x86_64 and `round` and `roundf` on wasm32, which the link resolves: an x86_64 Linux test executable inspected after the move, whose other code still rounds this way, took them from Rust's compiler builtins, not the C library (for wasm32 this was not checked). Not measured: macOS x86_64 (no x86_64 translation layer is installed on the development machine), Windows, x86_64 hardware, 32-bit targets. The test skips unless the inputs are supplied. `wai-jpegai-e2e.yml` supplies them and requires this test to run and pass on its Linux runner; that workflow has not yet run through the forge | `wai-rs/tests/jpegai_cross_platform.rs` |
|
||||
| `float_repro`, the decode's `exp` and rounding: the same bits over a fixed grid of inputs; `exp` within 2 ULPs of the host math library (1 ULP measured); rounding equal to the standard library's | library unit tests under the library-test step of the conformance workflow (`neural_int`), which has not yet run them on a runner. The pinned-bits and rounding tests never skip. The host comparison asserts the host library's accuracy as much as this module's. Its gate selects by target, not by host: it runs on macOS arm64, on Linux x86_64 and arm64 with glibc, and on WASI, and is reported as ignored, with the reason, elsewhere, Windows included. On the source of the change that added this row it was run against these hosts' `exp`, with a worst case of 1 ULP on each: Apple's math library on macOS arm64; the WASI C library under wasmtime; glibc 2.36 on Linux arm64, and glibc 2.35 in the Ubuntu 22.04 image the workflow's Linux arm64 runner used on 2026-09-28; glibc 2.36 on Linux x86_64 under user-mode emulation, against three of the four `exp` variants glibc chooses among by processor feature (the variant each run used follows from the features glibc's loader reported active): SSE2 (QEMU 7.0, whose emulated processor reports none of AVX, FMA, AVX2 and FMA4), AVX (QEMU 7.2, `-cpu max` with FMA and AVX2 masked by glibc's `hwcaps` tunable, which leaves AVX active) and FMA (QEMU 7.2, `-cpu max`). On a separate probe of 20 million inputs, the SSE2 and AVX variants returned the same bits, and the FMA variant different bits for some of them. Not compared: the FMA4 variant, which glibc chooses on a processor that has FMA4 but not both FMA and AVX2, and which neither emulator offers. Not run on x86_64 hardware or on other glibc versions, where the gate still runs it; a failure there would say that the host's `exp` and this one are more than 2 ULPs apart, not that `float_repro`'s bits changed | `wai-rs/src/float_repro.rs` tests |
|
||||
| The exact decode path calls no platform math library | two checks. Clippy's `disallowed_methods` (`wai-rs/clippy.toml`), denied in each module on the decode path, rejects any use of a listed float method, including `f64::exp(x)` and `.map(f64::exp)`. `tests/jpegai_decode_symbols.rs` lists the symbols of an executable that links only the decoder's entry points and weight loaders. It rejects any C math-library function the executable imports, and any it defines other than the round-to-integral functions (`floor`, `round`, `rint` and the like), whose results IEEE 754 fixes. It allows those definitions because a definition shows what the linker kept, not what the decode calls: linked with GNU ld 2.40 for x86_64 Linux, the debug executable defined the compiler builtins' `floor`, though nothing in it referred to `floor`. On the source of the change that added this row, both checks pass on macOS arm64, the symbol check in debug and release builds. The symbol check also passes, in debug builds, on Linux arm64 (the container above) and on Linux x86_64 under user-mode emulation, linked with GNU ld (reporting that `floor` as allowed) and with LLD. With a platform `exp` put back as `f64::exp(…)`, both fail on macOS arm64 while the decoder's unit tests still pass, and the symbol check fails on Linux x86_64 with GNU ld. With a float `%` put into the decode, the symbol check fails on macOS arm64 and clippy, which does not see operators, passes. The conformance workflow runs both on its unix legs; neither has run on a runner yet | `wai-rs/clippy.toml`, `wai-rs/tests/jpegai_decode_symbols.rs`, `.github/workflows/wai-conformance.yml` |
|
||||
| `wai.image.jpegai` scope: the reference codestream is accepted; copies with β, `model_id`, the colour transform or an LSBS flag changed are refused; a flush-to-zero or round-toward-zero environment is detected. The float ONNX path calls the same check | library unit tests over a committed copy of the reference codestream, which never skip and fall under the library-test step of the conformance workflow (`neural_int`); the environment test runs on arm64 only. Run on the four platforms above, on the source of the change that added the cross-platform row (the environment test on macOS arm64 and Linux arm64). The float path's refusal (β and `model_id` edits) is tested only with its ONNX bundle supplied (`jpegai_scope_ort.rs`, skips otherwise) | `wai-rs/src/jpegai_decode.rs` tests, `wai-rs/tests/vectors/jpegai/` |
|
||||
| SPEC §7's replicate rule: the registry class decides first, so an integer-declared prior over `wai.image.jpegai` or a float codec is refused; `wai.image.jpegai` is the only `NeuralEntropyExact` capability; the standard codecs sit in the classes SPEC §7 lists, and the tolerance and unclassified ones are not refused; an encoder that pins no prior (`wai wrap`) writes `replicate` only over the standard classes, for the capability and for any declared fallback, or, in the `WAI2` companion form, for the base payload's capability and each companion's | library unit tests (`codecs`, `container`) that never skip and fall under the same library-test step. The strict-decode test in `tests/neural_native_decode.rs` needs the `neural` feature; `wai-neural-gates.yml` runs it on the Linux runner, in an Ubuntu 24.04 image because the `neural` test binary does not link in that runner's default 22.04 image. The workflow has not run through the forge; its job passed under `forgejo-runner exec` on the Linux runner's host (Linux x86_64, 2026-09-27). The test that runs the `wai wrap` binary needs the `codecs` feature, which no CI job builds. Otherwise run so far on macOS arm64 only | `wai-rs/src/codecs/mod.rs`, `container.rs` tests; `wai-rs/tests/neural_native_decode.rs`, `cli_wrap_intent.rs` |
|
||||
| `wai.neural.int_hyper` and `wai.video.int_hyper` at full model scale | tests that skip when the exported model is absent; not run in CI | `wai-rs/tests/int_hyper_synth.rs`, `int_video_sota_keyframe.rs` |
|
||||
| `wai.neural.int_hyper` and `wai.video.int_hyper` at full model scale (the pretrained `mbt2018_mean`, exported to integers): the Rust decode reproduces the export tool's integer reference byte for byte (synthesis, hyper-decoder, the whole conditional bitstream, an eight-frame clip), and costs at most 1 dB against the same model's float decode on one pinned photograph | `wai-sota-vectors.yml` regenerates the vectors on the runner from pinned inputs (the photograph's pixels and the checkpoint's bytes by SHA-256, the Python stack by version) and runs both tests with `WAI_REQUIRE_SOTA_VECTORS=1`, under which a missing vector fails instead of skipping. The job also fails on any skip banner in the test log, and runs `require_vectors_rule.rs`, which checks that without the switch a variable naming absent or incomplete inputs fails as well. Rust is checked against the export of the same run on the same machine: the vectors come from float computation, and some of them differed between machines and between runs, so none of their hashes is pinned and none is compared across machines. The workflow has not run through the forge. Its job passed under `forgejo-runner exec` on the Linux runner's host (x86_64 hardware, 2026-09-27) in an earlier form that did not yet run `require_vectors_rule.rs`; its current steps have run on macOS arm64 only, outside a runner | `wai-rs/tests/int_hyper_synth.rs`, `int_video_sota_keyframe.rs`, `require_vectors_rule.rs`; `.github/workflows/wai-sota-vectors.yml` |
|
||||
| `wai.audio.int_codec`, `wai.splat.int_codec` | committed vectors that never skip, matched against an independent Python integer reference; their cargo features (`audio_int`, `splat_int`) are enabled by no CI job | `wai-rs/tests/int_audio.rs`, `int_splat.rs` |
|
||||
|
||||
So at `entropy-consistency`, WAI's JPEG AI symbol decode is demonstrated across
|
||||
|
|
@ -495,8 +501,9 @@ identity (x86_64 against arm64) is demonstrated for `int_synth`, integer video a
|
|||
`int_hyper` on committed small-scale vectors; of these, only `int_synth` has also
|
||||
been decoded on Windows on main. `wai.audio.int_codec` and `wai.splat.int_codec`
|
||||
reproduce a separate implementation's output, but no CI job runs them and no
|
||||
cross-architecture comparison exists. For full-scale models, identity rests on
|
||||
construction, not on a cross-machine comparison of decoded output.
|
||||
cross-architecture comparison exists. For full-scale `int_hyper`, identity rests on
|
||||
construction, not on a cross-machine comparison of decoded output: CI checks its
|
||||
decode against the export's integer reference on one machine per run.
|
||||
|
||||
For the JPEG AI exact decoder's pinned reconstruction (§5), identity across
|
||||
machines rests on construction, and on one comparison: a single codestream on the
|
||||
|
|
|
|||
460
wai/research/decode-throughput.md
Normal file
460
wai/research/decode-throughput.md
Normal file
|
|
@ -0,0 +1,460 @@
|
|||
# Integer decoder throughput — fast kernels, same bytes
|
||||
|
||||
> **Status.** Implementation note for the reference implementation (`wai-rs`),
|
||||
> informative, not normative. Nothing here changes a bitstream or a conformance
|
||||
> hash, and decoded outputs keep the same integers by the argument and tests of
|
||||
> §2. Every number below was measured on **one machine** under the method stated
|
||||
> in §3. Other machines, other loads and other models will give other numbers.
|
||||
|
||||
## 1. What changed
|
||||
|
||||
The integer decoders (`int_synth`, `int_hyper_synth`, `int_audio`, `int_splat`,
|
||||
`int_world`, the neural video keyframe, and the JPEG-AI synthesis in
|
||||
`jpegai_synthesis` / `int_transformer`) now run their hot kernels through
|
||||
`wai-rs/src/int_kernels.rs`:
|
||||
|
||||
- **Convolution and transposed convolution** — register-blocked tiles (4 output
|
||||
channels × 8 adjacent outputs held in registers across every tap), zero-padded
|
||||
inputs, and a polyphase formulation of the transposed convolution so every tap
|
||||
reads a contiguous run of inputs. On aarch64 the 32×32→64-bit
|
||||
multiply-accumulate uses NEON widening instructions (`std::arch`, one small
|
||||
submodule); every other target runs the same kernel as portable Rust.
|
||||
- **IGDN** — the same register blocking, with NEON on aarch64 when every
|
||||
`x²` fits in 32 bits, and a portable 64-bit path otherwise.
|
||||
- **JPEG-AI channel attention, LayerNorm and element-wise gates** — rows, positions
|
||||
and element chunks spread over threads; the attention sums narrowed from `i128`
|
||||
to `i64` where a bound proves the result identical.
|
||||
- **Threads** — `std::thread::scope` workers taking tiles from a shared queue. No
|
||||
new dependency. None on wasm (the browser sink runs the same kernels
|
||||
single-threaded) or under the model checker. The worker count is automatic
|
||||
(all cores, capped so each worker has enough work), explicit through
|
||||
`int_kernels::Policy`, or set process-wide with `WAI_INT_THREADS`.
|
||||
|
||||
The scalar convolution, transposed-convolution, IGDN, channel-attention (MDTA)
|
||||
and channel-LayerNorm kernels they replace are kept verbatim in `reference`
|
||||
modules beside each caller and stay selectable at run time:
|
||||
`int_kernels::with_policy(Policy::reference(), || decode(..))`. They are the
|
||||
executable specification the fast kernels are checked against. The JPEG-AI
|
||||
element-wise steps (the GDFN gate and the residual additions in
|
||||
`int_transformer`, the activations and gates in `jpegai_synthesis`) have no
|
||||
separate reference copy: they were rewritten in place onto `par_map` /
|
||||
`par_zip`, each output element is the same pure function of its own inputs as
|
||||
before, and under the reference policy they run on one worker.
|
||||
|
||||
## 2. Why the output cannot change
|
||||
|
||||
Integer arithmetic is exact, so re-partitioning the work cannot change a result
|
||||
as long as every output element is computed from the same terms, computed the
|
||||
same way, in the same accumulator type, and no step's overflow behaviour changes.
|
||||
The kernels are written to that rule (stated in full in the `int_kernels` module
|
||||
documentation):
|
||||
|
||||
1. **Work is split only across output elements.** A thread owns whole tiles; a
|
||||
SIMD lane owns one output element. No sum is split between threads or lanes.
|
||||
2. **Each element adds the same terms in the reference's order**, from the same
|
||||
starting value, apart from terms that are zero. The fast kernel may add terms
|
||||
the reference skips because they are zero (a zero input, a zero-padding
|
||||
sample) and may omit terms the reference adds although they are zero (a tap
|
||||
whose weight is zero for every channel of a block). Adding or omitting a zero
|
||||
changes no partial sum, and a zero product cannot overflow.
|
||||
3. **A narrower type only where a bound over the actual operands proves the same
|
||||
integer.** 32×32→64-bit products require both operands to have `|v| ≤ i32::MAX`.
|
||||
An `i64` accumulator replaces an `i128` one only when
|
||||
`terms · max|w| · max|x| ≤ i64::MAX`. IGDN and channel attention use the
|
||||
analogous bounds.
|
||||
4. **A convolution takes the fast path only when no step of the reference can
|
||||
overflow.** The bound `terms · max|w| · max|x|` must fit the reference's
|
||||
accumulator (`i64`, or `i128` where the reference accumulates in `i128`).
|
||||
Otherwise the reference itself runs. This matters because the reference
|
||||
transposed convolutions accumulate into border positions and then crop them.
|
||||
The fast kernel never computes those positions, so only the reference can
|
||||
reproduce an overflow there. Under the bound nothing can overflow, so the SIMD
|
||||
lanes, which wrap where overflow-checked scalar code traps, are exact too.
|
||||
IGDN, channel attention, LayerNorm and the element-wise gates compute every
|
||||
value the reference computes, narrowing only under their own bounds. The
|
||||
per-output requantisation is the same expression in both, evaluated at the
|
||||
same kept positions.
|
||||
5. **Anything the fast path does not handle goes to the reference**: malformed
|
||||
geometry, short buffers, shift amounts outside the narrowed range, padding
|
||||
whose padded input would exceed the fast path's allocation cap or be out of
|
||||
proportion to what the reference reads and writes (rule 6), and a transposed
|
||||
geometry whose sizes the reference's own arithmetic cannot form. That is an
|
||||
uncropped plane `cout·(ho + 2·pad)·(wo + 2·pad)` beyond the allocation cap
|
||||
(the int_hyper, int_audio and JPEG-AI references allocate it before cropping,
|
||||
sized in unchecked `usize`), or an int_synth output size that overflows its
|
||||
reference's `i64` arithmetic. Such input gets the reference's own behaviour,
|
||||
in either profile.
|
||||
6. **The fast path's own work cannot overflow and stays in proportion to the
|
||||
reference's.** Its buffer sizes and indices are bounded by the checks of
|
||||
rule 5, and its worker-count estimates are widened or saturating. The
|
||||
reference performs none of these steps, so none of them may be the step at
|
||||
which a fast decode panics. What it allocates and loops over is bounded the
|
||||
same way. The zero-padded input copy may hold at most 16 times the elements
|
||||
the reference reads and writes (input, weights and output). The largest
|
||||
ratio is 0.9 over the committed vectors and 1.0 over the benchmark shapes
|
||||
(both measured with the check logged), and 7.2 over every shape the
|
||||
randomised convolution tests can draw (enumerated from their parameter
|
||||
ranges with the layout arithmetic mirrored), so none of those changes
|
||||
path. Nothing grows with the stride. A transposed convolution builds tap
|
||||
tables only for the (row, column) phases that hold output in a tile, so there
|
||||
are at most as many tables as the tile has output elements per channel, and
|
||||
their taps together number at most `cin · kh · kw`. A regular convolution
|
||||
lays out at most as many column phases as its padded row has columns.
|
||||
|
||||
What these rules give, precisely:
|
||||
|
||||
- **Integers.** For every input the reference returns from, every policy returns
|
||||
the same integer for every output element. In a release build (overflow checks
|
||||
off) that includes inputs whose accumulation wraps, at a kept output position
|
||||
or only at a border position a reference transposed convolution computes and
|
||||
then crops.
|
||||
- **Panics in an overflow-checked (debug) build.** A fast policy panics on the
|
||||
inputs where the reference panics, including an overflow confined to a cropped
|
||||
border position and one in the reference's own size arithmetic (rule 5). Only
|
||||
the panic itself can differ: another element may fail first, and a panic
|
||||
raised on a spawned worker reaches the caller as the thread scope's own panic
|
||||
("a scoped thread panicked"), not the reference's arithmetic-overflow message.
|
||||
- **Outside the claim: an allocation that fails.** The two paths allocate
|
||||
different buffers. Besides the output both allocate, the fast path allocates
|
||||
the padded input copy, per-tile tap tables and per-row accumulators (each
|
||||
bounded by rule 6), and the int_hyper, int_audio and JPEG-AI references
|
||||
allocate a transposed convolution's uncropped plane. While the fast path
|
||||
runs, the padded copy, the output and that plane each hold at most the
|
||||
allocation cap (2³² elements on 64-bit targets, 2²⁸ on 32-bit ones). Still,
|
||||
where memory runs out, or on a 32-bit target where one buffer would exceed
|
||||
`isize::MAX` bytes (which Rust reports as a "capacity overflow" panic), one
|
||||
path can fail where the other does not.
|
||||
|
||||
What was checked (the claim rests on these, not on the argument alone):
|
||||
|
||||
- **Element-by-element comparison with the reference** over randomised shapes
|
||||
(strides, paddings, output paddings, groups, kernels smaller than the stride,
|
||||
inputs narrower than the stride) and operand ranges. The ranges cover the 32-bit
|
||||
path, the 64-bit path, the `i128` path, both sides of every narrowing bound, and
|
||||
zero-heavy inputs. Every comparison runs under one worker, two explicit
|
||||
multi-threaded tilings, and the automatic count. Files: the `kernel_equivalence`
|
||||
and `igdn_equivalence` test modules in each kernel's file.
|
||||
- **Overflow parity.** Three kinds of operand are tested: ones that overflow the
|
||||
reference's accumulator at a kept output position (the `int_hyper` transposed
|
||||
and regular convolutions, including rows wide enough for the SIMD lanes to
|
||||
engage), ones that overflow it only at a border position a transposed
|
||||
convolution computes and then crops (the `int_hyper` and `int_audio`
|
||||
transposed convolutions, `i64` accumulator, and the JPEG-AI one, `i128`), and
|
||||
operands beyond the bound (`int_synth`). With overflow checks on, the tests
|
||||
require every policy to panic where the reference panics; with them off, the
|
||||
same integers. All of these run in both profiles. The border case was missed
|
||||
in the kernels' first form, whose tests covered overflow at kept positions
|
||||
only. Review found it with a probe through the public
|
||||
`int_hyper_synth::run_net` and `int_transformer::conv_transpose`, and with
|
||||
the `int_audio` cropped-border test run against that form. With overflow
|
||||
checks on, the reference panicked in all three and every fast policy returned
|
||||
a value. In release, all policies returned the same integers. Rule 4's
|
||||
routing and the border tests were added then. With rule 4's routing deleted
|
||||
in a scratch copy of the sources, the three border tests fail in an
|
||||
overflow-checked build, as does the SIMD-width kept-position test.
|
||||
- **Rule 6.** A transposed geometry with padding and output padding far beyond
|
||||
the kernel (a 1×1 output whose padded input would have about 2⁴⁰ elements) is
|
||||
now left to the reference, and an explicit worker count of any size, up to
|
||||
`usize::MAX`, gives the same integers as one worker. Both tests fail on the
|
||||
kernels as they stood before rule 6: the geometry was accepted, and the worker
|
||||
count overflowed `4 · threads` in an overflow-checked build. The channel
|
||||
attention's worker-count estimate `heads · c² · hw` is now widened before it is
|
||||
multiplied, because on a 32-bit target it can exceed `usize` (for example
|
||||
128 channels per head over 2¹⁸ positions). That last change is argued, not
|
||||
tested: the inputs that reach it are too large for a unit test.
|
||||
- **Rule 6, stride and padding.** Before this rule's final form, the fast
|
||||
transposed kernel built one tap table for every (row, column) phase, `sh · sw`
|
||||
tables per tile, whether or not a phase held output. A 1×1 input upsampled by
|
||||
a huge stride is a 1×1 output that the reference computes in a handful of
|
||||
operations. Review found, with a probe through the public API, that the fast
|
||||
path panicked with "capacity overflow" at stride 2²⁹ (which fits a WHM1
|
||||
model's `u32` stride field and the JSON specs), and on the multiplication
|
||||
`sh · sw` at stride 2³² in an overflow-checked build. At 2¹⁶ the kernel
|
||||
committed memory for 2³² tables: the geometry test, run against that earlier
|
||||
kernel, was killed after more than a minute, and swap grew by 17 GiB
|
||||
while it ran. Until the tables were limited to phases holding output, the
|
||||
"Integers" statement above did not hold for such geometry under the default
|
||||
policy. Decoded bytes of valid streams were never affected. The regular
|
||||
convolution's padded row grew with the stride too, bounded only by the
|
||||
allocation cap. Now
|
||||
`tests/int_kernels_geometry.rs` runs these cases through the public decoders
|
||||
(`int_hyper_synth::run_net`, `int_synth::decode_from_spec` and the JPEG-AI
|
||||
`conv_transpose` / `conv_grouped`) under every policy, and checks each against
|
||||
a hand computation:
|
||||
- a 1×1 input upsampled by strides from 1 to `usize::MAX`;
|
||||
- a 5×5 input whose stride reaches or exceeds its padded width;
|
||||
- padding of 2¹², far beyond a 3×3 or 1×1 kernel.
|
||||
|
||||
`int_audio` has the same stride test beside its kernel. A randomised test
|
||||
compares both convolutions element by element at strides up to 16, beyond
|
||||
kernel and input. There, row bands hold fewer rows than the stride as well as
|
||||
more, and many phases receive no tap. Unit tests check that geometry whose
|
||||
padded copy is out of proportion (about 2²⁴ and 2²⁷ elements where the
|
||||
reference touches 11 and 6) goes to the reference.
|
||||
- **Rule 6, at the factor.** The committed vectors, the benchmark shapes and
|
||||
the randomised tests all stay below half the factor (the ratios above), so
|
||||
dedicated tests pin both sides of it. Geometry whose padded copy is exactly
|
||||
16 times what the reference touches is accepted, and the fast kernel runs.
|
||||
Geometry one element, two elements, one padding step or one stride step
|
||||
beyond it is refused, and the reference runs. Each refused shape is small
|
||||
and well-formed and is measured with the kernels' own layout arithmetic, and
|
||||
a sweep over about 7,400 small regular and transposed shapes checks that
|
||||
acceptance is exactly "padded ≤ 16 · touched", equality included. So nothing
|
||||
but the proportion refuses these shapes. Both paths give the same integers,
|
||||
so the output cannot show which one ran. Test builds therefore count
|
||||
fast-kernel entries per thread. The count shows the fast kernel running once
|
||||
at the factor and never beyond it, through each of the six dispatching
|
||||
functions: `int_hyper_synth` regular and transposed, `int_synth`,
|
||||
`int_audio`, and the JPEG-AI `conv_grouped` and `conv_transpose`. In each
|
||||
case the output equals the reference's and a hand computation. Three
|
||||
mutations were each run against these tests, and each made all five fail: a
|
||||
factor of 17, the proportion check removed, and `>` changed to `≥`.
|
||||
- **Rule 5, the reference's own size arithmetic.** Review found, again with a
|
||||
probe through the public API, that the fast geometry check bounded only what
|
||||
the fast path lays out, not two sizes the references compute in unchecked
|
||||
arithmetic. One is the uncropped plane of the int_hyper, int_audio and
|
||||
JPEG-AI transposed convolutions, in `usize`. The other is int_synth's output
|
||||
size, in `i64`. An input of two or more rows with a stride and padding of
|
||||
2⁶² overflows them. So does, within a WHM1 model's `u32` fields, a stride of
|
||||
2³² − 1 with padding of 2³¹ − 1. The reference then panics with overflow
|
||||
checks on. With them off, the uncropped-plane references index out of a
|
||||
wrapped buffer. Every fast policy returned a value, the mathematically correct
|
||||
one, so under the default policy a model the reference cannot decode decoded.
|
||||
No decoded byte changed for any input the reference returns from. The earlier
|
||||
geometry tests missed this because their huge-stride transposed cases all
|
||||
used a 1×1 input, where `(h − 1)·stride` is 0. Rule 5 now covers both sizes.
|
||||
Beyond the cap the reference runs and fails, or not, in its own way. A plane
|
||||
that is merely huge also goes to the reference, so the fast path does not
|
||||
decode what the reference could not allocate. int_synth's dispatch
|
||||
evaluates its reference's `i64` expression step by step, and runs the fast
|
||||
kernel only where that expression neither overflows nor disagrees with the
|
||||
fast geometry. `tests/int_kernels_geometry.rs` runs these inputs through
|
||||
`int_hyper_synth::run_net`, the JPEG-AI `conv_transpose` and
|
||||
`int_synth::decode_from_spec`, and requires the reference's outcome under
|
||||
every policy. That outcome is a panic, except for int_synth with overflow
|
||||
checks off, where the wrapped arithmetic still yields the correct image.
|
||||
`int_audio` has the 1-D form beside its kernel: a buffer size that overflows
|
||||
`usize`, and one no allocation can hold. A unit test pins the bound at the
|
||||
cap and one row past it, where nothing else refuses the geometry. Mutations:
|
||||
- With the plane bound removed, the int_hyper, JPEG-AI, int_audio and bound
|
||||
tests fail in both profiles. The int_synth tests still pass on the `i64`
|
||||
check alone.
|
||||
- With both removed, the int_synth geometry test fails with overflow checks
|
||||
on. With them off it passes, because every policy then returns the same
|
||||
image. The int_synth unit test, which also counts fast-kernel entries,
|
||||
fails in both profiles.
|
||||
- With the `i64` check removed and the bound kept, every test passes. At the
|
||||
current cap the bound alone keeps that expression from overflowing, and
|
||||
the check keeps int_synth's rule from depending on the cap's value.
|
||||
|
||||
The JPEG-AI and int_hyper u32-width cases were also run alone against the
|
||||
unbounded geometry, and failed there in both profiles.
|
||||
- **Every committed vector that reaches the kernels, under every policy**
|
||||
(`tests/int_kernels_vectors.rs`): int_hyper, int_audio, int_splat, int_world
|
||||
(one and two agents), the neural-keyframe video and the four
|
||||
`neural-conformance` int_synth vectors, plus the real-footage video. Each
|
||||
decodes byte-identically under the reference and each fast policy, and
|
||||
reproduces its frozen hash; the video checks also compare the frame count.
|
||||
The plain int_video vector (block motion compensation and rANS only) never
|
||||
reaches the kernels and is not in this test. The int_synth photo vector also
|
||||
still equals its independent integer-forward reference `ref_rgb.bin`.
|
||||
- **CI.** The `wai-conformance` workflow runs the lib tests with every learned
|
||||
integer codec's feature (so every `kernel_equivalence` and `igdn_equivalence`
|
||||
module, the overflow-parity tests and the rule-6 tests),
|
||||
`tests/int_kernels_vectors.rs` and `tests/int_kernels_geometry.rs`, once with
|
||||
overflow checks on and once in
|
||||
release, on each leg of its matrix (Linux x86_64 and arm64, macOS arm64,
|
||||
Windows x86_64). `wai-byte-equality` and `wai-formal-verification` also fire
|
||||
on a change to `int_kernels.rs`, which their crates `#[path]`-include.
|
||||
- **The tests can fail.** Deliberate defects were injected one at a time: an
|
||||
off-by-one tap offset, a wrong row band, a skipped SIMD tail, a swapped NEON
|
||||
lane, a bypassed narrowing or overflow gate, a missing IGDN clamp, a wrong
|
||||
phase in the strided pad, a mis-indexed attention row. The tests caught each
|
||||
one. Two initially survived, an IGDN sum-bound bypass and a bypassed SIMD
|
||||
overflow gate, and each led to a new test case that catches it. Separately,
|
||||
review found a real defect in the new 64-bit IGDN path: with an all-zero
|
||||
`gamma`, the square of a huge input was formed in `i64`, where the reference
|
||||
squares in `i128`. It was fixed, with a test that fails without the fix.
|
||||
- **wasm.** `wai-web` builds for `wasm32-unknown-unknown`. That build caught one
|
||||
real 32-bit defect in the new code (an allocation cap that overflowed a 32-bit
|
||||
`usize`), since fixed. The equivalence tests and every committed vector also
|
||||
pass on `wasm32-wasip1` under wasmtime 49, exercising the portable,
|
||||
single-threaded, 32-bit-`usize` path. Skipped there: tests that need
|
||||
unwinding or a temp directory, and one pre-existing rANS test whose data
|
||||
generator overflows a 32-bit `usize`.
|
||||
- **Formal verification unchanged.** `cargo kani` in `wai-verify` (Kani 0.67):
|
||||
every harness still verifies. The harnesses do not cover the convolution or
|
||||
IGDN kernels, fast or reference (SPEC §7, *machine-checked totality*). The
|
||||
fast kernels' equivalence is **tested, not proved**.
|
||||
|
||||
## 3. Method
|
||||
|
||||
- **Machine:** Apple M5 Max, 18 logical cores, macOS 27.0; rustc 1.98.0. macOS
|
||||
reports the cores as two performance levels (`hw.perflevelN.name`): 6 "Super"
|
||||
and 12 "Performance". It reports no efficiency cores. The harness prints the
|
||||
levels as the OS names them.
|
||||
- **Revision:** measured once rule 6 had its final form, on the kernels as
|
||||
they stand apart from one later routing check. That check, rule 5's bound on
|
||||
the reference's size arithmetic, adds a few checked multiplications per
|
||||
transposed convolution and sends no benchmarked shape down a different path:
|
||||
every benchmarked uncropped plane has fewer than 2²⁴ elements, against a cap
|
||||
of 2³², and every int_synth size is far below `i64::MAX`. That check was not
|
||||
itself timed. The other later changes are a test-build counter (which
|
||||
compiles to nothing outside tests), tests, comments and lint declarations.
|
||||
- **Load and runs:** the machine was shared with other workloads throughout,
|
||||
and the load average read between 6.2 and 10.5 at the start and end of the
|
||||
runs. Four complete runs of the benchmark and four of the per-layer profile
|
||||
were made back to back. From run to run, the best time for one case under
|
||||
one policy varied by 2% to 67%, mostly by 10% to 30%. The largest variations
|
||||
were on all cores, up to 67% for the all-core IGDN kernel with Q12
|
||||
activations. The tables give, for each case and policy, the best over all
|
||||
rounds of the four runs. They are best-case figures for this machine under
|
||||
this load, not typical ones. Each speed-up is the ratio of those best times,
|
||||
the same ratio the harness reports within one run. The range of speed-ups
|
||||
within single runs is given under the tables. Nothing here compares
|
||||
revisions.
|
||||
- **Build:** `wai-rs` release profile (`opt-level = 3`, LTO, `codegen-units = 1`),
|
||||
`--no-default-features --features neural_int,video_int,audio_int,world_int,splat_int`.
|
||||
- **Harness:** `src/bin/wai_decode_bench.rs`. For each case, three policies run
|
||||
**interleaved in one process on the same inputs**, rotating the order each
|
||||
round: the reference kernels, the fast kernels on one thread (`1T`), and the
|
||||
fast kernels with the automatic worker count (`auto`, up to 18 here). The
|
||||
harness reports the best time over the rounds (medians are in its output).
|
||||
Every run's output is hashed, and the harness aborts rather than report a time
|
||||
for an output that differs from the reference's. In all eight runs, every
|
||||
output under every policy matched the reference's.
|
||||
- **Inputs:** self-contained. Kernel and whole-decoder cases use deterministic
|
||||
pseudo-random weights at the real layer shapes. The int_hyper `g_s` is the
|
||||
mbt2018-mean quality-6 topology the export tool emits (M = 320, N = 192;
|
||||
int12 weights, `afrac` 12 / `gf` 16 / `shift` 24, like the committed models).
|
||||
The JPEG-AI synthesis uses the exact weight shapes of the reference model_y and
|
||||
model_uv nets, with 16-bit weights and Q16 activations. The `vector` cases are
|
||||
the committed conformance vectors, and not all of them are trained models:
|
||||
- **Trained:** int_synth photo, int_world, int_splat and int_audio. Their
|
||||
export tools train the float model, then quantize it. int_synth and int_audio
|
||||
were trained on captured content (a photograph, an audio clip); int_world and
|
||||
int_splat on procedurally generated content.
|
||||
- **Random weights:** the int_hyper tiny conditional codec. Its generator,
|
||||
`tools/wai_int_hyper_tiny_vector.py`, draws the integer weights from a seeded
|
||||
generator.
|
||||
- **No learned model:** int_video real footage, which is integer block motion
|
||||
compensation plus rANS.
|
||||
|
||||
So no trained int_hyper model was timed: the `g_s` row uses synthetic weights
|
||||
and the int_hyper vector uses random ones. Time depends on shape and on operand
|
||||
*width*, not on trained values, with one exception that is visible below:
|
||||
activation magnitude selects the IGDN arithmetic path.
|
||||
|
||||
## 4. Results
|
||||
|
||||
### Whole decoders
|
||||
|
||||
| decode | reference | fast, 1 thread | fast, auto | speed-up 1T / auto |
|
||||
|---|---:|---:|---:|---:|
|
||||
| int_hyper `g_s`, mbt2018-mean q6 topology (synthetic weights), 16² latent → 256² RGB | 2961 ms (0.022 MPix/s) | 299 ms (0.22 MPix/s) | 35.3 ms (1.86 MPix/s) | 9.9× / 84.0× |
|
||||
| JPEG-AI synthesis, luma + chroma nets (reference topology, synthetic weights), 256² | 5556 ms (0.012 MPix/s) | 452 ms (0.14 MPix/s) | 79.7 ms (0.82 MPix/s) | 12.3× / 69.7× |
|
||||
| int_synth trained photo vector, 128² | 14.3 ms | 5.25 ms | 3.69 ms | 2.7× / 3.9× |
|
||||
| int_world trained interactive clip vector, 9 frames | 28.5 ms (316 frames/s) | 3.82 ms (2356 frames/s) | 3.76 ms (2396 frames/s) | 7.5× / 7.6× |
|
||||
| int_splat trained attribute vector, 1024 Gaussians | 1.99 ms | 0.29 ms | 0.27 ms | 6.9× / 7.5× |
|
||||
| int_audio trained clip vector, 8192 samples | 3.94 ms | 0.93 ms | 0.89 ms | 4.2× / 4.4× |
|
||||
| int_hyper tiny conditional vector (random weights), 256 px | 14.7 µs | 7.9 µs | 7.8 µs | 1.9× / 1.9× |
|
||||
| int_video half-pel real-footage vector, 8 frames 96×96 (no learned model: motion compensation + rANS only) | 2.38 ms | 2.32 ms | 2.36 ms | 1.0× / 1.0× |
|
||||
|
||||
Within single runs, the `g_s` speed-ups ranged from 8.5× to 10.5× on one thread
|
||||
and from 71.5× to 84.0× on all cores. For the JPEG-AI synthesis the ranges were
|
||||
11.7× to 13.0× and 57.0× to 69.7×.
|
||||
|
||||
The automatic worker count follows each layer's work: one worker per 2²⁰
|
||||
multiply-accumulates, up to the core count. Logged during the vector tests on
|
||||
this machine (18 cores), it stayed at one to three workers per call for the
|
||||
int_hyper tiny, int_audio, int_splat, int_world and neural-keyframe video
|
||||
decodes, so `auto` adds little over one thread there. The int_synth photo layers
|
||||
are larger: `auto` ran 10 to 18 workers on its four transposed convolutions,
|
||||
which is why that row gains more from all cores. The int_video decode has no
|
||||
learned kernel. Its time is the rANS entropy decode, which is sequential by
|
||||
construction and was not changed.
|
||||
|
||||
### Kernels (layer shapes from the two decoders above, and one regular convolution)
|
||||
|
||||
| kernel | reference | fast, 1 thread | fast, auto |
|
||||
|---|---:|---:|---:|
|
||||
| convT 192→192, 5×5, stride 2, 64²→128² (int_hyper `g_s`) | 2.5 GMAC/s | 31.7 GMAC/s | 289 GMAC/s |
|
||||
| conv 192→192, 3×3, 64² (a regular convolution at the `g_s` width; not a layer of either decoder) | 4.0 | 31.7 | 273 |
|
||||
| IGDN c = 192, 128², Q12 activations (x² fits 32 bits) | 1.1 | 9.0 | 106 |
|
||||
| IGDN c = 192, 128², x² > 2³¹ (64-bit path) | 1.1 | 6.5 | 58 |
|
||||
| conv 128→128, 3×3, 64² (JPEG-AI CAB) | 1.9 | 32.0 | 264 |
|
||||
| conv 1×1 128→1024, 64² (JPEG-AI GDFN in) | 1.9 | 22.7 | 207 |
|
||||
| depthwise 3×3 ×1024, 64² (JPEG-AI GDFN) | 1.7 | 6.6 | 17.1 |
|
||||
| conv 128→128, 3×3, stride 2, 128²→64² | 1.8 | 27.8 | 197 |
|
||||
| convT 128→128, 3×3, stride 2, 64²→128² (reference accumulates in `i128`) | 1.4 | 29.6 | 257 |
|
||||
|
||||
### Where the time goes (profile)
|
||||
|
||||
`wai_decode_bench --layers` times each layer of the synthetic `g_s` on its real
|
||||
intermediate activations under each policy, so the time is attributed to layers
|
||||
exactly. The total is the sum of the layers' best times:
|
||||
|
||||
| `g_s` layer, 256² decode | work | reference | fast 1T | fast auto |
|
||||
|---|---|---:|---:|---:|
|
||||
| convT 320→192, 16²→32² | 0.39 GMAC | 159 ms | 14.1 ms | 2.02 ms |
|
||||
| IGDN, 32² (max \|x\| 36562) | 0.04 GMAC | 29.1 | 4.1 | 0.66 |
|
||||
| convT 192→192, 32²→64² | 0.94 GMAC | 380 | 30.7 | 3.5 |
|
||||
| IGDN, 64² (max \|x\| 42256) | 0.15 GMAC | 118 | 16.2 | 1.8 |
|
||||
| convT 192→192, 64²→128² | 3.77 GMAC | 1504 | 119 | 12.7 |
|
||||
| IGDN, 128² (max \|x\| 57793: x² > 2³¹) | 0.60 GMAC | 475 | 92.0 | 9.5 |
|
||||
| convT 192→3, 128²→256² | 0.24 GMAC | 93 | 10.0 | 2.65 |
|
||||
| **total** | | **2759** | **287** | **32.8** |
|
||||
|
||||
With the reference kernels, transposed convolution takes 77% of the decode and
|
||||
IGDN 23%. With the fast kernels the split is about 61/39 on one thread and
|
||||
63/37 on all cores. The transposed convolution 192→192 at 64²→128² is still the
|
||||
largest layer, on one thread and on all cores. The last IGDN is second. On one
|
||||
thread it is the slowest layer per MAC, at about 6.6 GMAC/s against about 31.6
|
||||
for that convolution. Its activations exceed the 32-bit-square range, so it
|
||||
takes the 64-bit path.
|
||||
|
||||
For the JPEG-AI synthesis, sampling the all-core decode showed the channel
|
||||
attention, the GDFN element-wise gate, the channel LayerNorm and the
|
||||
element-wise gates of the synthesis running serially behind the parallel
|
||||
convolutions. Those were parallelised. Sampled symbol attribution in an LTO
|
||||
build is approximate, so the kernel table is the better guide.
|
||||
|
||||
## 5. Limits
|
||||
|
||||
- **One shared machine, best of four runs.** The load average was 6.2–10.5
|
||||
throughout, and a case's best time moved by up to 67% between runs (§3). The
|
||||
tables are the best observed, not the typical, and a quieter machine or
|
||||
another one will give other numbers.
|
||||
- **One machine, one architecture measured.** The NEON paths are aarch64-only.
|
||||
x86_64 runs the portable kernels: it compiles, but it was **not executed** here,
|
||||
because this machine cannot run x86_64 binaries. wasm runs the portable kernels
|
||||
single-threaded. It was tested for identical output, not timed.
|
||||
- **No trained int_hyper model was timed.** The `g_s` row uses synthetic
|
||||
weights, and the int_hyper tiny vector uses random ones. Synthetic weights
|
||||
reproduce shapes and operand widths, not trained activations, and activation
|
||||
size decides the IGDN path (§4). So a trained int_hyper model, which was not
|
||||
measured here, may run faster or slower. The JPEG-AI row also uses synthetic
|
||||
weights (see below). The trained models that were timed are the small
|
||||
committed vectors int_synth photo, int_world, int_splat and int_audio.
|
||||
- **JPEG-AI:** the synthesis nets only. The entropy stage and the float
|
||||
post-filter chain are unchanged and were not timed. The real exported JPEG-AI
|
||||
weights were not available in this environment, so the real-weight end-to-end
|
||||
tests skipped. The JPEG-AI kernels' equivalence rests on the randomised kernel
|
||||
and whole-TAB tests plus the synthetic-weight decode being identical under every
|
||||
policy on every run.
|
||||
- **Not proved.** The equivalence is argued by construction and tested
|
||||
(including by injected defects), not machine-proved.
|
||||
|
||||
## Reproduce
|
||||
|
||||
```sh
|
||||
cd wai/wai-rs
|
||||
F=neural_int,video_int,audio_int,world_int,splat_int
|
||||
cargo run --release --no-default-features --features $F --bin wai_decode_bench # all cases
|
||||
cargo run --release --no-default-features --features $F --bin wai_decode_bench -- --layers # per-layer profile
|
||||
cargo test --no-default-features --features $F --lib --test int_kernels_vectors --test int_kernels_geometry # equivalence, overflow checks on
|
||||
cargo test --release --no-default-features --features $F --lib --test int_kernels_vectors --test int_kernels_geometry # equivalence, release
|
||||
```
|
||||
|
||||
The tables in §4 come from running the first two commands four times each and
|
||||
taking, per case and policy, the smallest of the four reported best times.
|
||||
13
wai/tools/jpegai-e2e-v1.sha256
Normal file
13
wai/tools/jpegai-e2e-v1.sha256
Normal file
|
|
@ -0,0 +1,13 @@
|
|||
f3b325c583c6b7613ccddc2d89bfaa50e0e8314ceecd8d82fc4df7fb159b6083 hs_uv.json
|
||||
1d193f15b5d8d1531ef9ffef8676904d76cfedbd1d2df931d581dfb26daa2688 hs_y.json
|
||||
2c5f6f12fcf9a36489553e0e02e95684c339d82bab64cb7bae76a4635f9ec56d hyp_uv.json
|
||||
32ab3a153e33ded902b1d8a6c413047634bc39e8b74525dfbb40eded0763433d hyp_y.json
|
||||
bd3aa0c23294a462895416eda38399468781163858229c386deaba58def5a1c3 icci.json
|
||||
562740aba18ba4dccdc007b3d54a669ecd18feac9d0bd1fd7a29e66a0ba498ee mcm0_y.json
|
||||
fbd4b7c05e81cd32df4db4e2f6dfa4d6e3fc2edb6c21150f80ef5ac04765228d mcm1_y.json
|
||||
b42d6fa96bf865dd2db47955aba14d8d65b9152c941be3ee3e92eae5643834a7 mcm2_y.json
|
||||
e06766d0a7bff16188a0e96297d6e498e28fbaf7850f68b0a29447a9391c9d9d mcm3_y.json
|
||||
5353b76f25e59f52d85adc896c9dc6b4aa8552589e3e4e1d97b8f29692e37b13 norm.json
|
||||
8ddf8fe9dd15569910ad5755a5495526df959835bcf86a259d06ab6445f02b4e synth_uv.json
|
||||
4f2c9ae807ea22197adaae0263c9d6d4d53ef7d4fccb81848b7bb4da9ef0ea69 synth_y.json
|
||||
99840e5f472423f5a1261c9ec219d12dc910e9623d752672555dfea69ea4a097 test256.bin
|
||||
25
wai/tools/requirements-sota-int.txt
Normal file
25
wai/tools/requirements-sota-int.txt
Normal file
|
|
@ -0,0 +1,25 @@
|
|||
# Pinned Python stack for regenerating the int_hyper integer-codec vectors (the
|
||||
# pretrained mbt2018_mean, exported to integers; tools/wai_sota_vectors.sh). The vectors are regenerated rather than committed,
|
||||
# so the export must run against a fixed stack: an upgrade that changes the
|
||||
# pretrained weights or the float analysis transform must arrive as a reviewed
|
||||
# change here, not silently on the next CI run.
|
||||
#
|
||||
# torch/torchvision: install from the CPU wheel index first on Linux, or pip
|
||||
# resolves the CUDA builds (several GB). `==2.14.0` matches `2.14.0+cpu`.
|
||||
# pip install torch==2.14.0 torchvision==0.29.0 --index-url https://download.pytorch.org/whl/cpu
|
||||
# pip install -r tools/requirements-sota-int.txt
|
||||
torch==2.14.0
|
||||
torchvision==0.29.0
|
||||
# BSD-3-Clause-Clear. Supplies the mbt2018_mean architecture and the URL of its
|
||||
# published checkpoint, mbt2018-mean-6-a19628ab.pth.tar. The checkpoint is NOT
|
||||
# verified by compressai or torch.hub (the zoo calls torch.hub without
|
||||
# check_hash), so wai_sota_int_synth_export.py pins its full SHA-256 and refuses
|
||||
# to load any other bytes.
|
||||
compressai==1.2.8
|
||||
# compressai 1.2.8 requires numpy<2.
|
||||
numpy==1.26.4
|
||||
pillow==12.3.0
|
||||
# Source image only: `skimage.data.astronaut()` ships inside the wheel and is a
|
||||
# NASA photograph documented there as "No known copyright restrictions, released
|
||||
# into the public domain" (the function's docstring).
|
||||
scikit-image==0.26.0
|
||||
|
|
@ -1,12 +1,13 @@
|
|||
#!/usr/bin/env python3
|
||||
"""Generate a TINY, self-contained `wai.neural.int_hyper` conformance vector.
|
||||
|
||||
The SOTA int_hyper vector needs a ~23 MB PTQ model, too large to commit — so the
|
||||
Rust conformance test SKIPS without it, and CI never exercises the decode. This
|
||||
emits a SMALL synthetic conditional codec (deterministic random integer weights)
|
||||
that drives the SAME decode path end to end — factorized z -> integer hyper-
|
||||
decoder h_s -> integer-sigma entropy bucketing -> integer rANS -> integer
|
||||
synthesis g_s -> RGB — with a committed expected FNV-1a. A few KB, so it commits
|
||||
The pretrained int_hyper vector needs a ~23 MB PTQ model, too large to commit —
|
||||
so it is regenerated per run (tools/wai_sota_vectors.sh, wai-sota-vectors.yml)
|
||||
and checked on one machine at a time. This emits a SMALL synthetic conditional
|
||||
codec (deterministic random integer weights) that drives the SAME decode path
|
||||
end to end — factorized z -> integer hyper-decoder h_s -> integer-sigma entropy
|
||||
bucketing -> integer rANS -> integer synthesis g_s -> RGB — with a committed
|
||||
expected FNV-1a. A few KB, so it commits
|
||||
and the int_hyper decode is CI-verifiable with no external artifacts.
|
||||
|
||||
python tools/wai_int_hyper_tiny_vector.py <out-dir>
|
||||
|
|
@ -20,7 +21,7 @@ import numpy as np
|
|||
from compressai._CXX import pmf_to_quantized_cdf
|
||||
from compressai.ans import RansEncoder
|
||||
|
||||
# Canonical integer-net runner + helpers (the same the SOTA export uses).
|
||||
# Canonical integer-net runner + helpers (the same the pretrained-model export uses).
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
from wai_sota_int_synth_export import run_int_net, det_build_index_vec, fnv1a, model_to_bin
|
||||
|
||||
|
|
|
|||
|
|
@ -1,17 +1,20 @@
|
|||
#!/usr/bin/env python3
|
||||
"""Encode a REAL 256x256 int_hyper-keyframe clip: the I-frame is the SOTA learned
|
||||
integer codec int_hyper (its byte-exact 256x256 reconstruction, dumped by the
|
||||
`int_hyper_decode_frame` Rust example), and the P-frames pan it with integer block
|
||||
motion compensation + integer-quantized residual. Skip-if-absent: needs the SOTA
|
||||
int_hyper vector dir (default /tmp/sota_int) holding frame0.rgb.
|
||||
"""Encode a REAL 256x256 int_hyper-keyframe clip: the I-frame is the integer
|
||||
learned codec int_hyper (the pretrained mbt2018_mean exported to integers; its
|
||||
byte-exact 256x256 reconstruction, dumped by the `int_hyper_decode_frame` Rust
|
||||
example), and the P-frames pan it with integer block motion compensation +
|
||||
integer-quantized residual. Needs the int_hyper vector dir (default
|
||||
/tmp/sota_int) holding frame0.rgb, and exits non-zero without it.
|
||||
tools/wai_sota_vectors.sh runs the whole chain.
|
||||
|
||||
cargo run --release --example int_hyper_decode_frame --features neural_int -- <dir> <dir>/frame0.rgb
|
||||
cargo run --release --no-default-features --features neural_int --example int_hyper_decode_frame -- <dir> <dir>/frame0.rgb
|
||||
python tools/wai_int_video_sota_keyframe_encode.py <dir>
|
||||
|
||||
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::decode_pframes and checks these, skipping when the dir is absent.
|
||||
int_video::decode_pframes and checks these (skipping loudly when they are absent,
|
||||
failing instead under WAI_REQUIRE_SOTA_VECTORS=1).
|
||||
"""
|
||||
import sys, os, json, base64
|
||||
import numpy as np
|
||||
|
|
@ -75,8 +78,9 @@ def main(d, prefix='video_sota'):
|
|||
global H, W
|
||||
fp = f'{d}/frame0.rgb'
|
||||
if not os.path.exists(fp):
|
||||
print(f'skip: {fp} absent (run the int_hyper_decode_frame example first)')
|
||||
return
|
||||
# A producer that exits 0 without producing lets every consumer downstream
|
||||
# skip in turn, and the whole chain reports green having checked nothing.
|
||||
sys.exit(f'error: {fp} absent (run the int_hyper_decode_frame example first)')
|
||||
raw = open(fp, 'rb').read()
|
||||
H = W = int(round((len(raw) / 3) ** 0.5)) # square frame inferred from the dump
|
||||
frame0 = np.frombuffer(raw, np.uint8).reshape(H, W, 3).astype(np.int64)
|
||||
|
|
|
|||
54
wai/tools/wai_jpegai_fetch_e2e_inputs.sh
Executable file
54
wai/tools/wai_jpegai_fetch_e2e_inputs.sh
Executable file
|
|
@ -0,0 +1,54 @@
|
|||
#!/usr/bin/env bash
|
||||
# Fetch the JPEG-AI whole-decoder test inputs from WAI's content-versioned models
|
||||
# root and check every byte against the SHA-256 manifest committed next to this
|
||||
# script (tools/jpegai-e2e-v1.sha256).
|
||||
#
|
||||
# tools/wai_jpegai_fetch_e2e_inputs.sh [out-dir] # default /tmp/jpegai_e2e
|
||||
#
|
||||
# Then run the tests that read them (tests/jpegai_e2e_full.rs,
|
||||
# tests/jpegai_reconstruction_binding.rs) with
|
||||
# WAI_JPEGAI_CODESTREAM=<out-dir>/test256.bin
|
||||
# WAI_JPEGAI_WEIGHTS_DIR=<out-dir>
|
||||
# WAI_JPEGAI_NORM=<out-dir>/norm.json
|
||||
#
|
||||
# What is fetched: the eleven integer/PTQ weight files the decoder runs
|
||||
# ({hs,hyp,synth}_{y,uv}.json, mcm{0..3}_y.json, icci.json), norm.json (normative
|
||||
# tables plus the reference decoder's RGB output for the test picture) and
|
||||
# test256.bin (the test codestream). They were extracted by
|
||||
# tools/wai_jpegai_build_e2e_inputs.py from the JPEG-AI reference-software assets,
|
||||
# which are not in this repository; the models root publishes the extraction under
|
||||
# an immutable versioned path (jpegai/v1/). The manifest pins those exact bytes, so
|
||||
# a changed or truncated file at the URL fails here instead of changing what the
|
||||
# tests check.
|
||||
#
|
||||
# WAI_JPEGAI_INPUTS_URL overrides the base URL (a mirror holding the same files).
|
||||
set -euo pipefail
|
||||
|
||||
here="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
manifest="$here/jpegai-e2e-v1.sha256"
|
||||
base="${WAI_JPEGAI_INPUTS_URL:-https://wai.transaction.science/models/jpegai/v1}"
|
||||
out="${1:-/tmp/jpegai_e2e}"
|
||||
|
||||
mkdir -p "$out"
|
||||
out="$(cd "$out" && pwd)"
|
||||
|
||||
if command -v sha256sum >/dev/null 2>&1; then
|
||||
sha() { sha256sum "$@"; }
|
||||
else
|
||||
sha() { shasum -a 256 "$@"; }
|
||||
fi
|
||||
|
||||
while read -r want name; do
|
||||
[ -n "$name" ] || continue
|
||||
curl --proto '=https' --tlsv1.2 -fsSL --retry 3 -o "$out/$name.part" "$base/$name"
|
||||
got="$(sha "$out/$name.part" | cut -d' ' -f1)"
|
||||
if [ "$got" != "$want" ]; then
|
||||
echo "error: $base/$name has sha256 $got, pinned $want" >&2
|
||||
rm -f "$out/$name.part"
|
||||
exit 1
|
||||
fi
|
||||
mv "$out/$name.part" "$out/$name"
|
||||
echo "ok $want $name"
|
||||
done < "$manifest"
|
||||
|
||||
echo "JPEG-AI e2e inputs verified in $out"
|
||||
|
|
@ -1,17 +1,58 @@
|
|||
#!/usr/bin/env python3
|
||||
"""Export a pretrained SOTA-class learned-codec SYNTHESIS as an INTEGER-exact
|
||||
decoder and prove it preserves quality: per-channel int12 weights + fixed-point
|
||||
activations + integer IGDN, all i64 — byte-identical on every machine
|
||||
(NeuralIntegerExact), at the float codec's PSNR.
|
||||
"""Export the pretrained mbt2018_mean mean-scale hyperprior (CompressAI zoo,
|
||||
quality 6, MSE) as an INTEGER-exact decoder and measure what that costs:
|
||||
per-channel int12 weights + fixed-point activations + integer IGDN, all i64 —
|
||||
byte-identical on every machine (NeuralIntegerExact). The tool reports the
|
||||
integer and float PSNR on the given image; tests/int_hyper_synth.rs requires the
|
||||
integer decode within 1 dB of the float decode.
|
||||
|
||||
python tools/wai_sota_int_synth_export.py <image.png> <out-dir>
|
||||
|
||||
Normally run through tools/wai_sota_vectors.sh, which pins the source image and
|
||||
the Python stack and produces the whole vector directory the Rust int_hyper tests
|
||||
read.
|
||||
"""
|
||||
import sys, os, json, hashlib
|
||||
import sys, os, io, json, hashlib
|
||||
import numpy as np, torch, warnings; warnings.filterwarnings('ignore')
|
||||
from compressai.zoo import mbt2018_mean
|
||||
from compressai.zoo import image as zoo_image
|
||||
from compressai.zoo.pretrained import load_pretrained
|
||||
|
||||
AFRAC = 12; WBITS = 12; SHIFT = 24; GF = 16
|
||||
|
||||
# The pretrained checkpoint the model-derived vectors come from (all but the source
|
||||
# image, source.png and source.rgb), pinned by its full SHA-256.
|
||||
# CompressAI's own `mbt2018_mean(pretrained=True)` fetches it through torch.hub
|
||||
# without checking any hash (and unpickles it with weights_only=False), so a changed
|
||||
# object at the URL, or a stale file in the torch.hub cache, would silently yield
|
||||
# different vectors that the Rust tests — comparing against Python output built
|
||||
# from the same checkpoint — would still pass. So the bytes are checked here, and
|
||||
# only then loaded, as tensors only (weights_only=True).
|
||||
CKPT_URL = "https://compressai.s3.amazonaws.com/models/v1/mbt2018-mean-6-a19628ab.pth.tar"
|
||||
CKPT_SHA256 = "a19628aba19732125e02caf5c38e233edbe2a641145a5b1779c51f4527f51dd6"
|
||||
|
||||
|
||||
def load_pinned_mbt2018_mean():
|
||||
"""mbt2018_mean quality 6 (MSE), built from checkpoint bytes whose SHA-256 is
|
||||
CKPT_SHA256 — or exit non-zero. Downloads into the torch.hub cache if absent."""
|
||||
zoo_url = zoo_image.model_urls["mbt2018-mean"]["mse"][6]
|
||||
if zoo_url != CKPT_URL:
|
||||
sys.exit(f"error: the installed compressai points mbt2018_mean q6 at {zoo_url}, "
|
||||
f"not the pinned {CKPT_URL}; review the new checkpoint before re-pinning")
|
||||
path = os.path.join(torch.hub.get_dir(), "checkpoints", os.path.basename(CKPT_URL))
|
||||
if not os.path.exists(path):
|
||||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||||
torch.hub.download_url_to_file(CKPT_URL, path, progress=False)
|
||||
with open(path, "rb") as f:
|
||||
blob = f.read()
|
||||
got = hashlib.sha256(blob).hexdigest()
|
||||
if got != CKPT_SHA256:
|
||||
sys.exit(f"error: {path} has sha256 {got}, pinned {CKPT_SHA256}. "
|
||||
"Delete the file to re-download; if the published checkpoint itself "
|
||||
"changed, re-pin only after reviewing it.")
|
||||
print(f"checkpoint {os.path.basename(path)}: sha256 {got} (pinned, verified)")
|
||||
state = torch.load(io.BytesIO(blob), map_location="cpu", weights_only=True)
|
||||
return zoo_image.model_architectures["mbt2018-mean"].from_state_dict(load_pretrained(state))
|
||||
|
||||
def q_weight_pc(w): # ConvTranspose [in,out,k,k] → per-OUT-channel int12 + scale
|
||||
s = np.abs(w).max(axis=(0,2,3), keepdims=True)/(2**(WBITS-1)-1) + 1e-12
|
||||
wi = np.round(w/s).clip(-(2**(WBITS-1)), 2**(WBITS-1)-1).astype(np.int64)
|
||||
|
|
@ -149,11 +190,27 @@ def model_to_bin(model):
|
|||
|
||||
def main(imgpath, outdir):
|
||||
from PIL import Image
|
||||
# The ENCODER side (g_a, h_a, the float h_s used for the reference) is float, so
|
||||
# its latent can depend on kernel choice and reduction order. Both are pinned
|
||||
# to make two runs of this tool on one machine emit identical vectors. That
|
||||
# held for repeated runs on an Apple Silicon Mac, but two runs on an x86_64
|
||||
# Linux runner wrote different cond.json and int_synth_sota.json, most likely
|
||||
# in psnr_float (`base` below), which the Rust tests read only as a quality
|
||||
# threshold; the latent and both rgb_hash values matched. The decode being
|
||||
# tested is integer and needs no such pinning.
|
||||
torch.use_deterministic_algorithms(True)
|
||||
torch.set_num_threads(1)
|
||||
# Verify the checkpoint before writing anything, so a refused checkpoint leaves
|
||||
# no partial vector directory behind.
|
||||
m = load_pinned_mbt2018_mean().eval()
|
||||
os.makedirs(outdir, exist_ok=True)
|
||||
img = Image.open(imgpath).convert('RGB')
|
||||
x = torch.from_numpy(np.array(img).astype(np.float32)/255.).permute(2,0,1).unsqueeze(0)
|
||||
src8 = np.array(img)
|
||||
x = torch.from_numpy(src8.astype(np.float32)/255.).permute(2,0,1).unsqueeze(0)
|
||||
h,w = x.shape[2]//64*64, x.shape[3]//64*64; x = x[:,:,:h,:w]
|
||||
m = mbt2018_mean(quality=6, pretrained=True).eval()
|
||||
# The exact pixels the codec saw (HWC u8), so a decoder can measure its own
|
||||
# reconstruction quality instead of trusting the PSNR this tool reports.
|
||||
np.ascontiguousarray(src8[:h, :w]).tofile(f'{outdir}/source.rgb')
|
||||
def psnr(a,b): return float(-10*np.log10(((a-b)**2).mean()+1e-12))
|
||||
with torch.no_grad():
|
||||
y = m.g_a(x); z = m.h_a(y); z_hat = torch.round(z)
|
||||
|
|
@ -199,7 +256,7 @@ def main(imgpath, outdir):
|
|||
for byte in rgb8.tobytes():
|
||||
h = ((h ^ byte) * 1099511628211) & 0xFFFFFFFFFFFFFFFF
|
||||
digest = format(h, '016x')
|
||||
json.dump({'afrac':AFRAC,'wbits':WBITS,'shift':SHIFT,'psnr_int':p_int,'rgb_hash':digest,
|
||||
json.dump({'afrac':AFRAC,'wbits':WBITS,'shift':SHIFT,'psnr_int':p_int,'psnr_float':base,'rgb_hash':digest,
|
||||
'out_c':int(rec.shape[0]),'out_h':int(rec.shape[1]),'out_w':int(rec.shape[2])},
|
||||
open(f'{outdir}/int_synth_sota.json','w'))
|
||||
json.dump(model, open(f'{outdir}/int_synth_sota_model.json','w'))
|
||||
|
|
@ -308,7 +365,7 @@ def main(imgpath, outdir):
|
|||
'n_z':int(Cz),'zh':int(z_sym.shape[1]),'zw':int(z_sym.shape[2]),
|
||||
'z_rans':list(z_bits),'rans':list(bits)},
|
||||
open(f'{outdir}/cond_bitstream.json','w'))
|
||||
json.dump({'rgb_hash':cond_hash,'psnr_cond':p_cond,
|
||||
json.dump({'rgb_hash':cond_hash,'psnr_cond':p_cond,'psnr_float':base,
|
||||
'out_c':int(rec_int.shape[0]),'out_h':int(rec_int.shape[1]),'out_w':int(rec_int.shape[2])},
|
||||
open(f'{outdir}/cond.json','w'))
|
||||
print(f'wrote conditional bitstream: {len(z_bits)} z + {len(bits)} y rANS bytes '
|
||||
|
|
|
|||
49
wai/tools/wai_sota_source_image.py
Normal file
49
wai/tools/wai_sota_source_image.py
Normal file
|
|
@ -0,0 +1,49 @@
|
|||
#!/usr/bin/env python3
|
||||
"""Write the pinned source photograph the int_hyper integer-codec vectors (the
|
||||
pretrained mbt2018_mean, exported to integers) are regenerated from.
|
||||
|
||||
python tools/wai_sota_source_image.py <out.png>
|
||||
|
||||
The image is a 256x256 crop of `skimage.data.astronaut()` — a photograph that
|
||||
ships inside the scikit-image wheel and is documented there as public domain
|
||||
(NASA; "No known copyright restrictions, released into the public domain"). It
|
||||
is pinned by the SHA-256 of the cropped RGB bytes, so a different image can never
|
||||
be substituted silently: a mismatch is an error, not a warning.
|
||||
|
||||
A real photograph, rather than a synthetic pattern, is what makes the quality
|
||||
check meaningful: the integer decode is checked against the same model's float
|
||||
decode on natural-image content, and must cost at most 1 dB there
|
||||
(tests/int_hyper_synth.rs). That bar is relative, on this one image; it says
|
||||
nothing about how the model compares with other codecs, or about other images.
|
||||
The 42.6 dB development figure for this codec was taken on a different
|
||||
photograph that is not in this repository.
|
||||
"""
|
||||
import hashlib
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from skimage import data
|
||||
|
||||
CROP = (slice(64, 320), slice(128, 384)) # rows, cols of the 512x512 original
|
||||
SHA256 = "29cb3aee1a55f9585474d056b65edb0aaebbb80371ba200e85030d63665c19d6"
|
||||
|
||||
|
||||
def main(out: str) -> int:
|
||||
rgb = np.ascontiguousarray(data.astronaut()[CROP])
|
||||
got = hashlib.sha256(rgb.tobytes()).hexdigest()
|
||||
if rgb.shape != (256, 256, 3) or got != SHA256:
|
||||
print(f"error: source crop is {rgb.shape} sha256 {got}; pinned (256, 256, 3) {SHA256}. "
|
||||
"The scikit-image sample changed — pin the new hash only after reviewing the image.",
|
||||
file=sys.stderr)
|
||||
return 1
|
||||
Image.fromarray(rgb).save(out)
|
||||
print(f"wrote {out}: 256x256 RGB, sha256(rgb) {got}")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if len(sys.argv) != 2:
|
||||
print(__doc__, file=sys.stderr)
|
||||
raise SystemExit(2)
|
||||
raise SystemExit(main(sys.argv[1]))
|
||||
71
wai/tools/wai_sota_vectors.sh
Executable file
71
wai/tools/wai_sota_vectors.sh
Executable file
|
|
@ -0,0 +1,71 @@
|
|||
#!/usr/bin/env bash
|
||||
# Regenerate the int_hyper integer-codec vectors (the pretrained mbt2018_mean,
|
||||
# exported to integers) that the Rust tests read from WAI_SOTA_INT_DIR
|
||||
# (tests/int_hyper_synth.rs, tests/int_video_sota_keyframe.rs).
|
||||
#
|
||||
# tools/wai_sota_vectors.sh [out-dir] # default /tmp/sota_int
|
||||
#
|
||||
# Needs the pinned Python stack (tools/requirements-sota-int.txt) on PATH as
|
||||
# `python` (or set PYTHON), and cargo. The pretrained mbt2018_mean checkpoint is
|
||||
# downloaded into the torch.hub cache on first use and checked against its pinned
|
||||
# SHA-256 (CKPT_SHA256 in wai_sota_int_synth_export.py) before it is loaded.
|
||||
#
|
||||
# The chain, each step consuming the one before:
|
||||
# 1. the pinned public-domain source photograph → source.png
|
||||
# 2. the integer export of mbt2018_mean → model, bitstream, hashes
|
||||
# 3. the Rust integer decode of that bitstream → frame0.rgb (the keyframe)
|
||||
# 4. the P-frame encoder panning that keyframe → video bitstream, hashes
|
||||
# Every step must succeed and every file the tests read must exist at the end, or
|
||||
# this exits non-zero: a partial directory would let a test skip instead of fail.
|
||||
#
|
||||
# The output is not portable, and part of it is not repeatable run to run on
|
||||
# every machine. Step 2 runs the float model (torch, CPU) to encode the
|
||||
# photograph, on the pinned stack, which each platform installs as its own build.
|
||||
# - Between an Apple Silicon Mac and an x86_64 Linux runner,
|
||||
# int_synth_sota_input.json (the synthesis-only latent, which the float
|
||||
# encoder and hyper-decoder compute) differed, and with it the rgb_hash in
|
||||
# int_synth_sota.json. Each side wrote the same latent in all of its runs.
|
||||
# - Two runs on that one runner wrote different cond.json and
|
||||
# int_synth_sota.json, most likely in psnr_float, the float decode's PSNR
|
||||
# (the runner's files were not kept, so the field is inferred). Repeated runs
|
||||
# on the Mac wrote identical files. Only the quality test reads psnr_float,
|
||||
# against thresholds (the 1 dB bar and a 30 dB floor).
|
||||
# - source.png differed between the Mac and the runner in its compressed bytes.
|
||||
# Step 1 writes it with the PNG encoder from the pinned pixels, before any
|
||||
# model runs, and its pixels (source.rgb) were identical, so this is the
|
||||
# encoder's output, not float computation.
|
||||
# The other files matched in these comparisons, which does not make them
|
||||
# portable. So check a decode against the vectors of the same run, and do not
|
||||
# pin a hash from these files for another run, on this machine or another: it
|
||||
# can differ with both decoders correct. The decode's cross-architecture check
|
||||
# uses the committed tiny vector (wai-rs/tests/vectors/int_hyper/), whose hashes
|
||||
# are pinned.
|
||||
set -euo pipefail
|
||||
|
||||
here="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
wai="$(dirname "$here")"
|
||||
out="${1:-/tmp/sota_int}"
|
||||
py="${PYTHON:-python}"
|
||||
|
||||
mkdir -p "$out"
|
||||
out="$(cd "$out" && pwd)"
|
||||
|
||||
"$py" "$here/wai_sota_source_image.py" "$out/source.png"
|
||||
"$py" "$here/wai_sota_int_synth_export.py" "$out/source.png" "$out"
|
||||
(cd "$wai/wai-rs" && cargo run --release --no-default-features --features neural_int \
|
||||
--example int_hyper_decode_frame -- "$out" "$out/frame0.rgb")
|
||||
"$py" "$here/wai_int_video_sota_keyframe_encode.py" "$out"
|
||||
|
||||
missing=0
|
||||
for f in source.rgb \
|
||||
int_synth_sota.json int_synth_sota_model.json int_synth_sota_model.bin int_synth_sota_input.json \
|
||||
int_hs.json int_hs_model.json int_hs_model.bin int_hs_input.json \
|
||||
cond.json cond_entropy.json cond_bitstream.json \
|
||||
frame0.rgb video_sota.json video_sota_bitstream.json; do
|
||||
if [ ! -s "$out/$f" ]; then
|
||||
echo "error: $out/$f was not produced" >&2
|
||||
missing=1
|
||||
fi
|
||||
done
|
||||
[ "$missing" -eq 0 ] || exit 1
|
||||
echo "int_hyper vectors complete in $out"
|
||||
|
|
@ -21,6 +21,12 @@ crate-type = ["lib", "cdylib", "staticlib"] # cdylib + staticlib for FFI/SDK
|
|||
[lints.clippy]
|
||||
disallowed_methods = "allow"
|
||||
|
||||
# `cfg(kani)` is set only by the model checker (wai-verify `#[path]`-includes the
|
||||
# integer sources: rans, int_entropy, int_transform, int_kernels). Declared so an
|
||||
# ordinary build or clippy run knows the name and does not warn on it.
|
||||
[lints.rust]
|
||||
unexpected_cfgs = { level = "warn", check-cfg = ["cfg(kani)"] }
|
||||
|
||||
[dependencies]
|
||||
# Container manifest. preserve_order so JSON round-trips byte-identical.
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
|
|
@ -519,6 +525,11 @@ name = "wai_bench"
|
|||
path = "src/bin/wai_bench.rs"
|
||||
required-features = ["codecs"]
|
||||
|
||||
[[bin]]
|
||||
name = "wai_decode_bench"
|
||||
path = "src/bin/wai_decode_bench.rs"
|
||||
required-features = ["neural_int", "video_int", "audio_int", "world_int", "splat_int"]
|
||||
|
||||
[[bin]]
|
||||
name = "wai_bench_diff"
|
||||
path = "src/bin/wai_bench_diff.rs"
|
||||
|
|
|
|||
|
|
@ -1,8 +1,8 @@
|
|||
//! Decode a SOTA `int_hyper` conditional bitstream from a directory and dump the
|
||||
//! Decode an `int_hyper` conditional bitstream from a directory and dump the
|
||||
//! reconstructed frame as raw RGB (HWC u8) — the byte-exact keyframe the
|
||||
//! `int_hyper`-keyframe video codec builds its P-frames on.
|
||||
//!
|
||||
//! cargo run --release --example int_hyper_decode_frame --features neural_int -- <dir> <out.rgb>
|
||||
//! cargo run --release --no-default-features --features neural_int --example int_hyper_decode_frame -- <dir> <out.rgb>
|
||||
//!
|
||||
//! `<dir>` is an int_hyper vector dir (e.g. `/tmp/sota_int`) holding
|
||||
//! `int_synth_sota_model.json`, `int_hs_model.json`, `cond_entropy.json`,
|
||||
|
|
|
|||
644
wai/wai-rs/src/bin/wai_decode_bench.rs
Normal file
644
wai/wai-rs/src/bin/wai_decode_bench.rs
Normal file
|
|
@ -0,0 +1,644 @@
|
|||
//! Decode throughput of WAI's integer decoders: the scalar reference kernels
|
||||
//! against the fast kernels, measured in one process on the same inputs.
|
||||
//!
|
||||
//! ```text
|
||||
//! wai_decode_bench [--rounds N] [--filter SUBSTRING] [--quick]
|
||||
//! wai_decode_bench --profile reference|1t|auto [--seconds S] [--filter SUBSTRING]
|
||||
//! wai_decode_bench --layers [--rounds N]
|
||||
//! ```
|
||||
//!
|
||||
//! For every case three policies run **interleaved** — the scalar reference kernels,
|
||||
//! the fast kernels on one thread, and the fast kernels on the automatic worker
|
||||
//! count — with the order rotated each round. Every run's output is hashed, and the
|
||||
//! benchmark aborts if any policy's output differs from the reference's in any
|
||||
//! round: it does not report a speed for a result that is not identical. Reported
|
||||
//! per policy: the best and the median wall time, and throughput from the best.
|
||||
//!
|
||||
//! Inputs are self-contained. The kernel and synthetic whole-decoder cases use
|
||||
//! deterministic pseudo-random weights at the real layer shapes (time depends on
|
||||
//! shape and on operand *width*, not on trained values, and the synthetic operands
|
||||
//! are kept at the widths the export tools emit: int12 WHM1 weights with Q12
|
||||
//! activations, 16-bit PTQ JPEG-AI weights with Q16 activations). The
|
||||
//! committed-vector cases decode the repository's conformance vectors, which are
|
||||
//! not all trained models:
|
||||
//!
|
||||
//! - **trained:** int_synth photo, int_world, int_splat and int_audio (their export
|
||||
//! tools train the float model, then quantize it; int_world and int_splat train
|
||||
//! on procedurally generated content);
|
||||
//! - **random weights:** the int_hyper tiny conditional codec
|
||||
//! (`tools/wai_int_hyper_tiny_vector.py` draws its integer weights from a seeded
|
||||
//! generator);
|
||||
//! - **no learned model:** int_video real footage (integer block motion
|
||||
//! compensation and rANS).
|
||||
//!
|
||||
//! So no trained int_hyper model is timed here: the `g_s` case uses synthetic
|
||||
//! weights and the int_hyper vector uses random ones.
|
||||
|
||||
use std::hint::black_box;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use serde_json::{json, Map, Value};
|
||||
use wai::int_hyper_synth::{IntLayer, IntModel};
|
||||
use wai::int_kernels::{available_threads, with_policy, Policy};
|
||||
|
||||
const VECTORS: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/tests/vectors");
|
||||
const NEURAL_CONF: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/../neural-conformance");
|
||||
|
||||
/// Deterministic generator (SplitMix64).
|
||||
struct Rng(u64);
|
||||
impl Rng {
|
||||
fn next(&mut self) -> u64 {
|
||||
self.0 = self.0.wrapping_add(0x9E37_79B9_7F4A_7C15);
|
||||
let mut z = self.0;
|
||||
z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
|
||||
z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
|
||||
z ^ (z >> 31)
|
||||
}
|
||||
/// Uniform in `[-r, r]`.
|
||||
fn signed(&mut self, r: i64) -> i64 {
|
||||
(self.next() % (2 * r as u64 + 1)) as i64 - r
|
||||
}
|
||||
fn vec(&mut self, n: usize, r: i64) -> Vec<i64> {
|
||||
(0..n).map(|_| self.signed(r)).collect()
|
||||
}
|
||||
}
|
||||
|
||||
fn fnv1a(bytes: &[u8]) -> u64 {
|
||||
bytes.iter().fold(0xcbf2_9ce4_8422_2325u64, |h, &b| (h ^ b as u64).wrapping_mul(0x0100_0000_01b3))
|
||||
}
|
||||
|
||||
/// A decode's output, hashed only after the timer stops.
|
||||
enum Out {
|
||||
I(Vec<i64>),
|
||||
B(Vec<u8>),
|
||||
H(u64),
|
||||
}
|
||||
fn digest(o: &Out) -> u64 {
|
||||
match o {
|
||||
Out::I(v) => v.iter().fold(0xcbf2_9ce4_8422_2325u64, |h, &x| (h ^ x as u64).wrapping_mul(0x0100_0000_01b3)),
|
||||
Out::B(v) => fnv1a(v),
|
||||
Out::H(h) => *h,
|
||||
}
|
||||
}
|
||||
|
||||
/// One benchmark case: `run` performs one decode and returns its output.
|
||||
struct Case {
|
||||
name: String,
|
||||
/// What one run produces, e.g. `(0.065536, "MPix")`.
|
||||
per_run: (f64, &'static str),
|
||||
/// Multiply-accumulates of the convolution/IGDN work, when meaningful.
|
||||
macs: Option<u64>,
|
||||
run: Box<dyn Fn() -> Out>,
|
||||
}
|
||||
|
||||
// ─────────────────────────── synthetic models ───────────────────────────
|
||||
|
||||
/// Per-channel int12 requant multiplier keeping a layer's output at the input's
|
||||
/// scale for `n` uniform-int12 terms (std ≈ 1182·√n) with `shift` bits.
|
||||
fn mult_for(n: usize, shift: u32) -> i64 {
|
||||
(((1u64 << shift) as f64) / (1182.0 * (n as f64).sqrt())).round().max(1.0) as i64
|
||||
}
|
||||
|
||||
fn convt_layer(g: &mut Rng, cin: usize, cout: usize, k: usize, stride: usize, shift: u32) -> IntLayer {
|
||||
let m = mult_for(cin * k * k / (stride * stride), shift);
|
||||
IntLayer::ConvT {
|
||||
out_ch: cout, k, stride, pad: k / 2, opad: stride - 1,
|
||||
weight: g.vec(cin * cout * k * k, 2047),
|
||||
mult: vec![m; cout],
|
||||
bias: g.vec(cout, 1 << 8),
|
||||
}
|
||||
}
|
||||
fn conv_layer(g: &mut Rng, cin: usize, cout: usize, k: usize, shift: u32) -> IntLayer {
|
||||
let m = mult_for(cin * k * k, shift);
|
||||
IntLayer::Conv { out_ch: cout, k, stride: 1, pad: k / 2, weight: g.vec(cout * cin * k * k, 2047), mult: vec![m; cout], bias: g.vec(cout, 1 << 8) }
|
||||
}
|
||||
/// IGDN with a small, mostly diagonal gamma (Q16) and beta = 1.0 (Q24).
|
||||
fn igdn_layer(g: &mut Rng, c: usize) -> IntLayer {
|
||||
let gamma = (0..c * c).map(|i| if i % (c + 1) == 0 { 1300 + g.signed(200) } else { (g.next() % 14) as i64 }).collect();
|
||||
IntLayer::Igdn { gamma, beta: vec![1 << 24; c] }
|
||||
}
|
||||
|
||||
/// The synthesis `g_s` of an mbt2018-mean-class codec at quality 6 (M = 320,
|
||||
/// N = 192) — the topology `tools/wai_sota_int_synth_export.py` exports — with
|
||||
/// synthetic int12 weights; afrac 12, gf 16, shift 24 as in the committed models.
|
||||
fn sota_gs() -> IntModel {
|
||||
let mut g = Rng(0x5074_0001);
|
||||
let (m, n) = (320, 192);
|
||||
let layers = vec![
|
||||
convt_layer(&mut g, m, n, 5, 2, 24), igdn_layer(&mut g, n),
|
||||
convt_layer(&mut g, n, n, 5, 2, 24), igdn_layer(&mut g, n),
|
||||
convt_layer(&mut g, n, n, 5, 2, 24), igdn_layer(&mut g, n),
|
||||
convt_layer(&mut g, n, 3, 5, 2, 24),
|
||||
];
|
||||
IntModel { afrac: 12, gf: 16, shift: 24, layers }
|
||||
}
|
||||
|
||||
/// JPEG-AI synthesis weights (model_y + model_uv, tools_0) at the reference
|
||||
/// topology's exact weight shapes, synthetic values: 16-bit per-tensor PTQ weights
|
||||
/// with `mult = round(sw · 2^31)`, LayerNorm γ = 1, temperature = 1, Q16 biases.
|
||||
fn jpegai_weights(luma: bool) -> wai::jpegai_synthesis::Weights {
|
||||
let mut g = Rng(if luma { 0x1A_0001 } else { 0x1A_0002 });
|
||||
let mut m = Map::new();
|
||||
// name, dims (ONNX layout: conv [cout, cin/g, k, k], convT [cin, cout, k, k]), fan-in
|
||||
// The transposed convs (weight [cin, cout, k, k]) — every other module is a
|
||||
// regular conv (weight [cout, cin/g, k, k]).
|
||||
const CONV_T: [&str; 5] = ["first_stage.1", "conv2_t", "CAB.upscale", "TAM.us_conv", "conv4_t"];
|
||||
let mut conv = |m: &mut Map<String, Value>, name: &str, dims: [usize; 4], fan_in: usize| {
|
||||
let n: usize = dims.iter().product();
|
||||
let amax = (3.0 / fan_in as f64).sqrt();
|
||||
let mult = (amax / 32767.0 * (1u64 << 31) as f64).round() as i64;
|
||||
m.insert(format!("{name}.weight"), json!({"kind": "conv", "w": g.vec(n, 32767), "dims": dims, "mult": mult, "shift": 31}));
|
||||
let cout = if CONV_T.contains(&name) { dims[1] } else { dims[0] };
|
||||
m.insert(format!("{name}.bias"), json!({"kind": "q16", "v": g.vec(cout, 1 << 12)}));
|
||||
};
|
||||
let q16 = |m: &mut Map<String, Value>, name: &str, v: Vec<i64>| {
|
||||
m.insert(name.to_string(), json!({"kind": "q16", "v": v}));
|
||||
};
|
||||
let (c, hid) = if luma { (128, 512) } else { (64, 256) };
|
||||
type ConvFn<'a> = dyn FnMut(&mut Map<String, Value>, &str, [usize; 4], usize) + 'a;
|
||||
let res = |m: &mut Map<String, Value>, conv: &mut ConvFn<'_>, p: &str, c: usize| {
|
||||
conv(m, &format!("{p}.conv1"), [c, c, 3, 3], c * 9);
|
||||
conv(m, &format!("{p}.conv2"), [c, c, 3, 3], c * 9);
|
||||
};
|
||||
if luma {
|
||||
res(&mut m, &mut conv, "first_stage.0.0", 160);
|
||||
conv(&mut m, "first_stage.1", [160, 128, 3, 3], 160 * 9 / 4);
|
||||
conv(&mut m, "first_stage.2.conv", [128, 32, 3, 3], 32 * 9);
|
||||
conv(&mut m, "first_stage.2.conv2", [128, 128, 1, 1], 128);
|
||||
conv(&mut m, "conv2_t", [128, 128, 3, 3], 128 * 9 / 4);
|
||||
} else {
|
||||
conv(&mut m, "first_stage.conv1", [48, 256, 3, 3], 256 * 9);
|
||||
conv(&mut m, "conv2_t", [144, 64, 3, 3], 144 * 9 / 4);
|
||||
}
|
||||
for p in ["CAB.residual_trunk.0", "CAB.residual_trunk.1", "CAB.residual_mask2.0", "CAB.residual_mask3.0"] {
|
||||
res(&mut m, &mut conv, p, c);
|
||||
}
|
||||
conv(&mut m, "CAB.subscale", [c, c, 3, 3], c * 9);
|
||||
conv(&mut m, "CAB.upscale", [c, c, 3, 3], c * 9 / 4);
|
||||
conv(&mut m, "iact2.conv", [c, c / 4, 3, 3], c / 4 * 9);
|
||||
conv(&mut m, "iact2.conv2", [c, c, 1, 1], c);
|
||||
conv(&mut m, "conv3_t", [4 * c, c, 1, 1], c);
|
||||
if luma {
|
||||
conv(&mut m, "TAM.ds_conv", [128, 128, 3, 3], 128 * 9);
|
||||
conv(&mut m, "TAM.us_conv", [128, 128, 3, 3], 128 * 9 / 4);
|
||||
}
|
||||
for t in ["TAM.TABs.0", "TAM.TABs.1"] {
|
||||
q16(&mut m, &format!("{t}.prep_data.norm1.weight"), vec![1 << 16; c]);
|
||||
q16(&mut m, &format!("{t}.prep_data.norm1.bias"), vec![0; c]);
|
||||
conv(&mut m, &format!("{t}.prep_data.conv1"), [3 * c, c, 1, 1], c);
|
||||
conv(&mut m, &format!("{t}.prep_data.conv2"), [3 * c, 1, 3, 3], 9);
|
||||
q16(&mut m, &format!("{t}.attn.temperature"), vec![1 << 16; 4]);
|
||||
conv(&mut m, &format!("{t}.attn.project_out"), [c, c, 1, 1], c);
|
||||
q16(&mut m, &format!("{t}.ffn.norm1.weight"), vec![1 << 16; c]);
|
||||
q16(&mut m, &format!("{t}.ffn.norm1.bias"), vec![0; c]);
|
||||
conv(&mut m, &format!("{t}.ffn.project_in"), [2 * hid, c, 1, 1], c);
|
||||
conv(&mut m, &format!("{t}.ffn.dwconv"), [2 * hid, 1, 3, 3], 9);
|
||||
conv(&mut m, &format!("{t}.ffn.project_out"), [c, hid, 1, 1], hid);
|
||||
}
|
||||
conv(&mut m, "iact3.conv", [c, c / 4, 3, 3], c / 4 * 9);
|
||||
conv(&mut m, "iact3.conv2", [c, c, 1, 1], c);
|
||||
if luma {
|
||||
conv(&mut m, "conv4_t", [128, 1, 3, 3], 128 * 9 / 4);
|
||||
} else {
|
||||
conv(&mut m, "conv4_t", [64, 8, 3, 3], 64 * 9 / 4);
|
||||
}
|
||||
wai::jpegai_synthesis::Weights::from_json(&Value::Object(m))
|
||||
}
|
||||
|
||||
// ─────────────────────────────── cases ───────────────────────────────
|
||||
|
||||
fn kernel_cases(cases: &mut Vec<Case>) {
|
||||
use wai::int_hyper_synth::run_net;
|
||||
use wai::int_transformer::{conv_grouped, conv_transpose};
|
||||
|
||||
// int_hyper g_s middle layer: convT 192→192, 5×5, stride 2, 64² → 128².
|
||||
{
|
||||
let mut g = Rng(0x4B_0001);
|
||||
let model = IntModel { afrac: 12, gf: 16, shift: 24, layers: vec![convt_layer(&mut g, 192, 192, 5, 2, 24)] };
|
||||
let x = g.vec(192 * 64 * 64, 3 << 12);
|
||||
cases.push(Case {
|
||||
name: "kernel · int_hyper convT 192→192 k5 s2, 64²→128²".into(),
|
||||
per_run: (128.0 * 128.0 / 1e6, "MPix"),
|
||||
macs: Some(192 * 192 * 25 * 64 * 64),
|
||||
run: Box::new(move || Out::I(run_net(&model, &x, 192, 64, 64).0)),
|
||||
});
|
||||
}
|
||||
// int_hyper h_s-style regular conv: 3×3, 192→192 at 64².
|
||||
{
|
||||
let mut g = Rng(0x4B_0002);
|
||||
let model = IntModel { afrac: 12, gf: 16, shift: 24, layers: vec![conv_layer(&mut g, 192, 192, 3, 24)] };
|
||||
let x = g.vec(192 * 64 * 64, 3 << 12);
|
||||
cases.push(Case {
|
||||
name: "kernel · int_hyper conv 192→192 k3 s1, 64²".into(),
|
||||
per_run: (64.0 * 64.0 / 1e6, "MPix"),
|
||||
macs: Some(192 * 192 * 9 * 64 * 64),
|
||||
run: Box::new(move || Out::I(run_net(&model, &x, 192, 64, 64).0)),
|
||||
});
|
||||
}
|
||||
// IGDN c = 192 at 128²: Q12 activations (32-bit term path) and wide ones (64-bit terms).
|
||||
for (label, xr, seed) in [("Q12 |x|<4", 4i64 << 12, 3u64), ("|x|² > 2³¹", 1 << 17, 4)] {
|
||||
let mut g = Rng(0x4B_0000 + seed);
|
||||
let IntLayer::Igdn { gamma, beta } = igdn_layer(&mut g, 192) else { unreachable!() };
|
||||
let x = g.vec(192 * 128 * 128, xr);
|
||||
cases.push(Case {
|
||||
name: format!("kernel · IGDN c=192, 128², {label}"),
|
||||
per_run: (128.0 * 128.0 / 1e6, "MPix"),
|
||||
macs: Some(192 * 192 * 128 * 128),
|
||||
run: Box::new(move || Out::I(wai::int_transform::igdn(&x, 192, 128 * 128, 12, &gamma, &beta, 16))),
|
||||
});
|
||||
}
|
||||
// JPEG-AI shapes (Q16 activations, 16-bit weights, shift 31).
|
||||
type Spec = (&'static str, usize, usize, usize, usize, usize, usize, usize, bool);
|
||||
// name, cin, cout, k, stride, groups, h(=w) in, transposed?
|
||||
let specs: [Spec; 5] = [
|
||||
("conv 128→128 k3 (CAB res block), 64²", 128, 128, 3, 1, 1, 64, 0, false),
|
||||
("conv 1×1 128→1024 (GDFN project_in), 64²", 128, 1024, 1, 1, 1, 64, 0, false),
|
||||
("depthwise 3×3 ×1024 (GDFN dwconv), 64²", 1024, 1024, 3, 1, 1024, 64, 0, false),
|
||||
("conv 128→128 k3 s2 (TAM.ds_conv), 128²→64²", 128, 128, 3, 2, 1, 128, 0, false),
|
||||
("convT 128→128 k3 s2 (TAM.us_conv), 64²→128²", 128, 128, 3, 2, 1, 64, 1, true),
|
||||
];
|
||||
for (i, &(label, cin, cout, k, stride, groups, hw, opad, transposed)) in specs.iter().enumerate() {
|
||||
let mut g = Rng(0x4C_0000 + i as u64);
|
||||
let cin_g = cin / groups;
|
||||
let x = g.vec(cin * hw * hw, 2 << 16);
|
||||
let wt = g.vec(cin_g * cout * k * k, 32767);
|
||||
let mult = vec![(((3.0 / (cin_g * k * k) as f64).sqrt() / 32767.0) * (1u64 << 31) as f64) as i64; cout];
|
||||
let bias = g.vec(cout, 1 << 12);
|
||||
let pad = k / 2;
|
||||
let ho = if transposed { (hw - 1) * stride + k + opad - 2 * pad } else { (hw + 2 * pad - k) / stride + 1 };
|
||||
let macs = if transposed { cin * cout * k * k * hw * hw } else { cin_g * cout * k * k * ho * ho } as u64;
|
||||
cases.push(Case {
|
||||
name: format!("kernel · JPEG-AI {label}"),
|
||||
per_run: ((ho * ho) as f64 / 1e6, "MPix"),
|
||||
macs: Some(macs),
|
||||
run: Box::new(move || {
|
||||
let out = if transposed {
|
||||
conv_transpose(&x, cin, hw, hw, &wt, cout, k, stride, pad, opad, &mult, &bias, 31).0
|
||||
} else {
|
||||
conv_grouped(&x, cin, hw, hw, &wt, cout, k, stride, pad, groups, &mult, &bias, 31).0
|
||||
};
|
||||
Out::I(out)
|
||||
}),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
fn decoder_cases(cases: &mut Vec<Case>, quick: bool) {
|
||||
// Synthetic mbt2018-mean-class g_s (quality 6 topology): 16² latent → 256² RGB.
|
||||
{
|
||||
let model = sota_gs();
|
||||
let mut g = Rng(0x5074_0002);
|
||||
let lat = if quick { 8 } else { 16 };
|
||||
let y = g.vec(320 * lat * lat, 3 << 12);
|
||||
let px = (16 * lat) * (16 * lat);
|
||||
cases.push(Case {
|
||||
name: format!("decoder · int_hyper g_s, mbt2018-mean q6 topology (synthetic weights), {}²", 16 * lat),
|
||||
per_run: (px as f64 / 1e6, "MPix"),
|
||||
macs: None,
|
||||
run: Box::new(move || Out::B(wai::int_hyper_synth::decode(&model, &y, 320, lat, lat).0)),
|
||||
});
|
||||
}
|
||||
// JPEG-AI synthesis (luma + chroma nets), 256² (the supported geometry).
|
||||
if !quick {
|
||||
let (wy, wuv) = (jpegai_weights(true), jpegai_weights(false));
|
||||
let mut g = Rng(0x1A_0003);
|
||||
let latent = g.vec(256 * 16 * 16, 1 << 16); // y‖uv latent, Q16
|
||||
cases.push(Case {
|
||||
name: "decoder · JPEG-AI synthesis Y+UV (reference topology, synthetic weights), 256²".into(),
|
||||
per_run: (256.0 * 256.0 / 1e6, "MPix"),
|
||||
macs: None,
|
||||
run: Box::new(move || {
|
||||
let mut y = wai::jpegai_synthesis::decode_y(&latent[..160 * 256], &wy);
|
||||
y.extend(wai::jpegai_synthesis::decode_uv(&latent, &wuv));
|
||||
Out::B(y)
|
||||
}),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
fn read_json(p: &str) -> Value {
|
||||
serde_json::from_slice(&std::fs::read(p).unwrap_or_else(|e| panic!("read {p}: {e}"))).unwrap()
|
||||
}
|
||||
|
||||
fn vector_cases(cases: &mut Vec<Case>) {
|
||||
use wai::int_hyper_synth::{decode_conditional, 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.
|
||||
{
|
||||
let d = format!("{VECTORS}/int_hyper");
|
||||
let g_s = IntModel::from_bytes(&std::fs::read(format!("{d}/tiny_g_s.bin")).unwrap()).unwrap();
|
||||
let h_s = IntModel::from_bytes(&std::fs::read(format!("{d}/tiny_h_s.bin")).unwrap()).unwrap();
|
||||
let ent = CondEntropy::from_value(&read_json(&format!("{d}/tiny_entropy.json")));
|
||||
let bs = CondBitstream::from_value(&read_json(&format!("{d}/tiny_bitstream.json")));
|
||||
let conf = read_json(&format!("{d}/tiny.json"));
|
||||
let px = conf["out_h"].as_u64().unwrap() * conf["out_w"].as_u64().unwrap();
|
||||
cases.push(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)),
|
||||
});
|
||||
}
|
||||
// int_splat: trained Gaussian-splat attribute decode.
|
||||
{
|
||||
let d = format!("{VECTORS}/int_splat");
|
||||
let gs = IntModel::from_bytes(&std::fs::read(format!("{d}/splat_gs.bin")).unwrap()).unwrap();
|
||||
let bs = read_json(&format!("{d}/splat.json"));
|
||||
let n = bs["n_gaussians"].as_u64().unwrap();
|
||||
cases.push(Case {
|
||||
name: format!("vector · int_splat trained attribute decode ({n} Gaussians, committed)"),
|
||||
per_run: (n as f64 / 1e3, "kGaussians"),
|
||||
macs: None,
|
||||
run: Box::new(move || Out::B(wai::int_splat::decode(&gs, &bs).unwrap().attr)),
|
||||
});
|
||||
}
|
||||
// int_world: trained transition + synthesis per generated frame.
|
||||
{
|
||||
let d = format!("{VECTORS}/int_world");
|
||||
let t = IntModel::from_bytes(&std::fs::read(format!("{d}/world_T.bin")).unwrap()).unwrap();
|
||||
let gs = IntModel::from_bytes(&std::fs::read(format!("{d}/world_gs.bin")).unwrap()).unwrap();
|
||||
let j = read_json(&format!("{d}/world.json"));
|
||||
let (c, lat) = (j["c"].as_u64().unwrap() as usize, j["lat"].as_u64().unwrap() as usize);
|
||||
let z0: Vec<i64> = j["z0"].as_array().unwrap().iter().map(|v| v.as_i64().unwrap()).collect();
|
||||
let actions: Vec<usize> = j["actions"].as_array().unwrap().iter().map(|v| v.as_u64().unwrap() as usize).collect();
|
||||
let frames = actions.len() + 1;
|
||||
cases.push(Case {
|
||||
name: format!("vector · int_world trained interactive clip ({frames} frames, committed)"),
|
||||
per_run: (frames as f64, "frames"),
|
||||
macs: None,
|
||||
run: Box::new(move || Out::H(wai::int_world::decode(&t, &gs, &z0, &actions, c, lat).clip_hash())),
|
||||
});
|
||||
}
|
||||
// int_audio: trained 1-D codec (entropy + transposed convs + IGDN).
|
||||
{
|
||||
let d = format!("{VECTORS}/int_audio");
|
||||
let model = wai::int_audio::AudioModel::from_bytes(&std::fs::read(format!("{d}/audio_model.bin")).unwrap()).unwrap();
|
||||
let bs = read_json(&format!("{d}/audio_bitstream.json"));
|
||||
let n = read_json(&format!("{d}/audio.json"))["n_samples"].as_u64().unwrap();
|
||||
cases.push(Case {
|
||||
name: format!("vector · int_audio trained clip ({n} samples, committed)"),
|
||||
per_run: (n as f64 / 1e3, "kSamples"),
|
||||
macs: None,
|
||||
run: Box::new(move || Out::H(wai::int_audio::pcm_fnv1a(&wai::int_audio::decode(&model, &bs).unwrap()))),
|
||||
});
|
||||
}
|
||||
// int_video real footage: integer block MC + rANS only (no learned model).
|
||||
{
|
||||
let bs = read_json(&format!("{VECTORS}/int_video/video_real_bitstream.json"));
|
||||
let conf = read_json(&format!("{VECTORS}/int_video/video_real.json"));
|
||||
let (h, w, nf) = (conf["h"].as_u64().unwrap(), conf["w"].as_u64().unwrap(), conf["frame_hashes"].as_array().unwrap().len());
|
||||
cases.push(Case {
|
||||
name: format!("vector · int_video half-pel real footage ({nf} frames {w}×{h}, committed; no learned model, MC + rANS only)"),
|
||||
per_run: (nf as f64, "frames"),
|
||||
macs: None,
|
||||
run: Box::new(move || {
|
||||
Out::B(wai::int_video::decode(&bs).unwrap().frames.concat())
|
||||
}),
|
||||
});
|
||||
}
|
||||
// int_synth trained real-photo vector.
|
||||
{
|
||||
let d = format!("{NEURAL_CONF}/vectors/photo");
|
||||
let payload = std::fs::read(format!("{d}/latent.bin")).unwrap();
|
||||
let spec = std::fs::read(format!("{d}/int_synth.json")).unwrap();
|
||||
cases.push(Case {
|
||||
name: "vector · int_synth trained photo, 128² (committed)".into(),
|
||||
per_run: (128.0 * 128.0 / 1e6, "MPix"),
|
||||
macs: None,
|
||||
run: Box::new(move || Out::B(wai::int_synth::decode_from_spec(&payload, &spec).unwrap().rgb)),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// ─────────────────────────────── layer mode ───────────────────────────────
|
||||
|
||||
fn layer_profile(policies: &[(&str, Policy); 3], rounds: usize) {
|
||||
use wai::int_hyper_synth::run_net;
|
||||
let model = sota_gs();
|
||||
let mut g = Rng(0x5074_0002);
|
||||
let (mut c, mut h, mut w) = (320usize, 16usize, 16usize);
|
||||
let mut x = g.vec(c * h * w, 3 << 12);
|
||||
println!("| g_s layer (256² decode) | work | reference best (ms) | fast 1T best (ms) | fast auto best (ms) |");
|
||||
println!("|---|---|---:|---:|---:|");
|
||||
let mut totals = [0f64; 3];
|
||||
for layer in &model.layers {
|
||||
let (desc, work) = match layer {
|
||||
IntLayer::ConvT { out_ch, k, stride, .. } => {
|
||||
let macs = c * out_ch * k * k * h * w;
|
||||
(format!("convT {c}→{out_ch} k{k} s{stride}, {h}²→{}²", h * stride), format!("{:.2} GMAC", macs as f64 / 1e9))
|
||||
}
|
||||
IntLayer::Igdn { .. } => {
|
||||
let xm = x.iter().map(|v| v.unsigned_abs()).max().unwrap_or(0);
|
||||
(format!("IGDN c={c}, {h}² (max|x| = {xm}{})", if xm * xm <= i32::MAX as u64 { ", x² fits i32" } else { ", x² > i32" }), format!("{:.2} GMAC", (c * c * h * w) as f64 / 1e9))
|
||||
}
|
||||
_ => ("other".into(), String::new()),
|
||||
};
|
||||
let one = IntModel { afrac: model.afrac, gf: model.gf, shift: model.shift, layers: vec![clone_layer(layer)] };
|
||||
let mut best = [Duration::MAX; 3];
|
||||
let mut outs: Vec<Vec<i64>> = Vec::new();
|
||||
for r in 0..rounds {
|
||||
for k in 0..3 {
|
||||
let idx = (k + r) % 3;
|
||||
let t0 = Instant::now();
|
||||
let (o, oc, oh, ow) = with_policy(policies[idx].1, || black_box(run_net(&one, &x, c, h, w)));
|
||||
best[idx] = best[idx].min(t0.elapsed());
|
||||
if r == 0 && k == 0 {
|
||||
outs.push(o.clone());
|
||||
}
|
||||
if o != outs[0] {
|
||||
eprintln!("MISMATCH in layer {desc}");
|
||||
std::process::exit(1);
|
||||
}
|
||||
if r == rounds - 1 && k == 2 {
|
||||
(c, h, w) = (oc, oh, ow);
|
||||
}
|
||||
}
|
||||
}
|
||||
for i in 0..3 {
|
||||
totals[i] += ms(best[i]);
|
||||
}
|
||||
println!("| {desc} | {work} | {:.2} | {:.2} | {:.2} |", ms(best[0]), ms(best[1]), ms(best[2]));
|
||||
x = outs.pop().unwrap();
|
||||
}
|
||||
println!("| sum of layers | | {:.2} | {:.2} | {:.2} |", totals[0], totals[1], totals[2]);
|
||||
}
|
||||
|
||||
fn clone_layer(l: &IntLayer) -> IntLayer {
|
||||
match l {
|
||||
IntLayer::ConvT { out_ch, k, stride, pad, opad, weight, mult, bias } => IntLayer::ConvT {
|
||||
out_ch: *out_ch, k: *k, stride: *stride, pad: *pad, opad: *opad, weight: weight.clone(), mult: mult.clone(), bias: bias.clone(),
|
||||
},
|
||||
IntLayer::Conv { out_ch, k, stride, pad, weight, mult, bias } => IntLayer::Conv {
|
||||
out_ch: *out_ch, k: *k, stride: *stride, pad: *pad, weight: weight.clone(), mult: mult.clone(), bias: bias.clone(),
|
||||
},
|
||||
IntLayer::Igdn { gamma, beta } => IntLayer::Igdn { gamma: gamma.clone(), beta: beta.clone() },
|
||||
IntLayer::Relu => IntLayer::Relu,
|
||||
IntLayer::Leaky { slope_q, slope_frac } => IntLayer::Leaky { slope_q: *slope_q, slope_frac: *slope_frac },
|
||||
}
|
||||
}
|
||||
|
||||
// ─────────────────────────────── harness ───────────────────────────────
|
||||
|
||||
fn sh(cmd: &str, args: &[&str]) -> Option<String> {
|
||||
let o = std::process::Command::new(cmd).args(args).output().ok()?;
|
||||
o.status.success().then(|| String::from_utf8_lossy(&o.stdout).trim().to_string())
|
||||
}
|
||||
|
||||
/// The core types as the OS reports them, e.g. `6 Super + 12 Performance`: each
|
||||
/// macOS performance level's logical-core count and its own name
|
||||
/// (`hw.perflevelN.name`), in the OS's order. Names are reported as given, not
|
||||
/// inferred from the level's index, since what level 0 and 1 are differs between
|
||||
/// chips. `None` where the OS does not report levels.
|
||||
fn core_levels() -> Option<String> {
|
||||
let n: usize = sh("sysctl", &["-n", "hw.nperflevels"])?.parse().ok()?;
|
||||
let parts: Vec<String> = (0..n)
|
||||
.map(|i| {
|
||||
let cpus = sh("sysctl", &["-n", &format!("hw.perflevel{i}.logicalcpu")])?;
|
||||
let name = sh("sysctl", &["-n", &format!("hw.perflevel{i}.name")]).unwrap_or_else(|| format!("level-{i}"));
|
||||
Some(format!("{cpus} {name}"))
|
||||
})
|
||||
.collect::<Option<_>>()?;
|
||||
(!parts.is_empty()).then(|| parts.join(" + "))
|
||||
}
|
||||
|
||||
fn machine() -> String {
|
||||
let cpu = sh("sysctl", &["-n", "machdep.cpu.brand_string"])
|
||||
.or_else(|| {
|
||||
std::fs::read_to_string("/proc/cpuinfo").ok()?.lines().find(|l| l.starts_with("model name")).map(|l| l.split(':').nth(1).unwrap_or("").trim().to_string())
|
||||
})
|
||||
.unwrap_or_else(|| "unknown CPU".into());
|
||||
let levels = core_levels().map(|l| format!(" ({l})")).unwrap_or_default();
|
||||
format!(
|
||||
"{cpu}, {} logical cores{levels}; {}-{}; {} build",
|
||||
std::thread::available_parallelism().map(|n| n.get()).unwrap_or(0),
|
||||
std::env::consts::OS,
|
||||
std::env::consts::ARCH,
|
||||
if cfg!(debug_assertions) { "DEBUG (numbers not meaningful)" } else { "release" },
|
||||
)
|
||||
}
|
||||
|
||||
fn loadavg() -> String {
|
||||
sh("sysctl", &["-n", "vm.loadavg"])
|
||||
.map(|s| s.trim_matches(|c| c == '{' || c == '}' || c == ' ').to_string())
|
||||
.or_else(|| std::fs::read_to_string("/proc/loadavg").ok().map(|s| s.split_whitespace().take(3).collect::<Vec<_>>().join(" ")))
|
||||
.unwrap_or_else(|| "n/a".into())
|
||||
}
|
||||
|
||||
fn ms(d: Duration) -> f64 {
|
||||
d.as_secs_f64() * 1e3
|
||||
}
|
||||
|
||||
fn main() {
|
||||
let args: Vec<String> = std::env::args().collect();
|
||||
let arg = |k: &str| args.iter().position(|a| a == k).and_then(|i| args.get(i + 1)).cloned();
|
||||
let rounds_arg: Option<usize> = arg("--rounds").and_then(|s| s.parse().ok());
|
||||
let filter = arg("--filter");
|
||||
let quick = args.iter().any(|a| a == "--quick");
|
||||
|
||||
println!("# wai_decode_bench");
|
||||
println!();
|
||||
println!("- machine: {}", machine());
|
||||
println!("- fast-kernel automatic worker count: {}", available_threads());
|
||||
println!("- load average at start: {}", loadavg());
|
||||
println!("- method: per case, reference / fast 1-thread / fast auto-threads run interleaved (order rotated per round), same process, same inputs; best and median wall time over the rounds; outputs hashed every run and required identical across all three.");
|
||||
println!();
|
||||
|
||||
let mut cases = Vec::new();
|
||||
kernel_cases(&mut cases);
|
||||
decoder_cases(&mut cases, quick);
|
||||
vector_cases(&mut cases);
|
||||
if let Some(f) = &filter {
|
||||
cases.retain(|c| c.name.contains(f.as_str()));
|
||||
}
|
||||
|
||||
let policies = [("reference", Policy::reference()), ("fast 1T", Policy::fast(1)), ("fast auto", Policy::fast(0))];
|
||||
|
||||
// Layer mode: time every layer of the synthetic mbt2018-class g_s on its real
|
||||
// intermediate activations, per policy (same interleaving and identity check).
|
||||
// This attributes decode time to layers exactly, without a sampler.
|
||||
if args.iter().any(|a| a == "--layers") {
|
||||
layer_profile(&policies, rounds_arg.unwrap_or(5));
|
||||
return;
|
||||
}
|
||||
|
||||
// Profiling mode: run the selected cases under ONE policy for a fixed time so an
|
||||
// external sampler (e.g. `sample <pid>`) can attribute the time. No comparison,
|
||||
// no table — this mode measures nothing by itself.
|
||||
if let Some(which) = arg("--profile") {
|
||||
let p = match which.as_str() {
|
||||
"reference" => Policy::reference(),
|
||||
"1t" => Policy::fast(1),
|
||||
_ => Policy::fast(0),
|
||||
};
|
||||
let secs: f64 = arg("--seconds").and_then(|s| s.parse().ok()).unwrap_or(10.0);
|
||||
println!("profiling mode: {p:?}, {secs} s per case, pid {}", std::process::id());
|
||||
for case in &cases {
|
||||
let t0 = Instant::now();
|
||||
let mut n = 0;
|
||||
while t0.elapsed().as_secs_f64() < secs {
|
||||
let out = with_policy(p, || black_box((case.run)()));
|
||||
black_box(digest(&out));
|
||||
n += 1;
|
||||
}
|
||||
println!("{}: {n} runs", case.name);
|
||||
}
|
||||
return;
|
||||
}
|
||||
println!("| case | rounds | reference best / median (ms) | fast 1T best / median (ms) | fast auto best / median (ms) | throughput ref → 1T → auto | speed-up 1T / auto |");
|
||||
println!("|---|---:|---:|---:|---:|---|---:|");
|
||||
for case in &cases {
|
||||
// Warm-up: one run per policy (also fixes the reference hash).
|
||||
let mut t_ref = Duration::ZERO;
|
||||
let mut want = 0u64;
|
||||
for (i, (_, p)) in policies.iter().enumerate() {
|
||||
let t0 = Instant::now();
|
||||
let out = with_policy(*p, || black_box((case.run)()));
|
||||
let dt = t0.elapsed();
|
||||
let h = digest(&out);
|
||||
drop(out);
|
||||
if i == 0 {
|
||||
want = h;
|
||||
t_ref = dt;
|
||||
} else if h != want {
|
||||
eprintln!("MISMATCH in {}: {:?} gave {h:016x}, reference {want:016x}", case.name, p);
|
||||
std::process::exit(1);
|
||||
}
|
||||
}
|
||||
let rounds = rounds_arg.unwrap_or_else(|| ((4.0 / t_ref.as_secs_f64().max(1e-6)).ceil() as usize).clamp(3, 25));
|
||||
let mut times: Vec<Vec<Duration>> = vec![Vec::new(); policies.len()];
|
||||
for r in 0..rounds {
|
||||
for k in 0..policies.len() {
|
||||
let idx = (k + r) % policies.len();
|
||||
let (_, p) = policies[idx];
|
||||
let t0 = Instant::now();
|
||||
let out = with_policy(p, || black_box((case.run)()));
|
||||
times[idx].push(t0.elapsed());
|
||||
let h = digest(&out);
|
||||
drop(out);
|
||||
if h != want {
|
||||
eprintln!("MISMATCH in {}: {:?} gave {h:016x}, reference {want:016x}", case.name, p);
|
||||
std::process::exit(1);
|
||||
}
|
||||
}
|
||||
}
|
||||
let stat = |v: &mut Vec<Duration>| {
|
||||
v.sort();
|
||||
(v[0], v[v.len() / 2])
|
||||
};
|
||||
let s: Vec<(Duration, Duration)> = times.iter_mut().map(stat).collect();
|
||||
let (amount, unit) = case.per_run;
|
||||
let tp = |d: Duration| amount / d.as_secs_f64();
|
||||
let mac = case.macs.map(|m| format!(" ({:.2} → {:.2} → {:.2} GMAC/s)", m as f64 / s[0].0.as_secs_f64() / 1e9, m as f64 / s[1].0.as_secs_f64() / 1e9, m as f64 / s[2].0.as_secs_f64() / 1e9)).unwrap_or_default();
|
||||
println!(
|
||||
"| {} | {rounds} | {:.2} / {:.2} | {:.2} / {:.2} | {:.2} / {:.2} | {:.3} → {:.3} → {:.3} {unit}/s{mac} | {:.1}× / {:.1}× |",
|
||||
case.name,
|
||||
ms(s[0].0), ms(s[0].1), ms(s[1].0), ms(s[1].1), ms(s[2].0), ms(s[2].1),
|
||||
tp(s[0].0), tp(s[1].0), tp(s[2].0),
|
||||
s[0].0.as_secs_f64() / s[1].0.as_secs_f64(),
|
||||
s[0].0.as_secs_f64() / s[2].0.as_secs_f64(),
|
||||
);
|
||||
}
|
||||
println!();
|
||||
println!("- load average at end: {}", loadavg());
|
||||
println!("- all outputs identical across the three policies in every run");
|
||||
}
|
||||
|
|
@ -49,7 +49,7 @@ pub const CAP_NEURAL_WAVTOKENIZER: &str = "wai.neural.wavtokenizer"; // ultra-lo
|
|||
pub const CAP_NEURAL_BMSHJ2018: &str = "wai.neural.bmshj2018"; // bmshj2018-factorized image codec (2018 baseline)
|
||||
pub const CAP_NEURAL_MBT2018_MEAN: &str = "wai.neural.mbt2018_mean"; // mean-scale HYPERPRIOR image codec (Gaussian-conditional entropy) — decoder IMPLEMENTED in this crate (neural/hyperprior.rs); the conditional-entropy family, stronger than factorized bmshj2018; sink-supplied ONNX (hyper_synthesis + synthesis)
|
||||
pub const CAP_NEURAL_INT_SYNTH: &str = "wai.neural.int_synth"; // INTEGER fixed-point synthesis (neural/int_synth.rs): all-i64 transposed-conv decode, no float/ONNX — bit-identical on every machine BY CONSTRUCTION. The first capability that can soundly back intent=replicate cross-hardware (determinism IS integer-exact). Weights are sink-supplied (int_synth.json from a QAT model)
|
||||
pub const CAP_NEURAL_INT_HYPER: &str = "wai.neural.int_hyper"; // INTEGER fixed-point CONDITIONAL codec (int_hyper_synth.rs): the full Gaussian-conditional decode integer end-to-end — factorized rANS-decode z (z_q = s·2^afrac + median_q) → integer hyper-decoder h_s → σ,μ → integer-σ bucket (int_entropy) → integer rANS-decode residual → y_hat = r·2^afrac + μ → integer synthesis g_s. Both bitstreams entropy-coded, no raw latents. No float/ONNX anywhere, so the WHOLE bitstream→image decode is bit-identical on every machine BY CONSTRUCTION — int_synth's sibling that adds the conditional entropy. SOTA quality (PTQ of mbt2018_mean, ~42.6 dB). Tables sink-supplied
|
||||
pub const CAP_NEURAL_INT_HYPER: &str = "wai.neural.int_hyper"; // INTEGER fixed-point CONDITIONAL codec (int_hyper_synth.rs): the full Gaussian-conditional decode integer end-to-end — factorized rANS-decode z (z_q = s·2^afrac + median_q) → integer hyper-decoder h_s → σ,μ → integer-σ bucket (int_entropy) → integer rANS-decode residual → y_hat = r·2^afrac + μ → integer synthesis g_s. Both bitstreams entropy-coded, no raw latents. No float/ONNX anywhere, so the WHOLE bitstream→image decode is bit-identical on every machine BY CONSTRUCTION — int_synth's sibling that adds the conditional entropy. Quality: PTQ of the pretrained mbt2018_mean (quality 6); CI (wai-sota-vectors.yml) requires the integer decode within 1 dB of the same model's float decode on a pinned test photograph — measured cost there 0.19 dB (35.985 vs 36.175 dB, Apple Silicon Mac); image-dependent, see int_hyper_synth.rs. Tables sink-supplied
|
||||
pub const CAP_NEURAL_ELIC: &str = "wai.neural.elic"; // checkerboard CONTEXT-MODEL image codec (neural/context.rs) — the JPEG-AI/MLIC++ entropy class: hyperprior + 2-pass spatial context (anchor → context net → non-anchor). Decoder IMPLEMENTED; sink-supplied ONNX (hyper_synthesis + context + synthesis). WAI-native, NOT the ISO JPEG AI bitstream
|
||||
pub const CAP_NEURAL_MLICPP: &str = "wai.neural.mlicpp"; // MLIC++ (−13.39% BD-rate vs VVC/VTM-17): the multi-reference entropy model — channel-autoregressive slices + checkerboard + global attention. Decoder IMPLEMENTED (neural/mlicpp.rs): the slice/checkerboard decode orchestration; the attention nets are sink-supplied ONNX (hyper_synthesis + entropy + synthesis)
|
||||
pub const CAP_NEURAL_MLICV2: &str = "wai.neural.mlicv2"; // MLICv2 / MLICv2+ (2025): CURRENT SOTA learned image, −20.46% BD-rate vs VTM-17 (vs MLIC++'s −13.39%). SAME multi-reference channel-autoregressive decode orchestration as mlicpp — the gains are in the MetaFormer transform + enhanced entropy/global-context nets, all sink-supplied ONNX. Decoder route IMPLEMENTED (shares neural/mlicpp.rs)
|
||||
|
|
|
|||
|
|
@ -105,43 +105,67 @@ impl AudioModel {
|
|||
}
|
||||
}
|
||||
|
||||
/// One integer 1-D transposed conv with per-output-channel requant
|
||||
/// `((raw·mult)>>shift)+bias`. `x` is `[cin, t]`; `weight` is `[cin][cout][k]`.
|
||||
/// Returns `([cout, to], to)`. Pure i64 (i128 only for the requant multiply) —
|
||||
/// mirrors the export byte-for-byte.
|
||||
/// One integer 1-D transposed conv with per-output-channel requant — see
|
||||
/// [`reference::conv_t_1d`] for the definition. Runs the fast kernel
|
||||
/// ([`crate::int_kernels`], as a 1-row 2-D transposed conv) unless the current
|
||||
/// [`crate::int_kernels::Policy`] selects the reference; both give the same integer
|
||||
/// for every output element.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn conv_t_1d(
|
||||
x: &[i64], cin: usize, t: usize, weight: &[i64], cout: usize, k: usize,
|
||||
stride: usize, pad: usize, opad: usize, mult: &[i64], bias: &[i64], shift: u32,
|
||||
) -> (Vec<i64>, usize) {
|
||||
let to = (t - 1) * stride + k + opad - 2 * pad;
|
||||
let tp = to + 2 * pad;
|
||||
let mut raw = vec![0i64; cout * tp];
|
||||
for ic in 0..cin {
|
||||
let xc = &x[ic * t..ic * t + t];
|
||||
for oc in 0..cout {
|
||||
let wk = &weight[(ic * cout + oc) * k..(ic * cout + oc) * k + k];
|
||||
let plane = oc * tp;
|
||||
for (it, &v) in xc.iter().enumerate() {
|
||||
if v == 0 {
|
||||
continue;
|
||||
}
|
||||
let base = plane + it * stride;
|
||||
for kx in 0..k {
|
||||
raw[base + kx] += wk[kx] * v;
|
||||
use crate::int_kernels::{conv_acc_i64, policy, ConvGeom, KernelPath};
|
||||
if policy().path == KernelPath::Fast && mult.len() >= cout && bias.len() >= cout
|
||||
&& let Some(g) = ConvGeom::transposed(cin, 1, t, cout, 1, k, 1, stride, 0, pad, 0, opad, x.len(), weight.len())
|
||||
&& let Some(out) = conv_acc_i64(&g, x, weight, |oc, raw| ((raw as i128 * mult[oc] as i128) >> shift) as i64 + bias[oc])
|
||||
{
|
||||
return (out, g.wo);
|
||||
}
|
||||
reference::conv_t_1d(x, cin, t, weight, cout, k, stride, pad, opad, mult, bias, shift)
|
||||
}
|
||||
|
||||
/// The scalar reference kernel the fast path is checked against, kept verbatim.
|
||||
#[doc(hidden)]
|
||||
pub mod reference {
|
||||
/// One integer 1-D transposed conv with per-output-channel requant
|
||||
/// `((raw·mult)>>shift)+bias`. `x` is `[cin, t]`; `weight` is `[cin][cout][k]`.
|
||||
/// Returns `([cout, to], to)`. Pure i64 (i128 only for the requant multiply) —
|
||||
/// mirrors the export byte-for-byte.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn conv_t_1d(
|
||||
x: &[i64], cin: usize, t: usize, weight: &[i64], cout: usize, k: usize,
|
||||
stride: usize, pad: usize, opad: usize, mult: &[i64], bias: &[i64], shift: u32,
|
||||
) -> (Vec<i64>, usize) {
|
||||
let to = (t - 1) * stride + k + opad - 2 * pad;
|
||||
let tp = to + 2 * pad;
|
||||
let mut raw = vec![0i64; cout * tp];
|
||||
for ic in 0..cin {
|
||||
let xc = &x[ic * t..ic * t + t];
|
||||
for oc in 0..cout {
|
||||
let wk = &weight[(ic * cout + oc) * k..(ic * cout + oc) * k + k];
|
||||
let plane = oc * tp;
|
||||
for (it, &v) in xc.iter().enumerate() {
|
||||
if v == 0 {
|
||||
continue;
|
||||
}
|
||||
let base = plane + it * stride;
|
||||
for kx in 0..k {
|
||||
raw[base + kx] += wk[kx] * v;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
let mut out = vec![0i64; cout * to];
|
||||
for oc in 0..cout {
|
||||
let (m, b) = (mult[oc] as i128, bias[oc]);
|
||||
for ot in 0..to {
|
||||
let r = raw[oc * tp + ot + pad] as i128;
|
||||
out[oc * to + ot] = ((r * m) >> shift) as i64 + b;
|
||||
let mut out = vec![0i64; cout * to];
|
||||
for oc in 0..cout {
|
||||
let (m, b) = (mult[oc] as i128, bias[oc]);
|
||||
for ot in 0..to {
|
||||
let r = raw[oc * tp + ot + pad] as i128;
|
||||
out[oc * to + ot] = ((r * m) >> shift) as i64 + b;
|
||||
}
|
||||
}
|
||||
(out, to)
|
||||
}
|
||||
(out, to)
|
||||
}
|
||||
|
||||
/// Run the integer synthesis `g_s` on the fixed-point latent `x` (`[c, t]` at
|
||||
|
|
@ -261,3 +285,136 @@ fn b64_decode(s: &str) -> Result<Vec<u8>, &'static str> {
|
|||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// The fast 1-D transposed conv against the scalar reference, element by element.
|
||||
#[cfg(test)]
|
||||
mod kernel_equivalence {
|
||||
use super::*;
|
||||
use crate::int_kernels::testkit::{Rng, FAST};
|
||||
use crate::int_kernels::with_policy;
|
||||
|
||||
#[test]
|
||||
fn conv_t_1d_equals_reference() {
|
||||
let mut g = Rng(0xA0D_0001);
|
||||
let mut ran = 0;
|
||||
for case in 0..500 {
|
||||
let k = 1 + g.below(9);
|
||||
let stride = 1 + g.below(5);
|
||||
let pad = g.below(k);
|
||||
let opad = if stride > 1 { g.below(stride) } else { 0 };
|
||||
let t = 1 + g.below(40);
|
||||
let (cin, cout) = (1 + g.below(6), 1 + g.below(6));
|
||||
if (t - 1) * stride + k + opad <= 2 * pad {
|
||||
continue;
|
||||
}
|
||||
let (xr, wr) = [(1 << 15, 2047), (1 << 40, 1 << 10)][case % 2];
|
||||
let x = g.vec(cin * t, xr, 30);
|
||||
let wt = g.vec(cin * cout * k, wr, 5);
|
||||
let mult: Vec<i64> = (0..cout).map(|_| 1 + g.signed(1 << 16).abs()).collect();
|
||||
let bias = g.vec(cout, 1 << 16, 0);
|
||||
let want = reference::conv_t_1d(&x, cin, t, &wt, cout, k, stride, pad, opad, &mult, &bias, 24);
|
||||
for p in FAST {
|
||||
let got = with_policy(p, || conv_t_1d(&x, cin, t, &wt, cout, k, stride, pad, opad, &mult, &bias, 24));
|
||||
assert_eq!(got, want, "conv_t_1d {p:?} k={k} s={stride} p={pad} op={opad} t={t}");
|
||||
}
|
||||
ran += 1;
|
||||
}
|
||||
assert!(ran > 250, "too few valid cases ({ran})");
|
||||
}
|
||||
|
||||
/// Overflow only at a position the reference accumulates into and then crops:
|
||||
/// a length-1 input, kernel 3, pad 1 keeps only the middle of the reference's
|
||||
/// 3-sample raw buffer (`2·2^40·1 = 2^41`); both outer samples are
|
||||
/// `2·2^40·2^22 = 2^63` and overflow `i64`. The fast path must panic where the
|
||||
/// reference does and otherwise give the same 2^41.
|
||||
#[test]
|
||||
fn overflow_only_at_a_cropped_border_behaves_like_reference() {
|
||||
use crate::int_kernels::testkit::overflow_checks;
|
||||
use crate::int_kernels::Policy;
|
||||
let (cin, cout, k) = (2usize, 1usize, 3usize);
|
||||
let x = vec![1i64 << 40; cin];
|
||||
let wt: Vec<i64> = (0..cin * cout).flat_map(|_| [1i64 << 22, 1, 1 << 22]).collect();
|
||||
let (mult, bias) = (vec![1i64; cout], vec![0i64; cout]);
|
||||
let run = |p: Policy| {
|
||||
std::panic::catch_unwind(|| with_policy(p, || conv_t_1d(&x, cin, 1, &wt, cout, k, 1, 1, 0, &mult, &bias, 0))).ok()
|
||||
};
|
||||
let want = run(Policy::reference());
|
||||
if overflow_checks() {
|
||||
assert!(want.is_none(), "the reference overflows at a cropped sample in a checked build");
|
||||
} else {
|
||||
assert_eq!(want, Some((vec![1i64 << 41], 1)));
|
||||
}
|
||||
for p in FAST {
|
||||
assert_eq!(run(p), want, "fast must match the reference, panic included ({p:?})");
|
||||
}
|
||||
}
|
||||
|
||||
/// A one-sample input upsampled by any stride is one output sample (k = 3,
|
||||
/// pad 1): a handful of operations for the reference. The fast path must return
|
||||
/// the same sample without overflowing, or allocating or looping per stride step.
|
||||
#[test]
|
||||
fn a_one_sample_input_upsampled_by_any_stride() {
|
||||
use crate::int_kernels::Policy;
|
||||
let (x, wt) = ([5i64], [1i64, 2, 3]);
|
||||
// usize::MAX first: a kernel allocating a table per phase fails there at once
|
||||
// ("capacity overflow") instead of committing memory for 2^29 tables later.
|
||||
for s in [usize::MAX, 1, 2, 1 << 13, 1 << 16, 1 << 29] {
|
||||
let run = |p: Policy| with_policy(p, || conv_t_1d(&x, 1, 1, &wt, 1, 3, s, 1, 0, &[1], &[0], 0));
|
||||
let want = run(Policy::reference());
|
||||
assert_eq!(want, (vec![10], 1), "stride {s}: the middle of 5·[1, 2, 3]");
|
||||
for p in FAST {
|
||||
assert_eq!(run(p), want, "stride {s} {p:?}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Several input samples upsampled by a stride and cropped by padding so large
|
||||
/// that the reference's uncropped buffer `cout · (to + 2·pad)` cannot be sized:
|
||||
/// three samples, kernel 1, stride = pad = 2^62, one output sample. With two
|
||||
/// output channels the size overflows `usize` (the reference panics with
|
||||
/// overflow checks on, and with them off indexes out of the wrapped buffer);
|
||||
/// with one it is 2^63 + 1 samples, more than any allocation can hold ("capacity
|
||||
/// overflow"). The reference panics in every build, so every policy must, and
|
||||
/// the fast kernel must not be entered.
|
||||
#[cfg(target_pointer_width = "64")]
|
||||
#[test]
|
||||
fn several_samples_whose_uncropped_buffer_cannot_be_sized() {
|
||||
use crate::int_kernels::testkit::fast_conv_runs;
|
||||
use crate::int_kernels::Policy;
|
||||
let s = 1usize << 62;
|
||||
for cout in [2usize, 1] {
|
||||
let (x, wt) = ([1i64, 2, 3], vec![7i64; cout]);
|
||||
let (mult, bias) = (vec![1i64; cout], vec![0i64; cout]);
|
||||
let run = |p: Policy| {
|
||||
fast_conv_runs(|| std::panic::catch_unwind(|| with_policy(p, || conv_t_1d(&x, 1, 3, &wt, cout, 1, s, s, 0, &mult, &bias, 0))).ok())
|
||||
};
|
||||
assert_eq!(run(Policy::reference()), (None, 0), "{cout} channels: the reference panics");
|
||||
for p in FAST {
|
||||
assert_eq!(run(p), (None, 0), "{cout} channels {p:?}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Geometry whose padded input would be more than `PAD_FACTOR` (16) times what
|
||||
/// the reference reads and writes is refused by the fast path, and the reference
|
||||
/// runs: the fast kernel is not entered and the samples are the reference's. At
|
||||
/// exactly the factor the fast kernel runs. A one-sample input, kernel 1, stride
|
||||
/// 1, cropped by `pad` with output padding `2·pad` restoring one output sample
|
||||
/// that no tap reaches: the padded input is `cin · (1 + pad)` samples against
|
||||
/// `cin + cin + 1` the reference touches (input, weights, output).
|
||||
#[test]
|
||||
fn padding_beyond_the_pad_factor_runs_the_reference() {
|
||||
use crate::int_kernels::testkit::fast_conv_runs;
|
||||
// (channels, pad, expected fast-kernel runs): 1 · 48 = 16 · 3 and
|
||||
// 2 · 40 = 16 · 5 at the factor; 1 · 49 and 2 · 41 beyond it.
|
||||
for (cin, pad, runs) in [(1usize, 47usize, 1u64), (1, 48, 0), (2, 39, 1), (2, 40, 0)] {
|
||||
let (x, wt) = (vec![5i64; cin], vec![3i64; cin]);
|
||||
let want = reference::conv_t_1d(&x, cin, 1, &wt, 1, 1, 1, pad, 2 * pad, &[3], &[7], 1);
|
||||
assert_eq!(want, (vec![7], 1), "reference, {cin} channels, pad {pad}");
|
||||
for p in FAST {
|
||||
let got = fast_conv_runs(|| with_policy(p, || conv_t_1d(&x, cin, 1, &wt, 1, 1, 1, pad, 2 * pad, &[3], &[7], 1)));
|
||||
assert_eq!(got, (want.clone(), runs), "{cin} channels, pad {pad} {p:?}");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,14 +1,26 @@
|
|||
//! Integer-exact synthesis of a SOTA-class learned codec (per-channel int12
|
||||
//! Integer-exact synthesis of a pretrained learned image codec (per-channel int12
|
||||
//! weights + integer IGDN) — the decode side of a `NeuralIntegerExact`
|
||||
//! conditional codec at SOTA quality. Pure i64, no float, no `ort`: the same
|
||||
//! conditional codec, run on the pretrained mbt2018_mean mean-scale hyperprior
|
||||
//! exported to integers. Pure i64, no float, no `ort`: the same
|
||||
//! discipline as `int_synth`/`int_entropy`, so the reconstruction is
|
||||
//! byte-identical on every machine.
|
||||
//!
|
||||
//! It runs the model the export tool emits (`tools/wai_sota_int_synth_export.py`):
|
||||
//! a stack of integer transposed-convolutions (per-output-channel `(mult,shift)`
|
||||
//! requant) and inverse-GDN activations ([`crate::int_transform::igdn`]). On a
|
||||
//! real pretrained `mbt2018_mean` it reconstructs at 42.6 dB — within ~0.6 dB of
|
||||
//! the float baseline, fully integer.
|
||||
//! requant) and inverse-GDN activations ([`crate::int_transform::igdn`]).
|
||||
//!
|
||||
//! Quality is image-dependent, so it is measured on a pinned image rather than
|
||||
//! quoted. On the public-domain test photograph `tools/wai_sota_vectors.sh`
|
||||
//! regenerates the vectors from (a 256×256 crop of scikit-image's `astronaut`),
|
||||
//! measured on an Apple Silicon Mac: the integer decode of the pretrained
|
||||
//! `mbt2018_mean`'s full conditional bitstream scores 35.985 dB (8-bit PSNR of this
|
||||
//! decoder's output against the source pixels, `tests/int_hyper_synth.rs`), against
|
||||
//! 36.175 dB for the same model's float decode (the export tool's PSNR of its float
|
||||
//! reconstruction) — a cost of 0.19 dB. CI (`wai-sota-vectors.yml`) requires that
|
||||
//! cost to stay within 1 dB on every run. It is a relative bar on one image and does
|
||||
//! not compare the model with any other codec. An earlier development figure of
|
||||
//! 42.6 dB (43.25 dB float), for the integer synthesis alone, was measured on a
|
||||
//! photograph that is not in the repository and cannot be reproduced from it.
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
|
|
@ -23,129 +35,172 @@ fn i64s(v: &Value) -> Vec<i64> {
|
|||
v.as_array().map(|a| a.iter().filter_map(|n| n.as_i64()).collect()).unwrap_or_default()
|
||||
}
|
||||
|
||||
/// One integer transposed conv + per-output-channel requant `((raw·mult)>>shift)
|
||||
/// + bias`. `x` is `[cin, h·w]`; `weight` is `[cin][cout][k][k]`. Returns
|
||||
/// `([cout, ho·wo], ho, wo)`. Pure i64 (the accumulation stays well inside i64).
|
||||
/// One integer transposed conv + per-output-channel requant
|
||||
/// `((raw·mult)>>shift) + bias`. `x` is `[cin, h·w]`; `weight` is `[cin][cout][k][k]`.
|
||||
/// Returns `([cout, ho·wo], ho, wo)`. Runs [`reference::conv_t`] or the fast kernel
|
||||
/// ([`crate::int_kernels`]) per the current [`crate::int_kernels::Policy`]; both give
|
||||
/// the same integer for every output element.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn conv_t(
|
||||
x: &[i64], cin: usize, h: usize, w: usize, weight: &[i64], cout: usize, k: usize,
|
||||
stride: usize, pad: usize, opad: usize, mult: &[i64], bias: &[i64], shift: u32,
|
||||
) -> (Vec<i64>, usize, usize) {
|
||||
let ho = (h - 1) * stride + k + opad - 2 * pad;
|
||||
let wo = (w - 1) * stride + k + opad - 2 * pad;
|
||||
let (hp, wp) = (ho + 2 * pad, wo + 2 * pad);
|
||||
let kk = k * k;
|
||||
let mut raw = vec![0i64; cout * hp * wp];
|
||||
// Loop order ic → oc → spatial → kernel: for a fixed (ic, oc) the k×k weight
|
||||
// block is loaded once and every scatter lands in the single `oc` output plane,
|
||||
// so that plane (and the input row) stays hot in cache. Integer accumulation is
|
||||
// order-independent, so this is byte-identical to the naive oc-inner order.
|
||||
for ic in 0..cin {
|
||||
let xc = &x[ic * h * w..ic * h * w + h * w];
|
||||
for oc in 0..cout {
|
||||
let wb = (ic * cout + oc) * kk;
|
||||
let wk = &weight[wb..wb + kk];
|
||||
let plane = oc * hp * wp;
|
||||
for iy in 0..h {
|
||||
let ry0 = iy * stride;
|
||||
let xrow = &xc[iy * w..iy * w + w];
|
||||
for (ix, &v) in xrow.iter().enumerate() {
|
||||
if v == 0 {
|
||||
continue;
|
||||
}
|
||||
let base = plane + ry0 * wp + ix * stride;
|
||||
for ky in 0..k {
|
||||
let dst = base + ky * wp;
|
||||
let wky = ky * k;
|
||||
for kx in 0..k {
|
||||
raw[dst + kx] += wk[wky + kx] * v;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
use crate::int_kernels::{conv_acc_i64, policy, ConvGeom, KernelPath};
|
||||
if policy().path == KernelPath::Fast && mult.len() >= cout && bias.len() >= cout
|
||||
&& let Some(g) = ConvGeom::transposed(cin, h, w, cout, k, k, stride, stride, pad, pad, opad, opad, x.len(), weight.len())
|
||||
&& let Some(out) = conv_acc_i64(&g, x, weight, |oc, raw| ((raw as i128 * mult[oc] as i128) >> shift) as i64 + bias[oc])
|
||||
{
|
||||
return (out, g.ho, g.wo);
|
||||
}
|
||||
let mut out = vec![0i64; cout * ho * wo];
|
||||
for oc in 0..cout {
|
||||
let (m, b) = (mult[oc] as i128, bias[oc]);
|
||||
for oy in 0..ho {
|
||||
for ox in 0..wo {
|
||||
let r = raw[(oc * hp + oy + pad) * wp + ox + pad] as i128;
|
||||
out[(oc * ho + oy) * wo + ox] = ((r * m) >> shift) as i64 + b;
|
||||
}
|
||||
}
|
||||
}
|
||||
(out, ho, wo)
|
||||
reference::conv_t(x, cin, h, w, weight, cout, k, stride, pad, opad, mult, bias, shift)
|
||||
}
|
||||
|
||||
/// One integer **regular** conv (stride/pad, no output-pad) + per-output-channel
|
||||
/// requant — the stride-1 `Conv2d` that closes the hyper-decoder (`M·3/2 → M·2`,
|
||||
/// emitting σ‖μ). `weight` is `[cout][cin][k][k]`. Mirrors the export's
|
||||
/// `conv_int` byte-for-byte. Pure i64 (i128 only for the requant multiply).
|
||||
/// requant — see [`reference::conv`]. Dispatches like [`conv_t`].
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn conv(
|
||||
x: &[i64], cin: usize, h: usize, w: usize, weight: &[i64], cout: usize, k: usize,
|
||||
stride: usize, pad: usize, mult: &[i64], bias: &[i64], shift: u32,
|
||||
) -> (Vec<i64>, usize, usize) {
|
||||
let ho = (h + 2 * pad - k) / stride + 1;
|
||||
let wo = (w + 2 * pad - k) / stride + 1;
|
||||
let kk = k * k;
|
||||
let mut out = vec![0i64; cout * ho * wo]; // raw sums first; requantized below
|
||||
// Loop order oc → ic → kernel → spatial: the weight tap is a hoisted scalar
|
||||
// and the inner `ox` loop streams one output row against one input row (great
|
||||
// locality, vectorizable). For stride 1 the valid `ox` span is precomputed so
|
||||
// the inner loop is branch-free. Integer accumulation is order-independent, so
|
||||
// this is byte-identical to the naive output-stationary order.
|
||||
for oc in 0..cout {
|
||||
let plane = oc * ho * wo;
|
||||
use crate::int_kernels::{conv_acc_i64, policy, ConvGeom, KernelPath};
|
||||
if policy().path == KernelPath::Fast && mult.len() >= cout && bias.len() >= cout
|
||||
&& let Some(g) = ConvGeom::conv(cin, h, w, cout, 1, k, k, stride, stride, pad, pad, x.len(), weight.len())
|
||||
&& let Some(out) = conv_acc_i64(&g, x, weight, |oc, raw| ((raw as i128 * mult[oc] as i128) >> shift) as i64 + bias[oc])
|
||||
{
|
||||
return (out, g.ho, g.wo);
|
||||
}
|
||||
reference::conv(x, cin, h, w, weight, cout, k, stride, pad, mult, bias, shift)
|
||||
}
|
||||
|
||||
/// The scalar reference kernels — the executable specification the fast kernels
|
||||
/// in [`crate::int_kernels`] are checked against, kept verbatim. Selected for the
|
||||
/// whole decode by `int_kernels::with_policy(Policy::reference(), ..)`.
|
||||
#[doc(hidden)]
|
||||
pub mod reference {
|
||||
/// One integer transposed conv + per-output-channel requant `((raw·mult)>>shift)
|
||||
/// + bias`. `x` is `[cin, h·w]`; `weight` is `[cin][cout][k][k]`. Returns
|
||||
/// `([cout, ho·wo], ho, wo)`. Pure i64 (the accumulation stays well inside i64).
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn conv_t(
|
||||
x: &[i64], cin: usize, h: usize, w: usize, weight: &[i64], cout: usize, k: usize,
|
||||
stride: usize, pad: usize, opad: usize, mult: &[i64], bias: &[i64], shift: u32,
|
||||
) -> (Vec<i64>, usize, usize) {
|
||||
let ho = (h - 1) * stride + k + opad - 2 * pad;
|
||||
let wo = (w - 1) * stride + k + opad - 2 * pad;
|
||||
let (hp, wp) = (ho + 2 * pad, wo + 2 * pad);
|
||||
let kk = k * k;
|
||||
let mut raw = vec![0i64; cout * hp * wp];
|
||||
// Loop order ic → oc → spatial → kernel: for a fixed (ic, oc) the k×k weight
|
||||
// block is loaded once and every scatter lands in the single `oc` output plane,
|
||||
// so that plane (and the input row) stays hot in cache. Integer accumulation is
|
||||
// order-independent, so this is byte-identical to the naive oc-inner order.
|
||||
for ic in 0..cin {
|
||||
let xb = ic * h * w;
|
||||
let wb = (oc * cin + ic) * kk;
|
||||
for ky in 0..k {
|
||||
for kx in 0..k {
|
||||
let wv = weight[wb + ky * k + kx];
|
||||
if wv == 0 {
|
||||
continue;
|
||||
}
|
||||
// valid output columns: input ix = ox·stride + kx − pad ∈ [0, w)
|
||||
let (ox_lo, ox_hi) = if stride == 1 {
|
||||
(pad.saturating_sub(kx), (w + pad).saturating_sub(kx).min(wo))
|
||||
} else {
|
||||
(0, wo)
|
||||
};
|
||||
for oy in 0..ho {
|
||||
let iy = oy * stride + ky;
|
||||
if iy < pad || iy >= pad + h {
|
||||
let xc = &x[ic * h * w..ic * h * w + h * w];
|
||||
for oc in 0..cout {
|
||||
let wb = (ic * cout + oc) * kk;
|
||||
let wk = &weight[wb..wb + kk];
|
||||
let plane = oc * hp * wp;
|
||||
for iy in 0..h {
|
||||
let ry0 = iy * stride;
|
||||
let xrow = &xc[iy * w..iy * w + w];
|
||||
for (ix, &v) in xrow.iter().enumerate() {
|
||||
if v == 0 {
|
||||
continue;
|
||||
}
|
||||
let xrow = xb + (iy - pad) * w;
|
||||
let orow = plane + oy * wo;
|
||||
if stride == 1 {
|
||||
// over [ox_lo, ox_hi): ox + kx ≥ pad, so ox+kx−pad ≥ 0.
|
||||
for ox in ox_lo..ox_hi {
|
||||
out[orow + ox] += wv * x[xrow + ox + kx - pad];
|
||||
}
|
||||
} else {
|
||||
for ox in 0..wo {
|
||||
let ix = ox * stride + kx;
|
||||
if ix < pad || ix >= pad + w {
|
||||
continue;
|
||||
}
|
||||
out[orow + ox] += wv * x[xrow + ix - pad];
|
||||
let base = plane + ry0 * wp + ix * stride;
|
||||
for ky in 0..k {
|
||||
let dst = base + ky * wp;
|
||||
let wky = ky * k;
|
||||
for kx in 0..k {
|
||||
raw[dst + kx] += wk[wky + kx] * v;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
for oc in 0..cout {
|
||||
let (m, b) = (mult[oc] as i128, bias[oc]);
|
||||
for v in &mut out[oc * ho * wo..(oc + 1) * ho * wo] {
|
||||
*v = ((*v as i128 * m) >> shift) as i64 + b;
|
||||
let mut out = vec![0i64; cout * ho * wo];
|
||||
for oc in 0..cout {
|
||||
let (m, b) = (mult[oc] as i128, bias[oc]);
|
||||
for oy in 0..ho {
|
||||
for ox in 0..wo {
|
||||
let r = raw[(oc * hp + oy + pad) * wp + ox + pad] as i128;
|
||||
out[(oc * ho + oy) * wo + ox] = ((r * m) >> shift) as i64 + b;
|
||||
}
|
||||
}
|
||||
}
|
||||
(out, ho, wo)
|
||||
}
|
||||
|
||||
/// One integer **regular** conv (stride/pad, no output-pad) + per-output-channel
|
||||
/// requant — the stride-1 `Conv2d` that closes the hyper-decoder (`M·3/2 → M·2`,
|
||||
/// emitting σ‖μ). `weight` is `[cout][cin][k][k]`. Mirrors the export's
|
||||
/// `conv_int` byte-for-byte. Pure i64 (i128 only for the requant multiply).
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn conv(
|
||||
x: &[i64], cin: usize, h: usize, w: usize, weight: &[i64], cout: usize, k: usize,
|
||||
stride: usize, pad: usize, mult: &[i64], bias: &[i64], shift: u32,
|
||||
) -> (Vec<i64>, usize, usize) {
|
||||
let ho = (h + 2 * pad - k) / stride + 1;
|
||||
let wo = (w + 2 * pad - k) / stride + 1;
|
||||
let kk = k * k;
|
||||
let mut out = vec![0i64; cout * ho * wo]; // raw sums first; requantized below
|
||||
// Loop order oc → ic → kernel → spatial: the weight tap is a hoisted scalar
|
||||
// and the inner `ox` loop streams one output row against one input row (great
|
||||
// locality, vectorizable). For stride 1 the valid `ox` span is precomputed so
|
||||
// the inner loop is branch-free. Integer accumulation is order-independent, so
|
||||
// this is byte-identical to the naive output-stationary order.
|
||||
for oc in 0..cout {
|
||||
let plane = oc * ho * wo;
|
||||
for ic in 0..cin {
|
||||
let xb = ic * h * w;
|
||||
let wb = (oc * cin + ic) * kk;
|
||||
for ky in 0..k {
|
||||
for kx in 0..k {
|
||||
let wv = weight[wb + ky * k + kx];
|
||||
if wv == 0 {
|
||||
continue;
|
||||
}
|
||||
// valid output columns: input ix = ox·stride + kx − pad ∈ [0, w)
|
||||
let (ox_lo, ox_hi) = if stride == 1 {
|
||||
(pad.saturating_sub(kx), (w + pad).saturating_sub(kx).min(wo))
|
||||
} else {
|
||||
(0, wo)
|
||||
};
|
||||
for oy in 0..ho {
|
||||
let iy = oy * stride + ky;
|
||||
if iy < pad || iy >= pad + h {
|
||||
continue;
|
||||
}
|
||||
let xrow = xb + (iy - pad) * w;
|
||||
let orow = plane + oy * wo;
|
||||
if stride == 1 {
|
||||
// over [ox_lo, ox_hi): ox + kx ≥ pad, so ox+kx−pad ≥ 0.
|
||||
for ox in ox_lo..ox_hi {
|
||||
out[orow + ox] += wv * x[xrow + ox + kx - pad];
|
||||
}
|
||||
} else {
|
||||
for ox in 0..wo {
|
||||
let ix = ox * stride + kx;
|
||||
if ix < pad || ix >= pad + w {
|
||||
continue;
|
||||
}
|
||||
out[orow + ox] += wv * x[xrow + ix - pad];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
for oc in 0..cout {
|
||||
let (m, b) = (mult[oc] as i128, bias[oc]);
|
||||
for v in &mut out[oc * ho * wo..(oc + 1) * ho * wo] {
|
||||
*v = ((*v as i128 * m) >> shift) as i64 + b;
|
||||
}
|
||||
}
|
||||
(out, ho, wo)
|
||||
}
|
||||
(out, ho, wo)
|
||||
}
|
||||
|
||||
/// A parsed integer conv net — the typed form [`run_net`] consumes, built from
|
||||
|
|
@ -478,3 +533,251 @@ pub fn decode_conditional(
|
|||
// 3. integer synthesis g_s → pixels.
|
||||
Ok(decode(g_s, &y_hat, bs.m, sh, sw))
|
||||
}
|
||||
|
||||
/// 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
|
||||
/// than the stride.
|
||||
#[cfg(test)]
|
||||
mod kernel_equivalence {
|
||||
use super::*;
|
||||
use crate::int_kernels::testkit::{Rng, FAST};
|
||||
use crate::int_kernels::{with_policy, Policy};
|
||||
|
||||
/// (input range, weight range) per case: fits-i32, beyond-i32, mixed.
|
||||
const RANGES: [(i64, i64); 3] = [(1 << 14, 2047), (1 << 40, 1 << 10), (1 << 20, 1 << 20)];
|
||||
|
||||
#[test]
|
||||
fn transposed_conv_equals_reference() {
|
||||
let mut g = Rng(0xC0FF_EE01);
|
||||
let mut ran = 0;
|
||||
for case in 0..600 {
|
||||
let k = 1 + g.below(5);
|
||||
let stride = 1 + g.below(3);
|
||||
let pad = g.below(k);
|
||||
let opad = if stride > 1 { g.below(stride) } else { 0 };
|
||||
let (h, w) = (1 + g.below(7), 1 + g.below(9));
|
||||
let (cin, cout) = (1 + g.below(6), 1 + g.below(9));
|
||||
if (h - 1) * stride + k + opad <= 2 * pad || (w - 1) * stride + k + opad <= 2 * pad {
|
||||
continue;
|
||||
}
|
||||
let (xr, wr) = RANGES[case % 3];
|
||||
let x = g.vec(cin * h * w, xr, [0, 50, 90][case % 3]);
|
||||
let wt = g.vec(cin * cout * k * k, wr, 10);
|
||||
let mult: Vec<i64> = (0..cout).map(|_| 1 + g.signed(1 << 20).abs()).collect();
|
||||
let bias = g.vec(cout, 1 << 20, 0);
|
||||
let shift = 16 + g.below(8) as u32;
|
||||
let want = reference::conv_t(&x, cin, h, w, &wt, cout, k, stride, pad, opad, &mult, &bias, shift);
|
||||
for p in FAST {
|
||||
let got = with_policy(p, || conv_t(&x, cin, h, w, &wt, cout, k, stride, pad, opad, &mult, &bias, shift));
|
||||
assert_eq!(got, want, "conv_t {p:?} k={k} s={stride} p={pad} op={opad} {cin}x{h}x{w}->{cout}");
|
||||
}
|
||||
ran += 1;
|
||||
}
|
||||
assert!(ran > 300, "too few valid cases ({ran})");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn regular_conv_equals_reference() {
|
||||
let mut g = Rng(0xC0FF_EE02);
|
||||
let mut ran = 0;
|
||||
for case in 0..600 {
|
||||
let k = 1 + g.below(5);
|
||||
let stride = 1 + g.below(3);
|
||||
let pad = g.below(k);
|
||||
let (h, w) = (1 + g.below(9), 1 + g.below(11));
|
||||
let (cin, cout) = (1 + g.below(6), 1 + g.below(9));
|
||||
if h + 2 * pad < k || w + 2 * pad < k {
|
||||
continue;
|
||||
}
|
||||
let (xr, wr) = RANGES[case % 3];
|
||||
let x = g.vec(cin * h * w, xr, [0, 50, 90][case % 3]);
|
||||
let wt = g.vec(cin * cout * k * k, wr, 10);
|
||||
let mult: Vec<i64> = (0..cout).map(|_| 1 + g.signed(1 << 20).abs()).collect();
|
||||
let bias = g.vec(cout, 1 << 20, 0);
|
||||
let shift = 16 + g.below(8) as u32;
|
||||
let want = reference::conv(&x, cin, h, w, &wt, cout, k, stride, pad, &mult, &bias, shift);
|
||||
for p in FAST {
|
||||
let got = with_policy(p, || conv(&x, cin, h, w, &wt, cout, k, stride, pad, &mult, &bias, shift));
|
||||
assert_eq!(got, want, "conv {p:?} k={k} s={stride} p={pad} {cin}x{h}x{w}->{cout}");
|
||||
}
|
||||
ran += 1;
|
||||
}
|
||||
assert!(ran > 300, "too few valid cases ({ran})");
|
||||
}
|
||||
|
||||
/// Strides up to 16, beyond the kernel and the input. The transposed kernel's
|
||||
/// row bands then hold fewer rows than the stride (a tap table only for each row
|
||||
/// phase present) as well as more (rows wrapping onto the same phases), with
|
||||
/// bands starting anywhere in a phase; many phases receive no tap at all. The
|
||||
/// regular kernel's stride often exceeds its padded row.
|
||||
#[test]
|
||||
fn wide_strides_equal_reference() {
|
||||
let mut g = Rng(0xC0FF_EE03);
|
||||
let (mut ran_t, mut ran_c) = (0, 0);
|
||||
for case in 0..300 {
|
||||
let k = 1 + g.below(5);
|
||||
let stride = 1 + g.below(16);
|
||||
let pad = g.below(k);
|
||||
let opad = g.below(stride);
|
||||
let (h, w) = (1 + g.below(4), 1 + g.below(5));
|
||||
let (cin, cout) = (1 + g.below(4), 1 + g.below(6));
|
||||
let (xr, wr) = RANGES[case % 3];
|
||||
let x = g.vec(cin * h * w, xr, [0, 50, 90][case % 3]);
|
||||
let wt = g.vec(cin * cout * k * k, wr, 10);
|
||||
let mult: Vec<i64> = (0..cout).map(|_| 1 + g.signed(1 << 20).abs()).collect();
|
||||
let bias = g.vec(cout, 1 << 20, 0);
|
||||
let shift = 16 + g.below(8) as u32;
|
||||
if (h - 1) * stride + k + opad > 2 * pad && (w - 1) * stride + k + opad > 2 * pad {
|
||||
let want = reference::conv_t(&x, cin, h, w, &wt, cout, k, stride, pad, opad, &mult, &bias, shift);
|
||||
for p in FAST {
|
||||
let got = with_policy(p, || conv_t(&x, cin, h, w, &wt, cout, k, stride, pad, opad, &mult, &bias, shift));
|
||||
assert_eq!(got, want, "conv_t {p:?} k={k} s={stride} p={pad} op={opad} {cin}x{h}x{w}->{cout}");
|
||||
}
|
||||
ran_t += 1;
|
||||
}
|
||||
if h + 2 * pad >= k && w + 2 * pad >= k {
|
||||
let want = reference::conv(&x, cin, h, w, &wt, cout, k, stride, pad, &mult, &bias, shift);
|
||||
for p in FAST {
|
||||
let got = with_policy(p, || conv(&x, cin, h, w, &wt, cout, k, stride, pad, &mult, &bias, shift));
|
||||
assert_eq!(got, want, "conv {p:?} k={k} s={stride} p={pad} {cin}x{h}x{w}->{cout}");
|
||||
}
|
||||
ran_c += 1;
|
||||
}
|
||||
}
|
||||
assert!(ran_t > 150 && ran_c > 150, "too few valid cases ({ran_t} transposed, {ran_c} regular)");
|
||||
}
|
||||
|
||||
/// Operands whose `i64` accumulation overflows: the reference panics in an
|
||||
/// overflow-checked build and wraps otherwise, and the fast kernel must do the
|
||||
/// same — a panic where it panics, the identical wrapped integers where it wraps.
|
||||
#[test]
|
||||
fn overflowing_accumulation_behaves_like_reference() {
|
||||
let (cin, h, w, cout, k) = (4usize, 3usize, 3usize, 2usize, 3usize);
|
||||
let x = vec![i32::MAX as i64; cin * h * w];
|
||||
let wt = vec![i32::MAX as i64; cin * cout * k * k];
|
||||
let (mult, bias) = (vec![1i64; cout], vec![0i64; cout]);
|
||||
let run = |p: Policy| {
|
||||
std::panic::catch_unwind(|| {
|
||||
with_policy(p, || conv_t(&x, cin, h, w, &wt, cout, k, 2, 1, 1, &mult, &bias, 0))
|
||||
})
|
||||
};
|
||||
let want = run(Policy::reference()).ok();
|
||||
for p in FAST {
|
||||
let got = run(p).ok();
|
||||
if crate::int_kernels::testkit::overflow_checks() {
|
||||
assert!(want.is_none() && got.is_none(), "both must panic on overflow in a checked build ({p:?})");
|
||||
} else {
|
||||
assert!(want.is_some(), "an unchecked build wraps rather than panics");
|
||||
assert_eq!(got, want, "both must wrap identically ({p:?})");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The same with rows wide enough (16 outputs) for the SIMD lanes to engage: the
|
||||
/// lanes wrap instead of trapping, so they must not run where an accumulator can
|
||||
/// overflow — in a checked build the fast path must still panic like the reference.
|
||||
#[test]
|
||||
fn overflowing_accumulation_in_simd_width_rows_behaves_like_reference() {
|
||||
let (cin, h, w, cout, k) = (4usize, 3usize, 16usize, 4usize, 3usize);
|
||||
let x = vec![i32::MAX as i64; cin * h * w];
|
||||
let wt = vec![i32::MAX as i64; cin * cout * k * k];
|
||||
let (mult, bias) = (vec![1i64; cout], vec![0i64; cout]);
|
||||
let run = |p: Policy| {
|
||||
std::panic::catch_unwind(|| with_policy(p, || conv(&x, cin, h, w, &wt, cout, k, 1, 1, &mult, &bias, 0))).ok()
|
||||
};
|
||||
let want = run(Policy::reference());
|
||||
for p in FAST {
|
||||
let got = run(p);
|
||||
if crate::int_kernels::testkit::overflow_checks() {
|
||||
assert!(want.is_none() && got.is_none(), "both must panic on overflow in a checked build ({p:?})");
|
||||
} else {
|
||||
assert_eq!(got, want, "both must wrap identically ({p:?})");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Overflow only at a position the reference accumulates into and then crops.
|
||||
/// A 1×1 input with a 3×3 kernel, stride 1 and pad 1 keeps only the centre of
|
||||
/// the reference's 3×3 raw buffer. The centre is `2·2^40·1 = 2^41`. Each corner
|
||||
/// is `2·2^40·2^22 = 2^63` and overflows `i64` on its second term, which the
|
||||
/// fast kernel never computes. So the reference panics in an overflow-checked
|
||||
/// build and the fast path must too (these operands are outside its no-overflow
|
||||
/// bound, so it hands them to the reference). Without the checks both give 2^41.
|
||||
#[test]
|
||||
fn overflow_only_at_a_cropped_border_behaves_like_reference() {
|
||||
use crate::int_kernels::testkit::overflow_checks;
|
||||
let (cin, cout, k) = (2usize, 1usize, 3usize);
|
||||
let x = vec![1i64 << 40; cin];
|
||||
let mut wt = vec![0i64; cin * cout * k * k];
|
||||
for tap in wt.chunks_mut(k * k) {
|
||||
for c in [0, 2, 6, 8] {
|
||||
tap[c] = 1 << 22;
|
||||
}
|
||||
tap[4] = 1;
|
||||
}
|
||||
let (mult, bias) = (vec![1i64; cout], vec![0i64; cout]);
|
||||
let run = |p: Policy| {
|
||||
std::panic::catch_unwind(|| with_policy(p, || conv_t(&x, cin, 1, 1, &wt, cout, k, 1, 1, 0, &mult, &bias, 0))).ok()
|
||||
};
|
||||
let want = run(Policy::reference());
|
||||
if overflow_checks() {
|
||||
assert!(want.is_none(), "the reference overflows at a cropped corner in a checked build");
|
||||
} else {
|
||||
assert_eq!(want, Some((vec![1i64 << 41], 1, 1)), "the kept centre is unaffected by the wrapped corners");
|
||||
}
|
||||
for p in FAST {
|
||||
assert_eq!(run(p), want, "fast must match the reference, panic included ({p:?})");
|
||||
}
|
||||
}
|
||||
|
||||
/// Malformed geometry goes to the reference, so it fails exactly as before.
|
||||
#[test]
|
||||
fn malformed_geometry_is_left_to_the_reference() {
|
||||
let x = vec![1i64; 4];
|
||||
let short_w = vec![1i64; 3]; // needs 1·1·3·3
|
||||
let r = std::panic::catch_unwind(|| conv_t(&x, 1, 2, 2, &short_w, 1, 3, 1, 0, 0, &[1], &[0], 0));
|
||||
let q = std::panic::catch_unwind(|| {
|
||||
with_policy(Policy::reference(), || conv_t(&x, 1, 2, 2, &short_w, 1, 3, 1, 0, 0, &[1], &[0], 0))
|
||||
});
|
||||
assert_eq!(r.is_err(), q.is_err());
|
||||
}
|
||||
|
||||
/// Geometry whose padded input would be more than `PAD_FACTOR` times what the
|
||||
/// reference reads and writes is refused by the fast path, and the reference
|
||||
/// runs: the fast kernel is not entered and the integers are the reference's.
|
||||
/// At exactly the factor the fast kernel runs. The shapes and their sizes are
|
||||
/// worked in `int_kernels`' `a_padded_input_beyond_pad_factor_is_left_to_the_reference`.
|
||||
#[test]
|
||||
fn padding_beyond_the_pad_factor_runs_the_reference() {
|
||||
use crate::int_kernels::testkit::fast_conv_runs;
|
||||
let (mult, bias, shift) = ([3i64], [7i64], 1u32);
|
||||
// Regular, 1×1 kernel, pad 5: (channels, input side, stride, expected
|
||||
// fast-kernel runs, output). Only output (1, 1) can read the input.
|
||||
for (cin, h, stride, runs, out) in [
|
||||
(1usize, 2usize, 6usize, 1u64, [7i64, 7, 7, ((10 * 4 * 3) >> 1) + 7]), // at the factor
|
||||
(1, 2, 7, 0, [7; 4]), // beyond it
|
||||
(3, 3, 7, 0, [7, 7, 7, (((10 * 9 + 20 * 18 + 30 * 27) * 3) >> 1) + 7]), // beyond it
|
||||
] {
|
||||
let x: Vec<i64> = (1..=(cin * h * h) as i64).collect();
|
||||
let wt: Vec<i64> = (1..=cin as i64).map(|c| 10 * c).collect();
|
||||
let want = reference::conv(&x, cin, h, h, &wt, 1, 1, stride, 5, &mult, &bias, shift);
|
||||
assert_eq!(want, (out.to_vec(), 2, 2), "reference, {cin}×{h}² stride {stride}");
|
||||
for p in FAST {
|
||||
let got = fast_conv_runs(|| with_policy(p, || conv(&x, cin, h, h, &wt, 1, 1, stride, 5, &mult, &bias, shift)));
|
||||
assert_eq!(got, (want.clone(), runs), "{cin}×{h}² stride {stride} {p:?}");
|
||||
}
|
||||
}
|
||||
// Transposed, 1×1 input and kernel, stride 1: (channels, pad, output pad,
|
||||
// expected fast-kernel runs, output side). No tap reaches the kept output.
|
||||
for (cin, pad, opad, runs, side) in [(2usize, 6usize, 13usize, 1u64, 2usize), (2, 7, 15, 0, 2), (1, 6, 12, 0, 1)] {
|
||||
let (x, wt) = (vec![5i64; cin], vec![3i64; cin]);
|
||||
let want = reference::conv_t(&x, cin, 1, 1, &wt, 1, 1, 1, pad, opad, &mult, &bias, shift);
|
||||
assert_eq!(want, (vec![7; side * side], side, side), "reference, pad {pad} opad {opad}");
|
||||
for p in FAST {
|
||||
let got = fast_conv_runs(|| with_policy(p, || conv_t(&x, cin, 1, 1, &wt, 1, 1, 1, pad, opad, &mult, &bias, shift)));
|
||||
assert_eq!(got, (want.clone(), runs), "pad {pad} opad {opad} {p:?}");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
1775
wai/wai-rs/src/int_kernels.rs
Normal file
1775
wai/wai-rs/src/int_kernels.rs
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -89,9 +89,11 @@ struct SynthSpec {
|
|||
layers: Vec<LayerSpec>,
|
||||
}
|
||||
|
||||
/// One integer transposed convolution + bias + requantize + optional ReLU.
|
||||
/// `input` is `[in_ch, hin, win]` row-major `i64`. Returns `([out_ch, hout,
|
||||
/// wout] i64, hout, wout)`. Pure integer arithmetic — the whole point.
|
||||
/// One integer transposed convolution + bias + requantize + optional ReLU — see
|
||||
/// [`reference::conv_transpose2d`] for the definition. Runs the fast kernel
|
||||
/// ([`crate::int_kernels`]) unless the current [`crate::int_kernels::Policy`] selects
|
||||
/// the reference; both give the same integer for every output element, and every
|
||||
/// error case is left to the reference so errors are reported exactly as before.
|
||||
fn conv_transpose2d(
|
||||
input: &[i64],
|
||||
cin: usize,
|
||||
|
|
@ -99,74 +101,132 @@ fn conv_transpose2d(
|
|||
win: usize,
|
||||
l: &LayerSpec,
|
||||
) -> Result<(Vec<i64>, usize, usize), IntSynthError> {
|
||||
use crate::int_kernels::{conv_acc_i64, policy, ConvGeom, KernelPath};
|
||||
let (k, s, p, op) = (l.k, l.stride, l.pad, l.out_pad);
|
||||
let hout_i = (hin as i64 - 1) * s as i64 + k as i64 + op as i64 - 2 * p as i64;
|
||||
let wout_i = (win as i64 - 1) * s as i64 + k as i64 + op as i64 - 2 * p as i64;
|
||||
if hout_i <= 0 || wout_i <= 0 {
|
||||
return Err(IntSynthError::InvalidPayload(format!(
|
||||
"int_synth layer produces non-positive output {hout_i}x{wout_i}")));
|
||||
}
|
||||
let (hout, wout) = (hout_i as usize, wout_i as usize);
|
||||
if l.weight.len() != cin * l.out_ch * k * k {
|
||||
return Err(IntSynthError::InvalidPayload(format!(
|
||||
"int_synth weight len {} != {}*{}*{}*{}", l.weight.len(), cin, l.out_ch, k, k)));
|
||||
}
|
||||
if l.bias.len() != l.out_ch {
|
||||
return Err(IntSynthError::InvalidPayload(format!(
|
||||
"int_synth bias len {} != out_ch {}", l.bias.len(), l.out_ch)));
|
||||
if policy().path == KernelPath::Fast
|
||||
&& Some(l.weight.len()) == cin.checked_mul(l.out_ch).and_then(|v| v.checked_mul(k)).and_then(|v| v.checked_mul(k))
|
||||
&& l.bias.len() == l.out_ch
|
||||
&& let Some(g) = ConvGeom::transposed(cin, hin, win, l.out_ch, k, k, s, s, p, p, op, op, input.len(), l.weight.len())
|
||||
&& reference_output_side(hin, l) == Some(g.ho)
|
||||
&& reference_output_side(win, l) == Some(g.wo)
|
||||
{
|
||||
let wt: Vec<i64> = l.weight.iter().map(|&v| v as i64).collect();
|
||||
let out = conv_acc_i64(&g, input, &wt, |oc, v| {
|
||||
let mut x = (v + l.bias[oc] as i64) >> l.shift;
|
||||
if l.relu && x < 0 {
|
||||
x = 0;
|
||||
}
|
||||
x
|
||||
});
|
||||
// `None`: operands outside the kernel's no-overflow bound — the reference runs.
|
||||
if let Some(out) = out {
|
||||
return Ok((out, g.ho, g.wo));
|
||||
}
|
||||
}
|
||||
reference::conv_transpose2d(input, cin, hin, win, l)
|
||||
}
|
||||
|
||||
let mut out = vec![0i64; l.out_ch * hout * wout];
|
||||
for ic in 0..cin {
|
||||
for iy in 0..hin {
|
||||
for ix in 0..win {
|
||||
let v = input[(ic * hin + iy) * win + ix];
|
||||
if v == 0 {
|
||||
continue;
|
||||
}
|
||||
for oc in 0..l.out_ch {
|
||||
let wbase = (ic * l.out_ch + oc) * k * k;
|
||||
let obase = oc * hout * wout;
|
||||
for ky in 0..k {
|
||||
let ty = iy * s + ky;
|
||||
if ty < p {
|
||||
continue;
|
||||
}
|
||||
let oy = ty - p;
|
||||
if oy >= hout {
|
||||
continue;
|
||||
}
|
||||
for kx in 0..k {
|
||||
let tx = ix * s + kx;
|
||||
if tx < p {
|
||||
/// One output side exactly as [`reference::conv_transpose2d`] computes it — in
|
||||
/// `i64`, from the same `as` casts, in the same order — or `None` where that
|
||||
/// expression overflows `i64` (the reference then panics with overflow checks on
|
||||
/// and wraps with them off) or is negative. The fast path runs only where this
|
||||
/// equals the side it computed in checked `usize` arithmetic, so it never returns
|
||||
/// where the reference's size arithmetic fails. `ConvGeom::transposed`'s bound on
|
||||
/// the uncropped plane already keeps every term far below `i64::MAX`; this check
|
||||
/// states the rule where the reference's arithmetic is, independent of that cap.
|
||||
fn reference_output_side(n: usize, l: &LayerSpec) -> Option<usize> {
|
||||
let v = (n as i64)
|
||||
.checked_sub(1)?
|
||||
.checked_mul(l.stride as i64)?
|
||||
.checked_add(l.k as i64)?
|
||||
.checked_add(l.out_pad as i64)?
|
||||
.checked_sub(2i64.checked_mul(l.pad as i64)?)?;
|
||||
usize::try_from(v).ok()
|
||||
}
|
||||
|
||||
/// The scalar reference kernel the fast path is checked against, kept verbatim.
|
||||
mod reference {
|
||||
use super::{IntSynthError, LayerSpec};
|
||||
|
||||
/// One integer transposed convolution + bias + requantize + optional ReLU.
|
||||
/// `input` is `[in_ch, hin, win]` row-major `i64`. Returns `([out_ch, hout,
|
||||
/// wout] i64, hout, wout)`. Pure integer arithmetic — the whole point.
|
||||
pub(super) fn conv_transpose2d(
|
||||
input: &[i64],
|
||||
cin: usize,
|
||||
hin: usize,
|
||||
win: usize,
|
||||
l: &LayerSpec,
|
||||
) -> Result<(Vec<i64>, usize, usize), IntSynthError> {
|
||||
let (k, s, p, op) = (l.k, l.stride, l.pad, l.out_pad);
|
||||
let hout_i = (hin as i64 - 1) * s as i64 + k as i64 + op as i64 - 2 * p as i64;
|
||||
let wout_i = (win as i64 - 1) * s as i64 + k as i64 + op as i64 - 2 * p as i64;
|
||||
if hout_i <= 0 || wout_i <= 0 {
|
||||
return Err(IntSynthError::InvalidPayload(format!(
|
||||
"int_synth layer produces non-positive output {hout_i}x{wout_i}")));
|
||||
}
|
||||
let (hout, wout) = (hout_i as usize, wout_i as usize);
|
||||
if l.weight.len() != cin * l.out_ch * k * k {
|
||||
return Err(IntSynthError::InvalidPayload(format!(
|
||||
"int_synth weight len {} != {}*{}*{}*{}", l.weight.len(), cin, l.out_ch, k, k)));
|
||||
}
|
||||
if l.bias.len() != l.out_ch {
|
||||
return Err(IntSynthError::InvalidPayload(format!(
|
||||
"int_synth bias len {} != out_ch {}", l.bias.len(), l.out_ch)));
|
||||
}
|
||||
|
||||
let mut out = vec![0i64; l.out_ch * hout * wout];
|
||||
for ic in 0..cin {
|
||||
for iy in 0..hin {
|
||||
for ix in 0..win {
|
||||
let v = input[(ic * hin + iy) * win + ix];
|
||||
if v == 0 {
|
||||
continue;
|
||||
}
|
||||
for oc in 0..l.out_ch {
|
||||
let wbase = (ic * l.out_ch + oc) * k * k;
|
||||
let obase = oc * hout * wout;
|
||||
for ky in 0..k {
|
||||
let ty = iy * s + ky;
|
||||
if ty < p {
|
||||
continue;
|
||||
}
|
||||
let ox = tx - p;
|
||||
if ox >= wout {
|
||||
let oy = ty - p;
|
||||
if oy >= hout {
|
||||
continue;
|
||||
}
|
||||
out[obase + oy * wout + ox] +=
|
||||
l.weight[wbase + ky * k + kx] as i64 * v;
|
||||
for kx in 0..k {
|
||||
let tx = ix * s + kx;
|
||||
if tx < p {
|
||||
continue;
|
||||
}
|
||||
let ox = tx - p;
|
||||
if ox >= wout {
|
||||
continue;
|
||||
}
|
||||
out[obase + oy * wout + ox] +=
|
||||
l.weight[wbase + ky * k + kx] as i64 * v;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Bias, requantize (arithmetic shift — deterministic), integer ReLU.
|
||||
for oc in 0..l.out_ch {
|
||||
let b = l.bias[oc] as i64;
|
||||
let base = oc * hout * wout;
|
||||
for v in &mut out[base..base + hout * wout] {
|
||||
let mut x = (*v + b) >> l.shift;
|
||||
if l.relu && x < 0 {
|
||||
x = 0;
|
||||
// Bias, requantize (arithmetic shift — deterministic), integer ReLU.
|
||||
for oc in 0..l.out_ch {
|
||||
let b = l.bias[oc] as i64;
|
||||
let base = oc * hout * wout;
|
||||
for v in &mut out[base..base + hout * wout] {
|
||||
let mut x = (*v + b) >> l.shift;
|
||||
if l.relu && x < 0 {
|
||||
x = 0;
|
||||
}
|
||||
*v = x;
|
||||
}
|
||||
*v = x;
|
||||
}
|
||||
Ok((out, hout, wout))
|
||||
}
|
||||
Ok((out, hout, wout))
|
||||
}
|
||||
|
||||
/// Decode an int_synth payload against the `int_synth.json` sibling of
|
||||
|
|
@ -321,3 +381,139 @@ mod tests {
|
|||
let _ = std::fs::remove_dir_all(&dir);
|
||||
}
|
||||
}
|
||||
|
||||
/// The fast transposed conv against the scalar reference, element by element —
|
||||
/// results AND errors.
|
||||
#[cfg(test)]
|
||||
mod kernel_equivalence {
|
||||
use super::*;
|
||||
use crate::int_kernels::testkit::{Rng, FAST};
|
||||
use crate::int_kernels::with_policy;
|
||||
|
||||
#[test]
|
||||
fn conv_transpose2d_equals_reference() {
|
||||
let mut g = Rng(0x5E7_0001);
|
||||
for case in 0..500 {
|
||||
let k = 1 + g.below(5);
|
||||
let stride = 1 + g.below(3);
|
||||
let pad = g.below(k + 1);
|
||||
let out_pad = if stride > 1 { g.below(stride) } else { 0 };
|
||||
let (h, w) = (1 + g.below(6), 1 + g.below(8));
|
||||
let (cin, out_ch) = (1 + g.below(5), 1 + g.below(7));
|
||||
let mut weight: Vec<i32> = g.vec(cin * out_ch * k * k, 1 << 14, 10).into_iter().map(|v| v as i32).collect();
|
||||
let mut bias: Vec<i32> = g.vec(out_ch, 1 << 18, 0).into_iter().map(|v| v as i32).collect();
|
||||
// exercise the error paths too: wrong weight / bias length
|
||||
match case % 25 {
|
||||
0 => { weight.pop(); }
|
||||
1 => { bias.push(0); }
|
||||
_ => {}
|
||||
}
|
||||
let l = LayerSpec {
|
||||
in_ch: cin, out_ch, k, stride, pad, out_pad, weight, bias,
|
||||
shift: g.below(12) as u32, relu: g.below(2) == 1,
|
||||
};
|
||||
let x = g.vec(cin * h * w, 127, 40); // int8 latents, sparse
|
||||
let want = reference::conv_transpose2d(&x, cin, h, w, &l);
|
||||
for p in FAST {
|
||||
let got = with_policy(p, || conv_transpose2d(&x, cin, h, w, &l));
|
||||
assert_eq!(got, want, "int_synth {p:?} case {case}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Operands beyond the fast kernel's no-overflow bound, laid out so the only
|
||||
/// overflowing sums would be at border positions (a 1×1 input, 3×3 kernel,
|
||||
/// pad 1, corner taps `2^22` on inputs `2^40`). This reference skips positions
|
||||
/// outside the output rather than accumulating and cropping them, so it forms
|
||||
/// no such sum and returns the kept centre `2·2^40 = 2^41` in every build. The
|
||||
/// fast path hands these operands to the reference and must agree.
|
||||
#[test]
|
||||
fn operands_beyond_the_bound_follow_the_reference() {
|
||||
let (cin, out_ch, k) = (2usize, 1usize, 3usize);
|
||||
let mut weight = vec![0i32; cin * out_ch * k * k];
|
||||
for tap in weight.chunks_mut(k * k) {
|
||||
for c in [0, 2, 6, 8] {
|
||||
tap[c] = 1 << 22;
|
||||
}
|
||||
tap[4] = 1;
|
||||
}
|
||||
let l = LayerSpec { in_ch: cin, out_ch, k, stride: 1, pad: 1, out_pad: 0, weight, bias: vec![0], shift: 0, relu: false };
|
||||
let x = vec![1i64 << 40; cin];
|
||||
let want = reference::conv_transpose2d(&x, cin, 1, 1, &l);
|
||||
assert_eq!(want, Ok((vec![1i64 << 41], 1, 1)), "the reference forms no border sum");
|
||||
for p in FAST {
|
||||
assert_eq!(with_policy(p, || conv_transpose2d(&x, cin, 1, 1, &l)), want, "{p:?}");
|
||||
}
|
||||
}
|
||||
|
||||
/// Geometry whose padded input would be more than `PAD_FACTOR` times what the
|
||||
/// reference reads and writes is refused by the fast path, and the reference
|
||||
/// runs: the fast kernel is not entered and the result is the reference's. At
|
||||
/// exactly the factor the fast kernel runs. The shapes and their sizes are worked
|
||||
/// in `int_kernels`' `a_padded_input_beyond_pad_factor_is_left_to_the_reference`.
|
||||
#[test]
|
||||
fn padding_beyond_the_pad_factor_runs_the_reference() {
|
||||
use crate::int_kernels::testkit::fast_conv_runs;
|
||||
// 1×1 input and kernel, stride 1: (channels, pad, output pad, expected
|
||||
// fast-kernel runs, output side). No tap reaches the kept output, so each
|
||||
// output is the bias, shifted.
|
||||
for (cin, pad, out_pad, runs, side) in [(2usize, 6usize, 13usize, 1u64, 2usize), (2, 7, 15, 0, 2), (1, 6, 12, 0, 1)] {
|
||||
let l = LayerSpec {
|
||||
in_ch: cin, out_ch: 1, k: 1, stride: 1, pad, out_pad, weight: vec![3; cin], bias: vec![14], shift: 1, relu: false,
|
||||
};
|
||||
let x = vec![5i64; cin];
|
||||
let want = reference::conv_transpose2d(&x, cin, 1, 1, &l);
|
||||
assert_eq!(want, Ok((vec![7; side * side], side, side)), "reference, pad {pad} out_pad {out_pad}");
|
||||
for p in FAST {
|
||||
let got = fast_conv_runs(|| with_policy(p, || conv_transpose2d(&x, cin, 1, 1, &l)));
|
||||
assert_eq!(got, (want.clone(), runs), "pad {pad} out_pad {out_pad} {p:?}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The reference sizes its output in `i64`; the fast path runs only where that
|
||||
/// arithmetic neither overflows nor disagrees with its own `usize` geometry.
|
||||
/// `reference_output_side` evaluates the reference's expression step by step,
|
||||
/// and here refuses every input on which the reference's arithmetic overflows.
|
||||
/// Through the dispatch, the fast kernel is not entered for them, and every
|
||||
/// policy has the reference's outcome: a panic with overflow checks on, and the
|
||||
/// wrapped (here still correct) size with them off.
|
||||
#[test]
|
||||
fn the_reference_output_size_is_checked_in_its_own_arithmetic() {
|
||||
use crate::int_kernels::testkit::{fast_conv_runs, overflow_checks};
|
||||
use crate::int_kernels::Policy;
|
||||
let spec = |stride, pad, out_pad| LayerSpec {
|
||||
in_ch: 1, out_ch: 1, k: 3, stride, pad, out_pad, weight: (1..=9).collect(), bias: vec![0], shift: 0, relu: false,
|
||||
};
|
||||
// Ordinary: (4 − 1)·2 + 3 + 1 − 2·1 = 8.
|
||||
assert_eq!(reference_output_side(4, &spec(2, 1, 1)), Some(8));
|
||||
// Not positive: 0·1 + 3 + 0 − 2·2 = −1.
|
||||
assert_eq!(reference_output_side(1, &spec(1, 2, 0)), None);
|
||||
#[cfg(target_pointer_width = "64")]
|
||||
{
|
||||
let s = 1usize << 62;
|
||||
// (3 − 1)·2^62 = 2^63: the product overflows i64.
|
||||
assert_eq!(reference_output_side(3, &spec(s, s, 0)), None);
|
||||
// 2^62 + 3 + 2^62: the second addition overflows i64, although the
|
||||
// wrapped result would be the true size, 3.
|
||||
assert_eq!(reference_output_side(2, &spec(s, s, s)), None);
|
||||
// 2·2^62 overflows i64 (0 + 3 + 0 − 2^63 would be negative anyway).
|
||||
assert_eq!(reference_output_side(1, &spec(1, s, 0)), None);
|
||||
|
||||
// Through the dispatch: a 2×2 input, stride = pad = output pad = 2^62,
|
||||
// a 3×3 output that only input (1, 1) reaches, through every tap.
|
||||
let l = spec(s, s, s);
|
||||
let x = [1i64, 2, 3, 4];
|
||||
let run = |p: Policy| fast_conv_runs(|| std::panic::catch_unwind(|| with_policy(p, || conv_transpose2d(&x, 1, 2, 2, &l))).ok());
|
||||
let want = run(Policy::reference());
|
||||
if overflow_checks() {
|
||||
assert_eq!(want, (None, 0), "the reference overflows i64 sizing its output");
|
||||
} else {
|
||||
assert_eq!(want, (Some(Ok(((1..=9).map(|v| 4 * v).collect(), 3, 3))), 0));
|
||||
}
|
||||
for p in FAST {
|
||||
assert_eq!(run(p), want, "{p:?}");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -420,28 +420,53 @@ pub fn multi_head_attention(
|
|||
/// The sqrt is exact integer `isqrt` — no float, no Newton — so the op is
|
||||
/// byte-identical on every machine (the missing piece to make a real CNN learned
|
||||
/// codec's synthesis `NeuralIntegerExact`).
|
||||
///
|
||||
/// Runs the fast kernel ([`crate::int_kernels`]) unless the current
|
||||
/// [`crate::int_kernels::Policy`] selects [`reference::igdn`]; both give the same
|
||||
/// integer for every output element.
|
||||
pub fn igdn(x: &[i64], c: usize, hw: usize, frac: u32, gamma: &[i64], beta: &[i64], gf: u32) -> Vec<i64> {
|
||||
let mut out = vec![0i64; c * hw];
|
||||
for p in 0..hw {
|
||||
// x_j² at 2^(2·frac)
|
||||
let mut x2 = vec![0i128; c];
|
||||
for j in 0..c {
|
||||
let xj = x[j * hw + p] as i128;
|
||||
x2[j] = xj * xj;
|
||||
}
|
||||
for i in 0..c {
|
||||
// norm_i = beta_i + Σ_j gamma_ij·x_j² at 2^(2·frac)
|
||||
let mut norm = beta[i] as i128;
|
||||
let grow = i * c;
|
||||
if crate::int_kernels::policy().path == crate::int_kernels::KernelPath::Fast
|
||||
&& let Some(out) = crate::int_kernels::igdn(x, c, hw, frac, gamma, beta, gf)
|
||||
{
|
||||
return out;
|
||||
}
|
||||
reference::igdn(x, c, hw, frac, gamma, beta, gf)
|
||||
}
|
||||
|
||||
/// The scalar reference kernel the fast IGDN is checked against, kept verbatim.
|
||||
#[doc(hidden)]
|
||||
pub mod reference {
|
||||
/// Integer-exact inverse GDN (IGDN) — the divisive-normalization activation of
|
||||
/// the CNN learned-codec synthesis (bmshj2018 / mbt2018 / cheng2020 `g_s`):
|
||||
/// `y_i = x_i · sqrt(beta_i + Σ_j gamma_ij · x_j²)`, per spatial position across
|
||||
/// the `c` channels. `x` is `[c, hw]` channel-major at `2^frac`; `gamma` is
|
||||
/// `c·c` (row-major, `[out_i][in_j]`) at `2^gf`; `beta` is `c` at `2^(2·frac)`.
|
||||
/// The sqrt is exact integer `isqrt` — no float, no Newton — so the op is
|
||||
/// byte-identical on every machine (the missing piece to make a real CNN learned
|
||||
/// codec's synthesis `NeuralIntegerExact`).
|
||||
pub fn igdn(x: &[i64], c: usize, hw: usize, frac: u32, gamma: &[i64], beta: &[i64], gf: u32) -> Vec<i64> {
|
||||
let mut out = vec![0i64; c * hw];
|
||||
for p in 0..hw {
|
||||
// x_j² at 2^(2·frac)
|
||||
let mut x2 = vec![0i128; c];
|
||||
for j in 0..c {
|
||||
norm += (gamma[grow + j] as i128 * x2[j]) >> gf;
|
||||
let xj = x[j * hw + p] as i128;
|
||||
x2[j] = xj * xj;
|
||||
}
|
||||
for i in 0..c {
|
||||
// norm_i = beta_i + Σ_j gamma_ij·x_j² at 2^(2·frac)
|
||||
let mut norm = beta[i] as i128;
|
||||
let grow = i * c;
|
||||
for j in 0..c {
|
||||
norm += (gamma[grow + j] as i128 * x2[j]) >> gf;
|
||||
}
|
||||
// sqrt(norm@2^(2frac)) = isqrt(norm) at 2^frac
|
||||
let s = (norm.max(0) as u128).isqrt() as i64;
|
||||
out[i * hw + p] = ((x[i * hw + p] as i128 * s as i128) >> frac) as i64;
|
||||
}
|
||||
// sqrt(norm@2^(2frac)) = isqrt(norm) at 2^frac
|
||||
let s = (norm.max(0) as u128).isqrt() as i64;
|
||||
out[i * hw + p] = ((x[i * hw + p] as i128 * s as i128) >> frac) as i64;
|
||||
}
|
||||
out
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Per-token layer norm of a `seq×d` row-major matrix: [`layernorm`] applied
|
||||
|
|
@ -1065,3 +1090,64 @@ mod tests {
|
|||
assert_eq!(out, enc.forward(&xi, seq), "encoder must be deterministic");
|
||||
}
|
||||
}
|
||||
|
||||
/// The fast IGDN against the scalar reference, element by element, in each of its
|
||||
/// arithmetic modes: 32×32→64-bit terms, `i64` terms, the `i128` fallback, and the
|
||||
/// `i64`-sum / `i128`-output-product combination.
|
||||
#[cfg(test)]
|
||||
mod igdn_equivalence {
|
||||
use super::*;
|
||||
use crate::int_kernels::testkit::{Rng, FAST};
|
||||
use crate::int_kernels::with_policy;
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn check(g: &mut Rng, cases: usize, xr: i64, gr: i64, br: i64, gf_lo: u32, signed_gamma: bool, label: &str) {
|
||||
for _ in 0..cases {
|
||||
let c = 1 + g.below(12);
|
||||
let hw = 1 + g.below(150);
|
||||
let frac = 8 + g.below(8) as u32;
|
||||
let gf = gf_lo + g.below(8) as u32;
|
||||
let x = g.vec(c * hw, xr, 10);
|
||||
let gamma: Vec<i64> =
|
||||
g.vec(c * c, gr, 20).into_iter().map(|v| if signed_gamma { v } else { v.abs() }).collect();
|
||||
let beta: Vec<i64> = g.vec(c, br, 0).into_iter().map(i64::abs).collect();
|
||||
let want = reference::igdn(&x, c, hw, frac, &gamma, &beta, gf);
|
||||
for p in FAST {
|
||||
let got = with_policy(p, || igdn(&x, c, hw, frac, &gamma, &beta, gf));
|
||||
assert_eq!(got, want, "igdn[{label}] {p:?} c={c} hw={hw} frac={frac} gf={gf}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn igdn_equals_reference_in_every_mode() {
|
||||
let mut g = Rng(0x16D_0001);
|
||||
// Q12 activations, the committed models' magnitudes (32×32→64 terms)
|
||||
check(&mut g, 120, 1 << 15, 1 << 18, 1 << 26, 16, false, "narrow");
|
||||
// x² beyond i32, terms still in i64; signed gamma (floor shifts of negatives)
|
||||
check(&mut g, 120, 1 << 20, 1 << 16, 1 << 30, 16, true, "i64");
|
||||
// gamma·x² beyond i64: the i128 path
|
||||
check(&mut g, 120, 1 << 30, 1 << 24, 1 << 40, 16, true, "i128");
|
||||
// every operand fits i32 but the SUM of terms does not fit i64: the sum
|
||||
// bound, not the operand widths, must send this to the i128 path
|
||||
check(&mut g, 120, 46340, i32::MAX as i64, 1 << 62, 0, false, "sum-bound");
|
||||
// |x| ≥ 2^31 with gamma ≤ 1: i64 sums, i128 output product
|
||||
check(&mut g, 60, (1i64 << 31) + (1 << 29), 1, 1 << 20, 30, false, "wide-x");
|
||||
// an all-zero gamma says nothing about x: x² beyond i64 must still not be
|
||||
// formed in i64 (the reference squares in i128)
|
||||
check(&mut g, 40, 1i64 << 40, 0, 1 << 20, 16, false, "zero-gamma-huge-x");
|
||||
}
|
||||
|
||||
/// A position whose norm is negative (negative beta) exercises the `max(0)` clamp.
|
||||
#[test]
|
||||
fn igdn_negative_norm_clamps_like_reference() {
|
||||
let (c, hw) = (3usize, 5usize);
|
||||
let x: Vec<i64> = (0..c * hw).map(|i| (i as i64 - 7) * 311).collect();
|
||||
let gamma = vec![1i64; c * c];
|
||||
let beta = vec![-(1i64 << 40), 5, -3];
|
||||
let want = reference::igdn(&x, c, hw, 12, &gamma, &beta, 16);
|
||||
for p in FAST {
|
||||
assert_eq!(with_policy(p, || igdn(&x, c, hw, 12, &gamma, &beta, 16)), want, "{p:?}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -15,6 +15,8 @@
|
|||
// On the JPEG AI exact decode path: no platform math-library call (see clippy.toml).
|
||||
#![deny(clippy::disallowed_methods)]
|
||||
|
||||
use crate::int_kernels::par_zip;
|
||||
|
||||
/// Fractional bits of the Q16 fixed-point format.
|
||||
pub const F: u32 = 16;
|
||||
/// 1.0 in Q16.
|
||||
|
|
@ -141,116 +143,214 @@ pub fn dot_q16(a: &[i64], b: &[i64]) -> i64 {
|
|||
((acc + (1 << (F - 1))) >> F) as i64
|
||||
}
|
||||
|
||||
/// Grouped / depthwise integer convolution — the one op the dense conv in
|
||||
/// [`crate::int_hyper_synth::run_net`] lacks (5 DWConv modules in the Restormer
|
||||
/// synthesis need it). `groups` divides both `cin` and `cout`; output channel
|
||||
/// `oc` (in group `g = oc/(cout/groups)`) sums only over its group's input
|
||||
/// channels. Depthwise ⇔ `groups == cin == cout` (weight `[cout, 1, k, k]`).
|
||||
/// Per-output-channel requant `((raw·mult) >> shift) + bias`, matching the dense
|
||||
/// conv exactly. Integer accumulation ⇒ byte-identical on every machine.
|
||||
/// Grouped / depthwise integer convolution — see [`reference::conv_grouped`] for
|
||||
/// the definition. Runs the fast kernel ([`crate::int_kernels`]) unless the current
|
||||
/// [`crate::int_kernels::Policy`] selects the reference; both give the same integer
|
||||
/// for every output element.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn conv_grouped(
|
||||
x: &[i64], cin: usize, h: usize, w: usize, weight: &[i64], cout: usize, k: usize,
|
||||
stride: usize, pad: usize, groups: usize, mult: &[i64], bias: &[i64], shift: u32,
|
||||
) -> (Vec<i64>, usize, usize) {
|
||||
let ho = (h + 2 * pad - k) / stride + 1;
|
||||
let wo = (w + 2 * pad - k) / stride + 1;
|
||||
let kk = k * k;
|
||||
let (cin_g, cout_g) = (cin / groups, cout / groups);
|
||||
let mut out = vec![0i64; cout * ho * wo];
|
||||
for oc in 0..cout {
|
||||
let g = oc / cout_g;
|
||||
let plane = oc * ho * wo;
|
||||
for icl in 0..cin_g {
|
||||
let ic = g * cin_g + icl;
|
||||
let xb = ic * h * w;
|
||||
let wb = (oc * cin_g + icl) * kk;
|
||||
for ky in 0..k {
|
||||
for kx in 0..k {
|
||||
let wv = weight[wb + ky * k + kx];
|
||||
if wv == 0 {
|
||||
continue;
|
||||
}
|
||||
for oy in 0..ho {
|
||||
let iy = oy * stride + ky;
|
||||
if iy < pad || iy >= pad + h {
|
||||
continue;
|
||||
}
|
||||
let xrow = xb + (iy - pad) * w;
|
||||
let orow = plane + oy * wo;
|
||||
for ox in 0..wo {
|
||||
let ix = ox * stride + kx;
|
||||
if ix < pad || ix >= pad + w {
|
||||
continue;
|
||||
}
|
||||
out[orow + ox] += wv * x[xrow + ix - pad];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
use crate::int_kernels::{conv_acc_i64, policy, ConvGeom, KernelPath};
|
||||
if policy().path == KernelPath::Fast && mult.len() >= cout && bias.len() >= cout
|
||||
&& let Some(g) = ConvGeom::conv(cin, h, w, cout, groups, k, k, stride, stride, pad, pad, x.len(), weight.len())
|
||||
&& let Some(out) = conv_acc_i64(&g, x, weight, |oc, raw| ((raw as i128 * mult[oc] as i128) >> shift) as i64 + bias[oc])
|
||||
{
|
||||
return (out, g.ho, g.wo);
|
||||
}
|
||||
for oc in 0..cout {
|
||||
let (m, b) = (mult[oc] as i128, bias[oc]);
|
||||
for v in &mut out[oc * ho * wo..(oc + 1) * ho * wo] {
|
||||
*v = ((*v as i128 * m) >> shift) as i64 + b;
|
||||
}
|
||||
}
|
||||
(out, ho, wo)
|
||||
reference::conv_grouped(x, cin, h, w, weight, cout, k, stride, pad, groups, mult, bias, shift)
|
||||
}
|
||||
|
||||
/// Transposed (fractionally-strided) integer convolution — the 5 upsamplers in the
|
||||
/// Restormer synthesis (`first_stage.1`, `conv2_t`, `CAB.upscale`, `TAM.us_conv`,
|
||||
/// `conv4_t`), all group=1. Weight `[cin, cout, k, k]` (ONNX ConvTranspose layout).
|
||||
/// `ho = (h-1)·stride + k + opad - 2·pad`. Per-oc requant `((raw·mult)>>shift)+bias`.
|
||||
/// Same convention as `int_hyper_synth::conv_t`. Byte-identical on every machine.
|
||||
/// Transposed integer convolution — see [`reference::conv_transpose`] for the
|
||||
/// definition. Dispatches like [`conv_grouped`]; the reference accumulates in
|
||||
/// `i128`, which the fast kernel narrows to `i64` only when a bound over the actual
|
||||
/// operands proves no partial sum can leave `i64`.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn conv_transpose(
|
||||
x: &[i64], cin: usize, h: usize, w: usize, weight: &[i64], cout: usize, k: usize,
|
||||
stride: usize, pad: usize, opad: usize, mult: &[i64], bias: &[i64], shift: u32,
|
||||
) -> (Vec<i64>, usize, usize) {
|
||||
let ho = (h - 1) * stride + k + opad - 2 * pad;
|
||||
let wo = (w - 1) * stride + k + opad - 2 * pad;
|
||||
let (hp, wp) = (ho + 2 * pad, wo + 2 * pad);
|
||||
let kk = k * k;
|
||||
let mut raw = vec![0i128; cout * hp * wp];
|
||||
for ic in 0..cin {
|
||||
let xc = &x[ic * h * w..ic * h * w + h * w];
|
||||
for oc in 0..cout {
|
||||
let wk = &weight[(ic * cout + oc) * kk..(ic * cout + oc) * kk + kk];
|
||||
let plane = oc * hp * wp;
|
||||
for iy in 0..h {
|
||||
for (ix, &v) in xc[iy * w..iy * w + w].iter().enumerate() {
|
||||
if v == 0 {
|
||||
continue;
|
||||
use crate::int_kernels::{conv_acc_i128, policy, ConvGeom, KernelPath};
|
||||
if policy().path == KernelPath::Fast && mult.len() >= cout && bias.len() >= cout
|
||||
&& let Some(g) = ConvGeom::transposed(cin, h, w, cout, k, k, stride, stride, pad, pad, opad, opad, x.len(), weight.len())
|
||||
&& let Some(out) = conv_acc_i128(&g, x, weight, |oc, r| ((r * mult[oc] as i128) >> shift) as i64 + bias[oc])
|
||||
{
|
||||
return (out, g.ho, g.wo);
|
||||
}
|
||||
reference::conv_transpose(x, cin, h, w, weight, cout, k, stride, pad, opad, mult, bias, shift)
|
||||
}
|
||||
|
||||
/// The scalar reference kernels the fast paths are checked against, kept verbatim.
|
||||
#[doc(hidden)]
|
||||
pub mod reference {
|
||||
use super::{dot_q16, l2norm, layernorm, softmax, F};
|
||||
|
||||
/// MDTA — Restormer transposed **channel** attention (`TAM.TABs.*.attn`), the
|
||||
/// core learned-attention unit. `q,k,v` are `[heads][c][hw]` Q16; `temp` is the
|
||||
/// learned per-head scale (Q16). Per head: L2-norm(Q,K) over HW → `Qn·Knᵀ·temp`
|
||||
/// (a `c×c` channel-attention matrix) → softmax → `·V`. Byte-exact.
|
||||
pub fn mdta(
|
||||
q: &[Vec<Vec<i64>>],
|
||||
k: &[Vec<Vec<i64>>],
|
||||
v: &[Vec<Vec<i64>>],
|
||||
temp: &[i64],
|
||||
eps_q32: u128,
|
||||
) -> Vec<Vec<Vec<i64>>> {
|
||||
let (heads, c, hw) = (q.len(), q[0].len(), q[0][0].len());
|
||||
let mut out = vec![vec![vec![0i64; hw]; c]; heads];
|
||||
for h in 0..heads {
|
||||
let qn: Vec<Vec<i64>> = (0..c).map(|i| l2norm(&q[h][i], eps_q32)).collect();
|
||||
let kn: Vec<Vec<i64>> = (0..c).map(|i| l2norm(&k[h][i], eps_q32)).collect();
|
||||
for i in 0..c {
|
||||
let scores: Vec<i64> = (0..c)
|
||||
.map(|j| ((dot_q16(&qn[i], &kn[j]) as i128 * temp[h] as i128 + (1 << (F - 1))) >> F) as i64)
|
||||
.collect();
|
||||
let attn = softmax(&scores);
|
||||
for x in 0..hw {
|
||||
let mut acc: i128 = 0;
|
||||
for j in 0..c {
|
||||
acc += attn[j] as i128 * v[h][j][x] as i128;
|
||||
}
|
||||
let base = plane + iy * stride * wp + ix * stride;
|
||||
for ky in 0..k {
|
||||
for kx in 0..k {
|
||||
raw[base + ky * wp + kx] += wk[ky * k + kx] as i128 * v as i128;
|
||||
out[h][i][x] = ((acc + (1 << (F - 1))) >> F) as i64;
|
||||
}
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Channel-wise LayerNorm for a `[c,h,w]` (CHW) Q16 tensor: normalize the `c`-vector
|
||||
/// at each spatial location (Restormer's `norm1`). `gamma`,`beta` are `[c]` Q16.
|
||||
pub fn layernorm_chw(x: &[i64], c: usize, hw: usize, gamma: &[i64], beta: &[i64], eps_q32: u128) -> Vec<i64> {
|
||||
let mut out = vec![0i64; c * hw];
|
||||
let mut col = vec![0i64; c];
|
||||
for p in 0..hw {
|
||||
for ch in 0..c {
|
||||
col[ch] = x[ch * hw + p];
|
||||
}
|
||||
let n = layernorm(&col, gamma, beta, eps_q32);
|
||||
for ch in 0..c {
|
||||
out[ch * hw + p] = n[ch];
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Grouped / depthwise integer convolution — the one op the dense conv in
|
||||
/// [`crate::int_hyper_synth::run_net`] lacks (5 DWConv modules in the Restormer
|
||||
/// synthesis need it). `groups` divides both `cin` and `cout`; output channel
|
||||
/// `oc` (in group `g = oc/(cout/groups)`) sums only over its group's input
|
||||
/// channels. Depthwise ⇔ `groups == cin == cout` (weight `[cout, 1, k, k]`).
|
||||
/// Per-output-channel requant `((raw·mult) >> shift) + bias`, matching the dense
|
||||
/// conv exactly. Integer accumulation ⇒ byte-identical on every machine.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn conv_grouped(
|
||||
x: &[i64], cin: usize, h: usize, w: usize, weight: &[i64], cout: usize, k: usize,
|
||||
stride: usize, pad: usize, groups: usize, mult: &[i64], bias: &[i64], shift: u32,
|
||||
) -> (Vec<i64>, usize, usize) {
|
||||
let ho = (h + 2 * pad - k) / stride + 1;
|
||||
let wo = (w + 2 * pad - k) / stride + 1;
|
||||
let kk = k * k;
|
||||
let (cin_g, cout_g) = (cin / groups, cout / groups);
|
||||
let mut out = vec![0i64; cout * ho * wo];
|
||||
for oc in 0..cout {
|
||||
let g = oc / cout_g;
|
||||
let plane = oc * ho * wo;
|
||||
for icl in 0..cin_g {
|
||||
let ic = g * cin_g + icl;
|
||||
let xb = ic * h * w;
|
||||
let wb = (oc * cin_g + icl) * kk;
|
||||
for ky in 0..k {
|
||||
for kx in 0..k {
|
||||
let wv = weight[wb + ky * k + kx];
|
||||
if wv == 0 {
|
||||
continue;
|
||||
}
|
||||
for oy in 0..ho {
|
||||
let iy = oy * stride + ky;
|
||||
if iy < pad || iy >= pad + h {
|
||||
continue;
|
||||
}
|
||||
let xrow = xb + (iy - pad) * w;
|
||||
let orow = plane + oy * wo;
|
||||
for ox in 0..wo {
|
||||
let ix = ox * stride + kx;
|
||||
if ix < pad || ix >= pad + w {
|
||||
continue;
|
||||
}
|
||||
out[orow + ox] += wv * x[xrow + ix - pad];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
let mut out = vec![0i64; cout * ho * wo];
|
||||
for oc in 0..cout {
|
||||
let (m, b) = (mult[oc] as i128, bias[oc]);
|
||||
for oy in 0..ho {
|
||||
for ox in 0..wo {
|
||||
let r = raw[(oc * hp + oy + pad) * wp + ox + pad];
|
||||
out[(oc * ho + oy) * wo + ox] = ((r * m) >> shift) as i64 + b;
|
||||
for oc in 0..cout {
|
||||
let (m, b) = (mult[oc] as i128, bias[oc]);
|
||||
for v in &mut out[oc * ho * wo..(oc + 1) * ho * wo] {
|
||||
*v = ((*v as i128 * m) >> shift) as i64 + b;
|
||||
}
|
||||
}
|
||||
(out, ho, wo)
|
||||
}
|
||||
|
||||
/// Transposed (fractionally-strided) integer convolution — the 5 upsamplers in the
|
||||
/// Restormer synthesis (`first_stage.1`, `conv2_t`, `CAB.upscale`, `TAM.us_conv`,
|
||||
/// `conv4_t`), all group=1. Weight `[cin, cout, k, k]` (ONNX ConvTranspose layout).
|
||||
/// `ho = (h-1)·stride + k + opad - 2·pad`. Per-oc requant `((raw·mult)>>shift)+bias`.
|
||||
/// Same convention as `int_hyper_synth::conv_t`. Byte-identical on every machine.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn conv_transpose(
|
||||
x: &[i64], cin: usize, h: usize, w: usize, weight: &[i64], cout: usize, k: usize,
|
||||
stride: usize, pad: usize, opad: usize, mult: &[i64], bias: &[i64], shift: u32,
|
||||
) -> (Vec<i64>, usize, usize) {
|
||||
let ho = (h - 1) * stride + k + opad - 2 * pad;
|
||||
let wo = (w - 1) * stride + k + opad - 2 * pad;
|
||||
let (hp, wp) = (ho + 2 * pad, wo + 2 * pad);
|
||||
let kk = k * k;
|
||||
let mut raw = vec![0i128; cout * hp * wp];
|
||||
for ic in 0..cin {
|
||||
let xc = &x[ic * h * w..ic * h * w + h * w];
|
||||
for oc in 0..cout {
|
||||
let wk = &weight[(ic * cout + oc) * kk..(ic * cout + oc) * kk + kk];
|
||||
let plane = oc * hp * wp;
|
||||
for iy in 0..h {
|
||||
for (ix, &v) in xc[iy * w..iy * w + w].iter().enumerate() {
|
||||
if v == 0 {
|
||||
continue;
|
||||
}
|
||||
let base = plane + iy * stride * wp + ix * stride;
|
||||
for ky in 0..k {
|
||||
for kx in 0..k {
|
||||
raw[base + ky * wp + kx] += wk[ky * k + kx] as i128 * v as i128;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
let mut out = vec![0i64; cout * ho * wo];
|
||||
for oc in 0..cout {
|
||||
let (m, b) = (mult[oc] as i128, bias[oc]);
|
||||
for oy in 0..ho {
|
||||
for ox in 0..wo {
|
||||
let r = raw[(oc * hp + oy + pad) * wp + ox + pad];
|
||||
out[(oc * ho + oy) * wo + ox] = ((r * m) >> shift) as i64 + b;
|
||||
}
|
||||
}
|
||||
}
|
||||
(out, ho, wo)
|
||||
}
|
||||
(out, ho, wo)
|
||||
}
|
||||
|
||||
/// MDTA — Restormer transposed **channel** attention (`TAM.TABs.*.attn`), the
|
||||
/// core learned-attention unit. `q,k,v` are `[heads][c][hw]` Q16; `temp` is the
|
||||
/// learned per-head scale (Q16). Per head: L2-norm(Q,K) over HW → `Qn·Knᵀ·temp`
|
||||
/// (a `c×c` channel-attention matrix) → softmax → `·V`. Byte-exact.
|
||||
///
|
||||
/// Runs [`mdta_fast`] (rows in parallel, `i64` accumulation where a bound proves it
|
||||
/// exact) unless the current [`crate::int_kernels::Policy`] selects
|
||||
/// [`reference::mdta`]; both give the same integer for every output element.
|
||||
pub fn mdta(
|
||||
q: &[Vec<Vec<i64>>],
|
||||
k: &[Vec<Vec<i64>>],
|
||||
|
|
@ -258,43 +358,132 @@ pub fn mdta(
|
|||
temp: &[i64],
|
||||
eps_q32: u128,
|
||||
) -> Vec<Vec<Vec<i64>>> {
|
||||
let (heads, c, hw) = (q.len(), q[0].len(), q[0][0].len());
|
||||
let mut out = vec![vec![vec![0i64; hw]; c]; heads];
|
||||
for h in 0..heads {
|
||||
let qn: Vec<Vec<i64>> = (0..c).map(|i| l2norm(&q[h][i], eps_q32)).collect();
|
||||
let kn: Vec<Vec<i64>> = (0..c).map(|i| l2norm(&k[h][i], eps_q32)).collect();
|
||||
for i in 0..c {
|
||||
let scores: Vec<i64> = (0..c)
|
||||
.map(|j| ((dot_q16(&qn[i], &kn[j]) as i128 * temp[h] as i128 + (1 << (F - 1))) >> F) as i64)
|
||||
.collect();
|
||||
let attn = softmax(&scores);
|
||||
for x in 0..hw {
|
||||
let mut acc: i128 = 0;
|
||||
for j in 0..c {
|
||||
acc += attn[j] as i128 * v[h][j][x] as i128;
|
||||
}
|
||||
out[h][i][x] = ((acc + (1 << (F - 1))) >> F) as i64;
|
||||
}
|
||||
}
|
||||
use crate::int_kernels::{policy, KernelPath};
|
||||
if policy().path == KernelPath::Fast
|
||||
&& let Some(out) = mdta_fast(q, k, v, temp, eps_q32)
|
||||
{
|
||||
return out;
|
||||
}
|
||||
out
|
||||
reference::mdta(q, k, v, temp, eps_q32)
|
||||
}
|
||||
|
||||
/// Channel-wise LayerNorm for a `[c,h,w]` (CHW) Q16 tensor: normalize the `c`-vector
|
||||
/// at each spatial location (Restormer's `norm1`). `gamma`,`beta` are `[c]` Q16.
|
||||
pub fn layernorm_chw(x: &[i64], c: usize, hw: usize, gamma: &[i64], beta: &[i64], eps_q32: u128) -> Vec<i64> {
|
||||
let mut out = vec![0i64; c * hw];
|
||||
let mut col = vec![0i64; c];
|
||||
for p in 0..hw {
|
||||
for ch in 0..c {
|
||||
col[ch] = x[ch * hw + p];
|
||||
}
|
||||
let n = layernorm(&col, gamma, beta, eps_q32);
|
||||
for ch in 0..c {
|
||||
out[ch * hw + p] = n[ch];
|
||||
}
|
||||
/// [`reference::mdta`] with every `(head, row)` computed independently in parallel —
|
||||
/// each output row depends only on its own head's inputs, so the partition cannot
|
||||
/// change it. The two `i128` accumulations are narrowed to `i64` only under a bound
|
||||
/// over the actual operands (`terms · max|a| · max|b| ≤ i64::MAX`), under which no
|
||||
/// partial sum can overflow and the `i64` sum is the same integer. `None` for
|
||||
/// malformed input (ragged rows, short `temp`), which the reference then handles.
|
||||
fn mdta_fast(
|
||||
q: &[Vec<Vec<i64>>],
|
||||
k: &[Vec<Vec<i64>>],
|
||||
v: &[Vec<Vec<i64>>],
|
||||
temp: &[i64],
|
||||
eps_q32: u128,
|
||||
) -> Option<Vec<Vec<Vec<i64>>>> {
|
||||
use crate::int_kernels::{max_abs, par_for_each, threads_for};
|
||||
let heads = q.len();
|
||||
if heads == 0 || k.len() != heads || v.len() != heads || temp.len() < heads || q[0].is_empty() {
|
||||
return None;
|
||||
}
|
||||
out
|
||||
let (c, hw) = (q[0].len(), q[0][0].len());
|
||||
let well = (0..heads).all(|h| {
|
||||
q[h].len() == c && k[h].len() == c && v[h].len() == c && q[h].iter().chain(&k[h]).chain(&v[h]).all(|r| r.len() == hw)
|
||||
});
|
||||
if !well {
|
||||
return None;
|
||||
}
|
||||
// Widened (and saturating) before multiplying: on a 32-bit target the product
|
||||
// can exceed `usize`, and an overflow-checked build would then panic here, in
|
||||
// arithmetic the reference never performs. The value only sizes the worker count.
|
||||
let threads = threads_for(
|
||||
(heads as u64).saturating_mul(c as u64).saturating_mul(c as u64).saturating_mul(hw as u64).saturating_mul(2),
|
||||
);
|
||||
const HALF: i128 = 1 << (F - 1);
|
||||
|
||||
// 1. L2-normalised Q and K rows.
|
||||
let mut qn = vec![vec![Vec::new(); c]; heads];
|
||||
let mut kn = vec![vec![Vec::new(); c]; heads];
|
||||
let mut items: Vec<(&Vec<i64>, &mut Vec<i64>)> = Vec::with_capacity(2 * heads * c);
|
||||
for (src, dst) in q.iter().zip(qn.iter_mut()).chain(k.iter().zip(kn.iter_mut())) {
|
||||
items.extend(src.iter().zip(dst.iter_mut()));
|
||||
}
|
||||
par_for_each(items, threads, |(src, dst)| *dst = l2norm(src, eps_q32));
|
||||
|
||||
// Per-head bounds: dot products (hw terms) and attn·V (c terms, |attn| ≤ its max).
|
||||
let rows_max = |rows: &[Vec<i64>]| rows.iter().map(|r| max_abs(r)).max().unwrap_or(0) as u128;
|
||||
let dot_i64: Vec<bool> = (0..heads)
|
||||
.map(|h| (hw as u128).saturating_mul(rows_max(&qn[h])).saturating_mul(rows_max(&kn[h])) <= i64::MAX as u128)
|
||||
.collect();
|
||||
let v_max: Vec<u128> = (0..heads).map(|h| rows_max(&v[h])).collect();
|
||||
|
||||
// 2. Every (head, row): scores → softmax → attn·V.
|
||||
let mut out = vec![vec![vec![0i64; hw]; c]; heads];
|
||||
let rows: Vec<(usize, usize, &mut Vec<i64>)> =
|
||||
out.iter_mut().enumerate().flat_map(|(h, r)| r.iter_mut().enumerate().map(move |(i, row)| (h, i, row))).collect();
|
||||
par_for_each(rows, threads, |(h, i, row)| {
|
||||
let dot = |a: &[i64], b: &[i64]| -> i64 {
|
||||
if dot_i64[h] {
|
||||
let mut acc = 0i64;
|
||||
for (&x, &y) in a.iter().zip(b) {
|
||||
acc += x * y;
|
||||
}
|
||||
((acc as i128 + HALF) >> F) as i64
|
||||
} else {
|
||||
dot_q16(a, b)
|
||||
}
|
||||
};
|
||||
let scores: Vec<i64> =
|
||||
(0..c).map(|j| ((dot(&qn[h][i], &kn[h][j]) as i128 * temp[h] as i128 + HALF) >> F) as i64).collect();
|
||||
let attn = softmax(&scores);
|
||||
if (c as u128).saturating_mul(max_abs(&attn) as u128).saturating_mul(v_max[h]) <= i64::MAX as u128 {
|
||||
let mut acc = vec![0i64; hw];
|
||||
for (j, &a) in attn.iter().enumerate() {
|
||||
if a != 0 {
|
||||
for (o, &vv) in acc.iter_mut().zip(&v[h][j]) {
|
||||
*o += a * vv;
|
||||
}
|
||||
}
|
||||
}
|
||||
for (d, &s) in row.iter_mut().zip(&acc) {
|
||||
*d = ((s as i128 + HALF) >> F) as i64;
|
||||
}
|
||||
} else {
|
||||
let mut acc = vec![0i128; hw];
|
||||
for (j, &a) in attn.iter().enumerate() {
|
||||
for (o, &vv) in acc.iter_mut().zip(&v[h][j]) {
|
||||
*o += a as i128 * vv as i128;
|
||||
}
|
||||
}
|
||||
for (d, &s) in row.iter_mut().zip(&acc) {
|
||||
*d = ((s + HALF) >> F) as i64;
|
||||
}
|
||||
}
|
||||
});
|
||||
Some(out)
|
||||
}
|
||||
|
||||
/// Channel-wise LayerNorm for a `[c,h,w]` (CHW) Q16 tensor — see
|
||||
/// [`reference::layernorm_chw`]. Under the fast policy the positions are spread over
|
||||
/// workers (each position is normalised by the same [`layernorm`] call either way).
|
||||
pub fn layernorm_chw(x: &[i64], c: usize, hw: usize, gamma: &[i64], beta: &[i64], eps_q32: u128) -> Vec<i64> {
|
||||
use crate::int_kernels::{par_for_each, policy, position_tiles, threads_for, KernelPath};
|
||||
if policy().path == KernelPath::Fast && c > 0 && hw > 0 && x.len() >= c * hw && gamma.len() >= c && beta.len() >= c {
|
||||
let mut out = vec![0i64; c * hw];
|
||||
let threads = threads_for((c * hw) as u64 * 32);
|
||||
par_for_each(position_tiles(&mut out, hw, 64), threads, |(p0, mut dst)| {
|
||||
let mut col = vec![0i64; c];
|
||||
for pp in 0..dst[0].len() {
|
||||
for ch in 0..c {
|
||||
col[ch] = x[ch * hw + p0 + pp];
|
||||
}
|
||||
for (ch, v) in layernorm(&col, gamma, beta, eps_q32).into_iter().enumerate() {
|
||||
dst[ch][pp] = v;
|
||||
}
|
||||
}
|
||||
});
|
||||
return out;
|
||||
}
|
||||
reference::layernorm_chw(x, c, hw, gamma, beta, eps_q32)
|
||||
}
|
||||
|
||||
/// Weights for one GDFN (Restormer gated-Dconv feed-forward). All Q16 / int.
|
||||
|
|
@ -320,16 +509,10 @@ pub fn gdfn(x: &[i64], wt: &GdfnWeights, eps_q32: u128) -> Vec<i64> {
|
|||
let (h2, _, _) = conv_grouped(&y, c, h, w, &wt.pin_w, 2 * hid, 1, 1, 0, 1, &wt.pin_m, &wt.pin_b, wt.shift);
|
||||
let (hd, _, _) = conv_grouped(&h2, 2 * hid, h, w, &wt.dw_w, 2 * hid, 3, 1, 1, 2 * hid, &wt.dw_m, &wt.dw_b, wt.shift);
|
||||
// split channels into a (gate) and b (value), elementwise Elu(a)⊙b
|
||||
let mut g = vec![0i64; hid * hw];
|
||||
for ch in 0..hid {
|
||||
for p in 0..hw {
|
||||
let a = hd[ch * hw + p];
|
||||
let b = hd[(hid + ch) * hw + p];
|
||||
g[ch * hw + p] = ((elu(a, S) as i128 * b as i128) >> F) as i64;
|
||||
}
|
||||
}
|
||||
// (element-wise: each output computed from its own inputs, spread over workers)
|
||||
let g = par_zip(&hd[..hid * hw], &hd[hid * hw..2 * hid * hw], |a, b| ((elu(a, S) as i128 * b as i128) >> F) as i64);
|
||||
let (o, _, _) = conv_grouped(&g, hid, h, w, &wt.pout_w, c, 1, 1, 0, 1, &wt.pout_m, &wt.pout_b, wt.shift);
|
||||
(0..c * hw).map(|i| x[i] + o[i]).collect()
|
||||
par_zip(&x[..c * hw], &o[..c * hw], |a, b| a + b)
|
||||
}
|
||||
|
||||
/// Weights for one MDTA block (Restormer transposed-attention half of a TAB).
|
||||
|
|
@ -367,7 +550,7 @@ pub fn mdta_block(x: &[i64], wt: &MdtaBlockWeights, eps_q32: u128) -> Vec<i64> {
|
|||
}
|
||||
}
|
||||
let (po, _, _) = conv_grouped(&o_chw, c, h, w, &wt.po_w, c, 1, 1, 0, 1, &wt.po_m, &wt.po_b, wt.shift);
|
||||
(0..c * hw).map(|i| x[i] + po[i]).collect()
|
||||
par_zip(&x[..c * hw], &po[..c * hw], |a, b| a + b)
|
||||
}
|
||||
|
||||
/// ONNX **DepthToSpace** (DCR mode) — the synthesis tail that folds channels into
|
||||
|
|
@ -860,3 +1043,245 @@ mod tests {
|
|||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The fast convolutions against the scalar reference, element by element —
|
||||
/// including grouped/depthwise and strided convs, and the transposed conv on both
|
||||
/// sides of the proved `i64`-narrowing bound (the reference accumulates in `i128`).
|
||||
#[cfg(test)]
|
||||
mod kernel_equivalence {
|
||||
use super::*;
|
||||
use crate::int_kernels::testkit::{Rng, FAST};
|
||||
use crate::int_kernels::with_policy;
|
||||
|
||||
#[test]
|
||||
fn grouped_conv_equals_reference() {
|
||||
let mut g = Rng(0x6E0_0001);
|
||||
let mut ran = 0;
|
||||
for case in 0..600 {
|
||||
let k = [1usize, 3, 3, 5][g.below(4)];
|
||||
let stride = 1 + g.below(2);
|
||||
let pad = g.below(k / 2 + 1);
|
||||
let (h, w) = (1 + g.below(10), 1 + g.below(12));
|
||||
if h + 2 * pad < k || w + 2 * pad < k {
|
||||
continue;
|
||||
}
|
||||
let base = 1 + g.below(4);
|
||||
let groups = [1usize, 2, base][g.below(3)];
|
||||
let (cin, cout) = (groups * base, groups * (1 + g.below(3)));
|
||||
// the reference accumulates in i64: keep ≤ 100 terms × |product| inside it
|
||||
let (xr, wr) = [(4 << 16, 1 << 15), (1 << 40, 1 << 12), (1 << 34, 1 << 18)][case % 3];
|
||||
let x = g.vec(cin * h * w, xr, [0, 30, 80][case % 3]);
|
||||
let wt = g.vec(cout * (cin / groups) * k * k, wr, 5);
|
||||
let mult: Vec<i64> = (0..cout).map(|_| 1 + g.signed(1 << 16).abs()).collect();
|
||||
let bias = g.vec(cout, 1 << 16, 0);
|
||||
let shift = 16 + g.below(16) as u32;
|
||||
let want = reference::conv_grouped(&x, cin, h, w, &wt, cout, k, stride, pad, groups, &mult, &bias, shift);
|
||||
for p in FAST {
|
||||
let got = with_policy(p, || conv_grouped(&x, cin, h, w, &wt, cout, k, stride, pad, groups, &mult, &bias, shift));
|
||||
assert_eq!(got, want, "conv_grouped {p:?} k={k} s={stride} p={pad} g={groups} {cin}x{h}x{w}->{cout}");
|
||||
}
|
||||
ran += 1;
|
||||
}
|
||||
assert!(ran > 300, "too few valid cases ({ran})");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn transposed_conv_equals_reference_across_the_narrowing_bound() {
|
||||
let mut g = Rng(0x6E0_0002);
|
||||
let mut ran = 0;
|
||||
for case in 0..600 {
|
||||
let k = [1usize, 2, 3, 4, 5][g.below(5)];
|
||||
let stride = 1 + g.below(3);
|
||||
let pad = g.below(k);
|
||||
let opad = if stride > 1 { g.below(stride) } else { 0 };
|
||||
let (h, w) = (1 + g.below(6), 1 + g.below(8));
|
||||
let (cin, cout) = (1 + g.below(5), 1 + g.below(6));
|
||||
if (h - 1) * stride + k + opad <= 2 * pad || (w - 1) * stride + k + opad <= 2 * pad {
|
||||
continue;
|
||||
}
|
||||
// Q16 activations × 16-bit weights (narrowable, fits i32); beyond i32 but
|
||||
// narrowable; products beyond i64 (the i128 path).
|
||||
let (xr, wr, shift) = [(4 << 16, 1 << 15, 16u32), (1 << 38, 1 << 12, 24), (1 << 40, 1 << 30, 44)][case % 3];
|
||||
let x = g.vec(cin * h * w, xr, [0, 40, 80][case % 3]);
|
||||
let wt = g.vec(cin * cout * k * k, wr, 5);
|
||||
let mult: Vec<i64> = (0..cout).map(|_| 1 + g.signed(1 << 16).abs()).collect();
|
||||
let bias = g.vec(cout, 1 << 16, 0);
|
||||
let want = reference::conv_transpose(&x, cin, h, w, &wt, cout, k, stride, pad, opad, &mult, &bias, shift);
|
||||
for p in FAST {
|
||||
let got = with_policy(p, || conv_transpose(&x, cin, h, w, &wt, cout, k, stride, pad, opad, &mult, &bias, shift));
|
||||
assert_eq!(got, want, "conv_transpose {p:?} k={k} s={stride} p={pad} op={opad} {cin}x{h}x{w}->{cout}");
|
||||
}
|
||||
ran += 1;
|
||||
}
|
||||
assert!(ran > 300, "too few valid cases ({ran})");
|
||||
}
|
||||
|
||||
/// MDTA: rows in parallel and bounded `i64` accumulation vs the `i128` reference,
|
||||
/// including V magnitudes that force the `i128` attn·V path.
|
||||
#[test]
|
||||
fn mdta_equals_reference() {
|
||||
use crate::int_kernels::Policy;
|
||||
let mut g = Rng(0x6E0_0003);
|
||||
for case in 0..150 {
|
||||
let (heads, c, hw) = (1 + g.below(4), 1 + g.below(10), 1 + g.below(70));
|
||||
let vr = [1i64 << 18, 1 << 52][case % 2];
|
||||
let mk = |g: &mut Rng, r: i64| -> Vec<Vec<Vec<i64>>> { (0..heads).map(|_| (0..c).map(|_| g.vec(hw, r, 10)).collect()).collect() };
|
||||
let (q, k, v) = (mk(&mut g, 1 << 18), mk(&mut g, 1 << 18), mk(&mut g, vr));
|
||||
let temp: Vec<i64> = (0..heads).map(|_| (1 << 14) + g.below(3 << 16) as i64).collect();
|
||||
let eps = [0u128, 42_949][case % 2];
|
||||
let want = with_policy(Policy::reference(), || mdta(&q, &k, &v, &temp, eps));
|
||||
assert_eq!(want, reference::mdta(&q, &k, &v, &temp, eps));
|
||||
for p in FAST {
|
||||
assert_eq!(with_policy(p, || mdta(&q, &k, &v, &temp, eps)), want, "mdta {p:?} case {case}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn layernorm_chw_equals_reference() {
|
||||
let mut g = Rng(0x6E0_0004);
|
||||
for _ in 0..150 {
|
||||
let (c, hw) = (1 + g.below(40), 1 + g.below(300));
|
||||
let x = g.vec(c * hw, 4 << 16, 5);
|
||||
let gamma = g.vec(c, 2 << 16, 0);
|
||||
let beta = g.vec(c, 1 << 15, 0);
|
||||
let eps = g.below(1 << 20) as u128;
|
||||
let want = reference::layernorm_chw(&x, c, hw, &gamma, &beta, eps);
|
||||
for p in FAST {
|
||||
assert_eq!(with_policy(p, || layernorm_chw(&x, c, hw, &gamma, &beta, eps)), want, "{p:?} c={c} hw={hw}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A whole Restormer TAB (MDTA block + GDFN: LayerNorm, 1×1 and depthwise convs,
|
||||
/// channel attention, ELU gate, residuals) under every policy.
|
||||
#[test]
|
||||
fn tab_block_equals_reference() {
|
||||
use crate::int_kernels::Policy;
|
||||
let mut g = Rng(0x6E0_0005);
|
||||
for _ in 0..12 {
|
||||
let heads = 1 + g.below(3);
|
||||
let c = heads * (2 + g.below(4));
|
||||
let hid = 2 + g.below(8);
|
||||
let (h, w) = (3 + g.below(9), 3 + g.below(11));
|
||||
let m = |g: &mut Rng, n: usize| (0..n).map(|_| 20_000 + g.below(60_000) as i64).collect::<Vec<_>>();
|
||||
let attn = MdtaBlockWeights {
|
||||
c, heads, h, w,
|
||||
ln_g: g.vec(c, 2 << 16, 0), ln_b: g.vec(c, 1 << 14, 0),
|
||||
qkv_w: g.vec(3 * c * c, 32767, 5), qkv_m: m(&mut g, 3 * c), qkv_b: g.vec(3 * c, 1 << 14, 0),
|
||||
dw_w: g.vec(3 * c * 9, 32767, 5), dw_m: m(&mut g, 3 * c), dw_b: g.vec(3 * c, 1 << 14, 0),
|
||||
temp: (0..heads).map(|_| (1 << 15) + g.below(1 << 17) as i64).collect(),
|
||||
po_w: g.vec(c * c, 32767, 5), po_m: m(&mut g, c), po_b: g.vec(c, 1 << 14, 0),
|
||||
shift: 31,
|
||||
};
|
||||
let ffn = GdfnWeights {
|
||||
c, hidden: hid, h, w,
|
||||
ln_g: g.vec(c, 2 << 16, 0), ln_b: g.vec(c, 1 << 14, 0),
|
||||
pin_w: g.vec(2 * hid * c, 32767, 5), pin_m: m(&mut g, 2 * hid), pin_b: g.vec(2 * hid, 1 << 14, 0),
|
||||
dw_w: g.vec(2 * hid * 9, 32767, 5), dw_m: m(&mut g, 2 * hid), dw_b: g.vec(2 * hid, 1 << 14, 0),
|
||||
pout_w: g.vec(c * hid, 32767, 5), pout_m: m(&mut g, c), pout_b: g.vec(c, 1 << 14, 0),
|
||||
shift: 31,
|
||||
};
|
||||
let x = g.vec(c * h * w, 3 << 16, 0);
|
||||
let eps = 42_949u128;
|
||||
let want = with_policy(Policy::reference(), || tab(&x, &attn, &ffn, eps));
|
||||
for p in FAST {
|
||||
assert_eq!(with_policy(p, || tab(&x, &attn, &ffn, eps)), want, "tab {p:?} c={c} heads={heads} {h}x{w}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Operands placed exactly at, and just past, the bound `terms·max|w|·max|x| ≤
|
||||
/// i64::MAX` that decides `i64` vs `i128` accumulation; every sum is at its
|
||||
/// extreme (all terms the same sign). Both sides must equal the `i128` reference.
|
||||
#[test]
|
||||
fn narrowing_bound_edges_equal_reference() {
|
||||
let (cin, h, w, cout, k) = (2usize, 2usize, 2usize, 1usize, 2usize);
|
||||
let terms = (cin * k * k) as i64; // stride 1: every tap reaches the centre
|
||||
let wv = 1i64 << 20;
|
||||
let at = i64::MAX / terms / wv; // terms·wv·at ≤ i64::MAX
|
||||
for xv in [at, at + 1, -at, -(at + 1)] {
|
||||
let x = vec![xv; cin * h * w];
|
||||
let wt = vec![wv; cin * cout * k * k];
|
||||
let want = reference::conv_transpose(&x, cin, h, w, &wt, cout, k, 1, 0, 0, &[1], &[0], 8);
|
||||
for p in FAST {
|
||||
let got = with_policy(p, || conv_transpose(&x, cin, h, w, &wt, cout, k, 1, 0, 0, &[1], &[0], 8));
|
||||
assert_eq!(got, want, "bound edge x={xv} {p:?}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Overflow of the reference's `i128` accumulator only at a position it
|
||||
/// accumulates into and then crops. A 1×1 input with a 3×3 kernel, stride 1 and
|
||||
/// pad 1 keeps only the centre of the 3×3 raw buffer: `8·2^62·1 = 2^65`, which
|
||||
/// requantises (`>> 4`) to 2^61. Each corner sums eight terms `2^62·2^62 = 2^124`
|
||||
/// and reaches `2^127` on the eighth, one past `i128::MAX`. The fast path must
|
||||
/// panic where the reference does and otherwise give the same 2^61.
|
||||
#[test]
|
||||
fn i128_overflow_only_at_a_cropped_border_behaves_like_reference() {
|
||||
use crate::int_kernels::testkit::overflow_checks;
|
||||
use crate::int_kernels::Policy;
|
||||
let (cin, cout, k) = (8usize, 1usize, 3usize);
|
||||
let x = vec![1i64 << 62; cin];
|
||||
let mut wt = vec![0i64; cin * cout * k * k];
|
||||
for tap in wt.chunks_mut(k * k) {
|
||||
for c in [0, 2, 6, 8] {
|
||||
tap[c] = 1 << 62;
|
||||
}
|
||||
tap[4] = 1;
|
||||
}
|
||||
let run = |p: Policy| {
|
||||
std::panic::catch_unwind(|| with_policy(p, || conv_transpose(&x, cin, 1, 1, &wt, cout, k, 1, 1, 0, &[1], &[0], 4))).ok()
|
||||
};
|
||||
let want = run(Policy::reference());
|
||||
if overflow_checks() {
|
||||
assert!(want.is_none(), "the reference overflows i128 at a cropped corner in a checked build");
|
||||
} else {
|
||||
assert_eq!(want, Some((vec![1i64 << 61], 1, 1)));
|
||||
}
|
||||
for p in FAST {
|
||||
assert_eq!(run(p), want, "fast must match the reference, panic included ({p:?})");
|
||||
}
|
||||
}
|
||||
|
||||
/// Geometry whose padded input would be more than `PAD_FACTOR` times what the
|
||||
/// reference reads and writes is refused by the fast path, and the reference
|
||||
/// runs: the fast kernel is not entered and the integers are the reference's.
|
||||
/// At exactly the factor the fast kernel runs. The shapes and their sizes are
|
||||
/// worked in `int_kernels`' `a_padded_input_beyond_pad_factor_is_left_to_the_reference`.
|
||||
#[test]
|
||||
fn padding_beyond_the_pad_factor_runs_the_reference() {
|
||||
use crate::int_kernels::testkit::fast_conv_runs;
|
||||
let (mult, bias, shift) = ([3i64], [7i64], 1u32);
|
||||
// Regular, 1×1 kernel, pad 5, one group: (channels, input side, stride,
|
||||
// expected fast-kernel runs, output). Only output (1, 1) can read the input.
|
||||
for (cin, h, stride, runs, out) in [
|
||||
(1usize, 2usize, 6usize, 1u64, [7i64, 7, 7, ((10 * 4 * 3) >> 1) + 7]), // at the factor
|
||||
(1, 2, 7, 0, [7; 4]), // beyond it
|
||||
(3, 3, 7, 0, [7, 7, 7, (((10 * 9 + 20 * 18 + 30 * 27) * 3) >> 1) + 7]), // beyond it
|
||||
] {
|
||||
let x: Vec<i64> = (1..=(cin * h * h) as i64).collect();
|
||||
let wt: Vec<i64> = (1..=cin as i64).map(|c| 10 * c).collect();
|
||||
let want = reference::conv_grouped(&x, cin, h, h, &wt, 1, 1, stride, 5, 1, &mult, &bias, shift);
|
||||
assert_eq!(want, (out.to_vec(), 2, 2), "reference, {cin}×{h}² stride {stride}");
|
||||
for p in FAST {
|
||||
let got =
|
||||
fast_conv_runs(|| with_policy(p, || conv_grouped(&x, cin, h, h, &wt, 1, 1, stride, 5, 1, &mult, &bias, shift)));
|
||||
assert_eq!(got, (want.clone(), runs), "{cin}×{h}² stride {stride} {p:?}");
|
||||
}
|
||||
}
|
||||
// Transposed (the reference accumulates in `i128`), 1×1 input and kernel,
|
||||
// stride 1: (channels, pad, output pad, expected fast-kernel runs, output
|
||||
// side). No tap reaches the kept output.
|
||||
for (cin, pad, opad, runs, side) in [(2usize, 6usize, 13usize, 1u64, 2usize), (2, 7, 15, 0, 2), (1, 6, 12, 0, 1)] {
|
||||
let (x, wt) = (vec![5i64; cin], vec![3i64; cin]);
|
||||
let want = reference::conv_transpose(&x, cin, 1, 1, &wt, 1, 1, 1, pad, opad, &mult, &bias, shift);
|
||||
assert_eq!(want, (vec![7; side * side], side, side), "reference, pad {pad} opad {opad}");
|
||||
for p in FAST {
|
||||
let got =
|
||||
fast_conv_runs(|| with_policy(p, || conv_transpose(&x, cin, 1, 1, &wt, 1, 1, 1, pad, opad, &mult, &bias, shift)));
|
||||
assert_eq!(got, (want.clone(), runs), "pad {pad} opad {opad} {p:?}");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -30,6 +30,7 @@
|
|||
// On the JPEG AI exact decode path: no platform math-library call (see clippy.toml).
|
||||
#![deny(clippy::disallowed_methods)]
|
||||
|
||||
use crate::int_kernels::{par_map, par_zip};
|
||||
use crate::int_transformer as it;
|
||||
use std::collections::HashMap;
|
||||
|
||||
|
|
@ -93,11 +94,14 @@ impl Weights {
|
|||
}
|
||||
}
|
||||
|
||||
// The element-wise helpers compute each output element with the same pure
|
||||
// function of its own inputs whatever the partition; `par_map`/`par_zip` only spread
|
||||
// the elements over workers (one worker under the reference policy).
|
||||
fn relu(x: &[i64]) -> Vec<i64> {
|
||||
x.iter().map(|&v| if v < 0 { 0 } else { v }).collect()
|
||||
par_map(x, |v| if v < 0 { 0 } else { v })
|
||||
}
|
||||
fn relu6_of(x: &[i64]) -> Vec<i64> {
|
||||
x.iter().map(|&v| it::clip(v, 0, 6 * S)).collect()
|
||||
par_map(x, |v| it::clip(v, 0, 6 * S))
|
||||
}
|
||||
/// conv1 → ReLU → conv2 → + input residual block (`first_stage.0.0`, `CAB.*`).
|
||||
fn res_block(p: &str, x: &[i64], c: usize, h: usize, w: usize, wt: &Weights) -> Vec<i64> {
|
||||
|
|
@ -127,15 +131,15 @@ fn depth_to_space_crd(x: &[i64], c: usize, h: usize, w: usize, b: usize) -> (Vec
|
|||
(out, co, ho, wo)
|
||||
}
|
||||
fn emul(a: &[i64], b: &[i64]) -> Vec<i64> {
|
||||
a.iter().zip(b).map(|(&x, &y)| ((x as i128 * y as i128) >> F) as i64).collect()
|
||||
par_zip(a, b, |x, y| ((x as i128 * y as i128) >> F) as i64)
|
||||
}
|
||||
fn eadd(a: &[i64], b: &[i64]) -> Vec<i64> {
|
||||
a.iter().zip(b).map(|(&x, &y)| x + y).collect()
|
||||
par_zip(a, b, |x, y| x + y)
|
||||
}
|
||||
/// Residual multiplicative gate `x · (f(x) + 1)` (the exact form: the ONNX gate
|
||||
/// `Add`s carry const 1.0 with no activation). Near-identity when f≈0.
|
||||
fn gate(x: &[i64], modulation: &[i64]) -> Vec<i64> {
|
||||
let one_plus: Vec<i64> = modulation.iter().map(|&v| v + S).collect();
|
||||
let one_plus: Vec<i64> = par_map(modulation, |v| v + S);
|
||||
emul(x, &one_plus)
|
||||
}
|
||||
|
||||
|
|
@ -232,7 +236,7 @@ fn decode_y_planes(latent: &[i64], wt: &Weights) -> Vec<i64> {
|
|||
let m2 = res_block("CAB.residual_mask2.0", &s0, 128, sh, sw, wt);
|
||||
let m3 = res_block("CAB.residual_mask3.0", &m2, 128, sh, sw, wt);
|
||||
let (u, _, _) = wt.convt("CAB.upscale", &m3, 128, sh, sw, 3, 2, 1, 1); // ⇑2 → [128,64,64]
|
||||
u.iter().map(|&v| it::sigmoid(v)).collect::<Vec<_>>()
|
||||
par_map(&u, it::sigmoid)
|
||||
};
|
||||
x = eadd(&xc, &emul(&trunk, &mask));
|
||||
// iact2 gated: ReLU6 feeds the conv path, gate uses RAW x
|
||||
|
|
@ -309,7 +313,7 @@ fn decode_uv_planes(latent: &[i64], wt: &Weights) -> Vec<i64> {
|
|||
let m2 = res_block("CAB.residual_mask2.0", &s0, c, sh, sw, wt);
|
||||
let m3 = res_block("CAB.residual_mask3.0", &m2, c, sh, sw, wt);
|
||||
let (u, _, _) = wt.convt("CAB.upscale", &m3, c, sh, sw, 3, 2, 1, 1); // ⇑2
|
||||
u.iter().map(|&v| it::sigmoid(v)).collect::<Vec<_>>()
|
||||
par_map(&u, it::sigmoid)
|
||||
};
|
||||
let mut x = eadd(&xc, &emul(&trunk, &mask));
|
||||
// iact2 gated: ReLU6 feeds the conv path, gate uses RAW x
|
||||
|
|
|
|||
|
|
@ -31,6 +31,16 @@ pub mod ffi;
|
|||
#[cfg(feature = "neural_int")]
|
||||
pub mod int_synth;
|
||||
|
||||
// Fast kernels (convolution, transposed convolution, IGDN, channel attention)
|
||||
// for the integer decoders, plus the run-time policy that selects them or the
|
||||
// scalar reference kernels. Output is the same integer for every element either
|
||||
// way, and an overflow-checked build panics on the same inputs (the message can
|
||||
// differ); an allocation that fails is outside that claim. Argued and tested, not
|
||||
// machine-proved: the module docs state the claim precisely. Pure std — no dependency — and single-threaded
|
||||
// on wasm.
|
||||
#[cfg(feature = "neural_int")]
|
||||
pub mod int_kernels;
|
||||
|
||||
// Integer Gaussian-conditional entropy — the reusable scale→bucket primitive
|
||||
// that makes a conditional codec's entropy decode hardware-deterministic (the
|
||||
// integer half of the integer-entropy contract). Pure i64, no `ort`, like
|
||||
|
|
@ -54,9 +64,10 @@ pub mod rans;
|
|||
#[cfg(feature = "neural_int")]
|
||||
pub mod int_transform;
|
||||
|
||||
// Integer-exact synthesis of a SOTA-class learned codec (per-channel int12 convs
|
||||
// + integer IGDN) — the decode side of a NeuralIntegerExact conditional codec at
|
||||
// SOTA quality. Pure i64, no `ort`. Runs the model the export tool emits.
|
||||
// Integer-exact synthesis of a pretrained learned image codec (per-channel int12
|
||||
// convs + integer IGDN) — the decode side of a NeuralIntegerExact conditional
|
||||
// codec (the pretrained mbt2018_mean, exported to integers; quality measured in
|
||||
// int_hyper_synth.rs). Pure i64, no `ort`. Runs the model the export tool emits.
|
||||
#[cfg(feature = "neural_int")]
|
||||
pub mod int_hyper_synth;
|
||||
|
||||
|
|
|
|||
|
|
@ -1,32 +1,64 @@
|
|||
//! The Rust integer synthesis of a SOTA-class learned codec reproduces the
|
||||
//! Python export tool's reconstruction BYTE-FOR-BYTE — proving the decode is
|
||||
//! cross-language / cross-hardware exact at SOTA quality (mbt2018_mean, 42.6 dB,
|
||||
//! fully integer). The model + input + expected FNV-1a checksum are emitted by
|
||||
//! `tools/wai_sota_int_synth_export.py` into a directory pointed to by
|
||||
//! `WAI_SOTA_INT_DIR` (default `/tmp/sota_int`); skipped if absent.
|
||||
//! The Rust integer decode of `wai.neural.int_hyper` — the pretrained mbt2018_mean
|
||||
//! mean-scale hyperprior (CompressAI zoo, quality 6), exported to integers —
|
||||
//! reproduces the Python export tool's reconstruction BYTE-FOR-BYTE, and costs at
|
||||
//! most 1 dB against the same model's float decode on the photograph the vectors
|
||||
//! were built from. The decode is all-integer by construction; its
|
||||
//! cross-architecture byte-equality is checked on the committed tiny vector
|
||||
//! (`tests/vectors/int_hyper/`, `wai-byte-equality.yml`), not on these. The
|
||||
//! pretrained-model tests below take their expected hashes from the export run
|
||||
//! that wrote the vector directory, not from the repository (only the tiny
|
||||
//! vector's test reads a committed hash). That export encodes with the float model,
|
||||
//! and its files are not portable: the synthesis-only test's latent and its
|
||||
//! `rgb_hash` were measured to differ between an Apple Silicon Mac and an x86_64
|
||||
//! Linux runner, and two runs on that runner wrote different `cond.json` and
|
||||
//! `int_synth_sota.json` (most likely in `psnr_float`, which only the quality
|
||||
//! test reads, against thresholds). None of these hashes is pinned for use by
|
||||
//! another run (`tools/wai_sota_vectors.sh`).
|
||||
//!
|
||||
//! The model + bitstream + expected FNV-1a checksums are regenerated by
|
||||
//! `tools/wai_sota_vectors.sh` (pretrained mbt2018_mean exported to integers by
|
||||
//! `tools/wai_sota_int_synth_export.py`) into `WAI_SOTA_INT_DIR` (default
|
||||
//! `/tmp/sota_int`). CI (`wai-sota-vectors.yml`) regenerates them from a pinned
|
||||
//! public-domain photograph and runs this file with `WAI_REQUIRE_SOTA_VECTORS=1`,
|
||||
//! so an absent vector fails there; a local run without them skips loudly. A
|
||||
//! `WAI_SOTA_INT_DIR` naming a directory that does not exist, or one missing a
|
||||
//! file a test reads, fails either way: only the default may be absent and skip.
|
||||
//!
|
||||
//! Quality is image-dependent: the 42.6 dB development figure for this codec was
|
||||
//! measured on a photograph that is not in the repository. What is asserted here,
|
||||
//! on whatever image produced the vectors, is the integer decode's cost against the
|
||||
//! float codec on that same image. Nothing here compares the model with any other
|
||||
//! codec.
|
||||
#![cfg(feature = "neural_int")]
|
||||
|
||||
use std::path::PathBuf;
|
||||
#[path = "support/require_vectors.rs"]
|
||||
mod require_vectors;
|
||||
|
||||
use require_vectors::{load_all, load_json, InputDir};
|
||||
use wai::int_hyper_synth::{decode, fnv1a, IntModel};
|
||||
|
||||
fn dir() -> PathBuf {
|
||||
PathBuf::from(std::env::var("WAI_SOTA_INT_DIR").unwrap_or_else(|_| "/tmp/sota_int".into()))
|
||||
const PRODUCE: &str = "tools/wai_sota_vectors.sh <dir> (then WAI_SOTA_INT_DIR=<dir>, default /tmp/sota_int)";
|
||||
|
||||
/// `WAI_SOTA_INT_DIR`, else `/tmp/sota_int`. Set to a directory that does not
|
||||
/// exist, or to one missing a file the test reads, it fails the test instead of
|
||||
/// letting it skip.
|
||||
fn dir() -> InputDir {
|
||||
require_vectors::env_dir("int_hyper_synth", "WAI_SOTA_INT_DIR", "/tmp/sota_int", PRODUCE)
|
||||
}
|
||||
|
||||
fn jload(p: &std::path::Path) -> Option<serde_json::Value> {
|
||||
serde_json::from_slice(&std::fs::read(p).ok()?).ok()
|
||||
/// Every named file from the vector dir, or `None` after reporting what is absent.
|
||||
fn inputs<const N: usize>(test: &str, d: &InputDir, names: [&str; N]) -> Option<[serde_json::Value; N]> {
|
||||
load_all(test, d, names, PRODUCE)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rust_integer_synthesis_reproduces_the_export_hash() {
|
||||
let d = dir();
|
||||
let (Some(conf), Some(model), Some(input)) = (
|
||||
jload(&d.join("int_synth_sota.json")),
|
||||
jload(&d.join("int_synth_sota_model.json")),
|
||||
jload(&d.join("int_synth_sota_input.json")),
|
||||
let Some([conf, model, input]) = inputs(
|
||||
"int_hyper_synth::rust_integer_synthesis_reproduces_the_export_hash",
|
||||
&d,
|
||||
["int_synth_sota.json", "int_synth_sota_model.json", "int_synth_sota_input.json"],
|
||||
) else {
|
||||
eprintln!("skipping SOTA int-synth: artifacts not in {} (run tools/wai_sota_int_synth_export.py)", d.display());
|
||||
return;
|
||||
};
|
||||
|
||||
|
|
@ -44,19 +76,18 @@ fn rust_integer_synthesis_reproduces_the_export_hash() {
|
|||
|
||||
let got = format!("{:016x}", fnv1a(&rgb));
|
||||
let want = conf["rgb_hash"].as_str().unwrap();
|
||||
println!("SOTA integer synthesis: {oh}x{ow}, fnv1a {got} (psnr {:.2} dB)", conf["psnr_int"].as_f64().unwrap_or(0.0));
|
||||
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");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rust_integer_hyper_decoder_reproduces_the_export_hash() {
|
||||
let d = dir();
|
||||
let (Some(conf), Some(model), Some(input)) = (
|
||||
jload(&d.join("int_hs.json")),
|
||||
jload(&d.join("int_hs_model.json")),
|
||||
jload(&d.join("int_hs_input.json")),
|
||||
let Some([conf, model, input]) = inputs(
|
||||
"int_hyper_synth::rust_integer_hyper_decoder_reproduces_the_export_hash",
|
||||
&d,
|
||||
["int_hs.json", "int_hs_model.json", "int_hs_input.json"],
|
||||
) else {
|
||||
eprintln!("skipping SOTA int h_s: artifacts not in {}", d.display());
|
||||
return;
|
||||
};
|
||||
let (c, h, w) = (
|
||||
|
|
@ -81,18 +112,18 @@ fn decodes_the_committed_tiny_conformance_vector() {
|
|||
// A SMALL synthetic int_hyper codec, committed under tests/vectors/int_hyper/,
|
||||
// drives the full decode path (factorized z → integer h_s → integer-σ entropy
|
||||
// → integer g_s) with NO external artifacts — so CI verifies the int_hyper
|
||||
// capability without the 23 MB SOTA model. Both the JSON and the compact
|
||||
// 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};
|
||||
let d = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/vectors/int_hyper");
|
||||
let conf = jload(&d.join("tiny.json")).expect("committed tiny.json");
|
||||
let ent = CondEntropy::from_value(&jload(&d.join("tiny_entropy.json")).unwrap());
|
||||
let bs = CondBitstream::from_value(&jload(&d.join("tiny_bitstream.json")).unwrap());
|
||||
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());
|
||||
let bs = CondBitstream::from_value(&load_json(&d.join("tiny_bitstream.json")).unwrap());
|
||||
let want = conf["rgb_hash"].as_str().unwrap();
|
||||
let (oh, ow) = (conf["out_h"].as_u64().unwrap() as usize, conf["out_w"].as_u64().unwrap() as usize);
|
||||
|
||||
let g_j = IntModel::from_json(&jload(&d.join("tiny_g_s.json")).unwrap());
|
||||
let h_j = IntModel::from_json(&jload(&d.join("tiny_h_s.json")).unwrap());
|
||||
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");
|
||||
assert_eq!((h, w), (oh, ow));
|
||||
assert_eq!(format!("{:016x}", fnv1a(&rgb)), want, "tiny int_hyper (json model) must reproduce the committed hash");
|
||||
|
|
@ -113,14 +144,12 @@ fn rust_decodes_the_full_conditional_bitstream_to_the_export_image() {
|
|||
// identical on every machine). This is the integer-entropy contract assembled.
|
||||
use wai::int_hyper_synth::{decode_conditional, CondBitstream, CondEntropy};
|
||||
let d = dir();
|
||||
let (Some(conf), Some(g_s), Some(h_s), Some(ent), Some(bs)) = (
|
||||
jload(&d.join("cond.json")),
|
||||
jload(&d.join("int_synth_sota_model.json")),
|
||||
jload(&d.join("int_hs_model.json")),
|
||||
jload(&d.join("cond_entropy.json")),
|
||||
jload(&d.join("cond_bitstream.json")),
|
||||
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(
|
||||
test,
|
||||
&d,
|
||||
["cond.json", "int_synth_sota_model.json", "int_hs_model.json", "cond_entropy.json", "cond_bitstream.json"],
|
||||
) else {
|
||||
eprintln!("skipping conditional bitstream: artifacts not in {}", d.display());
|
||||
return;
|
||||
};
|
||||
let (g_s, h_s) = (IntModel::from_json(&g_s), IntModel::from_json(&h_s));
|
||||
|
|
@ -137,9 +166,12 @@ fn rust_decodes_the_full_conditional_bitstream_to_the_export_image() {
|
|||
|
||||
// The COMPACT BINARY model (~½ the size) must decode bit-identically to the
|
||||
// JSON — same codec, smaller download (what makes the in-browser demo practical).
|
||||
if let (Ok(g_bin), Ok(h_bin)) =
|
||||
(std::fs::read(d.join("int_synth_sota_model.bin")), std::fs::read(d.join("int_hs_model.bin")))
|
||||
{
|
||||
// The export writes both, so an absent binary is an incomplete vector dir.
|
||||
let (gb, hb) = (d.join("int_synth_sota_model.bin"), d.join("int_hs_model.bin"));
|
||||
if !gb.exists() || !hb.exists() {
|
||||
d.absent(&format!("{test} (binary model)"), &format!("int_synth_sota_model.bin / int_hs_model.bin not in {}", d.display()), PRODUCE);
|
||||
} else {
|
||||
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");
|
||||
|
|
@ -148,3 +180,54 @@ fn rust_decodes_the_full_conditional_bitstream_to_the_export_image() {
|
|||
println!("compact binary model ({} + {} bytes) decodes byte-identically", g_bin.len(), h_bin.len());
|
||||
}
|
||||
}
|
||||
|
||||
/// 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 (Apple Silicon Mac); 0.62 dB for the integer synthesis alone
|
||||
/// in development, on a photograph that is not in the repository. A larger gap
|
||||
/// means the quantization has regressed, even when every byte-exactness check
|
||||
/// above still passes (they compare the integer decode only with itself).
|
||||
const MAX_INTEGER_COST_DB: f64 = 1.0;
|
||||
/// A sanity floor for the float codec on the source image: below it the pretrained
|
||||
/// checkpoint is not coding natural content at all, and a matching integer decode
|
||||
/// of it proves nothing. A sanity floor only, not a quality claim.
|
||||
const MIN_FLOAT_PSNR_DB: f64 = 30.0;
|
||||
|
||||
#[test]
|
||||
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};
|
||||
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(
|
||||
test,
|
||||
&d,
|
||||
["cond.json", "int_synth_sota_model.json", "int_hs_model.json", "cond_entropy.json", "cond_bitstream.json"],
|
||||
) else {
|
||||
return;
|
||||
};
|
||||
let src_path = d.join("source.rgb");
|
||||
let Some(psnr_float) = conf.get("psnr_float").and_then(|v| v.as_f64()) else {
|
||||
d.absent(test, "cond.json predates psnr_float (regenerate the vectors)", PRODUCE);
|
||||
return;
|
||||
};
|
||||
if !src_path.exists() {
|
||||
d.absent(test, &format!("source.rgb not in {}", d.display()), PRODUCE);
|
||||
return;
|
||||
}
|
||||
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))
|
||||
.expect("conditional decode");
|
||||
assert_eq!(source.len(), rgb.len(), "source.rgb is not the {h}x{w} RGB image the bitstream decodes to");
|
||||
|
||||
let sse: f64 = rgb.iter().zip(&source).map(|(a, b)| (f64::from(*a) - f64::from(*b)).powi(2)).sum();
|
||||
let mse = sse / rgb.len() as f64;
|
||||
let psnr = 10.0 * (255.0f64 * 255.0 / mse).log10();
|
||||
let cost = psnr_float - psnr;
|
||||
println!("integer decode {psnr:.3} dB vs float codec {psnr_float:.3} dB on the same {h}x{w} image (cost {cost:.3} dB)");
|
||||
assert!(psnr_float >= MIN_FLOAT_PSNR_DB, "float codec {psnr_float:.2} dB < {MIN_FLOAT_PSNR_DB} dB: the checkpoint or source is wrong");
|
||||
assert!(cost <= MAX_INTEGER_COST_DB, "integer decode costs {cost:.2} dB against the float codec (limit {MAX_INTEGER_COST_DB} dB)");
|
||||
}
|
||||
|
|
|
|||
293
wai/wai-rs/tests/int_kernels_geometry.rs
Normal file
293
wai/wai-rs/tests/int_kernels_geometry.rs
Normal file
|
|
@ -0,0 +1,293 @@
|
|||
//! Geometry at the edges of what the fast integer kernels lay out, run through the
|
||||
//! public API under the scalar reference kernels AND under each fast-kernel policy
|
||||
//! (serial, explicit multi-threaded tilings, automatic):
|
||||
//!
|
||||
//! * a 1×1 input upsampled by strides from 1 up to `usize::MAX`, through the
|
||||
//! transposed convolutions of int_hyper (`run_net`), int_synth
|
||||
//! (`decode_from_spec`) and the JPEG-AI synthesis (`conv_transpose`);
|
||||
//! * a regular convolution whose stride reaches or exceeds its padded input
|
||||
//! (`run_net`, `conv_grouped`);
|
||||
//! * padding far beyond the kernel, with output padding restoring a small output;
|
||||
//! * an input of two or more rows and columns whose stride and padding make the
|
||||
//! reference's own size arithmetic overflow (int_hyper, int_synth, JPEG-AI).
|
||||
//!
|
||||
//! The reference computes each of the first three in a handful of operations. Every
|
||||
//! policy must return the same integers (checked against a hand computation), and
|
||||
//! the fast path must not overflow, or allocate or loop in proportion to the stride
|
||||
//! or the padding, where the reference does not. In the last, the reference panics
|
||||
//! (or, for int_synth with overflow checks off, returns), and every policy must do
|
||||
//! the same. The int_audio 1-D transposed convolution has the
|
||||
//! same stride check beside its kernel (`int_audio` unit tests). CI runs this file
|
||||
//! beside `int_kernels_vectors` on every `wai-conformance` leg, with overflow checks
|
||||
//! on and off.
|
||||
#![cfg(feature = "neural_int")]
|
||||
|
||||
use serde_json::json;
|
||||
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)];
|
||||
|
||||
/// A 3×3 kernel `[[1, 2, 3], [4, 5, 6], [7, 8, 9]]`.
|
||||
const KERNEL: [i64; 9] = [1, 2, 3, 4, 5, 6, 7, 8, 9];
|
||||
|
||||
/// Run under every policy; all results must be identical. Returns the result.
|
||||
fn same_under_every_policy<T: PartialEq + std::fmt::Debug>(what: &str, f: impl Fn() -> T) -> T {
|
||||
let want = with_policy(POLICIES[0], &f);
|
||||
for p in &POLICIES[1..] {
|
||||
let got = with_policy(*p, &f);
|
||||
assert!(got == want, "{what}: result under {p:?} differs from the reference kernels ({got:?} vs {want:?})");
|
||||
}
|
||||
want
|
||||
}
|
||||
|
||||
/// Whether this build traps on integer overflow. Detected by performing one, since
|
||||
/// `cfg!(debug_assertions)` only usually coincides with `overflow-checks`.
|
||||
#[cfg(target_pointer_width = "64")]
|
||||
fn overflow_checks() -> bool {
|
||||
std::panic::catch_unwind(|| std::hint::black_box(i64::MAX) + std::hint::black_box(1)).is_err()
|
||||
}
|
||||
|
||||
/// Run under every policy, catching a panic. The outcome (the value, or that the
|
||||
/// call panicked) must be the same under every policy. Returns the reference
|
||||
/// kernels' outcome, `None` for a panic. The panic message is not compared: a fast
|
||||
/// policy may fail at another element first, or on a worker thread.
|
||||
fn same_outcome_under_every_policy<T: PartialEq + std::fmt::Debug>(what: &str, f: impl Fn() -> T) -> Option<T> {
|
||||
let run = |p: Policy| std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| with_policy(p, &f))).ok();
|
||||
let want = run(POLICIES[0]);
|
||||
for p in &POLICIES[1..] {
|
||||
let got = run(*p);
|
||||
assert!(got == want, "{what}: outcome under {p:?} differs from the reference kernels' ({got:?} vs {want:?}; None = panicked)");
|
||||
}
|
||||
want
|
||||
}
|
||||
|
||||
/// Strides from ordinary to the largest a `usize` holds. 2^29 fits the `u32` stride
|
||||
/// field of a WHM1 model; from 2^32 (64-bit targets) the product of the two strides
|
||||
/// no longer fits in 64 bits.
|
||||
///
|
||||
/// 2^29 comes first, deliberately. A kernel that allocates one tap table per
|
||||
/// (row, column) phase — as the fast transposed kernel did before it built tables
|
||||
/// only for phases holding output — fails there at once (the table count exceeds
|
||||
/// what any allocation can hold: "capacity overflow"). Placed after 2^16 it would
|
||||
/// first commit memory for 2^32 tables (hundreds of GiB) and drive the machine into
|
||||
/// swap before failing.
|
||||
fn strides() -> Vec<usize> {
|
||||
#[allow(unused_mut)] // pushed to on 64-bit targets only
|
||||
let mut s = vec![1 << 29, 1, 2, 3, 1 << 13, 1 << 16, usize::MAX];
|
||||
#[cfg(target_pointer_width = "64")]
|
||||
s.push(1 << 32);
|
||||
s
|
||||
}
|
||||
|
||||
/// A one-layer int_hyper model.
|
||||
fn hyper_model(layer: serde_json::Value) -> wai::int_hyper_synth::IntModel {
|
||||
wai::int_hyper_synth::IntModel::from_json(&json!({ "afrac": 0, "gf": 0, "shift": 0, "layers": [layer] }))
|
||||
}
|
||||
|
||||
/// An int_synth payload: a 1×1 output image from one int8 latent (C = 1, S = 1).
|
||||
fn synth_payload(latent: i8) -> Vec<u8> {
|
||||
let mut p = Vec::new();
|
||||
p.extend(1u32.to_le_bytes()); // H
|
||||
p.extend(1u32.to_le_bytes()); // W
|
||||
p.extend(1u16.to_le_bytes()); // C
|
||||
p.extend([0u8, 0]); // S_log2, reserved
|
||||
p.push(latent as u8);
|
||||
p
|
||||
}
|
||||
|
||||
/// A one-layer int_synth spec: 1 → 3 channels, 3×3 kernel `1..=27` (channel `oc`
|
||||
/// holds `9·oc + KERNEL`), shift 0, no ReLU.
|
||||
fn synth_spec(stride: usize, pad: usize, out_pad: usize, bias: [i32; 3]) -> Vec<u8> {
|
||||
let weight: Vec<i32> = (1..=27).collect();
|
||||
json!({ "layers": [{
|
||||
"in_ch": 1, "out_ch": 3, "k": 3, "stride": stride, "pad": pad, "out_pad": out_pad,
|
||||
"weight": weight, "bias": bias, "shift": 0, "relu": false
|
||||
}]})
|
||||
.to_string()
|
||||
.into_bytes()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_1x1_input_upsampled_by_any_stride() {
|
||||
// k = 3, pad 1: the output is 1×1 and receives only the centre tap, 5·x.
|
||||
for s in strides() {
|
||||
let got = same_under_every_policy(&format!("JPEG-AI conv_transpose, stride {s}"), || {
|
||||
wai::int_transformer::conv_transpose(&[5], 1, 1, 1, &KERNEL, 1, 3, s, 1, 0, &[1], &[0], 0)
|
||||
});
|
||||
assert_eq!(got, (vec![25], 1, 1), "JPEG-AI conv_transpose, stride {s}");
|
||||
|
||||
let model = hyper_model(json!({
|
||||
"type": "convT", "out_ch": 1, "k": 3, "stride": s, "pad": 1, "opad": 0,
|
||||
"weight": KERNEL, "mult": [1], "bias": [0]
|
||||
}));
|
||||
let got = same_under_every_policy(&format!("int_hyper convT, stride {s}"), || {
|
||||
wai::int_hyper_synth::run_net(&model, &[5], 1, 1, 1)
|
||||
});
|
||||
assert_eq!(got, (vec![25], 1, 1, 1), "int_hyper convT, stride {s}");
|
||||
|
||||
let (payload, spec) = (synth_payload(5), synth_spec(s, 1, 0, [0, 0, 0]));
|
||||
let got = same_under_every_policy(&format!("int_synth, stride {s}"), || {
|
||||
wai::int_synth::decode_from_spec(&payload, &spec).map(|img| img.rgb)
|
||||
});
|
||||
// centre taps 5, 14, 23 of the three output channels, times 5
|
||||
assert_eq!(got, Ok(vec![25, 70, 115]), "int_synth, stride {s}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_regular_conv_whose_stride_exceeds_its_padded_input() {
|
||||
// 5×5 input 1..=25, k = 3. Pad 1 (7×7 padded): from stride 5 on, one output,
|
||||
// the taps that meet x[0], x[1], x[5], x[6]: 5·1 + 6·2 + 8·6 + 9·7. Pad 2 (9×9
|
||||
// padded): from stride 7 on, one output, only tap (2, 2) meets x[0]: 9·1. At
|
||||
// stride usize::MAX, pad 2 puts two leading column phases before the first input
|
||||
// column, whose position must be found without overflowing.
|
||||
let x: Vec<i64> = (1..=25).collect();
|
||||
for (pad, want, first) in [(1usize, 5 + 6 * 2 + 8 * 6 + 9 * 7, 5usize), (2, 9, 7)] {
|
||||
#[allow(unused_mut)] // pushed to on 64-bit targets only
|
||||
let mut ss = vec![first, 7, 8, 9, 1 << 13, 1 << 16, 1 << 29, usize::MAX];
|
||||
#[cfg(target_pointer_width = "64")]
|
||||
ss.push(1 << 32);
|
||||
for s in ss {
|
||||
let got = same_under_every_policy(&format!("JPEG-AI conv_grouped, pad {pad}, stride {s}"), || {
|
||||
wai::int_transformer::conv_grouped(&x, 1, 5, 5, &KERNEL, 1, 3, s, pad, 1, &[1], &[0], 0)
|
||||
});
|
||||
assert_eq!(got, (vec![want], 1, 1), "JPEG-AI conv_grouped, pad {pad}, stride {s}");
|
||||
|
||||
let model = hyper_model(json!({
|
||||
"type": "conv", "out_ch": 1, "k": 3, "stride": s, "pad": pad,
|
||||
"weight": KERNEL, "mult": [1], "bias": [0]
|
||||
}));
|
||||
let got = same_under_every_policy(&format!("int_hyper conv, pad {pad}, stride {s}"), || {
|
||||
wai::int_hyper_synth::run_net(&model, &x, 1, 5, 5)
|
||||
});
|
||||
assert_eq!(got, (vec![want], 1, 1, 1), "int_hyper conv, pad {pad}, stride {s}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn padding_far_beyond_the_kernel() {
|
||||
let p = 1usize << 12;
|
||||
// Transposed: pad 2^12 with output padding 2^13 − 2 restores a 1×1 output that
|
||||
// no tap reaches (it sits 2^12 rows into the uncropped plane), so it is the bias.
|
||||
let (payload, spec) = (synth_payload(5), synth_spec(1, p, 2 * p - 2, [10, 20, 30]));
|
||||
let got = same_under_every_policy("int_synth, pad 2^12", || {
|
||||
wai::int_synth::decode_from_spec(&payload, &spec).map(|img| img.rgb)
|
||||
});
|
||||
assert_eq!(got, Ok(vec![10, 20, 30]));
|
||||
// Regular: a 1×1 input, 1×1 kernel, pad 2^12, stride 2^13 — a 2×2 output whose
|
||||
// samples all fall in the padding, so each is the bias.
|
||||
let got = same_under_every_policy("JPEG-AI conv_grouped, pad 2^12", || {
|
||||
wai::int_transformer::conv_grouped(&[5], 1, 1, 1, &[3], 1, 1, 2 * p, p, 1, &[1], &[7], 0)
|
||||
});
|
||||
assert_eq!(got, (vec![7; 4], 2, 2));
|
||||
let model = hyper_model(json!({
|
||||
"type": "conv", "out_ch": 1, "k": 1, "stride": 2 * p, "pad": p, "weight": [3], "mult": [1], "bias": [7]
|
||||
}));
|
||||
let got = same_under_every_policy("int_hyper conv, pad 2^12", || wai::int_hyper_synth::run_net(&model, &[5], 1, 1, 1));
|
||||
assert_eq!(got, (vec![7; 4], 1, 2, 2));
|
||||
}
|
||||
|
||||
/// An int_synth payload: an `h`×`w` output image from a 2×2 int8 latent (C = 1,
|
||||
/// S = 2), latents in row-major order.
|
||||
#[cfg(target_pointer_width = "64")]
|
||||
fn synth_payload_2x2(h: u32, w: u32, latents: [i8; 4]) -> Vec<u8> {
|
||||
let mut p = Vec::new();
|
||||
p.extend(h.to_le_bytes());
|
||||
p.extend(w.to_le_bytes());
|
||||
p.extend(1u16.to_le_bytes()); // C
|
||||
p.extend([1u8, 0]); // S_log2 = 1, reserved
|
||||
p.extend(latents.map(|v| v as u8));
|
||||
p
|
||||
}
|
||||
|
||||
// Transposed convolutions of an input with two or more rows and columns, whose
|
||||
// stride and padding make the reference's own size arithmetic overflow. The
|
||||
// int_hyper, int_audio and JPEG-AI references size an uncropped plane
|
||||
// `cout·(ho + 2·pad)·(wo + 2·pad)` in unchecked `usize` arithmetic; the int_synth
|
||||
// reference computes the output size in unchecked `i64`. With overflow checks on,
|
||||
// the reference panics on that arithmetic. With them off it wraps and, for the
|
||||
// uncropped plane, then indexes out of the wrapped buffer. A fast policy must have
|
||||
// the same outcome: it may not return a value where the reference panics. A 1×1
|
||||
// input cannot show this, since `(h − 1)·stride` is 0. The int_audio 1-D form is
|
||||
// tested beside its kernel.
|
||||
|
||||
/// A 3×3 input, and a 2×2 one with fields that fit a WHM1 model's `u32`.
|
||||
const X: [i64; 9] = [1, 2, 3, 4, 5, 6, 7, 8, 9];
|
||||
/// Stride 2^32 − 1 and pad 2^31 − 1: a 2×2 output from a 2×2 input, k = 1 (or a
|
||||
/// 4×4 one, k = 3), under an uncropped plane of 2^32 × 2^32. On a 64-bit target
|
||||
/// the plane size overflows; on a 32-bit one the stride arithmetic already does.
|
||||
const S32: usize = u32::MAX as usize;
|
||||
const P32: usize = (u32::MAX >> 1) as usize;
|
||||
|
||||
#[test]
|
||||
fn jpegai_conv_transpose_where_the_reference_plane_size_overflows() {
|
||||
// 3×3 input, k = 1, stride = pad = 2^62: a 1×1 output, but an uncropped plane
|
||||
// of (2^63 + 1)² elements. The reference panics in both profiles.
|
||||
#[cfg(target_pointer_width = "64")]
|
||||
{
|
||||
let s = 1usize << 62;
|
||||
let got = same_outcome_under_every_policy("JPEG-AI conv_transpose, 3×3, stride = pad = 2^62", || {
|
||||
wai::int_transformer::conv_transpose(&X, 1, 3, 3, &[7], 1, 1, s, s, 0, &[1], &[0], 0)
|
||||
});
|
||||
assert_eq!(got, None, "the reference panics");
|
||||
}
|
||||
let got = same_outcome_under_every_policy("JPEG-AI conv_transpose, 2×2, u32 stride and pad", || {
|
||||
wai::int_transformer::conv_transpose(&X[..4], 1, 2, 2, &[7], 1, 1, S32, P32, 0, &[1], &[0], 0)
|
||||
});
|
||||
assert_eq!(got, None, "u32 fields: the reference panics");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn int_hyper_convt_where_the_reference_plane_size_overflows() {
|
||||
// As above, through the int_hyper decoder.
|
||||
let convt = |s: usize, p: usize| {
|
||||
hyper_model(json!({
|
||||
"type": "convT", "out_ch": 1, "k": 1, "stride": s, "pad": p, "opad": 0,
|
||||
"weight": [7], "mult": [1], "bias": [0]
|
||||
}))
|
||||
};
|
||||
#[cfg(target_pointer_width = "64")]
|
||||
{
|
||||
let model = convt(1 << 62, 1 << 62);
|
||||
let got = same_outcome_under_every_policy("int_hyper convT, 3×3, stride = pad = 2^62", || {
|
||||
wai::int_hyper_synth::run_net(&model, &X, 1, 3, 3)
|
||||
});
|
||||
assert_eq!(got, None, "the reference panics");
|
||||
}
|
||||
let model = convt(S32, P32);
|
||||
let got = same_outcome_under_every_policy("int_hyper convT, 2×2, u32 stride and pad", || {
|
||||
wai::int_hyper_synth::run_net(&model, &X[..4], 1, 2, 2)
|
||||
});
|
||||
assert_eq!(got, None, "u32 fields: the reference panics");
|
||||
}
|
||||
|
||||
#[cfg(target_pointer_width = "64")]
|
||||
#[test]
|
||||
fn int_synth_where_the_reference_output_size_overflows_i64() {
|
||||
// 2×2 latent, k = 3, stride = pad = out_pad = 2^62: a 3×3 output. The reference
|
||||
// computes (2 − 1)·2^62 + 3 + 2^62 − 2·2^62 in `i64`, which overflows (2^63 + 3)
|
||||
// although every `usize` step fits. With overflow checks on it panics; with them
|
||||
// off the wrapped sum is still 3 and the decode returns. Only latent (1, 1)
|
||||
// reaches the output, through every tap.
|
||||
let s = 1usize << 62;
|
||||
let (payload, spec) = (synth_payload_2x2(3, 3, [1, 2, 3, 4]), synth_spec(s, s, s, [0, 0, 0]));
|
||||
let got = same_outcome_under_every_policy("int_synth, 2×2 latent, stride = pad = out_pad = 2^62", || {
|
||||
wai::int_synth::decode_from_spec(&payload, &spec).map(|img| img.rgb)
|
||||
});
|
||||
if overflow_checks() {
|
||||
assert_eq!(got, None, "the reference panics on its i64 output size");
|
||||
} else {
|
||||
// rgb[(3y + x)·3 + c] = 4 · weight[c][y][x] = 4 · (9c + 3y + x + 1)
|
||||
let want: Vec<u8> = (0..9).flat_map(|px| (0..3).map(move |c| (4 * (9 * c + px + 1)) as u8)).collect();
|
||||
assert_eq!(got, Some(Ok(want)), "the wrapped output size is 3");
|
||||
}
|
||||
|
||||
// The u32 fields: no `i64` step overflows, and int_synth's reference allocates
|
||||
// no uncropped plane, so it returns. A 4×4 output that no tap reaches: the bias.
|
||||
let (payload, spec) = (synth_payload_2x2(4, 4, [1, 2, 3, 4]), synth_spec(S32, P32, 0, [10, 20, 30]));
|
||||
let got = same_outcome_under_every_policy("int_synth, 2×2 latent, u32 stride and pad", || {
|
||||
wai::int_synth::decode_from_spec(&payload, &spec).map(|img| img.rgb)
|
||||
});
|
||||
assert_eq!(got, Some(Ok([10u8, 20, 30].repeat(16))), "u32 fields: the bias alone");
|
||||
}
|
||||
168
wai/wai-rs/tests/int_kernels_vectors.rs
Normal file
168
wai/wai-rs/tests/int_kernels_vectors.rs
Normal file
|
|
@ -0,0 +1,168 @@
|
|||
//! Every committed integer-decoder vector whose decode reaches the fast kernels
|
||||
//! (int_hyper, int_audio, int_splat, int_world one and two agents, the
|
||||
//! neural-keyframe video, the int_synth corpus), plus the real-footage video,
|
||||
//! decoded under the scalar reference kernels AND under each fast-kernel policy
|
||||
//! (serial, explicit multi-threaded tilings, automatic): the outputs must be
|
||||
//! byte-identical to each other and must reproduce the vector's frozen hash. The
|
||||
//! plain `int_video` vector (block motion compensation + rANS, no learned layer)
|
||||
//! never reaches the kernels and is left to `tests/int_video.rs` and the
|
||||
//! `byte-exact-conformance` goldens. This is the whole-decoder half of the
|
||||
//! fast-vs-reference equivalence check; the per-kernel half (randomised shapes and
|
||||
//! operand ranges, element by element) lives beside each kernel. CI runs both
|
||||
//! halves on every platform leg of `wai-conformance`, with overflow checks on and
|
||||
//! off.
|
||||
#![cfg(all(feature = "neural_int", feature = "video_int", feature = "audio_int", feature = "world_int", feature = "splat_int"))]
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
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)];
|
||||
|
||||
fn vroot() -> PathBuf {
|
||||
PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/vectors")
|
||||
}
|
||||
fn jload(p: &Path) -> serde_json::Value {
|
||||
serde_json::from_slice(&std::fs::read(p).unwrap_or_else(|e| panic!("read {}: {e}", p.display()))).unwrap()
|
||||
}
|
||||
|
||||
/// Decode under every policy; all outputs must be identical. Returns the output.
|
||||
fn same_under_every_policy<T: PartialEq + std::fmt::Debug>(what: &str, f: impl Fn() -> T) -> T {
|
||||
let want = with_policy(POLICIES[0], &f);
|
||||
for p in &POLICIES[1..] {
|
||||
let got = with_policy(*p, &f);
|
||||
assert!(got == want, "{what}: output under {p:?} differs from the reference kernels");
|
||||
}
|
||||
want
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn int_hyper_tiny_conditional_codec() {
|
||||
use wai::int_hyper_synth::{decode_conditional, 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());
|
||||
assert_eq!(format!("{:016x}", fnv1a(&rgb)), want);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn int_audio_clip() {
|
||||
use wai::int_audio::{decode, pcm_fnv1a, AudioModel};
|
||||
let d = vroot().join("int_audio");
|
||||
let model = AudioModel::from_bytes(&std::fs::read(d.join("audio_model.bin")).unwrap()).unwrap();
|
||||
let bs = jload(&d.join("audio_bitstream.json"));
|
||||
let pcm = same_under_every_policy("int_audio", || decode(&model, &bs).unwrap());
|
||||
assert_eq!(format!("{:016x}", pcm_fnv1a(&pcm)), jload(&d.join("audio.json"))["pcm_fnv1a"].as_str().unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn int_splat_cloud() {
|
||||
use wai::int_hyper_synth::IntModel;
|
||||
use wai::int_splat::decode;
|
||||
let d = vroot().join("int_splat");
|
||||
let gs = IntModel::from_bytes(&std::fs::read(d.join("splat_gs.bin")).unwrap()).unwrap();
|
||||
let bs = jload(&d.join("splat.json"));
|
||||
let attr = same_under_every_policy("int_splat", || decode(&gs, &bs).unwrap().attr);
|
||||
let hash = wai::int_hyper_synth::fnv1a(&attr);
|
||||
assert_eq!(format!("{hash:016x}"), bs["attr_hash"].as_str().unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn int_world_clips() {
|
||||
use wai::int_hyper_synth::IntModel;
|
||||
use wai::int_world::{decode, decode_planes, ACTIONS};
|
||||
// single agent
|
||||
let d = vroot().join("int_world");
|
||||
let t = IntModel::from_bytes(&std::fs::read(d.join("world_T.bin")).unwrap()).unwrap();
|
||||
let gs = IntModel::from_bytes(&std::fs::read(d.join("world_gs.bin")).unwrap()).unwrap();
|
||||
let j = jload(&d.join("world.json"));
|
||||
let (c, lat) = (j["c"].as_u64().unwrap() as usize, j["lat"].as_u64().unwrap() as usize);
|
||||
let z0: Vec<i64> = j["z0"].as_array().unwrap().iter().map(|v| v.as_i64().unwrap()).collect();
|
||||
let actions: Vec<usize> = j["actions"].as_array().unwrap().iter().map(|v| v.as_u64().unwrap() as usize).collect();
|
||||
let frames = same_under_every_policy("int_world", || decode(&t, &gs, &z0, &actions, c, lat).frames);
|
||||
let clip = wai::int_world::WorldClip { frames, frame_px: 0 };
|
||||
assert_eq!(format!("{:016x}", clip.clip_hash()), j["clip_hash"].as_str().unwrap());
|
||||
|
||||
// two agents
|
||||
let d = vroot().join("int_world2");
|
||||
let t = IntModel::from_bytes(&std::fs::read(d.join("world2_T.bin")).unwrap()).unwrap();
|
||||
let gs = IntModel::from_bytes(&std::fs::read(d.join("world2_gs.bin")).unwrap()).unwrap();
|
||||
let j = jload(&d.join("world2.json"));
|
||||
let (c, lat) = (j["c"].as_u64().unwrap() as usize, j["lat"].as_u64().unwrap() as usize);
|
||||
let one = 1i64 << j["afrac"].as_u64().unwrap();
|
||||
let z0: Vec<i64> = j["z0"].as_array().unwrap().iter().map(|v| v.as_i64().unwrap()).collect();
|
||||
let rows: Vec<Vec<i64>> = j["pairs"]
|
||||
.as_array()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.map(|p| {
|
||||
let (a, b) = (ACTIONS[p[0].as_u64().unwrap() as usize], ACTIONS[p[1].as_u64().unwrap() as usize]);
|
||||
vec![a.0 * one, a.1 * one, b.0 * one, b.1 * one]
|
||||
})
|
||||
.collect();
|
||||
let frames = same_under_every_policy("int_world2", || decode_planes(&t, &gs, &z0, &rows, c, lat).frames);
|
||||
let want = j["frame_hashes"].as_array().unwrap();
|
||||
assert_eq!(frames.len(), want.len());
|
||||
for (f, w) in frames.iter().zip(want) {
|
||||
assert_eq!(format!("{:016x}", wai::int_hyper_synth::fnv1a(f)), w.as_str().unwrap());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn int_video_neural_keyframe_and_real_footage() {
|
||||
use wai::int_hyper_synth::{decode_conditional, 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")));
|
||||
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 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();
|
||||
decode_pframes(&f0, h, w, &vbs).unwrap().frames
|
||||
});
|
||||
let want = jload(&iv.join("video_nk.json"));
|
||||
let want = want["frame_hashes"].as_array().unwrap();
|
||||
assert_eq!(frames.len(), want.len(), "neural-keyframe video: frame count");
|
||||
for (f, w) in frames.iter().zip(want) {
|
||||
assert_eq!(format!("{:016x}", fnv1a(f)), w.as_str().unwrap());
|
||||
}
|
||||
// integer block-MC only (no learned kernels) — unaffected, checked for completeness
|
||||
let rbs = jload(&iv.join("video_real_bitstream.json"));
|
||||
let frames = same_under_every_policy("real-footage video", || decode(&rbs).unwrap().frames);
|
||||
let want = jload(&iv.join("video_real.json"));
|
||||
let want = want["frame_hashes"].as_array().unwrap();
|
||||
assert_eq!(frames.len(), want.len(), "real-footage video: frame count");
|
||||
for (f, w) in frames.iter().zip(want) {
|
||||
assert_eq!(format!("{:016x}", fnv1a(f)), w.as_str().unwrap());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn neural_conformance_int_synth_corpus() {
|
||||
use wai::int_synth::decode_from_spec;
|
||||
let dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../neural-conformance");
|
||||
let manifest = jload(&dir.join("expected_hashes.json"));
|
||||
let mut n = 0;
|
||||
for (name, spec) in manifest.as_object().unwrap() {
|
||||
let vd = dir.join("vectors").join(name);
|
||||
let payload = std::fs::read(vd.join("latent.bin")).unwrap();
|
||||
let model = std::fs::read(vd.join("int_synth.json")).unwrap();
|
||||
let rgb = same_under_every_policy(name, || decode_from_spec(&payload, &model).unwrap().rgb);
|
||||
let mut h = blake3::Hasher::new();
|
||||
h.update(b"wai:int-synth-rgb\x01");
|
||||
h.update(&rgb);
|
||||
let got: String = h.finalize().as_bytes().iter().map(|b| format!("{b:02x}")).collect();
|
||||
assert_eq!(got, spec["rgb_hash"].as_str().unwrap(), "{name}");
|
||||
if let Ok(r) = std::fs::read(vd.join("ref_rgb.bin")) {
|
||||
assert_eq!(rgb, r, "{name}: must equal the python integer-forward reference");
|
||||
}
|
||||
n += 1;
|
||||
}
|
||||
assert!(n >= 4, "expected the four committed int_synth vectors, found {n}");
|
||||
}
|
||||
|
|
@ -1,12 +1,19 @@
|
|||
//! `wai.video.int_hyper` at SOTA scale — a real 256×256 clip whose I-frame is the
|
||||
//! full SOTA `int_hyper` learned codec (42.6 dB photo) decoded byte-exact, panned
|
||||
//! into P-frames by integer motion compensation. Reads the SOTA int_hyper vector +
|
||||
//! the panned-clip bitstream from `WAI_SOTA_INT_DIR` (default `/tmp/sota_int`);
|
||||
//! skips if absent (the SOTA model is too large to commit, like the int_hyper SOTA
|
||||
//! test). The committed 16×16 `int_video_keyframe` test is the never-skip companion.
|
||||
//! `wai.video.int_hyper` on the full pretrained model — a real 256×256 clip whose
|
||||
//! I-frame is the byte-exact `int_hyper` decode of the pretrained mbt2018_mean
|
||||
//! (exported to integers), panned into P-frames by integer motion compensation.
|
||||
//! Reads the int_hyper vector + the panned-clip bitstream from `WAI_SOTA_INT_DIR`
|
||||
//! (default `/tmp/sota_int`), regenerated by `tools/wai_sota_vectors.sh` (the
|
||||
//! model is too large to commit). CI (`wai-sota-vectors.yml`) regenerates them and
|
||||
//! runs this with `WAI_REQUIRE_SOTA_VECTORS=1`, so an absent vector fails there;
|
||||
//! locally it skips loudly. A `WAI_SOTA_INT_DIR` naming a directory that does not
|
||||
//! exist, or one missing a file this test reads, fails either way. The committed
|
||||
//! 16×16 `int_video_keyframe` test is the never-skip companion.
|
||||
#![cfg(all(feature = "neural_int", feature = "video_int"))]
|
||||
|
||||
use std::path::PathBuf;
|
||||
#[path = "support/require_vectors.rs"]
|
||||
mod require_vectors;
|
||||
|
||||
use require_vectors::{env_dir, load_all};
|
||||
|
||||
use wai::int_hyper_synth::{decode_conditional, CondBitstream, CondEntropy, IntModel};
|
||||
use wai::int_video::{decode_pframes, fnv1a};
|
||||
|
|
@ -21,19 +28,22 @@ fn hex16(s: &str) -> [u8; 8] {
|
|||
|
||||
#[test]
|
||||
fn sota_int_hyper_keyframe_clip_decodes_byte_for_byte() {
|
||||
let d = PathBuf::from(std::env::var("WAI_SOTA_INT_DIR").unwrap_or_else(|_| "/tmp/sota_int".into()));
|
||||
let jl = |n: &str| -> Option<serde_json::Value> {
|
||||
serde_json::from_slice(&std::fs::read(d.join(n)).ok()?).ok()
|
||||
};
|
||||
let (Some(g), Some(hh), Some(e), Some(b), Some(vb), Some(conf)) = (
|
||||
jl("int_synth_sota_model.json"), jl("int_hs_model.json"), jl("cond_entropy.json"),
|
||||
jl("cond_bitstream.json"), jl("video_sota_bitstream.json"), jl("video_sota.json"),
|
||||
const TEST: &str = "int_video_sota_keyframe::sota_int_hyper_keyframe_clip_decodes_byte_for_byte";
|
||||
const PRODUCE: &str = "tools/wai_sota_vectors.sh <dir> (then WAI_SOTA_INT_DIR=<dir>, default /tmp/sota_int)";
|
||||
let d = env_dir(TEST, "WAI_SOTA_INT_DIR", "/tmp/sota_int", PRODUCE);
|
||||
let Some([g, hh, e, b, vb, conf]) = load_all(
|
||||
TEST,
|
||||
&d,
|
||||
[
|
||||
"int_synth_sota_model.json", "int_hs_model.json", "cond_entropy.json",
|
||||
"cond_bitstream.json", "video_sota_bitstream.json", "video_sota.json",
|
||||
],
|
||||
PRODUCE,
|
||||
) else {
|
||||
eprintln!("skipping SOTA int_hyper-keyframe video: artifacts not in {} (run the export, the int_hyper_decode_frame example, then tools/wai_int_video_sota_keyframe_encode.py)", d.display());
|
||||
return;
|
||||
};
|
||||
|
||||
// frame 0: the SOTA int_hyper learned I-frame (byte-exact 256×256 photo).
|
||||
// frame 0: the int_hyper learned I-frame (byte-exact 256×256 photo).
|
||||
let g_s = IntModel::from_json(&g);
|
||||
let h_s = IntModel::from_json(&hh);
|
||||
let ent = CondEntropy::from_value(&e);
|
||||
|
|
@ -44,7 +54,7 @@ fn sota_int_hyper_keyframe_clip_decodes_byte_for_byte() {
|
|||
let clip = decode_pframes(&frame0, h, w, &vb).expect("p-frames");
|
||||
let want = conf["frame_hashes"].as_array().unwrap();
|
||||
assert_eq!(clip.frames.len(), want.len());
|
||||
assert_eq!(format!("{:016x}", fnv1a(&frame0)), want[0].as_str().unwrap(), "I-frame = SOTA int_hyper decode");
|
||||
assert_eq!(format!("{:016x}", fnv1a(&frame0)), want[0].as_str().unwrap(), "I-frame = int_hyper decode");
|
||||
|
||||
let mut all = Vec::new();
|
||||
for (i, frame) in clip.frames.iter().enumerate() {
|
||||
|
|
@ -53,6 +63,6 @@ fn sota_int_hyper_keyframe_clip_decodes_byte_for_byte() {
|
|||
all.extend_from_slice(&hex16(&got));
|
||||
}
|
||||
assert_eq!(format!("{:016x}", fnv1a(&all)), conf["clip_hash"].as_str().unwrap(), "clip hash");
|
||||
println!("SOTA int_hyper-keyframe video: {} frames {h}x{w}, I-frame={} (learned), clip {}",
|
||||
println!("int_hyper-keyframe video (pretrained model): {} frames {h}x{w}, I-frame={} (learned), clip {}",
|
||||
clip.frames.len(), &want[0].as_str().unwrap()[..8], conf["clip_hash"].as_str().unwrap());
|
||||
}
|
||||
|
|
|
|||
|
|
@ -25,8 +25,21 @@
|
|||
//! WAI_JPEGAI_WEIGHTS_DIR = dir of {hs,hyp}_{y,uv}.json, mcm{0..3}_y.json,
|
||||
//! synth_{y,uv}.json, icci.json
|
||||
//! WAI_JPEGAI_NORM = norm.json (see wai_jpegai_build_e2e_inputs.py)
|
||||
//!
|
||||
//! These inputs are not in the repository. They were extracted by
|
||||
//! `tools/wai_jpegai_build_e2e_inputs.py` from JPEG-AI reference-software assets
|
||||
//! (ISO/IEC 6048-1; kept locally in `wai/jpegai-vm-assets`, not committed), and the
|
||||
//! extraction is published under WAI's versioned models root (`jpegai/v1/`).
|
||||
//! `tools/wai_jpegai_fetch_e2e_inputs.sh` fetches it and checks every file against
|
||||
//! the SHA-256 manifest `tools/jpegai-e2e-v1.sha256`; CI (`wai-jpegai-e2e.yml`) does
|
||||
//! that and runs this file with `WAI_REQUIRE_SOTA_VECTORS=1`. Without the three
|
||||
//! variables it skips loudly; with the switch it fails instead; a variable naming a
|
||||
//! missing path, or a weights directory missing a file, fails either way.
|
||||
#![cfg(feature = "neural_int")]
|
||||
|
||||
#[path = "support/require_vectors.rs"]
|
||||
mod require_vectors;
|
||||
|
||||
use wai::jpegai_decode::{decode, DecoderWeights, NormTables, PlaneWeights};
|
||||
use wai::jpegai_icci::{Conv, IcciNet, IcciPath, ResBlock};
|
||||
use wai::jpegai_synthesis::{IntNet, Weights};
|
||||
|
|
@ -80,12 +93,14 @@ fn path_of(v: &serde_json::Value) -> IcciPath {
|
|||
|
||||
#[test]
|
||||
fn codestream_to_rgb_matches_reference_png() {
|
||||
let (Ok(cs), Ok(dir), Ok(nf)) = (std::env::var("WAI_JPEGAI_CODESTREAM"),
|
||||
std::env::var("WAI_JPEGAI_WEIGHTS_DIR"), std::env::var("WAI_JPEGAI_NORM")) else {
|
||||
eprintln!("skip: set WAI_JPEGAI_CODESTREAM + WAI_JPEGAI_WEIGHTS_DIR + WAI_JPEGAI_NORM");
|
||||
let Some([cs, dir, nf]) = require_vectors::env_paths(
|
||||
"jpegai_e2e_full::codestream_to_rgb_matches_reference_png",
|
||||
["WAI_JPEGAI_CODESTREAM", "WAI_JPEGAI_WEIGHTS_DIR", "WAI_JPEGAI_NORM"],
|
||||
"tools/wai_jpegai_fetch_e2e_inputs.sh <dir> (pinned by tools/jpegai-e2e-v1.sha256)",
|
||||
) else {
|
||||
return;
|
||||
};
|
||||
let Ok(data) = std::fs::read(&cs) else { eprintln!("skip: no codestream"); return };
|
||||
let data = std::fs::read(&cs).unwrap_or_else(|e| panic!("read {cs}: {e}"));
|
||||
let nv = jload(std::path::Path::new(&nf));
|
||||
|
||||
let iw = jload(&std::path::Path::new(&dir).join("icci.json"));
|
||||
|
|
|
|||
|
|
@ -4,33 +4,62 @@
|
|||
//! rvs); the model bundle supplies only checkpoint constants. Compared against
|
||||
//! the reference decoder's PNG — the same bar as the staged `jpegai_e2e` test,
|
||||
//! but through the dispatchable library entry point.
|
||||
//!
|
||||
//! This is the whole-decode test behind the dispatchable `wai.image.jpegai` ONNX
|
||||
//! path, so a missing input follows the headline-claim rule (support/
|
||||
//! require_vectors.rs): without its inputs it skips loudly, and with
|
||||
//! `WAI_REQUIRE_SOTA_VECTORS=1` it fails. Its inputs — the float model bundle, the
|
||||
//! traced nets under `model16/`, the extracted model constants and the reference
|
||||
//! PNG — exist only in the local JPEG-AI asset directory
|
||||
//! (`wai/jpegai-vm-assets`, exported from the ISO/IEC 6048-1 reference software;
|
||||
//! not in the repository), and the test also needs the `neural` (ONNX Runtime) and
|
||||
//! `codecs` features. No CI job runs it; it is verified only on a machine holding
|
||||
//! those assets. `WAI_JPEGAI_ASSETS` points it at the directory from elsewhere; a
|
||||
//! directory it names must exist and hold every input, or the test fails with or
|
||||
//! without the switch. Unset or empty, the default below is used, and only that
|
||||
//! default may be absent and skip.
|
||||
#![cfg(all(feature = "neural", feature = "codecs"))]
|
||||
|
||||
use std::path::PathBuf;
|
||||
#[path = "support/require_vectors.rs"]
|
||||
mod require_vectors;
|
||||
|
||||
use require_vectors::InputDir;
|
||||
use wai::neural::jpegai::{decode, JpegAiModel};
|
||||
|
||||
fn assets() -> PathBuf {
|
||||
PathBuf::from("../jpegai-vm-assets")
|
||||
const TEST: &str = "jpegai_live::jpegai_live_decode_matches_reference";
|
||||
const PRODUCE: &str = "the local JPEG-AI asset directory wai/jpegai-vm-assets (exported from the \
|
||||
ISO/IEC 6048-1 reference software; not in the repository), or WAI_JPEGAI_ASSETS=<dir>";
|
||||
|
||||
/// `WAI_JPEGAI_ASSETS`, else `../jpegai-vm-assets` beside the crate. A variable that
|
||||
/// names a directory that does not exist, or one missing an input, is a
|
||||
/// misconfigured run and fails.
|
||||
fn assets() -> InputDir {
|
||||
require_vectors::env_dir(TEST, "WAI_JPEGAI_ASSETS", "../jpegai-vm-assets", PRODUCE)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn jpegai_live_decode_matches_reference() {
|
||||
let a = assets();
|
||||
// Needs the model bundle (incl. the extracted model-constant files) + the
|
||||
// 256×256 traced nets + the reference PNG. Skipped if archived out-of-band.
|
||||
for need in [
|
||||
// 256×256 traced nets + the reference PNG.
|
||||
let missing: Vec<&str> = [
|
||||
"e2e/test256.bin",
|
||||
"e2e/test256_dec.png",
|
||||
"jpegai_factorized_cdfs.json",
|
||||
"jpegai_scaler_vec.json",
|
||||
"jpegai_dequant_consts.json",
|
||||
"model16/CCS_SGMM",
|
||||
] {
|
||||
if !a.join(need).exists() {
|
||||
eprintln!("skipping jpegai live decode: {need} not present");
|
||||
return;
|
||||
}
|
||||
]
|
||||
.into_iter()
|
||||
.filter(|need| !a.join(need).exists())
|
||||
.collect();
|
||||
if !missing.is_empty() {
|
||||
a.absent(
|
||||
TEST,
|
||||
&format!("{} not in {}", missing.join(", "), a.display()),
|
||||
PRODUCE,
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
let model = JpegAiModel::load(&a).expect("load jpeg-ai model bundle");
|
||||
|
|
|
|||
|
|
@ -6,14 +6,36 @@
|
|||
//! It exercises the property the binding rests on: the exact decoder is deterministic,
|
||||
//! so the binding taken over one decode verifies a second decode, and a single flipped
|
||||
//! pixel breaks it. Both decodes run in this process; that the same pixels come out on
|
||||
//! another machine rests on the conditions in `extensions/determinism-tiers.md` §5–§6,
|
||||
//! not on this test. (`wai.image.jpegai` is registered at entropy-consistency; the
|
||||
//! pinned reconstruction is a property of the exact decoder.)
|
||||
//! another machine rests on the conditions and measurements in
|
||||
//! `extensions/determinism-tiers.md` §5–§6. This test compares with other machines
|
||||
//! only through `WAI_JPEGAI_EXPECT_SHA256` (below). (`wai.image.jpegai` is registered
|
||||
//! at entropy-consistency; the pinned reconstruction is a property of the exact
|
||||
//! decoder.)
|
||||
//!
|
||||
//! env: WAI_JPEGAI_CODESTREAM, WAI_JPEGAI_WEIGHTS_DIR, WAI_JPEGAI_NORM
|
||||
//! (the bundle from `tools/wai_jpegai_build_e2e_inputs.py`; see jpegai_e2e_full).
|
||||
//! (the bundle from `tools/wai_jpegai_fetch_e2e_inputs.sh`; see jpegai_e2e_full).
|
||||
//! WAI_JPEGAI_EXPECT_SHA256 (optional) = the SHA-256 the decoded pixels must
|
||||
//! have. Set, it also compares this machine's decode with a value decoded on
|
||||
//! other machines, which the two in-process decodes cannot do. That value rests
|
||||
//! on a measurement, not on the registry: the codec registry classes
|
||||
//! `wai.image.jpegai` as `NeuralEntropyExact`, a conservative class that
|
||||
//! guarantees only the entropy decode (the decoded symbols and their entropy
|
||||
//! parameters) to be byte-identical across machines, not the reconstruction.
|
||||
//! The exact decoder's pinned output for the `jpegai/v1` inputs was measured
|
||||
//! identical on four platforms: macOS arm64, wasm32-wasip2, Linux arm64, and
|
||||
//! Linux x86_64 under user-mode emulation (`extensions/determinism-tiers.md`
|
||||
//! §6). `wai-jpegai-e2e.yml` pins that value, so a mismatch there is a departure
|
||||
//! from that measurement.
|
||||
//!
|
||||
//! The inputs are not in the repository; see jpegai_e2e_full for where they come
|
||||
//! from and how CI fetches them. Without the three variables this test skips
|
||||
//! loudly; with `WAI_REQUIRE_SOTA_VECTORS=1` it fails instead; a variable naming a
|
||||
//! missing path, or a weights directory missing a file, fails either way.
|
||||
#![cfg(all(feature = "neural_int", feature = "provenance"))]
|
||||
|
||||
#[path = "support/require_vectors.rs"]
|
||||
mod require_vectors;
|
||||
|
||||
use wai::jpegai_decode::{decode, DecoderWeights, NormTables, PlaneWeights};
|
||||
use wai::jpegai_icci::IcciNet;
|
||||
use wai::jpegai_synthesis::{IntNet, Weights};
|
||||
|
|
@ -25,15 +47,14 @@ fn jload(p: &std::path::Path) -> serde_json::Value {
|
|||
|
||||
#[test]
|
||||
fn real_decode_to_reconstruction_binding() {
|
||||
let (Ok(cs), Ok(dir), Ok(nf)) = (
|
||||
std::env::var("WAI_JPEGAI_CODESTREAM"),
|
||||
std::env::var("WAI_JPEGAI_WEIGHTS_DIR"),
|
||||
std::env::var("WAI_JPEGAI_NORM"),
|
||||
let Some([cs, dir, nf]) = require_vectors::env_paths(
|
||||
"jpegai_reconstruction_binding::real_decode_to_reconstruction_binding",
|
||||
["WAI_JPEGAI_CODESTREAM", "WAI_JPEGAI_WEIGHTS_DIR", "WAI_JPEGAI_NORM"],
|
||||
"tools/wai_jpegai_fetch_e2e_inputs.sh <dir> (pinned by tools/jpegai-e2e-v1.sha256)",
|
||||
) else {
|
||||
eprintln!("skip: set WAI_JPEGAI_CODESTREAM + WAI_JPEGAI_WEIGHTS_DIR + WAI_JPEGAI_NORM");
|
||||
return;
|
||||
};
|
||||
let Ok(data) = std::fs::read(&cs) else { eprintln!("skip: no codestream"); return };
|
||||
let data = std::fs::read(&cs).unwrap_or_else(|e| panic!("read {cs}: {e}"));
|
||||
|
||||
// --- load the model bundle (same shape as jpegai_e2e_full) ---
|
||||
let dpath = std::path::Path::new(&dir);
|
||||
|
|
@ -102,4 +123,15 @@ fn real_decode_to_reconstruction_binding() {
|
|||
eprintln!(" assertion = {}", ReconstructionBinding::C2PA_LABEL);
|
||||
eprintln!(" sha256 = {}", binding.hash_hex());
|
||||
eprintln!(" json = {}", binding.to_json());
|
||||
|
||||
// Against other machines: the pixels this machine decoded must hash to the value
|
||||
// measured from other machines' decodes of the same inputs (see the header).
|
||||
if let Ok(want) = std::env::var("WAI_JPEGAI_EXPECT_SHA256") {
|
||||
assert_eq!(
|
||||
binding.hash_hex(),
|
||||
want.trim().to_ascii_lowercase(),
|
||||
"the decoded pixels differ from the pinned decode of these inputs"
|
||||
);
|
||||
eprintln!(" matches the pinned WAI_JPEGAI_EXPECT_SHA256");
|
||||
}
|
||||
}
|
||||
|
|
|
|||
300
wai/wai-rs/tests/require_vectors_rule.rs
Normal file
300
wai/wai-rs/tests/require_vectors_rule.rs
Normal file
|
|
@ -0,0 +1,300 @@
|
|||
//! The missing-input rule of `support/require_vectors.rs`, checked in the
|
||||
//! environment it is about.
|
||||
//!
|
||||
//! The vector tests' CI jobs always set `WAI_REQUIRE_SOTA_VECTORS=1`, so those runs
|
||||
//! never exercise what the helpers do WITHOUT the switch: that a variable naming a
|
||||
//! test's inputs still makes them required, and that a test skips only when no
|
||||
//! variable names what is missing (an unset variable, or an absent or incomplete
|
||||
//! default location). These tests check that rule directly. Each case starts this
|
||||
//! test binary again as a child process running the ignored `probe` test, with the
|
||||
//! switch and the input variables set exactly as the case needs (and every other
|
||||
//! one removed, so the parent's environment cannot leak in), then checks whether
|
||||
//! the child passed or failed and what it printed.
|
||||
//!
|
||||
//! Their inputs are scratch directories and small files that each case creates
|
||||
//! under the system temporary directory, so these tests never skip.
|
||||
|
||||
#[path = "support/require_vectors.rs"]
|
||||
mod require_vectors;
|
||||
|
||||
use std::path::PathBuf;
|
||||
use std::process::Command;
|
||||
|
||||
/// Which helper the probe runs: `dir` (`env_dir` + `load_all`, as the int_hyper
|
||||
/// tests use them) or `paths` (`env_paths`, as the JPEG-AI tests use it).
|
||||
const MODE: &str = "WAI_RV_PROBE";
|
||||
/// The default directory the probe passes to `env_dir`.
|
||||
const DEFAULT: &str = "WAI_RV_PROBE_DEFAULT";
|
||||
/// The directory variable the probe reads, as `int_hyper_synth` does.
|
||||
const DIR_VAR: &str = "WAI_SOTA_INT_DIR";
|
||||
/// The two path variables the probe reads, as the JPEG-AI tests read three.
|
||||
const PATH_VAR: &str = "WAI_JPEGAI_CODESTREAM";
|
||||
const PATH_VAR2: &str = "WAI_JPEGAI_NORM";
|
||||
/// The one input file the `dir` probe needs.
|
||||
const INPUT: &str = "cond.json";
|
||||
|
||||
/// Runs one helper the way a vector test does and prints what it returned. It is
|
||||
/// only meaningful in a child started by [`run`], with the environment [`run`]
|
||||
/// gives it; run any other way, it returns without doing anything.
|
||||
#[test]
|
||||
#[ignore = "a probe that the other tests in this file run in a child process"]
|
||||
fn probe() {
|
||||
let Ok(mode) = std::env::var(MODE) else {
|
||||
return;
|
||||
};
|
||||
let loaded = match mode.as_str() {
|
||||
"dir" => {
|
||||
let default = std::env::var(DEFAULT).expect("the default directory");
|
||||
let d = require_vectors::env_dir("probe", DIR_VAR, &default, "the probe");
|
||||
require_vectors::load_all("probe", &d, [INPUT], "the probe").is_some()
|
||||
}
|
||||
"paths" => require_vectors::env_paths("probe", [PATH_VAR, PATH_VAR2], "the probe").is_some(),
|
||||
other => panic!("unknown probe mode {other}"),
|
||||
};
|
||||
println!(
|
||||
"probe: {}",
|
||||
if loaded {
|
||||
"inputs loaded"
|
||||
} else {
|
||||
"returned without inputs"
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
/// What a probe run did.
|
||||
struct Outcome {
|
||||
passed: bool,
|
||||
/// The child's stdout followed by its stderr.
|
||||
out: String,
|
||||
}
|
||||
|
||||
impl Outcome {
|
||||
/// The probe passed because it found nothing to check: the loud skip.
|
||||
fn assert_skipped_loudly(&self) {
|
||||
assert!(self.passed, "expected a skip, got a failure:\n{}", self.out);
|
||||
assert!(
|
||||
self.out.contains("probe: returned without inputs"),
|
||||
"{}",
|
||||
self.out
|
||||
);
|
||||
// The banner CI's backstop looks for, so a skip is never read as a pass.
|
||||
assert!(
|
||||
self.out.contains("SKIPPED") && self.out.contains("NOT VERIFIED"),
|
||||
"{}",
|
||||
self.out
|
||||
);
|
||||
}
|
||||
|
||||
fn assert_failed_with(&self, why: &str) {
|
||||
assert!(
|
||||
!self.passed,
|
||||
"expected a failure naming {why:?}, got a pass:\n{}",
|
||||
self.out
|
||||
);
|
||||
assert!(
|
||||
self.out.contains(why),
|
||||
"the failure does not say {why:?}:\n{}",
|
||||
self.out
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Runs the probe in a child process in `mode`, with `env` set and the switch and
|
||||
/// every input variable otherwise removed.
|
||||
fn run(mode: &str, env: &[(&str, &str)]) -> Outcome {
|
||||
let exe = std::env::current_exe().expect("this test binary");
|
||||
let mut cmd = Command::new(exe);
|
||||
cmd.args([
|
||||
"probe",
|
||||
"--exact",
|
||||
"--ignored",
|
||||
"--nocapture",
|
||||
"--test-threads=1",
|
||||
]);
|
||||
for v in [require_vectors::SWITCH, MODE, DEFAULT, DIR_VAR, PATH_VAR, PATH_VAR2] {
|
||||
cmd.env_remove(v);
|
||||
}
|
||||
cmd.env(MODE, mode);
|
||||
for (k, v) in env {
|
||||
cmd.env(k, v);
|
||||
}
|
||||
let o = cmd.output().expect("start the probe");
|
||||
let out = format!(
|
||||
"{}{}",
|
||||
String::from_utf8_lossy(&o.stdout),
|
||||
String::from_utf8_lossy(&o.stderr)
|
||||
);
|
||||
// A filter that matched nothing would also exit 0.
|
||||
assert!(
|
||||
out.contains("running 1 test"),
|
||||
"the probe did not run:\n{out}"
|
||||
);
|
||||
Outcome {
|
||||
passed: o.status.success(),
|
||||
out,
|
||||
}
|
||||
}
|
||||
|
||||
/// A scratch directory unique to one test case, removed when dropped.
|
||||
struct Scratch(PathBuf);
|
||||
|
||||
impl Scratch {
|
||||
fn new(case: &str) -> Scratch {
|
||||
let p = std::env::temp_dir().join(format!(
|
||||
"wai-require-vectors-rule-{}-{case}",
|
||||
std::process::id()
|
||||
));
|
||||
let _ = std::fs::remove_dir_all(&p);
|
||||
std::fs::create_dir_all(&p).expect("create the scratch directory");
|
||||
Scratch(p)
|
||||
}
|
||||
|
||||
/// A path inside the scratch directory, not created.
|
||||
fn at(&self, name: &str) -> String {
|
||||
self.0
|
||||
.join(name)
|
||||
.to_str()
|
||||
.expect("a UTF-8 temp path")
|
||||
.to_owned()
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for Scratch {
|
||||
fn drop(&mut self) {
|
||||
let _ = std::fs::remove_dir_all(&self.0);
|
||||
}
|
||||
}
|
||||
|
||||
// ---- without the switch ----
|
||||
|
||||
#[test]
|
||||
fn without_the_switch_a_named_directory_that_does_not_exist_fails() {
|
||||
let s = Scratch::new("dir-missing");
|
||||
let (named, default) = (s.at("not-there"), s.at("default"));
|
||||
run("dir", &[(DIR_VAR, &named), (DEFAULT, &default)]).assert_failed_with("is not a directory");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn without_the_switch_a_named_path_that_is_a_file_fails_as_a_directory() {
|
||||
let s = Scratch::new("dir-is-file");
|
||||
let (named, default) = (s.at("a-file"), s.at("default"));
|
||||
std::fs::write(&named, b"{}").unwrap();
|
||||
run("dir", &[(DIR_VAR, &named), (DEFAULT, &default)]).assert_failed_with("is not a directory");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn without_the_switch_a_named_empty_directory_fails() {
|
||||
// The case this file was written for: the directory exists, so only the
|
||||
// missing file can fail the run.
|
||||
let s = Scratch::new("dir-empty");
|
||||
let (named, default) = (s.at("vectors"), s.at("default"));
|
||||
std::fs::create_dir(&named).unwrap();
|
||||
run("dir", &[(DIR_VAR, &named), (DEFAULT, &default)])
|
||||
.assert_failed_with(&format!("{DIR_VAR} names this directory"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn without_the_switch_an_unset_or_empty_directory_variable_skips_on_an_absent_default() {
|
||||
let s = Scratch::new("dir-default");
|
||||
let default = s.at("default");
|
||||
run("dir", &[(DEFAULT, &default)]).assert_skipped_loudly();
|
||||
run("dir", &[(DIR_VAR, ""), (DEFAULT, &default)]).assert_skipped_loudly();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn without_the_switch_an_incomplete_default_directory_skips() {
|
||||
let s = Scratch::new("dir-default-empty");
|
||||
let default = s.at("default");
|
||||
std::fs::create_dir(&default).unwrap();
|
||||
run("dir", &[(DEFAULT, &default)]).assert_skipped_loudly();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_named_directory_holding_the_input_loads_it() {
|
||||
// The positive control: without it, a probe that always failed would make the
|
||||
// failure cases above pass for the wrong reason.
|
||||
let s = Scratch::new("dir-full");
|
||||
let (named, default) = (s.at("vectors"), s.at("default"));
|
||||
std::fs::create_dir(&named).unwrap();
|
||||
std::fs::write(PathBuf::from(&named).join(INPUT), b"{}").unwrap();
|
||||
let o = run("dir", &[(DIR_VAR, &named), (DEFAULT, &default)]);
|
||||
assert!(
|
||||
o.passed && o.out.contains("probe: inputs loaded"),
|
||||
"{}",
|
||||
o.out
|
||||
);
|
||||
assert!(!o.out.contains("NOT VERIFIED"), "{}", o.out);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn without_the_switch_a_path_variable_naming_nothing_fails() {
|
||||
let s = Scratch::new("path-missing");
|
||||
let other = s.at("norm.json");
|
||||
std::fs::write(&other, b"{}").unwrap();
|
||||
run("paths", &[(PATH_VAR, &s.at("not-there")), (PATH_VAR2, &other)])
|
||||
.assert_failed_with("does not exist");
|
||||
// env_paths has no default: an empty value names no existing path.
|
||||
run("paths", &[(PATH_VAR, ""), (PATH_VAR2, &other)]).assert_failed_with("does not exist");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn without_the_switch_a_path_variable_naming_nothing_fails_when_another_is_unset() {
|
||||
// The other variable being unset would end the test as a skip; the one that
|
||||
// names a missing path must fail it first.
|
||||
let s = Scratch::new("path-missing-partial");
|
||||
run("paths", &[(PATH_VAR, &s.at("not-there"))]).assert_failed_with("does not exist");
|
||||
run("paths", &[(PATH_VAR2, &s.at("not-there"))]).assert_failed_with("does not exist");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn without_the_switch_an_unset_path_variable_skips() {
|
||||
run("paths", &[]).assert_skipped_loudly();
|
||||
// One set to a file that exists, the other unset: nothing named is missing.
|
||||
let s = Scratch::new("path-partial");
|
||||
let file = s.at("codestream.bin");
|
||||
std::fs::write(&file, b"x").unwrap();
|
||||
run("paths", &[(PATH_VAR, &file)]).assert_skipped_loudly();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn path_variables_naming_files_that_exist_load_them() {
|
||||
// The positive control for `env_paths`, as the directory one is for `env_dir`.
|
||||
let s = Scratch::new("path-full");
|
||||
let (a, b) = (s.at("codestream.bin"), s.at("norm.json"));
|
||||
std::fs::write(&a, b"x").unwrap();
|
||||
std::fs::write(&b, b"{}").unwrap();
|
||||
let o = run("paths", &[(PATH_VAR, &a), (PATH_VAR2, &b)]);
|
||||
assert!(
|
||||
o.passed && o.out.contains("probe: inputs loaded"),
|
||||
"{}",
|
||||
o.out
|
||||
);
|
||||
assert!(!o.out.contains("NOT VERIFIED"), "{}", o.out);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn switch_values_that_leave_it_off_behave_as_unset() {
|
||||
let s = Scratch::new("switch-off");
|
||||
let default = s.at("default");
|
||||
for off in ["", "0", "false", "no", "FALSE", " no "] {
|
||||
run(
|
||||
"dir",
|
||||
&[(require_vectors::SWITCH, off), (DEFAULT, &default)],
|
||||
)
|
||||
.assert_skipped_loudly();
|
||||
}
|
||||
}
|
||||
|
||||
// ---- with the switch ----
|
||||
|
||||
#[test]
|
||||
fn with_the_switch_an_absent_default_fails() {
|
||||
let s = Scratch::new("switch-on");
|
||||
let default = s.at("default");
|
||||
for on in ["1", "true", "yes"] {
|
||||
run("dir", &[(require_vectors::SWITCH, on), (DEFAULT, &default)])
|
||||
.assert_failed_with("a missing input is a failure, not a skip");
|
||||
run("paths", &[(require_vectors::SWITCH, on)])
|
||||
.assert_failed_with("a missing input is a failure, not a skip");
|
||||
}
|
||||
}
|
||||
204
wai/wai-rs/tests/support/require_vectors.rs
Normal file
204
wai/wai-rs/tests/support/require_vectors.rs
Normal file
|
|
@ -0,0 +1,204 @@
|
|||
//! What a test does when the out-of-tree inputs behind a headline decode claim are
|
||||
//! absent.
|
||||
//!
|
||||
//! Some decode claims are tested against inputs too large, or not redistributable,
|
||||
//! to commit: the `int_hyper` vectors (the pretrained mbt2018_mean, exported to
|
||||
//! integers) and the JPEG-AI whole-decoder bundle. A test that `return`s when its
|
||||
//! inputs are missing is reported as *passed*, so a run with nothing to check looks
|
||||
//! exactly like a run that checked everything. These helpers remove that ambiguity:
|
||||
//!
|
||||
//! - `WAI_REQUIRE_SOTA_VECTORS` unset (or `0`/`false`/`no`): the test still skips,
|
||||
//! but writes a banner to the process's real stderr — past libtest's output
|
||||
//! capture, so it shows in an ordinary `cargo test` run, not only under
|
||||
//! `--nocapture`.
|
||||
//! - `WAI_REQUIRE_SOTA_VECTORS=1`: a missing input is a test failure naming the
|
||||
//! tool that produces it. A run that claims to verify these decodes sets it.
|
||||
//!
|
||||
//! Either way, with or without the switch, a variable that names a test's inputs
|
||||
//! makes them required: pointing a test at inputs and having it pass because they
|
||||
//! are not there is the failure this module exists to prevent. So the test fails
|
||||
//! when a variable read through [`env_paths`] names a path that does not exist,
|
||||
//! and when a directory variable read through [`env_dir`] (`WAI_SOTA_INT_DIR`,
|
||||
//! `WAI_JPEGAI_ASSETS`) names something that is not a directory, or a directory
|
||||
//! missing a file the test reads ([`InputDir::absent`]). The JPEG-AI tests that
|
||||
//! take a directory through [`env_paths`] fail the same way on a missing file,
|
||||
//! because their loaders do not skip. So without the switch a test skips only when
|
||||
//! no variable names what is missing: when a variable read through [`env_paths`]
|
||||
//! is unset (the ones that are set must still name existing paths), or when the
|
||||
//! variable of [`env_dir`] is unset or empty and the default location it falls
|
||||
//! back to is absent or incomplete. [`env_paths`] has no default, so an empty value
|
||||
//! there names no existing path and fails. `tests/require_vectors_rule.rs` checks
|
||||
//! this rule by running these helpers in child processes, with the switch unset and
|
||||
//! set.
|
||||
//!
|
||||
//! Where each covered input comes from, and who verifies it:
|
||||
//!
|
||||
//! | test | inputs | produced by | verified in CI |
|
||||
//! |---|---|---|---|
|
||||
//! | `int_hyper_synth` (pretrained-model cases) | `WAI_SOTA_INT_DIR` | `tools/wai_sota_vectors.sh` | yes — `wai-sota-vectors.yml` regenerates and requires them |
|
||||
//! | `int_video_sota_keyframe` | `WAI_SOTA_INT_DIR` | `tools/wai_sota_vectors.sh` | yes — same job |
|
||||
//! | `jpegai_e2e_full` | `WAI_JPEGAI_{CODESTREAM,WEIGHTS_DIR,NORM}` | `tools/wai_jpegai_fetch_e2e_inputs.sh` (the models root's `jpegai/v1/`, pinned by `tools/jpegai-e2e-v1.sha256`) | yes — `wai-jpegai-e2e.yml` fetches, checks and requires them |
|
||||
//! | `jpegai_reconstruction_binding` | same, plus optional `WAI_JPEGAI_EXPECT_SHA256` | same | yes — same job, with the decoded-pixel hash pinned |
|
||||
//! | `jpegai_live` | `WAI_JPEGAI_ASSETS` (default `../jpegai-vm-assets`) | the local JPEG-AI asset directory (not in the repository) | no — local only; also needs the `neural` + `codecs` features |
|
||||
//!
|
||||
//! Tests that still skip silently on absent inputs and are NOT covered by this
|
||||
//! switch — an open gap, listed here so it stays visible:
|
||||
//! `jpegai_*_calibration.rs`, `jpegai_e2e.rs`, `jpegai_e2e_{chain,chroma,dequant,
|
||||
//! frontend,hyperlatent,synthesis}.rs`, `jpegai_mans_golden.rs`,
|
||||
//! `jpegai_postfilter_chain.rs` (per-stage and staged JPEG-AI fixtures traced from
|
||||
//! the reference software), `neural_native_decode.rs` (float ONNX decoders, models
|
||||
//! and `/tmp` fixtures built per machine), and the `WAI_JPEGAI_*` unit tests in
|
||||
//! `src/jpegai_synthesis.rs` and `src/neural/mod.rs`.
|
||||
|
||||
use std::io::Write;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
pub const SWITCH: &str = "WAI_REQUIRE_SOTA_VECTORS";
|
||||
|
||||
/// Whether this run has declared that the out-of-tree inputs must be present.
|
||||
pub fn required() -> bool {
|
||||
match std::env::var(SWITCH) {
|
||||
Ok(v) => !matches!(v.trim().to_ascii_lowercase().as_str(), "" | "0" | "false" | "no"),
|
||||
Err(_) => false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Call at the point a test would otherwise return early for lack of inputs.
|
||||
/// Panics when [`required`]; otherwise announces the skip where it cannot be
|
||||
/// missed. The caller returns (or skips the sub-check) after this.
|
||||
pub fn absent(test: &str, missing: &str, produce: &str) {
|
||||
let msg = format!("{test}: NOT VERIFIED — {missing}\n produce with: {produce}");
|
||||
if required() {
|
||||
panic!("{msg}\n {SWITCH} is set, so a missing input is a failure, not a skip");
|
||||
}
|
||||
// Deliberately not eprintln!: libtest captures the print macros of a passing
|
||||
// test, which is how skips became invisible in the first place.
|
||||
// wai-sota-vectors.yml and wai-jpegai-e2e.yml fail on "SKIPPED" or
|
||||
// "NOT VERIFIED" in their test log, so keep both words in this banner.
|
||||
let _ = writeln!(
|
||||
std::io::stderr(),
|
||||
"\n!!! SKIPPED — {msg}\n set {SWITCH}=1 to make a missing input fail\n"
|
||||
);
|
||||
}
|
||||
|
||||
/// Read a JSON input. `None` only when the file does not exist; a file that exists
|
||||
/// but cannot be read or parsed is a broken input and fails the test either way.
|
||||
#[allow(dead_code)]
|
||||
pub fn load_json(p: &Path) -> Option<serde_json::Value> {
|
||||
if !p.exists() {
|
||||
return None;
|
||||
}
|
||||
let bytes = std::fs::read(p).unwrap_or_else(|e| panic!("read {}: {e}", p.display()));
|
||||
Some(serde_json::from_slice(&bytes).unwrap_or_else(|e| panic!("parse {}: {e}", p.display())))
|
||||
}
|
||||
|
||||
/// Every named JSON file from `dir`, or `None` after [`InputDir::absent`] names
|
||||
/// the ones missing (which fails instead when a variable named `dir`).
|
||||
#[allow(dead_code)]
|
||||
pub fn load_all<const N: usize>(
|
||||
test: &str,
|
||||
dir: &InputDir,
|
||||
names: [&str; N],
|
||||
produce: &str,
|
||||
) -> Option<[serde_json::Value; N]> {
|
||||
let loaded = names.map(|n| load_json(&dir.join(n)));
|
||||
let missing: Vec<&str> =
|
||||
names.iter().zip(&loaded).filter(|(_, v)| v.is_none()).map(|(n, _)| *n).collect();
|
||||
if !missing.is_empty() {
|
||||
dir.absent(test, &format!("{} not in {}", missing.join(", "), dir.display()), produce);
|
||||
return None;
|
||||
}
|
||||
Some(loaded.map(Option::unwrap))
|
||||
}
|
||||
|
||||
/// The environment variables a test reads its input paths from. `None` (after
|
||||
/// [`absent`]) when any is unset. A variable that IS set but names a path that does
|
||||
/// not exist is a misconfigured run — pointing a test at a file and having it pass
|
||||
/// because the file is missing is the failure this module exists to prevent — so
|
||||
/// that fails regardless of the switch, and whether or not the other variables are
|
||||
/// set.
|
||||
#[allow(dead_code)]
|
||||
pub fn env_paths<const N: usize>(test: &str, vars: [&str; N], produce: &str) -> Option<[String; N]> {
|
||||
let got = vars.map(|v| std::env::var(v).ok());
|
||||
// The set ones first: a variable naming a missing path fails even when another
|
||||
// variable the test reads is unset, which would otherwise end the test as a skip.
|
||||
for (v, p) in vars.iter().zip(&got) {
|
||||
if let Some(p) = p {
|
||||
assert!(Path::new(p).exists(), "{test}: {v}={p} does not exist (produce with: {produce})");
|
||||
}
|
||||
}
|
||||
let unset: Vec<&str> =
|
||||
vars.iter().zip(&got).filter(|(_, g)| g.is_none()).map(|(v, _)| *v).collect();
|
||||
if !unset.is_empty() {
|
||||
absent(test, &format!("{} not set", unset.join(" + ")), produce);
|
||||
return None;
|
||||
}
|
||||
Some(got.map(Option::unwrap))
|
||||
}
|
||||
|
||||
/// The directory a test reads its inputs from, and whether a variable named it.
|
||||
/// Made by [`env_dir`].
|
||||
#[allow(dead_code)]
|
||||
pub struct InputDir {
|
||||
path: PathBuf,
|
||||
/// The variable that named `path`; `None` when `path` is the default.
|
||||
named_by: Option<String>,
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
impl InputDir {
|
||||
pub fn path(&self) -> &Path {
|
||||
&self.path
|
||||
}
|
||||
|
||||
pub fn join(&self, name: impl AsRef<Path>) -> PathBuf {
|
||||
self.path.join(name)
|
||||
}
|
||||
|
||||
pub fn display(&self) -> std::path::Display<'_> {
|
||||
self.path.display()
|
||||
}
|
||||
|
||||
/// [`absent`] for an input that belongs in this directory. When a variable
|
||||
/// named the directory, the run pointed the test at its inputs, so a missing
|
||||
/// one fails whether or not the switch is set. Only the default directory may
|
||||
/// be absent or incomplete and skip.
|
||||
pub fn absent(&self, test: &str, missing: &str, produce: &str) {
|
||||
if let Some(var) = &self.named_by {
|
||||
panic!(
|
||||
"{test}: NOT VERIFIED — {missing}\n produce with: {produce}\n \
|
||||
{var} names this directory, so a missing input is a failure with or \
|
||||
without {SWITCH}"
|
||||
);
|
||||
}
|
||||
absent(test, missing, produce);
|
||||
}
|
||||
}
|
||||
|
||||
impl AsRef<Path> for InputDir {
|
||||
fn as_ref(&self) -> &Path {
|
||||
&self.path
|
||||
}
|
||||
}
|
||||
|
||||
/// The directory a test reads its inputs from: the variable `var` when it is set
|
||||
/// and non-empty, else `default`. Empty counts as unset, as it does for
|
||||
/// `tools/wai_sota_vectors.sh`'s own `${1:-/tmp/sota_int}`, so the script and the
|
||||
/// tests fall back to the same place. A variable that IS set names the test's
|
||||
/// inputs: it fails the test, regardless of the switch, when it does not name a
|
||||
/// directory (as in [`env_paths`]), and the returned [`InputDir`] fails it when a
|
||||
/// file is missing from that directory. The default may be absent or incomplete:
|
||||
/// the test then goes through [`absent`].
|
||||
#[allow(dead_code)]
|
||||
pub fn env_dir(test: &str, var: &str, default: &str, produce: &str) -> InputDir {
|
||||
match std::env::var(var) {
|
||||
Ok(p) if !p.is_empty() => {
|
||||
assert!(
|
||||
Path::new(&p).is_dir(),
|
||||
"{test}: {var}={p} is not a directory (produce with: {produce})"
|
||||
);
|
||||
InputDir { path: PathBuf::from(p), named_by: Some(var.to_owned()) }
|
||||
}
|
||||
_ => InputDir { path: PathBuf::from(default), named_by: None },
|
||||
}
|
||||
}
|
||||
|
|
@ -14,5 +14,10 @@ edition = "2024"
|
|||
license = "Apache-2.0"
|
||||
publish = false
|
||||
|
||||
# `cfg(kani)` is set by the model checker on the included sources. Declared so an
|
||||
# ordinary build or clippy run of this crate does not warn on it.
|
||||
[lints.rust]
|
||||
unexpected_cfgs = { level = "warn", check-cfg = ["cfg(kani)"] }
|
||||
|
||||
[dependencies]
|
||||
# none — the integer core this crate verifies is pure std.
|
||||
|
|
|
|||
|
|
@ -30,6 +30,13 @@ pub mod int_entropy;
|
|||
#[path = "../../wai-rs/src/int_transform.rs"]
|
||||
pub mod int_transform;
|
||||
|
||||
/// The fast kernels `int_transform::igdn` dispatches to. Present so the included
|
||||
/// sources compile (only IGDN uses it here, hence the dead-code allowance); no
|
||||
/// harness reaches it, and it spawns no threads under the model checker.
|
||||
#[allow(dead_code)]
|
||||
#[path = "../../wai-rs/src/int_kernels.rs"]
|
||||
pub mod int_kernels;
|
||||
|
||||
/// The pure decision policy behind `wai.video.ssai`. Dependency-free by
|
||||
/// construction, which is what lets the checker reach it without stubbing the
|
||||
/// crypto layer — a stand-in would close the goal without deciding it.
|
||||
|
|
|
|||
|
|
@ -30,7 +30,7 @@
|
|||
no float anywhere</b>. The <i>whole</i> path is integer: factorized <code>z</code> → integer
|
||||
hyper-decoder → integer-σ entropy bucketing → integer synthesis. Because of that (the integer-entropy
|
||||
contract) the reconstruction is byte-identical on every machine, so the page re-checks its FNV-1a against
|
||||
the hash the native reference produced — a SOTA-class learned codec verified client-side, no server.
|
||||
the hash the native reference produced — a pretrained learned codec verified client-side, no server.
|
||||
</p>
|
||||
<canvas id="c" width="512" height="512"></canvas>
|
||||
<div id="badge">loading model…</div>
|
||||
|
|
@ -42,12 +42,13 @@
|
|||
<span>native hash <b class="hash" id="want">—</b></span>
|
||||
</div>
|
||||
<p class="note">
|
||||
<b>Why this is the moat:</b> the SOTA learned-image family is float, so its entropy decode desyncs on
|
||||
~1e-5 of cross-platform drift in the scale σ — reproducible only on the same build, never receiptable.
|
||||
<code>int_hyper</code> is a real pretrained codec (post-training-quantized to integer, ~0.6 dB off
|
||||
the float baseline) whose entire decode is integer, so the hash above is portable: this browser, a native
|
||||
sink, and the conformance corpus all produce the same bytes. The wire carried a ~1.4 KB bitstream;
|
||||
the model is the installed codec. <b>Note:</b> the model files are large (PTQ weights) and not committed —
|
||||
<b>Why the whole decode is integer:</b> a float learned codec's entropy decode can desync on ~1e-5 of
|
||||
cross-platform drift in the scale σ, so its output is reproducible only on the same build.
|
||||
<code>int_hyper</code> is the pretrained mbt2018_mean mean-scale hyperprior, post-training-quantized to
|
||||
integer (CI requires its decode within 1 dB of the same model's float decode on a pinned test
|
||||
photograph), and its entire decode is integer, so the hash above is portable: this browser, a native
|
||||
sink, and the conformance corpus all produce the same bytes. The wire carries only the bitstream (its
|
||||
size is shown above); the model is the installed codec. <b>Note:</b> the model files are large (PTQ weights) and not committed —
|
||||
generate them with <code>tools/wai_sota_int_synth_export.py</code> into <code>demo/int_hyper/</code>.
|
||||
</p>
|
||||
<script type="module" src="int_hyper.js"></script>
|
||||
|
|
|
|||
|
|
@ -3,7 +3,8 @@
|
|||
// Rust that decodes natively, compiled to wasm32, with NO ML runtime (no
|
||||
// onnxruntime-web, no WebGPU, no float). Then verify the decoded image's FNV-1a
|
||||
// against the native reference hash, in-page — the integer-entropy contract: a
|
||||
// SOTA-class learned codec whose reconstruction hash is portable.
|
||||
// pretrained learned codec (mbt2018_mean, exported to integers) whose
|
||||
// reconstruction hash is portable.
|
||||
//
|
||||
// The model files (g_s/h_s, PTQ weights) are large and NOT committed; place the
|
||||
// export tool's output under demo/int_hyper/ (the .bin models are ~½ the JSON).
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
// Reproducible proof: the wasm32 FULL CONDITIONAL codec decode is byte-identical
|
||||
// to the native reference — a SOTA-class learned image codec (PTQ of mbt2018_mean)
|
||||
// to the native reference — a pretrained learned image codec (PTQ of mbt2018_mean)
|
||||
// whose whole bitstream→image decode (factorized z, integer h_s, integer-σ
|
||||
// entropy, integer g_s) runs in the browser engine with NO ML runtime, and whose
|
||||
// RGB verifies against the native FNV-1a hash. The integer-entropy contract,
|
||||
|
|
@ -59,6 +59,6 @@ const bin = check('binary model', wai.decodeIntHyperBin(
|
|||
|
||||
const pass = json && bin;
|
||||
console.log(pass
|
||||
? '\n✓ wasm32 conditional-codec decode (JSON + compact-binary model) is byte-identical to the native reference — a SOTA learned codec decoded in-browser, zero JS ML runtime.'
|
||||
? '\n✓ wasm32 conditional-codec decode (JSON + compact-binary model) is byte-identical to the native reference — a pretrained learned codec decoded in-browser, zero JS ML runtime.'
|
||||
: '\n✗ mismatch');
|
||||
process.exit(pass ? 0 : 1);
|
||||
|
|
|
|||
|
|
@ -60,7 +60,7 @@ pub fn decode_int_synth(payload: &[u8], spec_json: &[u8]) -> Result<DecodedImage
|
|||
/// `cond_entropy.json` = the entropy tables); WAI carries the bitstream
|
||||
/// (`cond_bitstream.json`). The WHOLE decode — factorized z, integer h_s,
|
||||
/// integer-σ entropy, integer g_s — is pure i64, no float, no ONNX, so the RGB
|
||||
/// is byte-identical to the native reference on every machine: a SOTA-class
|
||||
/// is byte-identical to the native reference on every machine: a pretrained
|
||||
/// learned codec decoding deterministically in-page, zero JS ML runtime.
|
||||
#[wasm_bindgen(js_name = decodeIntHyper)]
|
||||
pub fn decode_int_hyper(
|
||||
|
|
|
|||
Loading…
Reference in a new issue