From 4b3b3d1728b95b5acc2a7e9300303d9da4d9a76d Mon Sep 17 00:00:00 2001 From: Kevaundray Wedderburn Date: Thu, 8 Oct 2026 11:35:55 +0100 Subject: [PATCH 1/5] Fuse ARM dense-grid final extension with mixed accumulation --- crates/pcs/src/whir/sumcheck/first_pass.rs | 137 ++++++++++++++++++++- 1 file changed, 131 insertions(+), 6 deletions(-) diff --git a/crates/pcs/src/whir/sumcheck/first_pass.rs b/crates/pcs/src/whir/sumcheck/first_pass.rs index cfca183bf..7e9ad83d2 100644 --- a/crates/pcs/src/whir/sumcheck/first_pass.rs +++ b/crates/pcs/src/whir/sumcheck/first_pass.rs @@ -52,8 +52,13 @@ const ROW: usize = 8; /// a multilinear's `inf` is the sum of its `0` and `1`. #[inline(always)] fn extend_grid(grid: &mut [T; GRID], add: impl Fn(&T, &T) -> T) { + extend_grid_digits::(grid, R, add); +} + +#[inline(always)] +fn extend_grid_digits(grid: &mut [T; GRID], digits: usize, add: impl Fn(&T, &T) -> T) { let mut stride = 1; - for i in 0..R { + for i in 0..digits { for &high in &LANE_IN_GRID[..1 << (R - 1 - i)] { let base = 3 * stride * high; for at in base..base + stride { @@ -64,6 +69,28 @@ fn extend_grid(grid: &mut [T; GRID], add: impl Fn(&T, & } } +/// Finish each row's extension at the point where its products are consumed. +#[inline(always)] +fn accumulate_grid( + fg: &mut [[u64; ROW]; GRID], + bg: &mut [WeightRow; GRID], + acc: &mut [ProductRow; GRID], +) { + #[cfg(all(target_arch = "aarch64", target_feature = "aes"))] + if R > 0 { + extend_grid_digits::<_, R>(fg, R - 1, |a, b| std::array::from_fn(|i| a[i] ^ b[i])); + extend_grid_digits::<_, R>(bg, R - 1, WeightRow::add); + // SAFETY: PMULL is enabled at compile time; the helper stays within the grid. + unsafe { lanes::finish_grid::(fg, bg, acc) }; + return; + } + extend_grid::<_, R>(fg, |a, b| std::array::from_fn(|i| a[i] ^ b[i])); + extend_grid::<_, R>(bg, WeightRow::add); + for (a, (k, w)) in acc.iter_mut().zip(fg.iter().zip(bg.iter())).take(3usize.pow(R as u32)) { + a.mul_acc(w, k); + } +} + /// A row of `ROW` weights, laid out so that both qwords of every 128-bit lane meet a word: /// lane `j` of `lo` is `[c0, c1]` of weight `2j`, of `hi` the same of weight `2j+1`, and /// `c2` holds every weight's last coefficient in order. @@ -300,6 +327,8 @@ mod avx512 { ))] #[cfg_attr(all(target_feature = "vpclmulqdq", target_feature = "avx512f"), allow(dead_code))] mod lanes { + #[cfg(target_arch = "aarch64")] + use super::GRID; use super::{F64, F192, LaneWeight, ProductRow, ROW, WeightRow}; #[cfg(target_arch = "aarch64")] use core::arch::aarch64::*; @@ -513,6 +542,54 @@ mod lanes { } } + /// Fuse the final extension digit with its three product updates, without storing the infinity rows. + #[cfg(all(target_arch = "aarch64", target_feature = "aes"))] + #[inline(always)] + pub(super) unsafe fn finish_grid( + fg: &[[u64; ROW]; GRID], + bg: &[WeightRow; GRID], + acc: &mut [ProductRow; GRID], + ) { + #[inline(always)] + unsafe fn update(acc: &mut ProductRow, c: usize, k: Neon, lo: Neon, hi: Neon, c2: Neon) { + // SAFETY: c names a complete pair in each row, and PMULL is enabled by the module's caller. + unsafe { + let s0 = acc.0[0].as_mut_ptr().add(c); + let s1 = acc.0[1].as_mut_ptr().add(c); + let s2 = acc.0[2].as_mut_ptr().add(c); + Neon::load(s0).xor3(lo.mul::<0x00>(k), hi.mul::<0x10>(k)).store(s0); + Neon::load(s1).xor3(lo.mul::<0x01>(k), hi.mul::<0x11>(k)).store(s1); + Neon::load(s2).xor3(c2.mul::<0x00>(k), c2.mul::<0x11>(k)).store(s2); + } + } + let stride = 3usize.pow((R - 1) as u32); + for at in 0..stride { + for c in (0..ROW).step_by(2) { + // SAFETY: R is in 1..=PRECOMPUTED_ROUNDS, and c indexes two words in a ROW-word row. + unsafe { + let k0 = Neon::load(fg[at].as_ptr().add(c)); + let k1 = Neon::load(fg[at + stride].as_ptr().add(c)); + let lo0 = Neon::load(bg[at].lo.as_ptr().add(c)); + let lo1 = Neon::load(bg[at + stride].lo.as_ptr().add(c)); + let hi0 = Neon::load(bg[at].hi.as_ptr().add(c)); + let hi1 = Neon::load(bg[at + stride].hi.as_ptr().add(c)); + let c20 = Neon::load(bg[at].c2.as_ptr().add(c)); + let c21 = Neon::load(bg[at + stride].c2.as_ptr().add(c)); + update(&mut acc[at], c, k0, lo0, hi0, c20); + update(&mut acc[at + stride], c, k1, lo1, hi1, c21); + update( + &mut acc[at + 2 * stride], + c, + k0.xor(k1), + lo0.xor(lo1), + hi0.xor(hi1), + c20.xor(c21), + ); + } + } + } + } + /// The products of the AVX-512 `fold_acc`, a register at a time, each coefficient packed from scalars. /// /// # Safety @@ -760,11 +837,7 @@ fn grid_pass_with( } bg[LANE_IN_GRID[l]] = WeightRow::pack(&bs[l][x..x + width]); } - extend_grid::<_, R>(fg, |a, b| std::array::from_fn(|i| a[i] ^ b[i])); - extend_grid::<_, R>(bg, WeightRow::add); - for (a, (k, w)) in acc.iter_mut().zip(fg.iter().zip(bg.iter())).take(points) { - a.mul_acc(w, k); - } + accumulate_grid::(fg, bg, acc); } }; @@ -935,6 +1008,58 @@ mod tests { use crate::whir::INITIAL_FOLDING_FACTOR; use primitives::test_util::Rng; + #[test] + fn grid_extension_products_match_scalar_definition() { + fn scalar_mul(mut a: u64, mut b: u64) -> u64 { + let mut p = 0; + for _ in 0..64 { + if b & 1 != 0 { + p ^= a; + } + let carry = a >> 63; + a <<= 1; + if carry != 0 { + a ^= 0x1b; + } + b >>= 1; + } + p + } + let mut rng = Rng::new(0x671D); + for rounds in 1..=PRECOMPUTED_ROUNDS { + for block in [1, 2, 4, 8, 16] { + for lanes in [1, (1 << rounds) - 1, 1 << rounds, (1 << rounds) + 3] { + let f: Vec = (0..block * lanes).map(|_| F64(rng.next_u64())).collect(); + let weight = rng.ext_vec(f.len()); + let got = initial_rounds(&f, block, rounds, &Basis::Dense(weight.clone())); + for (point, &sum) in got.grid.iter().enumerate() { + let mut want = F192::ZERO; + for start in (0..lanes).step_by(1 << rounds) { + for x in 0..block { + let (mut k, mut e) = (0, F192::ZERO); + for lane in 0..(1 << rounds).min(lanes - start) { + let included = (0..rounds).all(|bit| { + let digit = point / 3usize.pow(bit as u32) % 3; + digit == 2 || digit == (lane >> bit) & 1 + }); + if included { + k ^= f[(start + lane) * block + x].0; + e += weight[(start + lane) * block + x]; + } + } + want += F192::new(scalar_mul(e.c0, k), scalar_mul(e.c1, k), scalar_mul(e.c2, k)); + } + } + assert_eq!( + sum, want, + "rounds={rounds}, block={block}, lanes={lanes}, point={point}" + ); + } + } + } + } + } + #[test] fn the_kept_weight_is_the_filled_one() { // Invariant: writing the weight out leaves the sums alone and keeps every word of every lane, partial groups From ceb165c833f16df100776168f84849fbe5e7ec2c Mon Sep 17 00:00:00 2001 From: Kevaundray Wedderburn Date: Thu, 8 Oct 2026 11:39:10 +0100 Subject: [PATCH 2/5] diagnostic: compare captured grid fusion on native hardware --- .github/workflows/bench.yml | 179 ++--------- bins/leanvm/src/tracked.rs | 14 + crates/pcs/src/whir/sumcheck/first_pass.rs | 4 + .../pcs/src/whir/sumcheck/grid_diagnostic.rs | 289 ++++++++++++++++++ scripts/grid-experiment.py | 105 +++++++ 5 files changed, 438 insertions(+), 153 deletions(-) create mode 100644 crates/pcs/src/whir/sumcheck/grid_diagnostic.rs create mode 100755 scripts/grid-experiment.py diff --git a/.github/workflows/bench.yml b/.github/workflows/bench.yml index 2a941128b..2316e4120 100644 --- a/.github/workflows/bench.yml +++ b/.github/workflows/bench.yml @@ -1,172 +1,45 @@ name: Bench - -# Proves a PR's merge commit and its base (the merge commit's first parent) in turns on one -# runner, so the comparison is free of which machine the job landed on: one job per testbed -# and benchmark, each building both sides and running them `ROUNDS` times in the order base, -# PR, PR, base, ... Each job saves its runs as the artifact `ab--`, and -# bench-comment.yml, on the default branch, posts what changed on the PR. It gates nothing. -# A job is one CLI benchmark, or one `benches/` target with all the benchmarks it prints. -# scripts/ab.sh does all of it, and runs the same comparison on a developer's machine. -# -# The testbeds are the self-hosted runners `size-attester7870-x64` and `size-attester7870-arm64` -# (with `self-hosted-ghr`, an ephemeral VM per job; x86-64 an Intel Xeon Gold 6548N, 32 vCPUs, -# 64 GB; arm64 an Arm Neoverse-V2, 16 vCPUs, 64 GB): base and PR build and run on one machine, so -# each is built for `target-cpu=native` (on x86-64 AVX-512, VPCLMULQDQ and GFNI included), and -# its cache is keyed on the CPU model so a native build never lands on another CPU. A fork has no -# such runners, so in a fork's own repository the testbeds fall back to GitHub-hosted runners with -# a pinned ISA, x86-64 built for haswell and arm64 for neoverse-n2 (a fork's PR to this repository -# runs here, on the self-hosted runners): `target-cpu=native` would build for whichever CPU the -# job lands on, and a cached build then dies of SIGILL on an older one (see rust.yml). -# -# Run from the Actions tab (`workflow_dispatch`), it compares any two refs the same way: `head` -# (default: the branch it is run on) against `base` (default: the head's first parent). Such a -# run leaves its artifacts only; bench-comment.yml posts for PRs alone. - on: - pull_request: - # Not `edited`: a description edit would restart, and so cancel, the PR's run. - types: [ opened, reopened, synchronize ] workflow_dispatch: inputs: base: - description: "Base: a branch, tag or commit (default: the head's first parent)" - required: false + description: Paired upstream baseline + default: f52bd991c8623894e7b03a2e50bede6c248ea1ad head: - description: "Head: a branch, tag or commit (default: the branch this runs on)" + description: Diagnostic head required: false - + phase: + description: screen or final + default: screen permissions: contents: read - -concurrency: - group: bench-${{ github.ref }}-${{ inputs.base }}-${{ inputs.head }} - cancel-in-progress: true - -env: - CARGO_TERM_COLOR: always - RUST_BACKTRACE: 1 - ROUNDS: 5 - + actions: read jobs: - ab: - name: ${{ matrix.benchmark.name }} (${{ matrix.testbed.name }}) - runs-on: ${{ matrix.testbed.runner }} - timeout-minutes: 60 + grid: strategy: fail-fast: false matrix: - testbed: - - name: x86-64 - runner: ${{ github.repository == 'leanEthereum/leanVM' && fromJSON('["size-attester7870-x64", "self-hosted-ghr"]') || 'ubuntu-latest' }} - rustflags: ${{ github.repository == 'leanEthereum/leanVM' && '-C target-cpu=native' || '-C target-cpu=haswell' }} - # The NEON kernels are chosen at compile time on `aes` (PMULL) and `sha3` (EOR3), which - # a generic aarch64 target does not enable: the native build enables them where the CPU - # has them, the fork's neoverse-n2 pin names them, and the job fails early without them. - - name: arm64 - runner: ${{ github.repository == 'leanEthereum/leanVM' && fromJSON('["size-attester7870-arm64", "self-hosted-ghr"]') || 'ubuntu-24.04-arm' }} - rustflags: ${{ github.repository == 'leanEthereum/leanVM' && '-C target-cpu=native' || '-C target-cpu=neoverse-n2 -C target-feature=+aes,+sha3' }} - # What the CPU must have to run that build. - cpu: ${{ github.repository == 'leanEthereum/leanVM' && 'aes pmull sha3' || 'aes pmull sha3 sve2' }} - # The cases of `proven()` and `proven_trees()` in bins/leanvm/src/tracked.rs, run by the - # CLI (`-thread`, on T threads: the CLI proves it, a tree's leaf included, with - # `LEANVM_NUM_THREADS=T`, unpinned to cores since the job has its runner's VM to itself and - # base and PR alternate within it; a tree's case, `Nto1` for N proofs combined into 1, - # prints two benchmarks: `-first`, a first-level node, the RISC-V verifier in rows over N - # leaf proofs, and `-node`, a higher node, the recursion verifier in rows over N child - # proofs), then the `benches/` targets scripts/ab.sh names, each run with `--json`. - benchmark: - - name: fibonacci-asm-2000000 - - name: hash-50000 - - name: leanxmss-100 - - name: leanxmss-100-1thread - - name: leanxmss-100-4thread - - name: leanxmss-100-8thread - - name: leansphincs-26 - - name: falcon-7 - - name: stateproof-5 - - name: shielded-258 - - name: aggregate-leanxmss-100-2to1 - - name: aggregate-leanxmss-100-4to1 - - name: aggregate-leanxmss-100-2to1-1thread - - name: aggregate-leanxmss-100-2to1-4thread - - name: aggregate-leanxmss-100-2to1-8thread - - name: aggregate-leanxmss-100-4to1-1thread - - name: aggregate-leanxmss-100-4to1-4thread - - name: aggregate-leanxmss-100-4to1-8thread - - name: flock-hash-batch-262144 - - name: flock-arithmetic-batch - - name: pcs-throughput - - name: blake2s-batch - - name: kernels + include: + - arch: x86-64 + runner: size-attester7870-x64 + - arch: arm64 + runner: size-attester7870-arm64 + runs-on: [self-hosted-ghr, '${{ matrix.runner }}'] + timeout-minutes: 90 env: - RUSTFLAGS: ${{ matrix.testbed.rustflags }} - BENCHMARK: ${{ matrix.benchmark.name }} + BASE: '${{ inputs.base }}' + PHASE: '${{ inputs.phase }}' + GH_TOKEN: '${{ github.token }}' steps: - uses: actions/checkout@v4 with: - ref: ${{ inputs.head }} - # The head (a PR's merge commit) and its first parent, the default base. - fetch-depth: 2 - - uses: dtolnay/rust-toolchain@stable - # What the build can use, so the PR comment's numbers can be read against the hardware. - # The model goes into the cache key: a native build is only reused on the CPU it was built for. - - name: CPU - id: cpu - run: | - lscpu | grep -E '^(Vendor ID|Model name|CPU\(s\)|Thread\(s\) per core|Core\(s\) per (socket|cluster)|L1d cache|L2 cache|L3 cache):' - free -g | grep -E '^Mem:' - uname -srvm - echo "Page size (bytes): $(getconf PAGESIZE)" - for path in \ - /sys/kernel/mm/transparent_hugepage/{enabled,defrag,hpage_pmd_size} \ - /sys/devices/system/cpu/cpufreq/policy*/{scaling_driver,scaling_governor,scaling_cur_freq,cpuinfo_cur_freq,cpuinfo_min_freq,cpuinfo_max_freq}; do - if value=$(cat "$path" 2>/dev/null); then - printf '%s: %s\n' "$path" "$value" - fi - done - if [ "$RUNNER_ARCH" = ARM64 ]; then - flags=$(grep -m1 -E '^Features' /proc/cpuinfo) - wanted="asimd aes pmull sha3 sve sve2 svepmull svebitperm i8mm bf16" - else - flags=$(grep -m1 -E '^flags' /proc/cpuinfo) - wanted="avx2 avx512f avx512bw avx512vl avx512vbmi avx512ifma vpclmulqdq gfni" - fi - for f in $wanted; do - if grep -qw "$f" <<<"$flags"; then echo "$f: yes"; else echo "$f: no"; fi - done - echo "model=-$(lscpu | sed -n 's/^Model name: *//p' | head -n1 | tr -c 'A-Za-z0-9\n' '-')" >> "$GITHUB_OUTPUT" - # The build must enable the NEON kernels' features, and the CPU must have what the build - # enables, or the comparison measures the fallback or dies of SIGILL. - - name: NEON kernels - if: runner.arch == 'ARM64' - env: - CPU_FEATURES: ${{ matrix.testbed.cpu }} - run: | - read -ra flags <<<"$RUSTFLAGS" - cfg=$(rustc --print cfg "${flags[@]}") - for f in aes sha3; do grep -qx "target_feature=\"$f\"" <<<"$cfg" || { echo "::error::$RUSTFLAGS does not enable $f"; exit 1; }; done - features=$(grep -m1 -E '^Features' /proc/cpuinfo) - for f in $CPU_FEATURES; do grep -qw "$f" <<<"$features" || { echo "::error::no $f on this runner"; exit 1; }; done - # One cache per testbed and CPU model; RUSTFLAGS is part of its key. - - uses: Swatinem/rust-cache@v2 - with: - key: ${{ matrix.testbed.name }}${{ steps.cpu.outputs.model }} - - name: Prove the PR and its base in turns - env: - PR: ${{ github.event.number }} - TESTBED: ${{ matrix.testbed.name }} - BASE: ${{ inputs.base }} - run: | - args=(--ci --testbed "$TESTBED" --out ab) - if [ -n "$PR" ]; then args+=(--pr "$PR"); fi - if [ -n "$BASE" ]; then - git fetch --quiet --no-tags --depth=1 origin "$BASE" - args+=(--base "$(git rev-parse FETCH_HEAD)") - fi - scripts/ab.sh "${args[@]}" "$BENCHMARK" + ref: '${{ inputs.head }}' + fetch-depth: 0 + - uses: dtolnay/rust-toolchain@1.99.0 + - name: Paired dense-grid experiment + run: python3 scripts/grid-experiment.py - uses: actions/upload-artifact@v4 + if: always() with: - name: ab-${{ matrix.testbed.name }}-${{ matrix.benchmark.name }} - path: | - ab/ab-${{ matrix.testbed.name }}-${{ matrix.benchmark.name }}/ab.json - ab/ab-${{ matrix.testbed.name }}-${{ matrix.benchmark.name }}/trace-*.log + name: grid-${{ matrix.arch }} + path: grid-evidence/ diff --git a/bins/leanvm/src/tracked.rs b/bins/leanvm/src/tracked.rs index 1598ba0d6..4d3e33228 100644 --- a/bins/leanvm/src/tracked.rs +++ b/bins/leanvm/src/tracked.rs @@ -371,6 +371,7 @@ fn proved(case: &Case, prover: &Prover, plan: Plan) -> Vec<(String, Metric)> { let (verified, verify_time) = Plan::new(VERIFY_PASSES, 0).measure_quiet(|_| workload.program.verify(output, &proof)); verified.expect("an honest proof verifies"); + diagnostic_save(&proof.to_bytes()); measures( &time, proof.to_bytes().len(), @@ -417,6 +418,7 @@ fn proved_tree( ) -> Vec<(String, Vec<(String, Metric)>)> { eprintln!("{}", tree.name); let ProvenRun { proof, output, .. } = tree.leaf.prove(leaf_prover); + diagnostic_save(&proof.to_bytes()); let built = tree.tree(LeafShape::of(&proof).expect("an honest announcement"), prover.rate()); // The leaf's stages. bench::take_stages(); @@ -450,6 +452,7 @@ fn proved_node( let peak_memory = bench::peak_rss_bytes(); let (verified, verify_time) = Plan::new(VERIFY_PASSES, 0).measure_quiet(|_| tree.verify(&proof, outputs)); verified.expect("an honest tree proof verifies"); + diagnostic_save(&proof.to_bytes()); let report = measures( &time, proof.to_bytes().len(), @@ -543,3 +546,14 @@ fn tree_table(trees: &[(Aggregation, [(Kind, CircuitStats); 2])]) -> String { } table } + +fn diagnostic_save(bytes: &[u8]) { + use std::sync::atomic::{AtomicUsize, Ordering}; + static NEXT: AtomicUsize = AtomicUsize::new(0); + if let Some(directory) = std::env::var_os("GRID_PROOFS") { + let directory = std::path::PathBuf::from(directory); + std::fs::create_dir_all(&directory).unwrap(); + let index = NEXT.fetch_add(1, Ordering::Relaxed); + std::fs::write(directory.join(format!("{index}.bin")), bytes).unwrap(); + } +} diff --git a/crates/pcs/src/whir/sumcheck/first_pass.rs b/crates/pcs/src/whir/sumcheck/first_pass.rs index 7e9ad83d2..bb058bc61 100644 --- a/crates/pcs/src/whir/sumcheck/first_pass.rs +++ b/crates/pcs/src/whir/sumcheck/first_pass.rs @@ -1129,3 +1129,7 @@ mod tests { check::("neon"); } } + +#[cfg(test)] +#[path = "grid_diagnostic.rs"] +mod grid_diagnostic; diff --git a/crates/pcs/src/whir/sumcheck/grid_diagnostic.rs b/crates/pcs/src/whir/sumcheck/grid_diagnostic.rs new file mode 100644 index 000000000..abc61eaf3 --- /dev/null +++ b/crates/pcs/src/whir/sumcheck/grid_diagnostic.rs @@ -0,0 +1,289 @@ +use super::*; + +#[test] +#[ignore = "requires ARM_ATTRIBUTION_BASIS_DIR from a non-timed staged production proof"] +fn captured_basis_components() { + use std::hint::black_box; + use std::io::{BufReader, Read, Seek, SeekFrom}; + use std::time::Instant; + + const SAMPLE_ROWS: usize = 64; + const ROWS_PER_GROUP: usize = SAMPLE_ROWS / 2; + struct Input { + k: [[u64; ROW]; GROUP], + e: [[F192; ROW]; GROUP], + } + struct Expanded { + k: [[u64; ROW]; GRID], + e: [WeightRow; GRID], + } + + fn read_window(file: &mut BufReader, offset: usize) -> [u8; N] { + file.seek(SeekFrom::Start(offset as u64)).expect("seek captured row"); + let mut bytes = [0; N]; + file.read_exact(&mut bytes).expect("read captured row"); + bytes + } + fn word(bytes: &[u8]) -> u64 { + u64::from_le_bytes(bytes.try_into().expect("eight-byte coefficient")) + } + fn unpack(row: &WeightRow, x: usize) -> F192 { + let pair = if x.is_multiple_of(2) { &row.lo } else { &row.hi }; + F192::new(pair[x & !1], pair[x | 1], row.c2[x]) + } + // Independent bit-serial K product, including reduction by x^64+x^4+x^3+x+1. + fn scalar_mul(mut a: u64, mut b: u64) -> u64 { + let mut product = 0; + for _ in 0..64 { + if b & 1 != 0 { + product ^= a; + } + let carry = a >> 63; + a <<= 1; + if carry != 0 { + a ^= 0x1b; + } + b >>= 1; + } + product + } + + #[inline(never)] + fn pack_k(inputs: &[Input], expanded: &mut [Expanded]) { + for (input, grid) in inputs.iter().zip(expanded) { + for (lane, &at) in LANE_IN_GRID.iter().enumerate() { + for (dst, &src) in grid.k[at].iter_mut().zip(&input.k[lane]) { + *dst = src; + } + } + } + } + #[inline(never)] + fn pack_e(inputs: &[Input], expanded: &mut [Expanded]) { + for (input, grid) in inputs.iter().zip(expanded) { + for (lane, &at) in LANE_IN_GRID.iter().enumerate() { + grid.e[at] = WeightRow::pack(&input.e[lane]); + } + } + } + #[inline(never)] + fn extend_k(expanded: &mut [Expanded]) { + for grid in expanded { + extend_grid::<_, PRECOMPUTED_ROUNDS>(&mut grid.k, |a, b| std::array::from_fn(|i| a[i] ^ b[i])); + } + } + #[inline(never)] + fn extend_e(expanded: &mut [Expanded]) { + for grid in expanded { + extend_grid::<_, PRECOMPUTED_ROUNDS>(&mut grid.e, WeightRow::add); + } + } + #[inline(never)] + fn products(expanded: &[Expanded], acc: &mut [ProductRow; GRID]) { + for grid in expanded { + for (acc, (k, e)) in acc.iter_mut().zip(grid.k.iter().zip(&grid.e)) { + acc.mul_acc(e, k); + } + } + } + #[inline(never)] + fn sums(acc: &[ProductRow; GRID], out: &mut [F192; GRID]) { + for (out, acc) in out.iter_mut().zip(acc) { + *out = acc.sum(); + } + } + fn measure(label: &str, iterations: usize, units: usize, unit: &str, mut run: impl FnMut()) { + for _ in 0..8 { + run(); + } + for sample in 0..5 { + let start = Instant::now(); + for _ in 0..iterations { + run(); + } + let elapsed_ns = start.elapsed().as_nanos(); + let ns_per_sweep = elapsed_ns as f64 / iterations as f64; + eprintln!( + "basis_component phase={label} sample={sample} iterations={iterations} elapsed_ns={elapsed_ns} \ + ns_per_sweep={ns_per_sweep:.3} units_per_sweep={units} unit={unit} ns_per_unit={:.3}", + ns_per_sweep / units as f64, + ); + } + } + + let directory = std::path::PathBuf::from( + std::env::var_os("ARM_ATTRIBUTION_BASIS_DIR").expect("captured production Basis directory"), + ); + let shape = std::fs::read(directory.join("shape.bin")).expect("read Basis shape"); + assert_eq!(shape.len(), 24, "three little-endian u64 shape words"); + let (n, block, initial_k) = (word(&shape[..8]) as usize, word(&shape[8..16]) as usize, word(&shape[16..]) as usize); + assert_eq!(PRECOMPUTED_ROUNDS, 4); + assert!(initial_k >= PRECOMPUTED_ROUNDS); + assert!(block.is_power_of_two() && block / ROW >= ROWS_PER_GROUP); + assert!(n.is_multiple_of(block) && n / block >= 2 * GROUP, "two complete real lane groups"); + let witness_file = std::fs::File::open(directory.join("witness.bin")).expect("open Basis witness"); + let weight_file = std::fs::File::open(directory.join("weight.bin")).expect("open Basis weight"); + assert_eq!(witness_file.metadata().unwrap().len(), (n * size_of::()) as u64); + assert_eq!(weight_file.metadata().unwrap().len(), (n * size_of::()) as u64); + let (mut witness, mut weight) = (BufReader::new(witness_file), BufReader::new(weight_file)); + let mut inputs = Vec::with_capacity(SAMPLE_ROWS); + let mut offsets = Vec::with_capacity(SAMPLE_ROWS); + for group in 0..2 { + for sample in 0..ROWS_PER_GROUP { + let x = sample * (block / ROW - 1) / (ROWS_PER_GROUP - 1) * ROW; + offsets.push((group, x)); + let mut input = Input { + k: [[0; ROW]; GROUP], + e: [[F192::ZERO; ROW]; GROUP], + }; + for lane in 0..GROUP { + let at = (group * GROUP + lane) * block + x; + let ks = read_window::<{ ROW * 8 }>(&mut witness, at * 8); + let es = read_window::<{ ROW * 24 }>(&mut weight, at * 24); + for i in 0..ROW { + input.k[lane][i] = word(&ks[i * 8..(i + 1) * 8]); + let e = &es[i * 24..(i + 1) * 24]; + input.e[lane][i] = F192::new(word(&e[..8]), word(&e[8..16]), word(&e[16..])); + } + } + inputs.push(input); + } + } + drop((witness, weight)); + let mut expanded: Vec<_> = inputs.iter().map(|_| Expanded { + k: [[0; ROW]; GRID], + e: [WeightRow::default(); GRID], + }).collect(); + pack_k(&inputs, &mut expanded); + pack_e(&inputs, &mut expanded); + // Reconstruct each Boolean lane's ternary index without LANE_IN_GRID. + for (input, grid) in inputs.iter().zip(&expanded) { + for lane in 0..GROUP { + let at = (0..PRECOMPUTED_ROUNDS).map(|bit| ((lane >> bit) & 1) * 3usize.pow(bit as u32)).sum::(); + assert_eq!(grid.k[at], input.k[lane]); + for x in 0..ROW { + assert_eq!(unpack(&grid.e[at], x), input.e[lane][x]); + } + } + } + extend_k(&mut expanded); + extend_e(&mut expanded); + let mut expected = [F192::ZERO; GRID]; + for (input, grid) in inputs.iter().zip(&expanded) { + for (point, expected_sum) in expected.iter_mut().enumerate() { + let mut point_sum = F192::ZERO; + for x in 0..ROW { + let (mut k, mut e) = (0u64, F192::ZERO); + for lane in 0..GROUP { + let included = (0..PRECOMPUTED_ROUNDS).all(|bit| { + let digit = point / 3usize.pow(bit as u32) % 3; + digit == 2 || digit == (lane >> bit) & 1 + }); + if included { + k ^= input.k[lane][x]; + e += input.e[lane][x]; + } + } + assert_eq!(grid.k[point][x], k, "K extension point={point}, offset={x}"); + assert_eq!(unpack(&grid.e[point], x), e, "E extension point={point}, offset={x}"); + point_sum += F192::new(scalar_mul(e.c0, k), scalar_mul(e.c1, k), scalar_mul(e.c2, k)); + } + let mut product = ProductRow::default(); + product.mul_acc(&grid.e[point], &grid.k[point]); + assert_eq!(product.sum(), point_sum, "mixed products and reduction point={point}"); + *expected_sum += point_sum; + } + } + let mut acc = [ProductRow::default(); GRID]; + let mut output = [F192::ZERO; GRID]; + products(&expanded, &mut acc); + sums(&acc, &mut output); + assert_eq!(output, expected, "aggregate across captured rows"); + eprintln!( + "basis_component_reference checked=true words={n} block={block} initial_k={initial_k} rows={SAMPLE_ROWS} \ + input_bytes={} expanded_bytes={} accumulator_bytes={} offsets={offsets:?}", + inputs.len() * size_of::(), expanded.len() * size_of::(), size_of_val(&acc), + ); + // The expanded batch intentionally changes cache/scheduling from the fused production loop. + // All buffers are allocated before measurement; product accumulation is XOR, so odd repeats retain the sum. + measure("pack_k", 513, SAMPLE_ROWS, "input_row", || { + pack_k(black_box(&inputs), black_box(&mut expanded)); + }); + measure("pack_e", 513, SAMPLE_ROWS, "input_row", || { + pack_e(black_box(&inputs), black_box(&mut expanded)); + }); + measure("extend_k", 513, SAMPLE_ROWS, "input_row", || { + extend_k(black_box(&mut expanded)); + }); + measure("extend_e", 513, SAMPLE_ROWS, "input_row", || { + extend_e(black_box(&mut expanded)); + }); + acc.fill(ProductRow::default()); + measure("products", 513, SAMPLE_ROWS * GRID, "grid_point", || { + products(black_box(&expanded), black_box(&mut acc)); + }); + sums(&acc, &mut output); + assert_eq!(output, expected, "odd timed product repetitions preserve the reference sum"); + measure("final_sums", 4097, GRID, "grid_point", || { + sums(black_box(&acc), black_box(&mut output)); + }); + assert_eq!(output, expected, "timed final reductions preserve the reference sum"); + #[inline(never)] + fn fused(inputs: &[Input], grid: &mut Expanded, acc: &mut [ProductRow; GRID]) { + for input in inputs { + for (lane, &at) in LANE_IN_GRID.iter().enumerate() { + grid.k[at] = input.k[lane]; + grid.e[at] = WeightRow::pack(&input.e[lane]); + } + #[cfg(leanvm_grid_candidate)] + accumulate_grid::<4>(&mut grid.k, &mut grid.e, acc); + #[cfg(not(leanvm_grid_candidate))] + { + extend_grid::<_, 4>(&mut grid.k, |a, b| std::array::from_fn(|i| a[i] ^ b[i])); + extend_grid::<_, 4>(&mut grid.e, WeightRow::add); + for (a, (k, e)) in acc.iter_mut().zip(grid.k.iter().zip(&grid.e)) { + a.mul_acc(e, k); + } + } + } + } + acc.fill(ProductRow::default()); + fused(&inputs, &mut expanded[0], &mut acc); + sums(&acc, &mut output); + assert_eq!(output, expected, "fused captured grid"); + acc.fill(ProductRow::default()); + measure("fused", 513, SAMPLE_ROWS, "input_row", || { + fused(black_box(&inputs), black_box(&mut expanded[0]), black_box(&mut acc)); + }); + sums(&acc, &mut output); + assert_eq!(output, expected, "timed fused captured grid"); + black_box(output); +} + +#[test] +#[ignore = "requires captured production witness and weights"] +fn captured_dense_pass() { + let directory = std::path::PathBuf::from(std::env::var_os("ARM_ATTRIBUTION_BASIS_DIR").unwrap()); + let shape = std::fs::read(directory.join("shape.bin")).unwrap(); + let word = |b: &[u8]| u64::from_le_bytes(b.try_into().unwrap()) as usize; + let (n, block, initial_k) = (word(&shape[..8]), word(&shape[8..16]), word(&shape[16..])); + let f: Vec<_> = std::fs::read(directory.join("witness.bin")).unwrap().chunks_exact(8) + .map(|b| F64(u64::from_le_bytes(b.try_into().unwrap()))).collect(); + let b = Basis::Dense(std::fs::read(directory.join("weight.bin")).unwrap().chunks_exact(24) + .map(|b| F192::new(word(&b[..8]) as u64, word(&b[8..16]) as u64, word(&b[16..]) as u64)).collect()); + assert_eq!(f.len(), n); + let mut result = None; + for sample in 0..6 { + let start = std::time::Instant::now(); + let grid = initial_rounds(std::hint::black_box(&f), block, initial_k, std::hint::black_box(&b)).grid; + eprintln!("dense_pass sample={sample} elapsed_ns={} grid={grid:?}", start.elapsed().as_nanos()); + if let Some(expected) = &result { + assert_eq!(&grid, expected); + } + result = Some(grid); + } + if let Some(path) = std::env::var_os("GRID_RESULT") { + let bytes: Vec = result.unwrap().iter().flat_map(|e| [e.c0, e.c1, e.c2].into_iter().flat_map(u64::to_le_bytes)).collect(); + std::fs::write(path, bytes).unwrap(); + } +} diff --git a/scripts/grid-experiment.py b/scripts/grid-experiment.py new file mode 100755 index 000000000..017dba4d8 --- /dev/null +++ b/scripts/grid-experiment.py @@ -0,0 +1,105 @@ +#!/usr/bin/env python3 +"""Temporary actual-hardware grid experiment, removed before shipping.""" +import fcntl +import hashlib +import json +import os +from pathlib import Path +import shutil +import subprocess +import time + +ROOT = Path.cwd() +OUT = ROOT / 'grid-evidence' +OUT.mkdir(exist_ok=True) +BASE = os.environ.get('BASE', 'f52bd991c8623894e7b03a2e50bede6c248ea1ad') +PRODUCTION = '4b3b3d17' +HEAD = subprocess.check_output(['git', 'rev-parse', 'HEAD'], text=True).strip() +SCOPE = ['systemd-run', '--user', '--scope', '-q', '-p', 'MemoryMax=16G', '-p', 'MemorySwapMax=0'] +ENV = dict(os.environ, CARGO_BUILD_JOBS='4', CARGO_TERM_COLOR='never') +ENV.pop('LEANVM_NUM_THREADS', None) +FLAGS = '-C target-cpu=native --check-cfg=cfg(leanvm_grid_candidate)' +records = [] + +def run(args, name, cwd=ROOT, env=None): + with (OUT / name).open('w') as log: + p = subprocess.run(list(map(str, args)), cwd=cwd, env=env or ENV, stdout=log, stderr=subprocess.STDOUT) + if p.returncode: + print((OUT / name).read_text()[-16000:], flush=True) + raise RuntimeError(f'{name}: exit {p.returncode}') + +def measured(args, name, env): + with open('/tmp/leanvm-bench.lock', 'a') as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + deadline = time.monotonic() + 900 + while os.getloadavg()[0] >= 12: + if time.monotonic() > deadline: + raise RuntimeError('load never below 12') + time.sleep(5) + before = os.getloadavg() + run(SCOPE + args, name, env=env) + after = os.getloadavg() + record = dict(name=name, before=before, after=after, accepted=max(before[0], after[0]) <= 20) + records.append(record) + (OUT / 'loads.json').write_text(json.dumps(records, indent=2)) + if not record['accepted']: + raise RuntimeError(f'{name}: overloaded sample discarded') + +for command, name in [(['rustc', '-vV'], 'compiler.txt'), (['lscpu'], 'cpu.txt'), (['rustc', '--print', 'cfg', '-C', 'target-cpu=native'], 'cfg.txt')]: + run(command, name) +(OUT / 'metadata.json').write_text(json.dumps(dict(base=BASE, head=HEAD, production=PRODUCTION, flags=FLAGS, affinity=sorted(os.sched_getaffinity(0)), phase=os.getenv('PHASE', 'screen')), indent=2)) +base = Path('/tmp/grid-base') +run(['git', 'worktree', 'add', '--detach', base, BASE], 'checkout.log') +patch = subprocess.check_output(['git', 'diff', PRODUCTION, HEAD, '--', 'crates/pcs/src/whir/sumcheck/first_pass.rs', 'crates/pcs/src/whir/sumcheck/grid_diagnostic.rs', 'bins/leanvm/src/tracked.rs']) +(OUT / 'diagnostic.patch').write_bytes(patch) +run(['git', 'apply', OUT / 'diagnostic.patch'], 'apply.log', cwd=base) +# Use the original production input captures, not a newly generated or synthetic workload. +capture = Path(os.getenv('CAPTURE_DIR', '/tmp/grid-capture/basis-inputs')) +if not capture.exists(): + run(['gh', 'run', 'download', '37755153088', '--repo', 'leanEthereum/leanVM', '--name', 'attribution-arm64', '--dir', '/tmp/grid-capture'], 'capture-download.log') +for name, size in [('witness.bin', 81788928), ('weight.bin', 245366784), ('shape.bin', 24)]: + data = (capture / name).read_bytes() + assert len(data) == size + with (OUT / 'capture-hashes.txt').open('a') as receipt: + receipt.write(f'{name} {len(data)} {hashlib.sha256(data).hexdigest()}\n') +executables = {} +tests = {} +for side, cwd in [('base', base), ('head', ROOT)]: + target = Path('/tmp/grid-target-' + side) + env = dict(ENV, CARGO_TARGET_DIR=str(target), RUSTFLAGS=FLAGS + (' --cfg leanvm_grid_candidate' if side == 'head' else '')) + run(SCOPE + ['cargo', 'build', '--release', '-p', 'leanvm-cli'], f'build-{side}.log', cwd, env) + run(SCOPE + ['cargo', 'test', '--release', '-p', 'pcs', '--lib', '--no-run', '--message-format=json'], f'test-build-{side}.log', cwd, env) + messages = [json.loads(line) for line in (OUT / f'test-build-{side}.log').read_text().splitlines() if line.startswith('{')] + tests[side] = next(m['executable'] for m in messages if m.get('reason') == 'compiler-artifact' and m.get('executable') and m.get('profile', {}).get('test')) + executables[side] = target / 'release/leanvm' + run(SCOPE + [tests[side], 'whir::sumcheck::first_pass::tests', '--test-threads=2'], f'correctness-{side}.log') + run(['objdump', '-d', '-C', executables[side]], f'assembly-{side}.txt') + run(['nm', '-C', executables[side]], f'symbols-{side}.txt') + measured([tests[side], 'grid_diagnostic::captured_basis_components', '--ignored', '--nocapture', '--test-threads=1'], f'components-{side}.log', dict(ENV, ARM_ATTRIBUTION_BASIS_DIR=str(capture), LEANVM_NUM_THREADS='1')) + +for workers in ['1', '4', '8', 'default']: + env = dict(ENV, ARM_ATTRIBUTION_BASIS_DIR=str(capture)) + if workers != 'default': + env['LEANVM_NUM_THREADS'] = workers + for pair in range(5): + for side in (['base', 'head'] if pair % 2 == 0 else ['head', 'base']): + name = f'dense-{workers}-{pair}-{side}' + measured([tests[side], 'grid_diagnostic::captured_dense_pass', '--ignored', '--nocapture', '--test-threads=1'], name + '.log', dict(env, GRID_RESULT=str(OUT / (name + '.bin')))) + assert (OUT / f'dense-{workers}-{pair}-base.bin').read_bytes() == (OUT / f'dense-{workers}-{pair}-head.bin').read_bytes() + cases = [('leanxmss-100', 1)] + if os.getenv('PHASE') == 'final': + cases += [('leanxmss-100', 2), ('aggregate-leanxmss-100-2to1', 1), ('aggregate-leanxmss-100-2to1', 2)] + for case, rate in cases: + for pair in range(5): + for side in (['base', 'head'] if pair % 2 == 0 else ['head', 'base']): + name = f'{case}-rate{rate}-{workers}-{pair}-{side}' + proofs = OUT / (name + '-proofs') + measured([executables[side], 'bench', '--only', case, '--log-inv-rate', str(rate), '--repeat', '1', '--cooldown', '0', '--tracing'], name + '.log', dict(env, GRID_PROOFS=str(proofs))) + a = OUT / f'{case}-rate{rate}-{workers}-{pair}-base-proofs' + b = OUT / f'{case}-rate{rate}-{workers}-{pair}-head-proofs' + assert sorted(p.name for p in a.iterdir()) == sorted(p.name for p in b.iterdir()) + for proof in a.iterdir(): + assert proof.read_bytes() == (b / proof.name).read_bytes(), proof + with (OUT / 'proof-equality.jsonl').open('a') as receipt: + receipt.write(json.dumps(dict(case=case, rate=rate, workers=workers, pair=pair, file=proof.name, bytes=proof.stat().st_size, sha256=hashlib.sha256(proof.read_bytes()).hexdigest())) + '\n') +print('All captured-grid references, paired grids, production verification and literal paired proof bytes passed.', flush=True) From 5186f1c86a97037c59640b79cc11eb0ea1af3333 Mon Sep 17 00:00:00 2001 From: Kevaundray Wedderburn Date: Thu, 8 Oct 2026 11:42:56 +0100 Subject: [PATCH 3/5] diagnostic: screen register-resident dense-grid sums --- crates/pcs/src/whir/sumcheck/first_pass.rs | 64 ++++++++++++++++++- .../pcs/src/whir/sumcheck/grid_diagnostic.rs | 62 +++++++++++++----- scripts/grid-experiment.py | 39 +++++++---- 3 files changed, 135 insertions(+), 30 deletions(-) diff --git a/crates/pcs/src/whir/sumcheck/first_pass.rs b/crates/pcs/src/whir/sumcheck/first_pass.rs index bb058bc61..7d52b4e3a 100644 --- a/crates/pcs/src/whir/sumcheck/first_pass.rs +++ b/crates/pcs/src/whir/sumcheck/first_pass.rs @@ -543,7 +543,7 @@ mod lanes { } /// Fuse the final extension digit with its three product updates, without storing the infinity rows. - #[cfg(all(target_arch = "aarch64", target_feature = "aes"))] + #[cfg(all(target_arch = "aarch64", target_feature = "aes", not(leanvm_grid_registers)))] #[inline(always)] pub(super) unsafe fn finish_grid( fg: &[[u64; ROW]; GRID], @@ -590,6 +590,68 @@ mod lanes { } } + #[cfg(all(target_arch = "aarch64", target_feature = "aes", leanvm_grid_registers))] + #[inline(always)] + pub(super) unsafe fn finish_grid( + fg: &[[u64; ROW]; GRID], + bg: &[WeightRow; GRID], + acc: &mut [ProductRow; GRID], + ) { + #[inline(always)] + fn update(s: (Neon, Neon, Neon), k: Neon, lo: Neon, hi: Neon, c2: Neon) -> (Neon, Neon, Neon) { + ( + s.0.xor3(lo.mul::<0x00>(k), hi.mul::<0x10>(k)), + s.1.xor3(lo.mul::<0x01>(k), hi.mul::<0x11>(k)), + s.2.xor3(c2.mul::<0x00>(k), c2.mul::<0x11>(k)), + ) + } + #[inline(always)] + unsafe fn load(acc: &ProductRow) -> (Neon, Neon, Neon) { + // SAFETY: every coefficient has a complete pair; only the first pair holds this path's sum. + unsafe { + ( + Neon::load(acc.0[0].as_ptr()), + Neon::load(acc.0[1].as_ptr()), + Neon::load(acc.0[2].as_ptr()), + ) + } + } + #[inline(always)] + unsafe fn store(s: (Neon, Neon, Neon), acc: &mut ProductRow) { + // SAFETY: as for load; remaining pairs stay zero. + unsafe { + s.0.store(acc.0[0].as_mut_ptr()); + s.1.store(acc.0[1].as_mut_ptr()); + s.2.store(acc.0[2].as_mut_ptr()); + } + } + let stride = 3usize.pow((R - 1) as u32); + for at in 0..stride { + // SAFETY: R is in 1..=PRECOMPUTED_ROUNDS; the three points and each pair lie within their rows. + unsafe { + let mut s0 = load(&acc[at]); + let mut s1 = load(&acc[at + stride]); + let mut s2 = load(&acc[at + 2 * stride]); + for c in (0..ROW).step_by(2) { + let k0 = Neon::load(fg[at].as_ptr().add(c)); + let k1 = Neon::load(fg[at + stride].as_ptr().add(c)); + let lo0 = Neon::load(bg[at].lo.as_ptr().add(c)); + let lo1 = Neon::load(bg[at + stride].lo.as_ptr().add(c)); + let hi0 = Neon::load(bg[at].hi.as_ptr().add(c)); + let hi1 = Neon::load(bg[at + stride].hi.as_ptr().add(c)); + let c20 = Neon::load(bg[at].c2.as_ptr().add(c)); + let c21 = Neon::load(bg[at + stride].c2.as_ptr().add(c)); + s0 = update(s0, k0, lo0, hi0, c20); + s1 = update(s1, k1, lo1, hi1, c21); + s2 = update(s2, k0.xor(k1), lo0.xor(lo1), hi0.xor(hi1), c20.xor(c21)); + } + store(s0, &mut acc[at]); + store(s1, &mut acc[at + stride]); + store(s2, &mut acc[at + 2 * stride]); + } + } + } + /// The products of the AVX-512 `fold_acc`, a register at a time, each coefficient packed from scalars. /// /// # Safety diff --git a/crates/pcs/src/whir/sumcheck/grid_diagnostic.rs b/crates/pcs/src/whir/sumcheck/grid_diagnostic.rs index abc61eaf3..265aee453 100644 --- a/crates/pcs/src/whir/sumcheck/grid_diagnostic.rs +++ b/crates/pcs/src/whir/sumcheck/grid_diagnostic.rs @@ -116,11 +116,18 @@ fn captured_basis_components() { ); let shape = std::fs::read(directory.join("shape.bin")).expect("read Basis shape"); assert_eq!(shape.len(), 24, "three little-endian u64 shape words"); - let (n, block, initial_k) = (word(&shape[..8]) as usize, word(&shape[8..16]) as usize, word(&shape[16..]) as usize); + let (n, block, initial_k) = ( + word(&shape[..8]) as usize, + word(&shape[8..16]) as usize, + word(&shape[16..]) as usize, + ); assert_eq!(PRECOMPUTED_ROUNDS, 4); assert!(initial_k >= PRECOMPUTED_ROUNDS); assert!(block.is_power_of_two() && block / ROW >= ROWS_PER_GROUP); - assert!(n.is_multiple_of(block) && n / block >= 2 * GROUP, "two complete real lane groups"); + assert!( + n.is_multiple_of(block) && n / block >= 2 * GROUP, + "two complete real lane groups" + ); let witness_file = std::fs::File::open(directory.join("witness.bin")).expect("open Basis witness"); let weight_file = std::fs::File::open(directory.join("weight.bin")).expect("open Basis weight"); assert_eq!(witness_file.metadata().unwrap().len(), (n * size_of::()) as u64); @@ -150,16 +157,21 @@ fn captured_basis_components() { } } drop((witness, weight)); - let mut expanded: Vec<_> = inputs.iter().map(|_| Expanded { - k: [[0; ROW]; GRID], - e: [WeightRow::default(); GRID], - }).collect(); + let mut expanded: Vec<_> = inputs + .iter() + .map(|_| Expanded { + k: [[0; ROW]; GRID], + e: [WeightRow::default(); GRID], + }) + .collect(); pack_k(&inputs, &mut expanded); pack_e(&inputs, &mut expanded); // Reconstruct each Boolean lane's ternary index without LANE_IN_GRID. for (input, grid) in inputs.iter().zip(&expanded) { for lane in 0..GROUP { - let at = (0..PRECOMPUTED_ROUNDS).map(|bit| ((lane >> bit) & 1) * 3usize.pow(bit as u32)).sum::(); + let at = (0..PRECOMPUTED_ROUNDS) + .map(|bit| ((lane >> bit) & 1) * 3usize.pow(bit as u32)) + .sum::(); assert_eq!(grid.k[at], input.k[lane]); for x in 0..ROW { assert_eq!(unpack(&grid.e[at], x), input.e[lane][x]); @@ -202,7 +214,9 @@ fn captured_basis_components() { eprintln!( "basis_component_reference checked=true words={n} block={block} initial_k={initial_k} rows={SAMPLE_ROWS} \ input_bytes={} expanded_bytes={} accumulator_bytes={} offsets={offsets:?}", - inputs.len() * size_of::(), expanded.len() * size_of::(), size_of_val(&acc), + inputs.len() * size_of::(), + expanded.len() * size_of::(), + size_of_val(&acc), ); // The expanded batch intentionally changes cache/scheduling from the fused production loop. // All buffers are allocated before measurement; product accumulation is XOR, so odd repeats retain the sum. @@ -223,7 +237,10 @@ fn captured_basis_components() { products(black_box(&expanded), black_box(&mut acc)); }); sums(&acc, &mut output); - assert_eq!(output, expected, "odd timed product repetitions preserve the reference sum"); + assert_eq!( + output, expected, + "odd timed product repetitions preserve the reference sum" + ); measure("final_sums", 4097, GRID, "grid_point", || { sums(black_box(&acc), black_box(&mut output)); }); @@ -267,23 +284,38 @@ fn captured_dense_pass() { let shape = std::fs::read(directory.join("shape.bin")).unwrap(); let word = |b: &[u8]| u64::from_le_bytes(b.try_into().unwrap()) as usize; let (n, block, initial_k) = (word(&shape[..8]), word(&shape[8..16]), word(&shape[16..])); - let f: Vec<_> = std::fs::read(directory.join("witness.bin")).unwrap().chunks_exact(8) - .map(|b| F64(u64::from_le_bytes(b.try_into().unwrap()))).collect(); - let b = Basis::Dense(std::fs::read(directory.join("weight.bin")).unwrap().chunks_exact(24) - .map(|b| F192::new(word(&b[..8]) as u64, word(&b[8..16]) as u64, word(&b[16..]) as u64)).collect()); + let f: Vec<_> = std::fs::read(directory.join("witness.bin")) + .unwrap() + .chunks_exact(8) + .map(|b| F64(u64::from_le_bytes(b.try_into().unwrap()))) + .collect(); + let b = Basis::Dense( + std::fs::read(directory.join("weight.bin")) + .unwrap() + .chunks_exact(24) + .map(|b| F192::new(word(&b[..8]) as u64, word(&b[8..16]) as u64, word(&b[16..]) as u64)) + .collect(), + ); assert_eq!(f.len(), n); let mut result = None; for sample in 0..6 { let start = std::time::Instant::now(); let grid = initial_rounds(std::hint::black_box(&f), block, initial_k, std::hint::black_box(&b)).grid; - eprintln!("dense_pass sample={sample} elapsed_ns={} grid={grid:?}", start.elapsed().as_nanos()); + eprintln!( + "dense_pass sample={sample} elapsed_ns={} grid={grid:?}", + start.elapsed().as_nanos() + ); if let Some(expected) = &result { assert_eq!(&grid, expected); } result = Some(grid); } if let Some(path) = std::env::var_os("GRID_RESULT") { - let bytes: Vec = result.unwrap().iter().flat_map(|e| [e.c0, e.c1, e.c2].into_iter().flat_map(u64::to_le_bytes)).collect(); + let bytes: Vec = result + .unwrap() + .iter() + .flat_map(|e| [e.c0, e.c1, e.c2].into_iter().flat_map(u64::to_le_bytes)) + .collect(); std::fs::write(path, bytes).unwrap(); } } diff --git a/scripts/grid-experiment.py b/scripts/grid-experiment.py index 017dba4d8..0fafa1693 100755 --- a/scripts/grid-experiment.py +++ b/scripts/grid-experiment.py @@ -13,12 +13,13 @@ OUT = ROOT / 'grid-evidence' OUT.mkdir(exist_ok=True) BASE = os.environ.get('BASE', 'f52bd991c8623894e7b03a2e50bede6c248ea1ad') -PRODUCTION = '4b3b3d17' +PRODUCTION = '4b3b3d1728b95b5acc2a7e9300303d9da4d9a76d' +DIAGNOSTIC = 'ceb165c833f16df100776168f84849fbe5e7ec2c' HEAD = subprocess.check_output(['git', 'rev-parse', 'HEAD'], text=True).strip() SCOPE = ['systemd-run', '--user', '--scope', '-q', '-p', 'MemoryMax=16G', '-p', 'MemorySwapMax=0'] ENV = dict(os.environ, CARGO_BUILD_JOBS='4', CARGO_TERM_COLOR='never') ENV.pop('LEANVM_NUM_THREADS', None) -FLAGS = '-C target-cpu=native --check-cfg=cfg(leanvm_grid_candidate)' +FLAGS = '-C target-cpu=native --check-cfg=cfg(leanvm_grid_candidate) --check-cfg=cfg(leanvm_grid_registers)' records = [] def run(args, name, cwd=ROOT, env=None): @@ -47,10 +48,13 @@ def measured(args, name, env): for command, name in [(['rustc', '-vV'], 'compiler.txt'), (['lscpu'], 'cpu.txt'), (['rustc', '--print', 'cfg', '-C', 'target-cpu=native'], 'cfg.txt')]: run(command, name) +if os.uname().machine == 'aarch64': + cfg = (OUT / 'cfg.txt').read_text() + assert 'target_feature="aes"' in cfg and 'target_feature="sha3"' in cfg (OUT / 'metadata.json').write_text(json.dumps(dict(base=BASE, head=HEAD, production=PRODUCTION, flags=FLAGS, affinity=sorted(os.sched_getaffinity(0)), phase=os.getenv('PHASE', 'screen')), indent=2)) base = Path('/tmp/grid-base') run(['git', 'worktree', 'add', '--detach', base, BASE], 'checkout.log') -patch = subprocess.check_output(['git', 'diff', PRODUCTION, HEAD, '--', 'crates/pcs/src/whir/sumcheck/first_pass.rs', 'crates/pcs/src/whir/sumcheck/grid_diagnostic.rs', 'bins/leanvm/src/tracked.rs']) +patch = subprocess.check_output(['git', 'diff', PRODUCTION, DIAGNOSTIC, '--', 'crates/pcs/src/whir/sumcheck/first_pass.rs', 'crates/pcs/src/whir/sumcheck/grid_diagnostic.rs', 'bins/leanvm/src/tracked.rs']) (OUT / 'diagnostic.patch').write_bytes(patch) run(['git', 'apply', OUT / 'diagnostic.patch'], 'apply.log', cwd=base) # Use the original production input captures, not a newly generated or synthetic workload. @@ -64,9 +68,14 @@ def measured(args, name, env): receipt.write(f'{name} {len(data)} {hashlib.sha256(data).hexdigest()}\n') executables = {} tests = {} -for side, cwd in [('base', base), ('head', ROOT)]: +variants = ['base', 'head', 'registers'] +for side in variants: + cwd = base if side == 'base' else ROOT target = Path('/tmp/grid-target-' + side) - env = dict(ENV, CARGO_TARGET_DIR=str(target), RUSTFLAGS=FLAGS + (' --cfg leanvm_grid_candidate' if side == 'head' else '')) + extra = '' if side == 'base' else ' --cfg leanvm_grid_candidate' + if side == 'registers': + extra += ' --cfg leanvm_grid_registers' + env = dict(ENV, CARGO_TARGET_DIR=str(target), RUSTFLAGS=FLAGS + extra) run(SCOPE + ['cargo', 'build', '--release', '-p', 'leanvm-cli'], f'build-{side}.log', cwd, env) run(SCOPE + ['cargo', 'test', '--release', '-p', 'pcs', '--lib', '--no-run', '--message-format=json'], f'test-build-{side}.log', cwd, env) messages = [json.loads(line) for line in (OUT / f'test-build-{side}.log').read_text().splitlines() if line.startswith('{')] @@ -82,24 +91,26 @@ def measured(args, name, env): if workers != 'default': env['LEANVM_NUM_THREADS'] = workers for pair in range(5): - for side in (['base', 'head'] if pair % 2 == 0 else ['head', 'base']): + for side in (variants if pair % 2 == 0 else list(reversed(variants))): name = f'dense-{workers}-{pair}-{side}' measured([tests[side], 'grid_diagnostic::captured_dense_pass', '--ignored', '--nocapture', '--test-threads=1'], name + '.log', dict(env, GRID_RESULT=str(OUT / (name + '.bin')))) - assert (OUT / f'dense-{workers}-{pair}-base.bin').read_bytes() == (OUT / f'dense-{workers}-{pair}-head.bin').read_bytes() + for side in variants[1:]: + assert (OUT / f'dense-{workers}-{pair}-base.bin').read_bytes() == (OUT / f'dense-{workers}-{pair}-{side}.bin').read_bytes() cases = [('leanxmss-100', 1)] if os.getenv('PHASE') == 'final': cases += [('leanxmss-100', 2), ('aggregate-leanxmss-100-2to1', 1), ('aggregate-leanxmss-100-2to1', 2)] for case, rate in cases: for pair in range(5): - for side in (['base', 'head'] if pair % 2 == 0 else ['head', 'base']): + for side in (variants if pair % 2 == 0 else list(reversed(variants))): name = f'{case}-rate{rate}-{workers}-{pair}-{side}' proofs = OUT / (name + '-proofs') measured([executables[side], 'bench', '--only', case, '--log-inv-rate', str(rate), '--repeat', '1', '--cooldown', '0', '--tracing'], name + '.log', dict(env, GRID_PROOFS=str(proofs))) a = OUT / f'{case}-rate{rate}-{workers}-{pair}-base-proofs' - b = OUT / f'{case}-rate{rate}-{workers}-{pair}-head-proofs' - assert sorted(p.name for p in a.iterdir()) == sorted(p.name for p in b.iterdir()) - for proof in a.iterdir(): - assert proof.read_bytes() == (b / proof.name).read_bytes(), proof - with (OUT / 'proof-equality.jsonl').open('a') as receipt: - receipt.write(json.dumps(dict(case=case, rate=rate, workers=workers, pair=pair, file=proof.name, bytes=proof.stat().st_size, sha256=hashlib.sha256(proof.read_bytes()).hexdigest())) + '\n') + for side in variants[1:]: + b = OUT / f'{case}-rate{rate}-{workers}-{pair}-{side}-proofs' + assert sorted(p.name for p in a.iterdir()) == sorted(p.name for p in b.iterdir()) + for proof in a.iterdir(): + assert proof.read_bytes() == (b / proof.name).read_bytes(), proof + with (OUT / 'proof-equality.jsonl').open('a') as receipt: + receipt.write(json.dumps(dict(case=case, rate=rate, workers=workers, pair=pair, side=side, file=proof.name, bytes=proof.stat().st_size, sha256=hashlib.sha256(proof.read_bytes()).hexdigest())) + '\n') print('All captured-grid references, paired grids, production verification and literal paired proof bytes passed.', flush=True) From 2872ebe0b0ae403213082c4e9f8fe72cff002266 Mon Sep 17 00:00:00 2001 From: Kevaundray Wedderburn Date: Thu, 8 Oct 2026 12:19:19 +0100 Subject: [PATCH 4/5] Keep register-resident ARM grid fusion and remove diagnostics --- .github/workflows/bench.yml | 179 ++++++++-- bins/leanvm/src/tracked.rs | 14 - crates/pcs/src/whir/sumcheck/first_pass.rs | 64 +--- .../pcs/src/whir/sumcheck/grid_diagnostic.rs | 321 ------------------ scripts/grid-experiment.py | 116 ------- 5 files changed, 164 insertions(+), 530 deletions(-) delete mode 100644 crates/pcs/src/whir/sumcheck/grid_diagnostic.rs delete mode 100755 scripts/grid-experiment.py diff --git a/.github/workflows/bench.yml b/.github/workflows/bench.yml index 2316e4120..2a941128b 100644 --- a/.github/workflows/bench.yml +++ b/.github/workflows/bench.yml @@ -1,45 +1,172 @@ name: Bench + +# Proves a PR's merge commit and its base (the merge commit's first parent) in turns on one +# runner, so the comparison is free of which machine the job landed on: one job per testbed +# and benchmark, each building both sides and running them `ROUNDS` times in the order base, +# PR, PR, base, ... Each job saves its runs as the artifact `ab--`, and +# bench-comment.yml, on the default branch, posts what changed on the PR. It gates nothing. +# A job is one CLI benchmark, or one `benches/` target with all the benchmarks it prints. +# scripts/ab.sh does all of it, and runs the same comparison on a developer's machine. +# +# The testbeds are the self-hosted runners `size-attester7870-x64` and `size-attester7870-arm64` +# (with `self-hosted-ghr`, an ephemeral VM per job; x86-64 an Intel Xeon Gold 6548N, 32 vCPUs, +# 64 GB; arm64 an Arm Neoverse-V2, 16 vCPUs, 64 GB): base and PR build and run on one machine, so +# each is built for `target-cpu=native` (on x86-64 AVX-512, VPCLMULQDQ and GFNI included), and +# its cache is keyed on the CPU model so a native build never lands on another CPU. A fork has no +# such runners, so in a fork's own repository the testbeds fall back to GitHub-hosted runners with +# a pinned ISA, x86-64 built for haswell and arm64 for neoverse-n2 (a fork's PR to this repository +# runs here, on the self-hosted runners): `target-cpu=native` would build for whichever CPU the +# job lands on, and a cached build then dies of SIGILL on an older one (see rust.yml). +# +# Run from the Actions tab (`workflow_dispatch`), it compares any two refs the same way: `head` +# (default: the branch it is run on) against `base` (default: the head's first parent). Such a +# run leaves its artifacts only; bench-comment.yml posts for PRs alone. + on: + pull_request: + # Not `edited`: a description edit would restart, and so cancel, the PR's run. + types: [ opened, reopened, synchronize ] workflow_dispatch: inputs: base: - description: Paired upstream baseline - default: f52bd991c8623894e7b03a2e50bede6c248ea1ad + description: "Base: a branch, tag or commit (default: the head's first parent)" + required: false head: - description: Diagnostic head + description: "Head: a branch, tag or commit (default: the branch this runs on)" required: false - phase: - description: screen or final - default: screen + permissions: contents: read - actions: read + +concurrency: + group: bench-${{ github.ref }}-${{ inputs.base }}-${{ inputs.head }} + cancel-in-progress: true + +env: + CARGO_TERM_COLOR: always + RUST_BACKTRACE: 1 + ROUNDS: 5 + jobs: - grid: + ab: + name: ${{ matrix.benchmark.name }} (${{ matrix.testbed.name }}) + runs-on: ${{ matrix.testbed.runner }} + timeout-minutes: 60 strategy: fail-fast: false matrix: - include: - - arch: x86-64 - runner: size-attester7870-x64 - - arch: arm64 - runner: size-attester7870-arm64 - runs-on: [self-hosted-ghr, '${{ matrix.runner }}'] - timeout-minutes: 90 + testbed: + - name: x86-64 + runner: ${{ github.repository == 'leanEthereum/leanVM' && fromJSON('["size-attester7870-x64", "self-hosted-ghr"]') || 'ubuntu-latest' }} + rustflags: ${{ github.repository == 'leanEthereum/leanVM' && '-C target-cpu=native' || '-C target-cpu=haswell' }} + # The NEON kernels are chosen at compile time on `aes` (PMULL) and `sha3` (EOR3), which + # a generic aarch64 target does not enable: the native build enables them where the CPU + # has them, the fork's neoverse-n2 pin names them, and the job fails early without them. + - name: arm64 + runner: ${{ github.repository == 'leanEthereum/leanVM' && fromJSON('["size-attester7870-arm64", "self-hosted-ghr"]') || 'ubuntu-24.04-arm' }} + rustflags: ${{ github.repository == 'leanEthereum/leanVM' && '-C target-cpu=native' || '-C target-cpu=neoverse-n2 -C target-feature=+aes,+sha3' }} + # What the CPU must have to run that build. + cpu: ${{ github.repository == 'leanEthereum/leanVM' && 'aes pmull sha3' || 'aes pmull sha3 sve2' }} + # The cases of `proven()` and `proven_trees()` in bins/leanvm/src/tracked.rs, run by the + # CLI (`-thread`, on T threads: the CLI proves it, a tree's leaf included, with + # `LEANVM_NUM_THREADS=T`, unpinned to cores since the job has its runner's VM to itself and + # base and PR alternate within it; a tree's case, `Nto1` for N proofs combined into 1, + # prints two benchmarks: `-first`, a first-level node, the RISC-V verifier in rows over N + # leaf proofs, and `-node`, a higher node, the recursion verifier in rows over N child + # proofs), then the `benches/` targets scripts/ab.sh names, each run with `--json`. + benchmark: + - name: fibonacci-asm-2000000 + - name: hash-50000 + - name: leanxmss-100 + - name: leanxmss-100-1thread + - name: leanxmss-100-4thread + - name: leanxmss-100-8thread + - name: leansphincs-26 + - name: falcon-7 + - name: stateproof-5 + - name: shielded-258 + - name: aggregate-leanxmss-100-2to1 + - name: aggregate-leanxmss-100-4to1 + - name: aggregate-leanxmss-100-2to1-1thread + - name: aggregate-leanxmss-100-2to1-4thread + - name: aggregate-leanxmss-100-2to1-8thread + - name: aggregate-leanxmss-100-4to1-1thread + - name: aggregate-leanxmss-100-4to1-4thread + - name: aggregate-leanxmss-100-4to1-8thread + - name: flock-hash-batch-262144 + - name: flock-arithmetic-batch + - name: pcs-throughput + - name: blake2s-batch + - name: kernels env: - BASE: '${{ inputs.base }}' - PHASE: '${{ inputs.phase }}' - GH_TOKEN: '${{ github.token }}' + RUSTFLAGS: ${{ matrix.testbed.rustflags }} + BENCHMARK: ${{ matrix.benchmark.name }} steps: - uses: actions/checkout@v4 with: - ref: '${{ inputs.head }}' - fetch-depth: 0 - - uses: dtolnay/rust-toolchain@1.99.0 - - name: Paired dense-grid experiment - run: python3 scripts/grid-experiment.py + ref: ${{ inputs.head }} + # The head (a PR's merge commit) and its first parent, the default base. + fetch-depth: 2 + - uses: dtolnay/rust-toolchain@stable + # What the build can use, so the PR comment's numbers can be read against the hardware. + # The model goes into the cache key: a native build is only reused on the CPU it was built for. + - name: CPU + id: cpu + run: | + lscpu | grep -E '^(Vendor ID|Model name|CPU\(s\)|Thread\(s\) per core|Core\(s\) per (socket|cluster)|L1d cache|L2 cache|L3 cache):' + free -g | grep -E '^Mem:' + uname -srvm + echo "Page size (bytes): $(getconf PAGESIZE)" + for path in \ + /sys/kernel/mm/transparent_hugepage/{enabled,defrag,hpage_pmd_size} \ + /sys/devices/system/cpu/cpufreq/policy*/{scaling_driver,scaling_governor,scaling_cur_freq,cpuinfo_cur_freq,cpuinfo_min_freq,cpuinfo_max_freq}; do + if value=$(cat "$path" 2>/dev/null); then + printf '%s: %s\n' "$path" "$value" + fi + done + if [ "$RUNNER_ARCH" = ARM64 ]; then + flags=$(grep -m1 -E '^Features' /proc/cpuinfo) + wanted="asimd aes pmull sha3 sve sve2 svepmull svebitperm i8mm bf16" + else + flags=$(grep -m1 -E '^flags' /proc/cpuinfo) + wanted="avx2 avx512f avx512bw avx512vl avx512vbmi avx512ifma vpclmulqdq gfni" + fi + for f in $wanted; do + if grep -qw "$f" <<<"$flags"; then echo "$f: yes"; else echo "$f: no"; fi + done + echo "model=-$(lscpu | sed -n 's/^Model name: *//p' | head -n1 | tr -c 'A-Za-z0-9\n' '-')" >> "$GITHUB_OUTPUT" + # The build must enable the NEON kernels' features, and the CPU must have what the build + # enables, or the comparison measures the fallback or dies of SIGILL. + - name: NEON kernels + if: runner.arch == 'ARM64' + env: + CPU_FEATURES: ${{ matrix.testbed.cpu }} + run: | + read -ra flags <<<"$RUSTFLAGS" + cfg=$(rustc --print cfg "${flags[@]}") + for f in aes sha3; do grep -qx "target_feature=\"$f\"" <<<"$cfg" || { echo "::error::$RUSTFLAGS does not enable $f"; exit 1; }; done + features=$(grep -m1 -E '^Features' /proc/cpuinfo) + for f in $CPU_FEATURES; do grep -qw "$f" <<<"$features" || { echo "::error::no $f on this runner"; exit 1; }; done + # One cache per testbed and CPU model; RUSTFLAGS is part of its key. + - uses: Swatinem/rust-cache@v2 + with: + key: ${{ matrix.testbed.name }}${{ steps.cpu.outputs.model }} + - name: Prove the PR and its base in turns + env: + PR: ${{ github.event.number }} + TESTBED: ${{ matrix.testbed.name }} + BASE: ${{ inputs.base }} + run: | + args=(--ci --testbed "$TESTBED" --out ab) + if [ -n "$PR" ]; then args+=(--pr "$PR"); fi + if [ -n "$BASE" ]; then + git fetch --quiet --no-tags --depth=1 origin "$BASE" + args+=(--base "$(git rev-parse FETCH_HEAD)") + fi + scripts/ab.sh "${args[@]}" "$BENCHMARK" - uses: actions/upload-artifact@v4 - if: always() with: - name: grid-${{ matrix.arch }} - path: grid-evidence/ + name: ab-${{ matrix.testbed.name }}-${{ matrix.benchmark.name }} + path: | + ab/ab-${{ matrix.testbed.name }}-${{ matrix.benchmark.name }}/ab.json + ab/ab-${{ matrix.testbed.name }}-${{ matrix.benchmark.name }}/trace-*.log diff --git a/bins/leanvm/src/tracked.rs b/bins/leanvm/src/tracked.rs index 4d3e33228..1598ba0d6 100644 --- a/bins/leanvm/src/tracked.rs +++ b/bins/leanvm/src/tracked.rs @@ -371,7 +371,6 @@ fn proved(case: &Case, prover: &Prover, plan: Plan) -> Vec<(String, Metric)> { let (verified, verify_time) = Plan::new(VERIFY_PASSES, 0).measure_quiet(|_| workload.program.verify(output, &proof)); verified.expect("an honest proof verifies"); - diagnostic_save(&proof.to_bytes()); measures( &time, proof.to_bytes().len(), @@ -418,7 +417,6 @@ fn proved_tree( ) -> Vec<(String, Vec<(String, Metric)>)> { eprintln!("{}", tree.name); let ProvenRun { proof, output, .. } = tree.leaf.prove(leaf_prover); - diagnostic_save(&proof.to_bytes()); let built = tree.tree(LeafShape::of(&proof).expect("an honest announcement"), prover.rate()); // The leaf's stages. bench::take_stages(); @@ -452,7 +450,6 @@ fn proved_node( let peak_memory = bench::peak_rss_bytes(); let (verified, verify_time) = Plan::new(VERIFY_PASSES, 0).measure_quiet(|_| tree.verify(&proof, outputs)); verified.expect("an honest tree proof verifies"); - diagnostic_save(&proof.to_bytes()); let report = measures( &time, proof.to_bytes().len(), @@ -546,14 +543,3 @@ fn tree_table(trees: &[(Aggregation, [(Kind, CircuitStats); 2])]) -> String { } table } - -fn diagnostic_save(bytes: &[u8]) { - use std::sync::atomic::{AtomicUsize, Ordering}; - static NEXT: AtomicUsize = AtomicUsize::new(0); - if let Some(directory) = std::env::var_os("GRID_PROOFS") { - let directory = std::path::PathBuf::from(directory); - std::fs::create_dir_all(&directory).unwrap(); - let index = NEXT.fetch_add(1, Ordering::Relaxed); - std::fs::write(directory.join(format!("{index}.bin")), bytes).unwrap(); - } -} diff --git a/crates/pcs/src/whir/sumcheck/first_pass.rs b/crates/pcs/src/whir/sumcheck/first_pass.rs index 7d52b4e3a..0189e9bd7 100644 --- a/crates/pcs/src/whir/sumcheck/first_pass.rs +++ b/crates/pcs/src/whir/sumcheck/first_pass.rs @@ -9,6 +9,12 @@ //! drawn, and one more pass folds all `R` lane bits at once (the small-value //! precomputation of Bagad, Dao, Domb and Thaler, : //! the witness is in `K`, so every product is a mixed one). +//! +//! On PMULL targets, the grid's final extension digit is consumed directly by +//! mixed-product accumulation. Its zero, one and infinity sums stay in registers +//! across the offsets in a row, so neither the infinity weights nor intermediate +//! partial sums need a separate memory pass. Wider x86 kernels retain their +//! lane-parallel layout. use super::{ Basis, BasisFill, FIRST_PASS_PAR_THRESHOLD, INITIAL_BASIS_CHUNK, KEEP_WEIGHT_MAX_THREADS, PRECOMPUTED_ROUNDS, @@ -542,55 +548,11 @@ mod lanes { } } - /// Fuse the final extension digit with its three product updates, without storing the infinity rows. - #[cfg(all(target_arch = "aarch64", target_feature = "aes", not(leanvm_grid_registers)))] - #[inline(always)] - pub(super) unsafe fn finish_grid( - fg: &[[u64; ROW]; GRID], - bg: &[WeightRow; GRID], - acc: &mut [ProductRow; GRID], - ) { - #[inline(always)] - unsafe fn update(acc: &mut ProductRow, c: usize, k: Neon, lo: Neon, hi: Neon, c2: Neon) { - // SAFETY: c names a complete pair in each row, and PMULL is enabled by the module's caller. - unsafe { - let s0 = acc.0[0].as_mut_ptr().add(c); - let s1 = acc.0[1].as_mut_ptr().add(c); - let s2 = acc.0[2].as_mut_ptr().add(c); - Neon::load(s0).xor3(lo.mul::<0x00>(k), hi.mul::<0x10>(k)).store(s0); - Neon::load(s1).xor3(lo.mul::<0x01>(k), hi.mul::<0x11>(k)).store(s1); - Neon::load(s2).xor3(c2.mul::<0x00>(k), c2.mul::<0x11>(k)).store(s2); - } - } - let stride = 3usize.pow((R - 1) as u32); - for at in 0..stride { - for c in (0..ROW).step_by(2) { - // SAFETY: R is in 1..=PRECOMPUTED_ROUNDS, and c indexes two words in a ROW-word row. - unsafe { - let k0 = Neon::load(fg[at].as_ptr().add(c)); - let k1 = Neon::load(fg[at + stride].as_ptr().add(c)); - let lo0 = Neon::load(bg[at].lo.as_ptr().add(c)); - let lo1 = Neon::load(bg[at + stride].lo.as_ptr().add(c)); - let hi0 = Neon::load(bg[at].hi.as_ptr().add(c)); - let hi1 = Neon::load(bg[at + stride].hi.as_ptr().add(c)); - let c20 = Neon::load(bg[at].c2.as_ptr().add(c)); - let c21 = Neon::load(bg[at + stride].c2.as_ptr().add(c)); - update(&mut acc[at], c, k0, lo0, hi0, c20); - update(&mut acc[at + stride], c, k1, lo1, hi1, c21); - update( - &mut acc[at + 2 * stride], - c, - k0.xor(k1), - lo0.xor(lo1), - hi0.xor(hi1), - c20.xor(c21), - ); - } - } - } - } - - #[cfg(all(target_arch = "aarch64", target_feature = "aes", leanvm_grid_registers))] + /// Consume the final extension digit and keep each point's sums in registers across the row. + /// + /// Only the first pair of each coefficient stores this path's sum; the remaining pairs stay zero. + /// Reducing or combining a ProductRow therefore uses the same representation as the other paths. + #[cfg(all(target_arch = "aarch64", target_feature = "aes"))] #[inline(always)] pub(super) unsafe fn finish_grid( fg: &[[u64; ROW]; GRID], @@ -1191,7 +1153,3 @@ mod tests { check::("neon"); } } - -#[cfg(test)] -#[path = "grid_diagnostic.rs"] -mod grid_diagnostic; diff --git a/crates/pcs/src/whir/sumcheck/grid_diagnostic.rs b/crates/pcs/src/whir/sumcheck/grid_diagnostic.rs deleted file mode 100644 index 265aee453..000000000 --- a/crates/pcs/src/whir/sumcheck/grid_diagnostic.rs +++ /dev/null @@ -1,321 +0,0 @@ -use super::*; - -#[test] -#[ignore = "requires ARM_ATTRIBUTION_BASIS_DIR from a non-timed staged production proof"] -fn captured_basis_components() { - use std::hint::black_box; - use std::io::{BufReader, Read, Seek, SeekFrom}; - use std::time::Instant; - - const SAMPLE_ROWS: usize = 64; - const ROWS_PER_GROUP: usize = SAMPLE_ROWS / 2; - struct Input { - k: [[u64; ROW]; GROUP], - e: [[F192; ROW]; GROUP], - } - struct Expanded { - k: [[u64; ROW]; GRID], - e: [WeightRow; GRID], - } - - fn read_window(file: &mut BufReader, offset: usize) -> [u8; N] { - file.seek(SeekFrom::Start(offset as u64)).expect("seek captured row"); - let mut bytes = [0; N]; - file.read_exact(&mut bytes).expect("read captured row"); - bytes - } - fn word(bytes: &[u8]) -> u64 { - u64::from_le_bytes(bytes.try_into().expect("eight-byte coefficient")) - } - fn unpack(row: &WeightRow, x: usize) -> F192 { - let pair = if x.is_multiple_of(2) { &row.lo } else { &row.hi }; - F192::new(pair[x & !1], pair[x | 1], row.c2[x]) - } - // Independent bit-serial K product, including reduction by x^64+x^4+x^3+x+1. - fn scalar_mul(mut a: u64, mut b: u64) -> u64 { - let mut product = 0; - for _ in 0..64 { - if b & 1 != 0 { - product ^= a; - } - let carry = a >> 63; - a <<= 1; - if carry != 0 { - a ^= 0x1b; - } - b >>= 1; - } - product - } - - #[inline(never)] - fn pack_k(inputs: &[Input], expanded: &mut [Expanded]) { - for (input, grid) in inputs.iter().zip(expanded) { - for (lane, &at) in LANE_IN_GRID.iter().enumerate() { - for (dst, &src) in grid.k[at].iter_mut().zip(&input.k[lane]) { - *dst = src; - } - } - } - } - #[inline(never)] - fn pack_e(inputs: &[Input], expanded: &mut [Expanded]) { - for (input, grid) in inputs.iter().zip(expanded) { - for (lane, &at) in LANE_IN_GRID.iter().enumerate() { - grid.e[at] = WeightRow::pack(&input.e[lane]); - } - } - } - #[inline(never)] - fn extend_k(expanded: &mut [Expanded]) { - for grid in expanded { - extend_grid::<_, PRECOMPUTED_ROUNDS>(&mut grid.k, |a, b| std::array::from_fn(|i| a[i] ^ b[i])); - } - } - #[inline(never)] - fn extend_e(expanded: &mut [Expanded]) { - for grid in expanded { - extend_grid::<_, PRECOMPUTED_ROUNDS>(&mut grid.e, WeightRow::add); - } - } - #[inline(never)] - fn products(expanded: &[Expanded], acc: &mut [ProductRow; GRID]) { - for grid in expanded { - for (acc, (k, e)) in acc.iter_mut().zip(grid.k.iter().zip(&grid.e)) { - acc.mul_acc(e, k); - } - } - } - #[inline(never)] - fn sums(acc: &[ProductRow; GRID], out: &mut [F192; GRID]) { - for (out, acc) in out.iter_mut().zip(acc) { - *out = acc.sum(); - } - } - fn measure(label: &str, iterations: usize, units: usize, unit: &str, mut run: impl FnMut()) { - for _ in 0..8 { - run(); - } - for sample in 0..5 { - let start = Instant::now(); - for _ in 0..iterations { - run(); - } - let elapsed_ns = start.elapsed().as_nanos(); - let ns_per_sweep = elapsed_ns as f64 / iterations as f64; - eprintln!( - "basis_component phase={label} sample={sample} iterations={iterations} elapsed_ns={elapsed_ns} \ - ns_per_sweep={ns_per_sweep:.3} units_per_sweep={units} unit={unit} ns_per_unit={:.3}", - ns_per_sweep / units as f64, - ); - } - } - - let directory = std::path::PathBuf::from( - std::env::var_os("ARM_ATTRIBUTION_BASIS_DIR").expect("captured production Basis directory"), - ); - let shape = std::fs::read(directory.join("shape.bin")).expect("read Basis shape"); - assert_eq!(shape.len(), 24, "three little-endian u64 shape words"); - let (n, block, initial_k) = ( - word(&shape[..8]) as usize, - word(&shape[8..16]) as usize, - word(&shape[16..]) as usize, - ); - assert_eq!(PRECOMPUTED_ROUNDS, 4); - assert!(initial_k >= PRECOMPUTED_ROUNDS); - assert!(block.is_power_of_two() && block / ROW >= ROWS_PER_GROUP); - assert!( - n.is_multiple_of(block) && n / block >= 2 * GROUP, - "two complete real lane groups" - ); - let witness_file = std::fs::File::open(directory.join("witness.bin")).expect("open Basis witness"); - let weight_file = std::fs::File::open(directory.join("weight.bin")).expect("open Basis weight"); - assert_eq!(witness_file.metadata().unwrap().len(), (n * size_of::()) as u64); - assert_eq!(weight_file.metadata().unwrap().len(), (n * size_of::()) as u64); - let (mut witness, mut weight) = (BufReader::new(witness_file), BufReader::new(weight_file)); - let mut inputs = Vec::with_capacity(SAMPLE_ROWS); - let mut offsets = Vec::with_capacity(SAMPLE_ROWS); - for group in 0..2 { - for sample in 0..ROWS_PER_GROUP { - let x = sample * (block / ROW - 1) / (ROWS_PER_GROUP - 1) * ROW; - offsets.push((group, x)); - let mut input = Input { - k: [[0; ROW]; GROUP], - e: [[F192::ZERO; ROW]; GROUP], - }; - for lane in 0..GROUP { - let at = (group * GROUP + lane) * block + x; - let ks = read_window::<{ ROW * 8 }>(&mut witness, at * 8); - let es = read_window::<{ ROW * 24 }>(&mut weight, at * 24); - for i in 0..ROW { - input.k[lane][i] = word(&ks[i * 8..(i + 1) * 8]); - let e = &es[i * 24..(i + 1) * 24]; - input.e[lane][i] = F192::new(word(&e[..8]), word(&e[8..16]), word(&e[16..])); - } - } - inputs.push(input); - } - } - drop((witness, weight)); - let mut expanded: Vec<_> = inputs - .iter() - .map(|_| Expanded { - k: [[0; ROW]; GRID], - e: [WeightRow::default(); GRID], - }) - .collect(); - pack_k(&inputs, &mut expanded); - pack_e(&inputs, &mut expanded); - // Reconstruct each Boolean lane's ternary index without LANE_IN_GRID. - for (input, grid) in inputs.iter().zip(&expanded) { - for lane in 0..GROUP { - let at = (0..PRECOMPUTED_ROUNDS) - .map(|bit| ((lane >> bit) & 1) * 3usize.pow(bit as u32)) - .sum::(); - assert_eq!(grid.k[at], input.k[lane]); - for x in 0..ROW { - assert_eq!(unpack(&grid.e[at], x), input.e[lane][x]); - } - } - } - extend_k(&mut expanded); - extend_e(&mut expanded); - let mut expected = [F192::ZERO; GRID]; - for (input, grid) in inputs.iter().zip(&expanded) { - for (point, expected_sum) in expected.iter_mut().enumerate() { - let mut point_sum = F192::ZERO; - for x in 0..ROW { - let (mut k, mut e) = (0u64, F192::ZERO); - for lane in 0..GROUP { - let included = (0..PRECOMPUTED_ROUNDS).all(|bit| { - let digit = point / 3usize.pow(bit as u32) % 3; - digit == 2 || digit == (lane >> bit) & 1 - }); - if included { - k ^= input.k[lane][x]; - e += input.e[lane][x]; - } - } - assert_eq!(grid.k[point][x], k, "K extension point={point}, offset={x}"); - assert_eq!(unpack(&grid.e[point], x), e, "E extension point={point}, offset={x}"); - point_sum += F192::new(scalar_mul(e.c0, k), scalar_mul(e.c1, k), scalar_mul(e.c2, k)); - } - let mut product = ProductRow::default(); - product.mul_acc(&grid.e[point], &grid.k[point]); - assert_eq!(product.sum(), point_sum, "mixed products and reduction point={point}"); - *expected_sum += point_sum; - } - } - let mut acc = [ProductRow::default(); GRID]; - let mut output = [F192::ZERO; GRID]; - products(&expanded, &mut acc); - sums(&acc, &mut output); - assert_eq!(output, expected, "aggregate across captured rows"); - eprintln!( - "basis_component_reference checked=true words={n} block={block} initial_k={initial_k} rows={SAMPLE_ROWS} \ - input_bytes={} expanded_bytes={} accumulator_bytes={} offsets={offsets:?}", - inputs.len() * size_of::(), - expanded.len() * size_of::(), - size_of_val(&acc), - ); - // The expanded batch intentionally changes cache/scheduling from the fused production loop. - // All buffers are allocated before measurement; product accumulation is XOR, so odd repeats retain the sum. - measure("pack_k", 513, SAMPLE_ROWS, "input_row", || { - pack_k(black_box(&inputs), black_box(&mut expanded)); - }); - measure("pack_e", 513, SAMPLE_ROWS, "input_row", || { - pack_e(black_box(&inputs), black_box(&mut expanded)); - }); - measure("extend_k", 513, SAMPLE_ROWS, "input_row", || { - extend_k(black_box(&mut expanded)); - }); - measure("extend_e", 513, SAMPLE_ROWS, "input_row", || { - extend_e(black_box(&mut expanded)); - }); - acc.fill(ProductRow::default()); - measure("products", 513, SAMPLE_ROWS * GRID, "grid_point", || { - products(black_box(&expanded), black_box(&mut acc)); - }); - sums(&acc, &mut output); - assert_eq!( - output, expected, - "odd timed product repetitions preserve the reference sum" - ); - measure("final_sums", 4097, GRID, "grid_point", || { - sums(black_box(&acc), black_box(&mut output)); - }); - assert_eq!(output, expected, "timed final reductions preserve the reference sum"); - #[inline(never)] - fn fused(inputs: &[Input], grid: &mut Expanded, acc: &mut [ProductRow; GRID]) { - for input in inputs { - for (lane, &at) in LANE_IN_GRID.iter().enumerate() { - grid.k[at] = input.k[lane]; - grid.e[at] = WeightRow::pack(&input.e[lane]); - } - #[cfg(leanvm_grid_candidate)] - accumulate_grid::<4>(&mut grid.k, &mut grid.e, acc); - #[cfg(not(leanvm_grid_candidate))] - { - extend_grid::<_, 4>(&mut grid.k, |a, b| std::array::from_fn(|i| a[i] ^ b[i])); - extend_grid::<_, 4>(&mut grid.e, WeightRow::add); - for (a, (k, e)) in acc.iter_mut().zip(grid.k.iter().zip(&grid.e)) { - a.mul_acc(e, k); - } - } - } - } - acc.fill(ProductRow::default()); - fused(&inputs, &mut expanded[0], &mut acc); - sums(&acc, &mut output); - assert_eq!(output, expected, "fused captured grid"); - acc.fill(ProductRow::default()); - measure("fused", 513, SAMPLE_ROWS, "input_row", || { - fused(black_box(&inputs), black_box(&mut expanded[0]), black_box(&mut acc)); - }); - sums(&acc, &mut output); - assert_eq!(output, expected, "timed fused captured grid"); - black_box(output); -} - -#[test] -#[ignore = "requires captured production witness and weights"] -fn captured_dense_pass() { - let directory = std::path::PathBuf::from(std::env::var_os("ARM_ATTRIBUTION_BASIS_DIR").unwrap()); - let shape = std::fs::read(directory.join("shape.bin")).unwrap(); - let word = |b: &[u8]| u64::from_le_bytes(b.try_into().unwrap()) as usize; - let (n, block, initial_k) = (word(&shape[..8]), word(&shape[8..16]), word(&shape[16..])); - let f: Vec<_> = std::fs::read(directory.join("witness.bin")) - .unwrap() - .chunks_exact(8) - .map(|b| F64(u64::from_le_bytes(b.try_into().unwrap()))) - .collect(); - let b = Basis::Dense( - std::fs::read(directory.join("weight.bin")) - .unwrap() - .chunks_exact(24) - .map(|b| F192::new(word(&b[..8]) as u64, word(&b[8..16]) as u64, word(&b[16..]) as u64)) - .collect(), - ); - assert_eq!(f.len(), n); - let mut result = None; - for sample in 0..6 { - let start = std::time::Instant::now(); - let grid = initial_rounds(std::hint::black_box(&f), block, initial_k, std::hint::black_box(&b)).grid; - eprintln!( - "dense_pass sample={sample} elapsed_ns={} grid={grid:?}", - start.elapsed().as_nanos() - ); - if let Some(expected) = &result { - assert_eq!(&grid, expected); - } - result = Some(grid); - } - if let Some(path) = std::env::var_os("GRID_RESULT") { - let bytes: Vec = result - .unwrap() - .iter() - .flat_map(|e| [e.c0, e.c1, e.c2].into_iter().flat_map(u64::to_le_bytes)) - .collect(); - std::fs::write(path, bytes).unwrap(); - } -} diff --git a/scripts/grid-experiment.py b/scripts/grid-experiment.py deleted file mode 100755 index 0fafa1693..000000000 --- a/scripts/grid-experiment.py +++ /dev/null @@ -1,116 +0,0 @@ -#!/usr/bin/env python3 -"""Temporary actual-hardware grid experiment, removed before shipping.""" -import fcntl -import hashlib -import json -import os -from pathlib import Path -import shutil -import subprocess -import time - -ROOT = Path.cwd() -OUT = ROOT / 'grid-evidence' -OUT.mkdir(exist_ok=True) -BASE = os.environ.get('BASE', 'f52bd991c8623894e7b03a2e50bede6c248ea1ad') -PRODUCTION = '4b3b3d1728b95b5acc2a7e9300303d9da4d9a76d' -DIAGNOSTIC = 'ceb165c833f16df100776168f84849fbe5e7ec2c' -HEAD = subprocess.check_output(['git', 'rev-parse', 'HEAD'], text=True).strip() -SCOPE = ['systemd-run', '--user', '--scope', '-q', '-p', 'MemoryMax=16G', '-p', 'MemorySwapMax=0'] -ENV = dict(os.environ, CARGO_BUILD_JOBS='4', CARGO_TERM_COLOR='never') -ENV.pop('LEANVM_NUM_THREADS', None) -FLAGS = '-C target-cpu=native --check-cfg=cfg(leanvm_grid_candidate) --check-cfg=cfg(leanvm_grid_registers)' -records = [] - -def run(args, name, cwd=ROOT, env=None): - with (OUT / name).open('w') as log: - p = subprocess.run(list(map(str, args)), cwd=cwd, env=env or ENV, stdout=log, stderr=subprocess.STDOUT) - if p.returncode: - print((OUT / name).read_text()[-16000:], flush=True) - raise RuntimeError(f'{name}: exit {p.returncode}') - -def measured(args, name, env): - with open('/tmp/leanvm-bench.lock', 'a') as lock: - fcntl.flock(lock, fcntl.LOCK_EX) - deadline = time.monotonic() + 900 - while os.getloadavg()[0] >= 12: - if time.monotonic() > deadline: - raise RuntimeError('load never below 12') - time.sleep(5) - before = os.getloadavg() - run(SCOPE + args, name, env=env) - after = os.getloadavg() - record = dict(name=name, before=before, after=after, accepted=max(before[0], after[0]) <= 20) - records.append(record) - (OUT / 'loads.json').write_text(json.dumps(records, indent=2)) - if not record['accepted']: - raise RuntimeError(f'{name}: overloaded sample discarded') - -for command, name in [(['rustc', '-vV'], 'compiler.txt'), (['lscpu'], 'cpu.txt'), (['rustc', '--print', 'cfg', '-C', 'target-cpu=native'], 'cfg.txt')]: - run(command, name) -if os.uname().machine == 'aarch64': - cfg = (OUT / 'cfg.txt').read_text() - assert 'target_feature="aes"' in cfg and 'target_feature="sha3"' in cfg -(OUT / 'metadata.json').write_text(json.dumps(dict(base=BASE, head=HEAD, production=PRODUCTION, flags=FLAGS, affinity=sorted(os.sched_getaffinity(0)), phase=os.getenv('PHASE', 'screen')), indent=2)) -base = Path('/tmp/grid-base') -run(['git', 'worktree', 'add', '--detach', base, BASE], 'checkout.log') -patch = subprocess.check_output(['git', 'diff', PRODUCTION, DIAGNOSTIC, '--', 'crates/pcs/src/whir/sumcheck/first_pass.rs', 'crates/pcs/src/whir/sumcheck/grid_diagnostic.rs', 'bins/leanvm/src/tracked.rs']) -(OUT / 'diagnostic.patch').write_bytes(patch) -run(['git', 'apply', OUT / 'diagnostic.patch'], 'apply.log', cwd=base) -# Use the original production input captures, not a newly generated or synthetic workload. -capture = Path(os.getenv('CAPTURE_DIR', '/tmp/grid-capture/basis-inputs')) -if not capture.exists(): - run(['gh', 'run', 'download', '37755153088', '--repo', 'leanEthereum/leanVM', '--name', 'attribution-arm64', '--dir', '/tmp/grid-capture'], 'capture-download.log') -for name, size in [('witness.bin', 81788928), ('weight.bin', 245366784), ('shape.bin', 24)]: - data = (capture / name).read_bytes() - assert len(data) == size - with (OUT / 'capture-hashes.txt').open('a') as receipt: - receipt.write(f'{name} {len(data)} {hashlib.sha256(data).hexdigest()}\n') -executables = {} -tests = {} -variants = ['base', 'head', 'registers'] -for side in variants: - cwd = base if side == 'base' else ROOT - target = Path('/tmp/grid-target-' + side) - extra = '' if side == 'base' else ' --cfg leanvm_grid_candidate' - if side == 'registers': - extra += ' --cfg leanvm_grid_registers' - env = dict(ENV, CARGO_TARGET_DIR=str(target), RUSTFLAGS=FLAGS + extra) - run(SCOPE + ['cargo', 'build', '--release', '-p', 'leanvm-cli'], f'build-{side}.log', cwd, env) - run(SCOPE + ['cargo', 'test', '--release', '-p', 'pcs', '--lib', '--no-run', '--message-format=json'], f'test-build-{side}.log', cwd, env) - messages = [json.loads(line) for line in (OUT / f'test-build-{side}.log').read_text().splitlines() if line.startswith('{')] - tests[side] = next(m['executable'] for m in messages if m.get('reason') == 'compiler-artifact' and m.get('executable') and m.get('profile', {}).get('test')) - executables[side] = target / 'release/leanvm' - run(SCOPE + [tests[side], 'whir::sumcheck::first_pass::tests', '--test-threads=2'], f'correctness-{side}.log') - run(['objdump', '-d', '-C', executables[side]], f'assembly-{side}.txt') - run(['nm', '-C', executables[side]], f'symbols-{side}.txt') - measured([tests[side], 'grid_diagnostic::captured_basis_components', '--ignored', '--nocapture', '--test-threads=1'], f'components-{side}.log', dict(ENV, ARM_ATTRIBUTION_BASIS_DIR=str(capture), LEANVM_NUM_THREADS='1')) - -for workers in ['1', '4', '8', 'default']: - env = dict(ENV, ARM_ATTRIBUTION_BASIS_DIR=str(capture)) - if workers != 'default': - env['LEANVM_NUM_THREADS'] = workers - for pair in range(5): - for side in (variants if pair % 2 == 0 else list(reversed(variants))): - name = f'dense-{workers}-{pair}-{side}' - measured([tests[side], 'grid_diagnostic::captured_dense_pass', '--ignored', '--nocapture', '--test-threads=1'], name + '.log', dict(env, GRID_RESULT=str(OUT / (name + '.bin')))) - for side in variants[1:]: - assert (OUT / f'dense-{workers}-{pair}-base.bin').read_bytes() == (OUT / f'dense-{workers}-{pair}-{side}.bin').read_bytes() - cases = [('leanxmss-100', 1)] - if os.getenv('PHASE') == 'final': - cases += [('leanxmss-100', 2), ('aggregate-leanxmss-100-2to1', 1), ('aggregate-leanxmss-100-2to1', 2)] - for case, rate in cases: - for pair in range(5): - for side in (variants if pair % 2 == 0 else list(reversed(variants))): - name = f'{case}-rate{rate}-{workers}-{pair}-{side}' - proofs = OUT / (name + '-proofs') - measured([executables[side], 'bench', '--only', case, '--log-inv-rate', str(rate), '--repeat', '1', '--cooldown', '0', '--tracing'], name + '.log', dict(env, GRID_PROOFS=str(proofs))) - a = OUT / f'{case}-rate{rate}-{workers}-{pair}-base-proofs' - for side in variants[1:]: - b = OUT / f'{case}-rate{rate}-{workers}-{pair}-{side}-proofs' - assert sorted(p.name for p in a.iterdir()) == sorted(p.name for p in b.iterdir()) - for proof in a.iterdir(): - assert proof.read_bytes() == (b / proof.name).read_bytes(), proof - with (OUT / 'proof-equality.jsonl').open('a') as receipt: - receipt.write(json.dumps(dict(case=case, rate=rate, workers=workers, pair=pair, side=side, file=proof.name, bytes=proof.stat().st_size, sha256=hashlib.sha256(proof.read_bytes()).hexdigest())) + '\n') -print('All captured-grid references, paired grids, production verification and literal paired proof bytes passed.', flush=True) From 299bcef05c0ca716bdfecd086ebdecb6484f76b2 Mon Sep 17 00:00:00 2001 From: Kevaundray Wedderburn Date: Thu, 8 Oct 2026 12:54:22 +0100 Subject: [PATCH 5/5] Preserve the original non-PMULL grid loop and captured bound --- crates/pcs/src/whir/sumcheck/first_pass.rs | 23 ++++++++++++++++++++-- 1 file changed, 21 insertions(+), 2 deletions(-) diff --git a/crates/pcs/src/whir/sumcheck/first_pass.rs b/crates/pcs/src/whir/sumcheck/first_pass.rs index 0189e9bd7..37afa5dcc 100644 --- a/crates/pcs/src/whir/sumcheck/first_pass.rs +++ b/crates/pcs/src/whir/sumcheck/first_pass.rs @@ -58,9 +58,19 @@ const ROW: usize = 8; /// a multilinear's `inf` is the sum of its `0` and `1`. #[inline(always)] fn extend_grid(grid: &mut [T; GRID], add: impl Fn(&T, &T) -> T) { - extend_grid_digits::(grid, R, add); + let mut stride = 1; + for i in 0..R { + for &high in &LANE_IN_GRID[..1 << (R - 1 - i)] { + let base = 3 * stride * high; + for at in base..base + stride { + grid[at + 2 * stride] = add(&grid[at], &grid[at + stride]); + } + } + stride *= 3; + } } +#[cfg(all(target_arch = "aarch64", target_feature = "aes"))] #[inline(always)] fn extend_grid_digits(grid: &mut [T; GRID], digits: usize, add: impl Fn(&T, &T) -> T) { let mut stride = 1; @@ -76,13 +86,13 @@ fn extend_grid_digits(grid: &mut [T; GRID], digits: usi } /// Finish each row's extension at the point where its products are consumed. +#[cfg(all(target_arch = "aarch64", target_feature = "aes"))] #[inline(always)] fn accumulate_grid( fg: &mut [[u64; ROW]; GRID], bg: &mut [WeightRow; GRID], acc: &mut [ProductRow; GRID], ) { - #[cfg(all(target_arch = "aarch64", target_feature = "aes"))] if R > 0 { extend_grid_digits::<_, R>(fg, R - 1, |a, b| std::array::from_fn(|i| a[i] ^ b[i])); extend_grid_digits::<_, R>(bg, R - 1, WeightRow::add); @@ -861,7 +871,16 @@ fn grid_pass_with( } bg[LANE_IN_GRID[l]] = WeightRow::pack(&bs[l][x..x + width]); } + #[cfg(all(target_arch = "aarch64", target_feature = "aes"))] accumulate_grid::(fg, bg, acc); + #[cfg(not(all(target_arch = "aarch64", target_feature = "aes")))] + { + extend_grid::<_, R>(fg, |a, b| std::array::from_fn(|i| a[i] ^ b[i])); + extend_grid::<_, R>(bg, WeightRow::add); + for (a, (k, w)) in acc.iter_mut().zip(fg.iter().zip(bg.iter())).take(points) { + a.mul_acc(w, k); + } + } } };