From 30b8f23268a615508daad4da752d7626f8f61f37 Mon Sep 17 00:00:00 2001 From: Transaction Science Date: Fri, 16 Oct 2026 10:07:26 -0400 Subject: [PATCH 1/2] wai-quantum 0.6.2: the HIP member's matrix-instruction kernel A second HIP kernel for the gate (wmma_limbs.hip, HipMember::wmma), for the gfx11 architectures: every product formed by the wavefront-of-32 16x16x16 bf16 matrix instruction, accumulated in f32, in the instruction's own operand layout; one wavefront takes sixteen pairs in seven 16-column steps. The instruction's accumulation is not exact once a negative term enters a sum. Measured on gfx1151, on the instruction alone: 1*1 + (-1)*0 returns 1 - 2^-24, 1 - 1 returns -2^-24, 128 - 64 returns 64 - 2^-17; non-negative sums are exact. The CUDA kernel's signed top limb fails the probe there. So the limbs are offset-binary (x + 2^31, all unsigned) and the kernel removes the offset in 64-bit integers: e*a = (e+2^31)(a+2^31) - 2^31(e+a) - 2^62 mod 2^64. Own implementation id (hip-wmma-limbs8/...), HipError::Unsupported for other architectures, eight mutants (nine-bit limbs, the result rows in another layout, a wrong offset correction among them). Measured on gfx1151 (silicon): hip_wmma 9/9, 66020 probe pairs, all four pins, 22 qubits, 8 of 8 mutants; hip_member 9/9 and hip_sv 3/3 unchanged. With no HIP, both suites pass over no device. CI: a step of wai-quantum-hip and its absent leg. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/wai-quantum-hip.yml | 28 +- wai/CHANGELOG.md | 28 ++ wai/tools/quantum-meter/Cargo.lock | 2 +- wai/tools/quantum-wgsl-check/Cargo.lock | 2 +- wai/wai-pack/Cargo.lock | 2 +- wai/wai-quantum-qpu/Cargo.lock | 2 +- wai/wai-quantum-serve/Cargo.lock | 2 +- wai/wai-quantum/Cargo.lock | 2 +- wai/wai-quantum/Cargo.toml | 6 +- wai/wai-quantum/README.md | 2 +- wai/wai-quantum/src/quantum_fabric/hip.rs | 140 ++++++ .../src/quantum_fabric/wmma_limbs.hip | 267 +++++++++++ wai/wai-quantum/src/quantum_fabric_hip.rs | 65 ++- wai/wai-quantum/tests/hip_wmma.rs | 444 ++++++++++++++++++ wai/wai-rs/Cargo.lock | 2 +- wai/wai-web/Cargo.lock | 2 +- 16 files changed, 972 insertions(+), 24 deletions(-) create mode 100644 wai/wai-quantum/src/quantum_fabric/wmma_limbs.hip create mode 100644 wai/wai-quantum/tests/hip_wmma.rs diff --git a/.github/workflows/wai-quantum-hip.yml b/.github/workflows/wai-quantum-hip.yml index 760a6ba86..238a669d1 100644 --- a/.github/workflows/wai-quantum-hip.yml +++ b/.github/workflows/wai-quantum-hip.yml @@ -8,9 +8,11 @@ name: wai-quantum HIP member # exactly the declared device, agree with the CPU on the arithmetic probe and # on random states (out-of-range runs included), reproduce all four pinned # certificates (two gatesets, random-v1 and fabric-v1), and catch every -# kernel mutant; and the dense f64 state vector must come back with the CPU's -# bits on its pinned corpus and at 16 and 20 qubits, every state-vector -# mutant caught. +# kernel mutant; the matrix-instruction kernel, every product formed by the +# GPU's bf16 matrix instruction from offset-binary 8-bit limbs, must do all +# of the same under its own implementation id, with its own mutants caught; +# and the dense f64 state vector must come back with the CPU's bits on its +# pinned corpus and at 16 and 20 qubits, every state-vector mutant caught. # - hip-absent (a container with no GPU and no HIP libraries): the same build # finds no member and says why, rather than failing; the device tests run over # no member with the declaration `none`. @@ -106,6 +108,23 @@ jobs: n="$(grep -c "^hip certified: key=$key .* parity=bit-exact " "$report")" [ "$n" = 4 ] || { echo "::error::$n of 4 certificates were earned"; exit 1; } grep -qx "hip mutants: key=$key killed 6 of 6" "$report" || { echo "::error::a kernel mutant survived"; exit 1; } + - name: The matrix-instruction kernel, every product by the instruction + shell: bash + working-directory: src/wai/wai-quantum + run: | + set -euo pipefail + export PATH="$HOME/.cargo/bin:$PATH" + key="${WAI_QUANTUM_AMD_HIP_EXPECT:-gpu-hip-gfx1151}" + report="$RUNNER_TEMP/hip-wmma-report.txt"; : > "$report" + WAI_QUANTUM_HIP_EXPECT="$key" WAI_QUANTUM_HIP_REPORT="$report" CARGO_TARGET_DIR="$HOME/hw-target/wai-quantum-hip" nice -n 10 \ + cargo test --locked --release --features quantum_fabric_hip --test hip_wmma -- --test-threads=1 2>&1 | tee "$RUNNER_TEMP/wmma.log" + grep -qE '^test result: ok\. 9 passed; 0 failed' "$RUNNER_TEMP/wmma.log" || { echo "::error::not every matrix-instruction test ran and passed"; exit 1; } + cat "$report" + grep -qx "hip wmma devices: as declared ($key)" "$report" || { echo "::error::the devices found are not the declared ones"; exit 1; } + grep -q "^hip wmma member: key=$key .* implementation=hip-wmma-limbs8/[0-9a-f]\{16\} " "$report" || { echo "::error::the device ran no matrix-instruction member"; exit 1; } + n="$(grep -c "^hip wmma certified: key=$key .* implementation=hip-wmma-limbs8/[0-9a-f]\{16\} parity=bit-exact " "$report")" + [ "$n" = 4 ] || { echo "::error::$n of 4 matrix-instruction certificates were earned"; exit 1; } + grep -qx "hip wmma mutants: key=$key killed 8 of 8" "$report" || { echo "::error::a matrix-instruction mutant survived"; exit 1; } - name: The dense f64 state vector, bit for bit shell: bash working-directory: src/wai/wai-quantum @@ -149,5 +168,8 @@ jobs: WAI_QUANTUM_HIP_REPORT="$report" cargo test --locked --features quantum_fabric_hip --test hip_member -- --test-threads=1 2>&1 | tee "$RUNNER_TEMP/device.log" grep -qE '^test result: ok\. 9 passed; 0 failed' "$RUNNER_TEMP/device.log" || { echo "::error::not every device test ran and passed"; exit 1; } grep -qx 'hip devices: as declared (none)' "$report" || { echo "::error::the absent case was not established"; exit 1; } + WAI_QUANTUM_HIP_REPORT="$report" cargo test --locked --features quantum_fabric_hip --test hip_wmma -- --test-threads=1 2>&1 | tee "$RUNNER_TEMP/wmma.log" + grep -qE '^test result: ok\. 9 passed; 0 failed' "$RUNNER_TEMP/wmma.log" || { echo "::error::the matrix-instruction tests did not run over no device"; exit 1; } + grep -qx 'hip wmma devices: as declared (none)' "$report" || { echo "::error::the matrix-instruction absent case was not established"; exit 1; } cargo test --locked --features quantum_fabric_hip,quantum_sv --test hip_sv -- --test-threads=1 2>&1 | tee "$RUNNER_TEMP/sv.log" grep -qE '^test result: ok\. 3 passed; 0 failed' "$RUNNER_TEMP/sv.log" || { echo "::error::the state-vector tests did not run over no device"; exit 1; } diff --git a/wai/CHANGELOG.md b/wai/CHANGELOG.md index 831ec7541..9e52b129c 100644 --- a/wai/CHANGELOG.md +++ b/wai/CHANGELOG.md @@ -6,6 +6,34 @@ wire format may change and such changes are recorded here. ## v0.4.0 — current +### wai-quantum 0.6.2: the HIP member's matrix-instruction kernel + +Additive. Every pinned digest is unchanged. + +- **Every product by the bf16 matrix instruction** (`quantum_fabric::hip::WMMA_SOURCE`, + `wmma_limbs.hip`; `HipMember::wmma`): a second HIP kernel for the same gate, for the gfx11 + architectures (`DeviceIdentity::has_wmma`), whose wavefronts of 32 multiply bf16 at + 16 × 16 × 16 and accumulate in f32. One wavefront multiplies sixteen pairs' limbs by the + gate's limb weights in seven 16-column steps, in the instruction's own operand layout. A + device without it is `HipError::Unsupported` (new). Eight mutants (`WMMA_MUTANTS`): nine-bit + limbs, the result's rows in another instruction's layout, and a wrong offset correction + among them. +- **Offset-binary limbs, because of what the instruction does.** Its f32 accumulation is not + exact once a negative term enters a sum. Measured on gfx1151, on the instruction alone: + `1·1 + (−1)·0` returns `1 − 2^-24`, `1 − 1` returns `−2^-24`, `128 − 64` returns + `64 − 2^-17`; sums of non-negative terms come back exact. The CUDA kernel's signed top limb + therefore fails the probe on this GPU, off by one in a coefficient. So each value is taken + as `x + 2^31`, all four limbs unsigned, the instruction forms `(e + 2^31)·(a + 2^31)` + exactly, and the kernel removes the offset in 64-bit integers: + `e·a = (e + 2^31)·(a + 2^31) − 2^31·(e + a) − 2^62`, modulo `2^64`. The non-integer + results the limb program gave through one compiler's HIP build for this GPU (0.6.0) are + consistent with the same behaviour, though which instruction that build chose was not + checked. +- **Measured** on the integrated GPU (gfx1151, silicon): 9/9, 66020 probe pairs, all four + pinned certificates, circuits to 22 qubits, 8 of 8 mutants killed; the other kernel and the + state vector unchanged there. CI: a step of `wai-quantum-hip`; its absent leg runs the new + tests over no device. + ### wai-quantum 0.6.1: the CUDA member's tensor-core kernel, and the Metal member Additive. Every pinned digest is unchanged. diff --git a/wai/tools/quantum-meter/Cargo.lock b/wai/tools/quantum-meter/Cargo.lock index 2f4f18e60..c68e813e5 100644 --- a/wai/tools/quantum-meter/Cargo.lock +++ b/wai/tools/quantum-meter/Cargo.lock @@ -1017,7 +1017,7 @@ checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" [[package]] name = "wai-quantum" -version = "0.6.1" +version = "0.6.2" dependencies = [ "blake3", "ed25519-dalek", diff --git a/wai/tools/quantum-wgsl-check/Cargo.lock b/wai/tools/quantum-wgsl-check/Cargo.lock index 3ca21b7b3..b87c7e6a7 100644 --- a/wai/tools/quantum-wgsl-check/Cargo.lock +++ b/wai/tools/quantum-wgsl-check/Cargo.lock @@ -599,7 +599,7 @@ checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" [[package]] name = "wai-quantum" -version = "0.6.1" +version = "0.6.2" dependencies = [ "blake3", "ed25519-dalek", diff --git a/wai/wai-pack/Cargo.lock b/wai/wai-pack/Cargo.lock index 69503345d..150bbcbbb 100644 --- a/wai/wai-pack/Cargo.lock +++ b/wai/wai-pack/Cargo.lock @@ -713,7 +713,7 @@ dependencies = [ [[package]] name = "wai-quantum" -version = "0.6.1" +version = "0.6.2" dependencies = [ "blake3", "ed25519-dalek", diff --git a/wai/wai-quantum-qpu/Cargo.lock b/wai/wai-quantum-qpu/Cargo.lock index f9ac5d194..b701fef57 100644 --- a/wai/wai-quantum-qpu/Cargo.lock +++ b/wai/wai-quantum-qpu/Cargo.lock @@ -422,7 +422,7 @@ checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" [[package]] name = "wai-quantum" -version = "0.6.1" +version = "0.6.2" dependencies = [ "blake3", "ed25519-dalek", diff --git a/wai/wai-quantum-serve/Cargo.lock b/wai/wai-quantum-serve/Cargo.lock index 817232a82..a44393f33 100644 --- a/wai/wai-quantum-serve/Cargo.lock +++ b/wai/wai-quantum-serve/Cargo.lock @@ -599,7 +599,7 @@ dependencies = [ [[package]] name = "wai-quantum" -version = "0.6.1" +version = "0.6.2" dependencies = [ "blake3", "ed25519-dalek", diff --git a/wai/wai-quantum/Cargo.lock b/wai/wai-quantum/Cargo.lock index e5fb779b2..d8148e1bc 100644 --- a/wai/wai-quantum/Cargo.lock +++ b/wai/wai-quantum/Cargo.lock @@ -1079,7 +1079,7 @@ checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" [[package]] name = "wai-quantum" -version = "0.6.1" +version = "0.6.2" dependencies = [ "blake3", "block2", diff --git a/wai/wai-quantum/Cargo.toml b/wai/wai-quantum/Cargo.toml index b60b10869..c25a09607 100644 --- a/wai/wai-quantum/Cargo.toml +++ b/wai/wai-quantum/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "wai-quantum" -version = "0.6.1" +version = "0.6.2" edition = "2024" rust-version = "1.97" license = "Apache-2.0" @@ -230,6 +230,10 @@ required-features = ["quantum_fabric_metal"] name = "cuda_sv" required-features = ["quantum_fabric_cuda", "quantum_sv"] +[[test]] +name = "hip_wmma" +required-features = ["quantum_fabric_hip"] + [[test]] name = "hip_member" required-features = ["quantum_fabric_hip"] diff --git a/wai/wai-quantum/README.md b/wai/wai-quantum/README.md index a0ade27d6..d3263b076 100644 --- a/wai/wai-quantum/README.md +++ b/wai/wai-quantum/README.md @@ -107,7 +107,7 @@ never creates information. | `quantum_fabric_webgpu` | **the fabric's WebGPU member**, opt-in and not in `full`: the 32-bit-lane gate kernel in WGSL, one dispatch per gate, through Metal, Vulkan or DX12 natively and through a browser's WebGPU on `wasm32-unknown-unknown`; a second implementation applies runs of gates tile by tile, one dispatch per run, and is certified on its own. Each adapter is probed against the CPU's arithmetic in the process that uses it and certified bit-exact by reproducing the CPU reference's pinned corpus digests; a result that would leave 32 bits is refused, never returned. One of the features that reach outside the crate | | `quantum_fabric_cuda` | **the fabric's CUDA member**, opt-in and not in `full`, native targets only: the gate kernel in CUDA C with exact 64-bit integers, compiled at run time for each device's own architecture, one launch per gate. The CUDA driver and runtime compiler are loaded when first used, so a build with the feature runs on a host without CUDA and finds no member there, saying why. Each device is probed against the CPU's arithmetic in the process that uses it and certified bit-exact by reproducing the CPU reference's pinned corpus digests; a result that would leave 32 bits is refused, never returned. On a device of compute capability 8.0 or later, a **second kernel runs every product on the tensor cores** (`CudaMember::tensor`): each value as four 8-bit limbs in bf16, accumulated in f32, which holds every partial sum exactly; probed and certified under its own implementation id, with its own mutants. With `quantum_sv` too, the **dense f64 state vector** runs on the device bit for bit as on the CPU: the CPU's own launches (diagonal, one-qubit and grouped formulas), each an IEEE binary64 operation with no fused multiply-add, certified against a pinned corpus digest | | `quantum_fabric_metal` | **the fabric's Metal member**, opt-in and not in `full`, macOS: the gate kernel in the Metal shading language with exact 64-bit integers, compiled at run time by the operating system's own compiler, one serial dispatch per gate in one command buffer, the state in memory the GPU and the host share. Elsewhere the build finds no member and says why. Each device is probed against the CPU's arithmetic in the process that uses it and certified bit-exact by reproducing the CPU reference's pinned corpus digests, keyed by its GPU family number (`gpu-msl-`). On a GPU with SIMD-group matrices, a **second kernel forms every product with the matrix multiply** (`MetalMember::matrix`): each value as four 8-bit limbs in bf16, accumulated in f32, which holds every partial sum exactly; probed and certified under its own implementation id, with its own mutants | -| `quantum_fabric_hip` | **the fabric's HIP member**, opt-in and not in `full`, Linux: the CUDA member's kernel text, unchanged, compiled at run time by HIP's runtime compiler for each device's own architecture, one launch per gate. The runtime and its compiler are loaded when first used, so a build with the feature runs on a host without them and finds no member there, saying why; the device's properties are read in the runtime's frozen layout or the device is refused. Probed and certified as the CUDA member is, under its own implementation id, and with `quantum_sv` the dense f64 state vector runs bit for bit as on the CPU | +| `quantum_fabric_hip` | **the fabric's HIP member**, opt-in and not in `full`, Linux: the CUDA member's kernel text, unchanged, compiled at run time by HIP's runtime compiler for each device's own architecture, one launch per gate. The runtime and its compiler are loaded when first used, so a build with the feature runs on a host without them and finds no member there, saying why; the device's properties are read in the runtime's frozen layout or the device is refused. Probed and certified as the CUDA member is, under its own implementation id, and with `quantum_sv` the dense f64 state vector runs bit for bit as on the CPU. On a gfx11 GPU, a **second kernel forms every product with the bf16 matrix instruction** (`HipMember::wmma`): each value offset-binary, as four unsigned 8-bit limbs, because the instruction's accumulation is not exact once a negative term enters a sum; the offset is removed in 64-bit integers. Probed and certified under its own implementation id, with its own mutants | | `quantum_fabric_opencl` | **the fabric's OpenCL member**, opt-in and not in `full`, native targets only: the same exact 64-bit kernel in OpenCL C, built at run time by each device's own compiler, one launch per gate, on any OpenCL GPU or CPU (an OpenCL accelerator is refused: no fabric class names it). The ICD loader is opened when first used, so a build with the feature runs on a host without OpenCL and finds no member there, saying why. Each device is probed and certified as the CUDA member's are, and keyed by its vendor id (`opencl--`). With `quantum_sv` too, the dense f64 state vector runs bit for bit as on the CPU on every device with binary64 (contraction off); a device without it refuses, saying so | | `quantum_fabric_fpga` | **the fabric's FPGA member, simulated**, opt-in and not in `full`, native targets only: the gate datapath in Verilog (each 32×32 product from four 16-bit partial products, the size of an FPGA's multiplier blocks; a state RAM; one pair per two cycles), compiled at run time by Verilator into a cycle-accurate simulation driven over a pipe. Probed and certified as the other members are, and every certificate says it was earned in a simulator, so placement ranks it after silicon. A host without Verilator finds no member and says why | | `quantum_fabric_npu` | **the fabric's neural-processor member**, opt-in and not in `full`, macOS only: each 32-bit operand split into eight 4-bit limbs, every limb product and limb sum formed in fp16 on the platform's neural processor (each value an integer fp16 holds exactly, whatever order it is added in), the host assembling each product and applying the fixed-point arithmetic. The member exists only where the platform model runtime's compute plan places every operation of the model on the neural processor; a model it would run elsewhere is refused with the placement it reported. Probed and certified as every member is | diff --git a/wai/wai-quantum/src/quantum_fabric/hip.rs b/wai/wai-quantum/src/quantum_fabric/hip.rs index d3281a3da..74ac133ef 100644 --- a/wai/wai-quantum/src/quantum_fabric/hip.rs +++ b/wai/wai-quantum/src/quantum_fabric/hip.rs @@ -8,8 +8,31 @@ //! compiler, its options and the device, and the implementation id and the //! member id bind them. No device is touched here, so all of it is testable on //! any host; the HIP member itself is `quantum_fabric_hip`. +//! +//! **The matrix-instruction kernel** ([`WMMA_SOURCE`], `wmma_limbs.hip`) is a +//! second implementation of the same gate, for the architectures whose +//! wavefronts of 32 multiply bf16 at 16 × 16 × 16 and accumulate in f32 +//! ([`DeviceIdentity::has_wmma`]): every product is formed by that +//! instruction, from four 8-bit limbs a value, by the CUDA tensor-core +//! kernel's argument ([`cuda`](super::cuda)) in the instruction's own operand +//! layout. One wavefront multiplies sixteen pairs' limbs by the gate's limb +//! weights (16 × 112) in seven 16-column steps. +//! +//! One thing differs from that kernel, because the instruction's f32 +//! accumulation is **not exact once a negative term enters a sum**: measured +//! on gfx1151, `1·1 + (−1)·0` returns `1 − 2^-24` and `128 − 64` returns +//! `64 − 2^-17`, while sums of non-negative terms come back exact. So no +//! operand is negative: each value `x` is taken offset-binary, `x + 2^31`, all +//! four limbs unsigned, the instruction forms `(e + 2^31)·(a + 2^31)` exactly, +//! and the kernel removes the offset's terms in 64-bit integers: +//! `e·a = (e + 2^31)·(a + 2^31) − 2^31·(e + a) − 2^62`, modulo `2^64`. +//! +//! It has its own implementation id ([`wmma_implementation_id`]) and its own +//! mutants ([`WMMA_MUTANTS`]): one is nine-bit limbs, which bf16 cannot hold, +//! and one gets the offset's correction wrong. use super::FabricClass; +use super::wgsl::{Detector, Mutant}; pub use super::cuda::{MUTANTS, mutate}; /// The kernel, exactly as it is compiled: the CUDA kernel's text. @@ -143,6 +166,85 @@ impl DeviceIdentity { } } +/// The matrix-instruction kernel, `wmma_limbs.hip`, exactly as it is +/// compiled. +pub const WMMA_SOURCE: &str = include_str!("wmma_limbs.hip"); + +/// The matrix-instruction kernel's implementation name. +pub const WMMA_IMPLEMENTATION: &str = "hip-wmma-limbs8"; + +/// Version of the matrix-instruction kernel's host-side layout: +/// [`HOST_LAYOUT`]'s, except that a wavefront takes sixteen pairs (or probe +/// cases), so a launch has [`WMMA_THREADS_PER_ITEM`] threads per item. Part of +/// its implementation id. +pub const WMMA_HOST_LAYOUT: u32 = 1; + +/// Threads per pair (or probe case) of a matrix-instruction launch: 32 lanes +/// take 16. +pub const WMMA_THREADS_PER_ITEM: u64 = 2; + +const DOMAIN_WMMA_IMPLEMENTATION: &[u8] = b"wai:quantum-fabric-hip-wmma\x01"; + +/// `"hip-wmma-limbs8/<16 hex>"`: as [`implementation_id`], with the domain +/// `"wai:quantum-fabric-hip-wmma\x01"`, [`WMMA_SOURCE`] and +/// [`WMMA_HOST_LAYOUT`]. +pub fn wmma_implementation_id() -> String { + wmma_implementation_id_of(WMMA_SOURCE) +} + +/// [`wmma_implementation_id`] for another kernel text, such as a mutant's. +pub fn wmma_implementation_id_of(source: &str) -> String { + let mut o = DOMAIN_WMMA_IMPLEMENTATION.to_vec(); + put(&mut o, source.as_bytes()); + for option in COMPILE_OPTIONS { + put(&mut o, option.as_bytes()); + } + o.extend_from_slice(&WMMA_HOST_LAYOUT.to_be_bytes()); + format!("{WMMA_IMPLEMENTATION}/{}", hex(&blake3::hash(&o).as_bytes()[..8])) +} + +impl DeviceIdentity { + /// Whether the device runs the matrix-instruction kernel: a gfx11 + /// architecture, whose wavefronts of 32 have the 16 × 16 × 16 bf16 + /// instruction in the operand layout the kernel is written for. Other + /// architectures' matrix instructions differ in name or layout. + pub fn has_wmma(&self) -> bool { + self.arch_name().strip_prefix("gfx11").is_some_and(|rest| rest.len() == 2 && rest.bytes().all(|b| b.is_ascii_hexdigit())) + } +} + +/// The matrix-instruction kernel's mutants, as [`MUTANTS`] are the other's. +/// Each `find` occurs exactly once in [`WMMA_SOURCE`]. +pub const WMMA_MUTANTS: [Mutant; 8] = [ + // The rounding constant is half what it should be. + Mutant { name: "W1", find: "return (v + ((i64)1 << 29)) >> 30;", replace: "return (v + ((i64)1 << 28)) >> 30;", detector: Detector::Probe }, + // The product is shifted by one bit too few. + Mutant { name: "W2", find: "return (v + ((i64)1 << 29)) >> 30;", replace: "return (v + ((i64)1 << 29)) >> 29;", detector: Detector::Probe }, + // Nine-bit limbs: a limb above 2^8 has nine significant bits, which bf16 + // does not hold, so the products are not exact. + Mutant { name: "W3", find: "#define LIMB 8", replace: "#define LIMB 9", detector: Detector::Probe }, + // The result's rows are read in another instruction's layout (rows + // i and i + 8, not 2i and 2i + 1). + Mutant { name: "W4", find: "int row = 2 * i + half;", replace: "int row = i + 8 * half;", detector: Detector::Probe }, + // The control test is inverted. + Mutant { name: "W5", find: "if ((i & (u64)ctrl) == (u64)ctrl) {", replace: "if ((i & (u64)ctrl) != (u64)ctrl) {", detector: Detector::Certificate }, + // A sign of row 0's real part is wrong, on the imaginary part of the + // top-right entry (as C5). + Mutant { name: "W6", find: "fx(pr[2]) - fx(pr[3]);", replace: "fx(pr[2]) + fx(pr[3]);", detector: Detector::Certificate }, + // The range guard never fires. + Mutant { name: "W7", find: "if (!(fits32(o0r) && fits32(o0i) && fits32(o1r) && fits32(o1i))) {", replace: "if (false) {", detector: Detector::Range }, + // The offset's square is removed at half its value. + Mutant { name: "W8", find: "- ((u64)1 << 62));", replace: "- ((u64)1 << 61));", detector: Detector::Probe }, +]; + +/// [`WMMA_SOURCE`] with `m` applied. Panics unless `m.find` occurs exactly +/// once. +pub fn wmma_mutate(m: &Mutant) -> String { + let n = WMMA_SOURCE.matches(m.find).count(); + assert_eq!(n, 1, "mutant {}: `{}` occurs {n} times in the matrix-instruction kernel", m.name, m.find); + WMMA_SOURCE.replacen(m.find, m.replace, 1) +} + #[cfg(test)] mod tests { use super::*; @@ -199,4 +301,42 @@ mod tests { assert_ne!(implementation_id_of(&mutate(m)), id, "mutant {}", m.name); } } + + #[test] + fn the_wmma_kernel_has_its_own_id_and_mutants() { + let id = wmma_implementation_id(); + assert!(id.starts_with("hip-wmma-limbs8/") && id.len() == "hip-wmma-limbs8/".len() + 16, "{id}"); + assert_ne!(wmma_implementation_id_of(SOURCE).split('/').nth(1), implementation_id_of(SOURCE).split('/').nth(1), "the domains differ"); + let mut seen = std::collections::BTreeSet::new(); + for m in &WMMA_MUTANTS { + let edited = wmma_mutate(m); + assert_ne!(edited, WMMA_SOURCE, "mutant {} changes nothing", m.name); + assert_ne!(wmma_implementation_id_of(&edited), id, "mutant {}", m.name); + assert!(seen.insert(edited), "mutant {} duplicates another", m.name); + } + for d in [Detector::Probe, Detector::Certificate, Detector::Range] { + assert!(WMMA_MUTANTS.iter().any(|m| m.detector == d), "no matrix-instruction mutant for {d:?}"); + } + for entry in ["extern \"C\" __global__ void apply_gate(", "extern \"C\" __global__ void probe("] { + assert_eq!((SOURCE.matches(entry).count(), WMMA_SOURCE.matches(entry).count()), (1, 1), "{entry}"); + } + } + + #[test] + fn the_wmma_kernel_is_for_gfx11_architectures() { + for (arch, has) in [ + ("gfx1100", true), + ("gfx1103", true), + ("gfx1151", true), + ("gfx1151:xnack-", true), + ("gfx1030", false), + ("gfx1200", false), + ("gfx90a:sramecc+:xnack-", false), + ("gfx942", false), + ("gfx11", false), + ("gfx110", false), + ] { + assert_eq!(DeviceIdentity { arch: arch.into(), ..identity() }.has_wmma(), has, "{arch}"); + } + } } diff --git a/wai/wai-quantum/src/quantum_fabric/wmma_limbs.hip b/wai/wai-quantum/src/quantum_fabric/wmma_limbs.hip new file mode 100644 index 000000000..d4cd11944 --- /dev/null +++ b/wai/wai-quantum/src/quantum_fabric/wmma_limbs.hip @@ -0,0 +1,267 @@ +// wai-quantum fabric: the gate kernel for HIP devices whose matrix +// instructions multiply bf16 and accumulate in f32 at 16 × 16 × 16, a +// wavefront of 32 lanes at a time (WMMA, the gfx11 architectures). +// +// The same entry points as lanes64.cu, with the same parameters, so the +// member runs either one; here every product the gate needs is formed by the +// matrix instruction. +// +// The instruction's f32 accumulation is not exact once a negative term enters +// a sum: measured on gfx1151, 1·1 + (−1)·0 returns 1 − 2^-24 and 128 − 64 +// returns 64 − 2^-17, while sums of non-negative terms come back exact. So no +// operand here is negative. Each 32-bit value x is taken offset-binary, +// u = x + 2^31 (its bits with the top one flipped), and split into four +// unsigned limbs of 8 bits, which bf16 holds exactly. The product u_e·u_a is +// Σ c_s·2^(8s) with c_s = Σ_{i+j=s} e_i·a_j: seven coefficients of at most +// four limb products each, every limb product at most 255², every +// coefficient below 2^18 and never negative. So every value the instruction +// forms is an integer f32 holds exactly, in whatever order it adds them +// (quantum_fabric::tensor_engine has the argument and the layout). The kernel +// adds the coefficients, shifted into place, in 64-bit integers, which gives +// u_e·u_a, and removes the offset exactly: e·a = u_e·u_a − 2^31·(e + a) − +// 2^62, modulo 2^64, and e·a fits 64 bits. The limb weights put each of a +// gate's sixteen products per pair in its own seven columns of a 16 × 112 +// matrix; one wavefront multiplies a 16-pair tile of limbs by it in seven +// 16-column steps. The rest is lanes64.cu's arithmetic on those exact +// products: fx = (product + 2^29) >> 30, each row the sum of four, the range +// reported and the low 32 bits kept. +// +// The instruction's operands, lane by lane: lane l holds row l mod 16 of A +// and column l mod 16 of B, all sixteen of K each (lanes 16–31 repeat lanes +// 0–15); element i of its result is row 2i + l / 16, column l mod 16. +// +// Each thread owns its pair's two amplitudes for the launch, and a launch is +// one gate, so no thread reads what another writes within a launch. The block +// steps through its tiles together, so every barrier is reached by every +// thread. + +typedef long long i64; +typedef unsigned long long u64; +typedef unsigned int u32; +typedef unsigned short u16; + +// The instruction's operand and result types: sixteen bf16 values as their +// bits, and eight f32 values. +typedef short v16 __attribute__((ext_vector_type(16))); +typedef float v8f __attribute__((ext_vector_type(8))); + +// Wavefronts per block: the member launches 256 threads a block, in +// wavefronts of 32. +#define WAVES 8 + +// Bits per limb. Eight is what bf16 holds exactly; the radix is a mutant's +// one edit, and nine-bit limbs are not exact. +#define LIMB 8 + +// The offset that makes every operand non-negative: x + 2^31. +#define OFFSET 0x80000000u + +// The sixteen products, as (entry, component): rows 0r, 0i, 1r and 1i, four +// terms each, in the order lanes::exact_row adds them. +__constant__ int TERM_ENTRY[16] = {0, 1, 2, 3, 0, 1, 2, 3, 4, 5, 6, 7, 4, 5, 6, 7}; +__constant__ int TERM_COMP[16] = {0, 1, 2, 3, 1, 0, 3, 2, 0, 1, 2, 3, 1, 0, 3, 2}; + +// Limb j of x + 2^31: four unsigned limbs. +__device__ __forceinline__ int limb(int x, int j) { + return (int)((((u32)x ^ OFFSET) >> (LIMB * j)) & ((1u << LIMB) - 1)); +} + +// e·a from the instruction's (e + 2^31)·(a + 2^31): the offset's terms +// removed, modulo 2^64. +__device__ __forceinline__ i64 unoffset(i64 raw, int e, int a) { + return (i64)((u64)raw - (((u64)(i64)e + (u64)(i64)a) << 31) - ((u64)1 << 62)); +} + +// v as bf16 bits. An integer of at most 256 in size has at most 8 significant +// bits, so its f32 has zero low 16 bits and the truncation is exact. +__device__ __forceinline__ u16 bf16(int v) { + return (u16)(__float_as_uint((float)v) >> 16); +} + +// The exact fixed-point rounding of a product. +__device__ __forceinline__ i64 fx(i64 v) { + return (v + ((i64)1 << 29)) >> 30; +} + +__device__ __forceinline__ int low32(i64 v) { + return (int)(u32)(u64)v; +} + +__device__ __forceinline__ bool fits32(i64 v) { + return v >= -(i64)2147483648LL && v <= (i64)2147483647LL; +} + +// P[row][p] += product p of row `row`, for the first `rows` rows of the tile: +// A (16 × 16 limbs) times W (16 × 112), coefficient s of product p at column +// p·7 + s, each shifted into place. Called by all 32 lanes of a wavefront. +__device__ void products(const u16 (*A)[16], const u16 (*W)[112], i64 (*P)[16], int rows, int lane) { + int r = lane & 15, half = lane >> 4; + v16 a; + for (int k = 0; k < 16; k++) { + a[k] = (short)A[r][k]; + } + for (int nt = 0; nt < 7; nt++) { + int col = nt * 16 + r; + v16 b; + for (int k = 0; k < 16; k++) { + b[k] = (short)W[k][col]; + } + v8f c = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; + v8f d = __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32(a, b, c); + int p = col / 7, s = col % 7; + for (int i = 0; i < 8; i++) { + int row = 2 * i + half; + if (row >= rows) { + continue; + } + // An integer below 2^24 in size, so the conversion is exact; the + // shift and the add are in two's complement. + u64 v = (u64)(i64)d[i] << (LIMB * s); + atomicAdd((u64*)&P[row][p], v); + } + } +} + +// One gate on a state of `2 * pairs` lanes, one pair per lane of a +// wavefront's first half, sixteen pairs a tile. The parameters are +// lanes64.cu's. +extern "C" __global__ void apply_gate( + int2* __restrict__ amps, + u64 pairs, + unsigned int tbit, + unsigned int ctrl, + int m0, int m1, int m2, int m3, int m4, int m5, int m6, int m7, + unsigned int index, + unsigned int* __restrict__ report) +{ + __shared__ u16 W[16][112]; + __shared__ u16 A[WAVES][16][16]; + __shared__ i64 P[WAVES][16][16]; + if (blockDim.x != WAVES * 32 || warpSize != 32) { + return; + } + int m[8] = {m0, m1, m2, m3, m4, m5, m6, m7}; + // The gate's limb weights, once a block: row c·4 + j is limb j of + // component c; column p·7 + s takes limb s − j of product p's entry. + for (int x = threadIdx.x; x < 16 * 112; x += blockDim.x) { + int k = x / 112, col = x % 112; + int c = k / 4, j = k % 4, p = col / 7, i = col % 7 - j; + W[k][col] = (c == TERM_COMP[p] && i >= 0 && i <= 3) ? bf16(limb(m[TERM_ENTRY[p]], i)) : (u16)0; + } + __syncthreads(); + int w = threadIdx.x >> 5, lane = threadIdx.x & 31; + u64 tiles = (pairs + 15) / 16; + u64 stride = (u64)gridDim.x * WAVES; + for (u64 base = (u64)blockIdx.x * WAVES; base < tiles; base += stride) { + u64 tile = base + w; + bool live = tile < tiles; + bool active = false; + u64 i = 0, j = 0; + int2 a0 = make_int2(0, 0), a1 = make_int2(0, 0); + if (live && lane < 16) { + u64 p = tile * 16 + lane; + if (p < pairs) { + // The p-th index whose target bit is clear. + u64 low = p & ((u64)tbit - 1); + i = ((p - low) << 1) | low; + if ((i & (u64)ctrl) == (u64)ctrl) { + active = true; + j = i | (u64)tbit; + a0 = amps[i]; + a1 = amps[j]; + } + } + int comp[4] = {a0.x, a0.y, a1.x, a1.y}; + for (int c = 0; c < 4; c++) { + for (int jj = 0; jj < 4; jj++) { + A[w][lane][c * 4 + jj] = bf16(limb(comp[c], jj)); + } + } + for (int q = 0; q < 16; q++) { + P[w][lane][q] = 0; + } + } + __syncthreads(); + if (live) { + products(A[w], W, P[w], 16, lane); + } + __syncthreads(); + if (active) { + int comp[4] = {a0.x, a0.y, a1.x, a1.y}; + i64 pr[16]; + for (int q = 0; q < 16; q++) { + pr[q] = unoffset(P[w][lane][q], m[TERM_ENTRY[q]], comp[TERM_COMP[q]]); + } + i64 o0r = fx(pr[0]) - fx(pr[1]) + fx(pr[2]) - fx(pr[3]); + i64 o0i = fx(pr[4]) + fx(pr[5]) + fx(pr[6]) + fx(pr[7]); + i64 o1r = fx(pr[8]) - fx(pr[9]) + fx(pr[10]) - fx(pr[11]); + i64 o1i = fx(pr[12]) + fx(pr[13]) + fx(pr[14]) + fx(pr[15]); + if (!(fits32(o0r) && fits32(o0i) && fits32(o1r) && fits32(o1i))) { + atomicMin(&report[0], index); + } + amps[i] = make_int2(low32(o0r), low32(o0i)); + amps[j] = make_int2(low32(o1r), low32(o1i)); + } + __syncthreads(); + } +} + +// The arithmetic probe: for each case (a, b), the 64-bit product a·b from the +// matrix instruction and its fx, as (high, low) words. Sixteen cases a tile: +// product q takes b_q's limbs against component q % 4, and row q / 4 carries +// a_{4r}, …, a_{4r+3} (quantum_fabric::tensor_engine::probe_weights). +extern "C" __global__ void probe( + const int2* __restrict__ cases, + uint4* __restrict__ out, + u64 count) +{ + __shared__ u16 W[WAVES][16][112]; + __shared__ u16 A[WAVES][16][16]; + __shared__ i64 P[WAVES][4][16]; + if (blockDim.x != WAVES * 32 || warpSize != 32) { + return; + } + int w = threadIdx.x >> 5, lane = threadIdx.x & 31; + u64 tiles = (count + 15) / 16; + u64 stride = (u64)gridDim.x * WAVES; + for (u64 base = (u64)blockIdx.x * WAVES; base < tiles; base += stride) { + u64 tile = base + w; + bool live = tile < tiles; + u64 k0 = tile * 16; + if (live) { + for (int x = lane; x < 16 * 112; x += 32) { + int k = x / 112, col = x % 112; + int c = k / 4, j = k % 4, p = col / 7, i = col % 7 - j; + int b = k0 + p < count ? cases[k0 + p].y : 0; + W[w][k][col] = (c == p % 4 && i >= 0 && i <= 3) ? bf16(limb(b, i)) : (u16)0; + } + if (lane < 16) { + int r = lane / 4, c = lane % 4; + int a = k0 + lane < count ? cases[k0 + lane].x : 0; + for (int jj = 0; jj < 4; jj++) { + A[w][r][c * 4 + jj] = bf16(limb(a, jj)); + } + for (int rr = 4; rr < 16; rr++) { + A[w][rr][lane] = 0; + } + P[w][r][lane] = 0; + P[w][r][(lane + 4) % 16] = 0; + P[w][r][(lane + 8) % 16] = 0; + P[w][r][(lane + 12) % 16] = 0; + } + } + __syncthreads(); + if (live) { + products(A[w], W[w], P[w], 4, lane); + } + __syncthreads(); + if (live && lane < 16 && k0 + lane < count) { + int2 ab = cases[k0 + lane]; + i64 v = unoffset(P[w][lane / 4][lane], ab.y, ab.x); + u64 prod = (u64)v; + u64 f = (u64)fx(v); + out[k0 + lane] = make_uint4((u32)(prod >> 32), (u32)prod, (u32)(f >> 32), (u32)f); + } + __syncthreads(); + } +} diff --git a/wai/wai-quantum/src/quantum_fabric_hip.rs b/wai/wai-quantum/src/quantum_fabric_hip.rs index 660514195..d7862f135 100644 --- a/wai/wai-quantum/src/quantum_fabric_hip.rs +++ b/wai/wai-quantum/src/quantum_fabric_hip.rs @@ -17,6 +17,11 @@ //! one launch per gate, in op-log order; then the report word and the state //! come back. The kernel reports exactly when the reference's result leaves //! 32 bits; the member then returns [`HipError::OutOfRange`], never a state. +//! - **The matrix-instruction kernel** ([`HipMember::wmma`]): the same device +//! running [`hip::WMMA_SOURCE`], whose every product is formed by the GPU's +//! 16 × 16 × 16 bf16 matrix instruction from 8-bit limbs; probed before it is +//! returned, with its own implementation id. A device without the +//! instruction ([`DeviceIdentity::has_wmma`]) is [`HipError::Unsupported`]. //! - **Certification**, **placement**, **receipts** and **metering** are as //! for the CUDA member. //! @@ -80,6 +85,8 @@ pub enum HipError { Internal(&'static str), /// The device's arithmetic disagreed with the CPU's on this pair. Probe { a: i32, b: i32 }, + /// The device lacks what the kernel needs, and what. + Unsupported(String), } impl std::fmt::Display for HipError { @@ -94,6 +101,7 @@ impl std::fmt::Display for HipError { HipError::OutOfRange { gate } => write!(f, "the result of op {gate} leaves 32 bits"), HipError::Internal(why) => write!(f, "host defect: {why}"), HipError::Probe { a, b } => write!(f, "the device's arithmetic differs from the CPU's at ({a}, {b})"), + HipError::Unsupported(why) => write!(f, "the device cannot run this kernel: {why}"), } } } @@ -559,18 +567,13 @@ fn open(a: &'static Api, ordinal: c_int) -> Result { id: identity.member_id(), key: key.clone(), implementation: hip::implementation_id(), + spread: 1, + discovered: detail.clone(), detail, max_qubits, identity: identity.clone(), }; - let pairs = lanes::probe_pairs(4096, 0x5eed); - match member.probe(&pairs) { - Ok(out) => match lanes::first_mismatch(&pairs, &out) { - None => Ok(member), - Some((a, b)) => Err(refuse(HipError::Probe { a, b })), - }, - Err(e) => Err(refuse(e)), - } + member.probed().map_err(refuse) } /// One HIP device as a fabric member: the kernel compiled for it and loaded @@ -584,6 +587,12 @@ pub struct HipMember { id: String, key: String, implementation: String, + /// Threads a launch gives each pair or probe case: 1, or + /// [`hip::WMMA_THREADS_PER_ITEM`] for the matrix-instruction kernel. + spread: u64, + /// What discovery found, which [`Self::detail`] extends for a kernel that + /// says more. + discovered: String, detail: String, max_qubits: u8, } @@ -649,7 +658,7 @@ impl HipMember { let out = DeviceBuffer::upload(self.api, &vec![0u32; pairs.len() * 4])?; let count = pairs.len() as u64; let mut args = [arg(&cases.ptr), arg(&out.ptr), arg(&count)]; - launch(self.api, self.kernel.functions[1], count, &mut args)?; + launch(self.api, self.kernel.functions[1], count.saturating_mul(self.spread), &mut args)?; let words: Vec = out.download(self.api)?; Ok(words.as_chunks::<4>().0.iter().map(|w| ((w[0], w[1]), (w[2], w[3]))).collect()) } @@ -753,7 +762,7 @@ impl HipMember { ]; // The words were checked against the state above, and `amps` // holds `2·pairs` int2 lanes and `report` one word. - launch(self.api, self.kernel.functions[0], pairs, &mut args)?; + launch(self.api, self.kernel.functions[0], pairs.saturating_mul(self.spread), &mut args)?; } let bad: Vec = report.download(self.api)?; let words: Vec = amps.download(self.api)?; @@ -859,7 +868,41 @@ impl HipMember { #[doc(hidden)] pub fn with_source(&self, source: &str) -> Result { let kernel = compile(self.api, self.ordinal, &self.identity, source, &hip::COMPILE_OPTIONS, &["apply_gate", "probe"])?; - Ok(HipMember { kernel: Arc::new(kernel), implementation: hip::implementation_id_of(source), ..self.clone() }) + Ok(HipMember { kernel: Arc::new(kernel), implementation: hip::implementation_id_of(source), spread: 1, detail: self.discovered.clone(), ..self.clone() }) + } + + /// The same device running the matrix-instruction kernel + /// ([`hip::WMMA_SOURCE`]), compiled for it and probed in this call: the + /// member id is the device's, the implementation id the kernel's. A device + /// without the instruction ([`DeviceIdentity::has_wmma`]) is + /// [`HipError::Unsupported`]; one whose products disagree with the CPU's is + /// [`HipError::Probe`]. + pub fn wmma(&self) -> Result { + self.with_wmma_source(hip::WMMA_SOURCE)?.probed() + } + + /// The same device running `source` as the matrix-instruction kernel, + /// unprobed, under the implementation id that text earns: for the mutation + /// tests. + #[doc(hidden)] + pub fn with_wmma_source(&self, source: &str) -> Result { + if !self.identity.has_wmma() { + return Err(HipError::Unsupported(format!("{}: the matrix-instruction kernel is written for the gfx11 architectures' wavefront-of-32 bf16 instruction", self.identity.arch))); + } + let kernel = compile(self.api, self.ordinal, &self.identity, source, &hip::COMPILE_OPTIONS, &["apply_gate", "probe"])?; + let detail = format!("{}; every product by the matrix instruction (bf16 limbs of 8 bits, f32 accumulation)", self.discovered); + Ok(HipMember { kernel: Arc::new(kernel), implementation: hip::wmma_implementation_id_of(source), spread: hip::WMMA_THREADS_PER_ITEM, detail, ..self.clone() }) + } + + /// This member, once its kernel's arithmetic agrees with the CPU's on the + /// probe's pairs. + fn probed(self) -> Result { + let pairs = lanes::probe_pairs(4096, 0x5eed); + let out = self.probe(&pairs)?; + match lanes::first_mismatch(&pairs, &out) { + None => Ok(self), + Some((a, b)) => Err(HipError::Probe { a, b }), + } } } diff --git a/wai/wai-quantum/tests/hip_wmma.rs b/wai/wai-quantum/tests/hip_wmma.rs new file mode 100644 index 000000000..f7f96402f --- /dev/null +++ b/wai/wai-quantum/tests/hip_wmma.rs @@ -0,0 +1,444 @@ +//! The HIP member's matrix-instruction kernel on the devices of the machine +//! that runs this: every product the gate needs formed by the GPU's +//! 16 × 16 × 16 bf16 matrix instruction, from 8-bit limbs accumulated in f32, +//! held to the same tests as the other kernel (`hip_member`). +//! +//! The machine declares its devices in `WAI_QUANTUM_HIP_EXPECT`, as for +//! `hip_member`: their keys, comma-separated and sorted, or `none`. Every +//! declared device with the instruction (a gfx11 architecture) must run the +//! kernel and pass its probe; one without must be refused as unsupported. Each test prints what it established and, when +//! `WAI_QUANTUM_HIP_REPORT` names a file, appends the same lines to it. +//! +//! ```text +//! WAI_QUANTUM_HIP_EXPECT=gpu-hip-gfx1151 \ +//! cargo test --features quantum_fabric_hip --test hip_wmma -- --test-threads=1 +//! ``` +#![cfg(not(target_arch = "wasm32"))] + +use std::io::Write; +use std::sync::OnceLock; +use wai_quantum::ed25519_dalek::SigningKey; +use wai_quantum::quantum::{BaseGate, Circuit, Gate, Gateset, MAX_QUBITS, QuantumError}; +use wai_quantum::quantum_fabric::hip::{self, WMMA_MUTANTS}; +use wai_quantum::quantum_fabric::lanes::{self, GateWords, Lane}; +use wai_quantum::quantum_fabric::plan::{Program, Step}; +use wai_quantum::quantum_fabric::wgsl::{Detector, Mutant}; +use wai_quantum::quantum_fabric::{ + Corpus, CpuReference, FABRIC_DEFAULT, Fabric, FabricClass, FabricReceipt, MemberError, Parity, Presence, REFERENCE, apply_kernel, detect_local, place, ready, +}; +use wai_quantum::quantum_fabric_hip::{Discovery, HipError, HipMember, discover}; + +/// The CPU's corpus digests: random-v1 (seed 7, 200 cases) and fabric-v1 +/// (seed 7, 200 + 24 cases), per gateset. A bit-exact member reproduces them. +const PINS: [(Gateset, &str, &str); 2] = [ + (Gateset::V1, "d18454c639acfd9c7b66450f4ded1aac4ab46cc10a7bedbcf4958a96e204fb36", "41c726a0ac80f7defef92e15ecf893bdb7a3076b5e7f670fa0b7142e5c033df7"), + (Gateset::V2, "3dbde276c1095bf827faec71d5980e050f8a55a5ba7d97cd3bbdaec3f43a13b2", "9a84b55af68f1a8b07390b234cf562dd8d85cef3185a697ab85690c8790a7809"), +]; + +const RANDOM_V1: Corpus = Corpus::Random { seed: 7, cases: 200 }; + +static DISCOVERY: OnceLock = OnceLock::new(); +static WMMA: OnceLock> = OnceLock::new(); + +fn discovery() -> &'static Discovery { + DISCOVERY.get_or_init(discover) +} + +/// Prints `line` and appends it to the file `WAI_QUANTUM_HIP_REPORT` names. +fn report(line: &str) { + println!("{line}"); + if let Some(path) = std::env::var_os("WAI_QUANTUM_HIP_REPORT") { + let mut f = std::fs::OpenOptions::new().create(true).append(true).open(&path).expect("the report file opens"); + writeln!(f, "{line}").expect("the report line is written"); + } +} + +/// The declared device set, as written, if it is exactly what was found; +/// otherwise why not, with what was found and refused and how to declare it. +fn declared(d: &Discovery) -> Result { + let mut found: Vec<&str> = d.found.iter().map(|m| m.key()).collect(); + found.sort_unstable(); + let found_list = if found.is_empty() { "none".to_string() } else { found.join(",") }; + let refused: Vec = d.refused.iter().map(|r| format!("{} ({:?}): {}", r.key, r.device, r.error)).collect(); + let refused_list = if refused.is_empty() { "none".to_string() } else { refused.join("; ") }; + let unavailable = d.unavailable.as_deref().unwrap_or("no"); + let how = format!("found: {found_list}; refused: {refused_list}; unavailable: {unavailable}. Declare the devices this machine must have: WAI_QUANTUM_HIP_EXPECT={found_list}"); + let expect = match std::env::var("WAI_QUANTUM_HIP_EXPECT") { + Ok(v) if !v.is_empty() => v, + _ => return Err(format!("WAI_QUANTUM_HIP_EXPECT is not set. {how}")), + }; + let keys: Vec<&str> = if expect == "none" { Vec::new() } else { expect.split(',').collect() }; + let well_formed = keys.windows(2).all(|w| w[0] <= w[1]) && keys.iter().all(|k| !k.is_empty() && *k != "none" && k.trim() == *k); + if !well_formed { + return Err(format!("WAI_QUANTUM_HIP_EXPECT={expect:?} is not a sorted, comma-separated list of keys, nor none. {how}")); + } + for k in &keys { + if !found.contains(k) + && let Some(r) = d.refused.iter().find(|r| r.key == *k) + { + return Err(format!("the declared device {k} was refused: {}. {how}", r.error)); + } + } + if found != keys { + return Err(format!("the devices found are not the declared ones ({expect}). {how}")); + } + Ok(expect) +} + +/// The matrix-instruction members: one per declared device with the instruction, once +/// the declared devices are exactly the ones found. +fn members() -> &'static [HipMember] { + let d = discovery(); + if let Err(why) = declared(d) { + panic!("{why}"); + } + WMMA.get_or_init(|| { + d.found.iter().filter(|m| m.identity().has_wmma()).map(|m| m.wmma().unwrap_or_else(|e| panic!("{}: the matrix-instruction kernel: {e}", m.key()))).collect() + }) +} + +fn hex(b: &[u8]) -> String { + b.iter().map(|x| format!("{x:02x}")).collect() +} + +fn next(s: &mut u64) -> u64 { + *s = s.wrapping_add(0x9E37_79B9_7F4A_7C15); + let mut z = *s; + z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9); + z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB); + z ^ (z >> 31) +} + +/// A value in `[-bound, bound]`. +fn within(s: &mut u64, bound: i64) -> i32 { + ((next(s) % (2 * bound as u64 + 1)) as i64 - bound) as i32 +} + +/// A random circuit of `ops` gates on `n` qubits: every base gate, up to three +/// distinct controls, `P(1..=32)`. +fn random_circuit(n: u8, ops: usize, seed: u64, gs: Gateset) -> Circuit { + let bases = [BaseGate::I, BaseGate::X, BaseGate::Y, BaseGate::Z, BaseGate::H, BaseGate::S, BaseGate::Sdg, BaseGate::T, BaseGate::Tdg, BaseGate::P]; + let mut s = seed; + let mut c = Circuit::with_gateset(n, gs); + for _ in 0..ops { + let base = bases[(next(&mut s) % bases.len() as u64) as usize]; + let target = (next(&mut s) % u64::from(n)) as u8; + let mut controls = Vec::new(); + for _ in 0..next(&mut s) % 4 { + let q = (next(&mut s) % u64::from(n)) as u8; + if q != target && !controls.contains(&q) { + controls.push(q); + } + } + let param = if base == BaseGate::P { 1 + (next(&mut s) % 32) as u16 } else { 0 }; + c.ops.push(Gate { base, controls, target, param }); + } + c +} + +/// Five circuits the native simulator refuses, in each gateset. +fn refused_circuits() -> Vec { + let gate = |base, controls: Vec, target, param| Gate { base, controls, target, param }; + let bad = [ + (1, vec![gate(BaseGate::H, vec![], 0, 0), gate(BaseGate::P, vec![], 0, 0)]), + (1, vec![gate(BaseGate::P, vec![], 0, 33)]), + (1, vec![gate(BaseGate::X, vec![], 3, 0)]), + (2, vec![gate(BaseGate::X, vec![0], 0, 0)]), + (2, vec![gate(BaseGate::X, vec![5], 1, 0)]), + ]; + let mut out = Vec::new(); + for gs in Gateset::ALL { + for (n, ops) in bad.clone() { + let mut c = Circuit::with_gateset(n, gs); + c.ops = ops; + out.push(c); + } + } + out +} + +/// `Z` on qubit 0: on `[(0, 0), (i32::MIN, 0)]` its exact result is `2^31`, +/// which leaves 32 bits. +fn z_words() -> GateWords { + lanes::gate_words(&Gate { base: BaseGate::Z, controls: vec![], target: 0, param: 0 }, Gateset::V2) +} + +const Z_STATE: [Lane; 2] = [[0, 0], [i32::MIN, 0]]; + +#[test] +fn c0_devices_are_declared() { + let d = discovery(); + match declared(d) { + Ok(expect) => report(&format!("hip wmma devices: as declared ({expect})")), + Err(why) => panic!("{why}"), + } + for m in &d.found { + if !m.identity().has_wmma() { + assert!(matches!(m.wmma(), Err(HipError::Unsupported(_))), "{}", m.key()); + report(&format!("hip wmma unsupported: key={} arch {}", m.key(), m.identity().arch)); + } + } + for t in members() { + assert_eq!(t.class(), FabricClass::Gpu); + assert!(t.id().starts_with("local/gpu-hip-"), "{}", t.id()); + assert_eq!(t.implementation(), hip::wmma_implementation_id()); + assert_ne!(t.implementation(), hip::implementation_id()); + assert_eq!(t.identity().member_id(), t.id()); + let lanes64 = d.found.iter().find(|m| m.id() == t.id()).expect("the device's member"); + assert_eq!((lanes64.implementation(), lanes64.max_qubits()), (hip::implementation_id().as_str(), t.max_qubits())); + assert!(t.max_qubits() >= 20, "{}: holds only {} qubits", t.key(), t.max_qubits()); + assert!(t.detail().starts_with(lanes64.detail()) && t.detail().contains("matrix instruction"), "{}", t.detail()); + report(&format!("hip wmma member: key={} id={} implementation={} {}", t.key(), t.id(), t.implementation(), t.detail())); + } +} + +#[test] +fn c1_probe_matches_lanes() { + for m in members() { + let pairs = lanes::probe_pairs(1 << 16, 0x0cda); + let out = m.probe(&pairs).unwrap_or_else(|e| panic!("{}: {e}", m.key())); + assert_eq!(lanes::first_mismatch(&pairs, &out), None, "{}", m.key()); + report(&format!("hip wmma probe: key={} {} pairs agree", m.key(), pairs.len())); + } +} + +#[test] +fn c2_kernel_matches_run_words() { + for m in members() { + let mut s = 0x6b65_726e_656c_u64; + let (mut flagged, mut clean) = (0, 0); + for case in 0..2000 { + let n = 1 + (next(&mut s) % 12) as u32; + // Three regimes: the whole 32-bit range (nearly always out of + // range), small states and entries (never), and states and + // entries up to 2^30 (either). + let mode = case % 3; + let value = |s: &mut u64| -> i32 { + match mode { + 0 => next(s) as i32, + 1 => within(s, 1 << 22), + _ => within(s, 1 << 30), + } + }; + let entry = |s: &mut u64| -> u32 { + match mode { + 0 => next(s) as u32, + 1 => within(s, 1 << 29) as u32, + _ => within(s, 1 << 30) as u32, + } + }; + let state: Vec = (0..1usize << n).map(|_| [value(&mut s), value(&mut s)]).collect(); + let gates: Vec = (0..1 + next(&mut s) % 8) + .map(|_| { + let t = (next(&mut s) % u64::from(n)) as u32; + let ctrl = (0..n).filter(|&b| b != t && next(&mut s).is_multiple_of(4)).fold(0u32, |acc, b| acc | 1 << b); + let mut w = [0u32; 8]; + w.iter_mut().for_each(|x| *x = entry(&mut s)); + GateWords { m: w, tbit: 1 << t, ctrl } + }) + .collect(); + let want = lanes::run_words(&state, &gates).expect("valid shapes"); + let got = m.run_lanes(&state, &gates).unwrap_or_else(|e| panic!("{}: case {case}: {e}", m.key())); + assert_eq!(got.out_of_range, want.out_of_range, "{}: case {case}, n = {n}", m.key()); + assert!(got.amps == want.amps, "{}: case {case}, n = {n}: the lanes differ", m.key()); + if want.out_of_range.is_some() { + flagged += 1; + } else { + clean += 1; + } + } + assert!(flagged >= 100 && clean >= 100, "{}: {flagged} flagged, {clean} clean", m.key()); + report(&format!("hip wmma kernel: key={} ok", m.key())); + } +} + +#[test] +fn c3_certificates_reproduce_the_pins() { + let ms = members(); + for (gs, random_pin, fabric_pin) in PINS { + let cpu = lanes::certify_corpus(gs, "local/test", RANDOM_V1); + assert_eq!(hex(&cpu.corpus), random_pin, "{gs:?}"); + let other = Corpus::Random { seed: 8, cases: 300 }; + let cpu_other = lanes::certify_corpus(gs, "local/test", other); + for m in ms { + for (corpus, cases, pin) in [(RANDOM_V1, 200, random_pin), (FABRIC_DEFAULT, 224, fabric_pin)] { + let c = m.certify(gs, corpus).unwrap_or_else(|e| panic!("{} {gs:?} {}: {e}", m.key(), corpus.name())); + assert_eq!(c.parity, Parity::BitExact, "{} {gs:?} {}", m.key(), corpus.name()); + assert_eq!(c.cases, cases); + assert_eq!(hex(&c.corpus), pin, "{} {gs:?} {}", m.key(), corpus.name()); + assert_eq!((c.kernel.as_str(), c.implementation.as_str(), c.member.as_str()), (apply_kernel(gs), m.implementation(), m.id())); + assert_eq!(c.generator, corpus.generator(gs)); + report(&format!( + "hip wmma certified: key={} member={} kernel={} corpus={} implementation={} parity=bit-exact cases={} digest={}", + m.key(), + m.id(), + c.kernel, + corpus.name(), + c.implementation, + c.cases, + hex(&c.corpus) + )); + } + // A corpus with no pin: the member and the CPU agree on it. + let c = m.certify(gs, other).unwrap_or_else(|e| panic!("{} {gs:?} seed 8: {e}", m.key())); + assert_eq!((c.parity, c.cases), (Parity::BitExact, 300), "{} {gs:?}", m.key()); + assert_eq!(c.corpus, cpu_other.corpus, "{} {gs:?}", m.key()); + } + } +} + +#[test] +fn c4_wide_circuits_match_native() { + for m in members() { + let widest = m.max_qubits().min(22); + let mut s = 0x0077_6964_6500_u64; + for (k, n) in [12u8, 16, 18, 20, widest].into_iter().enumerate() { + for gs in Gateset::ALL { + let c = random_circuit(n, 40, 0x1000 + k as u64, gs); + let start = (next(&mut s) as usize) % (1usize << n); + let want = c.simulate_from(start).expect("the native run"); + let got = m.simulate_from(&c, start).unwrap_or_else(|e| panic!("{}: n = {n} {gs:?}: {e}", m.key())); + assert!(got == want, "{}: n = {n} {gs:?}: the states differ", m.key()); + } + } + // The transforms the conformance corpus carries. + for n in [3u8, 5, 8, 12] { + let qft = Circuit::qft(n); + assert!(m.simulate(&qft).unwrap() == qft.simulate().unwrap(), "{}: QFT-{n}", m.key()); + } + report(&format!("hip wmma wide: key={} up to {widest} qubits match native", m.key())); + } +} + +#[test] +fn c5_refusals_are_typed() { + for m in members() { + for c in refused_circuits() { + let native = c.simulate_from(0).unwrap_err(); + assert_eq!(m.simulate_from(&c, 0).unwrap_err(), HipError::Circuit(native), "{c:?}"); + } + assert_eq!(m.simulate(&Circuit::new(MAX_QUBITS + 1)).unwrap_err(), HipError::Circuit(QuantumError::TooManyQubits(MAX_QUBITS + 1))); + let small = m.with_max_qubits(4); + assert_eq!(small.simulate(&Circuit::new(5)).unwrap_err(), HipError::TooWide { n_qubits: 5, max: 4 }); + let wide = Program { n_qubits: 5, start: 0, gates: Vec::new(), steps: Vec::new() }; + assert_eq!(small.run_program(&wide).unwrap_err(), HipError::TooWide { n_qubits: 5, max: 4 }); + // A result that leaves 32 bits is reported at its op, and the lane + // holds its low 32 bits. + let r = m.run_lanes(&Z_STATE, &[z_words()]).unwrap(); + assert_eq!(r.out_of_range, Some(0), "{}", m.key()); + assert_eq!(r.amps[1][0] as u32, 0x8000_0000); + // Malformed words, a state the host cannot lay out and a tile step are + // host defects, never results. + let bad = GateWords { tbit: 3, ..z_words() }; + assert!(matches!(m.run_lanes(&[[0, 0]; 4], &[bad]), Err(HipError::Internal(_)))); + assert!(matches!(m.run_lanes(&[[0, 0]; 3], &[z_words()]), Err(HipError::Internal(_)))); + let tiled = Program { n_qubits: 2, start: 0, gates: Vec::new(), steps: vec![Step::Tile { first: 0, count: 1, size: 2 }] }; + assert!(matches!(m.run_program(&tiled), Err(HipError::Internal(_)))); + // The refusals that may still fall back to another member, do. + assert!(MemberError::from(HipError::TooWide { n_qubits: 5, max: 4 }).falls_back()); + report(&format!("hip wmma refusals: key={} typed", m.key())); + } +} + +/// Whether a member running `mutant`'s text is caught by the mutant's +/// detector. +fn caught(mm: &HipMember, mutant: &Mutant, key: &str) -> bool { + match mutant.detector { + Detector::Probe => { + let edges = lanes::probe_pairs(0, 0); + let out = mm.probe(&edges).expect("the mutant's probe runs"); + let differing = edges.iter().zip(&out).filter(|&(&(a, b), o)| *o != (lanes::mul_wide(a, b), lanes::fxmul_wide(a, b))).count(); + println!("hip wmma mutant {} on {key}: {differing} of 484 edge pairs differ", mutant.name); + differing >= 1 + } + Detector::Certificate => Gateset::ALL.into_iter().all(|gs| match mm.certify(gs, RANDOM_V1) { + Err(e) => { + println!("hip wmma mutant {} on {key}: {gs:?} certification refused: {e}", mutant.name); + true + } + Ok(c) => { + let cpu = lanes::certify_corpus(gs, "local/test", RANDOM_V1); + c.parity != Parity::BitExact || c.corpus != cpu.corpus + } + }), + Detector::Range => mm.run_lanes(&Z_STATE, &[z_words()]).map(|r| r.out_of_range) != Ok(Some(0)), + other => panic!("no matrix-instruction mutant uses the {other:?} detector"), + } +} + +#[test] +fn c6_mutants_are_killed() { + assert_eq!(lanes::probe_pairs(0, 0).len(), 484); + for m in members() { + let mut killed = 0; + for mutant in &WMMA_MUTANTS { + let mm = m.with_wmma_source(&hip::wmma_mutate(mutant)).unwrap_or_else(|e| panic!("{}: mutant {} does not compile: {e}", m.key(), mutant.name)); + assert_ne!(mm.implementation(), m.implementation()); + assert!(caught(&mm, mutant, m.key()), "{}: mutant {} survived", m.key(), mutant.name); + killed += 1; + } + assert_eq!(killed, WMMA_MUTANTS.len()); + report(&format!("hip wmma mutants: key={} killed {killed} of {}", m.key(), WMMA_MUTANTS.len())); + } +} + +#[test] +fn c7_receipt_placement_and_the_fabric() { + for m in members() { + let member = m.member(FABRIC_DEFAULT).unwrap_or_else(|e| panic!("{}: {e}", m.key())); + assert_eq!((member.id.as_str(), member.class, member.presence), (m.id(), FabricClass::Gpu, Presence::Detected)); + assert_eq!(member.kernels.len(), 2); + for (entry, gs) in member.kernels.iter().zip(Gateset::ALL) { + assert_eq!((entry.kernel.as_str(), entry.implementation.as_str()), (apply_kernel(gs), m.implementation())); + assert!(entry.energy.is_none()); + assert_eq!(entry.certificate.parity, Parity::BitExact, "{gs:?}"); + } + // A receipt over a member's run. + let qft = Circuit::qft(5); + let state = m.simulate(&qft).unwrap(); + assert_eq!(state, qft.simulate().unwrap()); + let entry = member.kernels.iter().find(|k| k.kernel == apply_kernel(qft.gateset)).unwrap(); + let signer = SigningKey::from_bytes(&[7; 32]); + let receipt = FabricReceipt::seal(&signer, "local/test", state.statevector_hash(), &member, entry, 0, None); + assert!(receipt.verify()); + assert!(receipt.certificate_matches(&entry.certificate)); + // Placement: beside the CPU reference with no energy figures, the + // reference; alone, the member. + let local = detect_local(); + for gs in Gateset::ALL { + let both = [local.clone(), member.clone()]; + let (placed, k) = place(&both, apply_kernel(gs), &Parity::BitExact).unwrap(); + assert_eq!((placed.id.as_str(), k.implementation.as_str()), (local.id.as_str(), REFERENCE)); + let (placed, k) = place(std::slice::from_ref(&member), apply_kernel(gs), &Parity::BitExact).unwrap(); + assert_eq!((placed.id.as_str(), k.implementation.as_str()), (m.id(), m.implementation())); + } + // On the fabric: the member alone runs the circuit; a pairing with + // another member's entries is refused. + let mut fabric = Fabric::new(); + fabric.add(member.clone(), m); + let run = ready(fabric.run(&qft, 0, &Parity::BitExact, &signer, "local/test")).expect("resolves").expect("runs"); + assert_eq!(run.state, state); + assert_eq!(run.receipt.member_id, m.id()); + assert!(run.receipt.verify()); + let mut mismatched = Fabric::new(); + mismatched.add(local.clone(), m); + assert!(ready(mismatched.run(&qft, 0, &Parity::BitExact, &signer, "local/test")).expect("resolves").is_err()); + let mut both = Fabric::new(); + both.add(local.clone(), CpuReference).add(member.clone(), m); + let run = ready(both.run(&qft, 0, &Parity::BitExact, &signer, "local/test")).expect("resolves").expect("runs"); + assert_eq!(run.state, state); + report(&format!("hip wmma receipt: key={} ok", m.key())); + } +} + +#[test] +fn c8_repeated_runs_agree() { + for m in members() { + let c = random_circuit(16, 60, 0xabc, Gateset::V2); + let first = m.simulate(&c).unwrap(); + for _ in 0..5 { + assert!(m.simulate(&c).unwrap() == first, "{}", m.key()); + } + report(&format!("hip wmma repeat: key={} ok", m.key())); + } +} diff --git a/wai/wai-rs/Cargo.lock b/wai/wai-rs/Cargo.lock index 537bb75e9..c1ec42dd9 100644 --- a/wai/wai-rs/Cargo.lock +++ b/wai/wai-rs/Cargo.lock @@ -4310,7 +4310,7 @@ version = "0.1.0" [[package]] name = "wai-quantum" -version = "0.6.1" +version = "0.6.2" dependencies = [ "blake3", "ed25519-dalek", diff --git a/wai/wai-web/Cargo.lock b/wai/wai-web/Cargo.lock index 67cf542a0..7e93c5ad2 100644 --- a/wai/wai-web/Cargo.lock +++ b/wai/wai-web/Cargo.lock @@ -1241,7 +1241,7 @@ version = "0.1.0" [[package]] name = "wai-quantum" -version = "0.6.1" +version = "0.6.2" dependencies = [ "blake3", "ed25519-dalek", -- 2.54.0 (Apple Git-157) From fb08994eeea3eef88db439d968f9dfcf82a829c1 Mon Sep 17 00:00:00 2001 From: Transaction Science Date: Fri, 16 Oct 2026 10:19:30 -0400 Subject: [PATCH 2/2] wai-quantum: the Metal member's tensor-operation kernel; M5 measured tensor_ops.metal (metal-tensor-ops-limbs8/..., MetalMember::tensor_ops): every product formed by the shading language's 4.0 tensor operation (matmul2d, bf16 accumulated in f32), which the platform maps to the GPU's matrix hardware where it has it. Each threadgroup owns a scratch slice in device memory (64x32 limbs, 32x128 weights, 64x128 f32 coefficients); a 64-pair tile is one multiply in four 32-column runs. Offset-binary limbs, as for the HIP kernel. Family 1007+ under OS 26+, else Unsupported. Eight mutants. Build records each kernel's language version and scratch. Measured, silicon: - family 1009 (M3 Ultra build host / macos runner): 9/9, 66020 probe pairs, four pins, 22 qubits, 8 of 8 mutants; - family 1010 (the only M5 in the fleet, the dev Mac; test binaries built on the build host, run once niced): all three Metal kernels 9/9, four pins each, 6/6, 7/7 and 8/8 mutants. With signed limbs instead, the tensor op agrees on all 66020 pairs on both generations. Off macOS all three suites pass over no device. CI: a step of wai-quantum-metal and its absent leg. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/wai-quantum-metal.yml | 23 +- wai/CHANGELOG.md | 24 +- wai/wai-quantum/Cargo.toml | 4 + wai/wai-quantum/README.md | 2 +- wai/wai-quantum/src/quantum_fabric/metal.rs | 132 +++++- .../src/quantum_fabric/tensor_ops.metal | 278 +++++++++++ wai/wai-quantum/src/quantum_fabric_metal.rs | 120 ++++- wai/wai-quantum/tests/metal_tensor_ops.rs | 446 ++++++++++++++++++ 8 files changed, 1003 insertions(+), 26 deletions(-) create mode 100644 wai/wai-quantum/src/quantum_fabric/tensor_ops.metal create mode 100644 wai/wai-quantum/tests/metal_tensor_ops.rs diff --git a/.github/workflows/wai-quantum-metal.yml b/.github/workflows/wai-quantum-metal.yml index 559064178..903ed74a7 100644 --- a/.github/workflows/wai-quantum-metal.yml +++ b/.github/workflows/wai-quantum-metal.yml @@ -10,7 +10,9 @@ name: wai-quantum Metal member # (two gatesets, random-v1 and fabric-v1), and catch every kernel mutant; and # the matrix kernel, every product formed by the GPU's SIMD-group matrix # multiply from 8-bit limbs, must do all of the same under its own -# implementation id, with its own mutants caught. +# implementation id, with its own mutants caught; and so must the +# tensor-operation kernel, every product formed by the shading language's +# tensor operation from offset-binary 8-bit limbs. # - metal-absent (a container with no Metal): the same build finds no member # and says why, rather than failing; the device tests run over no member with # the declaration `none`. @@ -103,6 +105,22 @@ jobs: n="$(grep -c "^metal matrix certified: key=$key .* implementation=metal-matrix-limbs8/[0-9a-f]\{16\} parity=bit-exact " "$report")" [ "$n" = 4 ] || { echo "::error::$n of 4 matrix certificates were earned"; exit 1; } grep -qx "metal matrix mutants: key=$key killed 7 of 7" "$report" || { echo "::error::a matrix mutant survived"; exit 1; } + - name: The tensor-operation kernel, every product by the tensor operation + shell: bash + working-directory: wai/wai-quantum + run: | + set -euo pipefail + report="$RUNNER_TEMP/metal-tensor-ops-report.txt"; : > "$report" + WAI_QUANTUM_METAL_REPORT="$report" nice -n 10 \ + cargo test --locked --release --features quantum_fabric_metal --test metal_tensor_ops -- --test-threads=1 2>&1 | tee "$RUNNER_TEMP/tensor-ops.log" + grep -qE '^test result: ok\. 9 passed; 0 failed' "$RUNNER_TEMP/tensor-ops.log" || { echo "::error::not every tensor-operation test ran and passed"; exit 1; } + cat "$report" + key="$WAI_QUANTUM_METAL_EXPECT" + grep -qx "metal tensor-ops devices: as declared ($key)" "$report" || { echo "::error::the devices found are not the declared ones"; exit 1; } + grep -q "^metal tensor-ops member: key=$key .* implementation=metal-tensor-ops-limbs8/[0-9a-f]\{16\} " "$report" || { echo "::error::the device ran no tensor-operation member"; exit 1; } + n="$(grep -c "^metal tensor-ops certified: key=$key .* implementation=metal-tensor-ops-limbs8/[0-9a-f]\{16\} parity=bit-exact " "$report")" + [ "$n" = 4 ] || { echo "::error::$n of 4 tensor-operation certificates were earned"; exit 1; } + grep -qx "metal tensor-ops mutants: key=$key killed 8 of 8" "$report" || { echo "::error::a tensor-operation mutant survived"; exit 1; } metal-absent: name: member · Metal absent (linux) @@ -131,9 +149,10 @@ jobs: cargo test --locked --features quantum_fabric_metal --lib -- quantum_fabric_metal:: quantum_fabric::metal:: 2>&1 | tee "$RUNNER_TEMP/unit.log" grep -qE '^test quantum_fabric_metal::tests::off_macos_nothing_is_found_and_discovery_says_why \.\.\. ok$' "$RUNNER_TEMP/unit.log" || { echo "::error::the off-platform discovery test did not run"; exit 1; } report="$RUNNER_TEMP/metal-report.txt"; : > "$report" - for t in metal_member metal_matrix; do + for t in metal_member metal_matrix metal_tensor_ops; do WAI_QUANTUM_METAL_REPORT="$report" cargo test --locked --features quantum_fabric_metal --test "$t" -- --test-threads=1 2>&1 | tee "$RUNNER_TEMP/$t.log" grep -qE '^test result: ok\. 9 passed; 0 failed' "$RUNNER_TEMP/$t.log" || { echo "::error::$t did not run over no device"; exit 1; } done grep -qx 'metal devices: as declared (none)' "$report" || { echo "::error::the absent case was not established"; exit 1; } grep -qx 'metal matrix devices: as declared (none)' "$report" || { echo "::error::the matrix absent case was not established"; exit 1; } + grep -qx 'metal tensor-ops devices: as declared (none)' "$report" || { echo "::error::the tensor-operation absent case was not established"; exit 1; } diff --git a/wai/CHANGELOG.md b/wai/CHANGELOG.md index 9e52b129c..6734836e3 100644 --- a/wai/CHANGELOG.md +++ b/wai/CHANGELOG.md @@ -6,10 +6,32 @@ wire format may change and such changes are recorded here. ## v0.4.0 — current -### wai-quantum 0.6.2: the HIP member's matrix-instruction kernel +### wai-quantum 0.6.2: the HIP member's matrix-instruction kernel, and the Metal member's tensor-operation kernel Additive. Every pinned digest is unchanged. +- **Every product by the shading language's tensor operation** + (`quantum_fabric::metal::TENSOR_OPS_SOURCE`, `tensor_ops.metal`; `MetalMember::tensor_ops`): + a third Metal kernel for the gate, through the language's version 4.0 tensor operation (a + bf16 matrix multiply accumulated in f32), which the platform maps to the GPU's matrix + hardware where it has it. The operation reads and writes device memory, so each threadgroup + owns a slice of scratch (`TENSOR_OPS_SCRATCH_BYTES`): a 64-pair tile is one multiply of + 64 × 32 limbs by 32 × 128 weights, in four 32-column runs. Its operands are offset-binary, as + for the HIP kernel below, so the result does not depend on how an accumulator treats a + negative term. A device or operating system without the operation + (`DeviceIdentity::has_tensor_ops`: GPU family 1007 or later, operating system 26 or later) + is `MetalError::Unsupported`. Eight mutants (`TENSOR_OPS_MUTANTS`), among them nine-bit + limbs, every run reading the first columns of the weights, and a wrong offset correction. +- **Measured** on two GPU generations, silicon: family 1009 (the macOS runner's) and family + 1010, the generation whose GPU cores carry matrix accelerators. On each: 9/9, 66020 probe + pairs, all four pinned certificates, circuits to 22 qubits, 8 of 8 mutants killed. On family + 1010 the integer and SIMD-group matrix kernels are bit-exact too (9/9 each, 6 of 6 and 7 of 7 + mutants). With signed limbs instead, the tensor operation still agrees on all 66020 pairs on + both generations: the offset is for the instruction below, not needed here. Where the + platform runs the operation is its own choice; it is not observed here. CI: a step of + `wai-quantum-metal` on family 1009; family 1010 was run by hand, the binaries built on the + build host. + - **Every product by the bf16 matrix instruction** (`quantum_fabric::hip::WMMA_SOURCE`, `wmma_limbs.hip`; `HipMember::wmma`): a second HIP kernel for the same gate, for the gfx11 architectures (`DeviceIdentity::has_wmma`), whose wavefronts of 32 multiply bf16 at diff --git a/wai/wai-quantum/Cargo.toml b/wai/wai-quantum/Cargo.toml index c25a09607..71aca7807 100644 --- a/wai/wai-quantum/Cargo.toml +++ b/wai/wai-quantum/Cargo.toml @@ -226,6 +226,10 @@ required-features = ["quantum_fabric_metal"] name = "metal_matrix" required-features = ["quantum_fabric_metal"] +[[test]] +name = "metal_tensor_ops" +required-features = ["quantum_fabric_metal"] + [[test]] name = "cuda_sv" required-features = ["quantum_fabric_cuda", "quantum_sv"] diff --git a/wai/wai-quantum/README.md b/wai/wai-quantum/README.md index d3263b076..7f2b776da 100644 --- a/wai/wai-quantum/README.md +++ b/wai/wai-quantum/README.md @@ -106,7 +106,7 @@ never creates information. | `quantum_fabric` | **the compute fabric**: members (the CPU and accelerators, quantum processors among them), each kernel implementation certified against the CPU reference in the process that uses it (bit-exact, within a stated tolerance, or statistical by classical fidelity over stated shots), with the certificate saying where it was earned (silicon, an emulator, or a simulator); placement by parity, provenance and metered energy; a signed receipt naming the member and implementation behind a result. Two kernels: apply (a circuit's state) and sample (its measurement counts), the latter with the CPU reference sampler, a corpus shaped for quantum processors, and a virtual quantum processor whose stated noise model lets every step of a processor's path run and be tested without one; its certificates say they were earned in a simulator. `quantum_fabric::stablehlo` writes the gate as a portable tensor program (StableHLO), every value a 64-bit integer that never overflows, for any tensor compiler to build for its own devices; the `stablehlo_gate` example writes it with a batch of inputs and the exact outputs a build must give, and `ci/tensor-program/check.sh` runs a build and holds it to them | | `quantum_fabric_webgpu` | **the fabric's WebGPU member**, opt-in and not in `full`: the 32-bit-lane gate kernel in WGSL, one dispatch per gate, through Metal, Vulkan or DX12 natively and through a browser's WebGPU on `wasm32-unknown-unknown`; a second implementation applies runs of gates tile by tile, one dispatch per run, and is certified on its own. Each adapter is probed against the CPU's arithmetic in the process that uses it and certified bit-exact by reproducing the CPU reference's pinned corpus digests; a result that would leave 32 bits is refused, never returned. One of the features that reach outside the crate | | `quantum_fabric_cuda` | **the fabric's CUDA member**, opt-in and not in `full`, native targets only: the gate kernel in CUDA C with exact 64-bit integers, compiled at run time for each device's own architecture, one launch per gate. The CUDA driver and runtime compiler are loaded when first used, so a build with the feature runs on a host without CUDA and finds no member there, saying why. Each device is probed against the CPU's arithmetic in the process that uses it and certified bit-exact by reproducing the CPU reference's pinned corpus digests; a result that would leave 32 bits is refused, never returned. On a device of compute capability 8.0 or later, a **second kernel runs every product on the tensor cores** (`CudaMember::tensor`): each value as four 8-bit limbs in bf16, accumulated in f32, which holds every partial sum exactly; probed and certified under its own implementation id, with its own mutants. With `quantum_sv` too, the **dense f64 state vector** runs on the device bit for bit as on the CPU: the CPU's own launches (diagonal, one-qubit and grouped formulas), each an IEEE binary64 operation with no fused multiply-add, certified against a pinned corpus digest | -| `quantum_fabric_metal` | **the fabric's Metal member**, opt-in and not in `full`, macOS: the gate kernel in the Metal shading language with exact 64-bit integers, compiled at run time by the operating system's own compiler, one serial dispatch per gate in one command buffer, the state in memory the GPU and the host share. Elsewhere the build finds no member and says why. Each device is probed against the CPU's arithmetic in the process that uses it and certified bit-exact by reproducing the CPU reference's pinned corpus digests, keyed by its GPU family number (`gpu-msl-`). On a GPU with SIMD-group matrices, a **second kernel forms every product with the matrix multiply** (`MetalMember::matrix`): each value as four 8-bit limbs in bf16, accumulated in f32, which holds every partial sum exactly; probed and certified under its own implementation id, with its own mutants | +| `quantum_fabric_metal` | **the fabric's Metal member**, opt-in and not in `full`, macOS: the gate kernel in the Metal shading language with exact 64-bit integers, compiled at run time by the operating system's own compiler, one serial dispatch per gate in one command buffer, the state in memory the GPU and the host share. Elsewhere the build finds no member and says why. Each device is probed against the CPU's arithmetic in the process that uses it and certified bit-exact by reproducing the CPU reference's pinned corpus digests, keyed by its GPU family number (`gpu-msl-`). On a GPU with SIMD-group matrices, a **second kernel forms every product with the matrix multiply** (`MetalMember::matrix`): each value as four 8-bit limbs in bf16, accumulated in f32, which holds every partial sum exactly; probed and certified under its own implementation id, with its own mutants. Under operating system 26 or later, a **third kernel forms every product with the shading language's tensor operation** (`MetalMember::tensor_ops`), which the platform maps to the GPU's matrix hardware where it has it: offset-binary 8-bit limbs in bf16, accumulated in f32, through per-threadgroup scratch | | `quantum_fabric_hip` | **the fabric's HIP member**, opt-in and not in `full`, Linux: the CUDA member's kernel text, unchanged, compiled at run time by HIP's runtime compiler for each device's own architecture, one launch per gate. The runtime and its compiler are loaded when first used, so a build with the feature runs on a host without them and finds no member there, saying why; the device's properties are read in the runtime's frozen layout or the device is refused. Probed and certified as the CUDA member is, under its own implementation id, and with `quantum_sv` the dense f64 state vector runs bit for bit as on the CPU. On a gfx11 GPU, a **second kernel forms every product with the bf16 matrix instruction** (`HipMember::wmma`): each value offset-binary, as four unsigned 8-bit limbs, because the instruction's accumulation is not exact once a negative term enters a sum; the offset is removed in 64-bit integers. Probed and certified under its own implementation id, with its own mutants | | `quantum_fabric_opencl` | **the fabric's OpenCL member**, opt-in and not in `full`, native targets only: the same exact 64-bit kernel in OpenCL C, built at run time by each device's own compiler, one launch per gate, on any OpenCL GPU or CPU (an OpenCL accelerator is refused: no fabric class names it). The ICD loader is opened when first used, so a build with the feature runs on a host without OpenCL and finds no member there, saying why. Each device is probed and certified as the CUDA member's are, and keyed by its vendor id (`opencl--`). With `quantum_sv` too, the dense f64 state vector runs bit for bit as on the CPU on every device with binary64 (contraction off); a device without it refuses, saying so | | `quantum_fabric_fpga` | **the fabric's FPGA member, simulated**, opt-in and not in `full`, native targets only: the gate datapath in Verilog (each 32×32 product from four 16-bit partial products, the size of an FPGA's multiplier blocks; a state RAM; one pair per two cycles), compiled at run time by Verilator into a cycle-accurate simulation driven over a pipe. Probed and certified as the other members are, and every certificate says it was earned in a simulator, so placement ranks it after silicon. A host without Verilator finds no member and says why | diff --git a/wai/wai-quantum/src/quantum_fabric/metal.rs b/wai/wai-quantum/src/quantum_fabric/metal.rs index fead29e44..e2f3e43e8 100644 --- a/wai/wai-quantum/src/quantum_fabric/metal.rs +++ b/wai/wai-quantum/src/quantum_fabric/metal.rs @@ -32,6 +32,19 @@ //! ([`matrix_implementation_id`]) and its own mutants ([`MATRIX_MUTANTS`]), //! one of which, nine-bit limbs, bf16 cannot hold. //! +//! **The tensor-operation kernel** ([`TENSOR_OPS_SOURCE`], `tensor_ops.metal`) +//! forms every product with the shading language's tensor operation (version +//! 4.0: a bf16 matrix multiply accumulated in f32), which the platform maps to +//! the GPU's matrix hardware where it has it. Its operands are offset-binary, +//! `x + 2^31`, so no term of any sum is negative, as for the HIP member's +//! matrix-instruction kernel ([`hip`](super::hip), which says why), and the +//! offset is removed in 64-bit integers. The operation reads and writes device +//! memory, so each threadgroup owns a slice of scratch +//! ([`TENSOR_OPS_SCRATCH_BYTES`]): a 64-pair tile is one multiply of +//! 64 × 32 limbs by 32 × 128 weights, in four 32-column runs. It has its own +//! implementation id ([`tensor_ops_implementation_id`]) and mutants +//! ([`TENSOR_OPS_MUTANTS`]). +//! //! Devices are named by number, not by vendor: [`DeviceIdentity::family`] is //! the highest of the platform's own GPU families (`MTLGPUFamily` values from //! 1001) the device reports. @@ -83,6 +96,9 @@ impl Shape { /// group takes eight pairs (four threads a pair) or sixteen cases (two /// threads a case). pub const MATRIX: Shape = Shape { threads: 128, per_pair: 4, per_case: 2 }; + /// The tensor-operation kernel's: 128 threads a threadgroup, which takes + /// 64 pairs (two threads a pair) or sixteen cases a tile. + pub const TENSOR_OPS: Shape = Shape { threads: 128, per_pair: 2, per_case: 8 }; /// Threadgroups for `items` at `per_item` threads each: at least one, at /// most [`GRID_MAX`]. @@ -100,6 +116,53 @@ pub const COMPILE_OPTIONS: [&str; 2] = ["language=3.1", "math=safe"]; /// The GPU family the matrix kernel needs at least: SIMD-group matrices. pub const MATRIX_MIN_FAMILY: u32 = 1007; +/// The tensor-operation kernel, `tensor_ops.metal`, exactly as it is +/// compiled. +pub const TENSOR_OPS_SOURCE: &str = include_str!("tensor_ops.metal"); + +/// The tensor-operation kernel's implementation name. +pub const TENSOR_OPS_IMPLEMENTATION: &str = "metal-tensor-ops-limbs8"; + +/// What the runtime compiler is given for the tensor-operation kernel: the +/// language version its tensors need, and safe math. +pub const TENSOR_OPS_COMPILE_OPTIONS: [&str; 2] = ["language=4.0", "math=safe"]; + +/// The operating system's major version from which the tensor operation is +/// there to compile. +pub const TENSOR_OPS_MIN_OS: u32 = 26; + +/// One threadgroup's scratch for the tensor-operation kernel, bytes: its +/// tile's limbs (64 × 32 bf16), the limb weights (32 × 128 bf16) and the +/// coefficients (64 × 128 f32), in that order. +pub const TENSOR_OPS_SCRATCH_BYTES: u64 = 64 * 32 * 2 + 32 * 128 * 2 + 64 * 128 * 4; + +/// Threadgroups per dispatch of the tensor-operation kernel, at most, so that +/// the scratch stays small; it strides over the rest. +pub const TENSOR_OPS_GROUP_MAX: u64 = 512; + +const DOMAIN_TENSOR_OPS_IMPLEMENTATION: &[u8] = b"wai:quantum-fabric-metal-tensor-ops\x01"; + +/// `"metal-tensor-ops-limbs8/<16 hex>"`: as [`implementation_id`], with the +/// domain `"wai:quantum-fabric-metal-tensor-ops\x01"`, [`TENSOR_OPS_SOURCE`], +/// [`TENSOR_OPS_COMPILE_OPTIONS`], [`Shape::TENSOR_OPS`], and the scratch +/// layout ([`TENSOR_OPS_SCRATCH_BYTES`], [`TENSOR_OPS_GROUP_MAX`]). +pub fn tensor_ops_implementation_id() -> String { + tensor_ops_implementation_id_of(TENSOR_OPS_SOURCE) +} + +/// [`tensor_ops_implementation_id`] for another kernel text, such as a +/// mutant's. +pub fn tensor_ops_implementation_id_of(source: &str) -> String { + id_with( + DOMAIN_TENSOR_OPS_IMPLEMENTATION, + TENSOR_OPS_IMPLEMENTATION, + source, + Shape::TENSOR_OPS, + &TENSOR_OPS_COMPILE_OPTIONS, + &[TENSOR_OPS_SCRATCH_BYTES, TENSOR_OPS_GROUP_MAX], + ) +} + const DOMAIN_IMPLEMENTATION: &[u8] = b"wai:quantum-fabric-metal\x01"; const DOMAIN_MATRIX_IMPLEMENTATION: &[u8] = b"wai:quantum-fabric-metal-matrix\x01"; const DOMAIN_MEMBER: &[u8] = b"wai:quantum-fabric-metal-member\x01"; @@ -115,13 +178,17 @@ fn put(o: &mut Vec, bytes: &[u8]) { } fn id_of(domain: &[u8], name: &str, source: &str, shape: Shape) -> String { + id_with(domain, name, source, shape, &COMPILE_OPTIONS, &[]) +} + +fn id_with(domain: &[u8], name: &str, source: &str, shape: Shape, options: &[&str], layout: &[u64]) -> String { let mut o = domain.to_vec(); put(&mut o, source.as_bytes()); - for option in COMPILE_OPTIONS { + for option in options { put(&mut o, option.as_bytes()); } o.extend_from_slice(&HOST_LAYOUT.to_be_bytes()); - for v in [shape.threads, shape.per_pair, shape.per_case] { + for v in [shape.threads, shape.per_pair, shape.per_case].iter().chain(layout) { o.extend_from_slice(&v.to_be_bytes()); } format!("{name}/{}", hex(&blake3::hash(&o).as_bytes()[..8])) @@ -211,6 +278,14 @@ impl DeviceIdentity { pub fn has_simdgroup_matrices(&self) -> bool { self.family >= MATRIX_MIN_FAMILY } + + /// Whether the device can run the tensor-operation kernel: family + /// [`MATRIX_MIN_FAMILY`] or later, under an operating system of major + /// version [`TENSOR_OPS_MIN_OS`] or later. + pub fn has_tensor_ops(&self) -> bool { + let major = self.os.split('.').next().and_then(|m| m.parse::().ok()).unwrap_or(0); + self.has_simdgroup_matrices() && major >= TENSOR_OPS_MIN_OS + } } /// The integer kernel's mutants: textual edits of [`SOURCE`] that a @@ -253,6 +328,33 @@ pub const MATRIX_MUTANTS: [Mutant; 7] = [ Mutant { name: "X7", find: "if (!fits32(o)) {", replace: "if (false) {", detector: Detector::Range }, ]; +/// The tensor-operation kernel's mutants. Each `find` occurs exactly once in +/// [`TENSOR_OPS_SOURCE`]. +pub const TENSOR_OPS_MUTANTS: [Mutant; 8] = [ + // The rounding constant is half what it should be. + Mutant { name: "N1", find: "return (v + (1L << 29)) >> 30;", replace: "return (v + (1L << 28)) >> 30;", detector: Detector::Probe }, + // The product is shifted by one bit too few. + Mutant { name: "N2", find: "return (v + (1L << 29)) >> 30;", replace: "return (v + (1L << 29)) >> 29;", detector: Detector::Probe }, + // Nine-bit limbs, which bf16 does not hold. + Mutant { name: "N3", find: "#define LIMB 8", replace: "#define LIMB 9", detector: Detector::Probe }, + // Every run multiplies the first 32 columns of the weights. + Mutant { name: "N4", find: "auto mB = W.slice(nt * 32, 0);", replace: "auto mB = W.slice(0, 0);", detector: Detector::Probe }, + // The control test is inverted. + Mutant { name: "N5", find: "if ((i & (ulong)g.ctrl) == (ulong)g.ctrl) {", replace: "if ((i & (ulong)g.ctrl) != (ulong)g.ctrl) {", detector: Detector::Certificate }, + // A sign of each real row's last term is wrong (as X6). + Mutant { name: "N6", find: "f0 - f1 + f2 - f3;", replace: "f0 - f1 + f2 + f3;", detector: Detector::Certificate }, + // The range guard never fires. + Mutant { name: "N7", find: "if (!fits32(o)) {", replace: "if (false) {", detector: Detector::Range }, + // The offset's square is removed at half its value. + Mutant { name: "N8", find: "- (1UL << 62));", replace: "- (1UL << 61));", detector: Detector::Probe }, +]; + +/// [`TENSOR_OPS_SOURCE`] with `m` applied. Panics unless `m.find` occurs +/// exactly once. +pub fn tensor_ops_mutate(m: &Mutant) -> String { + apply(TENSOR_OPS_SOURCE, m, "tensor-operation kernel") +} + fn apply(source: &str, m: &Mutant, what: &str) -> String { let n = source.matches(m.find).count(); assert_eq!(n, 1, "mutant {}: `{}` occurs {n} times in the {what}", m.name, m.find); @@ -325,7 +427,11 @@ mod tests { #[test] fn every_mutant_edits_its_kernel_once_and_names_distinct_edits() { - for (source, mutants, edit) in [(SOURCE, &MUTANTS[..], mutate as fn(&Mutant) -> String), (MATRIX_SOURCE, &MATRIX_MUTANTS[..], matrix_mutate)] { + for (source, mutants, edit) in [ + (SOURCE, &MUTANTS[..], mutate as fn(&Mutant) -> String), + (MATRIX_SOURCE, &MATRIX_MUTANTS[..], matrix_mutate), + (TENSOR_OPS_SOURCE, &TENSOR_OPS_MUTANTS[..], tensor_ops_mutate), + ] { let mut seen = std::collections::BTreeSet::new(); for m in mutants { let edited = edit(m); @@ -338,8 +444,26 @@ mod tests { } // The kernels' entry points are the same two. for entry in ["kernel void apply_gate(", "kernel void probe("] { - assert_eq!((SOURCE.matches(entry).count(), MATRIX_SOURCE.matches(entry).count()), (1, 1), "{entry}"); + assert_eq!((SOURCE.matches(entry).count(), MATRIX_SOURCE.matches(entry).count(), TENSOR_OPS_SOURCE.matches(entry).count()), (1, 1, 1), "{entry}"); } + let tid = tensor_ops_implementation_id(); + assert!(tid.starts_with("metal-tensor-ops-limbs8/") && tid.len() == "metal-tensor-ops-limbs8/".len() + 16, "{tid}"); + for m in &TENSOR_OPS_MUTANTS { + assert_ne!(tensor_ops_implementation_id_of(&tensor_ops_mutate(m)), tid, "mutant {}", m.name); + } + } + + #[test] + fn tensor_ops_need_family_1007_and_os_26() { + for (family, os, has) in [(1009, "27.0.1", true), (1010, "26.0", true), (1009, "15.6.1", false), (1006, "27.0.1", false), (1009, "", false)] { + assert_eq!(DeviceIdentity { family, os: os.into(), ..identity() }.has_tensor_ops(), has, "{family} {os}"); + } + // The scratch the shading language's struct lays out: 4096 + 8192 + 32768 bytes. + assert_eq!(TENSOR_OPS_SCRATCH_BYTES, 45056); + assert_eq!(Shape::TENSOR_OPS.groups(64, Shape::TENSOR_OPS.per_pair), 1); + assert_eq!(Shape::TENSOR_OPS.groups(65, Shape::TENSOR_OPS.per_pair), 2); + assert_eq!(Shape::TENSOR_OPS.groups(16, Shape::TENSOR_OPS.per_case), 1); + assert_eq!(Shape::TENSOR_OPS.groups(17, Shape::TENSOR_OPS.per_case), 2); } #[test] diff --git a/wai/wai-quantum/src/quantum_fabric/tensor_ops.metal b/wai/wai-quantum/src/quantum_fabric/tensor_ops.metal new file mode 100644 index 000000000..d0f466041 --- /dev/null +++ b/wai/wai-quantum/src/quantum_fabric/tensor_ops.metal @@ -0,0 +1,278 @@ +// wai-quantum fabric: the gate kernel for the platform's GPU whose every +// product is formed by the shading language's tensor operation (a bf16 matrix +// multiply accumulated in f32), the operation the platform maps to the GPU's +// matrix hardware where it has it. +// +// The same entry points as lanes64.metal, plus one buffer of scratch, so the +// member runs any of its kernels. As the matrix-instruction kernel for HIP +// devices does, every operand is non-negative, so the result does not depend +// on how an accumulator treats a negative term: each 32-bit value x is taken +// offset-binary, u = x + 2^31, and split into four unsigned limbs of 8 bits, +// which bf16 holds exactly. The product u_e·u_a is Σ c_s·2^(8s) with +// c_s = Σ_{i+j=s} e_i·a_j: seven coefficients of at most four limb products +// each, every limb product at most 255², every coefficient below 2^18 and +// never negative. So every value the operation forms is an integer f32 holds +// exactly, in whatever order it adds them (quantum_fabric::tensor_engine has +// the argument and the layout). The kernel adds the coefficients, shifted +// into place, in 64-bit integers, which gives u_e·u_a, and removes the offset +// exactly: e·a = u_e·u_a − 2^31·(e + a) − 2^62, modulo 2^64. +// +// The operation reads and writes tensors in device memory, so each +// threadgroup owns one slice of the scratch buffer: its tile of limbs +// (64 pairs × 32, the 16 limbs and 16 zeros, since the operation's K is a +// multiple of 32), the gate's limb weights (32 × 128: sixteen products of +// seven coefficients, and zeros), and the coefficients (64 × 128, f32). A tile +// is one multiply, in four 32-column runs. The rest is lanes64.metal's +// arithmetic on exact products, two threads a pair, each forming two rows. +// +// The threads of a pair read all four of its amplitude components before the +// tile's first barrier and write the rows that replace them after its last; a +// dispatch is one gate. So no thread reads what another writes within a +// dispatch. The threadgroup steps through its tiles together, so every +// barrier is reached by every thread. + +#include +#include +using namespace metal; +using namespace mpp::tensor_ops; + +// Bits per limb. Eight is what bf16 holds exactly; the radix is a mutant's +// one edit, and nine-bit limbs are not exact. +#define LIMB 8 + +// The offset that makes every operand non-negative: x + 2^31. +#define OFFSET 0x80000000u + +// Pairs (or probe rows) a tile: the multiply's M. +#define TILE 64 +// The multiply's K and N: 16 limbs and 112 coefficient columns, padded. +#define KP 32 +#define NP 128 + +typedef bfloat limb_t; + +// As lanes64.metal's. +struct Gate { + uint pairs_lo; + uint pairs_hi; + uint tbit; + uint ctrl; + int m[8]; + uint index; +}; + +// One threadgroup's scratch: the tile's limbs (TILE × KP), the limb weights +// (KP × NP), both bf16, and the coefficients (TILE × NP, f32), row-major. +struct Scratch { + limb_t a[TILE * KP]; + limb_t w[KP * NP]; + float c[TILE * NP]; +}; + +// The sixteen products, as (entry, component): rows 0r, 0i, 1r and 1i, four +// terms each, in the order lanes::exact_row adds them. Component c of a pair +// is a0.x, a0.y, a1.x, a1.y for c = 0, 1, 2, 3. +constant int TERM_ENTRY[16] = {0, 1, 2, 3, 0, 1, 2, 3, 4, 5, 6, 7, 4, 5, 6, 7}; +constant int TERM_COMP[16] = {0, 1, 2, 3, 1, 0, 3, 2, 0, 1, 2, 3, 1, 0, 3, 2}; + +// Limb j of x + 2^31: four unsigned limbs. +static inline int limb(int x, int j) { + return (int)((((uint)x ^ OFFSET) >> (LIMB * j)) & ((1u << LIMB) - 1)); +} + +// v as a matrix input. An integer of at most 256 in size has at most 8 +// significant bits, so bf16 holds it exactly. +static inline limb_t as_limb(int v) { + return limb_t(float(v)); +} + +// The exact fixed-point rounding of a product. +static inline long fx(long v) { + return (v + (1L << 29)) >> 30; +} + +// e·a from the operation's (e + 2^31)·(a + 2^31): the offset's terms removed, +// modulo 2^64. +static inline long unoffset(long raw, int e, int a) { + return (long)((ulong)raw - (((ulong)(long)e + (ulong)(long)a) << 31) - (1UL << 62)); +} + +static inline int low32(long v) { + return (int)(uint)(ulong)v; +} + +static inline bool fits32(long v) { + return v >= -2147483648L && v <= 2147483647L; +} + +// C = A · W by the tensor operation, in four 32-column runs. Called by every +// thread of the threadgroup. +static void multiply(device Scratch& s) { + auto A = tensor, tensor_inline>(s.a, dextents(KP, TILE)); + auto W = tensor, tensor_inline>(s.w, dextents(NP, KP)); + // Declared transposed: the operation then writes C[m][n] to c[m·NP + n]. + auto C = tensor, tensor_inline>(s.c, dextents(NP, TILE)); + constexpr auto desc = matmul2d_descriptor(TILE, 32, dynamic_length_v, false, false, false, matmul2d_descriptor::mode::multiply); + matmul2d> mm; + for (int nt = 0; nt < NP / 32; nt++) { + auto mA = A.slice(0, 0); + auto mB = W.slice(nt * 32, 0); + auto mC = C.slice(nt * 32, 0); + mm.run(mA, mB, mC); + } +} + +// Product q of row r of C: its seven coefficients, each an integer below +// 2^24 in size (so the conversion is exact), shifted into place and added in +// 64-bit two's complement. +static long assemble(device const float* C, int r, int q) { + ulong v = 0; + for (int s = 0; s < 7; s++) { + v += (ulong)(long)C[r * NP + q * 7 + s] << (LIMB * s); + } + return (long)v; +} + +// The limb weights of column `col`, row `k`, for product p's entry `e` +// (column p·7 + s takes limb s − j of it at row c·4 + j, c its component). +static inline limb_t weight(int k, int col, int c_of_p, int e) { + if (k >= 16 || col >= 112) { + return limb_t(0.0f); + } + int c = k / 4, j = k % 4, i = col % 7 - j; + return (c == c_of_p && i >= 0 && i <= 3) ? as_limb(limb(e, i)) : limb_t(0.0f); +} + +// One gate on a state of `2 * pairs` lanes, 64 pairs a tile, two threads a +// pair. The parameters are lanes64.metal's, with `amps` as words, and the +// scratch. +kernel void apply_gate( + device int* amps [[buffer(0)]], + constant Gate& g [[buffer(1)]], + device atomic_uint* report [[buffer(2)]], + device Scratch* scratch [[buffer(3)]], + uint tid [[thread_index_in_threadgroup]], + uint group [[threadgroup_position_in_grid]], + uint groups [[threadgroups_per_grid]], + uint width [[threads_per_threadgroup]]) +{ + if (width != 2 * TILE) { + return; + } + device Scratch& s = scratch[group]; + // The gate's limb weights and A's padding, once a threadgroup. + for (uint x = tid; x < KP * NP; x += width) { + int k = x / NP, col = x % NP, p = col / 7; + s.w[x] = col < 112 ? weight(k, col, TERM_COMP[p], g.m[TERM_ENTRY[p]]) : limb_t(0.0f); + } + for (uint x = tid; x < TILE * KP; x += width) { + s.a[x] = limb_t(0.0f); + } + threadgroup_barrier(mem_flags::mem_device); + ulong pairs = (ulong)g.pairs_lo | ((ulong)g.pairs_hi << 32); + ulong tiles = (pairs + TILE - 1) / TILE; + int r = tid / 2, side = tid % 2; + for (ulong tile = group; tile < tiles; tile += groups) { + ulong p = tile * TILE + r; + bool active = false; + ulong at[4] = {0, 0, 0, 0}; + int comp[4] = {0, 0, 0, 0}; + if (p < pairs) { + // The p-th index whose target bit is clear. + ulong low = p & ((ulong)g.tbit - 1); + ulong i = ((p - low) << 1) | low; + if ((i & (ulong)g.ctrl) == (ulong)g.ctrl) { + active = true; + ulong j = i | (ulong)g.tbit; + at[0] = 2 * i; + at[1] = 2 * i + 1; + at[2] = 2 * j; + at[3] = 2 * j + 1; + for (int c = 0; c < 4; c++) { + comp[c] = amps[at[c]]; + } + } + } + // This thread writes the limbs of components 2·side and 2·side + 1. + for (int h = 0; h < 2; h++) { + int c = side * 2 + h; + for (int jj = 0; jj < 4; jj++) { + s.a[r * KP + c * 4 + jj] = as_limb(limb(comp[c], jj)); + } + } + threadgroup_barrier(mem_flags::mem_device); + multiply(s); + threadgroup_barrier(mem_flags::mem_device); + if (active) { + // Rows 2·side and 2·side + 1 of the pair: row c is products 4c … 4c + 3. + for (int h = 0; h < 2; h++) { + int c = side * 2 + h; + long f[4]; + for (int t = 0; t < 4; t++) { + int q = c * 4 + t; + f[t] = fx(unoffset(assemble(s.c, r, q), g.m[TERM_ENTRY[q]], comp[TERM_COMP[q]])); + } + long f0 = f[0], f1 = f[1], f2 = f[2], f3 = f[3]; + long o = (c & 1) ? f0 + f1 + f2 + f3 : f0 - f1 + f2 - f3; + if (!fits32(o)) { + atomic_fetch_min_explicit(report, g.index, memory_order_relaxed); + } + amps[at[c]] = low32(o); + } + } + threadgroup_barrier(mem_flags::mem_device); + } +} + +// The arithmetic probe: for each case (a, b), the 64-bit product a·b from the +// tensor operation and its fx, as (high, low) words. Sixteen cases a tile: +// product q takes b_q's limbs against component q % 4, and row q / 4 carries +// a_{4r}, …, a_{4r+3} (quantum_fabric::tensor_engine::probe_weights); the +// tile's other rows are zero. +kernel void probe( + device const int2* cases [[buffer(0)]], + device uint4* out [[buffer(1)]], + constant uint2& count [[buffer(2)]], + device Scratch* scratch [[buffer(3)]], + uint tid [[thread_index_in_threadgroup]], + uint group [[threadgroup_position_in_grid]], + uint groups [[threadgroups_per_grid]], + uint width [[threads_per_threadgroup]]) +{ + if (width != 2 * TILE) { + return; + } + device Scratch& s = scratch[group]; + for (uint x = tid; x < TILE * KP; x += width) { + s.a[x] = limb_t(0.0f); + } + threadgroup_barrier(mem_flags::mem_device); + ulong n = (ulong)count.x | ((ulong)count.y << 32); + ulong tiles = (n + 15) / 16; + for (ulong tile = group; tile < tiles; tile += groups) { + ulong k0 = tile * 16; + for (uint x = tid; x < KP * NP; x += width) { + int k = x / NP, col = x % NP, p = col / 7; + int b = col < 112 && k0 + p < n ? cases[k0 + p].y : 0; + s.w[x] = col < 112 ? weight(k, col, p % 4, b) : limb_t(0.0f); + } + if (tid < 16) { + int r = tid / 4, c = tid % 4; + int a = k0 + tid < n ? cases[k0 + tid].x : 0; + for (int jj = 0; jj < 4; jj++) { + s.a[r * KP + c * 4 + jj] = as_limb(limb(a, jj)); + } + } + threadgroup_barrier(mem_flags::mem_device); + multiply(s); + threadgroup_barrier(mem_flags::mem_device); + if (tid < 16 && k0 + tid < n) { + int2 ab = cases[k0 + tid]; + long v = unoffset(assemble(s.c, tid / 4, tid), ab.y, ab.x); + ulong prod = (ulong)v; + ulong f = (ulong)fx(v); + out[k0 + tid] = uint4((uint)(prod >> 32), (uint)prod, (uint)(f >> 32), (uint)f); + } + threadgroup_barrier(mem_flags::mem_device); + } +} diff --git a/wai/wai-quantum/src/quantum_fabric_metal.rs b/wai/wai-quantum/src/quantum_fabric_metal.rs index 512083f39..0f60dc42c 100644 --- a/wai/wai-quantum/src/quantum_fabric_metal.rs +++ b/wai/wai-quantum/src/quantum_fabric_metal.rs @@ -16,6 +16,13 @@ //! returned, with its own implementation id, so its certificates are //! certificates of that path. A device below //! [`metal::MATRIX_MIN_FAMILY`] is [`MetalError::Unsupported`]. +//! - **The tensor-operation kernel** ([`MetalMember::tensor_ops`]): the same +//! device running [`metal::TENSOR_OPS_SOURCE`], whose every product is +//! formed by the shading language's tensor operation, which the platform +//! maps to the GPU's matrix hardware where it has it; probed before it is +//! returned, with its own implementation id. A device or operating system +//! without it ([`DeviceIdentity::has_tensor_ops`]) is +//! [`MetalError::Unsupported`]. //! - **Running** ([`MetalMember::run_program`]): the state stays in one //! buffer the device and the host share; one serial dispatch per gate, in //! op-log order, in one command buffer; then the report word and the state @@ -158,6 +165,31 @@ fn qubits_for(working_set: u64, max_buffer: u64) -> u8 { if lanes < 2 { 0 } else { (lanes.ilog2() as u8).min(MAX_QUBITS) } } +/// How a kernel is built and dispatched: its launch shape, the shading +/// language version it is compiled as, the scratch each threadgroup owns +/// (bytes, 0 for none) and the most threadgroups a dispatch takes. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[cfg_attr(not(target_os = "macos"), allow(dead_code))] +struct Build { + shape: Shape, + language: (u32, u32), + scratch: u64, + max_groups: u64, +} + +impl Build { + const LANES: Build = Build { shape: Shape::LANES, language: (3, 1), scratch: 0, max_groups: metal::GRID_MAX }; + const MATRIX: Build = Build { shape: Shape::MATRIX, language: (3, 1), scratch: 0, max_groups: metal::GRID_MAX }; + const TENSOR_OPS: Build = + Build { shape: Shape::TENSOR_OPS, language: (4, 0), scratch: metal::TENSOR_OPS_SCRATCH_BYTES, max_groups: metal::TENSOR_OPS_GROUP_MAX }; + + /// Threadgroups for `items` at `per_item` threads each. + #[cfg_attr(not(target_os = "macos"), allow(dead_code))] + fn groups(&self, items: u64, per_item: u64) -> u64 { + self.shape.groups(items, per_item).min(self.max_groups) + } +} + /// One gate as the kernels read it: 13 32-bit words ([`metal::HOST_LAYOUT`]). fn gate_words(pairs: u64, w: &GateWords, index: u32) -> [u32; 13] { let mut o = [0u32; 13]; @@ -417,7 +449,7 @@ impl MetalMember { /// require every mutant to be caught. #[doc(hidden)] pub fn with_source(&self, source: &str) -> Result { - let kernel = sys::compile(&self.gpu, source, Shape::LANES)?; + let kernel = sys::compile(&self.gpu, source, Build::LANES)?; Ok(MetalMember { kernel: Arc::new(kernel), implementation: metal::implementation_id_of(source), detail: self.discovered.clone(), ..self.clone() }) } @@ -441,11 +473,40 @@ impl MetalMember { metal::MATRIX_MIN_FAMILY ))); } - let kernel = sys::compile(&self.gpu, source, Shape::MATRIX)?; + let kernel = sys::compile(&self.gpu, source, Build::MATRIX)?; let detail = format!("{}; every product by the SIMD-group matrix multiply (bf16 limbs of 8 bits, f32 accumulation)", self.discovered); Ok(MetalMember { kernel: Arc::new(kernel), implementation: metal::matrix_implementation_id_of(source), detail, ..self.clone() }) } + /// The same device running the tensor-operation kernel + /// ([`metal::TENSOR_OPS_SOURCE`]), compiled for it and probed in this + /// call: the member id is the device's, the implementation id the + /// kernel's. A device or operating system without the operation + /// ([`DeviceIdentity::has_tensor_ops`]) is [`MetalError::Unsupported`]; + /// one whose products disagree with the CPU's is [`MetalError::Probe`]. + pub fn tensor_ops(&self) -> Result { + self.with_tensor_ops_source(metal::TENSOR_OPS_SOURCE)?.probed() + } + + /// The same device running `source` as the tensor-operation kernel, + /// unprobed, under the implementation id that text earns: for the + /// mutation tests. + #[doc(hidden)] + pub fn with_tensor_ops_source(&self, source: &str) -> Result { + if !self.identity.has_tensor_ops() { + return Err(MetalError::Unsupported(format!( + "GPU family {} under operating system {}; the tensor operation needs family {} or later and operating system {} or later", + self.identity.family, + self.identity.os, + metal::MATRIX_MIN_FAMILY, + metal::TENSOR_OPS_MIN_OS + ))); + } + let kernel = sys::compile(&self.gpu, source, Build::TENSOR_OPS)?; + let detail = format!("{}; every product by the tensor operation (bf16 offset-binary limbs of 8 bits, f32 accumulation)", self.discovered); + Ok(MetalMember { kernel: Arc::new(kernel), implementation: metal::tensor_ops_implementation_id_of(source), detail, ..self.clone() }) + } + /// This member, once its kernel's arithmetic agrees with the CPU's on the /// probe's pairs. fn probed(self) -> Result { @@ -482,9 +543,9 @@ impl Runner for MetalMember { /// kernels, and the calls on them. #[cfg(target_os = "macos")] mod sys { - use super::{Discovery, MetalError, MetalMember, Refusal, qubits_for}; + use super::{Build, Discovery, MetalError, MetalMember, Refusal, qubits_for}; use crate::quantum_energy::EnergyClass; - use crate::quantum_fabric::metal::{self, DeviceIdentity, Shape}; + use crate::quantum_fabric::metal::{self, DeviceIdentity}; use crate::quantum_fabric::{RunMeter, Window}; use objc2::rc::{Retained, autoreleasepool}; use objc2::runtime::ProtocolObject; @@ -502,11 +563,12 @@ mod sys { queue: Retained>, } - /// A kernel compiled for one device: its two entry points and its shape. + /// A kernel compiled for one device: its two entry points and how it is + /// dispatched. pub struct Kernel { apply: Retained>, probe: Retained>, - shape: Shape, + build: Build, } type Buffer = Retained>; @@ -559,7 +621,7 @@ mod sys { let refuse = |error: MetalError| Refusal { device: name.clone(), key: key.clone(), error }; let queue = device.newCommandQueue().ok_or_else(|| refuse(MetalError::Device("the device gave no command queue".into())))?; let gpu = Arc::new(Gpu { device, queue }); - let kernel = compile(&gpu, metal::SOURCE, Shape::LANES).map_err(refuse)?; + let kernel = compile(&gpu, metal::SOURCE, Build::LANES).map_err(refuse)?; let max_qubits = qubits_for(identity.working_set, identity.max_buffer); let detail = format!( "metal; device {name:?}; GPU family {}; registry id {:x}; unified memory {}; working set {} bytes; buffers up to {} bytes; operating system {}; up to {max_qubits} qubits; a metered run's window runs from the first upload to the last readback", @@ -579,12 +641,16 @@ mod sys { member.probed().map_err(refuse) } - /// `source` compiled for the device ([`metal::COMPILE_OPTIONS`]), its two - /// entry points as pipelines that take `shape.threads` a threadgroup. - pub fn compile(gpu: &Gpu, source: &str, shape: Shape) -> Result { + /// `source` compiled for the device as `build.language` with safe math + /// ([`metal::COMPILE_OPTIONS`], [`metal::TENSOR_OPS_COMPILE_OPTIONS`]), + /// its two entry points as pipelines that take `build.shape.threads` a + /// threadgroup. + pub fn compile(gpu: &Gpu, source: &str, build: Build) -> Result { + let shape = build.shape; + let (major, minor) = build.language; autoreleasepool(|_| { let options = MTLCompileOptions::new(); - options.setLanguageVersion(MTLLanguageVersion::Version3_1); + options.setLanguageVersion(MTLLanguageVersion(((major << 16) + minor) as usize)); options.setMathMode(MTLMathMode::Safe); let library = gpu.device.newLibraryWithSource_options_error(&NSString::from_str(source), Some(&options)).map_err(|e| MetalError::Compile(ns_error(&e)))?; let pipeline = |entry: &str| -> Result>, MetalError> { @@ -596,7 +662,7 @@ mod sys { } Ok(p) }; - Ok(Kernel { apply: pipeline("apply_gate")?, probe: pipeline("probe")?, shape }) + Ok(Kernel { apply: pipeline("apply_gate")?, probe: pipeline("probe")?, build }) }) } @@ -663,6 +729,16 @@ mod sys { unsafe { encoder.setBuffer_offset_atIndex(Some(b), 0, index) }; } + /// The scratch a dispatch of `groups` threadgroups needs at buffer index 3, + /// if the kernel takes any. + fn scratch(gpu: &Gpu, k: &Kernel, groups: u64) -> Result, MetalError> { + if k.build.scratch == 0 { + return Ok(None); + } + let bytes = usize::try_from(groups.saturating_mul(k.build.scratch)).map_err(|_| MetalError::Memory("the scratch".into()))?; + buffer(gpu, bytes).map(Some) + } + /// The probe of `cases` (interleaved (a, b) pairs): four words a case. pub fn probe(gpu: &Gpu, k: &Kernel, cases: &[i32]) -> Result, MetalError> { autoreleasepool(|_| { @@ -670,12 +746,17 @@ mod sys { let input = upload(gpu, cases)?; let out = buffer(gpu, n * 16)?; let count = [n as u32, ((n as u64) >> 32) as u32]; + let groups = k.build.groups(n as u64, k.build.shape.per_case); + let scratch = scratch(gpu, k, groups)?; submit(gpu, |e| { e.setComputePipelineState(&k.probe); bind(e, &input, 0); bind(e, &out, 1); bytes(e, &count, 2); - dispatch(e, k.shape.groups(n as u64, k.shape.per_case), k.shape.threads); + if let Some(s) = &scratch { + bind(e, s, 3); + } + dispatch(e, groups, k.build.shape.threads); })?; read(&out, n * 4) }) @@ -691,14 +772,18 @@ mod sys { let (state, bad) = autoreleasepool(|_| -> Result<(Vec, u32), MetalError> { let amps = upload(gpu, &init)?; let report = upload(gpu, &[u32::MAX])?; - let groups = k.shape.groups(pairs, k.shape.per_pair); + let groups = k.build.groups(pairs, k.build.shape.per_pair); + let scratch = scratch(gpu, k, groups)?; submit(gpu, |e| { e.setComputePipelineState(&k.apply); bind(e, &s, 0); bind(e, &report, 2); + if let Some(s) = &scratch { + bind(e, s, 3); + } for w in words { bytes(e, w, 1); - dispatch(e, groups, k.shape.threads); + dispatch(e, groups, k.build.shape.threads); } })?; let bad = read::(&report, 1)?[0]; @@ -712,10 +797,9 @@ mod sys { /// no values. #[cfg(not(target_os = "macos"))] mod sys { - use super::{Discovery, MetalError}; + use super::{Build, Discovery, MetalError}; use crate::quantum_energy::EnergyClass; use crate::quantum_fabric::RunMeter; - use crate::quantum_fabric::metal::Shape; pub enum Gpu {} pub enum Kernel {} @@ -724,7 +808,7 @@ mod sys { Discovery { found: Vec::new(), refused: Vec::new(), unavailable: Some("this platform has no Metal; the member runs on macOS".into()) } } - pub fn compile(gpu: &Gpu, _: &str, _: Shape) -> Result { + pub fn compile(gpu: &Gpu, _: &str, _: Build) -> Result { match *gpu {} } diff --git a/wai/wai-quantum/tests/metal_tensor_ops.rs b/wai/wai-quantum/tests/metal_tensor_ops.rs new file mode 100644 index 000000000..a5e33bd2f --- /dev/null +++ b/wai/wai-quantum/tests/metal_tensor_ops.rs @@ -0,0 +1,446 @@ +//! The Metal member's tensor-operation kernel on the devices of the machine +//! that runs this: every product the gate needs formed by the shading +//! language's tensor operation, from offset-binary 8-bit limbs in bf16 +//! accumulated in f32, held to the same tests as the other kernels +//! (`metal_member`). +//! +//! The machine declares its devices in `WAI_QUANTUM_METAL_EXPECT`, as for +//! `metal_member`: their keys, comma-separated and sorted, or `none`. Every +//! declared device with the tensor operation (GPU family 1007 or later, under +//! operating system 26 or later) must run the kernel and pass its probe; one +//! without must be refused as unsupported. Each test prints what it established and, when +//! `WAI_QUANTUM_METAL_REPORT` names a file, appends the same lines to it. +//! +//! ```text +//! WAI_QUANTUM_METAL_EXPECT=gpu-msl-1009 \ +//! cargo test --features quantum_fabric_metal --test metal_tensor_ops -- --test-threads=1 +//! ``` +#![cfg(not(target_arch = "wasm32"))] + +use std::io::Write; +use std::sync::OnceLock; +use wai_quantum::ed25519_dalek::SigningKey; +use wai_quantum::quantum::{BaseGate, Circuit, Gate, Gateset, MAX_QUBITS, QuantumError}; +use wai_quantum::quantum_fabric::metal::{self, TENSOR_OPS_MUTANTS}; +use wai_quantum::quantum_fabric::lanes::{self, GateWords, Lane}; +use wai_quantum::quantum_fabric::plan::{Program, Step}; +use wai_quantum::quantum_fabric::wgsl::{Detector, Mutant}; +use wai_quantum::quantum_fabric::{ + Corpus, CpuReference, FABRIC_DEFAULT, Fabric, FabricClass, FabricReceipt, MemberError, Parity, Presence, REFERENCE, apply_kernel, detect_local, place, ready, +}; +use wai_quantum::quantum_fabric_metal::{Discovery, MetalError, MetalMember, discover}; + +/// The CPU's corpus digests: random-v1 (seed 7, 200 cases) and fabric-v1 +/// (seed 7, 200 + 24 cases), per gateset. A bit-exact member reproduces them. +const PINS: [(Gateset, &str, &str); 2] = [ + (Gateset::V1, "d18454c639acfd9c7b66450f4ded1aac4ab46cc10a7bedbcf4958a96e204fb36", "41c726a0ac80f7defef92e15ecf893bdb7a3076b5e7f670fa0b7142e5c033df7"), + (Gateset::V2, "3dbde276c1095bf827faec71d5980e050f8a55a5ba7d97cd3bbdaec3f43a13b2", "9a84b55af68f1a8b07390b234cf562dd8d85cef3185a697ab85690c8790a7809"), +]; + +const RANDOM_V1: Corpus = Corpus::Random { seed: 7, cases: 200 }; + +static DISCOVERY: OnceLock = OnceLock::new(); +static TENSOR_OPS: OnceLock> = OnceLock::new(); + +fn discovery() -> &'static Discovery { + DISCOVERY.get_or_init(discover) +} + +/// Prints `line` and appends it to the file `WAI_QUANTUM_METAL_REPORT` names. +fn report(line: &str) { + println!("{line}"); + if let Some(path) = std::env::var_os("WAI_QUANTUM_METAL_REPORT") { + let mut f = std::fs::OpenOptions::new().create(true).append(true).open(&path).expect("the report file opens"); + writeln!(f, "{line}").expect("the report line is written"); + } +} + +/// The declared device set, as written, if it is exactly what was found; +/// otherwise why not, with what was found and refused and how to declare it. +fn declared(d: &Discovery) -> Result { + let mut found: Vec<&str> = d.found.iter().map(|m| m.key()).collect(); + found.sort_unstable(); + let found_list = if found.is_empty() { "none".to_string() } else { found.join(",") }; + let refused: Vec = d.refused.iter().map(|r| format!("{} ({:?}): {}", r.key, r.device, r.error)).collect(); + let refused_list = if refused.is_empty() { "none".to_string() } else { refused.join("; ") }; + let unavailable = d.unavailable.as_deref().unwrap_or("no"); + let how = format!("found: {found_list}; refused: {refused_list}; unavailable: {unavailable}. Declare the devices this machine must have: WAI_QUANTUM_METAL_EXPECT={found_list}"); + let expect = match std::env::var("WAI_QUANTUM_METAL_EXPECT") { + Ok(v) if !v.is_empty() => v, + _ => return Err(format!("WAI_QUANTUM_METAL_EXPECT is not set. {how}")), + }; + let keys: Vec<&str> = if expect == "none" { Vec::new() } else { expect.split(',').collect() }; + let well_formed = keys.windows(2).all(|w| w[0] <= w[1]) && keys.iter().all(|k| !k.is_empty() && *k != "none" && k.trim() == *k); + if !well_formed { + return Err(format!("WAI_QUANTUM_METAL_EXPECT={expect:?} is not a sorted, comma-separated list of keys, nor none. {how}")); + } + for k in &keys { + if !found.contains(k) + && let Some(r) = d.refused.iter().find(|r| r.key == *k) + { + return Err(format!("the declared device {k} was refused: {}. {how}", r.error)); + } + } + if found != keys { + return Err(format!("the devices found are not the declared ones ({expect}). {how}")); + } + Ok(expect) +} + +/// The tensor-operation members: one per declared device with the operation, once +/// the declared devices are exactly the ones found. +fn members() -> &'static [MetalMember] { + let d = discovery(); + if let Err(why) = declared(d) { + panic!("{why}"); + } + TENSOR_OPS.get_or_init(|| { + d.found.iter().filter(|m| m.identity().has_tensor_ops()).map(|m| m.tensor_ops().unwrap_or_else(|e| panic!("{}: the tensor-operation kernel: {e}", m.key()))).collect() + }) +} + +fn hex(b: &[u8]) -> String { + b.iter().map(|x| format!("{x:02x}")).collect() +} + +fn next(s: &mut u64) -> u64 { + *s = s.wrapping_add(0x9E37_79B9_7F4A_7C15); + let mut z = *s; + z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9); + z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB); + z ^ (z >> 31) +} + +/// A value in `[-bound, bound]`. +fn within(s: &mut u64, bound: i64) -> i32 { + ((next(s) % (2 * bound as u64 + 1)) as i64 - bound) as i32 +} + +/// A random circuit of `ops` gates on `n` qubits: every base gate, up to three +/// distinct controls, `P(1..=32)`. +fn random_circuit(n: u8, ops: usize, seed: u64, gs: Gateset) -> Circuit { + let bases = [BaseGate::I, BaseGate::X, BaseGate::Y, BaseGate::Z, BaseGate::H, BaseGate::S, BaseGate::Sdg, BaseGate::T, BaseGate::Tdg, BaseGate::P]; + let mut s = seed; + let mut c = Circuit::with_gateset(n, gs); + for _ in 0..ops { + let base = bases[(next(&mut s) % bases.len() as u64) as usize]; + let target = (next(&mut s) % u64::from(n)) as u8; + let mut controls = Vec::new(); + for _ in 0..next(&mut s) % 4 { + let q = (next(&mut s) % u64::from(n)) as u8; + if q != target && !controls.contains(&q) { + controls.push(q); + } + } + let param = if base == BaseGate::P { 1 + (next(&mut s) % 32) as u16 } else { 0 }; + c.ops.push(Gate { base, controls, target, param }); + } + c +} + +/// Five circuits the native simulator refuses, in each gateset. +fn refused_circuits() -> Vec { + let gate = |base, controls: Vec, target, param| Gate { base, controls, target, param }; + let bad = [ + (1, vec![gate(BaseGate::H, vec![], 0, 0), gate(BaseGate::P, vec![], 0, 0)]), + (1, vec![gate(BaseGate::P, vec![], 0, 33)]), + (1, vec![gate(BaseGate::X, vec![], 3, 0)]), + (2, vec![gate(BaseGate::X, vec![0], 0, 0)]), + (2, vec![gate(BaseGate::X, vec![5], 1, 0)]), + ]; + let mut out = Vec::new(); + for gs in Gateset::ALL { + for (n, ops) in bad.clone() { + let mut c = Circuit::with_gateset(n, gs); + c.ops = ops; + out.push(c); + } + } + out +} + +/// `Z` on qubit 0: on `[(0, 0), (i32::MIN, 0)]` its exact result is `2^31`, +/// which leaves 32 bits. +fn z_words() -> GateWords { + lanes::gate_words(&Gate { base: BaseGate::Z, controls: vec![], target: 0, param: 0 }, Gateset::V2) +} + +const Z_STATE: [Lane; 2] = [[0, 0], [i32::MIN, 0]]; + +#[test] +fn c0_devices_are_declared() { + let d = discovery(); + match declared(d) { + Ok(expect) => report(&format!("metal tensor-ops devices: as declared ({expect})")), + Err(why) => panic!("{why}"), + } + for m in &d.found { + if !m.identity().has_tensor_ops() { + assert!(matches!(m.tensor_ops(), Err(MetalError::Unsupported(_))), "{}", m.key()); + report(&format!("metal tensor-ops unsupported: key={} family {} os {}", m.key(), m.identity().family, m.identity().os)); + } + } + for t in members() { + assert_eq!(t.class(), FabricClass::Gpu); + assert!(t.id().starts_with("local/gpu-msl-"), "{}", t.id()); + assert_eq!(t.implementation(), metal::tensor_ops_implementation_id()); + assert_ne!(t.implementation(), metal::implementation_id()); + assert_eq!(t.identity().member_id(), t.id()); + let lanes64 = d.found.iter().find(|m| m.id() == t.id()).expect("the device's member"); + assert_eq!((lanes64.implementation(), lanes64.max_qubits()), (metal::implementation_id().as_str(), t.max_qubits())); + assert!(t.max_qubits() >= 20, "{}: holds only {} qubits", t.key(), t.max_qubits()); + assert!(t.detail().starts_with(lanes64.detail()) && t.detail().contains("tensor operation"), "{}", t.detail()); + report(&format!("metal tensor-ops member: key={} id={} implementation={} {}", t.key(), t.id(), t.implementation(), t.detail())); + } +} + +#[test] +fn c1_probe_matches_lanes() { + for m in members() { + let pairs = lanes::probe_pairs(1 << 16, 0x0cda); + let out = m.probe(&pairs).unwrap_or_else(|e| panic!("{}: {e}", m.key())); + assert_eq!(lanes::first_mismatch(&pairs, &out), None, "{}", m.key()); + report(&format!("metal tensor-ops probe: key={} {} pairs agree", m.key(), pairs.len())); + } +} + +#[test] +fn c2_kernel_matches_run_words() { + for m in members() { + let mut s = 0x6b65_726e_656c_u64; + let (mut flagged, mut clean) = (0, 0); + for case in 0..2000 { + let n = 1 + (next(&mut s) % 12) as u32; + // Three regimes: the whole 32-bit range (nearly always out of + // range), small states and entries (never), and states and + // entries up to 2^30 (either). + let mode = case % 3; + let value = |s: &mut u64| -> i32 { + match mode { + 0 => next(s) as i32, + 1 => within(s, 1 << 22), + _ => within(s, 1 << 30), + } + }; + let entry = |s: &mut u64| -> u32 { + match mode { + 0 => next(s) as u32, + 1 => within(s, 1 << 29) as u32, + _ => within(s, 1 << 30) as u32, + } + }; + let state: Vec = (0..1usize << n).map(|_| [value(&mut s), value(&mut s)]).collect(); + let gates: Vec = (0..1 + next(&mut s) % 8) + .map(|_| { + let t = (next(&mut s) % u64::from(n)) as u32; + let ctrl = (0..n).filter(|&b| b != t && next(&mut s).is_multiple_of(4)).fold(0u32, |acc, b| acc | 1 << b); + let mut w = [0u32; 8]; + w.iter_mut().for_each(|x| *x = entry(&mut s)); + GateWords { m: w, tbit: 1 << t, ctrl } + }) + .collect(); + let want = lanes::run_words(&state, &gates).expect("valid shapes"); + let got = m.run_lanes(&state, &gates).unwrap_or_else(|e| panic!("{}: case {case}: {e}", m.key())); + assert_eq!(got.out_of_range, want.out_of_range, "{}: case {case}, n = {n}", m.key()); + assert!(got.amps == want.amps, "{}: case {case}, n = {n}: the lanes differ", m.key()); + if want.out_of_range.is_some() { + flagged += 1; + } else { + clean += 1; + } + } + assert!(flagged >= 100 && clean >= 100, "{}: {flagged} flagged, {clean} clean", m.key()); + report(&format!("metal tensor-ops kernel: key={} ok", m.key())); + } +} + +#[test] +fn c3_certificates_reproduce_the_pins() { + let ms = members(); + for (gs, random_pin, fabric_pin) in PINS { + let cpu = lanes::certify_corpus(gs, "local/test", RANDOM_V1); + assert_eq!(hex(&cpu.corpus), random_pin, "{gs:?}"); + let other = Corpus::Random { seed: 8, cases: 300 }; + let cpu_other = lanes::certify_corpus(gs, "local/test", other); + for m in ms { + for (corpus, cases, pin) in [(RANDOM_V1, 200, random_pin), (FABRIC_DEFAULT, 224, fabric_pin)] { + let c = m.certify(gs, corpus).unwrap_or_else(|e| panic!("{} {gs:?} {}: {e}", m.key(), corpus.name())); + assert_eq!(c.parity, Parity::BitExact, "{} {gs:?} {}", m.key(), corpus.name()); + assert_eq!(c.cases, cases); + assert_eq!(hex(&c.corpus), pin, "{} {gs:?} {}", m.key(), corpus.name()); + assert_eq!((c.kernel.as_str(), c.implementation.as_str(), c.member.as_str()), (apply_kernel(gs), m.implementation(), m.id())); + assert_eq!(c.generator, corpus.generator(gs)); + report(&format!( + "metal tensor-ops certified: key={} member={} kernel={} corpus={} implementation={} parity=bit-exact cases={} digest={}", + m.key(), + m.id(), + c.kernel, + corpus.name(), + c.implementation, + c.cases, + hex(&c.corpus) + )); + } + // A corpus with no pin: the member and the CPU agree on it. + let c = m.certify(gs, other).unwrap_or_else(|e| panic!("{} {gs:?} seed 8: {e}", m.key())); + assert_eq!((c.parity, c.cases), (Parity::BitExact, 300), "{} {gs:?}", m.key()); + assert_eq!(c.corpus, cpu_other.corpus, "{} {gs:?}", m.key()); + } + } +} + +#[test] +fn c4_wide_circuits_match_native() { + for m in members() { + let widest = m.max_qubits().min(22); + let mut s = 0x0077_6964_6500_u64; + for (k, n) in [12u8, 16, 18, 20, widest].into_iter().enumerate() { + for gs in Gateset::ALL { + let c = random_circuit(n, 40, 0x1000 + k as u64, gs); + let start = (next(&mut s) as usize) % (1usize << n); + let want = c.simulate_from(start).expect("the native run"); + let got = m.simulate_from(&c, start).unwrap_or_else(|e| panic!("{}: n = {n} {gs:?}: {e}", m.key())); + assert!(got == want, "{}: n = {n} {gs:?}: the states differ", m.key()); + } + } + // The transforms the conformance corpus carries. + for n in [3u8, 5, 8, 12] { + let qft = Circuit::qft(n); + assert!(m.simulate(&qft).unwrap() == qft.simulate().unwrap(), "{}: QFT-{n}", m.key()); + } + report(&format!("metal tensor-ops wide: key={} up to {widest} qubits match native", m.key())); + } +} + +#[test] +fn c5_refusals_are_typed() { + for m in members() { + for c in refused_circuits() { + let native = c.simulate_from(0).unwrap_err(); + assert_eq!(m.simulate_from(&c, 0).unwrap_err(), MetalError::Circuit(native), "{c:?}"); + } + assert_eq!(m.simulate(&Circuit::new(MAX_QUBITS + 1)).unwrap_err(), MetalError::Circuit(QuantumError::TooManyQubits(MAX_QUBITS + 1))); + let small = m.with_max_qubits(4); + assert_eq!(small.simulate(&Circuit::new(5)).unwrap_err(), MetalError::TooWide { n_qubits: 5, max: 4 }); + let wide = Program { n_qubits: 5, start: 0, gates: Vec::new(), steps: Vec::new() }; + assert_eq!(small.run_program(&wide).unwrap_err(), MetalError::TooWide { n_qubits: 5, max: 4 }); + // A result that leaves 32 bits is reported at its op, and the lane + // holds its low 32 bits. + let r = m.run_lanes(&Z_STATE, &[z_words()]).unwrap(); + assert_eq!(r.out_of_range, Some(0), "{}", m.key()); + assert_eq!(r.amps[1][0] as u32, 0x8000_0000); + // Malformed words, a state the host cannot lay out and a tile step are + // host defects, never results. + let bad = GateWords { tbit: 3, ..z_words() }; + assert!(matches!(m.run_lanes(&[[0, 0]; 4], &[bad]), Err(MetalError::Internal(_)))); + assert!(matches!(m.run_lanes(&[[0, 0]; 3], &[z_words()]), Err(MetalError::Internal(_)))); + let tiled = Program { n_qubits: 2, start: 0, gates: Vec::new(), steps: vec![Step::Tile { first: 0, count: 1, size: 2 }] }; + assert!(matches!(m.run_program(&tiled), Err(MetalError::Internal(_)))); + // The refusals that may still fall back to another member, do. + assert!(MemberError::from(MetalError::TooWide { n_qubits: 5, max: 4 }).falls_back()); + report(&format!("metal tensor-ops refusals: key={} typed", m.key())); + } +} + +/// Whether a member running `mutant`'s text is caught by the mutant's +/// detector. +fn caught(mm: &MetalMember, mutant: &Mutant, key: &str) -> bool { + match mutant.detector { + Detector::Probe => { + let edges = lanes::probe_pairs(0, 0); + let out = mm.probe(&edges).expect("the mutant's probe runs"); + let differing = edges.iter().zip(&out).filter(|&(&(a, b), o)| *o != (lanes::mul_wide(a, b), lanes::fxmul_wide(a, b))).count(); + println!("metal tensor-ops mutant {} on {key}: {differing} of 484 edge pairs differ", mutant.name); + differing >= 1 + } + Detector::Certificate => Gateset::ALL.into_iter().all(|gs| match mm.certify(gs, RANDOM_V1) { + Err(e) => { + println!("metal tensor-ops mutant {} on {key}: {gs:?} certification refused: {e}", mutant.name); + true + } + Ok(c) => { + let cpu = lanes::certify_corpus(gs, "local/test", RANDOM_V1); + c.parity != Parity::BitExact || c.corpus != cpu.corpus + } + }), + Detector::Range => mm.run_lanes(&Z_STATE, &[z_words()]).map(|r| r.out_of_range) != Ok(Some(0)), + other => panic!("no tensor-operation mutant uses the {other:?} detector"), + } +} + +#[test] +fn c6_mutants_are_killed() { + assert_eq!(lanes::probe_pairs(0, 0).len(), 484); + for m in members() { + let mut killed = 0; + for mutant in &TENSOR_OPS_MUTANTS { + let mm = m.with_tensor_ops_source(&metal::tensor_ops_mutate(mutant)).unwrap_or_else(|e| panic!("{}: mutant {} does not compile: {e}", m.key(), mutant.name)); + assert_ne!(mm.implementation(), m.implementation()); + assert!(caught(&mm, mutant, m.key()), "{}: mutant {} survived", m.key(), mutant.name); + killed += 1; + } + assert_eq!(killed, TENSOR_OPS_MUTANTS.len()); + report(&format!("metal tensor-ops mutants: key={} killed {killed} of {}", m.key(), TENSOR_OPS_MUTANTS.len())); + } +} + +#[test] +fn c7_receipt_placement_and_the_fabric() { + for m in members() { + let member = m.member(FABRIC_DEFAULT).unwrap_or_else(|e| panic!("{}: {e}", m.key())); + assert_eq!((member.id.as_str(), member.class, member.presence), (m.id(), FabricClass::Gpu, Presence::Detected)); + assert_eq!(member.kernels.len(), 2); + for (entry, gs) in member.kernels.iter().zip(Gateset::ALL) { + assert_eq!((entry.kernel.as_str(), entry.implementation.as_str()), (apply_kernel(gs), m.implementation())); + assert!(entry.energy.is_none()); + assert_eq!(entry.certificate.parity, Parity::BitExact, "{gs:?}"); + } + // A receipt over a member's run. + let qft = Circuit::qft(5); + let state = m.simulate(&qft).unwrap(); + assert_eq!(state, qft.simulate().unwrap()); + let entry = member.kernels.iter().find(|k| k.kernel == apply_kernel(qft.gateset)).unwrap(); + let signer = SigningKey::from_bytes(&[7; 32]); + let receipt = FabricReceipt::seal(&signer, "local/test", state.statevector_hash(), &member, entry, 0, None); + assert!(receipt.verify()); + assert!(receipt.certificate_matches(&entry.certificate)); + // Placement: beside the CPU reference with no energy figures, the + // reference; alone, the member. + let local = detect_local(); + for gs in Gateset::ALL { + let both = [local.clone(), member.clone()]; + let (placed, k) = place(&both, apply_kernel(gs), &Parity::BitExact).unwrap(); + assert_eq!((placed.id.as_str(), k.implementation.as_str()), (local.id.as_str(), REFERENCE)); + let (placed, k) = place(std::slice::from_ref(&member), apply_kernel(gs), &Parity::BitExact).unwrap(); + assert_eq!((placed.id.as_str(), k.implementation.as_str()), (m.id(), m.implementation())); + } + // On the fabric: the member alone runs the circuit; a pairing with + // another member's entries is refused. + let mut fabric = Fabric::new(); + fabric.add(member.clone(), m); + let run = ready(fabric.run(&qft, 0, &Parity::BitExact, &signer, "local/test")).expect("resolves").expect("runs"); + assert_eq!(run.state, state); + assert_eq!(run.receipt.member_id, m.id()); + assert!(run.receipt.verify()); + let mut mismatched = Fabric::new(); + mismatched.add(local.clone(), m); + assert!(ready(mismatched.run(&qft, 0, &Parity::BitExact, &signer, "local/test")).expect("resolves").is_err()); + let mut both = Fabric::new(); + both.add(local.clone(), CpuReference).add(member.clone(), m); + let run = ready(both.run(&qft, 0, &Parity::BitExact, &signer, "local/test")).expect("resolves").expect("runs"); + assert_eq!(run.state, state); + report(&format!("metal tensor-ops receipt: key={} ok", m.key())); + } +} + +#[test] +fn c8_repeated_runs_agree() { + for m in members() { + let c = random_circuit(16, 60, 0xabc, Gateset::V2); + let first = m.simulate(&c).unwrap(); + for _ in 0..5 { + assert!(m.simulate(&c).unwrap() == first, "{}", m.key()); + } + report(&format!("metal tensor-ops repeat: key={} ok", m.key())); + } +} -- 2.54.0 (Apple Git-157)