open-standards/wai/research/decode-throughput.md
Transaction Science a8bf1c560c wai: the int_codec vector's README, and the statements about what it was trained on
wai-rs/tests/vectors/int_audio/README.md states the vector's provenance (a
signal the export tool generates; no audio input), the seed (1), the
command, the tool versions, and that two runs gave byte-identical files.

research/decode-throughput.md keeps its text (a dated record) and gains a
dated correction: the int_audio vector it timed was trained on, and coded,
a third-party recording, and has been replaced by one trained and coded on a
generated signal of the same shape. wai_decode_bench's doc comment now lists
int_audio among the vectors trained on generated content.

Co-Authored-By: Claude Opus 5.5 <[email protected]>
2026-10-07 10:09:20 -04:00

469 lines
31 KiB
Markdown
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

# 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:** an arm64 macOS laptop [model removed 2026-09-30], 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.
> **Correction (2026-10-07).** The int_audio vector timed here was trained on,
> and coded, a third-party recording, not content made for the repository. It
> has been replaced: the int_audio model and vector are now trained and coded on
> a signal `tools/wai_int_audio_export.py` generates from a seed. The model has
> the same shape (16 channels, 512 latent samples, 8192 output samples); the
> row's figures are for the earlier vector, and the current one's time can
> differ through its activations (the exception above) and the length of its
> rANS stream.
## 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.