Repository navigation
pcs: a cheaper commit NTT plan when few workers run - #464
Open
kevaundray wants to merge 9 commits into
Open
kevaundray wants to merge 9 commits into
kevaundray wants to merge 9 commits into
Conversation
- mul_lanes_avx512: the reduction's spill is one vpshufb lookup of the top
nibble and g(x) takes doublings, so fewer ops share the shift/shuffle
pipes with the unpacks (VPCLMULQDQ zmm issues every other cycle on Zen 4).
- radix-8 and radix-4 groups keep their rows in registers: each row is
loaded and stored once per group, not once per layer.
- run_layers puts the leftover one or two layers first (largest blocks)
and walks twiddles block to block with one XOR (TwiddleWalk) instead
of a subset sum per block; twiddles_radix8 is gone.
Byte-identical: fibonacci 2M rate 1 digest 87a8a891167e3b9a, leanxmss 400
744c5cfc65b1157d. cargo test --release -p pcs passes.
1T (LEANVM_NUM_THREADS=1, taskset -c 2), base / this, two runs each:
leanxmss 400: Commit 602, 601 -> 580, 572 ms; PCS open Commit L1-L4
380, 381 -> 349, 349 ms; Prove 5500, 5490 -> 5470, 5390 ms.
fibonacci 2M: Commit 1160, 1150 -> 1110 ms (other run disturbed, load 11);
L1-L4 798, 796 -> 746 ms.
Kernel probe, 48 lanes in L1: 0.343 -> 0.308 ns per lane-butterfly.
32T leanxmss 400: Commit 67.6, 66.9 -> 65.6, 71.7 ms (second run had every
stage slower, machine noise); proving unchanged within noise.
…s the lane-major stack in place - With one worker (gated on the planned worker count; 32 threads keep the old plan) a gathered group takes up to 64 adjacent residues: they share every twiddle, so the group is one transform of R*n lanes, its copies in and out are R*n-word runs (whole lines, so streaming stores stop splitting lines) and the 41- or 35-lane scalar tail disappears. LONE_GATHER_WORDS bounds the scratch (32 MiB). - encode_lane_major_with: commit no longer transposes the stack into the codeword. When the first gathered pass takes >= 8 residues, it reads each lane's 8-word lines straight from the lane-major message and turns 8x8 blocks around in registers (gather_lanes, AVX-512). Otherwise (32T, small domains, other targets) it transposes first as before. - transform_for takes the worker count, so tests pin the one-worker plan. Byte-identical (same digests). cargo test --release -p pcs passes; tests cover the lone plan for in-place, no-message and lane-major encodes. Measured together with the next commit (see there).
…ctor for 2-7 spare lanes
- gather_lanes walks one lane octet's lines front to back (8 lines per
lane per row group), which halves the lane-major read time at 1T.
- Rows ending in 2 to 7 lanes past a whole vector (fibonacci's 35 lanes)
take one masked vector in separate #[inline(never)] sweeps; folding the
mask into the common loop cost the whole-vector shapes about 5%.
Byte-identical: fibonacci 2M 87a8a891167e3b9a, leanxmss 400
744c5cfc65b1157d. cargo test --release -p pcs passes.
This and the previous two commits, 1T, base / this, A B A B:
leanxmss 400: Commit 606, 608 -> 484, 484 ms; L1-L4 378, 378 -> 360, 350;
Prove 5480, 5430 -> 5370, 5320 ms.
fibonacci 2M: Commit 1150, 1150 -> 835, 830 ms; L1-L4 795, 797 -> 731, 730;
Prove 9180, 9250 -> 8900, 8850 ms.
32T leanxmss 400 (--repeat 3): Commit 68.4, 67.5 -> 66.5, 67.3 ms; proving
0.556, 0.568 -> 0.559, 0.578 s (+-2 to 4% spreads): unchanged.
Probe (1T, xmss shape, commit alone): transpose 50 + gathered in 39 + out 50
ms became gathered in 45 + out 22 ms.
With one worker the deep pass takes fit3 rows whenever the plan has a
gathered pass: its L3 is its own, a deep layer costs about the same from
L3 as from L2, and the smaller gathered groups read the lane-major stack
faster. 32T planning is unchanged.
Probe (1T, encode + Merkle alone): leanxmss base commit 482 -> 459 ms,
its L1 extension commit 236 -> 231 ms, fibonacci's L1 493 -> 472 ms
(fibonacci's base commit already took fit3).
Byte-identical: fibonacci 2M 87a8a891167e3b9a, leanxmss 400
744c5cfc65b1157d. cargo test --release -p pcs passes.
All four commits, 1T, base / this, A B A B:
leanxmss 400: Commit 604, 607 -> 464, 523 ms (second run noisy);
L1-L4 378, 377 -> 342, 344 ms; Prove 5470, 5460 -> 5260, 5400 ms.
fibonacci 2M: Commit 1170, 1160 -> 837, 898 ms; L1-L4 794, 809 -> 706,
709 ms; Prove 10100 (disturbed), 9370 -> 8830, 8890 ms.
32T leanxmss 400: Commit 66.6, 71.4 -> 66.8, 68.7 ms; proving 0.558, 0.58
-> 0.572, 0.629 s (last run +-21.7%, load 6.8): no change beyond noise.
pcs throughput bench, 1T, PCS_LOG_N=22 (load 7 to 11): commit 47.9, 63.6
-> 43.0, 42.8 ms; encode alone 29.7, 33.7 -> 25.6, 24.7 ms.
The plan with wide gathered groups (up to 64 adjacent residues), the L3 deep pass and the lane-major stack read in place now runs whenever the transform is planned for at most FEW_WORKERS = 2 workers, not one alone. The tests plan for 1, 2 and 3 workers, so both sides of the gate are pinned to the scalar reference whatever the pool's size. Gate, measured on a Ryzen 9 7950X3D (workers pinned to the 32 MiB L3 CCD), Commit span of the base encode, one-worker gate / the plan at every count, two runs each: 2 workers (2 cores): leanxmss 400 307, 287 -> 281, 245 ms; leanda 305, 329 -> 263, 255 ms; fibonacci 2M 549 -> 467, 511 ms 4 workers (2 cores + SMT): leanxmss 220, 216 -> 232, 234 ms; fibonacci 406, 408 -> 455, 447 ms; leanda 235, 235 -> 238, 264 ms 4 workers (4 cores): leanxmss 148, 147 -> 164, 180 ms; leanda 157, 182 -> 193, 166 ms 8 workers (8 cores): leanxmss 92.5, 87.2 -> 132, 133 ms; leanda 116, 128 -> 141, 141 ms 32 workers, pcs throughput bench PCS_LOG_N=24 rate 1: commit 24.4, 24.0 -> 30.8, 30.9 ms; encode alone 16.5, 16.3 -> 24.6, 24.6 ms Three workers were not measured and keep the narrow plan. This branch against upstream 77077c6, Commit span (base encode), runs base -> this: 1 worker: leanxmss 601, 613, 613 -> 533, 472, 458 ms; fibonacci 1160, 1150, 1150 -> 846, 836, 850 ms; leanda 646, 663, 648 -> 484, 493, 497 ms 2 workers: leanxmss 279, 312, 312 -> 326, 260, 262 ms; fibonacci 550, 540, 540 -> 452, 461, 457 ms; leanda 335, 334, 336 -> 266, 281, 276 ms 4 workers (2 cores + SMT, the narrow plan): leanxmss 260, 259 -> 238, 238 ms; fibonacci 462, 473 -> 434, 429 ms; leanda 286, 257 -> 262, 260 ms 8 workers: leanxmss 92.7, 91.0 -> 90.3, 90.7 ms; fibonacci 171, 171 -> 163, 164 ms; leanda 98.7, 98.4 -> 99.2, 92.4 ms 32 workers: leanxmss 67.0, 67.3 -> 65.6, 64.9 ms; fibonacci 125, 120 -> 121, 119 ms; leanda 69.9, 71.3 -> 69.2, 69.9 ms Byte-identical proofs: leanxmss 400, fibonacci 2M and leanda 1 hash to the same digest on base and branch at every worker count. cargo test --release -p pcs passes, natively (AVX-512), with target-cpu=znver3 (the AVX2 arm) and on aarch64 under qemu.
Benchmarks: base against this PR, on one runner477ac5e merged into its base, against the base f7b03b5: 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). Proving time
Prover stages
Verifying time
Time per operation
Every resultProving time
Prover stages
Verifying time
Time per operation
Proof size
Peak memory
allocations
heap-peak
|
kevaundray
added a commit
to kevaundray/leanVM
that referenced
this pull request
Oct 6, 2026
# Conflicts: # crates/pcs/src/ntt/additive_ntt_f64.rs
# Conflicts: # crates/pcs/src/ntt/additive_ntt_f64.rs
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
With few workers the base commit's encode (
whir::commit, theCommitspan) is one of the largest stages of a proof: at one worker it takes 601 ms of leanXMSS-400's 5.46 s and 1.16 s of Fibonacci-2M's 9.2 s. Three things cost more than they need to:Change
All in
crates/pcs/src/ntt/additive_ntt_f64.rs, plus the call inwhir/commit.rs.For every worker count:
mul_lanes_avx512: the reduction's spill is onevpshufblookup of the top nibble (where AVX-512BW is present) and the left shifts ofg(x)are doublings, which run on the add ports. The field product is unchanged:p mod (x^64 + x^4 + x^3 + x + 1) = lo ^ g(hi ^ spill)withspill = n ^ n>>1 ^ n>>3forn = hi >> 60, andx + xisx << 1in each 64-bit lane.#[inline(never)]sweeps, so the whole-vector loop keeps its registers).run_layersruns the leftover one or two layers first, where blocks are fewest and largest, and walks the twiddles block to block (TwiddleWalk): a twiddle isΣ_j bit_j(block) · row[offset + j], so from blockbtob + 1it changes by the prefix sum of the entries underb's trailing ones and the zero above them, one addition instead of a subset sum per block.For up to
FEW_WORKERS = 2workers, the wide plan:La row's twiddle depends onrow >> (d - L)alone and every residue is below the group's step, so adjacent residues share every twiddle: the group is the same transform withR · nlanes, its copies areR · n-word runs, and the scalar tail disappears.whir::commitcallsencode_lane_major_with: when the first gathered pass takes at least eight residues, it reads each lane's eight-word lines straight from the lane-major stack and turns 8x8 blocks around in registers (gather_lanes, AVX-512), each lane octet's lines front to back. The layout istranspose_lane_major's (codeword lanetis stack blockn - 1 - t). Any other plan, and any target without AVX-512, transposes first as before.More workers keep the current plan: one residue a group, in their share of L2.
Why the gate is two workers
The wide plan trades L2 residency (a group of 64 residues spills into L3) for long runs and the in-place read. Measured on the same host and pinning as below (load below 15), with the plan at every worker count against the plan gated at one worker,
Commitspan of the base encode, two runs each:and at 32 workers the
pcsthroughput bench (PCS_LOG_N=24, rate 1/2) goes from 24.4, 24.0 to 30.8, 30.9 ms a commit. So the plan pays at one and two workers and loses from four on, with or without SMT siblings: the gate isworkers <= 2, which keeps a four-thread runner (two cores with SMT) on the narrow plan. Three workers were not measured and keep the narrow plan. The L3 deep pass goes with the wide groups: in the throughput bench, wide groups without it lose more at four and eight workers, and the L3 deep pass alone gains nothing there.Measurements
Ryzen 9 7950X3D (16 cores, SMT, Zen 4), upstream
77077c6bagainst this branch, built on the same commit, run alternately under one lock, every run with--tracing --repeat 3and a 1-minute load below 12 at its start. One worker isLEANVM_NUM_THREADS=1 taskset -c 2; two, four and eight workers are pinned to the CCD with the 32 MiB L3 (cores 10 and 11; cores 10 and 11 with their SMT siblings; cores 8 to 15); 32 workers are unpinned.Commitis the base encode's span,Provethe CLI's proving mean; each cell lists the runs, base -> branch.Commit:Prove:The first branch run of leanXMSS at one and two workers, and of Fibonacci at two, were disturbed (spreads of 8 to 41% within the run); the other runs agree. The extension commits of the opening (levels 1 to 4) gain too at one, two and four workers (leanXMSS level 1: 223 -> 204 ms at one worker, 91 -> 84 ms at four); at eight workers level 1 is 1.7 ms slower in all three workloads while levels 2 to 4 gain about as much.
pcsthroughput bench (BENCH_REPEAT=5, rate 1/2, same pinning), commit / the commit's encode alone, base -> branch:PCS_LOG_N=22PCS_LOG_N=24On
f774828dLocal runs: the same comparison after merging upstream
f774828d, the same host, pinning, lock and load rule (every run with a 1-minute load below 14 at its start and below 15 at its end), two alternating rounds of--tracing --repeat 3; base -> branch:CommitProveCommitProveThe encode saves about as many milliseconds as on
77077c6b, but upstream proves faster, so the share of proving it saves grew: at one worker leanXMSS 400 goes from 3.82 to 3.67 s (4%, against about 2% on77077c6b) and Fibonacci 2M from 6.81 to 6.35 s (7%, against 4%); at two workers 5% and 4%. 32 workers are unchanged. The first branch run of leanXMSS at one worker and the first base run at two had spreads of 7.5% and 15% within the run; the other runs agree.On
556206d7, older benchmark machinesThe bench bot's comment for this branch merged into
556206d7, against556206d7: each benchmark run in turns on one runner, base and branch five times each; the change is the median ratio, with the rounds' range.commitstage and proving, per testbed:commitstageThe one-thread runs gain the most, as the plan targets: leanXMSS 100's commit 5.85% on x86-64 and 2.79% on arm64, and the one-thread aggregation 1.7% to 2.0% of proving on x86-64.
pcs-commit-22is 44.77 -> 44.11 ms (-0.88%, -5.23% to -0.09%) on arm64 and 70.61 -> 70.24 ms (-0.53%, -1.26% to +1.94%) on x86-64 (EPYC 9V74), andntt-forward-2225.17 -> 24.95 ms (-1.48%) on arm64. leanXMSS 100 on the Xeon 6973P-C runner spread up to +35% within its rounds, and the bot does not flag it. The bot flags the aggregation's verifying time on x86-64 (+1.8% to +7.8%) and itsreduce-in-rowsstage there (+5.3% to +12.3%, about 0.2 ms, while on arm64 the same stage is 20% to 32% faster), and Fibonacci's verifying time on x86-64 (+1.2%): neither the verifiers nor the recursion's row reduction call intopcs::ntt, the only code this change touches besides the commit's call.Fitting this upstream:
lane_butterfliesstays on this PR's register kernels (columns_avx512), withtransposed_butterfly_regs_avx512besidebutterfly_regs_avx512on the samemul_lanes_avx512product; the lanes past the last whole vector go to the shared scalar tail.flock::zerocheck::ntt, sopcs::nttis the F64 NTT alone. It exportsRowSinkonly: with the commit reading the stack in place,transpose_lane_majoris called only inside the NTT, as the fallback ofencode_lane_major_with.encode_interleaved_in_place_withlost its last caller (the commit now callsencode_lane_major_with, nowpub(crate)) and is removed.write_onlycodeword is whatencode_lane_major_withgets, as before: it writes every word before it reads it. aggregation: leaves at rate 1/4, a two-scalar transcript step and stratified queries #523 proves aggregation leaves at rate 1/4, so the leaves' commits now encode with four replicas; the leaf and node proofs ofaggregate --n 100 --leaves 2are byte-identical to the base at 1, 2 and 32 workers.On
f7b03b5e, new benchmark machinesThe bench bot's comment for this branch merged into
f7b03b5e, againstf7b03b5e, on the new runners (both built withtarget-cpu=native): each benchmark run in turns on one runner, base and branch five times each; the change is the median ratio, with the rounds' range. The aggregation cases report nocommitstage, so they list proving alone; the kernels bench runs on the whole pool.x86-64, Xeon Gold 6548N (AVX-512):
commitstagepcs-commit-22ntt-forward-22arm64, Neoverse-V2:
commitstagepcs-commit-22ntt-forward-22The one-thread cases gain the most, as the plan targets: on the Xeon, leanXMSS 100's commit is 28.6% faster and its proving 6.0%, and the one-thread aggregation proves 5.0% to 8.9% faster; on the Neoverse-V2 the commit is 4.8% faster and the one-thread aggregation 0.8% to 1.2%. The 4- and 8-thread cases run the narrow plan (the gate is two workers), so what they gain comes from the parts for every worker count: on the Xeon those include the AVX-512 product and register kernels, and the commit is 12% to 13% faster at four, eight and all threads, with
ntt-forward-226.1% faster; on the Neoverse-V2, where the NEON kernels are unchanged, the commit is 1.6% faster at four threads and within its rounds' range at eight. The bot flags three rises: on arm64, leanxmss-100-8thread'spcs-openstage, 182.69 -> 183.67 ms (+1.90%, +0.53% to +3.14%), while its proving is not flagged, and the 4to1 8-thread node'sreduce-in-rows, 2.93 ms either way (+2.63%, +0.04% to +6.45%); on x86-64, Fibonacci'sexecute-program(+1.26%), which runs before any commit.How it composes with ntt-tasks (#501, now upstream): this branch decides where the deep pass starts (the L3 budget for one or two workers) and what the gathered passes take (adjacent residues, the in-place lane-major read); ntt-tasks' sizing then cuts the deep pass into tasks of whole sub-blocks, enough of them for 64 rows (
MIN_TASK_LOG = 6), and hands a task's rows to the Merkle tree inside the task unless one task covers the whole deep pass. Each sub-block runs this branch'srun_layers(leftover layers first, walked twiddles, register kernels) with its own block index, so the result is the same as one task a sub-block. A lane-major message never reaches a deep task: the first gathered pass reads it in place, or it is transposed into the first replica before any pass, so a deep task builds its sub-blocks from a row-major message only. Where the sub-blocks hold 64 rows or more, each task is one sub-block, as before.Correctness
Proof::to_bytes). After mergingf774828d,d936ad4f,556206d7andf7b03b5ethey match at 1, 2 and 32 workers, native (AVX-512) and with-C target-cpu=haswell; onf7b03b5eso do the leaf, first-level node and node proofs ofaggregate --n 100.cargo test --release -p pcspasses natively (AVX-512), with-C target-cpu=znver3(the AVX2 arm, VPCLMULQDQ without AVX-512), with the flag sets of CI's SIMD legs and with-C target-cpu=haswell, and on aarch64 (neoverse-n2,+aes,+sha3) under qemu.ARM
The NEON kernels are unchanged; the planner changes (leftover layer first, walked twiddles, the wide plan for one or two workers, which on aarch64 still transposes the stack first) apply there too. The aarch64 build compiles and passes clippy with CI's ARM flags and its NTT tests pass under qemu. CI's ARM bench on
556206d7(Neoverse-N2, above) has the commit 2.8% faster at one thread and 0.6% to 1.5% faster in the multi-threaded runs. Onf7b03b5e(Neoverse-V2, above) the commit is 4.8% faster at one thread, 1.6% at four and 1.2% at all threads.Checks
On
f7b03b5e:cargo fmt --all --check;cargo clippyallwith-C target-cpu=haswellon 1.97 and 1.99, native on 1.99, and with both x86-64-v3 flag sets of CI's SIMD legs; the system-allocator clippy step, the ARM64 clippy leg and the AVX-512 clippy step ofrust.yml;cargo docall;cargo test --release -p primitives -p pcs -p flocknatively (AVX-512), with-C target-cpu=haswelland with both x86-64-v3 flag sets, with the bus kernels' tests of those legs;cargo test --release -p pcs -p flockon aarch64 (neoverse-n2,+aes,+sha3) under qemu; the verifiers' end-to-end tests (cargo test --release -p leanvm_core --test verifiers, the Python verifier included); every CI job,cargo testallin the Test job included.