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

wai-ops/1-i32 on GPUs: WGSL kernels exact at acc32 and acc64, and the tier on the adapters a run showed equal (G1) #380

Merged
dcharlot merged 12 commits from wai/ops-gpu into main 2026-10-21 07:12:52 -04:00
Conversation 0 Commits 12 Files changed 37 +8695 -10
37 changed files with 8695 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

220
.github/workflows/wai-ops-gpu.yml vendored Normal file
View file

@ -0,0 +1,220 @@
name: wai-ops-gpu integer operators on GPUs
# wai-ops/1-i32 on GPUs (wai/wai-ops-gpu): the integer operator set's graph operators
# as WGSL kernels through wgpu, exact at acc32 (one 32-bit word per value) and at acc64
# (64- and 128-bit arithmetic built from 32-bit words), and the integer decoders with
# their graphs on the device (wai-rs `ops_gpu`).
#
# - metal (the `macos` runner, whose host has a GPU): the arithmetic probe on a
# million cases; every graph vector of wai/ops-conformance that the subset covers,
# at both widths, against the committed results and the interpreter; the shader
# mutants, every one of which the corpus must catch; every committed vector and
# refusal fixture of every capability that runs a wai-ops/1 graph, the jpegai/v1
# inputs included, against the decode on the CPU; and the corpus again in a
# headless browser's WebGPU, the kernels compiled by the browser's own WGSL
# compiler. WAI_OPS_GPU_EXPECT declares the one adapter the host must offer, and
# WAI_OPS_GPU_REQUIRE makes an absent adapter a failure, never a skip; each step
# then requires the report lines with their counts, so a filtered or empty run
# fails too.
# - checks (`linux`, a container with no GPU): the crate and its tests build, the
# device tests print SKIPPED and pass, and the same tests under
# WAI_OPS_GPU_REQUIRE=1 fail: an absent adapter is not a pass. The browser module
# builds for wasm32-unknown-unknown; the bench tool builds.
#
# The equality on other adapters (a discrete and an integrated GPU on Vulkan, an
# integrated GPU on Vulkan elsewhere, software Vulkan) was run by hand and is recorded in
# wai/extensions/determinism-tiers.md §6; this workflow runs no job on those hosts.
# One run per branch at a time; a push supersedes the run before it.
concurrency:
group: ${{ github.workflow }}-${{ github.ref }}
cancel-in-progress: true
on:
pull_request:
paths:
- 'wai/wai-ops-gpu/**'
- 'wai/wai-ops/**'
- 'wai/ops-conformance/**'
# The decoders whose graphs run on the device, and what they decode.
- 'wai/wai-rs/Cargo.toml'
- 'wai/wai-rs/Cargo.lock'
- 'wai/wai-rs/src/int_*.rs'
- 'wai/wai-rs/src/int_payload/**'
- 'wai/wai-rs/src/jpegai_*.rs'
- 'wai/wai-rs/tests/ops_gpu.rs'
- 'wai/wai-rs/tests/vectors/**'
- 'wai/int-payload-conformance/**'
- 'wai/neural-conformance/**'
- 'wai/tools/ops-gpu-bench/**'
- 'wai/tools/wai_jpegai_fetch_e2e_inputs.sh'
- 'wai/tools/jpegai-e2e-v1.sha256'
- 'rust-toolchain.toml'
- '.github/workflows/wai-ops-gpu.yml'
schedule:
- cron: '17 6 * * *'
jobs:
# Does this pull request's own diff touch the paths below? This forge matches a
# workflow's `paths:` against the push, so a merge of main into a branch fires the
# workflow whatever the pull request changes; every job here waits on this answer.
# A push to main, a schedule, or anything the script cannot determine answers
# `true` (scripts/pr-touches-paths.py).
relevant:
name: does this change touch it
runs-on: linux-arm64
outputs:
run: ${{ steps.gate.outputs.run }}
steps:
- uses: actions/checkout@v5
with:
fetch-depth: 0
- id: gate
run: |
set -euo pipefail
r=$(python3 scripts/pr-touches-paths.py .github/workflows/wai-ops-gpu.yml \
"${{ github.event.pull_request.base.sha }}" \
"${{ github.event.pull_request.head.sha }}")
echo "this change touches wai-ops-gpu: $r"
echo "run=$r" >> "$GITHUB_OUTPUT"
checks:
needs: relevant
if: needs.relevant.outputs.run == 'true'
name: builds · no adapter is not a pass (linux)
runs-on: linux
timeout-minutes: 60
steps:
- uses: actions/checkout@v5
- name: Install Rust 1.98.0 with the wasm32-unknown-unknown target
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
rustup target add --toolchain 1.98.0 wasm32-unknown-unknown
echo "$HOME/.cargo/bin" >> "$GITHUB_PATH"
# This container offers no adapter. The device tests must say so and pass; under
# WAI_OPS_GPU_REQUIRE=1 the same tests must fail. Were an adapter ever offered
# here, its results must then be equal instead.
- name: The crate's tests, and an absent adapter under WAI_OPS_GPU_REQUIRE
shell: bash
working-directory: wai/wai-ops-gpu
run: |
set -euo pipefail
log="$RUNNER_TEMP/ops-gpu.log"
rep="$RUNNER_TEMP/ops-gpu.report"
: > "$rep"
WAI_OPS_GPU_REPORT="$rep" cargo test --locked --release -- --nocapture --test-threads=1 2>&1 | tee "$log"
if grep -q '^ops-gpu corpus: key=' "$rep"; then
echo "this runner offered an adapter"
if grep -E '^ops-gpu corpus: key=' "$rep" | grep -vq ' mismatches=0 errors=0 '; then
echo "::error::an adapter here gave a result that is not the interpreter's"; exit 1
fi
else
grep -q '^ops-gpu corpus: SKIPPED' "$rep" || { echo "::error::the corpus test printed neither a result nor SKIPPED"; exit 1; }
if WAI_OPS_GPU_REQUIRE=1 cargo test --locked --release --test corpus --test arith --test mutants > "$RUNNER_TEMP/require.log" 2>&1; then
echo "::error::with no adapter, the device tests passed under WAI_OPS_GPU_REQUIRE=1"; exit 1
fi
grep -q 'WAI_OPS_GPU_REQUIRE is set' "$RUNNER_TEMP/require.log" \
|| { echo "::error::the tests failed under WAI_OPS_GPU_REQUIRE=1, but not for the absent adapter"; cat "$RUNNER_TEMP/require.log"; exit 1; }
echo "no adapter: SKIPPED without WAI_OPS_GPU_REQUIRE, a failure with it"
fi
- name: The browser module builds for wasm32-unknown-unknown
working-directory: wai/wai-ops-gpu/browser
run: cargo build --locked --release --target wasm32-unknown-unknown
- name: wai-rs builds with ops_gpu, and its GPU test target
working-directory: wai/wai-rs
run: cargo test --locked --no-default-features --features ops_gpu,video_int,audio_int,splat_int,world_int --test ops_gpu --no-run
- name: The bench tool builds
working-directory: wai/tools/ops-gpu-bench
run: cargo build --locked --release
metal:
needs: relevant
if: needs.relevant.outputs.run == 'true'
name: Metal GPU and browser WebGPU (macos)
runs-on: macos
timeout-minutes: 90
env:
WAI_OPS_GPU_EXPECT: gpu-metal-0
WAI_OPS_GPU_REQUIRE: '1'
steps:
- uses: actions/checkout@v5
# launchd's PATH omits cargo on the macOS host runner.
- name: Put cargo on PATH
shell: bash
run: echo "$HOME/.cargo/bin" >> "$GITHUB_PATH"
- name: The probe, the corpus at both widths, and the mutants, on the GPU
shell: bash
working-directory: wai/wai-ops-gpu
run: |
set -euo pipefail
log="$RUNNER_TEMP/ops-gpu.log"
export WAI_OPS_GPU_REPORT="$RUNNER_TEMP/ops-gpu.report"
: > "$WAI_OPS_GPU_REPORT"
nice -n 10 cargo test --locked --release -- --nocapture --test-threads=1 2>&1 | tee "$log"
k=gpu-metal-0
# Six tests in four binaries: the library's one, arith's two, corpus's two and
# mutants' one, none filtered, ignored or failed.
ran=$(awk '/^test result: ok\. [0-9]+ passed; 0 failed; 0 ignored; 0 measured; 0 filtered out/ { n += $4 } END { print n + 0 }' "$log")
[ "$ran" = 6 ] || { echo "::error::$ran tests ran and passed, not 6"; exit 1; }
! grep -q '^test result: FAILED' "$log" || { echo "::error::a test failed"; exit 1; }
grep -qxF "ops-gpu adapters: as declared ($k)" "$WAI_OPS_GPU_REPORT" || { echo "::error::the adapters found are not $k"; exit 1; }
grep -qE "^ops-gpu arith: key=$k .* cases=[0-9]+ words=[0-9]+ equal$" "$WAI_OPS_GPU_REPORT" || { echo "::error::no probe result for $k"; exit 1; }
grep -qE "^ops-gpu corpus: key=$k .* acc32 vectors=[1-9][0-9]* decodes=[1-9][0-9]*; acc64 vectors=[1-9][0-9]* decodes=[1-9][0-9]*; .* mismatches=0 errors=0 " "$WAI_OPS_GPU_REPORT" \
|| { echo "::error::no equal corpus run for $k"; exit 1; }
grep -qxF "ops-gpu mutants: key=$k killed 18 of 18" "$WAI_OPS_GPU_REPORT" || { echo "::error::a mutant survived on $k, or the mutant test did not run"; exit 1; }
- name: The jpegai/v1 inputs
shell: bash
run: bash wai/tools/wai_jpegai_fetch_e2e_inputs.sh "$RUNNER_TEMP/jpegai"
- name: Every capability that runs a wai-ops/1 graph, with its graphs on the GPU
shell: bash
working-directory: wai/wai-rs
run: |
set -euo pipefail
log="$RUNNER_TEMP/ops-gpu-rs.log"
rep="$RUNNER_TEMP/ops-gpu-rs.report"
: > "$rep"
J="$RUNNER_TEMP/jpegai"
WAI_OPS_GPU_REPORT="$rep" WAI_JPEGAI_CODESTREAM="$J/test256.bin" WAI_JPEGAI_WEIGHTS_DIR="$J" WAI_JPEGAI_NORM="$J/norm.json" WAI_REQUIRE_SOTA_VECTORS=1 \
nice -n 10 cargo test --locked --release --no-default-features --features ops_gpu,video_int,audio_int,splat_int,world_int \
--test ops_gpu -- --nocapture 2>&1 | tee "$log"
grep -q '^test result: ok\. 1 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out' "$log" || { echo "::error::the capability test did not run and pass"; exit 1; }
for cap in wai.neural.int_hyper wai.video.int_hyper wai.audio.int_codec wai.splat.int_codec wai.neural.int_synth \
wai.neural.int_mlicpp wai.world.model wai.image.jpegai; do
grep -qE "^ops-gpu capability: key=gpu-metal-0 .* cap=$cap cases=[1-9][0-9]* .* differ=0 errors=0$" "$rep" \
|| { echo "::error::$cap: no equal run on the GPU"; exit 1; }
done
- name: wasm-bindgen CLI 0.2.127, and the wasm32-unknown-unknown target
shell: bash
run: |
set -euo pipefail
rustup target add wasm32-unknown-unknown
cargo install --locked wasm-bindgen-cli --version 0.2.127
# Playwright and its headless browser go into the job's temporary directory, not
# the host's caches. `chromium` is Playwright's name for its engine.
- name: Playwright 1.60.0 and its headless browser
shell: bash
run: |
set -euo pipefail
pw="$RUNNER_TEMP/playwright"
mkdir -p "$pw"
npm install --prefix "$pw" --no-save --no-audit --no-fund [email protected]
PLAYWRIGHT_BROWSERS_PATH="$pw/browsers" "$pw/node_modules/.bin/playwright" install --only-shell chromium
echo "PLAYWRIGHT_HOST=$pw/" >> "$GITHUB_ENV"
echo "PLAYWRIGHT_BROWSERS_PATH=$pw/browsers" >> "$GITHUB_ENV"
- name: The corpus in a browser's WebGPU
shell: bash
working-directory: wai/wai-ops-gpu/browser
run: |
set -euo pipefail
nice -n 10 cargo build --locked --release --target wasm32-unknown-unknown
t="${CARGO_TARGET_DIR:-target}"
wasm-bindgen --target web --out-dir pkg --out-name wai_ops_gpu_browser "$t/wasm32-unknown-unknown/release/wai_ops_gpu_browser.wasm"
log="$RUNNER_TEMP/browser.log"
nice -n 10 node run.mjs 2>&1 | tee "$log"
grep -qE '^browser webgpu: key=gpu-webgpu-[0-9a-f]+ .* acc32 vectors=[1-9][0-9]* decodes=[1-9][0-9]*; acc64 vectors=[1-9][0-9]* decodes=[1-9][0-9]*; .* mismatches=0 errors=0$' "$log" \
|| { echo "::error::no equal corpus run on a GPU adapter in the browser"; exit 1; }

14
wai/CHANGELOG.md
View file

@ -6,6 +6,20 @@ wire format may change and such changes are recorded here.
## v0.4.0 — current
### `wai-ops/1-i32` on GPUs, and the tier where a run showed it
- **`wai/wai-ops-gpu`**, an operator implementation of `wai-ops/1-i32` (integer-operators §9.1) in WGSL compute kernels through wgpu: Metal, Vulkan and DX12 natively, and a browser's WebGPU. WGSL has no 64-bit integer, so the kernels use 32-bit words only:
- a graph whose certificate meets `acc32` runs with one word per value, its two `product` intermediates formed from 16-bit halves (§7.3);
- every other graph runs with each `acc` value in two words and each `wide` and `product` value in four, exact 64-bit and 128-bit arithmetic built from 32-bit words, an `acc64` implementation (§7.3);
- a graph with `layernorm_c`, `l2norm`, `softmax`, `sigmoid` or `elu` runs on the CPU, and `support` names the operator that sent it there.
The host makes the interpreter's run-time checks, in its order, before the device runs anything; a `check` tensor that leaves its range sets its operator's flag, and the first flagged operator refuses the decode with `field_range`. Each device passes an arithmetic probe against 128-bit host arithmetic before it runs a graph.
- **wai-rs `ops_gpu`** (off by default): `int_ops::run_graph` is now where every integer decoder runs a graph, and inside `int_ops::gpu::with_gpu` it runs each graph on the device, counting what ran where and why.
- **Tests that do not pass on a skip.** `wai-ops-gpu/tests/`: the probe on a million cases; every graph vector of `ops-conformance/` that the subset covers, at both widths, against the committed results and the interpreter; eighteen kernel edits, each of which the corpus must catch; and the 32-bit kernels shown to differ on graphs their certificate keeps from them. `wai-rs/tests/ops_gpu.rs`: every committed vector and refusal fixture of every capability that runs a graph, with its graphs on the device, against the CPU decode, graph by graph. `wai-ops-gpu/browser/` runs the corpus in a page's WebGPU. With no adapter each prints SKIPPED, and fails under `WAI_OPS_GPU_REQUIRE`.
- **determinism-tiers §5** now meets `decode-equivalence` for a decoder's `wai-ops/1-i32` graphs on the GPUs and backends a run showed equal, and lists them; §6 records the runs. Nothing is claimed for another adapter, and nothing for an NPU.
- **`tools/ops-gpu-bench`** times each capability's decode with its graphs on the CPU and on each adapter, and reads the energy counters the host has.
- **CI:** `wai-ops-gpu.yml` runs the Metal and browser legs on the `macos` runner, and requires on a runner with no GPU that the device tests skip without `WAI_OPS_GPU_REQUIRE` and fail with it.
### wai-quantum 0.6.12: cheaper contraction trees, and the cut-rank table on the GPU
0.6.11 was not published on its own: this release carries it to the package registry, and `wai-quantum-cli` 0.1.0 requires it, since the command line's `metal` feature turns on `quantum_cutrank_metal`.

5
wai/SPEC.md
View file

@ -1635,8 +1635,9 @@ WebAssembly, or a neural decoder.
**What the exact tier covers.** `decode-equivalence` covers the medium a sink emits and
can hash ([determinism-tiers](extensions/determinism-tiers.md) §2), on the fabrics
determinism-tiers §5 lists: native CPU and WebAssembly in this revision, with no GPU or
NPU decoder. It never covers output the application cannot read. A sink MUST NOT state
determinism-tiers §5 lists: native CPU and WebAssembly in this revision, and, for a
decoder's `wai-ops/1-i32` graphs, the GPUs and backends listed there, with no NPU
decoder. It never covers output the application cannot read. A sink MUST NOT state
a tier, an output digest, a reconstruction binding or an energy binding for output
decoded inside a protected media path, and a `reconstruct-session` claim
(jwp-receipts §2.4.12), whose every presentation carries the SHA-256 of its output,

46
wai/extensions/determinism-tiers.md
View file

@ -516,10 +516,41 @@ shows, by fabric:
(`wasm-tools validate`), which CI does not check. A WebAssembly decoder at this
tier MUST NOT use relaxed-SIMD instructions, whose results the WebAssembly
specification lets engines choose among, and MUST NOT depend on the bits of a NaN.
- **GPU and NPU.** No decoder of these capabilities runs on either in this revision,
and the tier is not claimed there. A network evaluated in floating point on a GPU or
NPU does not meet it (§1). An integer GPU kernel could; until a run shows its output
equal to the reference's, none is claimed.
- **GPU.** A decoder's `wai-ops/1-i32` graphs ([integer-operators](integer-operators.md)
§9.1) run on a GPU through `wai/wai-ops-gpu`: WGSL compute kernels in 32-bit words, at
the `acc32` widths for a graph whose certificate meets that profile, and otherwise in
64-bit and 128-bit arithmetic built from 32-bit words, an `acc64` implementation
(integer-operators §7.3). The host makes the interpreter's run-time checks before the
device runs anything, and a `check` tensor that leaves its range refuses the decode at
its operator. The driver's stage operators stay on the CPU, and so does a graph with an
operator outside `wai-ops/1-i32`. The tier is met for those graphs on each adapter and
backend below, where a run reproduced, at both widths, every graph vector of
integer-operators §14.1 and §14.2 that `wai-ops/1-i32` covers (§6):
- Metal, on the integrated GPU of an arm64 desktop (macOS 27.0.1);
- a browser's WebGPU on that same GPU, the kernels compiled by the browser's own WGSL
compiler (Chromium 148 and a Chromium-based browser at 155, headless);
- Vulkan, on a discrete GPU of PCI vendor 0x10de, device 0x28e1 (driver 595.84), and
on an integrated GPU of PCI vendor 0x8086, device 0xa7a0 (Mesa 25.2.8), in one Linux
x86_64 host;
- Vulkan, on an integrated GPU of PCI vendor 0x1002, device 0x1586 (Mesa 26.0.8),
Linux x86_64.
On each of these but the browser, the capability corpora (integer-operators §14.4)
also decoded with their graphs on the device: every committed vector and refusal
fixture of `wai.neural.int_hyper`, `wai.video.int_hyper`, `wai.audio.int_codec`,
`wai.splat.int_codec`, `wai.neural.int_synth` and `wai.neural.int_mlicpp` gave the
bytes, or the refusal code, of the decode on the CPU, and every graph the device ran
gave the interpreter's outputs. A software Vulkan driver (llvmpipe, on both Linux hosts)
gave the same results; it runs on the CPU. Nothing is claimed for an adapter, a backend
or a driver version that is not listed; DX12 was not run. `wai.image.jpegai`'s scale
graphs (`hs_y`, `hs_uv`, at `acc32`) and hyper decoders and context stages (`rec_y`,
`rec_uv`, at `acc64`) ran on each of these devices but the browser with the
interpreter's outputs, so its `entropy-consistency` claim (§5) holds with its scale
path there. Its synthesis graphs use operators outside `wai-ops/1-i32` and ran on the
CPU.
- **NPU.** No decoder of these capabilities runs on one in this revision, and the tier is
not claimed there. A network evaluated in floating point on a GPU or NPU does not meet
it (§1).
The tier covers the medium a sink emits and can hash (§2, *What a tier covers*). A
sink that hands a payload to a protected media path whose output the application
@ -564,6 +595,7 @@ task list has no Windows task after that day. The Windows results below are hist
| `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 run on the forge's Linux runner (x86_64 hardware) and passed on every run since 2026-09-28 | `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. Their frozen hashes are the output of an independent Python integer reference, computed when the vectors were exported (`tools/wai_int_audio_export.py`, `tools/wai_int_splat_export.py`); no Python runs in CI. The conformance workflow's fast-kernel steps enable `audio_int` and `splat_int` and, in debug and release builds, decode both vectors against those hashes under the reference kernels and every fast-kernel policy (`tests/int_kernels_vectors.rs`); that passed on the forge's Linux x86_64, Linux arm64 and macOS arm64 runners for the landing chain's head `7cfcecfa` (run 745, 2026-09-30). The same steps also name `tests/int_audio.rs` and `tests/int_splat.rs`, which check each vector against the same hash and against its own metadata (sample count; attribute count, tensor size and grid side); Cargo declares each with its feature as a required feature, so a step that names it without that feature fails instead of running nothing. This row records no forge run of those two | `wai-rs/tests/int_audio.rs`, `int_splat.rs`, `int_kernels_vectors.rs`; `.github/workflows/wai-conformance.yml` |
| `wai.neural.int_mlicpp` on its committed vector (synthetic seeded weights; a 64×64 image, 8 latent channels in 4 slices): the decode reproduces the export tool's independent integer reference byte for byte — the RGB, and every scale bucket and symbol of the latent stream in coding order (512 symbols over all 16 buckets of the scale table, 47 of the symbols coded by the rANS escape) — under the reference kernels and under the fast kernels on 1, 2, 3, 5 and all available threads; one integer unit changed in each of four chosen parameters, one on each path that feeds a value (a hyper-decoder weight, an entropy-network weight, a mean-head bias, a synthesis weight), changes the image; a header whose shapes the model cannot produce, or on which a network would form a buffer over the cap, is refused before anything is allocated (`tests/int_mlicpp_preflight.rs`, which counts allocations); damaged payloads are refused or decode, and never panic in an overflow-checked build | measured 2026-09-29 with rustc 1.98.0 on the source of the change that added this row: macOS arm64 native (the wai-rs tests, with overflow checks on and off; the `byte-exact-conformance` golden, release) and wasm32-wasip2 under wasmtime 49.0.1 (the same tests and golden, release; on wasm the fast kernels run single-threaded as portable scalar code, and `usize` is 32-bit). All pinned hashes matched on both. Not measured: x86_64 — the development machine has no x86_64 translation layer installed, and an x86_64-apple-darwin build of the golden compiles but cannot execute there — nor Linux or Windows. The golden runs on every leg of the byte-exact matrix (`wai-conformance.yml`, `wai-byte-equality.yml`) and `tests/int_mlicpp.rs` on every leg of `wai-conformance.yml` in both profiles. Both passed on the forge's Linux x86_64, Linux arm64 and macOS arm64 runners for the landing chain's head `7cfcecfa` (run 745, 2026-09-30); there is no Windows leg. The dispatch through `decode_envelope_strict` (the prior pinning the one-file model bundle, the strict gate) runs in `wai-neural-gates.yml`. The weights are synthetic, so nothing here measures compression quality | `wai-rs/tests/int_mlicpp.rs`, `wai-rs/tests/int_mlicpp_preflight.rs`, `byte-exact-conformance/src/int_mlicpp_golden.rs`, `wai-rs/tests/neural_native_decode.rs`; `tools/wai_int_mlicpp_export.py` |
| `wai-ops/1-i32` graphs on GPUs (`wai/wai-ops-gpu`, §5): each device's arithmetic agrees with 128-bit host arithmetic on every probe case; every graph vector of `ops-conformance/` that the subset covers gives the interpreter's outputs or refusal code, at `acc32` and at `acc64`; the integer decoders, with their graphs on the device, give the CPU decode's bytes or refusal codes on every committed vector and refusal fixture, and every graph the device runs gives the interpreter's outputs | measured 2026-10-20 and 2026-10-21 with rustc 1.98.0 and wgpu 30.0.1, release builds, on the source of the change that added this row, on the adapters §5 lists: Metal on macOS 27.0.1 (arm64); headless Chromium 148, and a Chromium-based browser at 155, on that GPU; Vulkan on Linux x86_64 with the 0x10de and 0x8086 adapters (driver 595.84, Mesa 25.2.8) and, on another host, with the 0x1002 adapter (Mesa 26.0.8); and llvmpipe on both Linux hosts. On each, the probe (`tests/arith.rs`) agreed on every word of 1,051,757 cases, and the corpus (`tests/corpus.rs`) ran 813 vectors (1,104 decodes) on the 32-bit kernels and 1,426 (1,921 decodes) on the 64-bit ones, 212 of the decodes refusals, with no difference; 209 graphs use an operator outside the subset and ran on the CPU, and 95 vectors are refused at load. Run on the 32-bit kernels, 330 of the 613 graphs whose certificate does not meet `acc32` gave other results, so the width a graph runs at is the certificate's. Eighteen edits of one kernel expression each (a rounding shift or its addend off by one, a tie rule, a clamp bound moved by one, a dropped carry or cross term, a sign rule, a post bias, an index rule; `tests/mutants.rs`) were each caught by the corpus on every adapter, and an edited product by the probe. `wai-rs/tests/ops_gpu.rs` decoded the committed vectors and refusal fixtures of the six capabilities §5 names, the world renderer's two vectors and `wai.image.jpegai`'s `jpegai/v1` codestream, comparing every graph the device ran with the interpreter as well as each decode's bytes; none differed. The browser leg ran the corpus, not the decoders. `wai-ops-gpu.yml` runs the Metal and browser legs on the `macos` runner, where an absent adapter fails; this row records no forge run of it yet. No other workflow runs a GPU leg of these tests | `wai-ops-gpu/tests/`, `wai-ops-gpu/browser/`, `wai-rs/tests/ops_gpu.rs`; `.github/workflows/wai-ops-gpu.yml` |
So at `entropy-consistency`, WAI's JPEG AI symbol decode is demonstrated across
architectures (x86_64 against arm64) once the distribution indices are given; the
@ -581,7 +613,11 @@ decode against the export's integer reference on one machine per run.
`wai.neural.int_mlicpp`'s committed vector decodes to the same pinned bytes on macOS
arm64 and on wasm32-wasip2 (a 32-bit target), measured on the development machine;
it also passed on the byte-exact matrix's Linux x86_64, Linux arm64 and macOS arm64
legs in run 745.
legs in run 745. On GPUs, the committed vectors of `int_hyper` (and its video keyframe),
`int_codec` audio and splat, `int_synth` and `int_mlicpp` decode with their graphs on
the device to the CPU decode's bytes on each adapter §5 lists, measured by hand on small
vectors; the `wai-ops-gpu` workflow repeats it on one of them, with no forge run
recorded yet.
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

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

@ -14,7 +14,7 @@
>
> **What a reader can build from this text.** With the payload extensions it cites, an implementation of this text decodes `wai.video.int_motion`, `wai.audio.int_codec`, `wai.splat.int_codec`, `wai.neural.int_synth`, `wai.neural.int_hyper`, `wai.video.int_hyper`, `wai.neural.int_mlicpp` and `wai.{image,video,audio,splat}.int_refine`, and reproduces the integers of `wai.image.jpegai`'s integer stages: the symbols and scale indices it entropy-decodes, and what its integer networks give for their inputs (§13.6). The payload text four of these depend on, `entropy.json`, the `int_synth` payload and the `WIM1` payload with its `WMB1` bundle, is in [integer-payloads](integer-payloads.md), and §15.2 says where. The integer renderer of `wai.world.model` is informative (§13.5).
>
> Reference implementation: the `wai-ops` crate (`wai/wai-ops`): parse, lowering, interpreter, stage operators and bounds. Tool: `wai-verify bounds`. Conformance: `wai/ops-conformance/`. A capability's reference decoder decodes through this operator set from the revision that registers the capability's `ops` member (§11); in this revision, `wai.neural.int_synth`, `wai.splat.int_codec`, `wai.audio.int_codec`, `wai.neural.int_hyper`, `wai.video.int_hyper` and `wai.neural.int_mlicpp`, and the integer stages of `wai.image.jpegai` (§13.6): its scale path and entropy decodes on both of its reference decode paths, and its integer networks on its exact decoder; and, through their stage operators alone with no graph (§13.5), `wai.video.int_motion`, the P-frames of `wai.video.int_hyper`, and `wai.{image,video,audio,splat}.int_refine`. The integer renderer of `wai.world.model` renders through it too, by `world/1` (§13.5), and stays informative. Until then, a capability's reference implementation is the decoder its payload extension names, and §15.3 lists where this text and that decoder differ.
> Reference implementation: the `wai-ops` crate (`wai/wai-ops`): parse, lowering, interpreter, stage operators and bounds. Tool: `wai-verify bounds`. Conformance: `wai/ops-conformance/`. An operator implementation of `wai-ops/1-i32` for GPUs, at both profiles, is `wai/wai-ops-gpu` (WGSL); [determinism-tiers](determinism-tiers.md) §5 lists the adapters on which it has reproduced the vectors. A capability's reference decoder decodes through this operator set from the revision that registers the capability's `ops` member (§11); in this revision, `wai.neural.int_synth`, `wai.splat.int_codec`, `wai.audio.int_codec`, `wai.neural.int_hyper`, `wai.video.int_hyper` and `wai.neural.int_mlicpp`, and the integer stages of `wai.image.jpegai` (§13.6): its scale path and entropy decodes on both of its reference decode paths, and its integer networks on its exact decoder; and, through their stage operators alone with no graph (§13.5), `wai.video.int_motion`, the P-frames of `wai.video.int_hyper`, and `wai.{image,video,audio,splat}.int_refine`. The integer renderer of `wai.world.model` renders through it too, by `world/1` (§13.5), and stays informative. Until then, a capability's reference implementation is the decoder its payload extension names, and §15.3 lists where this text and that decoder differ.
>
> The keywords **MUST**, **MUST NOT**, **SHOULD** and **MAY** are as defined in RFC 2119 and RFC 8174.

1434
wai/tools/ops-gpu-bench/Cargo.lock generated Normal file
View file

File diff suppressed because it is too large Load diff

32
wai/tools/ops-gpu-bench/Cargo.toml Normal file
View file

@ -0,0 +1,32 @@
# Time and energy of the integer decoders with their graphs on the CPU and on a GPU
# (wai-rs `ops_gpu`), per capability, on the committed vectors.
#
# A SEPARATE crate with its own lockfile, as quantum-meter is: the energy sources it reads
# (this process's CPU and GPU counters on macOS, RAPL package energy and the NVIDIA
# driver's cumulative energy counter on Linux) belong with the tool that reads them, and
# never compile in wai-rs's own builds or tests.
[package]
name = "ops-gpu-bench"
version = "0.0.0"
edition = "2024"
rust-version = "1.97"
license = "Apache-2.0"
publish = false
[dependencies]
wai = { path = "../../wai-rs", default-features = false, features = ["ops_gpu", "video_int", "audio_int", "splat_int", "world_int"] }
wai-ops = { path = "../../wai-ops" }
serde_json = "1"
# On macOS on Apple Silicon: this process's CPU energy and its share of the GPU rail
# (IOReport), the span meter quantum-meter uses.
[target.'cfg(target_os = "macos")'.dependencies]
jouleclaw-energy = { path = "../../../jouleclaw/jouleclaw-rs/crates/jouleclaw-energy", default-features = false, features = ["ioreport"] }
# On Linux: the NVIDIA driver's library is opened at run time, so a host without it
# builds and reports no GPU energy.
[target.'cfg(target_os = "linux")'.dependencies]
libc = "0.2"
[profile.release]
debug = false

529
wai/tools/ops-gpu-bench/src/main.rs Normal file
View file

@ -0,0 +1,529 @@
//! `ops-gpu-bench`: time, and energy where this host has a counter for it, of one decode
//! of each committed vector of each capability that runs a `wai-ops/1` graph, with its
//! graphs on the CPU (the reference sink's default: the fast kernels on every core) and
//! on each GPU adapter found (`wai::int_ops::gpu`).
//!
//! ops-gpu-bench [--secs S] [--min-runs N] [--only <capability>]
//!
//! with `WAI_JPEGAI_CODESTREAM`, `WAI_JPEGAI_WEIGHTS_DIR` and `WAI_JPEGAI_NORM` naming the
//! `jpegai/v1` inputs to include `wai.image.jpegai`'s exact decoder.
//!
//! **Method.** For each vector and each path: two decodes to warm up (on a GPU these
//! compile nothing, but upload the graph's parameters once), then decodes back to back
//! until at least `S` seconds (default 3) and `N` decodes (default 5) have run. Each
//! decode's wall time is recorded; the line gives the count, the median and the 10th and
//! 90th percentiles. Every decode's output is compared to the CPU's first, byte for byte.
//! The energy window is the timed loop, one window per line, read only at its two ends:
//!
//! - macOS on Apple Silicon: this process's CPU energy and its share of the GPU rail
//! (`jouleclaw_energy::span::ProcessSpanMeter`, IOReport), which gives no GPU figure
//! where another process used the GPU over the window;
//! - Linux: the RAPL `package-0` counter (`/sys/class/powercap/intel-rapl:0`), the whole
//! package, every process on it included; and on a host with the NVIDIA driver, the
//! driver's cumulative energy counter for GPU 0 (`nvmlDeviceGetTotalEnergyConsumption`,
//! the board), read at the two ends only, since a counter read every ~100 ms has been
//! seen to over-read on this hardware. Before the decodes, one idle window of `S`
//! seconds gives each counter's idle power, printed with the lines.
//!
//! The printed lines are the measurement; nothing here turns them into a claim.
use std::collections::BTreeMap;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Instant;
use serde_json::Value;
use wai::int_hyper_synth::IntModel;
use wai::int_ops::gpu::{self, wai_ops_gpu};
use wai::int_payload::{IntDecoded, IntParams, decode_int};
use wai::pin::HeldBytes;
fn wai_dir() -> PathBuf {
PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../..")
}
fn read(rel: &str) -> Vec<u8> {
std::fs::read(wai_dir().join(rel)).unwrap_or_else(|e| panic!("{rel}: {e}"))
}
fn json(rel: &str) -> Value {
serde_json::from_slice(&read(rel)).unwrap_or_else(|e| panic!("{rel}: {e}"))
}
struct Case {
cap: &'static str,
name: String,
run: Box<dyn Fn() -> Vec<u8>>,
}
fn cap_static(c: &str) -> Option<&'static str> {
["wai.neural.int_hyper", "wai.video.int_hyper", "wai.audio.int_codec", "wai.splat.int_codec"].into_iter().find(|k| *k == c)
}
fn bytes_of(d: &IntDecoded) -> Vec<u8> {
match d {
IntDecoded::Image { rgb, .. } => rgb.clone(),
IntDecoded::Video { frames, .. } => frames.concat(),
IntDecoded::Pcm { samples, .. } => samples.iter().flat_map(|s| s.to_le_bytes()).collect(),
IntDecoded::Attr { bytes, .. } => bytes.clone(),
_ => panic!("an output form this tool does not know"),
}
}
fn cases() -> Vec<Case> {
let mut out = Vec::new();
// The payload vectors whose capability runs a graph (`wai.video.int_motion` runs none).
for e in json("int-payload-conformance/vectors.json")["vectors"].as_array().unwrap() {
let Some(cap) = cap_static(e["capability"].as_str().unwrap()) else { continue };
let payload = read(e["payload"].as_str().unwrap());
let set: BTreeMap<String, Vec<u8>> = e["set"].as_object().into_iter().flatten().map(|(k, v)| (k.clone(), read(v.as_str().unwrap()))).collect();
let run = move || {
let view: BTreeMap<String, &[u8]> = set.iter().map(|(k, v)| (k.clone(), v.as_slice())).collect();
let params = wai::int_payload::resolve(cap, None, &[HeldBytes::Set { listing: b"", files: &view }], None).unwrap_or_else(|_| IntParams::none());
bytes_of(&decode_int(cap, &payload, &params).expect("a committed vector decodes"))
};
out.push(Case { cap, name: e["name"].as_str().unwrap().into(), run: Box::new(run) });
}
for name in json("neural-conformance/expected_hashes.json").as_object().unwrap().keys() {
let dir = format!("neural-conformance/vectors/{name}");
let (payload, set) = (read(&format!("{dir}/latent.bin")), read(&format!("{dir}/int_synth.json")));
out.push(Case { cap: "wai.neural.int_synth", name: name.clone(), run: Box::new(move || wai::int_synth::decode_from_spec(&payload, &set).unwrap().rgb) });
}
{
const DIR: &str = "wai-rs/tests/vectors/int_mlicpp";
let m = wai::int_mlicpp::IntMlicModel::from_bundle(&read(&format!("{DIR}/model.wmb"))).unwrap();
let p = wai::container::Wai::from_bytes(&read(&format!("{DIR}/stream.wai"))).unwrap().payload;
out.push(Case { cap: "wai.neural.int_mlicpp", name: "int_mlicpp".into(), run: Box::new(move || wai::int_mlicpp::decode(&m, &p).unwrap().rgb) });
}
for (name, dir, js, t, g) in [("world", "int_world", "world.json", "world_T.bin", "world_gs.bin"), ("world2", "int_world2", "world2.json", "world2_T.bin", "world2_gs.bin")] {
let j = json(&format!("wai-rs/tests/vectors/{dir}/{js}"));
let tm = IntModel::from_bytes(&read(&format!("wai-rs/tests/vectors/{dir}/{t}"))).unwrap();
let gm = IntModel::from_bytes(&read(&format!("wai-rs/tests/vectors/{dir}/{g}"))).unwrap();
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(|x| x.as_i64().unwrap()).collect();
let afrac = j["afrac"].as_u64().unwrap() as u32;
let q = |d: i64| wai_ops::stage::dequantize(d, afrac, 0).unwrap();
let mv = |a: &Value| wai::int_world::ACTIONS[a.as_u64().unwrap() as usize];
let rows: Vec<Vec<i64>> = match j.get("pairs") {
Some(p) => p.as_array().unwrap().iter().map(|p| {
let ((xa, ya), (xb, yb)) = (mv(&p[0]), mv(&p[1]));
vec![q(xa), q(ya), q(xb), q(yb)]
}).collect(),
None => j["actions"].as_array().unwrap().iter().map(|a| {
let (dx, dy) = mv(a);
vec![q(dx), q(dy), 0]
}).collect(),
};
out.push(Case { cap: "wai.world.model", name: name.into(), run: Box::new(move || wai::int_world::decode_planes(&tm, &gm, &z0, &rows, c, lat).unwrap().frames.concat()) });
}
let env = |k: &str| std::env::var(k).ok().filter(|v| !v.is_empty());
if let (Some(cs), Some(dir), Some(nf)) = (env("WAI_JPEGAI_CODESTREAM"), env("WAI_JPEGAI_WEIGHTS_DIR"), env("WAI_JPEGAI_NORM")) {
let (data, weights, norm) = jpegai::load(&cs, &dir, &nf);
out.push(Case { cap: "wai.image.jpegai", name: "jpegai/v1 test256".into(), run: Box::new(move || wai::jpegai_decode::decode(&data, &weights, &norm).unwrap().data) });
} else {
println!("ops-gpu bench: wai.image.jpegai not measured: WAI_JPEGAI_CODESTREAM, WAI_JPEGAI_WEIGHTS_DIR and WAI_JPEGAI_NORM are not all set");
}
out
}
mod jpegai {
use wai::jpegai_decode::{DecoderWeights, NormTables, PlaneWeights};
use wai::jpegai_icci::{Conv, IcciNet, IcciPath, ResBlock};
use wai::jpegai_synthesis::{IntNet, Weights};
fn jload(p: &std::path::Path) -> serde_json::Value {
serde_json::from_slice(&std::fs::read(p).unwrap()).unwrap()
}
fn w_of(dir: &str, name: &str) -> Weights {
Weights::from_json(&jload(&std::path::Path::new(dir).join(name))).unwrap()
}
fn int_of(dir: &str, name: &str) -> IntNet {
IntNet::from_json(&jload(&std::path::Path::new(dir).join(name))).unwrap()
}
fn i64v(v: &serde_json::Value, k: &str) -> Vec<i64> {
v[k].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect()
}
fn i32v(v: &serde_json::Value, k: &str) -> Vec<i32> {
v[k].as_array().unwrap().iter().map(|x| x.as_i64().unwrap() as i32).collect()
}
fn f3(v: &serde_json::Value, k: &str) -> [f32; 3] {
let a = v[k].as_array().unwrap();
[a[0].as_f64().unwrap() as f32, a[1].as_f64().unwrap() as f32, a[2].as_f64().unwrap() as f32]
}
fn i3(v: &serde_json::Value, k: &str) -> [i64; 3] {
let a = v[k].as_array().unwrap();
[a[0].as_i64().unwrap(), a[1].as_i64().unwrap(), a[2].as_i64().unwrap()]
}
fn conv_of(c: &serde_json::Value) -> Conv {
let g = |k: &str| c[k].as_u64().unwrap() as usize;
Conv {
w: c["w"].as_array().unwrap().iter().map(|x| x.as_f64().unwrap() as f32).collect(),
b: c["b"].as_array().unwrap().iter().map(|x| x.as_f64().unwrap() as f32).collect(),
cin: g("cin"),
cout: g("cout"),
kh: g("kh"),
kw: g("kw"),
ph: g("ph"),
pw: g("pw"),
}
}
fn path_of(v: &serde_json::Value) -> IcciPath {
IcciPath {
conv_first: conv_of(&v["conv_first"]),
blocks: v["blocks"].as_array().unwrap().iter().map(|b| ResBlock { conv1: conv_of(&b["c1"]), conv2: conv_of(&b["c2"]), scale: b["scale"].as_f64().unwrap() as f32 }).collect(),
conv_last: conv_of(&v["conv_last"]),
scale: v["scale"].as_f64().unwrap() as f32,
}
}
pub fn load(cs: &str, dir: &str, nf: &str) -> (Vec<u8>, DecoderWeights, NormTables) {
let data = std::fs::read(cs).unwrap();
let nv = jload(std::path::Path::new(nf));
let iw = jload(&std::path::Path::new(dir).join("icci.json"));
let icci = IcciNet { luma: path_of(&iw["luma"]), chroma: path_of(&iw["chroma"]), conv_last_v: conv_of(&iw["conv_last_v"]), scale_v: iw["scale_v"].as_f64().unwrap() as f32 };
let weights = DecoderWeights {
luma: PlaneWeights { hyper_scale: int_of(dir, "hs_y.json"), hyper_decoder: w_of(dir, "hyp_y.json"), synthesis: w_of(dir, "synth_y.json") },
chroma: PlaneWeights { hyper_scale: int_of(dir, "hs_uv.json"), hyper_decoder: w_of(dir, "hyp_uv.json"), synthesis: w_of(dir, "synth_uv.json") },
mcm_y: [w_of(dir, "mcm0_y.json"), w_of(dir, "mcm1_y.json"), w_of(dir, "mcm2_y.json"), w_of(dir, "mcm3_y.json")],
icci,
};
let norm = NormTables {
pmf: nv["pmf"].as_array().unwrap().iter().map(|r| r.as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect()).collect(),
bounds: i64v(&nv, "bounds"),
per_channel: nv["per_channel"].as_u64().unwrap() as usize,
cdfs_y: i64v(&nv, "cdfs_y"),
cdfs_uv: i64v(&nv, "cdfs_uv"),
gain_y: i64v(&nv, "gain_y"),
gain_uv: i64v(&nv, "gain_uv"),
lsbs_t0: i32v(&nv, "lsbs_t0"),
lsbs_t1: i32v(&nv, "lsbs_t1"),
lef_thr: i3(&nv, "lef_thr0"),
lef_mag: f3(&nv, "lef_mag0"),
};
(data, weights, norm)
}
}
// ---------------------------------------------------------------- energy, Linux
#[cfg(target_os = "linux")]
mod meter {
use std::time::{Duration, Instant};
/// The RAPL package-0 counter, µJ, with its wrap range.
pub struct Rapl {
range: u64,
}
const RAPL: &str = "/sys/class/powercap/intel-rapl:0";
impl Rapl {
pub fn open() -> Result<Rapl, String> {
let name = std::fs::read_to_string(format!("{RAPL}/name")).map_err(|e| format!("{RAPL}/name: {e}"))?;
if name.trim() != "package-0" {
return Err(format!("{RAPL} is {:?}, not package-0", name.trim()));
}
let range = std::fs::read_to_string(format!("{RAPL}/max_energy_range_uj")).map_err(|e| e.to_string())?.trim().parse().map_err(|e| format!("{e}"))?;
let r = Rapl { range };
r.read()?;
Ok(r)
}
pub fn read(&self) -> Result<u64, String> {
std::fs::read_to_string(format!("{RAPL}/energy_uj")).map_err(|e| format!("{RAPL}/energy_uj: {e}"))?.trim().parse().map_err(|e| format!("{e}"))
}
pub fn delta(&self, a: u64, b: u64) -> u64 {
if b >= a { b - a } else { b + self.range + 1 - a }
}
}
/// The NVIDIA driver's cumulative energy counter for GPU 0, mJ, through its library
/// opened at run time.
pub struct Nvml {
dev: *mut core::ffi::c_void,
energy: unsafe extern "C" fn(*mut core::ffi::c_void, *mut u64) -> i32,
power: unsafe extern "C" fn(*mut core::ffi::c_void, *mut u32) -> i32,
}
impl Nvml {
pub fn open() -> Result<Nvml, String> {
// SAFETY: dlopen/dlsym with NUL-terminated names; each symbol is cast to the
// signature NVML's header gives it, and is called only after a non-null check.
unsafe {
let h = libc::dlopen(c"libnvidia-ml.so.1".as_ptr(), libc::RTLD_NOW);
if h.is_null() {
return Err("libnvidia-ml.so.1 is not on this host".into());
}
let sym = |n: &core::ffi::CStr| {
let p = libc::dlsym(h, n.as_ptr());
if p.is_null() { Err(format!("no symbol {n:?}")) } else { Ok(p) }
};
let init: unsafe extern "C" fn() -> i32 = core::mem::transmute(sym(c"nvmlInit_v2")?);
let by_index: unsafe extern "C" fn(u32, *mut *mut core::ffi::c_void) -> i32 = core::mem::transmute(sym(c"nvmlDeviceGetHandleByIndex_v2")?);
let energy = core::mem::transmute::<*mut core::ffi::c_void, unsafe extern "C" fn(*mut core::ffi::c_void, *mut u64) -> i32>(sym(c"nvmlDeviceGetTotalEnergyConsumption")?);
let power = core::mem::transmute::<*mut core::ffi::c_void, unsafe extern "C" fn(*mut core::ffi::c_void, *mut u32) -> i32>(sym(c"nvmlDeviceGetPowerUsage")?);
if init() != 0 {
return Err("nvmlInit_v2 failed".into());
}
let mut dev = core::ptr::null_mut();
if by_index(0, &mut dev) != 0 {
return Err("no GPU 0".into());
}
let n = Nvml { dev, energy, power };
n.read()?;
Ok(n)
}
}
/// The cumulative energy, mJ.
pub fn read(&self) -> Result<u64, String> {
let mut v = 0u64;
// SAFETY: a device handle NVML gave, and a valid out-pointer.
let r = unsafe { (self.energy)(self.dev, &mut v) };
if r == 0 { Ok(v) } else { Err(format!("nvmlDeviceGetTotalEnergyConsumption returned {r}")) }
}
/// The board's power draw now, mW (read only outside the windows).
pub fn power_mw(&self) -> Option<u32> {
let mut v = 0u32;
// SAFETY: as above.
(unsafe { (self.power)(self.dev, &mut v) } == 0).then_some(v)
}
}
pub struct Meters {
pub rapl: Result<Rapl, String>,
pub nvml: Result<Nvml, String>,
}
pub struct Start {
t: Instant,
rapl: Option<u64>,
nvml: Option<u64>,
}
impl Meters {
pub fn open() -> Meters {
Meters { rapl: Rapl::open(), nvml: Nvml::open() }
}
pub fn describe(&self) -> String {
format!(
"rapl={} nvml={}",
self.rapl.as_ref().map_or_else(|e| format!("unavailable ({e})"), |_| "package-0".into()),
self.nvml.as_ref().map_or_else(|e| format!("unavailable ({e})"), |_| "gpu0 total energy".into())
)
}
pub fn start(&self) -> Start {
let nvml = self.nvml.as_ref().ok().and_then(|n| n.read().ok());
let rapl = self.rapl.as_ref().ok().and_then(|r| r.read().ok());
Start { t: Instant::now(), rapl, nvml }
}
/// `(seconds, package J, GPU board J)`.
pub fn finish(&self, s: Start) -> (f64, Option<f64>, Option<f64>) {
let rapl = self.rapl.as_ref().ok().and_then(|r| r.read().ok());
let nvml = self.nvml.as_ref().ok().and_then(|n| n.read().ok());
let secs = s.t.elapsed().as_secs_f64();
let pkg = match (&self.rapl, s.rapl, rapl) {
(Ok(r), Some(a), Some(b)) => Some(r.delta(a, b) as f64 / 1e6),
_ => None,
};
let gpu = match (s.nvml, nvml) {
(Some(a), Some(b)) if b >= a => Some((b - a) as f64 / 1e3),
_ => None,
};
(secs, pkg, gpu)
}
/// One idle window: `(package W, GPU board W, the board's draw now in W)`.
pub fn idle(&self, d: Duration) -> (Option<f64>, Option<f64>, Option<f64>) {
let s = self.start();
std::thread::sleep(d);
let (secs, pkg, gpu) = self.finish(s);
let draw = self.nvml.as_ref().ok().and_then(|n| n.power_mw()).map(|mw| mw as f64 / 1e3);
(pkg.map(|j| j / secs), gpu.map(|j| j / secs), draw)
}
}
}
// ---------------------------------------------------------------- energy, macOS
#[cfg(target_os = "macos")]
mod meter {
use jouleclaw_energy::process_gpu::GpuSnapshot;
use jouleclaw_energy::span::{ProcessSpanMeter, SpanEnergy, SpanSource, SpanStart, SpanUses};
pub struct Meters {
pub span: ProcessSpanMeter,
}
pub struct Start {
rail: Option<GpuSnapshot>,
span: SpanStart,
}
impl Meters {
/// The span meter; `gpu_load` puts some GPU work on the device from this process
/// for its liveness probe.
pub fn open(gpu_load: impl FnMut()) -> Meters {
Meters { span: ProcessSpanMeter::new(gpu_load) }
}
pub fn describe(&self) -> String {
format!("process span meter: metered={} fallback={}", self.span.is_metered(), self.span.fallback_name())
}
pub fn start(&self) -> Start {
let rail = self.span.gpu().snapshot().ok().filter(|s| s.rail_current);
Start { rail, span: self.span.start() }
}
/// The span's figure, and the GPU rail over the window: every GPU client's
/// energy and the GPU's idle draw, with this process's share of every client's
/// GPU time. The rail is an upper bound on this process's GPU energy, never its
/// figure.
pub fn finish(&self, s: Start) -> String {
let span = match self.span.finish(s.span, SpanUses::CpuGpu) {
Ok(e) => describe(&e),
Err(e) => format!("energy=none ({e})"),
};
let end = self.span.gpu().snapshot().ok().filter(|s| s.rail_current);
let rail = match (s.rail, end) {
(Some(a), Some(b)) => {
let j: u64 = b.rails_nj.iter().filter_map(|(n, e)| a.rails_nj.iter().find(|(m, _)| m == n).map(|(_, x)| e.saturating_sub(*x))).sum();
let pid = std::process::id() as i32;
let delta = |pick: &dyn Fn(i32) -> bool| -> u64 {
b.clients.iter().filter(|c| pick(c.pid)).map(|c| c.gpu_ns.saturating_sub(a.clients.iter().find(|x| x.id == c.id).map_or(0, |x| x.gpu_ns))).sum()
};
let (own, all) = (delta(&|p| p == pid), delta(&|_| true));
format!(" gpu_rail_window_j={:.4} own_gpu_time_share={:.3}", j as f64 / 1e9, if all > 0 { own as f64 / all as f64 } else { 0.0 })
}
_ => " gpu_rail_window_j=none".into(),
};
span + &rail
}
}
fn describe(e: &SpanEnergy) -> String {
match &e.source {
SpanSource::Process { cpu_uj, gpu } => format!(
"energy_j={:.4} process_cpu_j={:.4} process_gpu_j={:.4} provenance={:?} coverage={:?}",
e.uj as f64 / 1e6,
*cpu_uj as f64 / 1e6,
gpu.uj as f64 / 1e6,
e.provenance,
e.coverage
),
SpanSource::CpuOnly { cpu_uj, why } => format!("energy=partial process_cpu_j={:.4} process_gpu=none ({why}) coverage={:?}", *cpu_uj as f64 / 1e6, e.coverage),
SpanSource::Fallback { meter, why } => format!("energy=none: fell back to {meter} ({why})"),
}
}
}
#[cfg(not(any(target_os = "linux", target_os = "macos")))]
mod meter {
pub struct Meters;
impl Meters {
pub fn start(&self) {}
}
}
fn percentile(sorted: &[f64], p: f64) -> f64 {
let i = ((sorted.len() - 1) as f64 * p).round() as usize;
sorted[i]
}
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 secs: f64 = arg("--secs").map_or(3.0, |s| s.parse().expect("--secs"));
let min_runs: usize = arg("--min-runs").map_or(5, |s| s.parse().expect("--min-runs"));
let only = arg("--only");
let d = wai_ops_gpu::blocking::discover();
for (a, why) in &d.refused {
println!("ops-gpu bench: adapter refused: key={} {} reason={why:?}", a.key(), a.describe());
}
let gpus: Vec<Arc<wai_ops_gpu::Gpu>> = d.found.into_iter().map(Arc::new).collect();
let cases: Vec<Case> = cases().into_iter().filter(|c| only.as_deref().is_none_or(|o| o == c.cap)).collect();
let threads = std::thread::available_parallelism().map_or(0, |n| n.get());
println!("ops-gpu bench: host_threads={threads} secs={secs} min_runs={min_runs} cases={} adapters={}", cases.len(), gpus.len());
#[cfg(target_os = "linux")]
let meters = {
let m = meter::Meters::open();
let (pkg, gpu, draw) = m.idle(std::time::Duration::from_secs_f64(secs));
println!(
"ops-gpu bench: meters: {} idle_package_w={} idle_gpu_board_w={} gpu_draw_now_w={}",
m.describe(),
pkg.map_or("-".into(), |w| format!("{w:.2}")),
gpu.map_or("-".into(), |w| format!("{w:.2}")),
draw.map_or("-".into(), |w| format!("{w:.2}"))
);
m
};
#[cfg(not(any(target_os = "linux", target_os = "macos")))]
let meters = meter::Meters;
#[cfg(target_os = "macos")]
let meters = {
let probe = cases.iter().find(|c| c.cap == "wai.neural.int_hyper").or(cases.first()).map(|c| &c.run);
let g = gpus.first().cloned();
let m = meter::Meters::open(|| {
if let (Some(g), Some(run)) = (&g, probe) {
for _ in 0..20 {
let _ = gpu::with_gpu(g, run);
}
}
});
println!("ops-gpu bench: meters: {}", m.describe());
m
};
for c in &cases {
let want = (c.run)();
let mut paths: Vec<(String, Option<Arc<wai_ops_gpu::Gpu>>)> = vec![("cpu".into(), None)];
for g in &gpus {
paths.push((format!("gpu:{}", g.adapter().key()), Some(Arc::clone(g))));
}
for (label, g) in paths {
let once = || -> (Vec<u8>, u64) {
match &g {
None => ((c.run)(), 0),
Some(g) => {
let (o, st) = gpu::with_gpu(g, || (c.run)());
assert!(st.errors.is_empty(), "{}: device errors {:?}", c.name, st.errors);
(o, st.on_gpu())
}
}
};
for _ in 0..2 {
assert!(once().0 == want, "{} {}: the output differs from the CPU's", c.cap, c.name);
}
let mut times = Vec::new();
let mut graphs = 0u64;
let t0 = Instant::now();
let start = meters.start();
while times.len() < min_runs || t0.elapsed().as_secs_f64() < secs {
let t = Instant::now();
let (o, n) = once();
times.push(t.elapsed().as_secs_f64() * 1e3);
graphs += n;
assert!(o == want, "{} {}: the output differs from the CPU's", c.cap, c.name);
}
#[cfg(target_os = "linux")]
let energy = {
let (w, pkg, board) = meters.finish(start);
let per = |j: Option<f64>| j.map_or("-".into(), |j| format!("{:.4}", j / times.len() as f64));
format!("window_s={w:.2} package_j_per_decode={} gpu_board_j_per_decode={}", per(pkg), per(board))
};
#[cfg(target_os = "macos")]
let energy = format!("window_s={:.2} {} (window totals, over runs={})", t0.elapsed().as_secs_f64(), meters.finish(start), times.len());
#[cfg(not(any(target_os = "linux", target_os = "macos")))]
let energy = {
let () = start;
"energy=none".to_string()
};
let mut s = times.clone();
s.sort_by(|a, b| a.partial_cmp(b).unwrap());
println!(
"ops-gpu bench: cap={} vector={} path={} runs={} graphs_on_gpu_per_decode={} median_ms={:.3} p10_ms={:.3} p90_ms={:.3} {energy}",
c.cap,
c.name,
label,
times.len(),
graphs / times.len() as u64,
percentile(&s, 0.5),
percentile(&s, 0.1),
percentile(&s, 0.9),
);
}
}
}

1077
wai/wai-ops-gpu/Cargo.lock generated Normal file
View file

File diff suppressed because it is too large Load diff

33
wai/wai-ops-gpu/Cargo.toml Normal file
View file

@ -0,0 +1,33 @@
# `wai-ops/1-i32` on GPUs: the integer operator set's graph operators as WGSL compute
# kernels through wgpu (Metal, Vulkan, DX12 natively; a browser's WebGPU on
# wasm32-unknown-unknown), exact at both accumulator profiles, `acc32` in one 32-bit
# word per value and `acc64` in exact 64- and 128-bit arithmetic built from 32-bit
# words (extensions/integer-operators.md §7.3).
#
# A SEPARATE crate with its own lockfile, so that `wai-ops` stays `no_std` with no
# dependencies and `wai-rs` builds without wgpu by default; `wai-rs`'s `ops_gpu`
# feature is the only way it enters a decode.
[package]
name = "wai-ops-gpu"
version = "0.1.0"
edition = "2024"
rust-version = "1.97"
license = "Apache-2.0"
publish = false
description = "WAI integer operator set wai-ops/1-i32 on GPUs: WGSL kernels through wgpu, exact at acc32 and acc64"
[lib]
name = "wai_ops_gpu"
path = "src/lib.rs"
[dependencies]
wai-ops = { path = "../wai-ops" }
[target.'cfg(not(target_arch = "wasm32"))'.dependencies]
wgpu = { version = "30.0.1", default-features = false, features = ["std", "wgsl", "metal", "vulkan", "dx12"] }
[target.'cfg(all(target_arch = "wasm32", target_os = "unknown"))'.dependencies]
wgpu = { version = "30.0.1", default-features = false, features = ["std", "wgsl", "webgpu"] }
[profile.release]
debug = false

1
wai/wai-ops-gpu/browser/.gitignore vendored Normal file
View file

@ -0,0 +1 @@
pkg/

1077
wai/wai-ops-gpu/browser/Cargo.lock generated Normal file
View file

File diff suppressed because it is too large Load diff

24
wai/wai-ops-gpu/browser/Cargo.toml Normal file
View file

@ -0,0 +1,24 @@
# The ops-conformance graph vectors through a browser's WebGPU: wai-ops-gpu's kernels
# compiled by the browser's own WGSL compiler, held to the committed expectations and to
# the interpreter compiled to wasm32 in the same page. Built for wasm32-unknown-unknown
# and loaded by index.html; run.mjs serves it and drives a headless browser.
[package]
name = "wai-ops-gpu-browser"
version = "0.0.0"
edition = "2024"
rust-version = "1.97"
license = "Apache-2.0"
publish = false
[lib]
crate-type = ["cdylib", "rlib"]
[dependencies]
wai-ops = { path = "../../wai-ops" }
wai-ops-gpu = { path = ".." }
# Pinned to the release the wasm-bindgen CLI must match.
wasm-bindgen = "=0.2.127"
wasm-bindgen-futures = "=0.4.77"
[profile.release]
debug = false

28
wai/wai-ops-gpu/browser/index.html Normal file
View file

@ -0,0 +1,28 @@
<!doctype html>
<meta charset="utf-8">
<title>wai-ops-gpu in a browser</title>
<p id="status">running</p>
<pre id="report"></pre>
<script type="module">
// Loads every graph-vector file of ops-conformance/ (the list run.mjs serves as
// files.json), runs them through wai-ops-gpu on this browser's WebGPU, and leaves the
// report in window.opsGpu for run.mjs to read.
import init, { run_corpus } from "./pkg/wai_ops_gpu_browser.js";
window.opsGpu = { done: false, report: "" };
try {
await init();
const files = await (await fetch("files.json")).json();
let bundle = "";
for (const f of files) {
const r = await fetch(`/ops-conformance/${f}`);
if (!r.ok) throw new Error(`${f}: ${r.status}`);
bundle += `### ${f}\n` + await r.text();
}
window.opsGpu.report = await run_corpus(bundle);
} catch (e) {
window.opsGpu.report = `browser webgpu: page error: ${e}`;
}
document.getElementById("report").textContent = window.opsGpu.report;
document.getElementById("status").textContent = "done";
window.opsGpu.done = true;
</script>

92
wai/wai-ops-gpu/browser/run.mjs Normal file
View file

@ -0,0 +1,92 @@
// The ops-conformance graph vectors through a browser's WebGPU.
//
// Build the page's module first:
//
// cd wai/wai-ops-gpu/browser
// cargo build --locked --release --target wasm32-unknown-unknown
// wasm-bindgen --target web --out-dir pkg --out-name wai_ops_gpu_browser \
// target/wasm32-unknown-unknown/release/wai_ops_gpu_browser.wasm
//
// then, with Playwright installed under PLAYWRIGHT_HOST:
//
// [WAI_OPS_GPU_REQUIRE=1] [WAI_OPS_GPU_BROWSER_CHANNEL=chrome] \
// PLAYWRIGHT_HOST=<dir with node_modules/playwright>/ node wai/wai-ops-gpu/browser/run.mjs
//
// It serves this directory and ../../ops-conformance on 127.0.0.1, opens index.html in
// a headless browser and prints the page's report. It fails on a page error, a
// mismatch or a device error, and, under WAI_OPS_GPU_REQUIRE=1, when the page found no
// adapter or the report has no `browser webgpu: key=` line with mismatches=0 errors=0.
// Headless, macOS offers a page a hardware adapter only when the browser draws through
// Metal (`--use-angle=metal`); WAI_OPS_GPU_BROWSER_ARGS (space-separated) replaces the
// launch arguments.
import http from "node:http";
import fs from "node:fs";
import path from "node:path";
import { fileURLToPath } from "node:url";
import { createRequire } from "node:module";
const HERE = path.dirname(fileURLToPath(import.meta.url));
const CORPUS = path.resolve(HERE, "../../ops-conformance");
const REQUIRE = process.env.WAI_OPS_GPU_REQUIRE === "1";
const PORT = Number(process.env.PORT || 8797);
const TIMEOUT_MS = Number(process.env.WAI_OPS_GPU_TIMEOUT_MS || 20 * 60 * 1000);
const ARGS = process.env.WAI_OPS_GPU_BROWSER_ARGS !== undefined
? process.env.WAI_OPS_GPU_BROWSER_ARGS.split(/\s+/).filter(Boolean)
: (process.platform === "darwin" ? ["--use-angle=metal"] : []).concat(["--enable-unsafe-webgpu"]);
const { chromium } = createRequire(process.env.PLAYWRIGHT_HOST || HERE + "/")("playwright");
const files = ["vectors", "graphs"].flatMap((d) =>
fs.readdirSync(path.join(CORPUS, d)).filter((f) => f.endsWith(".jsonl")).sort().map((f) => `${d}/${f}`));
const ROOTS = { "/ops-conformance/": CORPUS, "/browser/": HERE };
const MIME = { ".js": "text/javascript", ".wasm": "application/wasm", ".html": "text/html", ".json": "application/json", ".jsonl": "text/plain" };
const server = http.createServer((req, res) => {
const url = decodeURIComponent(req.url.split("?")[0]);
if (url === "/browser/files.json") {
res.writeHead(200, { "content-type": "application/json" });
return res.end(JSON.stringify(files));
}
const prefix = Object.keys(ROOTS).find((p) => url.startsWith(p));
const root = prefix && ROOTS[prefix];
const f = root && path.resolve(root, "." + url.slice(prefix.length - 1));
if (!f || !f.startsWith(root + path.sep) || !fs.existsSync(f) || !fs.statSync(f).isFile()) {
res.writeHead(404);
return res.end("not found");
}
res.writeHead(200, { "content-type": MIME[path.extname(f)] || "application/octet-stream" });
fs.createReadStream(f).pipe(res);
});
let browser;
try {
if (!fs.existsSync(path.join(HERE, "pkg", "wai_ops_gpu_browser.js"))) {
console.error("no pkg/wai_ops_gpu_browser.js; build it as this file's header says");
process.exit(2);
}
await new Promise((r) => server.listen(PORT, "127.0.0.1", r));
const channel = process.env.WAI_OPS_GPU_BROWSER_CHANNEL || undefined;
browser = await chromium.launch({ args: ARGS, channel });
console.log(`browser: ${channel || "chromium"} ${browser.version()} args=${ARGS.join(" ")}`);
const page = await browser.newPage();
const errs = [];
page.on("pageerror", (e) => errs.push(String(e)));
await page.goto(`http://127.0.0.1:${PORT}/browser/index.html`, { waitUntil: "domcontentloaded", timeout: 30000 });
await page.waitForFunction(() => window.opsGpu && window.opsGpu.done === true, null, { timeout: TIMEOUT_MS, polling: 1000 });
const report = await page.evaluate(() => window.opsGpu.report);
const info = await page.evaluate(async () => {
if (!navigator.gpu) return "navigator.gpu is absent";
const a = await navigator.gpu.requestAdapter();
if (!a) return "requestAdapter gave none";
const i = a.info || {};
return `vendor=${i.vendor} architecture=${i.architecture} device=${i.device} description=${i.description}`;
});
console.log(`adapter info (navigator.gpu): ${info}`);
console.log(report);
for (const e of errs) console.log(`page error: ${e}`);
const ok = report.split("\n").some((l) => /^browser webgpu: key=\S+ .* mismatches=0 errors=0$/.test(l));
const bad = errs.length > 0 || /MISMATCH|DEVICE ERROR|page error|did not parse/.test(report) ||
report.split("\n").some((l) => /^browser webgpu: key=/.test(l) && !/mismatches=0 errors=0$/.test(l));
if (bad || (REQUIRE && !ok)) process.exitCode = 1;
} finally {
if (browser) await browser.close();
server.close();
}

134
wai/wai-ops-gpu/browser/src/lib.rs Normal file
View file

@ -0,0 +1,134 @@
//! The graph vectors of `ops-conformance/` on a browser's WebGPU adapter: the same
//! kernels as natively, compiled by the browser's own WGSL compiler, each decode held to
//! the committed expectation and to the interpreter running in the page (wasm32).
//!
//! [`run_corpus`] takes every vector file as one text, each file after a line
//! `### <name>`, and returns the report: one line per adapter, `browser webgpu: …`,
//! then a line per mismatch or device error (at most 20), or `browser webgpu: no adapter
//! (<why>)`.
use wai_ops::json::{self, Value, unhex};
use wai_ops::vectors::read_tensor;
use wai_ops::{Profile, Sequence, Tensor};
use wai_ops_gpu::{GpuSequence, Width, support};
use wasm_bindgen::prelude::*;
#[derive(Clone, Debug, PartialEq, Eq)]
enum Expect {
Outputs(Vec<Tensor>),
Refuse(String),
}
struct Decode {
inputs: Vec<Tensor>,
scalars: Vec<i64>,
expect: Expect,
}
struct Vector {
id: String,
graph: Vec<u8>,
decodes: Vec<Decode>,
}
fn parse(text: &str) -> Result<Vec<Vector>, String> {
let mut out = Vec::new();
for line in text.split_terminator('\n') {
if line.starts_with("### ") || line.is_empty() {
continue;
}
let v = json::parse(line.as_bytes()).map_err(|e| format!("JSON at byte {}", e.at))?;
let s = |k: &str| v.get(k).and_then(Value::as_str).map(str::to_string).ok_or(format!("no {k}"));
let id = s("id")?;
let graph = unhex(&s("graph")?).ok_or("graph hex")?;
let mut decodes = Vec::new();
for d in v.get("decodes").and_then(Value::as_arr).ok_or("decodes")? {
let inputs = d.get("inputs").and_then(Value::as_arr).ok_or("inputs")?.iter().map(read_tensor).collect::<Result<Vec<_>, _>>()?;
let scalars = d.get("scalars").and_then(Value::as_arr).ok_or("scalars")?.iter().map(|x| x.as_dec_i64().ok_or("scalar")).collect::<Result<Vec<_>, _>>()?;
let e = d.get("expect").ok_or("expect")?;
let expect = match e.get("refuse") {
Some(c) => Expect::Refuse(c.as_str().ok_or("refuse")?.to_string()),
None => Expect::Outputs(e.get("outputs").and_then(Value::as_arr).ok_or("outputs")?.iter().map(read_tensor).collect::<Result<Vec<_>, _>>()?),
};
decodes.push(Decode { inputs, scalars, expect });
}
out.push(Vector { id, graph, decodes });
}
Ok(out)
}
/// Run every vector on every adapter the page is offered.
#[wasm_bindgen]
pub async fn run_corpus(bundle: String) -> String {
let vs = match parse(&bundle) {
Ok(v) => v,
Err(e) => return format!("browser webgpu: the corpus did not parse: {e}"),
};
let d = wai_ops_gpu::discover().await;
let mut lines = Vec::new();
for (a, why) in &d.refused {
lines.push(format!("browser webgpu: adapter refused: key={} {} reason={why:?}", a.key(), a.describe()));
}
if d.found.is_empty() {
lines.push(format!("browser webgpu: no adapter ({} refused)", d.refused.len()));
return lines.join("\n");
}
for gpu in &d.found {
let (mut v32, mut d32, mut v64, mut d64, mut cpu, mut refused, mut refusals) = (0u32, 0u32, 0u32, 0u32, 0u32, 0u32, 0u32);
let mut bad: Vec<String> = Vec::new();
let (mut errors, mut mismatches) = (0u32, 0u32);
for v in &vs {
let Ok(l) = wai_ops::graph::load(&v.graph, Profile::Acc64) else {
refused += 1;
continue;
};
if support(&l).is_err() {
cpu += 1;
continue;
}
let mut widths = vec![Width::Acc64];
if l.meets(Profile::Acc32) {
widths.push(Width::Acc32);
}
for w in widths {
if w == Width::Acc32 { v32 += 1 } else { v64 += 1 }
let mut seq = GpuSequence::new();
let mut iseq = Sequence::new();
for (k, dd) in v.decodes.iter().enumerate() {
if w == Width::Acc32 { d32 += 1 } else { d64 += 1 }
refusals += matches!(dd.expect, Expect::Refuse(_)) as u32;
let interp = match wai_ops::interp::run(&l, &dd.inputs, &dd.scalars, &mut iseq) {
Ok(t) => Expect::Outputs(t),
Err(r) => Expect::Refuse(r.code.as_str().to_string()),
};
let got = match gpu.run_at(&l, w, &dd.inputs, &dd.scalars, &mut seq).await {
Ok(Ok(t)) => Expect::Outputs(t),
Ok(Err(r)) => Expect::Refuse(r.code.as_str().to_string()),
Err(e) => {
errors += 1;
if bad.len() < 20 {
bad.push(format!("browser webgpu: DEVICE ERROR {} decode {k} at {}: {e}", v.id, w.as_str()));
}
break;
}
};
if got != dd.expect || interp != dd.expect {
mismatches += 1;
if bad.len() < 20 {
bad.push(format!("browser webgpu: MISMATCH {} decode {k} at {} (the interpreter in the page {} the committed result)", v.id, w.as_str(), if interp == dd.expect { "gives" } else { "does NOT give" }));
}
break;
}
}
}
}
let a = gpu.adapter();
lines.push(format!(
"browser webgpu: key={} {} acc32 vectors={v32} decodes={d32}; acc64 vectors={v64} decodes={d64}; refusals={refusals}; cpu={cpu} load_refused={refused} mismatches={mismatches} errors={errors}",
a.key(),
a.describe(),
));
lines.extend(bad);
}
lines.join("\n")
}

534
wai/wai-ops-gpu/src/device.rs Normal file
View file

@ -0,0 +1,534 @@
//! Adapters and devices: discovery, the pipelines, error scopes, and waiting.
//!
//! [`discover`] opens every adapter of the platform's primary GPU interfaces (Metal,
//! Vulkan, DX12; a browser's WebGPU on `wasm32-unknown-unknown`), compiles the kernels
//! on it and runs the arithmetic probe ([`crate::probe`]) against the host. An adapter
//! that cannot host the kernels, or whose arithmetic disagrees with the host's on any
//! word, is listed with the reason and never runs a graph.
use std::borrow::Cow;
use std::collections::HashMap;
use std::num::NonZeroU64;
use std::sync::{Arc, Mutex, MutexGuard};
use core::future::Future;
use core::pin::Pin;
use core::task::{Context, Poll, Waker};
use crate::kernels::{OP_ENTRIES, REC_WORDS, Sources};
/// How long a native wait on one submission may take.
pub const POLL_TIMEOUT: core::time::Duration = core::time::Duration::from_secs(120);
/// Why the device did not give a result. A graph or a decode that this executor does
/// not run is [`GpuError::Fallback`]; everything else is the device's.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum GpuError {
/// The graph or the decode is outside what this executor runs (`crate::Fallback`):
/// the caller runs it on the CPU and says so.
Fallback(crate::Fallback),
/// The device reported an error, was lost, or a wait failed.
Device(String),
/// A native wait did not finish within [`POLL_TIMEOUT`].
Timeout,
/// A native future was still pending after one poll.
Pending,
}
impl core::fmt::Display for GpuError {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
GpuError::Fallback(r) => write!(f, "runs on the CPU: {r}"),
GpuError::Device(m) => write!(f, "device: {m}"),
GpuError::Timeout => write!(f, "a wait on the device timed out"),
GpuError::Pending => write!(f, "a native future was pending after one poll"),
}
}
}
impl std::error::Error for GpuError {}
/// What the platform reports about an adapter, as a run records it.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Adapter {
/// `metal`, `vulkan`, `dx12` or `webgpu`.
pub backend: String,
pub name: String,
pub vendor: u32,
pub device: u32,
/// `discrete`, `integrated`, `virtual`, `cpu` or `other`.
pub device_type: String,
pub driver: String,
pub driver_info: String,
/// The granted storage-binding and buffer limits, in bytes.
pub max_storage_binding: u64,
pub max_buffer: u64,
}
impl Adapter {
/// `gpu-<backend>-<vendor>` or, for a software adapter, `swgpu-<backend>-<vendor>`,
/// the vendor as lowercase hexadecimal.
pub fn key(&self) -> String {
let class = if self.device_type == "cpu" { "swgpu" } else { "gpu" };
format!("{class}-{}-{:x}", self.backend, self.vendor)
}
/// Whether the adapter is a GPU: not a software rasteriser.
pub fn is_gpu(&self) -> bool {
self.device_type != "cpu"
}
/// One line for a report: backend, adapter, type, driver.
pub fn describe(&self) -> String {
format!(
"backend={} adapter={:?} type={} vendor={:#x} device={:#x} driver={:?} driver_info={:?}",
self.backend, self.name, self.device_type, self.vendor, self.device, self.driver, self.driver_info
)
}
}
fn backend_label(b: wgpu::Backend) -> Option<&'static str> {
match b {
wgpu::Backend::Metal => Some("metal"),
wgpu::Backend::Vulkan => Some("vulkan"),
wgpu::Backend::Dx12 => Some("dx12"),
wgpu::Backend::BrowserWebGpu => Some("webgpu"),
_ => None,
}
}
fn device_type_label(t: wgpu::DeviceType) -> &'static str {
match t {
wgpu::DeviceType::DiscreteGpu => "discrete",
wgpu::DeviceType::IntegratedGpu => "integrated",
wgpu::DeviceType::VirtualGpu => "virtual",
wgpu::DeviceType::Cpu => "cpu",
wgpu::DeviceType::Other => "other",
}
}
#[derive(Default)]
pub(crate) struct Sink {
messages: Vec<String>,
lost: Option<String>,
}
pub(crate) fn lock<T>(m: &Mutex<T>) -> MutexGuard<'_, T> {
m.lock().unwrap_or_else(|p| p.into_inner())
}
/// The compiled kernels on one device.
pub(crate) struct Pipelines {
pub layout: wgpu::BindGroupLayout,
/// `ops32`'s entry points, in [`OP_ENTRIES`] order.
pub ops32: Vec<wgpu::ComputePipeline>,
/// `ops64`'s.
pub ops64: Vec<wgpu::ComputePipeline>,
pub mv: wgpu::ComputePipeline,
pub probe: wgpu::ComputePipeline,
}
/// One adapter opened as a device, its kernels compiled and its arithmetic probed.
pub struct Gpu {
pub(crate) device: wgpu::Device,
pub(crate) queue: wgpu::Queue,
pub(crate) pipes: Pipelines,
pub(crate) adapter: Adapter,
pub(crate) stride: u64,
pub(crate) dummy: wgpu::Buffer,
pub(crate) sink: Arc<Mutex<Sink>>,
pub(crate) cache: Mutex<Vec<crate::prepare::Cached>>,
pub(crate) pool: Mutex<HashMap<u64, Vec<wgpu::Buffer>>>,
pub(crate) label: String,
}
impl core::fmt::Debug for Gpu {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("Gpu").field("adapter", &self.adapter).field("kernels", &self.label).finish()
}
}
/// What discovery found: the adapters that passed, and why each other one did not.
#[derive(Debug, Default)]
pub struct Discovery {
pub found: Vec<Gpu>,
pub refused: Vec<(Adapter, String)>,
}
fn instance() -> wgpu::Instance {
let mut d = wgpu::InstanceDescriptor::new_without_display_handle();
d.backends = wgpu::Backends::PRIMARY;
d.flags = wgpu::InstanceFlags::empty();
d.backend_options.dx12.shader_compiler = wgpu::Dx12Compiler::Fxc;
wgpu::Instance::new(d)
}
/// Every adapter the platform offers, each opened and probed in this process.
pub async fn discover() -> Discovery {
discover_with(&Sources::shipped(), "shipped").await
}
/// [`discover`], with the kernels compiled from `src` and labelled `label` (a test's
/// mutant). The probe runs `src.probe`.
pub async fn discover_with(src: &Sources, label: &str) -> Discovery {
let mut d = Discovery::default();
if wgpu::Instance::enabled_backend_features().is_empty() {
return d;
}
let inst = instance();
#[cfg(not(target_arch = "wasm32"))]
let adapters = inst.enumerate_adapters(wgpu::Backends::PRIMARY).await;
#[cfg(target_arch = "wasm32")]
let adapters: Vec<wgpu::Adapter> = inst.request_adapter(&wgpu::RequestAdapterOptions::default()).await.into_iter().collect();
for a in adapters {
match open(a, src, label).await {
Ok(g) => d.found.push(g),
Err(e) => d.refused.push(*e),
}
}
d
}
fn identity(info: &wgpu::AdapterInfo, limits: &wgpu::Limits) -> Adapter {
Adapter {
backend: backend_label(info.backend).map_or_else(|| format!("{:?}", info.backend).to_lowercase(), str::to_string),
name: info.name.clone(),
vendor: info.vendor,
device: info.device,
device_type: device_type_label(info.device_type).into(),
driver: info.driver.clone(),
driver_info: info.driver_info.clone(),
max_storage_binding: limits.max_storage_buffer_binding_size,
max_buffer: limits.max_buffer_size,
}
}
async fn open(adapter: wgpu::Adapter, src: &Sources, label: &str) -> Result<Gpu, Box<(Adapter, String)>> {
let info = adapter.get_info();
let offered = adapter.limits();
let mut who = identity(&info, &offered);
if backend_label(info.backend).is_none() {
return Err(Box::new((who, format!("backend {:?} is not one the kernels run on", info.backend))));
}
if !adapter.get_downlevel_capabilities().flags.contains(wgpu::DownlevelFlags::COMPUTE_SHADERS) {
return Err(Box::new((who, "no compute shaders".into())));
}
// The WebGPU default limits (eight storage buffers per stage; the kernels bind five),
// with the two buffer limits raised to the adapter's.
let mut short: Option<String> = None;
wgpu::Limits::default().check_limits_with_fail_fn(&offered, true, |name, needed, has| {
short.get_or_insert_with(|| format!("limit {name}: the kernels need {needed}, the adapter offers {has}"));
});
if let Some(why) = short {
return Err(Box::new((who, why)));
}
let required_limits = wgpu::Limits {
max_storage_buffer_binding_size: offered.max_storage_buffer_binding_size,
max_buffer_size: offered.max_buffer_size,
..wgpu::Limits::default()
};
let descriptor = wgpu::DeviceDescriptor { label: Some("wai-ops-gpu"), required_features: wgpu::Features::empty(), required_limits, ..Default::default() };
let (device, queue) = match adapter.request_device(&descriptor).await {
Ok(dq) => dq,
Err(e) => return Err(Box::new((who, format!("request_device: {e}")))),
};
let sink = Arc::new(Mutex::new(Sink::default()));
let errors = Arc::clone(&sink);
device.on_uncaptured_error(Arc::new(move |e: wgpu::Error| lock(&errors).messages.push(e.to_string())));
let lost = Arc::clone(&sink);
device.set_device_lost_callback(move |reason, message| {
let mut s = lock(&lost);
if s.lost.is_none() {
s.lost = Some(format!("device lost ({reason:?}): {message}"));
}
});
let pipes = match pipelines(&device, src).await {
Ok(p) => p,
Err(e) => return Err(Box::new((who, format!("the kernels did not compile: {e}")))),
};
let granted = device.limits();
who.max_storage_binding = granted.max_storage_buffer_binding_size;
who.max_buffer = granted.max_buffer_size;
let align = u64::from(granted.min_uniform_buffer_offset_alignment).max(4);
let stride = ((REC_WORDS as u64 * 4).div_ceil(align)) * align;
let dummy = device.create_buffer(&wgpu::BufferDescriptor { label: Some("wai-ops-gpu dummy"), size: 16, usage: wgpu::BufferUsages::STORAGE, mapped_at_creation: false });
let gpu = Gpu {
device,
queue,
pipes,
adapter: who,
stride,
dummy,
sink,
cache: Mutex::new(Vec::new()),
pool: Mutex::new(HashMap::new()),
label: label.to_string(),
};
match crate::probe::check(&gpu).await {
Ok(()) => Ok(gpu),
Err(why) => Err(Box::new((gpu.adapter.clone(), why))),
}
}
/// The three error scopes a step of work runs inside, popped innermost first.
pub(crate) struct Scopes {
internal: wgpu::ErrorScopeGuard,
validation: wgpu::ErrorScopeGuard,
memory: wgpu::ErrorScopeGuard,
}
impl Scopes {
pub(crate) fn push(device: &wgpu::Device) -> Scopes {
let memory = device.push_error_scope(wgpu::ErrorFilter::OutOfMemory);
let validation = device.push_error_scope(wgpu::ErrorFilter::Validation);
let internal = device.push_error_scope(wgpu::ErrorFilter::Internal);
Scopes { internal, validation, memory }
}
pub(crate) async fn close(self) -> Result<(), GpuError> {
let Scopes { internal, validation, memory } = self;
let internal = internal.pop();
let validation = validation.pop();
let memory = memory.pop();
let (internal, validation, memory) = (internal.await, validation.await, memory.await);
match memory.or(validation).or(internal) {
Some(e) => Err(GpuError::Device(e.to_string())),
None => Ok(()),
}
}
}
async fn pipelines(device: &wgpu::Device, src: &Sources) -> Result<Pipelines, GpuError> {
let storage = |binding: u32, read_only: bool| wgpu::BindGroupLayoutEntry {
binding,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Buffer { ty: wgpu::BufferBindingType::Storage { read_only }, has_dynamic_offset: false, min_binding_size: None },
count: None,
};
let scopes = Scopes::push(device);
let layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
label: Some("wai-ops-gpu"),
entries: &[
wgpu::BindGroupLayoutEntry {
binding: 0,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Buffer { ty: wgpu::BufferBindingType::Uniform, has_dynamic_offset: false, min_binding_size: NonZeroU64::new(REC_WORDS as u64 * 4) },
count: None,
},
storage(1, true),
storage(2, true),
storage(3, true),
storage(4, false),
storage(5, false),
],
});
let pl = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { label: Some("wai-ops-gpu"), bind_group_layouts: &[Some(&layout)], immediate_size: 0 });
let module = |label: &str, text: &str| device.create_shader_module(wgpu::ShaderModuleDescriptor { label: Some(label), source: wgpu::ShaderSource::Wgsl(Cow::Owned(text.to_string())) });
let m32 = module("wai-ops-gpu ops32", &src.ops32);
let m64 = module("wai-ops-gpu ops64", &src.ops64);
let mmv = module("wai-ops-gpu move", &src.mv);
let mpr = module("wai-ops-gpu probe", &src.probe);
let pipe = |m: &wgpu::ShaderModule, entry: &str| {
device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
label: Some(entry),
layout: Some(&pl),
module: m,
entry_point: Some(entry),
compilation_options: Default::default(),
cache: None,
})
};
let ops32 = OP_ENTRIES.iter().map(|e| pipe(&m32, e)).collect();
let ops64 = OP_ENTRIES.iter().map(|e| pipe(&m64, e)).collect();
let mv = pipe(&mmv, "mv");
let probe = pipe(&mpr, "probe");
scopes.close().await?;
Ok(Pipelines { layout, ops32, ops64, mv, probe })
}
impl Gpu {
/// The adapter, as the platform reports it.
pub fn adapter(&self) -> &Adapter {
&self.adapter
}
/// The label of the kernels this device runs: `shipped`, or a test's mutant.
pub fn kernels(&self) -> &str {
&self.label
}
/// An error if the device has been lost, or reported anything outside a scope since
/// the last check.
pub(crate) fn settled(&self) -> Result<(), GpuError> {
let mut s = lock(&self.sink);
if let Some(why) = &s.lost {
return Err(GpuError::Device(why.clone()));
}
if s.messages.is_empty() {
return Ok(());
}
Err(GpuError::Device(std::mem::take(&mut s.messages).join("; ")))
}
pub(crate) fn buffer(&self, label: &str, size: u64, usage: wgpu::BufferUsages) -> wgpu::Buffer {
self.device.create_buffer(&wgpu::BufferDescriptor { label: Some(label), size, usage, mapped_at_creation: false })
}
/// A storage buffer of at least `size` bytes from the pool, or a new one.
pub(crate) fn take(&self, size: u64) -> wgpu::Buffer {
let size = size.div_ceil(256).max(1) * 256;
if let Some(b) = lock(&self.pool).get_mut(&size).and_then(Vec::pop) {
return b;
}
self.buffer("wai-ops-gpu tensor", size, wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC | wgpu::BufferUsages::COPY_DST)
}
/// Give buffers back to the pool, keeping at most 64 of each size.
pub(crate) fn give(&self, bufs: Vec<wgpu::Buffer>) {
let mut p = lock(&self.pool);
for b in bufs {
let v = p.entry(b.size()).or_default();
if v.len() < 64 {
v.push(b);
}
}
}
/// Read `len` bytes of `staging` (MAP_READ) once submission `index` has run.
pub(crate) async fn read(&self, staging: &wgpu::Buffer, len: u64, _index: wgpu::SubmissionIndex) -> Result<Vec<u8>, GpuError> {
let slot = Arc::new(Mutex::new(MapSlot::default()));
let filled = Arc::clone(&slot);
staging.map_async(wgpu::MapMode::Read, 0..len, move |r| {
let waker = {
let mut s = lock(&filled);
s.result = Some(r);
s.waker.take()
};
if let Some(w) = waker {
w.wake();
}
});
Mapped {
slot,
#[cfg(not(target_arch = "wasm32"))]
device: self.device.clone(),
#[cfg(not(target_arch = "wasm32"))]
sink: Arc::clone(&self.sink),
#[cfg(not(target_arch = "wasm32"))]
index: _index,
}
.await?;
let out = {
let view = staging.get_mapped_range(0..len).map_err(|e| GpuError::Device(e.to_string()))?;
view.to_vec()
};
staging.unmap();
Ok(out)
}
}
#[derive(Default)]
struct MapSlot {
result: Option<Result<(), wgpu::BufferAsyncError>>,
waker: Option<Waker>,
}
/// Resolves when `map_async`'s callback has run. Natively it waits on its submission
/// inside the device poll, so it is ready on the first poll; in a browser the event
/// loop runs the callback, which wakes it.
struct Mapped {
slot: Arc<Mutex<MapSlot>>,
#[cfg(not(target_arch = "wasm32"))]
device: wgpu::Device,
#[cfg(not(target_arch = "wasm32"))]
sink: Arc<Mutex<Sink>>,
#[cfg(not(target_arch = "wasm32"))]
index: wgpu::SubmissionIndex,
}
impl Future for Mapped {
type Output = Result<(), GpuError>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
#[cfg(not(target_arch = "wasm32"))]
if let Err(e) = wait(&self.device, &self.sink, self.index.clone()) {
return Poll::Ready(Err(e));
}
let mut s = lock(&self.slot);
match s.result.take() {
Some(r) => Poll::Ready(r.map_err(|e| GpuError::Device(e.to_string()))),
None => {
s.waker = Some(cx.waker().clone());
Poll::Pending
}
}
}
}
#[cfg(not(target_arch = "wasm32"))]
fn wait(device: &wgpu::Device, sink: &Mutex<Sink>, index: wgpu::SubmissionIndex) -> Result<(), GpuError> {
let poll = || device.poll(wgpu::PollType::Wait { submission_index: Some(index), timeout: Some(POLL_TIMEOUT) });
match std::panic::catch_unwind(std::panic::AssertUnwindSafe(poll)) {
Ok(Ok(_)) => Ok(()),
Ok(Err(wgpu::PollError::Timeout)) => Err(GpuError::Timeout),
Ok(Err(e)) => Err(GpuError::Device(e.to_string())),
Err(payload) => {
let message = payload.downcast_ref::<String>().cloned().or_else(|| payload.downcast_ref::<&str>().map(|m| m.to_string())).unwrap_or_else(|| "no message".into());
let why = format!("the device poll failed: {message}");
let mut s = lock(sink);
if s.lost.is_none() {
s.lost = Some(why.clone());
}
Err(GpuError::Device(why))
}
}
}
/// The calls for a native caller without an async runtime. Natively every future of
/// the GPU interface is already resolved and a buffer map completes inside the device
/// poll that waits for it, so one poll resolves each call.
#[cfg(not(target_arch = "wasm32"))]
pub mod blocking {
use super::GpuError;
use core::future::Future;
use core::task::{Context, Poll, Waker};
/// Polls `f` once with a waker that does nothing; a future still pending is
/// [`GpuError::Pending`], never a spin.
pub fn resolve<F: Future>(f: F) -> Result<F::Output, GpuError> {
let mut f = core::pin::pin!(f);
match f.as_mut().poll(&mut Context::from_waker(Waker::noop())) {
Poll::Ready(v) => Ok(v),
Poll::Pending => Err(GpuError::Pending),
}
}
/// [`super::discover`].
pub fn discover() -> super::Discovery {
resolve(super::discover()).unwrap_or_default()
}
/// [`super::discover_with`].
pub fn discover_with(src: &crate::kernels::Sources, label: &str) -> super::Discovery {
resolve(super::discover_with(src, label)).unwrap_or_default()
}
/// [`crate::Gpu::run`].
pub fn run(
g: &crate::Gpu, l: &wai_ops::Loaded, inputs: &[wai_ops::Tensor], scalars: &[i64], seq: &mut crate::GpuSequence,
) -> Result<Result<Vec<wai_ops::Tensor>, wai_ops::Refusal>, GpuError> {
resolve(g.run(l, inputs, scalars, seq))?
}
/// [`crate::Gpu::run_at`].
pub fn run_at(
g: &crate::Gpu, l: &wai_ops::Loaded, width: crate::Width, inputs: &[wai_ops::Tensor], scalars: &[i64], seq: &mut crate::GpuSequence,
) -> Result<Result<Vec<wai_ops::Tensor>, wai_ops::Refusal>, GpuError> {
resolve(g.run_at(l, width, inputs, scalars, seq))?
}
/// [`crate::probe::run`].
pub fn probe(g: &crate::Gpu, recs: &[crate::probe::ProbeIn]) -> Result<Vec<[u32; crate::probe::OUT_WORDS]>, GpuError> {
resolve(crate::probe::run(g, recs))?
}
}

480
wai/wai-ops-gpu/src/exec.rs Normal file
View file

@ -0,0 +1,480 @@
//! One decode of a loaded graph on the device.
//!
//! The host makes every check the interpreter makes before evaluation, in its order and
//! with its codes (integer-operators §4, run steps 1–3; `wai_ops::interp::run`): the
//! data inputs and the state, each input's channels and sizes, the scalars, the
//! run-time plan, then each input's policy. The operators then run on the device, in
//! graph order, in one submission, each output's policy after its operator: a `clamp`
//! in place, a `check` setting its operator's flag. The outputs and the flags come back
//! in one read. The first operator whose flag is set refuses the decode (`field_range`),
//! as the interpreter would at that operator; operators after it ran speculatively and
//! their values are discarded (§6). The state of a sequence (§10.3) is the
//! interpreter's: a refused decode changes nothing.
use std::sync::Arc;
use wai_ops::graph::{Operand, Params, Policy, opcode::*};
use wai_ops::plan::run_plan;
use wai_ops::{Code, Loaded, Refusal, Site, Tensor};
use crate::device::{Gpu, GpuError, Scopes};
use crate::kernels::{Entry, REC_WORDS, Rec, WG, f, grid};
use crate::prepare::{Fallback, Prepared, Width, support};
/// The state of a sequence of decodes on the device (§10.3): the tensors carried into
/// the state inputs, or nothing before the first decode.
#[derive(Clone, Debug, Default)]
pub struct GpuSequence {
state: Option<Vec<Tensor>>,
}
impl GpuSequence {
pub fn new() -> GpuSequence {
GpuSequence { state: None }
}
}
/// A decode's outcome, or the device's error.
pub type RunResult = Result<Result<Vec<Tensor>, Refusal>, GpuError>;
fn apply_policy(p: Policy, lo: i64, hi: i64, d: &mut [i64]) -> bool {
match p {
Policy::Derived => true,
Policy::Clamp => {
for v in d.iter_mut() {
*v = (*v).clamp(lo, hi);
}
true
}
Policy::Check => d.iter().all(|&v| lo <= v && v <= hi),
}
}
/// Which pipeline a dispatch runs.
#[derive(Clone, Copy)]
enum Pipe {
Op(Entry),
Move,
}
/// One dispatch: its pipeline, record, and the tensors (or the dummy) it binds.
struct Dispatch {
pipe: Pipe,
rec: Rec,
a: Option<u32>,
b: Option<u32>,
m: Option<usize>,
o: u32,
}
/// The clamp bounds the 32-bit kernels read: `[lo, hi]` held to i32, which leaves the
/// result on an i32 value unchanged.
fn bounds32(lo: i64, hi: i64) -> (i64, i64) {
(lo.clamp(i32::MIN as i64, i32::MAX as i64), hi.clamp(i32::MIN as i64, i32::MAX as i64))
}
impl Gpu {
/// One decode of `l`, at the width [`support`] chooses. A graph [`support`] sends to
/// the CPU is [`GpuError::Fallback`]; the caller runs it there.
pub async fn run(&self, l: &Loaded, inputs: &[Tensor], scalars: &[i64], seq: &mut GpuSequence) -> RunResult {
let width = support(l).map_err(GpuError::Fallback)?;
self.run_inner(l, width, inputs, scalars, seq).await
}
/// One decode at `width`. `Width::Acc64` runs any graph of `wai-ops/1-i32`, since
/// every graph that loads meets `acc64`; `Width::Acc32` only one whose certificate
/// meets `acc32` ([`Fallback::Profile`] otherwise).
pub async fn run_at(&self, l: &Loaded, width: Width, inputs: &[Tensor], scalars: &[i64], seq: &mut GpuSequence) -> RunResult {
support(l).map_err(GpuError::Fallback)?;
if width == Width::Acc32 && !l.meets(wai_ops::Profile::Acc32) {
return Err(GpuError::Fallback(Fallback::Profile));
}
self.run_inner(l, width, inputs, scalars, seq).await
}
/// [`Gpu::run_at`] without the profile check: the 32-bit kernels on a graph its
/// certificate does not admit to them. For tests that show why the check exists.
#[doc(hidden)]
pub async fn run_unchecked(&self, l: &Loaded, width: Width, inputs: &[Tensor], scalars: &[i64], seq: &mut GpuSequence) -> RunResult {
support(l).map_err(GpuError::Fallback)?;
self.run_inner(l, width, inputs, scalars, seq).await
}
async fn run_inner(&self, l: &Loaded, width: Width, inputs: &[Tensor], scalars: &[i64], seq: &mut GpuSequence) -> RunResult {
let g = l.graph();
// Run steps 1–3 and the inputs' policies, as the interpreter makes them.
let n_data = g.inputs.iter().filter(|i| i.kind == 0).count();
if inputs.len() != n_data {
return Ok(Err(Refusal::new(Code::ModelMismatch, Site::Input(inputs.len().min(n_data) as u32))));
}
let first = inputs.first().map(|t| (t.h, t.w));
let mut data = inputs.iter();
let mut all: Vec<Tensor> = Vec::with_capacity(g.inputs.len());
let mut s = 0usize;
for rec in &g.inputs {
if rec.kind == 0 {
all.push(data.next().expect("one data input per data record").clone());
} else {
let t = match &seq.state {
Some(st) => st[s].clone(),
None => {
let (h0, w0) = first.expect("a state input's first decode takes the first data input's sizes");
Tensor::filled(g.tensors[rec.tensor as usize].channels, h0, w0, rec.init)
}
};
all.push(t);
s += 1;
}
}
for (k, t) in all.iter().enumerate() {
if t.data.len() as u64 != t.c as u64 * t.h as u64 * t.w as u64 {
return Ok(Err(Refusal::new(Code::ModelMismatch, Site::Input(k as u32))));
}
}
let shapes: Vec<(u64, u64, u64)> = all.iter().map(|t| (t.c as u64, t.h as u64, t.w as u64)).collect();
let hw = match run_plan(g, &shapes, scalars) {
Ok(h) => h,
Err(r) => return Ok(Err(r)),
};
for (k, (rec, t)) in g.inputs.iter().zip(all.iter_mut()).enumerate() {
let tr = &g.tensors[rec.tensor as usize];
if !apply_policy(tr.policy, tr.lo, tr.hi, &mut t.data) {
return Ok(Err(Refusal::new(Code::FieldRange, Site::Input(k as u32))));
}
}
let prep = self.prepared(l, width).await?;
let ew = width.words();
let elems = |t: u32| g.tensors[t as usize].channels as u64 * hw[t as usize].0 * hw[t as usize].1;
// Every tensor's buffer: an input's, or an operator's output's.
let mut bufs: Vec<Option<wgpu::Buffer>> = (0..g.tensors.len()).map(|_| None).collect();
let max = self.adapter.max_storage_binding;
let release = |b: Vec<Option<wgpu::Buffer>>| self.give(b.into_iter().flatten().collect());
for t in 0..g.tensors.len() as u32 {
let bytes = elems(t) * ew as u64 * 4;
if bytes > max {
release(bufs);
return Err(GpuError::Fallback(Fallback::Size(format!("tensor {t}'s {bytes} bytes exceed the adapter's storage binding of {max}"))));
}
}
let produced: Vec<bool> = {
let mut p = vec![false; g.tensors.len()];
for rec in &g.inputs {
p[rec.tensor as usize] = true;
}
for op in &g.ops {
for &t in &op.outs {
p[t as usize] = true;
}
}
p
};
for (t, b) in bufs.iter_mut().enumerate() {
if produced[t] {
*b = Some(self.take(elems(t as u32).max(1) * ew as u64 * 4));
}
}
let scopes = Scopes::push(&self.device);
// The inputs, after their policies.
for (rec, t) in g.inputs.iter().zip(&all) {
let mut raw = Vec::with_capacity(t.data.len() * ew as usize * 4);
for &v in &t.data {
match width {
Width::Acc32 => {
// An input of an acc32 graph has a declared range that fits i32,
// and its policy has held it there.
debug_assert!(i32::try_from(v).is_ok());
raw.extend_from_slice(&(v as i32).to_le_bytes());
}
Width::Acc64 => raw.extend_from_slice(&v.to_le_bytes()),
}
}
self.queue.write_buffer(bufs[rec.tensor as usize].as_ref().expect("an input has a buffer"), 0, &raw);
}
let ds = match dispatches(g, &prep, width, &hw, scalars) {
Ok(d) => d,
Err(fb) => {
drop(scopes);
release(bufs);
return Err(GpuError::Fallback(fb));
}
};
let uni = self.buffer("wai-ops-gpu records", self.stride * ds.len().max(1) as u64, wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST);
let mut raw = vec![0u8; (self.stride * ds.len().max(1) as u64) as usize];
for (k, d) in ds.iter().enumerate() {
let at = k * self.stride as usize;
raw[at..at + REC_WORDS * 4].copy_from_slice(&d.rec.bytes());
}
self.queue.write_buffer(&uni, 0, &raw);
let n_flags = g.ops.len().max(1) as u64;
let flags = self.buffer("wai-ops-gpu flags", n_flags * 4, wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC | wgpu::BufferUsages::COPY_DST);
let mut binds = Vec::with_capacity(ds.len());
for (k, d) in ds.iter().enumerate() {
let tb = |t: Option<u32>| t.map_or(&self.dummy, |t| bufs[t as usize].as_ref().expect("a bound tensor has a buffer"));
let mb = d.m.and_then(|i| prep.ops[i].buf.as_ref()).unwrap_or(&self.dummy);
binds.push(self.device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some("wai-ops-gpu dispatch"),
layout: &self.pipes.layout,
entries: &[
wgpu::BindGroupEntry {
binding: 0,
resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding { buffer: &uni, offset: k as u64 * self.stride, size: core::num::NonZeroU64::new(REC_WORDS as u64 * 4) }),
},
wgpu::BindGroupEntry { binding: 1, resource: tb(d.a).as_entire_binding() },
wgpu::BindGroupEntry { binding: 2, resource: tb(d.b).as_entire_binding() },
wgpu::BindGroupEntry { binding: 3, resource: mb.as_entire_binding() },
wgpu::BindGroupEntry { binding: 4, resource: tb(Some(d.o)).as_entire_binding() },
wgpu::BindGroupEntry { binding: 5, resource: flags.as_entire_binding() },
],
}));
}
// The read: every output, then the flags.
let mut offsets = Vec::with_capacity(g.outputs.len());
let mut total = 0u64;
for &t in &g.outputs {
offsets.push(total);
total += elems(t) * ew as u64 * 4;
}
let flags_at = total;
total += n_flags * 4;
let staging = self.buffer("wai-ops-gpu staging", total, wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST);
let mut enc = self.device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("wai-ops-gpu decode") });
enc.clear_buffer(&flags, 0, None);
{
let mut pass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor { label: Some("wai-ops-gpu decode"), timestamp_writes: None });
for (d, bg) in ds.iter().zip(&binds) {
let pipe = match (d.pipe, width) {
(Pipe::Op(e), Width::Acc32) => &self.pipes.ops32[e as usize],
(Pipe::Op(e), Width::Acc64) => &self.pipes.ops64[e as usize],
(Pipe::Move, _) => &self.pipes.mv,
};
let (gx, gy, _) = grid(d.rec.0[f::N], WG).expect("the grid was checked when the record was made");
pass.set_pipeline(pipe);
pass.set_bind_group(0, bg, &[]);
pass.dispatch_workgroups(gx, gy, 1);
}
}
for (k, &t) in g.outputs.iter().enumerate() {
let bytes = elems(t) * ew as u64 * 4;
enc.copy_buffer_to_buffer(bufs[t as usize].as_ref().expect("an output has a buffer"), 0, &staging, offsets[k], bytes);
}
enc.copy_buffer_to_buffer(&flags, 0, &staging, flags_at, n_flags * 4);
let index = self.queue.submit([enc.finish()]);
if let Err(e) = scopes.close().await {
release(bufs);
return Err(e);
}
let bytes = self.read(&staging, total, index).await;
release(bufs);
let bytes = bytes?;
self.settled()?;
let word = |at: usize| u32::from_le_bytes([bytes[at], bytes[at + 1], bytes[at + 2], bytes[at + 3]]);
for i in 0..g.ops.len() {
if word(flags_at as usize + 4 * i) != 0 {
return Ok(Err(Refusal::new(Code::FieldRange, Site::Op(i as u32))));
}
}
let mut outputs = Vec::with_capacity(g.outputs.len());
for (k, &t) in g.outputs.iter().enumerate() {
let n = elems(t) as usize;
let at = offsets[k] as usize;
let data: Vec<i64> = match width {
Width::Acc32 => (0..n).map(|e| word(at + 4 * e) as i32 as i64).collect(),
Width::Acc64 => (0..n).map(|e| (word(at + 8 * e) as u64 | (word(at + 8 * e + 4) as u64) << 32) as i64).collect(),
};
let (h, w) = hw[t as usize];
outputs.push(Tensor::new(g.tensors[t as usize].channels, h as u32, w as u32, data));
}
let next: Vec<Tensor> = g.inputs.iter().filter(|i| i.kind == 1).map(|i| outputs[i.src as usize].clone()).collect();
seq.state = Some(next);
Ok(Ok(outputs))
}
}
/// Every dispatch of one decode, in order: each operator's, then its outputs'
/// policies. A grid past the dispatch limit is a [`Fallback::Size`].
fn dispatches(g: &wai_ops::Graph, prep: &Arc<Prepared>, width: Width, hw: &[(u64, u64)], sc: &[i64]) -> Result<Vec<Dispatch>, Fallback> {
let ew = width.words();
let mut out: Vec<Dispatch> = Vec::new();
let ch = |t: u32| g.tensors[t as usize].channels;
let sz = |t: u32| (hw[t as usize].0 as u32, hw[t as usize].1 as u32);
let mut push = |pipe: Pipe, mut rec: Rec, a: Option<u32>, b: Option<u32>, m: Option<usize>, o: u32| -> Result<(), Fallback> {
let n = rec.0[f::N];
let (_, _, row) = grid(n, WG).ok_or_else(|| Fallback::Size(format!("{n} threads exceed one dispatch's grid")))?;
rec.set(f::ROW, row).set(f::EW, ew);
out.push(Dispatch { pipe, rec, a, b, m, o });
Ok(())
};
for (i, op) in g.ops.iter().enumerate() {
let x = op.ins[0];
let (c, (h, w)) = (ch(x), sz(x));
let o0 = op.outs[0];
let (co, (ho, wo)) = (ch(o0), sz(o0));
let po = &prep.ops[i].offs;
let mut r = Rec::new();
r.set(f::C, c).set(f::H, h).set(f::W, w).set(f::CO, co).set(f::HO, ho).set(f::WO, wo);
r.set(f::W_OFF, po.w_off).set(f::M_OFF, po.m_off).set(f::S_OFF, po.s_off).set(f::B_OFF, po.b_off);
r.set(f::W_WORDS, po.w_words).set(f::ROUND, po.round).set(f::BIAS, po.bias).set(f::MK, po.mk).set(f::SK, po.sk);
let out_n = co * ho * wo;
let in_n = c * h * w;
let mv = |r: Rec, mode: u32| {
let mut q = r;
q.set(f::MODE, mode);
q
};
match (&op.p, op.code) {
(Params::Conv { kh, kw, sh, sw, ph, pw, groups, .. }, _) => {
r.set(f::N, out_n).set(f::KH, *kh).set(f::KW, *kw).set(f::SH, *sh).set(f::SW, *sw).set(f::PH, *ph).set(f::PW, *pw);
r.set(f::CG, c / groups).set(f::OG, co / groups);
push(Pipe::Op(Entry::Conv), r, Some(x), None, Some(i), o0)?;
}
(Params::ConvT { kh, kw, sh, sw, ph, pw, .. }, _) => {
r.set(f::N, out_n).set(f::KH, *kh).set(f::KW, *kw).set(f::SH, *sh).set(f::SW, *sw).set(f::PH, *ph).set(f::PW, *pw);
push(Pipe::Op(Entry::ConvT), r, Some(x), None, Some(i), o0)?;
}
(Params::Rescale(_), _) => {
r.set(f::N, in_n);
push(Pipe::Op(Entry::Rescale), r, Some(x), None, Some(i), o0)?;
}
(Params::Chscale { scalar, shift, .. }, _) => {
let row = sc[*scalar as usize] as u32;
r.set(f::N, in_n).set(f::F, *shift as u32).set(f::M_OFF, po.m_off + 2 * row * c);
push(Pipe::Op(Entry::Chscale), r, Some(x), None, Some(i), o0)?;
}
(Params::Mul(_), _) => {
r.set(f::N, in_n).set(f::C2, ch(op.ins[1]));
push(Pipe::Op(Entry::Mul), r, Some(x), Some(op.ins[1]), Some(i), o0)?;
}
(Params::Gram { groups, .. }, _) => {
let d = c / groups;
r.set(f::N, c * d).set(f::K0, d).set(f::K1, h * w);
push(Pipe::Op(Entry::Gram), r, Some(x), Some(op.ins[1]), Some(i), o0)?;
}
(Params::Apply { groups, .. }, _) => {
let v = op.ins[1];
let (vh, vw) = sz(v);
let d = c / groups;
r.set(f::N, ch(v) * vh * vw).set(f::K0, d).set(f::K1, vh * vw);
push(Pipe::Op(Entry::Apply), r, Some(x), Some(v), Some(i), o0)?;
}
(Params::Igdn { gf, frac, .. }, _) => {
r.set(f::N, in_n).set(f::F, *gf).set(f::K0, *frac).set(f::K1, h * w);
push(Pipe::Op(Entry::Igdn), r, Some(x), None, Some(i), o0)?;
}
(Params::Lut16 { .. }, _) => {
r.set(f::N, in_n);
push(Pipe::Op(Entry::Lut16), r, Some(x), None, Some(i), o0)?;
}
(_, ADD) => {
r.set(f::N, in_n).set(f::MODE, 4);
push(Pipe::Op(Entry::Ew), r, Some(x), Some(op.ins[1]), None, o0)?;
}
(_, CHUNK_ADD) => {
r.set(f::N, out_n).set(f::MODE, 5).set(f::K0, out_n);
push(Pipe::Op(Entry::Ew), r, Some(x), None, None, o0)?;
}
(Params::Clamp { lo, hi }, _) => {
let (lo, hi) = if width == Width::Acc32 { bounds32(*lo, *hi) } else { (*lo, *hi) };
r.set(f::N, in_n).set(f::MODE, 0).set64(f::LO_X, lo).set64(f::HI_X, hi);
push(Pipe::Op(Entry::Ew), r, Some(x), None, None, o0)?;
}
(_, RELU) => {
r.set(f::N, in_n).set(f::MODE, 1);
push(Pipe::Op(Entry::Ew), r, Some(x), None, None, o0)?;
}
(_, ABS) => {
r.set(f::N, in_n).set(f::MODE, 2);
push(Pipe::Op(Entry::Ew), r, Some(x), None, None, o0)?;
}
(Params::Leaky { a, f: sh }, _) => {
r.set(f::N, in_n).set(f::MODE, 3).set64(f::A_X, *a).set(f::F, *sh);
push(Pipe::Op(Entry::Ew), r, Some(x), None, None, o0)?;
}
(_, CONCAT) => {
let mut at = 0u32;
for &t in &op.ins {
let (th, tw) = sz(t);
let n = ch(t) * th * tw;
let mut q = mv(r, 0);
q.set(f::N, n).set(f::SRC_OFF, 0).set(f::DST_OFF, at);
push(Pipe::Move, q, Some(t), None, None, o0)?;
at += n;
}
}
(Params::Split { start, count }, _) => {
let mut q = mv(r, 0);
q.set(f::N, count * h * w).set(f::SRC_OFF, start * h * w).set(f::DST_OFF, 0);
push(Pipe::Move, q, Some(x), None, None, o0)?;
}
(Params::Shuffle { r: rr }, PIXEL_SHUFFLE) => {
let mut q = mv(r, 1);
q.set(f::N, out_n).set(f::R, *rr);
push(Pipe::Move, q, Some(x), None, None, o0)?;
}
(Params::Shuffle { r: rr }, _) => {
let mut q = mv(r, 2);
q.set(f::N, out_n).set(f::R, *rr);
push(Pipe::Move, q, Some(x), None, None, o0)?;
}
(Params::PadRep { .. }, _) => {
let mut q = mv(r, 3);
q.set(f::N, out_n);
push(Pipe::Move, q, Some(x), None, None, o0)?;
}
(_, CROP_TO) => {
let mut q = mv(r, 4);
q.set(f::N, out_n).set(f::Y0, 0).set(f::X0, 0);
push(Pipe::Move, q, Some(x), None, None, o0)?;
}
(Params::Crop { y0, x0, .. }, _) => {
let v = |o: &Operand| match *o {
Operand::Const(k) => k,
Operand::Scalar(s) => sc[s as usize] as u32,
};
let mut q = mv(r, 4);
q.set(f::N, out_n).set(f::Y0, v(y0)).set(f::X0, v(x0));
push(Pipe::Move, q, Some(x), None, None, o0)?;
}
(Params::Subpos { order }, SUBPOS_SPLIT) => {
for (k, &ord) in order.iter().enumerate() {
let ok = op.outs[k];
let mut q = mv(r, 5);
q.set(f::N, out_n).set(f::K0, (ord >> 1) as u32).set(f::K1, (ord & 1) as u32);
push(Pipe::Move, q, Some(x), None, None, ok)?;
}
}
(Params::Subpos { order }, _) => {
for (k, &ord) in order.iter().enumerate() {
let mut q = mv(r, 6);
q.set(f::N, in_n).set(f::K0, (ord >> 1) as u32).set(f::K1, (ord & 1) as u32);
push(Pipe::Move, q, Some(op.ins[k]), None, None, o0)?;
}
}
(Params::Fill { kind, value }, _) => {
let v = if *kind == 0 { *value } else { sc[*value as usize] };
let mut q = mv(r, 7);
q.set(f::N, out_n).set64(f::LO_X, v);
push(Pipe::Move, q, None, None, None, o0)?;
}
_ => return Err(Fallback::Operator { index: i as u32, name: wai_ops::graph::op_name(op.code).unwrap_or("?") }),
}
// The outputs' policies, in output order.
for &t in &op.outs {
let tr = &g.tensors[t as usize];
let mode = match tr.policy {
Policy::Derived => continue,
Policy::Clamp => 0,
Policy::Check => 1,
};
let (lo, hi) = if width == Width::Acc32 { bounds32(tr.lo, tr.hi) } else { (tr.lo, tr.hi) };
let (th, tw) = sz(t);
let mut q = Rec::new();
q.set(f::N, ch(t) * th * tw).set(f::MODE, mode).set(f::FLAG, i as u32).set64(f::LO_X, lo).set64(f::HI_X, hi);
push(Pipe::Op(Entry::Policy), q, None, None, None, t)?;
}
}
Ok(out)
}

166
wai/wai-ops-gpu/src/kernels.rs Normal file
View file

@ -0,0 +1,166 @@
//! The kernels' source and the record each dispatch reads.
//!
//! Three shader modules, each the shared arithmetic (`wgsl/arith.wgsl`) and bindings
//! (`wgsl/common.wgsl`) followed by its own entry points:
//!
//! - `ops32`: the graph operators of `wai-ops/1-i32` at the `acc32` widths, one word per
//! element (`wgsl/ops32.wgsl`);
//! - `ops64`: the same operators at the `acc64` widths, two words per element
//! (`wgsl/ops64.wgsl`);
//! - `move`: the data-movement operators for both, on elements of one or two words
//! (`wgsl/move.wgsl`).
//!
//! `probe` (`wgsl/probe.wgsl`) evaluates every arithmetic function on records the host
//! gives, so a device's arithmetic can be held to the host's before it runs a graph.
pub const ARITH: &str = include_str!("wgsl/arith.wgsl");
pub const COMMON: &str = include_str!("wgsl/common.wgsl");
pub const OPS32: &str = include_str!("wgsl/ops32.wgsl");
pub const OPS64: &str = include_str!("wgsl/ops64.wgsl");
pub const MOVE: &str = include_str!("wgsl/move.wgsl");
pub const PROBE: &str = include_str!("wgsl/probe.wgsl");
/// The text of one module: the arithmetic, the bindings, then `body`.
pub fn module(body: &str) -> String {
let mut s = String::with_capacity(ARITH.len() + COMMON.len() + body.len() + 2);
s.push_str(ARITH);
s.push('\n');
s.push_str(COMMON);
s.push('\n');
s.push_str(body);
s
}
/// The sources of the four modules, in the order [`Sources::ALL`] names them, so a test
/// can run a device on an edited text (a mutant) and see the edit caught.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Sources {
pub ops32: String,
pub ops64: String,
pub mv: String,
pub probe: String,
}
impl Sources {
pub fn shipped() -> Sources {
Sources { ops32: module(OPS32), ops64: module(OPS64), mv: module(MOVE), probe: module(PROBE) }
}
}
/// Words in one record (`struct Params` of `wgsl/common.wgsl`).
pub const REC_WORDS: usize = 48;
/// The fields of the record, by word index, in `struct Params`'s order.
pub mod f {
pub const N: usize = 0;
pub const ROW: usize = 1;
pub const C: usize = 2;
pub const H: usize = 3;
pub const W: usize = 4;
pub const CO: usize = 5;
pub const HO: usize = 6;
pub const WO: usize = 7;
pub const KH: usize = 8;
pub const KW: usize = 9;
pub const SH: usize = 10;
pub const SW: usize = 11;
pub const PH: usize = 12;
pub const PW: usize = 13;
pub const CG: usize = 14;
pub const OG: usize = 15;
pub const ROUND: usize = 16;
pub const BIAS: usize = 17;
pub const MK: usize = 18;
pub const SK: usize = 19;
pub const W_OFF: usize = 20;
pub const M_OFF: usize = 21;
pub const S_OFF: usize = 22;
pub const B_OFF: usize = 23;
pub const W_WORDS: usize = 24;
pub const MODE: usize = 25;
pub const F: usize = 26;
pub const FLAG: usize = 27;
pub const LO_X: usize = 28;
pub const LO_Y: usize = 29;
pub const HI_X: usize = 30;
pub const HI_Y: usize = 31;
pub const A_X: usize = 32;
pub const A_Y: usize = 33;
pub const Y0: usize = 34;
pub const X0: usize = 35;
pub const C2: usize = 36;
pub const R: usize = 37;
pub const EW: usize = 38;
pub const SRC_OFF: usize = 39;
pub const DST_OFF: usize = 40;
pub const K0: usize = 41;
pub const K1: usize = 42;
}
/// One dispatch's record.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Rec(pub [u32; REC_WORDS]);
impl Rec {
pub fn new() -> Rec {
Rec([0; REC_WORDS])
}
pub fn set(&mut self, field: usize, v: u32) -> &mut Rec {
self.0[field] = v;
self
}
/// A 64-bit value in two fields, low word first.
pub fn set64(&mut self, lo_field: usize, v: i64) -> &mut Rec {
self.0[lo_field] = v as u32;
self.0[lo_field + 1] = (v >> 32) as u32;
self
}
pub fn bytes(&self) -> [u8; REC_WORDS * 4] {
let mut b = [0u8; REC_WORDS * 4];
for (k, w) in self.0.iter().enumerate() {
b[4 * k..4 * k + 4].copy_from_slice(&w.to_le_bytes());
}
b
}
}
impl Default for Rec {
fn default() -> Rec {
Rec::new()
}
}
/// The entry points of `ops32` and `ops64`, which share their names.
pub const OP_ENTRIES: [&str; 11] = ["conv", "conv_t", "rescale", "chscale", "mul", "gram", "apply", "igdn", "lut16", "ew", "policy"];
/// An entry point of `ops32`/`ops64` by its index in [`OP_ENTRIES`].
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Entry {
Conv = 0,
ConvT = 1,
Rescale = 2,
Chscale = 3,
Mul = 4,
Gram = 5,
Apply = 6,
Igdn = 7,
Lut16 = 8,
Ew = 9,
Policy = 10,
}
/// Threads per workgroup of every entry point but `probe`'s.
pub const WG: u32 = 256;
/// Threads per workgroup of `probe`.
pub const WG_PROBE: u32 = 64;
/// Workgroups along x, at most (the WebGPU per-dimension limit).
pub const X_MAX: u32 = 65535;
/// The grid for `n` threads in workgroups of `wg`: `(x, y, row)`, with `row = x·wg`
/// threads per grid row. `None` past the per-dimension limit on y.
pub fn grid(n: u32, wg: u32) -> Option<(u32, u32, u32)> {
let groups = n.div_ceil(wg).max(1);
let x = groups.min(X_MAX);
let y = groups.div_ceil(x);
(y <= X_MAX).then_some((x, y, x * wg))
}

42
wai/wai-ops-gpu/src/lib.rs Normal file
View file

@ -0,0 +1,42 @@
//! `wai-ops/1-i32` on GPUs: an operator implementation of the integer operator set
//! (`extensions/integer-operators.md`) in WGSL compute kernels, run through `wgpu` on
//! Metal, Vulkan and DX12 natively and on a browser's WebGPU.
//!
//! The kernels use 32-bit integer words only, since WGSL has no 64-bit integer type.
//! Which widths a graph needs is its certificate's (§7, §8):
//!
//! - a graph that meets `acc32` runs on 32-bit kernels: one word per value, and its
//! two `product` intermediates formed exactly from the 16-bit halves of their
//! operands (§7.3);
//! - every other graph of `wai-ops/1-i32` meets `acc64` (no graph loads otherwise) and
//! runs on 64-bit kernels that form each `acc` value in two words and each `wide`
//! and `product` value in four, exact 64-bit and 128-bit arithmetic built from
//! 32-bit parts (§7.3).
//!
//! A graph with an operator outside `wai-ops/1-i32` runs on the CPU
//! ([`Fallback::Operator`]); [`support`] says so for each graph.
//!
//! Every output integer is the interpreter's (`wai_ops::interp`), and every refusal its
//! refusal, with its code and at its operator: the host makes the interpreter's checks
//! in the interpreter's order before the device runs anything, and the device flags a
//! `check` tensor that leaves its range. `tests/corpus.rs` holds this to every vector of
//! `ops-conformance/` that the subset covers, on every adapter found, and
//! `tests/mutants.rs` shows that an edited kernel is caught there.
//!
//! - [`device`]: discovery, devices, the pipelines and waiting ([`blocking`] natively);
//! - [`probe`]: the arithmetic probe each device passes before it runs a graph;
//! - [`prepare`]: where a graph runs, and its parameters on the device;
//! - [`exec`]: one decode;
//! - [`kernels`]: the shader sources and the record each dispatch reads.
pub mod device;
pub mod exec;
pub mod kernels;
pub mod prepare;
pub mod probe;
pub use device::{Adapter, Discovery, Gpu, GpuError, discover, discover_with};
#[cfg(not(target_arch = "wasm32"))]
pub use device::blocking;
pub use exec::{GpuSequence, RunResult};
pub use prepare::{Fallback, Width, outside_i32, support};

289
wai/wai-ops-gpu/src/prepare.rs Normal file
View file

@ -0,0 +1,289 @@
//! Which graphs run here, at which widths, and each graph's parameters on the device.
//!
//! [`support`] decides per graph, from its operators and its verdict alone:
//!
//! - a graph with an operator outside `wai-ops/1-i32` (`layernorm_c`, `l2norm`,
//! `softmax`, `sigmoid`, `elu`; integer-operators §9.1) runs on the CPU
//! ([`Fallback::Operator`], naming the first such operator);
//! - a graph whose certificate meets `acc32` runs on the 32-bit kernels
//! ([`Width::Acc32`]): every value it forms fits i32 except its two `product`
//! intermediates, which the kernels form in two words (§7.2, §7.3);
//! - every other graph, which meets `acc64` (no graph loads otherwise), runs on the
//! 64-bit kernels ([`Width::Acc64`]): two words per `acc` value, four per `wide` and
//! `product` value, which hold every value §7.1 allows that profile (§7.3: "a
//! fabric that emulates 64-bit and 128-bit arithmetic exactly is an `acc64`
//! implementation on that fabric").
//!
//! The parameters of each operator go to the device once per graph, as words: every
//! 64-bit value (a multiplier, a bias, a `γ`, a `β`, an entry of `chscale`'s table) in
//! two, low word first; a shift in one; the weights of a (transposed) convolution in
//! one each where they all fit i32, which the 32-bit kernels require and `acc32`
//! guarantees (§7.2, `weights`), and in two otherwise; `lut16`'s table one
//! sign-extended word per entry.
use std::sync::Arc;
use wai_ops::graph::{BiasKind, Graph, Op, Params, Round, Rq, opcode::*};
use wai_ops::{Loaded, Profile};
use crate::device::{Gpu, GpuError, Scopes, lock};
/// The kernels a graph runs on.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum Width {
/// `wai-ops/1-i32` at the `acc32` widths: one word per element.
Acc32,
/// `wai-ops/1-i32` at the `acc64` widths: two words per element.
Acc64,
}
impl Width {
pub fn as_str(self) -> &'static str {
match self {
Width::Acc32 => "acc32",
Width::Acc64 => "acc64",
}
}
/// Words per tensor element.
pub fn words(self) -> u32 {
match self {
Width::Acc32 => 1,
Width::Acc64 => 2,
}
}
}
/// Why a graph, or one decode of it, runs on the CPU.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Fallback {
/// The graph has an operator outside `wai-ops/1-i32` (§9.1).
Operator { index: u32, name: &'static str },
/// The graph forms no value past i64, but was asked to run on the 32-bit kernels,
/// which its certificate does not admit.
Profile,
/// A tensor, a parameter list or a grid of this decode is past what the adapter
/// binds or dispatches at once.
Size(String),
}
impl core::fmt::Display for Fallback {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
Fallback::Operator { index, name } => write!(f, "operator {index} ({name}) is outside wai-ops/1-i32"),
Fallback::Profile => write!(f, "the graph's certificate does not meet acc32"),
Fallback::Size(what) => write!(f, "{what}"),
}
}
}
/// The five operators outside `wai-ops/1-i32` (§9.1).
pub fn outside_i32(code: u8) -> bool {
matches!(code, LAYERNORM_C | L2NORM | SOFTMAX | SIGMOID | ELU)
}
/// Where `g` runs, and at which width (see the module documentation).
pub fn support(g: &Loaded) -> Result<Width, Fallback> {
for (i, op) in g.ops.iter().enumerate() {
if outside_i32(op.code) {
return Err(Fallback::Operator { index: i as u32, name: wai_ops::graph::op_name(op.code).unwrap_or("?") });
}
}
Ok(if g.meets(Profile::Acc32) { Width::Acc32 } else { Width::Acc64 })
}
/// Word offsets of one operator's parameters, and the kinds the kernels read.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub(crate) struct Offs {
pub w_off: u32,
pub m_off: u32,
pub s_off: u32,
pub b_off: u32,
pub w_words: u32,
pub round: u32,
pub bias: u32,
pub mk: u32,
pub sk: u32,
}
pub(crate) struct PrepOp {
pub buf: Option<wgpu::Buffer>,
pub offs: Offs,
}
pub(crate) struct Prepared {
pub ops: Vec<PrepOp>,
}
pub(crate) struct Cached {
graph: Graph,
width: Width,
prep: Arc<Prepared>,
}
/// Graphs whose parameters stay on the device, most recently used last.
const CACHE_GRAPHS: usize = 16;
pub(crate) fn round_code(r: Round) -> u32 {
match r {
Round::Floor => 0,
Round::HalfUp => 1,
Round::HalfEven => 2,
}
}
fn fits32(v: i64) -> bool {
i32::try_from(v).is_ok()
}
struct Words(Vec<u32>);
impl Words {
fn at(&self) -> u32 {
self.0.len() as u32
}
fn w64(&mut self, v: i64) {
self.0.push(v as u32);
self.0.push((v >> 32) as u32);
}
fn w32(&mut self, v: i64) {
self.0.push(v as u32);
}
/// An `RQ` block (§5.1): its multipliers, shifts and biases as stored.
fn rq(&mut self, rq: &Rq, o: &mut Offs) {
o.round = round_code(rq.round);
o.bias = match rq.bias_kind {
BiasKind::None => 0,
BiasKind::Pre => 1,
BiasKind::Post => 2,
};
o.mk = rq.m_kind as u32;
o.sk = rq.s_kind as u32;
o.m_off = self.at();
for &m in &rq.mult {
self.w64(m);
}
o.s_off = self.at();
for &s in &rq.shift {
self.w32(s as i64);
}
o.b_off = self.at();
for &b in &rq.bias {
self.w64(b);
}
}
}
/// What the 32-bit kernels require of an `RQ` block, which `acc32` guarantees (§7.2).
fn rq_fits32(rq: &Rq) -> bool {
rq.mult.iter().all(|&m| fits32(m)) && rq.shift.iter().all(|&s| s <= 63) && (rq.bias_kind != BiasKind::Post || rq.bias.iter().all(|&b| fits32(b)))
}
/// One operator's parameter words and offsets. `Err` for an operator whose parameters
/// break what the 32-bit kernels require, which no graph meeting `acc32` has.
fn op_words(op: &Op, width: Width) -> Result<(Vec<u32>, Offs), String> {
let mut w = Words(Vec::new());
let mut o = Offs::default();
let narrow = width == Width::Acc32;
let bad = |what: &str| Err(format!("operator {:#04x}: {what} under acc32", op.code));
match &op.p {
Params::Conv { rq, w: wt, .. } | Params::ConvT { rq, w: wt, .. } => {
let one = wt.iter().all(|&v| fits32(v));
if narrow && (!one || !rq_fits32(rq)) {
return bad("a weight or a requantisation parameter past its acc32 width");
}
o.w_off = w.at();
o.w_words = if one { 1 } else { 2 };
for &v in wt {
if one { w.w32(v) } else { w.w64(v) }
}
w.rq(rq, &mut o);
}
Params::Rescale(rq) | Params::Mul(rq) | Params::Gram { rq, .. } | Params::Apply { rq, .. } => {
if narrow && !rq_fits32(rq) {
return bad("a requantisation parameter past its acc32 width");
}
w.rq(rq, &mut o);
}
Params::Chscale { round, shift, m, .. } => {
if narrow && (!m.iter().all(|&v| fits32(v)) || *shift > 63) {
return bad("a table entry or the shift past its acc32 width");
}
o.round = round_code(*round);
o.m_off = w.at();
for &v in m {
w.w64(v);
}
}
Params::Lut16 { t, .. } => {
o.w_off = w.at();
for &v in t {
w.w32(v as i64);
}
}
Params::Igdn { gf, frac, gamma, beta, .. } => {
if narrow && (!gamma.iter().all(|&v| fits32(v)) || *gf > 63 || *frac > 63) {
return bad("a gamma or a shift past its acc32 width");
}
o.w_off = w.at();
for &v in gamma {
w.w64(v);
}
o.b_off = w.at();
for &v in beta {
w.w64(v);
}
}
Params::Leaky { a, f } if narrow && (!fits32(*a) || *f > 63) => return bad("the slope past its acc32 width"),
_ => {}
}
Ok((w.0, o))
}
impl Gpu {
/// The graph's parameters on the device, at `width`: from the cache when the same
/// graph (compared in full) was prepared at that width before.
pub(crate) async fn prepared(&self, l: &Loaded, width: Width) -> Result<Arc<Prepared>, GpuError> {
let g = l.graph();
{
let mut c = lock(&self.cache);
if let Some(k) = c.iter().position(|e| e.width == width && e.graph.ops.len() == g.ops.len() && e.graph.tensors.len() == g.tensors.len() && &e.graph == g) {
let e = c.remove(k);
let p = Arc::clone(&e.prep);
c.push(e);
return Ok(p);
}
}
let scopes = Scopes::push(&self.device);
let mut ops = Vec::with_capacity(g.ops.len());
for op in &g.ops {
let (words, offs) = op_words(op, width).map_err(GpuError::Device)?;
let bytes = words.len() as u64 * 4;
if bytes > self.adapter.max_storage_binding {
return Err(GpuError::Fallback(Fallback::Size(format!("an operator's {bytes} parameter bytes exceed the adapter's storage binding"))));
}
let buf = (!words.is_empty()).then(|| {
let b = self.buffer("wai-ops-gpu parameters", bytes, wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST);
let mut raw = Vec::with_capacity(words.len() * 4);
for x in &words {
raw.extend_from_slice(&x.to_le_bytes());
}
self.queue.write_buffer(&b, 0, &raw);
b
});
ops.push(PrepOp { buf, offs });
}
scopes.close().await?;
let prep = Arc::new(Prepared { ops });
let mut c = lock(&self.cache);
if c.len() >= CACHE_GRAPHS {
c.remove(0);
}
c.push(Cached { graph: g.clone(), width, prep: Arc::clone(&prep) });
Ok(prep)
}
/// Drop every prepared graph.
pub fn clear_cache(&self) {
lock(&self.cache).clear();
}
}

292
wai/wai-ops-gpu/src/probe.rs Normal file
View file

@ -0,0 +1,292 @@
//! The arithmetic probe: every function of `wgsl/arith.wgsl` on the device, against the
//! same functions in the host's 128-bit integers.
//!
//! [`check`] runs at discovery, on [`cases`]: the edges of every width pairwise, every
//! shift from 0 to 129 under each rounding mode, and SplitMix64 operands. A device whose
//! output differs from [`expect`] in any word is not opened. `tests/arith.rs` runs more
//! cases, and shows that an edited function is caught.
//!
//! The host model is written from the definitions in integer-operators §2 (`R_round`,
//! `isqrt`) and from modular arithmetic, never from the shader text: each result is the
//! exact value reduced modulo 2^64 or 2^128, which is what the shader's words hold.
use crate::device::{Gpu, GpuError, Scopes};
use crate::kernels::{Rec, WG_PROBE, f, grid};
/// Words of one probe record's input and output.
pub const IN_WORDS: usize = 12;
pub const OUT_WORDS: usize = 40;
/// One probe case: `a` and `b` as 128-bit values (the 64-bit functions read their low
/// halves, the 32-bit ones their low words), a shift and a rounding mode (0 floor, 1
/// half_up, 2 half_even).
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ProbeIn {
pub a: u128,
pub b: u128,
pub s: u32,
pub mode: u32,
}
fn words4(v: u128) -> [u32; 4] {
[v as u32, (v >> 32) as u32, (v >> 64) as u32, (v >> 96) as u32]
}
/// `R_round(p, s)` (§2) on the exact integer `p`.
fn round_exact(p: i128, s: u32, mode: u32) -> i128 {
if s == 0 {
return p;
}
let q = p >> s;
match mode {
0 => q,
1 => {
// ⌊(p + 2^(s−1)) / 2^s⌋ = q, or q + 1 when the remainder is at least half.
let rem = (p as u128) & ((1u128 << s) - 1);
if rem >= 1u128 << (s - 1) { q + 1 } else { q }
}
_ => {
let rem = (p as u128) & ((1u128 << s) - 1);
let h = 1u128 << (s - 1);
if rem > h || (rem == h && q & 1 == 1) { q + 1 } else { q }
}
}
}
/// The 128-bit modular `R_round` the shader forms: `half_up` adds `2^(s−1)` modulo
/// 2^128 before the shift. Where that sum does not wrap it is [`round_exact`].
fn round_mod128(p: i128, s: u32, mode: u32) -> i128 {
if s == 0 || mode != 1 {
return round_exact(p, s, mode);
}
p.wrapping_add(1i128 << (s - 1)) >> s
}
/// The 64-bit modular form, for `s` in 0..=63.
fn round_mod64(p: i64, s: u32, mode: u32) -> i64 {
if s == 0 || mode != 1 {
return round_exact(p as i128, s, mode) as i64;
}
p.wrapping_add(1i64 << (s - 1)) >> s
}
fn sar64(a: i64, s: u32) -> i64 {
if s >= 64 { a >> 63 } else { a >> s }
}
/// What the device must write for `c`.
pub fn expect(c: &ProbeIn) -> [u32; OUT_WORDS] {
let mut o = [0u32; OUT_WORDS];
let mut k = 0usize;
let put_w = |o: &mut [u32; OUT_WORDS], k: &mut usize, v: u64| {
o[*k] = v as u32;
o[*k + 1] = (v >> 32) as u32;
*k += 2;
};
let put_q = |o: &mut [u32; OUT_WORDS], k: &mut usize, v: u128| {
o[*k..*k + 4].copy_from_slice(&words4(v));
*k += 4;
};
let (a64, b64) = (c.a as u64 as i64, c.b as u64 as i64);
let (a32, b32) = (c.a as u32 as i32, c.b as u32 as i32);
let (a, b) = (c.a as i128, c.b as i128);
put_w(&mut o, &mut k, a64.wrapping_add(b64) as u64);
put_w(&mut o, &mut k, a64.wrapping_mul(b64) as u64);
put_w(&mut o, &mut k, (a32 as i64 * b32 as i64) as u64);
put_w(&mut o, &mut k, round_mod64(a64, c.s & 63, c.mode) as u64);
put_w(&mut o, &mut k, sar64(a64, c.s) as u64);
put_q(&mut o, &mut k, (a64 as i128 * b64 as i128) as u128);
put_q(&mut o, &mut k, a.wrapping_mul(b) as u128);
put_q(&mut o, &mut k, a.wrapping_add(b) as u128);
put_q(&mut o, &mut k, round_mod128(a, c.s & 127, c.mode) as u128);
put_w(&mut o, &mut k, (c.a & (u128::MAX >> 1)).isqrt() as u64);
o[k] = (c.a as u32).isqrt();
k += 1;
put_w(&mut o, &mut k, (c.a as u64).isqrt());
o[k] = if c.s >= 32 { (a32 >> 31) as u32 } else { (a32 >> c.s) as u32 };
k += 1;
o[k] = (a64 < b64) as u32 | (((c.a as u64) < (c.b as u64)) as u32) << 1 | ((c.a < c.b) as u32) << 2 | ((a < 0) as u32) << 3 | ((a64 < 0) as u32) << 4;
k += 1;
put_w(&mut o, &mut k, a64.wrapping_neg() as u64);
let hi = b64.wrapping_add(1i64 << 32);
put_w(&mut o, &mut k, a64.max(b64).min(hi) as u64);
debug_assert_eq!(k, 37);
o
}
struct SplitMix(u64);
impl SplitMix {
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)
}
}
/// Values at the edges of 32, 64 and 128 bits, signed and unsigned, and near them.
pub fn edges() -> Vec<u128> {
let mut e: Vec<u128> = Vec::new();
for v in [0i128, 1, -1, 2, -2, 3, -3, 0xFFFF, 0x10000, -0x10000, 0x7FFF_FFFF, -0x8000_0000, 0x8000_0000, 0xFFFF_FFFF, 0x1_0000_0000] {
e.push(v as u128);
}
for v in [i64::MAX as i128, i64::MIN as i128, i64::MAX as i128 - 1, i64::MIN as i128 + 1, u64::MAX as i128, 1i128 << 64, (1i128 << 62) + 7] {
e.push(v as u128);
}
for v in [i128::MAX, i128::MIN, i128::MAX - 1, i128::MIN + 1, 1i128 << 126, -(1i128 << 126), (1i128 << 96) - 1] {
e.push(v as u128);
}
e
}
/// The discovery probe's cases.
pub fn cases(random: usize, seed: u64) -> Vec<ProbeIn> {
let e = edges();
let mut v = Vec::new();
let mut k = 0u32;
for &a in &e {
for &b in &e {
v.push(ProbeIn { a, b, s: k % 130, mode: k % 3 });
k += 1;
}
}
// Every shift under every mode, on a positive and a negative tie and their
// neighbours.
for s in 0..130u32 {
for mode in 0..3u32 {
let sh = s.clamp(1, 127);
for p in [3i128 << (sh - 1), -(3i128 << (sh - 1)), 1i128 << (sh - 1), -(1i128 << (sh - 1)), (5i128 << (sh - 1)) + 1, -(5i128 << (sh - 1)) - 1] {
v.push(ProbeIn { a: p as u128, b: 0, s, mode });
}
}
}
let mut r = SplitMix(seed);
for i in 0..random {
let width = [32u32, 64, 96, 128][i % 4];
let mask = if width == 128 { u128::MAX } else { (1u128 << width) - 1 };
let a = ((r.next() as u128) << 64 | r.next() as u128) & mask;
let b = ((r.next() as u128) << 64 | r.next() as u128) & mask;
let sign = |x: u128| if r_bit(i) { x.wrapping_neg() } else { x };
v.push(ProbeIn { a: sign(a), b, s: (r.next() % 130) as u32, mode: (r.next() % 3) as u32 });
}
v
}
fn r_bit(i: usize) -> bool {
(i / 4) % 2 == 1
}
/// The device's words for each case.
pub async fn run(g: &Gpu, recs: &[ProbeIn]) -> Result<Vec<[u32; OUT_WORDS]>, GpuError> {
let n = recs.len() as u32;
let mut inw: Vec<u8> = Vec::with_capacity(recs.len() * IN_WORDS * 4);
for c in recs {
let mut w = [0u32; IN_WORDS];
w[..4].copy_from_slice(&words4(c.a));
w[4..8].copy_from_slice(&words4(c.b));
w[8] = c.s;
w[9] = c.mode;
for x in w {
inw.extend_from_slice(&x.to_le_bytes());
}
}
let out_bytes = (recs.len() * OUT_WORDS * 4) as u64;
let (gx, gy, row) = grid(n, WG_PROBE).ok_or_else(|| GpuError::Device("too many probe cases".into()))?;
let mut rec = Rec::new();
rec.set(f::N, n).set(f::ROW, row);
let scopes = Scopes::push(&g.device);
let uni = g.buffer("wai-ops-gpu probe record", g.stride, wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST);
g.queue.write_buffer(&uni, 0, &rec.bytes());
let ins = g.buffer("wai-ops-gpu probe in", inw.len() as u64, wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST);
g.queue.write_buffer(&ins, 0, &inw);
let out = g.buffer("wai-ops-gpu probe out", out_bytes, wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC);
let flags = g.buffer("wai-ops-gpu probe flags", 16, wgpu::BufferUsages::STORAGE);
let staging = g.buffer("wai-ops-gpu probe staging", out_bytes, wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST);
let bind = g.device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some("wai-ops-gpu probe"),
layout: &g.pipes.layout,
entries: &[
wgpu::BindGroupEntry { binding: 0, resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding { buffer: &uni, offset: 0, size: core::num::NonZeroU64::new(crate::kernels::REC_WORDS as u64 * 4) }) },
wgpu::BindGroupEntry { binding: 1, resource: ins.as_entire_binding() },
wgpu::BindGroupEntry { binding: 2, resource: g.dummy.as_entire_binding() },
wgpu::BindGroupEntry { binding: 3, resource: g.dummy.as_entire_binding() },
wgpu::BindGroupEntry { binding: 4, resource: out.as_entire_binding() },
wgpu::BindGroupEntry { binding: 5, resource: flags.as_entire_binding() },
],
});
let mut enc = g.device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("wai-ops-gpu probe") });
{
let mut pass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor { label: Some("wai-ops-gpu probe"), timestamp_writes: None });
pass.set_pipeline(&g.pipes.probe);
pass.set_bind_group(0, &bind, &[]);
pass.dispatch_workgroups(gx, gy, 1);
}
enc.copy_buffer_to_buffer(&out, 0, &staging, 0, out_bytes);
let index = g.queue.submit([enc.finish()]);
scopes.close().await?;
let bytes = g.read(&staging, out_bytes, index).await?;
g.settled()?;
Ok(bytes
.as_chunks::<{ OUT_WORDS * 4 }>()
.0
.iter()
.map(|c| {
let mut w = [0u32; OUT_WORDS];
for (k, b) in c.as_chunks::<4>().0.iter().enumerate() {
w[k] = u32::from_le_bytes(*b);
}
w
})
.collect())
}
/// The first case on which the device's words differ from [`expect`], with the word.
pub fn first_mismatch(recs: &[ProbeIn], got: &[[u32; OUT_WORDS]]) -> Option<(ProbeIn, usize)> {
for (c, g) in recs.iter().zip(got) {
let e = expect(c);
if let Some(k) = (0..OUT_WORDS).find(|&k| e[k] != g[k]) {
return Some((*c, k));
}
}
if recs.len() != got.len() && !recs.is_empty() {
return Some((recs[got.len().min(recs.len() - 1)], usize::MAX));
}
None
}
/// The discovery probe: `Err` with the first difference, or the device's error.
pub(crate) async fn check(g: &Gpu) -> Result<(), String> {
let recs = cases(4096, 0x5EED_0005);
let got = run(g, &recs).await.map_err(|e| format!("the probe did not run: {e}"))?;
match first_mismatch(&recs, &got) {
None => Ok(()),
Some((c, k)) => Err(format!("the probe's word {k} differs from the host's on {c:?}")),
}
}
#[cfg(test)]
mod tests {
use super::*;
/// The host model against the definitions on values where nothing wraps.
#[test]
fn the_host_model_is_the_definition_where_nothing_wraps() {
for s in 1..127u32 {
for p in [3i128 << (s - 1), -(3i128 << (s - 1)), 1i128 << (s - 1), -(1i128 << (s - 1))] {
for mode in 0..3 {
let want = match mode {
0 => wai_ops::interp::round(p, s, wai_ops::graph::Round::Floor),
1 => wai_ops::interp::round(p, s, wai_ops::graph::Round::HalfUp),
_ => wai_ops::interp::round(p, s, wai_ops::graph::Round::HalfEven),
};
if let Some(w) = want {
assert_eq!(round_mod128(p, s, mode), w, "p {p} s {s} mode {mode}");
}
}
}
}
assert_eq!(expect(&ProbeIn { a: 16, b: 0, s: 0, mode: 0 })[28], 4);
}
}

393
wai/wai-ops-gpu/src/wgsl/arith.wgsl Normal file
View file

@ -0,0 +1,393 @@
// The integer arithmetic of the kernels, in 32-bit words only.
//
// WGSL has no 64-bit integer. A 64-bit two's-complement value is a `W`, two words
// (x low, y high); a 128-bit one is a `Q`, four words (x lowest, w highest). Every
// operation here uses only what every WGSL implementation defines identically:
// `u32` addition, subtraction and multiplication, which wrap; bitwise operations;
// comparisons; `bitcast`; and shifts, whose amount WGSL takes modulo 32, so every
// shift of 32 or more is handled explicitly (integer-operators §6). No floating
// point, and no signed operation that could overflow: signed values are read from
// a word's bits only to compare or to shift right arithmetically.
//
// Each function's result is exact modulo 2^64 (W) or 2^128 (Q). A kernel uses one
// only where the graph's certificate places the exact value inside that width
// (§7.1, §8), so the modular result is the exact value. `tests/arith.rs` holds every
// function to 128-bit host arithmetic on edge cases and random operands, on the
// device itself.
alias W = vec2<u32>;
alias Q = vec4<u32>;
const ONES: u32 = 0xFFFFFFFFu;
// ------------------------------------------------------------------ 64 bits
fn w_add(a: W, b: W) -> W {
let lo = a.x + b.x;
return W(lo, a.y + b.y + select(0u, 1u, lo < a.x));
}
fn w_neg(a: W) -> W {
return w_add(W(~a.x, ~a.y), W(1u, 0u));
}
fn w_sub(a: W, b: W) -> W {
return w_add(a, w_neg(b));
}
fn w_isneg(a: W) -> bool {
return (a.y >> 31u) != 0u;
}
// The bits of an i32 sign-extended to 64.
fn w_sext(v: u32) -> W {
return W(v, select(0u, ONES, (v >> 31u) != 0u));
}
fn w_eq(a: W, b: W) -> bool {
return a.x == b.x && a.y == b.y;
}
// a < b, signed.
fn w_lt(a: W, b: W) -> bool {
let ah = bitcast<i32>(a.y);
let bh = bitcast<i32>(b.y);
return ah < bh || (ah == bh && a.x < b.x);
}
// a < b, unsigned.
fn w_ltu(a: W, b: W) -> bool {
return a.y < b.y || (a.y == b.y && a.x < b.x);
}
fn w_max(a: W, b: W) -> W {
return select(a, b, w_lt(a, b));
}
fn w_min(a: W, b: W) -> W {
return select(a, b, w_lt(b, a));
}
fn w_clamp(v: W, lo: W, hi: W) -> W {
return w_min(w_max(v, lo), hi);
}
// |a| as an unsigned 64-bit value: 2^63 for −2^63.
fn w_mag(a: W) -> W {
return select(a, w_neg(a), w_isneg(a));
}
// The unsigned product of two words, from their 16-bit halves: each partial
// product is below 2^32, and the carries of the middle sum are counted.
fn mul_wide(a: u32, b: u32) -> W {
let a0 = a & 0xFFFFu;
let a1 = a >> 16u;
let b0 = b & 0xFFFFu;
let b1 = b >> 16u;
let p00 = a0 * b0;
let p01 = a0 * b1;
let p10 = a1 * b0;
let p11 = a1 * b1;
let mid = p01 + p10;
let c1 = select(0u, 1u, mid < p01);
let mid2 = mid + (p00 >> 16u);
let c2 = select(0u, 1u, mid2 < mid);
let lo = (mid2 << 16u) | (p00 & 0xFFFFu);
let hi = p11 + (mid2 >> 16u) + ((c1 + c2) << 16u);
return W(lo, hi);
}
// The signed product of two i32 (given as their bits), exact in 64 bits.
fn mul_i32(a: u32, b: u32) -> W {
let na = (a >> 31u) != 0u;
let nb = (b >> 31u) != 0u;
let m = mul_wide(select(a, 0u - a, na), select(b, 0u - b, nb));
return select(m, w_neg(m), na != nb);
}
// a·b modulo 2^64: the exact product wherever it fits i64.
fn w_mul_lo(a: W, b: W) -> W {
let p = mul_wide(a.x, b.x);
return W(p.x, p.y + a.x * b.y + a.y * b.x);
}
// ⌊a / 2^s⌋ for any s (an arithmetic shift).
fn w_sar(a: W, s: u32) -> W {
let sg = bitcast<u32>(bitcast<i32>(a.y) >> 31u);
if (s == 0u) {
return a;
}
if (s < 32u) {
return W((a.x >> s) | (a.y << (32u - s)), bitcast<u32>(bitcast<i32>(a.y) >> s));
}
if (s < 64u) {
return W(bitcast<u32>(bitcast<i32>(a.y) >> (s - 32u)), sg);
}
return W(sg, sg);
}
// 2^k for k in 0..=63 (2^63 is the bit pattern of −2^63).
fn w_pow2(k: u32) -> W {
if (k < 32u) {
return W(1u << k, 0u);
}
return W(0u, 1u << (k - 32u));
}
// R_round(p, s) (§2) on 64 bits, for s in 0..=63, mode 0 floor, 1 half_up, 2
// half_even. half_up forms p + 2^(s−1), which rq.p's interval holds.
fn w_round(p: W, s: u32, mode: u32) -> W {
if (s == 0u) {
return p;
}
if (mode == 0u) {
return w_sar(p, s);
}
let h = w_pow2(s - 1u);
if (mode == 1u) {
return w_sar(w_add(p, h), s);
}
let q = w_sar(p, s);
let mask = w_add(w_pow2(s), W(ONES, ONES));
let rem = W(p.x & mask.x, p.y & mask.y);
let up = w_ltu(h, rem) || (w_eq(rem, h) && (q.x & 1u) == 1u);
return select(q, w_add(q, W(1u, 0u)), up);
}
// ⌊a / 2^s⌋ of an i32 (given as its bits) for any s.
fn sar32(a: u32, s: u32) -> u32 {
if (s >= 32u) {
return bitcast<u32>(bitcast<i32>(a) >> 31u);
}
return bitcast<u32>(bitcast<i32>(a) >> s);
}
// isqrt of a 32-bit value, digit by digit.
fn isqrt32(n0: u32) -> u32 {
var n = n0;
var res = 0u;
var bit = 1u << 30u;
for (var i = 0u; i < 16u; i++) {
if (bit <= n) {
break;
}
bit = bit >> 2u;
}
for (var i = 0u; i < 16u; i++) {
if (bit == 0u) {
break;
}
let t = res + bit;
if (n >= t) {
n = n - t;
res = (res >> 1u) + bit;
} else {
res = res >> 1u;
}
bit = bit >> 2u;
}
return res;
}
// isqrt of an unsigned 64-bit value.
fn w_isqrt(n0: W) -> W {
var n = n0;
var res = W(0u, 0u);
var bit = W(0u, 1u << 30u);
for (var i = 0u; i < 32u; i++) {
if (!w_ltu(n, bit)) {
break;
}
bit = W((bit.x >> 2u) | (bit.y << 30u), bit.y >> 2u);
}
for (var i = 0u; i < 32u; i++) {
if (bit.x == 0u && bit.y == 0u) {
break;
}
let t = w_add(res, bit);
let half = W((res.x >> 1u) | (res.y << 31u), res.y >> 1u);
if (!w_ltu(n, t)) {
n = w_sub(n, t);
res = w_add(half, bit);
} else {
res = half;
}
bit = W((bit.x >> 2u) | (bit.y << 30u), bit.y >> 2u);
}
return res;
}
// ------------------------------------------------------------------ 128 bits
fn q_add(a: Q, b: Q) -> Q {
let x = a.x + b.x;
let c0 = select(0u, 1u, x < a.x);
let y0 = a.y + b.y;
let y = y0 + c0;
let c1 = select(0u, 1u, y0 < a.y) + select(0u, 1u, y < y0);
let z0 = a.z + b.z;
let z = z0 + c1;
let c2 = select(0u, 1u, z0 < a.z) + select(0u, 1u, z < z0);
return Q(x, y, z, a.w + b.w + c2);
}
fn q_neg(a: Q) -> Q {
return q_add(Q(~a.x, ~a.y, ~a.z, ~a.w), Q(1u, 0u, 0u, 0u));
}
fn q_isneg(a: Q) -> bool {
return (a.w >> 31u) != 0u;
}
fn q_iszero(a: Q) -> bool {
return (a.x | a.y | a.z | a.w) == 0u;
}
fn q_eq(a: Q, b: Q) -> bool {
return a.x == b.x && a.y == b.y && a.z == b.z && a.w == b.w;
}
// a < b, unsigned.
fn q_ltu(a: Q, b: Q) -> bool {
if (a.w != b.w) {
return a.w < b.w;
}
if (a.z != b.z) {
return a.z < b.z;
}
if (a.y != b.y) {
return a.y < b.y;
}
return a.x < b.x;
}
// A W sign-extended to 128 bits.
fn q_from_w(a: W) -> Q {
let s = select(0u, ONES, w_isneg(a));
return Q(a.x, a.y, s, s);
}
// The unsigned product of two unsigned 64-bit values, exact in 128 bits.
fn q_mul_uu(a: W, b: W) -> Q {
let p00 = mul_wide(a.x, b.x);
let p01 = mul_wide(a.x, b.y);
let p10 = mul_wide(a.y, b.x);
let p11 = mul_wide(a.y, b.y);
var r = Q(p00.x, p00.y, p11.x, p11.y);
r = q_add(r, Q(0u, p01.x, p01.y, 0u));
r = q_add(r, Q(0u, p10.x, p10.y, 0u));
return r;
}
// The signed product of two i64, exact in 128 bits.
fn q_mul_ww(a: W, b: W) -> Q {
let m = q_mul_uu(w_mag(a), w_mag(b));
return select(m, q_neg(m), w_isneg(a) != w_isneg(b));
}
// a·b modulo 2^128: the exact product wherever it fits i128.
fn q_mul_lo(a: Q, b: Q) -> Q {
var r = q_mul_uu(a.xy, b.xy);
let m1 = w_mul_lo(a.xy, b.zw);
let m2 = w_mul_lo(a.zw, b.xy);
r = q_add(r, Q(0u, 0u, m1.x, m1.y));
r = q_add(r, Q(0u, 0u, m2.x, m2.y));
return r;
}
// Word i of a, or the fill word f past the top.
fn q_word(a: Q, i: u32, f: u32) -> u32 {
if (i == 0u) {
return a.x;
}
if (i == 1u) {
return a.y;
}
if (i == 2u) {
return a.z;
}
if (i == 3u) {
return a.w;
}
return f;
}
// a shifted right by s with fill word f (0 for a logical shift, the sign word for
// an arithmetic one), for any s.
fn q_shr_fill(a: Q, s: u32, f: u32) -> Q {
if (s >= 128u) {
return Q(f, f, f, f);
}
let k = s >> 5u;
let b = s & 31u;
var r = array<u32, 4>(0u, 0u, 0u, 0u);
for (var i = 0u; i < 4u; i++) {
let lo = q_word(a, i + k, f);
let hi = q_word(a, i + k + 1u, f);
r[i] = select((lo >> b) | (hi << (32u - b)), lo, b == 0u);
}
return Q(r[0], r[1], r[2], r[3]);
}
// ⌊a / 2^s⌋ for any s (an arithmetic shift).
fn q_sar(a: Q, s: u32) -> Q {
return q_shr_fill(a, s, bitcast<u32>(bitcast<i32>(a.w) >> 31u));
}
// 2^k for k in 0..=127.
fn q_pow2(k: u32) -> Q {
let b = 1u << (k & 31u);
let i = k >> 5u;
return Q(select(0u, b, i == 0u), select(0u, b, i == 1u), select(0u, b, i == 2u), select(0u, b, i == 3u));
}
// R_round(p, s) (§2) on 128 bits, for s in 0..=127.
fn q_round(p: Q, s: u32, mode: u32) -> Q {
if (s == 0u) {
return p;
}
if (mode == 0u) {
return q_sar(p, s);
}
let h = q_pow2(s - 1u);
if (mode == 1u) {
return q_sar(q_add(p, h), s);
}
let q = q_sar(p, s);
let mask = q_add(q_pow2(s), Q(ONES, ONES, ONES, ONES));
let rem = Q(p.x & mask.x, p.y & mask.y, p.z & mask.z, p.w & mask.w);
let up = q_ltu(h, rem) || (q_eq(rem, h) && (q.x & 1u) == 1u);
return select(q, q_add(q, Q(1u, 0u, 0u, 0u)), up);
}
// isqrt of a non-negative 128-bit value below 2^127. Where the value fits 64
// bits, the 64-bit form gives the same integer in fewer steps.
fn q_isqrt(n0: Q) -> W {
if (n0.z == 0u && n0.w == 0u) {
return w_isqrt(n0.xy);
}
var n = n0;
var res = Q(0u, 0u, 0u, 0u);
var bit = Q(0u, 0u, 0u, 1u << 30u);
for (var i = 0u; i < 64u; i++) {
if (!q_ltu(n, bit)) {
break;
}
bit = q_shr_fill(bit, 2u, 0u);
}
for (var i = 0u; i < 64u; i++) {
if (q_iszero(bit)) {
break;
}
let t = q_add(res, bit);
let half = q_shr_fill(res, 1u, 0u);
if (!q_ltu(n, t)) {
n = q_add(n, q_neg(t));
res = q_add(half, bit);
} else {
res = half;
}
bit = q_shr_fill(bit, 2u, 0u);
}
return res.xy;
}

39
wai/wai-ops-gpu/src/wgsl/common.wgsl Normal file
View file

@ -0,0 +1,39 @@
// The record every dispatch reads, and the bindings every kernel shares.
//
// One record per dispatch, in a uniform buffer: the host fills the fields an entry
// point names (`src/kernels.rs`, `Rec`, writes them in this order) and leaves the
// rest 0. Tensors and parameters are arrays of words: an element of a tensor is
// one word (i32) in the 32-bit kernels and two (i64, low word first) in the 64-bit
// ones; a parameter list keeps the layout `src/prepare.rs` gives it.
struct Params {
n: u32, row: u32, c: u32, h: u32,
w: u32, co: u32, ho: u32, wo: u32,
kh: u32, kw: u32, sh: u32, sw: u32,
ph: u32, pw: u32, cg: u32, og: u32,
round: u32, bias: u32, mk: u32, sk: u32,
w_off: u32, m_off: u32, s_off: u32, b_off: u32,
w_words: u32, mode: u32, f: u32, flag: u32,
lo_x: u32, lo_y: u32, hi_x: u32, hi_y: u32,
a_x: u32, a_y: u32, y0: u32, x0: u32,
c2: u32, r: u32, ew: u32, src_off: u32,
dst_off: u32, k0: u32, k1: u32, k2: u32,
k3: u32, k4: u32, k5: u32, k6: u32,
}
@group(0) @binding(0) var<uniform> P: Params;
// The operator's first input, and its second where it has one.
@group(0) @binding(1) var<storage, read> A: array<u32>;
@group(0) @binding(2) var<storage, read> B: array<u32>;
// The operator's parameters: weights, requantisation lists, tables.
@group(0) @binding(3) var<storage, read> M: array<u32>;
// The output tensor.
@group(0) @binding(4) var<storage, read_write> O: array<u32>;
// One word per operator: set when an element of a `check` tensor lies outside its
// range (integer-operators §3).
@group(0) @binding(5) var<storage, read_write> F: array<atomic<u32>>;
// The thread's index in a grid of rows of `P.row` threads.
fn tid(g: vec3<u32>) -> u32 {
return g.x + g.y * P.row;
}

83
wai/wai-ops-gpu/src/wgsl/move.wgsl Normal file
View file

@ -0,0 +1,83 @@
// The data-movement operators (integer-operators §5.10), on elements of `ew` words
// (1 or 2), so one kernel serves both widths. Each thread writes one output
// element, by mode, from the element its definition names:
//
// 0 copy O[dst_off + t] = A[src_off + t] (concat, split)
// 1 pixel_shuffle r
// 2 pixel_unshuffle r
// 3 pad_rep (c, min(y, h−1), min(x, w−1))
// 4 crop (c, y0 + y, x0 + x) (crop, crop_to)
// 5 subpos_split (c, 2y + k0, 2x + k1) (one output)
// 6 subpos_merge writes (c, 2y + k0, 2x + k1) of the output from (c, y, x) of
// one input: a thread per input element
// 7 fill every element is the value (lo_x, lo_y)
//
// (c, h, w) is the input's shape and (co, ho, wo) the output's.
fn copy_el(dst: u32, src: u32) {
O[dst * P.ew] = A[src * P.ew];
if (P.ew == 2u) {
O[dst * 2u + 1u] = A[src * 2u + 1u];
}
}
@compute @workgroup_size(256)
fn mv(@builtin(global_invocation_id) g: vec3<u32>) {
let t = tid(g);
if (t >= P.n) {
return;
}
switch P.mode {
case 0u: {
copy_el(P.dst_off + t, P.src_off + t);
}
case 1u: {
let x = t % P.wo;
let y = (t / P.wo) % P.ho;
let c = t / (P.wo * P.ho);
let r = P.r;
let src = ((c * r * r + (y % r) * r + x % r) * P.h + y / r) * P.w + x / r;
copy_el(t, src);
}
case 2u: {
let x = t % P.wo;
let y = (t / P.wo) % P.ho;
let co = t / (P.wo * P.ho);
let r = P.r;
let c = co / (r * r);
let k = co % (r * r);
let src = (c * P.h + y * r + k / r) * P.w + x * r + k % r;
copy_el(t, src);
}
case 3u: {
let x = t % P.wo;
let y = (t / P.wo) % P.ho;
let c = t / (P.wo * P.ho);
copy_el(t, (c * P.h + min(y, P.h - 1u)) * P.w + min(x, P.w - 1u));
}
case 4u: {
let x = t % P.wo;
let y = (t / P.wo) % P.ho;
let c = t / (P.wo * P.ho);
copy_el(t, (c * P.h + P.y0 + y) * P.w + P.x0 + x);
}
case 5u: {
let x = t % P.wo;
let y = (t / P.wo) % P.ho;
let c = t / (P.wo * P.ho);
copy_el(t, (c * P.h + 2u * y + P.k0) * P.w + 2u * x + P.k1);
}
case 6u: {
let x = t % P.w;
let y = (t / P.w) % P.h;
let c = t / (P.w * P.h);
copy_el((c * P.ho + 2u * y + P.k0) * P.wo + 2u * x + P.k1, t);
}
default: {
O[t * P.ew] = P.lo_x;
if (P.ew == 2u) {
O[t * 2u + 1u] = P.lo_y;
}
}
}
}

246
wai/wai-ops-gpu/src/wgsl/ops32.wgsl Normal file
View file

@ -0,0 +1,246 @@
// The operators of wai-ops/1-i32 at the `acc32` widths (integer-operators §7.1,
// §7.3): every tensor element and every `acc` and `wide` value in one word, and
// the two `product` intermediates, rq.p and igdn.gx2, in a W formed from the
// 16-bit halves of their two 32-bit operands. The rounding shift and the `post`
// bias act on that two-word value before it is narrowed.
//
// Sums and products of words are formed in `u32`, which wraps by definition: a
// graph that meets `acc32` forms no value outside i32 (§7.2), so the low word of
// each modular result is the exact value, and no signed overflow is ever asked of
// a translated shader. Only comparisons and right shifts read a word as i32. A
// `pre` bias is added in the same way: its low word added modulo 2^32 gives rq.t
// exactly, because rq.t fits i32.
// The low word of a 64-bit parameter at word offset k.
fn m32(k: u32) -> u32 {
return M[k];
}
fn lt32(a: u32, b: u32) -> bool {
return bitcast<i32>(a) < bitcast<i32>(b);
}
// rq_o(v) (§5.1) under acc32: the multiplier fits i32, the shift is at most 63 and
// a `post` bias fits i32 (§7.2), and rq.t, rq.q and rq.y fit i32.
fn rq(o: u32, v: u32) -> u32 {
var t = v;
if (P.bias == 1u) {
t = t + m32(P.b_off + 2u * o);
}
var p = w_sext(t);
if (P.mk != 0u) {
p = mul_i32(t, m32(P.m_off + select(0u, 2u * o, P.mk == 2u)));
}
let s = M[P.s_off + select(0u, o, P.sk == 2u)];
var y = w_round(p, s, P.round);
if (P.bias == 2u) {
y = w_add(y, w_sext(m32(P.b_off + 2u * o)));
}
return y.x;
}
@compute @workgroup_size(256)
fn conv(@builtin(global_invocation_id) g: vec3<u32>) {
let t = tid(g);
if (t >= P.n) {
return;
}
let x = t % P.wo;
let y = (t / P.wo) % P.ho;
let o = t / (P.wo * P.ho);
let base = (o / P.og) * P.cg;
var acc = 0u;
for (var ci = 0u; ci < P.cg; ci++) {
for (var i = 0u; i < P.kh; i++) {
let iy = y * P.sh + i;
if (iy < P.ph || iy - P.ph >= P.h) {
continue;
}
for (var j = 0u; j < P.kw; j++) {
let ix = x * P.sw + j;
if (ix < P.pw || ix - P.pw >= P.w) {
continue;
}
let wi = ((o * P.cg + ci) * P.kh + i) * P.kw + j;
let xi = ((base + ci) * P.h + iy - P.ph) * P.w + ix - P.pw;
acc = acc + M[P.w_off + wi] * A[xi];
}
}
}
O[t] = rq(o, acc);
}
@compute @workgroup_size(256)
fn conv_t(@builtin(global_invocation_id) g: vec3<u32>) {
let t = tid(g);
if (t >= P.n) {
return;
}
let x = t % P.wo;
let y = (t / P.wo) % P.ho;
let o = t / (P.wo * P.ho);
let yy = y + P.ph;
let xx = x + P.pw;
var acc = 0u;
for (var ci = 0u; ci < P.c; ci++) {
for (var i = yy % P.sh; i < P.kh && i <= yy; i += P.sh) {
let iy = (yy - i) / P.sh;
if (iy >= P.h) {
continue;
}
for (var j = xx % P.sw; j < P.kw && j <= xx; j += P.sw) {
let ix = (xx - j) / P.sw;
if (ix >= P.w) {
continue;
}
let wi = ((ci * P.co + o) * P.kh + i) * P.kw + j;
acc = acc + M[P.w_off + wi] * A[(ci * P.h + iy) * P.w + ix];
}
}
}
O[t] = rq(o, acc);
}
@compute @workgroup_size(256)
fn rescale(@builtin(global_invocation_id) g: vec3<u32>) {
let t = tid(g);
if (t >= P.n) {
return;
}
O[t] = rq(t / (P.h * P.w), A[t]);
}
// chscale under acc32: each entry of M fits i32 and the shift is at most 63.
@compute @workgroup_size(256)
fn chscale(@builtin(global_invocation_id) g: vec3<u32>) {
let t = tid(g);
if (t >= P.n) {
return;
}
let c = t / (P.h * P.w);
O[t] = w_round(mul_i32(A[t], m32(P.m_off + 2u * c)), P.f, P.round).x;
}
@compute @workgroup_size(256)
fn mul(@builtin(global_invocation_id) g: vec3<u32>) {
let t = tid(g);
if (t >= P.n) {
return;
}
let plane = P.h * P.w;
let bi = select(t % plane, t, P.c2 == P.c);
O[t] = rq(t / plane, A[t] * B[bi]);
}
@compute @workgroup_size(256)
fn gram(@builtin(global_invocation_id) g: vec3<u32>) {
let t = tid(g);
if (t >= P.n) {
return;
}
let d = P.k0;
let l = P.k1;
let ca = t / d;
let cb = (ca / d) * d + t % d;
var acc = 0u;
for (var p = 0u; p < l; p++) {
acc = acc + A[ca * l + p] * B[cb * l + p];
}
O[t] = rq(0u, acc);
}
@compute @workgroup_size(256)
fn apply(@builtin(global_invocation_id) g: vec3<u32>) {
let t = tid(g);
if (t >= P.n) {
return;
}
let d = P.k0;
let l = P.k1;
let row = t / l;
let p = t % l;
let g0 = (row / d) * d;
var acc = 0u;
for (var j = 0u; j < d; j++) {
acc = acc + A[row * d + j] * B[(g0 + j) * l + p];
}
O[t] = rq(0u, acc);
}
// igdn under acc32: x², each floored term, N, r and x·r fit i32; γ fits i32 and
// gf and frac are at most 63, so γ·x² is the one product formed in two words.
@compute @workgroup_size(256)
fn igdn(@builtin(global_invocation_id) g: vec3<u32>) {
let t = tid(g);
if (t >= P.n) {
return;
}
let l = P.k1;
let i = t / l;
let p = t % l;
var n = m32(P.b_off + 2u * i);
for (var j = 0u; j < P.c; j++) {
let xj = A[j * l + p];
let gx2 = mul_i32(m32(P.w_off + 2u * (i * P.c + j)), xj * xj);
n = n + w_sar(gx2, P.f).x;
}
let r = isqrt32(select(n, 0u, lt32(n, 0u)));
O[t] = sar32(A[t] * r, P.k0);
}
@compute @workgroup_size(256)
fn lut16(@builtin(global_invocation_id) g: vec3<u32>) {
let t = tid(g);
if (t >= P.n) {
return;
}
let v = clamp(bitcast<i32>(A[t]), -32768, 32767);
O[t] = M[P.w_off + bitcast<u32>(v + 32768)];
}
// lo and hi are the host's clamp bounds held to i32, which gives the same result
// on an i32 value (a bound past i32 is not reached by any i32).
@compute @workgroup_size(256)
fn ew(@builtin(global_invocation_id) g: vec3<u32>) {
let t = tid(g);
if (t >= P.n) {
return;
}
let v = A[t];
switch P.mode {
case 0u: {
O[t] = bitcast<u32>(clamp(bitcast<i32>(v), bitcast<i32>(P.lo_x), bitcast<i32>(P.hi_x)));
}
case 1u: {
O[t] = select(v, 0u, lt32(v, 0u));
}
case 2u: {
O[t] = select(v, 0u - v, lt32(v, 0u));
}
case 3u: {
O[t] = select(v, sar32(v * P.a_x, P.f), lt32(v, 0u));
}
case 4u: {
O[t] = v + B[t];
}
default: {
O[t] = v + A[t + P.k0];
}
}
}
@compute @workgroup_size(256)
fn policy(@builtin(global_invocation_id) g: vec3<u32>) {
let t = tid(g);
if (t >= P.n) {
return;
}
let v = bitcast<i32>(O[t]);
let lo = bitcast<i32>(P.lo_x);
let hi = bitcast<i32>(P.hi_x);
if (P.mode == 0u) {
O[t] = bitcast<u32>(clamp(v, lo, hi));
} else if (v < lo || v > hi) {
atomicOr(&F[P.flag], 1u);
}
}

292
wai/wai-ops-gpu/src/wgsl/ops64.wgsl Normal file
View file

@ -0,0 +1,292 @@
// The operators of wai-ops/1-i32 at the `acc64` widths (integer-operators §7.1):
// every tensor element and every `acc` value in a W (64 bits), every `wide` and
// `product` value in a Q (128 bits). A graph whose certificate meets `acc64` forms
// no value outside these widths (§8), so each modular result here is the exact
// integer of §5.
fn a64(i: u32) -> W {
return W(A[2u * i], A[2u * i + 1u]);
}
fn b64(i: u32) -> W {
return W(B[2u * i], B[2u * i + 1u]);
}
fn o64(i: u32) -> W {
return W(O[2u * i], O[2u * i + 1u]);
}
fn put64(i: u32, v: W) {
O[2u * i] = v.x;
O[2u * i + 1u] = v.y;
}
// A 64-bit parameter at word offset k.
fn m64(k: u32) -> W {
return W(M[k], M[k + 1u]);
}
// Weight i: one word (an i32) or two.
fn weight(i: u32) -> W {
if (P.w_words == 1u) {
return w_sext(M[P.w_off + i]);
}
return m64(P.w_off + 2u * i);
}
// rq_o(v) (§5.1): t = v + pre bias; p = t·mult, in 128 bits; R_round(p, shift);
// + post bias. rq.q fits i64, so its low two words are its value.
fn rq(o: u32, v: W) -> W {
var t = v;
if (P.bias == 1u) {
t = w_add(t, m64(P.b_off + 2u * o));
}
var p = q_from_w(t);
if (P.mk != 0u) {
p = q_mul_ww(t, m64(P.m_off + select(0u, 2u * o, P.mk == 2u)));
}
let s = M[P.s_off + select(0u, o, P.sk == 2u)];
var y = q_round(p, s, P.round).xy;
if (P.bias == 2u) {
y = w_add(y, m64(P.b_off + 2u * o));
}
return y;
}
// conv (§5.2), one output element per thread.
@compute @workgroup_size(256)
fn conv(@builtin(global_invocation_id) g: vec3<u32>) {
let t = tid(g);
if (t >= P.n) {
return;
}
let x = t % P.wo;
let y = (t / P.wo) % P.ho;
let o = t / (P.wo * P.ho);
let base = (o / P.og) * P.cg;
var acc = W(0u, 0u);
for (var ci = 0u; ci < P.cg; ci++) {
for (var i = 0u; i < P.kh; i++) {
let iy = y * P.sh + i;
if (iy < P.ph || iy - P.ph >= P.h) {
continue;
}
for (var j = 0u; j < P.kw; j++) {
let ix = x * P.sw + j;
if (ix < P.pw || ix - P.pw >= P.w) {
continue;
}
let wi = ((o * P.cg + ci) * P.kh + i) * P.kw + j;
let xi = ((base + ci) * P.h + iy - P.ph) * P.w + ix - P.pw;
acc = w_add(acc, w_mul_lo(weight(wi), a64(xi)));
}
}
}
put64(t, rq(o, acc));
}
// conv_transpose (§5.3) in its gather form: for output (o, y, x), the taps (i, j)
// with y + ph − i and x + pw − j non-negative multiples of the strides that land
// inside the input. Those are the taps of one residue, so every partial sum lies in
// that residue's interval, which `convT.acc` certifies.
@compute @workgroup_size(256)
fn conv_t(@builtin(global_invocation_id) g: vec3<u32>) {
let t = tid(g);
if (t >= P.n) {
return;
}
let x = t % P.wo;
let y = (t / P.wo) % P.ho;
let o = t / (P.wo * P.ho);
let yy = y + P.ph;
let xx = x + P.pw;
var acc = W(0u, 0u);
for (var ci = 0u; ci < P.c; ci++) {
for (var i = yy % P.sh; i < P.kh && i <= yy; i += P.sh) {
let iy = (yy - i) / P.sh;
if (iy >= P.h) {
continue;
}
for (var j = xx % P.sw; j < P.kw && j <= xx; j += P.sw) {
let ix = (xx - j) / P.sw;
if (ix >= P.w) {
continue;
}
let wi = ((ci * P.co + o) * P.kh + i) * P.kw + j;
let xi = (ci * P.h + iy) * P.w + ix;
acc = w_add(acc, w_mul_lo(weight(wi), a64(xi)));
}
}
}
put64(t, rq(o, acc));
}
// rescale (§5.4).
@compute @workgroup_size(256)
fn rescale(@builtin(global_invocation_id) g: vec3<u32>) {
let t = tid(g);
if (t >= P.n) {
return;
}
put64(t, rq(t / (P.h * P.w), a64(t)));
}
// chscale (§5.4): the host points m_off at row r of M, and gives the shift in f.
@compute @workgroup_size(256)
fn chscale(@builtin(global_invocation_id) g: vec3<u32>) {
let t = tid(g);
if (t >= P.n) {
return;
}
let c = t / (P.h * P.w);
let p = q_mul_ww(a64(t), m64(P.m_off + 2u * c));
put64(t, q_round(p, P.f, P.round).xy);
}
// mul (§5.4): `mul.p` = a·b fits i64; then rq, which has no bias here.
@compute @workgroup_size(256)
fn mul(@builtin(global_invocation_id) g: vec3<u32>) {
let t = tid(g);
if (t >= P.n) {
return;
}
let plane = P.h * P.w;
let o = t / plane;
let bi = select(t % plane, t, P.c2 == P.c);
put64(t, rq(o, w_mul_lo(a64(t), b64(bi))));
}
// gram (§5.9): out[g·d + i][0][j] = R(Σ A[g·d + i]·B[g·d + j], s); k0 = d, k1 = L.
@compute @workgroup_size(256)
fn gram(@builtin(global_invocation_id) g: vec3<u32>) {
let t = tid(g);
if (t >= P.n) {
return;
}
let d = P.k0;
let l = P.k1;
let ca = t / d;
let cb = (ca / d) * d + t % d;
var acc = W(0u, 0u);
for (var p = 0u; p < l; p++) {
acc = w_add(acc, w_mul_lo(a64(ca * l + p), b64(cb * l + p)));
}
put64(t, rq(0u, acc));
}
// apply (§5.9): out[g·d + i][y][x] = R(Σ_j A[g·d + i][0][j]·V[g·d + j][y][x], s).
@compute @workgroup_size(256)
fn apply(@builtin(global_invocation_id) g: vec3<u32>) {
let t = tid(g);
if (t >= P.n) {
return;
}
let d = P.k0;
let l = P.k1;
let row = t / l;
let p = t % l;
let g0 = (row / d) * d;
var acc = W(0u, 0u);
for (var j = 0u; j < d; j++) {
acc = w_add(acc, w_mul_lo(a64(row * d + j), b64((g0 + j) * l + p)));
}
put64(t, rq(0u, acc));
}
// igdn (§5.7), one output element per thread: N_i = β_i + Σ_j ⌊γ[i][j]·x_j² / 2^gf⌋,
// r_i = isqrt(max(0, N_i)), out = ⌊x_i·r_i / 2^frac⌋. γ at w_off, β at b_off, gf
// in f, frac in k0, L = H·W in k1.
@compute @workgroup_size(256)
fn igdn(@builtin(global_invocation_id) g: vec3<u32>) {
let t = tid(g);
if (t >= P.n) {
return;
}
let l = P.k1;
let i = t / l;
let p = t % l;
var n = q_from_w(m64(P.b_off + 2u * i));
for (var j = 0u; j < P.c; j++) {
let xj = a64(j * l + p);
let x2 = q_mul_ww(xj, xj);
let gx2 = q_mul_lo(q_from_w(m64(P.w_off + 2u * (i * P.c + j))), x2);
n = q_add(n, q_sar(gx2, P.f));
}
if (q_isneg(n)) {
n = Q(0u, 0u, 0u, 0u);
}
let r = q_isqrt(n);
let xr = q_mul_ww(a64(t), r);
put64(t, q_sar(xr, P.k0).xy);
}
// lut16 (§5.6): T[v + 32768], the table one sign-extended word per entry. A
// certified input lies in [−32767, 32767]; the index is held inside the table
// whatever the input, so a speculative run past a refused check reads no further.
@compute @workgroup_size(256)
fn lut16(@builtin(global_invocation_id) g: vec3<u32>) {
let t = tid(g);
if (t >= P.n) {
return;
}
let v = a64(t);
let lo = W(0xFFFF8000u, ONES);
let hi = W(32767u, 0u);
let k = w_clamp(v, lo, hi).x + 32768u;
put64(t, w_sext(M[P.w_off + k]));
}
// The element-wise operators (§5.5), by mode: 0 clamp(lo, hi), 1 relu, 2 abs,
// 3 leaky(a, f), 4 add, 5 chunk_add (k0 = the half's element count).
@compute @workgroup_size(256)
fn ew(@builtin(global_invocation_id) g: vec3<u32>) {
let t = tid(g);
if (t >= P.n) {
return;
}
let lo = W(P.lo_x, P.lo_y);
let hi = W(P.hi_x, P.hi_y);
switch P.mode {
case 0u: {
put64(t, w_clamp(a64(t), lo, hi));
}
case 1u: {
put64(t, w_max(a64(t), W(0u, 0u)));
}
case 2u: {
put64(t, w_mag(a64(t)));
}
case 3u: {
let v = a64(t);
if (w_isneg(v)) {
put64(t, w_sar(w_mul_lo(v, W(P.a_x, P.a_y)), P.f));
} else {
put64(t, v);
}
}
case 4u: {
put64(t, w_add(a64(t), b64(t)));
}
default: {
put64(t, w_add(a64(t), a64(t + P.k0)));
}
}
}
// A tensor's policy (§3), in place on O: mode 0 clamp to [lo, hi]; mode 1 check,
// which sets F[flag] for an element outside it.
@compute @workgroup_size(256)
fn policy(@builtin(global_invocation_id) g: vec3<u32>) {
let t = tid(g);
if (t >= P.n) {
return;
}
let lo = W(P.lo_x, P.lo_y);
let hi = W(P.hi_x, P.hi_y);
let v = o64(t);
if (P.mode == 0u) {
put64(t, w_clamp(v, lo, hi));
} else if (w_lt(v, lo) || w_lt(hi, v)) {
atomicOr(&F[P.flag], 1u);
}
}

58
wai/wai-ops-gpu/src/wgsl/probe.wgsl Normal file
View file

@ -0,0 +1,58 @@
// The arithmetic probe: each record of A (12 words: a as a Q, b as a Q, a shift
// s, a rounding mode, two spare) gives 40 words of O, the results of every
// function of arith.wgsl on it, in the order `src/probe.rs` reads them. The host
// computes the same functions in 128-bit integers and compares every word.
const REC_IN: u32 = 12u;
const REC_OUT: u32 = 40u;
fn put_w(o: ptr<function, u32>, v: W) {
O[*o] = v.x;
O[*o + 1u] = v.y;
*o = *o + 2u;
}
fn put_q(o: ptr<function, u32>, v: Q) {
O[*o] = v.x;
O[*o + 1u] = v.y;
O[*o + 2u] = v.z;
O[*o + 3u] = v.w;
*o = *o + 4u;
}
@compute @workgroup_size(64)
fn probe(@builtin(global_invocation_id) g: vec3<u32>) {
let t = tid(g);
if (t >= P.n) {
return;
}
let i = t * REC_IN;
let a = Q(A[i], A[i + 1u], A[i + 2u], A[i + 3u]);
let b = Q(A[i + 4u], A[i + 5u], A[i + 6u], A[i + 7u]);
let s = A[i + 8u];
let mode = A[i + 9u];
var o = t * REC_OUT;
put_w(&o, w_add(a.xy, b.xy));
put_w(&o, w_mul_lo(a.xy, b.xy));
put_w(&o, mul_i32(a.x, b.x));
put_w(&o, w_round(a.xy, s & 63u, mode));
put_w(&o, w_sar(a.xy, s));
put_q(&o, q_mul_ww(a.xy, b.xy));
put_q(&o, q_mul_lo(a, b));
put_q(&o, q_add(a, b));
put_q(&o, q_round(a, s & 127u, mode));
// isqrt of |a| with its top bit cleared: a non-negative value below 2^127.
put_w(&o, q_isqrt(Q(a.x, a.y, a.z, a.w & 0x7FFFFFFFu)));
O[o] = isqrt32(a.x);
o = o + 1u;
put_w(&o, w_isqrt(a.xy));
O[o] = sar32(a.x, s);
o = o + 1u;
O[o] = select(0u, 1u, w_lt(a.xy, b.xy)) | select(0u, 2u, w_ltu(a.xy, b.xy)) | select(0u, 4u, q_ltu(a, b)) | select(0u, 8u, q_isneg(a)) | select(0u, 16u, w_isneg(a.xy));
o = o + 1u;
put_w(&o, w_neg(a.xy));
put_w(&o, w_clamp(a.xy, b.xy, w_add(b.xy, W(0u, 1u))));
O[o] = 0u;
O[o + 1u] = 0u;
O[o + 2u] = 0u;
}

64
wai/wai-ops-gpu/tests/arith.rs Normal file
View file

@ -0,0 +1,64 @@
//! The kernels' arithmetic on the device (`wgsl/arith.wgsl`): every 64-bit and 128-bit
//! function, from 32-bit words, against the host's 128-bit integers, on the edges of
//! every width pairwise, every shift under every rounding mode, and a million SplitMix64
//! cases, on every adapter found. Then an edit of one function (a carry dropped from the
//! 32 × 32-bit product) is shown to be caught: discovery refuses the adapter.
use std::io::Write;
use wai_ops_gpu::probe::{self, ProbeIn};
use wai_ops_gpu::{blocking, kernels::Sources};
/// A result line, also appended to `WAI_OPS_GPU_REPORT` when it is set.
fn report(line: &str) {
println!("{line}");
if let Ok(p) = std::env::var("WAI_OPS_GPU_REPORT") {
let mut f = std::fs::OpenOptions::new().create(true).append(true).open(&p).unwrap();
writeln!(f, "{line}").unwrap();
}
}
#[test]
fn the_arithmetic_is_the_hosts_on_every_adapter() {
let d = blocking::discover();
if d.found.is_empty() {
println!("ops-gpu arith: SKIPPED: no adapter: {:?}", d.refused.iter().map(|(a, w)| (a.key(), w.clone())).collect::<Vec<_>>());
assert!(std::env::var_os("WAI_OPS_GPU_REQUIRE").is_none(), "WAI_OPS_GPU_REQUIRE is set and no adapter was found");
return;
}
let mut cases: Vec<ProbeIn> = probe::cases(0, 0);
for seed in 0..16u64 {
let c = probe::cases(65536, 0xA417_0000 + seed);
cases.extend_from_slice(&c[c.len() - 65536..]);
}
for g in &d.found {
let mut checked = 0usize;
for chunk in cases.chunks(1 << 16) {
let got = blocking::probe(g, chunk).unwrap_or_else(|e| panic!("{}: {e}", g.adapter().key()));
assert_eq!(probe::first_mismatch(chunk, &got), None, "{}", g.adapter().key());
checked += chunk.len();
}
report(&format!("ops-gpu arith: key={} {} cases={checked} words={} equal", g.adapter().key(), g.adapter().describe(), checked * probe::OUT_WORDS));
}
}
#[test]
fn an_edited_product_is_caught_by_the_probe() {
let d = blocking::discover();
if d.found.is_empty() {
println!("ops-gpu arith mutant: SKIPPED: no adapter");
assert!(std::env::var_os("WAI_OPS_GPU_REQUIRE").is_none(), "WAI_OPS_GPU_REQUIRE is set and no adapter was found");
return;
}
let mut src = Sources::shipped();
let from = "let hi = p11 + (mid2 >> 16u) + ((c1 + c2) << 16u);";
assert!(src.probe.contains(from));
src.probe = src.probe.replace(from, "let hi = p11 + (mid2 >> 16u) + (c2 << 16u);");
let m = blocking::discover_with(&src, "mutant: a carry dropped");
assert!(m.found.is_empty(), "an adapter passed the probe with a carry dropped");
for (a, why) in &m.refused {
assert!(why.contains("probe"), "{}: refused for another reason: {why}", a.key());
println!("ops-gpu arith mutant: key={} refused: {why}", a.key());
}
assert_eq!(m.refused.len(), d.found.len() + d.refused.len());
}

134
wai/wai-ops-gpu/tests/corpus.rs Normal file
View file

@ -0,0 +1,134 @@
//! Every graph vector of `ops-conformance/` that `wai-ops/1-i32` covers, on every adapter
//! this process finds: each decode's output integers, or its refusal code, must be the
//! committed ones, which are the CPU interpreter's (checked here too, decode by decode).
//! Each graph runs on the 64-bit kernels, and on the 32-bit kernels too where its
//! certificate meets `acc32`.
//!
//! **No adapter is not a pass.** With no adapter the test prints `ops-gpu corpus:
//! SKIPPED` and the reason, and fails if `WAI_OPS_GPU_REQUIRE` is set. With
//! `WAI_OPS_GPU_EXPECT` (a comma-separated list of adapter keys, `Adapter::key`), the
//! keys found must be exactly those. Each adapter's result is one line, `ops-gpu
//! corpus: key=…`, also appended to `WAI_OPS_GPU_REPORT` when it is set; a CI step
//! requires that line with its counts.
#[path = "support/corpus.rs"]
mod corpus;
use std::io::Write;
use wai_ops_gpu::blocking;
fn report(line: &str) {
println!("{line}");
if let Ok(p) = std::env::var("WAI_OPS_GPU_REPORT") {
let mut f = std::fs::OpenOptions::new().create(true).append(true).open(&p).unwrap_or_else(|e| panic!("{p}: {e}"));
writeln!(f, "{line}").unwrap();
}
}
#[test]
fn every_vector_the_subset_covers_is_the_interpreters_on_every_adapter() {
let d = blocking::discover();
for (a, why) in &d.refused {
report(&format!("ops-gpu adapter refused: key={} {} reason={why:?}", a.key(), a.describe()));
}
let mut keys: Vec<String> = d.found.iter().map(|g| g.adapter().key()).collect();
keys.sort();
if let Ok(want) = std::env::var("WAI_OPS_GPU_EXPECT") {
let mut want: Vec<String> = want.split(',').filter(|s| !s.is_empty()).map(str::to_string).collect();
want.sort();
assert_eq!(keys, want, "the adapters found are not the ones WAI_OPS_GPU_EXPECT declares");
report(&format!("ops-gpu adapters: as declared ({})", keys.join(",")));
}
if d.found.is_empty() {
let why = if d.refused.is_empty() { "the platform offers no adapter".to_string() } else { format!("{} adapter(s) refused, above", d.refused.len()) };
report(&format!("ops-gpu corpus: SKIPPED: no adapter ran the corpus: {why}"));
assert!(std::env::var_os("WAI_OPS_GPU_REQUIRE").is_none(), "WAI_OPS_GPU_REQUIRE is set and no adapter ran the corpus");
return;
}
let vs = corpus::vectors();
for gpu in &d.found {
let a = gpu.adapter();
let t0 = std::time::Instant::now();
let t = corpus::run_all(gpu, &vs, 20);
for m in &t.mismatches {
println!("MISMATCH {m}");
}
for e in &t.errors {
println!("DEVICE ERROR {e}");
}
report(&format!(
"ops-gpu corpus: key={} {} acc32 vectors={} decodes={}; acc64 vectors={} decodes={}; refusals={}; cpu={} load_refused={} mismatches={} errors={} secs={:.1}",
a.key(),
a.describe(),
t.vectors32,
t.decodes32,
t.vectors64,
t.decodes64,
t.refusals,
t.cpu,
t.load_refused,
t.mismatches.len(),
t.errors.len(),
t0.elapsed().as_secs_f64()
));
assert!(t.errors.is_empty(), "{}: device errors: {:?}", a.key(), t.errors);
assert!(t.mismatches.is_empty(), "{}: {} mismatches, the first: {:?}", a.key(), t.mismatches.len(), t.mismatches.first());
assert!(t.vectors64 > 1000 && t.vectors32 > 500, "{}: too few vectors ran ({} at acc64, {} at acc32)", a.key(), t.vectors64, t.vectors32);
}
}
/// Why the width is chosen by the certificate: the 32-bit kernels, run on the graphs whose
/// certificate does not meet `acc32` (which [`wai_ops_gpu::Gpu::run_at`] refuses, and
/// which only this hidden entry point allows), give results that differ from the
/// interpreter's on some of them. On every one of those the 64-bit kernels agree
/// (the test above).
#[test]
fn the_32_bit_kernels_are_not_exact_where_the_certificate_does_not_admit_them() {
use wai_ops::Profile;
use wai_ops_gpu::{GpuSequence, Width, support};
let d = blocking::discover();
if d.found.is_empty() {
println!("ops-gpu width: SKIPPED: no adapter");
assert!(std::env::var_os("WAI_OPS_GPU_REQUIRE").is_none(), "WAI_OPS_GPU_REQUIRE is set and no adapter was found");
return;
}
let vs = corpus::vectors();
for gpu in &d.found {
let (mut ran, mut differ, mut unrunnable) = (0u32, 0u32, 0u32);
for v in &vs {
let Ok(l) = wai_ops::graph::load(&v.graph, Profile::Acc64) else { continue };
if l.meets(Profile::Acc32) || support(&l).is_err() {
continue;
}
assert_eq!(
blocking::run_at(gpu, &l, Width::Acc32, &[], &[], &mut GpuSequence::new()).err(),
Some(wai_ops_gpu::GpuError::Fallback(wai_ops_gpu::Fallback::Profile)),
"{}: run_at refuses the 32-bit kernels for a graph that does not meet acc32",
v.id
);
ran += 1;
let mut seq = GpuSequence::new();
for dd in &v.decodes {
match wai_ops_gpu::blocking::resolve(gpu.run_unchecked(&l, Width::Acc32, &dd.inputs, &dd.scalars, &mut seq)) {
Ok(Ok(r)) => {
let got = match r {
Ok(t) => corpus::Expect::Outputs(t),
Err(r) => corpus::Expect::Refuse(r.code.as_str().to_string()),
};
if got != dd.expect {
differ += 1;
break;
}
}
_ => {
unrunnable += 1;
break;
}
}
}
}
println!("ops-gpu width: key={} acc64-only graphs={ran}: on the 32-bit kernels {differ} differ from the interpreter, {unrunnable} have parameters past those kernels' widths", gpu.adapter().key());
assert!(differ > 0, "{}: the 32-bit kernels agreed on every graph the certificate keeps from them", gpu.adapter().key());
}
}

95
wai/wai-ops-gpu/tests/mutants.rs Normal file
View file

@ -0,0 +1,95 @@
//! The corpus catches an edited kernel. Each mutant below changes one expression of
//! one shader module, the kind of slip an integer kernel makes: a rounding shift off by
//! one, a wrong saturation bound, a dropped carry or cross term, a sign rule, an index
//! rule. A device running the mutant must give at least one result of the corpus
//! (`tests/corpus.rs`) that differs from the committed one. The arithmetic probe is left
//! as shipped, so the device opens and the corpus alone has to catch the edit.
//!
//! Each adapter's result is one line, `ops-gpu mutants: key=… killed N of N`, appended
//! to `WAI_OPS_GPU_REPORT` when it is set.
#[path = "support/corpus.rs"]
mod corpus;
use std::io::Write;
use wai_ops_gpu::{blocking, kernels::Sources};
#[derive(Clone, Copy)]
enum Module {
Ops32,
Ops64,
Move,
}
/// `(name, module, the text, its replacement)`; the text must occur in the module.
const MUTANTS: &[(&str, Module, &str, &str)] = &[
("half_up addend off by one bit (64-bit rounding, acc32)", Module::Ops32, "let h = w_pow2(s - 1u);", "let h = w_pow2(s);"),
("half_up addend off by one bit (128-bit rounding, acc64)", Module::Ops64, "let h = q_pow2(s - 1u);", "let h = q_pow2(s);"),
("rounding shift one short (acc64)", Module::Ops64, "return q_sar(q_add(p, h), s);", "return q_sar(q_add(p, h), s - 1u);"),
("half_even ties to odd (acc64)", Module::Ops64, "(q_eq(rem, h) && (q.x & 1u) == 1u)", "(q_eq(rem, h) && (q.x & 1u) == 0u)"),
("half_even ties to odd (acc32)", Module::Ops32, "(w_eq(rem, h) && (q.x & 1u) == 1u)", "(w_eq(rem, h) && (q.x & 1u) == 0u)"),
("clamp operator saturates one below its top (acc64)", Module::Ops64, "put64(t, w_clamp(a64(t), lo, hi));", "put64(t, w_clamp(a64(t), lo, w_sub(hi, W(1u, 0u))));"),
("clamp policy saturates one above its bottom (acc32)", Module::Ops32, "O[t] = bitcast<u32>(clamp(v, lo, hi));", "O[t] = bitcast<u32>(clamp(v, lo + 1, hi));"),
("check policy misses values above its range (acc64)", Module::Ops64, "} else if (w_lt(v, lo) || w_lt(hi, v)) {", "} else if (w_lt(v, lo)) {"),
// In the 32-bit kernels the 32 × 32-bit product only ever sees magnitudes of at most
// 2^31, for which neither carry of its middle sum can be set (a0·b1 + a1·b0 stays
// below 2^32 − 2^17), so dropping one there is an equivalent mutant: it is dropped
// in the 64-bit kernels, whose products take full 32-bit words.
("product's middle carry dropped (acc64)", Module::Ops64, "let hi = p11 + (mid2 >> 16u) + ((c1 + c2) << 16u);", "let hi = p11 + (mid2 >> 16u) + (c2 << 16u);"),
("64-bit product's cross term dropped (acc64)", Module::Ops64, "return W(p.x, p.y + a.x * b.y + a.y * b.x);", "return W(p.x, p.y + a.x * b.y);"),
("signed product's sign rule (acc32)", Module::Ops32, "return select(m, w_neg(m), na != nb);", "return select(m, w_neg(m), na);"),
("post bias dropped (acc32)", Module::Ops32, "y = w_add(y, w_sext(m32(P.b_off + 2u * o)));", "y = y;"),
("transposed convolution's residue (acc64)", Module::Ops64, "for (var i = yy % P.sh; i < P.kh && i <= yy; i += P.sh) {", "for (var i = (yy + 1u) % P.sh; i < P.kh && i <= yy; i += P.sh) {"),
("convolution's bottom padding edge (acc32)", Module::Ops32, "if (iy < P.ph || iy - P.ph >= P.h) {", "if (iy < P.ph || iy - P.ph > P.h) {"),
("igdn term floored one bit further (acc64)", Module::Ops64, "n = q_add(n, q_sar(gx2, P.f));", "n = q_add(n, q_sar(gx2, P.f + 1u));"),
("leaky shift of 32 or more (acc32)", Module::Ops32, "if (s >= 32u) {\n return bitcast<u32>(bitcast<i32>(a) >> 31u);", "if (s > 32u) {\n return bitcast<u32>(bitcast<i32>(a) >> 31u);"),
("pixel_shuffle's sub-position order", Module::Move, "(y % r) * r + x % r", "(x % r) * r + y % r"),
("subpos_merge's column offset", Module::Move, "copy_el((c * P.ho + 2u * y + P.k0) * P.wo + 2u * x + P.k1, t);", "copy_el((c * P.ho + 2u * y + P.k1) * P.wo + 2u * x + P.k0, t);"),
];
fn report(line: &str) {
println!("{line}");
if let Ok(p) = std::env::var("WAI_OPS_GPU_REPORT") {
let mut f = std::fs::OpenOptions::new().create(true).append(true).open(&p).unwrap();
writeln!(f, "{line}").unwrap();
}
}
#[test]
fn every_mutant_is_caught_on_every_adapter() {
let d = blocking::discover();
if d.found.is_empty() {
report("ops-gpu mutants: SKIPPED: no adapter");
assert!(std::env::var_os("WAI_OPS_GPU_REQUIRE").is_none(), "WAI_OPS_GPU_REQUIRE is set and no adapter was found");
return;
}
let keys: Vec<String> = d.found.iter().map(|g| g.adapter().key()).collect();
let vs = corpus::vectors();
let mut killed = vec![0usize; keys.len()];
for &(name, module, from, to) in MUTANTS {
let mut src = Sources::shipped();
let text = match module {
Module::Ops32 => &mut src.ops32,
Module::Ops64 => &mut src.ops64,
Module::Move => &mut src.mv,
};
assert!(text.contains(from), "mutant {name:?}: its text is not in the module");
*text = text.replace(from, to);
let m = blocking::discover_with(&src, name);
assert_eq!(m.found.len(), keys.len(), "mutant {name:?}: the adapters did not all open: {:?}", m.refused.iter().map(|(_, w)| w).collect::<Vec<_>>());
for g in &m.found {
let k = keys.iter().position(|x| *x == g.adapter().key()).expect("the same adapters");
let t = corpus::run_all(g, &vs, 1);
let caught = !t.mismatches.is_empty() || !t.errors.is_empty();
println!("ops-gpu mutant: key={} {name:?}: {}", keys[k], t.mismatches.first().map_or_else(|| "SURVIVED".to_string(), |m| format!("caught: {m}")));
killed[k] += caught as usize;
}
}
for (k, key) in keys.iter().enumerate() {
report(&format!("ops-gpu mutants: key={key} killed {} of {}", killed[k], MUTANTS.len()));
}
for (k, key) in keys.iter().enumerate() {
assert_eq!(killed[k], MUTANTS.len(), "{key}: a mutant survived");
}
}

204
wai/wai-ops-gpu/tests/support/corpus.rs Normal file
View file

@ -0,0 +1,204 @@
//! The graph vectors of `ops-conformance/` (`vectors/` and `graphs/`), run on a device and
//! held to the committed expectations and to the CPU interpreter.
#![allow(dead_code)]
use std::path::PathBuf;
use wai_ops::json::{self, Value, unhex};
use wai_ops::vectors::read_tensor;
use wai_ops::{Loaded, Profile, Sequence, Tensor};
use wai_ops_gpu::{Gpu, GpuError, GpuSequence, Width, blocking, support};
pub fn conformance() -> PathBuf {
PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../ops-conformance")
}
/// Every graph-vector file: `vectors/*.jsonl`, then `graphs/*.jsonl`, each sorted by name.
pub fn files() -> Vec<(String, String)> {
let mut out = Vec::new();
for dir in ["vectors", "graphs"] {
let d = conformance().join(dir);
let mut names: Vec<String> = std::fs::read_dir(&d)
.unwrap_or_else(|e| panic!("{}: {e}", d.display()))
.filter_map(|e| e.ok().map(|e| e.file_name().to_string_lossy().into_owned()))
.filter(|n| n.ends_with(".jsonl"))
.collect();
names.sort();
for n in names {
let text = std::fs::read_to_string(d.join(&n)).unwrap();
out.push((format!("{dir}/{n}"), text));
}
}
assert!(out.len() > 30, "the corpus was not found");
out
}
/// A vector's expectation for one decode: the outputs, or a refusal code.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Expect {
Outputs(Vec<Tensor>),
Refuse(String),
}
pub struct Decode {
pub inputs: Vec<Tensor>,
pub scalars: Vec<i64>,
pub expect: Expect,
}
pub struct Vector {
pub file: String,
pub id: String,
pub graph: Vec<u8>,
pub load: String,
pub decodes: Vec<Decode>,
}
fn dec(v: &Value) -> i64 {
v.as_dec_i64().expect("an i64")
}
pub fn vectors() -> Vec<Vector> {
files().iter().flat_map(|(file, text)| parse(file, text)).collect()
}
/// The vectors of one file.
pub fn parse(file: &str, text: &str) -> Vec<Vector> {
let mut out = Vec::new();
for line in text.split_terminator('\n') {
let v = json::parse(line.as_bytes()).expect("a canonical line");
let id = v.get("id").and_then(Value::as_str).unwrap().to_string();
let graph = unhex(v.get("graph").and_then(Value::as_str).unwrap()).unwrap();
let load = v.get("load").and_then(Value::as_str).unwrap().to_string();
let decodes = v
.get("decodes")
.and_then(Value::as_arr)
.unwrap()
.iter()
.map(|d| {
let inputs = d.get("inputs").and_then(Value::as_arr).unwrap().iter().map(|t| read_tensor(t).unwrap()).collect();
let scalars = d.get("scalars").and_then(Value::as_arr).unwrap().iter().map(dec).collect();
let e = d.get("expect").unwrap();
let expect = match e.get("refuse") {
Some(c) => Expect::Refuse(c.as_str().unwrap().to_string()),
None => Expect::Outputs(e.get("outputs").and_then(Value::as_arr).unwrap().iter().map(|t| read_tensor(t).unwrap()).collect()),
};
Decode { inputs, scalars, expect }
})
.collect();
out.push(Vector { file: file.to_string(), id, graph, load, decodes });
}
out
}
/// What one adapter did with the corpus.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct Tally {
/// Vectors run, and their decodes, at each width.
pub vectors32: u32,
pub decodes32: u32,
pub vectors64: u32,
pub decodes64: u32,
/// Of those decodes, the ones whose expectation is a refusal.
pub refusals: u32,
/// Graphs with an operator outside wai-ops/1-i32: run on the CPU, not here.
pub cpu: u32,
/// Vectors whose load is refused: nothing for a device to run.
pub load_refused: u32,
/// Decodes whose device result differs from the expectation, with what differed.
pub mismatches: Vec<String>,
/// Device errors.
pub errors: Vec<String>,
}
fn got_of(r: Result<Result<Vec<Tensor>, wai_ops::Refusal>, GpuError>) -> Result<Expect, String> {
match r {
Ok(Ok(t)) => Ok(Expect::Outputs(t)),
Ok(Err(r)) => Ok(Expect::Refuse(r.code.as_str().to_string())),
Err(e) => Err(e.to_string()),
}
}
fn cpu_of(l: &Loaded, d: &Decode, seq: &mut Sequence) -> Expect {
match wai_ops::interp::run(l, &d.inputs, &d.scalars, seq) {
Ok(t) => Expect::Outputs(t),
Err(r) => Expect::Refuse(r.code.as_str().to_string()),
}
}
/// Run every vector on `gpu` at every width its graph admits (acc64 always; acc32 too
/// where its certificate meets it), stopping a width's sequence at its first mismatch.
/// `limit` caps the mismatches recorded.
pub fn run_all(gpu: &Gpu, vs: &[Vector], limit: usize) -> Tally {
let mut t = Tally::default();
for v in vs {
let l = match wai_ops::graph::load(&v.graph, Profile::Acc64) {
Ok(l) => l,
Err(r) => {
assert_eq!(r.code.as_str(), v.load, "{}: the load the vector records", v.id);
t.load_refused += 1;
continue;
}
};
assert_eq!(v.load, "ok", "{}", v.id);
if support(&l).is_err() {
t.cpu += 1;
continue;
}
let mut widths = vec![Width::Acc64];
if l.meets(Profile::Acc32) {
widths.push(Width::Acc32);
}
for w in widths {
let mut seq = GpuSequence::new();
let mut cpu = Sequence::new();
match w {
Width::Acc32 => t.vectors32 += 1,
Width::Acc64 => t.vectors64 += 1,
}
for (k, d) in v.decodes.iter().enumerate() {
let want = cpu_of(&l, d, &mut cpu);
assert_eq!(want, d.expect, "{}: decode {k}: the CPU interpreter's result is the committed one", v.id);
match w {
Width::Acc32 => t.decodes32 += 1,
Width::Acc64 => t.decodes64 += 1,
}
if matches!(d.expect, Expect::Refuse(_)) {
t.refusals += 1;
}
match got_of(blocking::run_at(gpu, &l, w, &d.inputs, &d.scalars, &mut seq)) {
Ok(got) if got == d.expect => {}
Ok(got) => {
if t.mismatches.len() < limit {
t.mismatches.push(format!("{} ({}) decode {k} at {}: {}", v.id, v.file, w.as_str(), diff(&d.expect, &got)));
}
break;
}
Err(e) => {
t.errors.push(format!("{} decode {k} at {}: {e}", v.id, w.as_str()));
break;
}
}
}
}
}
t
}
fn diff(want: &Expect, got: &Expect) -> String {
match (want, got) {
(Expect::Outputs(a), Expect::Outputs(b)) => {
for (k, (x, y)) in a.iter().zip(b).enumerate() {
if (x.c, x.h, x.w) != (y.c, y.h, y.w) {
return format!("output {k} shape {:?} vs {:?}", (x.c, x.h, x.w), (y.c, y.h, y.w));
}
if let Some(i) = (0..x.data.len()).find(|&i| x.data[i] != y.data[i]) {
return format!("output {k} element {i}: want {} got {}", x.data[i], y.data[i]);
}
}
"output count".into()
}
_ => format!("want {want:?} got {got:?}").chars().take(300).collect(),
}
}

9
wai/wai-rs/Cargo.lock generated
View file

@ -4298,6 +4298,7 @@ dependencies = [
"serde_json",
"sha2 0.10.9",
"wai-ops",
"wai-ops-gpu",
"wai-quantum",
"xz2",
"zstd",
@ -4308,6 +4309,14 @@ dependencies = [
name = "wai-ops"
version = "0.1.0"
[[package]]
name = "wai-ops-gpu"
version = "0.1.0"
dependencies = [
"wai-ops",
"wgpu",
]
[[package]]
name = "wai-quantum"
version = "0.6.12"

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

@ -108,6 +108,9 @@ wai-quantum = { version = "0.6.0", path = "../wai-quantum", optional = true, def
# capability's profile, the bounds certificate included (src/int_ops.rs). no_std, no
# dependencies of its own.
wai-ops = { path = "../wai-ops", optional = true }
# `wai-ops/1-i32` on GPUs (feature `ops_gpu`): WGSL kernels through wgpu. Optional and off
# by default, so a default build has no GPU interface in it.
wai-ops-gpu = { path = "../wai-ops-gpu", optional = true }
[features]
default = ["codecs"]
@ -146,6 +149,12 @@ legacy_int = ["neural_int"]
# arithmetic a 32-bit fabric runs, which the parity tests hold to the exact decode. For
# the integer GPU legs; no sink decodes through it by default.
i32_forms = ["neural_int", "wai-ops/i32_forms"]
# The integer decoders' graphs on a GPU (`int_ops::gpu`, `wai_ops_gpu`): `wai-ops/1-i32` in
# WGSL kernels through wgpu, at the width each graph's certificate admits. A decode runs
# its graphs there only inside `int_ops::gpu::with_gpu`; everywhere else, and for a graph
# with an operator outside `wai-ops/1-i32`, the interpreter runs them. Native only: on
# wasm32 the feature builds and changes nothing.
ops_gpu = ["neural_int", "dep:wai-ops-gpu"]
# Deterministic integer generative synthesis (`wai.generate.*`): byte-exact
# generation pinned to (model, prompt, seed, steps) — the create-side of the
# integer contract. Pure i64, no ONNX/float (serde_json is already a base dep),
@ -889,6 +898,15 @@ name = "int_splat_ops"
path = "tests/int_splat_ops.rs"
required-features = ["splat_int", "legacy_int"]
# The integer decoders with their graphs on a GPU (`int_ops::gpu`): every committed vector
# and refusal fixture of every capability that runs a `wai-ops/1` graph, against the decode
# on the CPU, on every adapter found. Needs a GPU (or a software adapter) to run anything;
# with none it prints SKIPPED, and fails under WAI_OPS_GPU_REQUIRE.
[[test]]
name = "ops_gpu"
path = "tests/ops_gpu.rs"
required-features = ["ops_gpu", "video_int", "audio_int", "splat_int", "world_int"]
[[test]]
name = "int_hyper_ops"
path = "tests/int_hyper_ops.rs"

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

@ -93,10 +93,139 @@ pub fn whm1_net(m: &IntModel) -> Result<Net, Refusal> {
/// One decode of a loaded graph with one data input and no state, through the
/// interpreter with this crate's kernels behind it ([`Host`]): its first output.
pub fn run1(g: &Loaded, x: Tensor, scalars: &[i64]) -> Result<Tensor, Refusal> {
let mut out = wai_ops::run_with(g, vec![x], scalars, &mut Sequence::new(), &Host)?;
let mut out = run_graph(g, vec![x], scalars)?;
Ok(out.swap_remove(0))
}
/// One decode of a loaded graph with no state: its outputs. Every integer decoder here
/// runs its graphs through this. It runs the interpreter with this crate's kernels
/// behind it ([`Host`]); with the `ops_gpu` feature, inside [`gpu::with_gpu`], it runs
/// the graph on that device instead wherever the device runs it, whose outputs and
/// refusals are the interpreter's (`wai_ops_gpu`), and on the interpreter where it does
/// not, counting why ([`gpu::Stats`]).
pub fn run_graph(g: &Loaded, inputs: Vec<Tensor>, scalars: &[i64]) -> Result<Vec<Tensor>, Refusal> {
#[cfg(all(feature = "ops_gpu", not(target_arch = "wasm32")))]
if let Some(r) = gpu::try_run(g, &inputs, scalars) {
return r;
}
wai_ops::run_with(g, inputs, scalars, &mut Sequence::new(), &Host)
}
/// The integer decoders' graphs on a GPU (`wai_ops_gpu`): `wai-ops/1-i32` in WGSL kernels,
/// at the width each graph's certificate admits (`extensions/integer-operators.md` §7.3).
/// A decode inside [`with_gpu`] runs each graph [`run_graph`] is given on the device,
/// except a graph with an operator outside `wai-ops/1-i32`, which runs on the interpreter;
/// [`Stats`] counts both, by width and by reason, and records any device error (the
/// graph then runs on the interpreter too). The outputs are the same bytes either way.
///
/// The device is the process's while [`with_gpu`] runs, whichever thread a decode runs
/// its graphs on; a second [`with_gpu`] waits for the first to end.
#[cfg(all(feature = "ops_gpu", not(target_arch = "wasm32")))]
pub mod gpu {
use std::collections::BTreeMap;
use std::sync::{Arc, Mutex, MutexGuard};
pub use wai_ops_gpu;
use wai_ops::{Loaded, Refusal, Tensor};
use wai_ops_gpu::{Gpu, GpuError, GpuSequence, Width};
/// What a run inside [`with_gpu`] did with its graphs.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct Stats {
/// Graph runs on the device at `acc32` and at `acc64`.
pub acc32: u64,
pub acc64: u64,
/// Graph runs on the interpreter, by reason.
pub cpu: BTreeMap<String, u64>,
/// Device errors, each followed by a run on the interpreter.
pub errors: Vec<String>,
/// Under [`with_gpu_compared`]: the device's runs also run on the interpreter,
/// and how many gave the interpreter's outputs or refusal, and how many did not.
pub compared: u64,
pub differed: u64,
}
impl Stats {
pub fn on_gpu(&self) -> u64 {
self.acc32 + self.acc64
}
pub fn on_cpu(&self) -> u64 {
self.cpu.values().sum()
}
}
struct State {
device: Option<Arc<Gpu>>,
compare: bool,
stats: Stats,
}
static SESSION: Mutex<()> = Mutex::new(());
static STATE: Mutex<State> =
Mutex::new(State { device: None, compare: false, stats: Stats { acc32: 0, acc64: 0, cpu: BTreeMap::new(), errors: Vec::new(), compared: 0, differed: 0 } });
fn lock<T>(m: &Mutex<T>) -> MutexGuard<'_, T> {
m.lock().unwrap_or_else(|p| p.into_inner())
}
/// Run `f` with the process's graphs on `g`: what it returns, and what its graphs did.
pub fn with_gpu<R>(g: &Arc<Gpu>, f: impl FnOnce() -> R) -> (R, Stats) {
session(g, false, f)
}
/// [`with_gpu`], with every graph the device runs also run on the interpreter and
/// the two results compared, output by output ([`Stats::compared`],
/// [`Stats::differed`]). The device's result is the one the decode uses. For tests.
pub fn with_gpu_compared<R>(g: &Arc<Gpu>, f: impl FnOnce() -> R) -> (R, Stats) {
session(g, true, f)
}
fn session<R>(g: &Arc<Gpu>, compare: bool, f: impl FnOnce() -> R) -> (R, Stats) {
let _session = lock(&SESSION);
{
let mut s = lock(&STATE);
s.device = Some(Arc::clone(g));
s.compare = compare;
s.stats = Stats::default();
}
let r = f();
let mut s = lock(&STATE);
s.device = None;
s.compare = false;
(r, std::mem::take(&mut s.stats))
}
/// The graph on the process's device, if it has one and the device runs the graph.
pub(super) fn try_run(g: &Loaded, inputs: &[Tensor], scalars: &[i64]) -> Option<Result<Vec<Tensor>, Refusal>> {
let (dev, compare) = {
let s = lock(&STATE);
(s.device.clone()?, s.compare)
};
let width = wai_ops_gpu::support(g);
let r = wai_ops_gpu::blocking::run(&dev, g, inputs, scalars, &mut GpuSequence::new());
let same = match (&r, compare) {
(Ok(got), true) => Some(*got == wai_ops::run_with(g, inputs.to_vec(), scalars, &mut wai_ops::Sequence::new(), &super::Host)),
_ => None,
};
let mut st = lock(&STATE);
let s = &mut st.stats;
if let Some(same) = same {
s.compared += 1;
s.differed += !same as u64;
}
match (&r, width) {
(Ok(_), Ok(Width::Acc32)) => s.acc32 += 1,
(Ok(_), _) => s.acc64 += 1,
(Err(GpuError::Fallback(why)), _) => *s.cpu.entry(why.to_string()).or_default() += 1,
(Err(e), _) => {
s.errors.push(e.to_string());
*s.cpu.entry("a device error".into()).or_default() += 1;
}
}
r.ok()
}
}
/// This crate's fast kernels, lent to the interpreter (see the module docs).
pub struct Host;

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

@ -378,7 +378,7 @@ pub fn load_networks_net(n: wai_ops::lower::jpegai_nets::JpegaiNets) -> Result<N
}
fn run_graph(g: &Loaded, inputs: Vec<Tensor>) -> Result<Tensor, Refusal> {
let mut out = wai_ops::run_with(g, inputs, &[], &mut wai_ops::Sequence::new(), &crate::int_ops::Host)?;
let mut out = crate::int_ops::run_graph(g, inputs, &[])?;
Ok(out.swap_remove(0))
}

386
wai/wai-rs/tests/ops_gpu.rs Normal file
View file

@ -0,0 +1,386 @@
//! The integer decoders with their graphs on a GPU (`int_ops::gpu`, `wai_ops_gpu`): every
//! committed vector and refusal fixture of every capability whose decode runs a
//! `wai-ops/1` graph decodes to the same bytes, or is refused with the same code, as the
//! decode whose graphs run on the CPU, on every adapter this process finds.
//!
//! - `int-payload-conformance/vectors.json`: its payload vectors and refusal fixtures
//! (`wai.neural.int_hyper`, `wai.video.int_hyper`, `wai.audio.int_codec`,
//! `wai.splat.int_codec`; `wai.video.int_motion` runs no graph, and is listed with no
//! graph runs);
//! - `neural-conformance/`: the `wai.neural.int_synth` vectors;
//! - `wai-rs/tests/vectors/int_mlicpp`: the `wai.neural.int_mlicpp` vector;
//! - `wai-rs/tests/vectors/int_world{,2}`: the integer renderer of `wai.world.model`
//! (informative, integer-operators §13.5);
//! - `wai.image.jpegai`'s exact decoder on the `jpegai/v1` inputs, when
//! `WAI_JPEGAI_CODESTREAM`, `WAI_JPEGAI_WEIGHTS_DIR` and `WAI_JPEGAI_NORM` name them
//! (`tools/wai_jpegai_fetch_e2e_inputs.sh`). Its scale graphs and hyper decoders run on
//! the device; its synthesis graphs use operators outside `wai-ops/1-i32` and run on the
//! CPU, which the line for it states.
//!
//! Each graph the device runs is also run on the interpreter, and its outputs compared
//! (`gpu::with_gpu_compared`): so the equality is shown graph by graph, and not only
//! through each decode's final bytes.
//!
//! **A skip is not a pass.** With no adapter the test prints `ops-gpu capabilities:
//! SKIPPED` and fails when `WAI_OPS_GPU_REQUIRE` is set; without the `jpegai/v1` inputs it
//! prints that `wai.image.jpegai` was not run, and fails when `WAI_REQUIRE_SOTA_VECTORS`
//! is set. A capability whose decodes ran no graph on the device fails. Each adapter's
//! result is one line per capability, `ops-gpu capability: key=…`, appended to
//! `WAI_OPS_GPU_REPORT` when it is set.
#![cfg(all(feature = "ops_gpu", feature = "video_int", feature = "audio_int", feature = "splat_int", feature = "world_int"))]
use std::collections::BTreeMap;
use std::io::Write;
use std::path::PathBuf;
use std::sync::Arc;
use serde_json::Value;
use wai::int_hyper_synth::IntModel;
use wai::int_ops::gpu::{self, Stats, wai_ops_gpu};
use wai::int_payload::{IntDecoded, IntParams, decode_int};
use wai::pin::HeldBytes;
fn wai() -> PathBuf {
PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("..")
}
fn read(rel: &str) -> Vec<u8> {
std::fs::read(wai().join(rel)).unwrap_or_else(|e| panic!("{rel}: {e}"))
}
fn json(rel: &str) -> Value {
serde_json::from_slice(&read(rel)).unwrap_or_else(|e| panic!("{rel}: {e}"))
}
fn report(line: &str) {
println!("{line}");
if let Ok(p) = std::env::var("WAI_OPS_GPU_REPORT") {
let mut f = std::fs::OpenOptions::new().create(true).append(true).open(&p).unwrap();
writeln!(f, "{line}").unwrap();
}
}
/// A decode's result as bytes: every output byte the decode gives, or its refusal code.
type Out = Result<Vec<u8>, String>;
fn bytes_of(d: &IntDecoded) -> Vec<u8> {
let mut b = Vec::new();
match d {
IntDecoded::Image { rgb, width, height } => {
b.extend_from_slice(&width.to_le_bytes());
b.extend_from_slice(&height.to_le_bytes());
b.extend_from_slice(rgb);
}
IntDecoded::Video { frames, width, height } => {
b.extend_from_slice(&width.to_le_bytes());
b.extend_from_slice(&height.to_le_bytes());
for f in frames {
b.extend_from_slice(f);
}
}
IntDecoded::Pcm { samples, sample_rate, channels } => {
b.extend_from_slice(&sample_rate.to_le_bytes());
b.extend_from_slice(&channels.to_le_bytes());
for s in samples {
b.extend_from_slice(&s.to_le_bytes());
}
}
IntDecoded::Attr { bytes, n_attr, s, n_gaussians, ranges } => {
for v in [n_attr, s, n_gaussians] {
b.extend_from_slice(&v.to_le_bytes());
}
for (lo, hi) in ranges {
b.extend_from_slice(&lo.to_bits().to_le_bytes());
b.extend_from_slice(&hi.to_bits().to_le_bytes());
}
b.extend_from_slice(bytes);
}
_ => panic!("an output form this test does not compare"),
}
b
}
/// One decode of the corpus: its capability, its name, and the decode.
struct Case {
cap: &'static str,
name: String,
refusal: bool,
run: Box<dyn Fn() -> Out>,
}
fn cap_static(c: &str) -> &'static str {
for k in [
"wai.video.int_motion", "wai.neural.int_hyper", "wai.video.int_hyper", "wai.audio.int_codec", "wai.splat.int_codec", "wai.neural.int_synth",
"wai.neural.int_mlicpp", "wai.world.model", "wai.image.jpegai",
] {
if k == c {
return k;
}
}
panic!("an unknown capability {c}")
}
fn payload_cases() -> Vec<Case> {
let idx = json("int-payload-conformance/vectors.json");
let mut out = Vec::new();
for (list, refusal) in [("vectors", false), ("refusals", true)] {
for e in idx[list].as_array().unwrap() {
let cap = cap_static(e["capability"].as_str().unwrap());
let payload = read(e["payload"].as_str().unwrap());
let set: BTreeMap<String, Vec<u8>> = e["set"].as_object().into_iter().flatten().map(|(k, v)| (k.clone(), read(v.as_str().unwrap()))).collect();
let run = move || -> Out {
let view: BTreeMap<String, &[u8]> = set.iter().map(|(k, v)| (k.clone(), v.as_slice())).collect();
let params = wai::int_payload::resolve(cap, None, &[HeldBytes::Set { listing: b"", files: &view }], None).unwrap_or_else(|_| IntParams::none());
decode_int(cap, &payload, &params).map(|d| bytes_of(&d)).map_err(|e| e.code().to_string())
};
out.push(Case { cap, name: e["name"].as_str().unwrap().to_string(), refusal, run: Box::new(run) });
}
}
out
}
fn synth_cases() -> Vec<Case> {
let hashes = json("neural-conformance/expected_hashes.json");
hashes
.as_object()
.unwrap()
.keys()
.map(|name| {
let dir = format!("neural-conformance/vectors/{name}");
let (payload, set) = (read(&format!("{dir}/latent.bin")), read(&format!("{dir}/int_synth.json")));
let run = move || -> Out { wai::int_synth::decode_from_spec(&payload, &set).map(|i| [i.width.to_le_bytes().to_vec(), i.height.to_le_bytes().to_vec(), i.rgb].concat()).map_err(|e| e.to_string()) };
Case { cap: "wai.neural.int_synth", name: name.clone(), refusal: false, run: Box::new(run) }
})
.collect()
}
fn mlicpp_cases() -> Vec<Case> {
const DIR: &str = "wai-rs/tests/vectors/int_mlicpp";
let m = wai::int_mlicpp::IntMlicModel::from_bundle(&read(&format!("{DIR}/model.wmb"))).unwrap();
let p = wai::container::Wai::from_bytes(&read(&format!("{DIR}/stream.wai"))).unwrap().payload;
let run = move || -> Out { wai::int_mlicpp::decode(&m, &p).map(|i| [(i.width as u64).to_le_bytes().to_vec(), (i.height as u64).to_le_bytes().to_vec(), i.rgb].concat()).map_err(|e| format!("{e:?}")) };
vec![Case { cap: "wai.neural.int_mlicpp", name: "int_mlicpp".into(), refusal: false, run: Box::new(run) }]
}
fn world_cases() -> Vec<Case> {
let mut out = Vec::new();
for (name, dir, js, t, g) in [("world", "int_world", "world.json", "world_T.bin", "world_gs.bin"), ("world2", "int_world2", "world2.json", "world2_T.bin", "world2_gs.bin")] {
let j = json(&format!("wai-rs/tests/vectors/{dir}/{js}"));
let tm = IntModel::from_bytes(&read(&format!("wai-rs/tests/vectors/{dir}/{t}"))).unwrap();
let gm = IntModel::from_bytes(&read(&format!("wai-rs/tests/vectors/{dir}/{g}"))).unwrap();
let c = j["c"].as_u64().unwrap() as usize;
let lat = j["lat"].as_u64().unwrap() as usize;
let z0: Vec<i64> = j["z0"].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect();
let afrac = j["afrac"].as_u64().unwrap() as u32;
let q = |d: i64| wai_ops::stage::dequantize(d, afrac, 0).unwrap();
let mv = |a: &Value| wai::int_world::ACTIONS[a.as_u64().unwrap() as usize];
let rows: Vec<Vec<i64>> = match j.get("pairs") {
Some(p) => p
.as_array()
.unwrap()
.iter()
.map(|p| {
let ((xa, ya), (xb, yb)) = (mv(&p[0]), mv(&p[1]));
vec![q(xa), q(ya), q(xb), q(yb)]
})
.collect(),
None => j["actions"]
.as_array()
.unwrap()
.iter()
.map(|a| {
let (dx, dy) = mv(a);
vec![q(dx), q(dy), 0]
})
.collect(),
};
let run = move || -> Out { wai::int_world::decode_planes(&tm, &gm, &z0, &rows, c, lat).map(|w| w.frames.concat()).map_err(|e| format!("{e:?}")) };
out.push(Case { cap: "wai.world.model", name: name.into(), refusal: false, run: Box::new(run) });
}
out
}
mod jpegai {
use wai::jpegai_decode::{DecoderWeights, NormTables, PlaneWeights};
use wai::jpegai_icci::{Conv, IcciNet, IcciPath, ResBlock};
use wai::jpegai_synthesis::{IntNet, Weights};
fn jload(p: &std::path::Path) -> serde_json::Value {
serde_json::from_slice(&std::fs::read(p).unwrap()).unwrap()
}
fn w_of(dir: &str, name: &str) -> Weights {
Weights::from_json(&jload(&std::path::Path::new(dir).join(name))).unwrap()
}
fn int_of(dir: &str, name: &str) -> IntNet {
IntNet::from_json(&jload(&std::path::Path::new(dir).join(name))).unwrap()
}
fn i64v(v: &serde_json::Value, k: &str) -> Vec<i64> {
v[k].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect()
}
fn i32v(v: &serde_json::Value, k: &str) -> Vec<i32> {
v[k].as_array().unwrap().iter().map(|x| x.as_i64().unwrap() as i32).collect()
}
fn f3(v: &serde_json::Value, k: &str) -> [f32; 3] {
let a = v[k].as_array().unwrap();
[a[0].as_f64().unwrap() as f32, a[1].as_f64().unwrap() as f32, a[2].as_f64().unwrap() as f32]
}
fn i3(v: &serde_json::Value, k: &str) -> [i64; 3] {
let a = v[k].as_array().unwrap();
[a[0].as_i64().unwrap(), a[1].as_i64().unwrap(), a[2].as_i64().unwrap()]
}
fn conv_of(c: &serde_json::Value) -> Conv {
let g = |k: &str| c[k].as_u64().unwrap() as usize;
Conv {
w: c["w"].as_array().unwrap().iter().map(|x| x.as_f64().unwrap() as f32).collect(),
b: c["b"].as_array().unwrap().iter().map(|x| x.as_f64().unwrap() as f32).collect(),
cin: g("cin"),
cout: g("cout"),
kh: g("kh"),
kw: g("kw"),
ph: g("ph"),
pw: g("pw"),
}
}
fn path_of(v: &serde_json::Value) -> IcciPath {
IcciPath {
conv_first: conv_of(&v["conv_first"]),
blocks: v["blocks"]
.as_array()
.unwrap()
.iter()
.map(|b| ResBlock { conv1: conv_of(&b["c1"]), conv2: conv_of(&b["c2"]), scale: b["scale"].as_f64().unwrap() as f32 })
.collect(),
conv_last: conv_of(&v["conv_last"]),
scale: v["scale"].as_f64().unwrap() as f32,
}
}
/// The `jpegai/v1` inputs, as `tests/jpegai_e2e_full.rs` loads them.
pub fn load(cs: &str, dir: &str, nf: &str) -> (Vec<u8>, DecoderWeights, NormTables) {
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"));
let icci = IcciNet { luma: path_of(&iw["luma"]), chroma: path_of(&iw["chroma"]), conv_last_v: conv_of(&iw["conv_last_v"]), scale_v: iw["scale_v"].as_f64().unwrap() as f32 };
let weights = DecoderWeights {
luma: PlaneWeights { hyper_scale: int_of(dir, "hs_y.json"), hyper_decoder: w_of(dir, "hyp_y.json"), synthesis: w_of(dir, "synth_y.json") },
chroma: PlaneWeights { hyper_scale: int_of(dir, "hs_uv.json"), hyper_decoder: w_of(dir, "hyp_uv.json"), synthesis: w_of(dir, "synth_uv.json") },
mcm_y: [w_of(dir, "mcm0_y.json"), w_of(dir, "mcm1_y.json"), w_of(dir, "mcm2_y.json"), w_of(dir, "mcm3_y.json")],
icci,
};
let norm = NormTables {
pmf: nv["pmf"].as_array().unwrap().iter().map(|r| r.as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect()).collect(),
bounds: i64v(&nv, "bounds"),
per_channel: nv["per_channel"].as_u64().unwrap() as usize,
cdfs_y: i64v(&nv, "cdfs_y"),
cdfs_uv: i64v(&nv, "cdfs_uv"),
gain_y: i64v(&nv, "gain_y"),
gain_uv: i64v(&nv, "gain_uv"),
lsbs_t0: i32v(&nv, "lsbs_t0"),
lsbs_t1: i32v(&nv, "lsbs_t1"),
lef_thr: i3(&nv, "lef_thr0"),
lef_mag: f3(&nv, "lef_mag0"),
};
(data, weights, norm)
}
}
fn jpegai_cases() -> (Vec<Case>, Option<String>) {
let env = |k: &str| std::env::var(k).ok().filter(|v| !v.is_empty());
let (Some(cs), Some(dir), Some(nf)) = (env("WAI_JPEGAI_CODESTREAM"), env("WAI_JPEGAI_WEIGHTS_DIR"), env("WAI_JPEGAI_NORM")) else {
return (vec![], Some("WAI_JPEGAI_CODESTREAM, WAI_JPEGAI_WEIGHTS_DIR and WAI_JPEGAI_NORM are not all set (tools/wai_jpegai_fetch_e2e_inputs.sh)".into()));
};
let (data, weights, norm) = jpegai::load(&cs, &dir, &nf);
let run = move || -> Out { wai::jpegai_decode::decode(&data, &weights, &norm).map(|i| [(i.h as u64).to_le_bytes().to_vec(), (i.w as u64).to_le_bytes().to_vec(), i.data.clone()].concat()).map_err(|r| format!("{:?}", r.code())) };
(vec![Case { cap: "wai.image.jpegai", name: "jpegai/v1 test256".into(), refusal: false, run: Box::new(run) }], None)
}
/// What one capability's cases did on one adapter.
#[derive(Default)]
struct Tally {
cases: u32,
refusals: u32,
stats: Stats,
differ: Vec<String>,
}
#[test]
fn every_capability_decodes_alike_with_its_graphs_on_every_adapter() {
let d = wai_ops_gpu::blocking::discover();
for (a, why) in &d.refused {
report(&format!("ops-gpu adapter refused: key={} {} reason={why:?}", a.key(), a.describe()));
}
if d.found.is_empty() {
report("ops-gpu capabilities: SKIPPED: no adapter");
assert!(std::env::var_os("WAI_OPS_GPU_REQUIRE").is_none(), "WAI_OPS_GPU_REQUIRE is set and no adapter was found");
return;
}
let mut cases = payload_cases();
cases.extend(synth_cases());
cases.extend(mlicpp_cases());
cases.extend(world_cases());
let (jp, missing) = jpegai_cases();
cases.extend(jp);
if let Some(why) = &missing {
report(&format!("ops-gpu capability: wai.image.jpegai NOT RUN: {why}"));
assert!(std::env::var_os("WAI_REQUIRE_SOTA_VECTORS").is_none(), "WAI_REQUIRE_SOTA_VECTORS is set and the jpegai/v1 inputs are absent");
}
// The CPU's result for every case, once.
let want: Vec<Out> = cases.iter().map(|c| (c.run)()).collect();
for (c, w) in cases.iter().zip(&want) {
assert_eq!(c.refusal, w.is_err(), "{} {}: a vector decodes and a refusal fixture is refused, on the CPU: {:?}", c.cap, c.name, w.as_ref().err());
}
for g in d.found.into_iter().map(Arc::new) {
let a = g.adapter().clone();
let mut by_cap: BTreeMap<&'static str, Tally> = BTreeMap::new();
for (c, w) in cases.iter().zip(&want) {
let (got, stats) = gpu::with_gpu_compared(&g, || (c.run)());
let t = by_cap.entry(c.cap).or_default();
t.cases += 1;
t.refusals += c.refusal as u32;
if &got != w {
t.differ.push(format!("{}: {}", c.name, match (w, &got) {
(Ok(x), Ok(y)) => format!("{} bytes vs {} bytes, first difference at {:?}", x.len(), y.len(), x.iter().zip(y).position(|(p, q)| p != q)),
_ => format!("{:?} vs {:?}", w.as_ref().map(|v| v.len()), got.as_ref().map(|v| v.len())),
}));
}
t.stats.acc32 += stats.acc32;
t.stats.acc64 += stats.acc64;
for (k, n) in stats.cpu {
*t.stats.cpu.entry(k).or_default() += n;
}
t.stats.errors.extend(stats.errors);
t.stats.compared += stats.compared;
t.stats.differed += stats.differed;
}
for (cap, t) in &by_cap {
let cpu: Vec<String> = t.stats.cpu.iter().map(|(k, n)| format!("{n}x {k}")).collect();
report(&format!(
"ops-gpu capability: key={} {} cap={cap} cases={} refusals={} graphs_gpu_acc32={} graphs_gpu_acc64={} graphs_cpu={} [{}] graphs_compared={} graphs_differ={} differ={} errors={}",
a.key(),
a.describe(),
t.cases,
t.refusals,
t.stats.acc32,
t.stats.acc64,
t.stats.on_cpu(),
cpu.join("; "),
t.stats.compared,
t.stats.differed,
t.differ.len(),
t.stats.errors.len()
));
}
for (cap, t) in &by_cap {
assert!(t.differ.is_empty(), "{} {cap}: {:?}", a.key(), t.differ);
assert!(t.stats.errors.is_empty(), "{} {cap}: device errors {:?}", a.key(), t.stats.errors);
assert_eq!(t.stats.differed, 0, "{} {cap}: a graph's outputs on the device differ from the interpreter's", a.key());
assert_eq!(t.stats.compared, t.stats.on_gpu(), "{} {cap}: every graph the device ran is compared", a.key());
if *cap != "wai.video.int_motion" {
assert!(t.stats.on_gpu() > 0, "{} {cap}: no graph ran on the device", a.key());
} else {
assert_eq!(t.stats.on_gpu() + t.stats.on_cpu(), 0, "{} {cap}: int_motion runs no graph", a.key());
}
if *cap != "wai.image.jpegai" {
assert_eq!(t.stats.on_cpu(), 0, "{} {cap}: every graph of this capability is in wai-ops/1-i32", a.key());
}
}
}
}