Skip to content

faster PCS - #650

Merged
TomWambsgans merged 16 commits into
mainfrom
opt/1-pcs
Oct 30, 2026
Merged

faster PCS#650
TomWambsgans merged 16 commits into
mainfrom
opt/1-pcs

Conversation

@TomWambsgans

@TomWambsgans TomWambsgans commented Oct 30, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

faster PCS

This PR rewrites the prover side of the PCS: how the witness is committed (Reed-Solomon encode plus Merkle tree) and how the WHIR opening runs. It ports the PCS work of gattaca-com/leanVM (keccak-opt branch), simplified and tuned so that it also helps on Apple silicon, not only on AVX-512.

This is a prover-side opti. No verifier change.

Results

aggregate --xmss 900 --log-inv-rate 1 and recursion --n 2 --xmss-per-leaf 900 --log-inv-rate 2, proving time against main, and peak memory of the aggregate:

machine aggregate recursion peak memory
M4 Max about -1% noise 8.2 -> 7.0 GiB
Zen 4 about -11% about -12% 8.2 -> 5.0 GiB
Zen 5 about -20% about -15% 8.2 -> 5.1 GiB

x86 gets much faster and uses about 40% less memory. The M4 keeps its speed and loses 15% of its memory.

The key ideas

1. Encode and hash in one pass over memory

The prover is limited by memory bandwidth, not by arithmetic, so the goal is to touch the large buffers as few times as possible.

  • No transposed copy of the witness. The witness is stored one column after the other ("lane-major"), but the encode wants it row by row. The first encode pass reads the columns directly and transposes them on the fly through a small buffer that fits in L2 (ntt/stage.rs).
  • Few sweeps of the codeword. The NTT layers are planned in two kinds (ntt/plan.rs). The far layers pair rows that are far apart, so they run as a few passes over the buffer, each on small groups of rows gathered into L2. The near layers only mix rows within small blocks, so they run as one last pass, a block at a time in L2.
  • Hash rows while they are hot. Each finished block of codeword rows is handed to the Merkle hasher right away, while it is still in cache, instead of hashing the whole codeword in a separate sweep. The tree keeps only its upper levels (CutTree); the opening rehashes the few leaf groups it needs.

2. Store the codeword, or recompute the rows the opening needs

The codeword is the largest buffer of the proof, often a gigabyte. After it is hashed, the opening only ever reads a few hundred of its rows.

Recomputing a row is cheap: the last NTT layers only mix rows inside one small block, so a final row depends only on its block as it was before those layers. Redoing a few hundred blocks is a tiny fraction of the encode.

So each machine chooses, through one constant, primitives::REBUILD_OVER_STORE:

x86 (and every target but Apple silicon) Apple silicon
writing a large buffer expensive: bandwidth bound, and a store first reads the line it overwrites cheap: a write-only pass costs about what a read does
what the commit keeps the codeword before its last layers, whose output is hashed but never written back; nothing at all above 1 GiB, encoding half of the cosets at a time the whole codeword
what the opening does redoes the last layers for the queried blocks, or rebuilds them from the witness copies the queried rows

Keeping less is also most of the memory saving on x86. One type, ntt::Codeword (Whole, BeforeDeep, Nothing), carries what an encode kept, and AdditiveNttF64::rows reads the opened rows from any of them.

3. The opening's first rounds run on small tables

The opening starts with a sumcheck whose first rounds fold the witness's lanes together. Main built the full weight vector, as large as the whole witness, and ran those rounds over it.

This PR writes the weight as a short sum of terms instead: each claim is an eq pattern over a few low bits times a table over the lanes, and the Keccak claims are folded through BitFold (GFNI byte matrices on x86, lookup tables elsewhere). The first rounds then run on two tables of a few hundred values per lane (stack_open/basis.rs). After the last of them, one fused pass rebuilds the folded witness and the weight together, block by block, while each block is in cache (whir::fold_all_lanes).

4. Wide field kernels that suit every target

  • F192x8 holds eight extension-field elements. On AVX-512 they are three registers of planes, so a product is a handful of carry-less multiplies with one reduction. On NEON the halves are split by hand so that products stay in registers.
  • Products accumulate unreduced (Prod8, Acc8) and are reduced once per sum, not once per term.
  • The sumcheck's fold and its next round message are computed in the same pass, staged in L1 and streamed out (whir::fold_and_msg_lsb).
  • Hashing gains AVX-512 overrides, including a single-permutation keccak_f held in five vectors (used wherever one state is hashed alone, such as the transcript and the tree's upper levels), and every AVX-512 memory move goes through one audited module (primitives::simd).

5. Smaller memory and threading details

  • Buffers that are fully written before being read skip their zero-fill (primitives::filled_vec, uninit_vec), and large x86 outputs that are not read back soon use streaming stores (primitives::stream).
  • On Linux the thread pool runs one worker per physical core the process may use (the cores behind its CPU affinity), rather than one per logical CPU.

Where to look

file what
crates/pcs/src/ntt.rs module map of the NTT: plan, passes, stage, openings
crates/pcs/src/ntt/encode.rs the encodes, and transform, the spine
crates/pcs/src/ntt/openings.rs reading rows back from a Codeword
crates/pcs/src/whir.rs commit, sumcheck prover, the WHIR levels
crates/pcs/src/stack_open/basis.rs the weight as terms for the first rounds
crates/primitives/src/field/f192x8.rs the eight-lane field type
crates/primitives/src/bit_fold.rs folding bit rows against weights (GFNI or tables)

Testing

  • cargo testall, cargo clippyall, cargo docall on the M4.
  • The same tests natively on Zen 4, and the x86 targets (AVX-512 with GFNI, Haswell) type-checked from the M4.
  • The pcs tests also pass with LEANVM_NUM_THREADS=1.

Notes

  • On an AVX2-only machine (no 512-bit carry-less multiply), storing the codeword as Apple silicon does measured 5% faster but used 1.5 GiB more memory; this PR keeps x86's choice there.
  • Credit: the PCS rewrite is gattaca's work, ported here with its authors as co-authors.

🤖 Generated with Claude Code

@github-actions

github-actions Bot commented Oct 30, 2026 •
edited
Loading

Copy link
Copy Markdown

Benchmarks: base against this PR, on one runner

8e7d480 merged into its base, against the base 8ba3bd9: each benchmark run in turns on one runner, base and PR 5 times each. A size is shown when it changed; a time or a peak memory when it moved the same way in every round and its median by at least 1% (the change is the median ratio, with the rounds' range).

Threads shows the named configuration; default means no thread count is encoded in the benchmark ID.

Keccak GKR

benchmark threads runner measure base this PR change
keccak-hash-batch-262144 default arm64 Proving time 4.00 s 3.83 s -4.42% (-5.21% to -3.70%)
Every result

Keccak GKR

benchmark threads runner CPU measure base this PR change
keccak-hash-batch-262144 default arm64 Neoverse-N2 Proving time 4.00 s 3.83 s -4.42% (-5.21% to -3.70%)
keccak-hash-batch-262144 default x86-64 AMD EPYC 7763 64-Core Processor Proving time 7.79 s 7.77 s -0.03% (-0.54% to +0.16%)

@TomWambsgans @alexX512
@vladimir-ea @claude
pcs: gattaca's PCS rewrite, tuned for Apple silicon
e1c794b 
Port of the PCS work on gattaca-com/leanVM's keccak-opt branch: the
lane-major base encode that hashes leaves as each block finishes, the
deferred-deep encode, the opening's Basis that precomputes
what makes its sumcheck almost free, the flat opened rows, and the
stack dropped after the lane rounds. It brings the primitives,
parallel and fiat_shamir changes it is built on: the eight-lane F192x8
type, Plain and the safe casts, and the pool counting physical cores.

Adapted for the M4 rather than taken as is:

- Off AVX-512, F192x8's product split its lanes with `[0, 4].map`, which
  LLVM leaves out of line, so every product passed both operands through
  memory; the halves are now split by hand. Off AVX-512 the two
  deferred sums are told apart by use: Prod8 keeps eight lanes for the
  kernels that reduce each, Acc8 one sum for those that reduce the
  total, so no kernel folds lanes it then needs or reaches unreachable!().
- Their target_feature refactor of the aarch64 field helpers (gf2_64,
  gf2_64x3, neon) is reverted to main's #[inline(always)] form: a
  target_feature fn is only an inline hint, and those helpers were left
  out of line in execution and the bus leaves on the M4.
- The commit rebuilds what an opening reads (the base codeword past
  1 GiB, the last layers of a codeword), which wins on
  x86, where the stores cost more traffic than the products; on Apple
  silicon a store-only pass is as cheap as a read, so there it keeps the
  codeword whole. One constant, primitives::REBUILD_OVER_STORE, decides.
- The cache-line assertion in Stream::copy_lines is x86-only, matching
  the non-temporal stores it guards, and the portable F192x8 is
  line-aligned, so a haswell build compiles and still streams.
- flux-profiler, their benchmark notes and the clippy lints they allowed
  are dropped or fixed, as is the code no branch of the stack reaches.
- primitives::filled_vec is an unsafe fn, its callers stating that the
  fill writes every element first; the pool's physical cores respect the
  process's CPU affinity on Linux; the AVX-512 NTT butterfly reduces
  through primitives' reduce_halves512 rather than a copy of it.

Simpler than the port, each cut measured on the M4, h1 and h3, in
interleaved pairs, and kept where it cost more than noise:

- Dropped: the repeat map (a later branch supplies it, and takes its
  consumers), the coset encoder of small extension codewords (a memory
  saving only), the lazily split OOD eq basis and the sparse low table
  of the opening's Basis, the late write-back and the deep-pass prefetch
  with the hook they needed through every NTT sweep, the NEON lane-pair
  kernel, the tunings for pools of at most four threads, the flat Rows
  of an opening, and ring_switch's private byte-table fold (BitFold's
  portable fold serves it). One 8x8 transpose and one portable module of
  F192x8 remain where there were three and thirty-one cfgs.
- Kept: the AVX-512 hash's own absorb and digest store (+1.0% on h1 and
  +3.0% on h3 without them), the L2 prefetch of the lane fold (+2.4% on
  h3), and the deferred-deep encode (+2.0% on h3's aggregate).

aggregate --xmss 900 --log-inv-rate 1 / recursion --n 2 --xmss-per-leaf
900 --log-inv-rate 2, proving time and peak memory, main -> this branch:
  M4 Max   0.878 -> 0.879 s (+0.2%), 8.18 -> 6.99 GiB / 0.301 -> 0.316 s (+4.6%), 2.32 -> 2.37 GiB
  h1 Zen 4 1.587 -> 1.414 s (-10.9%), 8.15 -> 4.74 GiB / 0.524 -> 0.459 s (-12.4%), 2.31 -> 2.27 GiB
  h3 Zen 5 0.518 -> 0.423 s (-18.4%), 8.23 -> 4.88 GiB / 0.209 -> 0.179 s (-14.3%), 2.40 -> 2.32 GiB
Interleaved runs, two to four per version, before the simplification.
Against that port, the simplified branch measured, aggregate / recursion:
  M4 Max   -1.1% / -0.4%, peak memory unchanged
  h1 Zen 4 -0.3% / +1.2%, 4.73 -> 4.99 GiB aggregate peak (no coset encoder)
  h3 Zen 5 -1.8% / -0.7%, 4.90 -> 5.07 GiB aggregate peak
On the M4, the aggregate
measured 0.980 s rebuilding as x86 does and 0.897 s keeping everything.
On h1 built for haswell (AVX2, 128-bit CLMUL), storing as Apple silicon
does measured 2.593 -> 2.448 s (-5.6%) on the xmss aggregate and
3.053 -> 2.911 s (-4.6%) on the sphincs one, at 4.66 -> 6.20 GiB and
5.19 -> 6.58 GiB of peak memory, recursion and hash unchanged; x86
keeps rebuilding.

The pool now counts physical cores. On main's code SMT had won on h1, so
it was re-measured on the full five-branch stack: proving time with every
SMT thread against the physical default, aggregate / recursion, positive
where SMT is slower:
  h1 Zen 4, 16 vs 8 workers:  +4.1% / -0.8%
  h3 Zen 5, 64 vs 32 workers: +1.1% / +9.3%

Co-Authored-By: alexX512 <alexey@gattaca.com>
Co-Authored-By: vladimir-ea <85992906+vladimir-ea@users.noreply.github.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
TomWambsgans and others added 14 commits October 30, 2026 21:37
@TomWambsgans @claude
review: fix the deep pass's coset offset and harden unsafe edges
a8c8cbf 
- The window keeps its first row rather than its first coset, so the deep
  pass and a gathered pass place a partial buffer by rows. The coset
  offset underflowed in debug builds whenever one deep block spanned
  several cosets; release wrapped back to the right index because the
  offset was zero.
- MerkleBuilder checks a block's shape before claiming it, so a caught
  panic never leaves a claimed block unwritten for finish to read.
- fold_high_k asserts a table of whole pairs; an odd length left its last
  outputs unwritten.
- The bounds checks of load_at, store_at and the simd loads and stores
  saturate rather than wrap, so a huge index fails the check.
- BitFold asserts the byte width it was built for, SendPtr is Send and
  Sync only for Send elements, Chunks::get asserts its index, the drain
  casts its scratch to lines through cast_slice, and the x86 SAFETY
  comments name the avx512bw requirement that makes them hold.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@TomWambsgans @claude
review: cut what the review found unused or written twice
6062df6 
- The extension codeword keeps its deep plan as an Option, as the base
  one does, and both read their opened rows through one kept_rows.
- RingWeight holds the claim's map and point only; the eq split and the
  bit folds that fill its values are built where the lift needs them
  (RingFold), and prove_prepare builds its eq split only when it folds
  the witness itself.
- whir_induce: the residual evaluator loses its always-empty prefix, both
  inducers share induce_sumcheck_enforced_sum, and a branch for a block of
  one row, which a rate of at least 2 never reaches, is gone.
- Layers::run always runs to the last layer, so it no longer takes it.
- Stage owns the lane-major message it fills from.
- The AVX-512 message kernels of F192x8 go through the type's own
  products, and the portable F192x8 and F64x8 share one even/odd split and
  one transpose.
- PrunedMerklePaths::prune takes the tree as a node function (the CutTree
  form); the flat-tree form lives in the one test that needs it.
- PaddingSpec, F192x8::from_coeffs, the four local PAR_THRESHOLDs and the
  pool's element-at-a-time map_collect go; map_collect is main's again.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@TomWambsgans @claude
review: correct the docs the cuts left stale
cfe5d8b 
Comments that named code since removed (the NEON pair kernel,
butterfly_zmm, functions "mirrored" from main), described behaviour the
code no longer has (a dense ring-switch weight, two cosets at a time past
the codeword threshold, SMT siblings sharing an L2, fuses_rows' condition,
when Encoder::rows returns a plan), or got a fact wrong (sixteen products
for eight, a q_flock ring, a core counted only when all its threads are
allowed, load_at in every hash backend). The garbled FOLD_STAGE doc is
rewritten and STREAM_FOLD_MIN no longer describes one machine's cache.
REBUILD_OVER_STORE is "every target but Apple silicon" in its doc and in
AGENTS.md, which also lists the pool's two new primitives, says the pool
counts the physical cores the process may use, and narrows the NEON
inlining rule to helpers that are #[inline(always)] already.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@TomWambsgans @claude
review: fold the NTT's wrappers into AdditiveNttF64, name its sizes, …
b55cb20 
…tidy docs

- AdditiveNttF64 is the twiddle table itself; the encodes, the layer
  sweeps (run_layers) and the reads of kept rows (rows) are its methods.
  Encoder, Layers, Openings and the forwarding facade.rs go, and the
  gathered pass writes a group back itself, so drain.rs goes too.
- Plan names its two deep sizes apart: log_sub is a sub-block (a block of
  layer deep_start), log_task a deep task (a run of them). reads_lanes and
  rebuilds_from_lanes replace the plan built twice by for_lanes, and the
  encode decides once whether the deep pass defers.
- The bounds checks of load_at, store_at and the simd loads and stores
  compare without overflow and without saturating arithmetic.
- The PCS throughput bench no longer times a witness copy as part of the
  open; the pool skips an allowed CPU that is offline rather than
  dropping its physical-core cap; the deferred-deep test asserts that some
  shape splits instead of which, so it passes on any pool; the projection
  test adds a shape whose claims sum eight blocks at a time.
- Docs: d, o and L become log_d, outer_log and layer; stage and parameter
  names match the code; misplaced and doubled comments move or go;
  RingFold and Stream each have one impl block.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@TomWambsgans @claude
review: one round-message kernel, the lift without a trait object, cl…
58bc0e1 
…earer names

- round_msg_lsb::<EVAL> replaces round_msg_lsb, round_msg_and_eval_lsb_ext
  and accumulate_msg: one chunked pass, with the inner product when asked.
- The lane rounds' lift is basis::Terms itself, its per-block body an impl
  Fn, so the single-implementation LaneLift trait, its LiftBlock box and
  the dyn call per block go.
- Sizes say what they count: the tables' n_terms, the committed lanes'
  lane_block, the lift's lift_len, and SumcheckProver::new takes the
  stack. The opening's basis builds its factors with point_factors and is
  built by initial_basis, apart from Terms::project.
- Prod8::reduce_each goes: it was reduce_lanes().to_array().

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@TomWambsgans @claude
review: one Codeword for what an encode kept
31cc311 
Encoded (a buffer that may be empty plus an optional plan), Kept,
L0Codeword, kept_rows and the Option<(AdditiveNttF64, Plan)> of both
commitments said one thing four ways. An encode now returns a Codeword,
Whole, BeforeDeep(words, plan) or Nothing(plan), and
AdditiveNttF64::rows reads any of them, from the message when nothing was
kept. ProverData and LigeroWitness each hold their ntt and Codeword, the
extension codeword is written as words, and an empty buffer with no plan
can no longer be built.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@TomWambsgans @claude
review: the NTT diagram names Codeword
c91bbb0 
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@TomWambsgans @claude
review: the last level's basis is the dense induce, each Codeword car…
ca64fe4 
…ries its plan

- With the residual prefix gone, induce_sumcheck_evaluate_at_residual
  evaluated sum_i w_i prod_j (1 + p_j (1 + W_j)) at boolean points only,
  which is the dense induce basis itself. The last level now calls
  induce_basis, the basis half of induce_sumcheck_poly, and the second
  implementation goes.
- Every Codeword carries the encode's Plan (Nothing also its lane count),
  so AdditiveNttF64::rows takes the codeword, the message for Nothing and
  the positions. encode_interleaved_with allocates and returns a Codeword
  as encode_lane_major does, so neither caller rebuilds one.
- Past KEEP_CODEWORD_MAX_BYTES, a codeword the opening cannot rebuild is
  kept before its deep layers rather than whole: the same memory, without
  the deep write-back. Its doc and the stack's say what is kept and when.
- Terms has one impl block and gives its lift's block length, which
  InitialBasis and SumcheckProver no longer store; Rebuild and
  finished_rows take the sub-block size from the plan.
- Wording left from the folded types ("the encoder", "Openings") and the
  whir_ntt_ext module title follow the code.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@TomWambsgans @claude
review: audited loads in the hash backends, plainer lane accessors, d…
b580212 
…ocs build

- The NEON hash lanes and the single-state AVX-512 Keccak load and store
  through load_at/store_at and the simd helpers, so no backend moves
  memory through a raw pointer of its own.
- F192x8::lane and set_lane on AVX-512 index its qwords through
  cast_slice instead of pointer arithmetic, the portable F64x8 is
  repr(transparent) as its Plain impl assumes, and Chunks is Copy for
  any T, as SendPtr already is.
- bit_fold's gfni module is private, and its register-to-byte map is one
  byte_of_reg rather than two closures.
- The pool's physical-core cap is documented as Linux-only, here and in
  AGENTS.md, and a test that compared perf_cores with itself goes.
- Public docs no longer link the private items the fold of the NTT's
  wrapper types made them reach, so rustdoc builds again.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@TomWambsgans @claude
review: the last level calls induce_sumcheck_poly; stale comments and…
6e2cabc 
… redundant allows go

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@TomWambsgans @claude
review: F192x8 and F64x8 keep only what this branch uses
4c51b2a 
Lane access, transposes, constants, the F64x8 arithmetic and its lift to
F192x8 were reached only by this branch's own test; the bus branch uses
them and brings them back. simd::to_words, the AVX-512 mul_words and the
portable transpose helper go with them, and the test checks what remains.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@TomWambsgans @claude
review: BitFold's register guard is a compile-time assert
a3f1dd0 
The GFNI fold reads its registers as initialized only if at least one
stage wrote them, which a row of 8 to 128 bytes guarantees. That is now a
const assert where the registers are read, and the per-call width checks
are debug_asserts again: on h1 (Zen 4) the runtime asserts measured
+0.5% on the xmss aggregate (0.0 +/- 0.1% for the write-back's
cast_slice and the bounds checks), in three interleaved rounds.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@TomWambsgans @claude
review: Chunks::get checks its contract in debug only; load_at bounds…
4b5dad4 
… by slicing

On h1 (Zen 4), three interleaved rounds each, against the tip before:
- Chunks::get is an unsafe fn whose contract is i < count(), which every
  caller's loop keeps; its release assert cost 0.6% on the xmss
  aggregate (-0.9/-0.2/-0.6 without it), so it is a debug_assert again.
- load_at and store_at slice to element i, then compare the rest's size,
  so the check cannot overflow and stays linear in i; the division form
  cost 0.5% on recursion against the original (-0.9/-0.4/-0.3).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@TomWambsgans @claude
review: the NEON hash lanes load through a fixed slice again
2977c3f 
Routing the two-word NEON loads and stores through load_at gave each a
length check the fixed two-word slice does not need: against the branch
before review, recursion measured +2.1% on the M4 (+1.1/+1.9/+1.2/+4.1)
with load_at, and +0.2% (+0.2/-0.2/+0.3/+0.6) with main's slice-then-vld1q
form, which keeps its bounds check and now states its SAFETY.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@TomWambsgans TomWambsgans changed the title PCS opti faster PCS Oct 30, 2026
@TomWambsgans
TomWambsgans marked this pull request as ready for review October 30, 2026 20:40
@TomWambsgans
TomWambsgans merged commit 02a78b3 into main Oct 30, 2026
8 checks passed
@TomWambsgans
TomWambsgans deleted the opt/1-pcs branch October 30, 2026 21:51
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@TomWambsgans