perf(sumcheck)!: redesign prover tables for runtime SIMD - #368
perf(sumcheck)!: redesign prover tables for runtime SIMD#368quangvdao wants to merge 72 commits into
Conversation
Add a coefficient-oriented constructor and specialize quartic slice packing so SIMD table transposes do not allocate temporary coefficient vectors. Extend the field benchmark with sumcheck fold and lazy tensor factor-pair layout controls.
Define the canonical coefficient-first evaluation table, binding order, runtime CPU selection, compact and sparse states, migration sequence, and measurable acceptance gates. Archive the compile-time packed-table proposal so there is one live implementation contract.
Add allocation-free ExtField coefficient primitives and a single-allocation coefficient-first table. Provide distinct row-preserving and multilinear binding-order constructors so dense and sparse callers cannot silently share the wrong row semantics.\n\nBREAKING CHANGE: ExtField implementations must provide from_base_fn and base_coefficient.
PR SummaryHigh Risk Overview
Protocol and tooling wire through Breaking: custom Reviewed by Cursor Bugbot for commit 261f975. Bugbot is set up for automated code reviews on this repo. Configure here. |
Documentation blast radius (advisory)These regions may need doc/spec/book updates based on changed paths. Changed files in this PR: 127
|
CI test timing
Run summary
Slowest testsNo JUnit data available for this run. Regressions vs mainNo per-test regressions above the threshold. New slow testsNo new tests ≥30s vs main baseline. |
Accumulate long fp32-extension small products in the full product accumulator without changing short-kernel arithmetic. Group stage-1 b8 prefix terms by their shared suffix weight to eliminate per-quad extension products.
# Conflicts: # crates/akita-prover/src/protocol/sumcheck/akita_stage3/mod.rs
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 1 potential issue.
Reviewed by Cursor Bugbot for commit a181130. Configure here.
Drop imports that became unused after the main branch merge so strict no-default Clippy remains clean.
Add one canonical four-product packed operation and reuse the existing exact fp32 backend accumulators in tensor projection maps. This removes three reductions per output coefficient and cuts the pinned Ice Lake dense-factor total by about 40 percent.
PCS Profile Benchmark
13 of 13 profiles passed. Times are medians of Each sample verifies the same proof first with the configured multi-threaded pool and then with one thread. Both timings reuse the same verifier setup. Merge-base comparisons are available for Benchmark shards
Public opening statements
One-hot profiles generate deterministic witnesses with one Direct evaluates the public setup contribution during Stage 2. Recursive carries the same check through a Stage 3 setup-product sumcheck. Both modes execute the complete fold schedule and terminal verification. The chunked profiles Generated profiles may select different A, B, and D ring dimensions at different fold levels. The short profile names omit those dimensions. Each sample generates deterministic witnesses and opening points, prepares setup, commits, proves, serializes the proof, checks its size, prepares verifier setup, and verifies the claimed openings. It does not test malformed proofs. Phase time
Memory and setup size
Proof size and protocol shape
Deltas are shown only for profiles with a matching merge-base case. Negative is smaller or faster. Terminal response components
The Detailed schedule and proof-size breakdowns by fold level are available in the uploaded |
Partition fp32 and fp64 folds and product rounds across Rayon workers while preserving runtime-selected field kernels. Keep fused packed tensor traversal for one worker and use parallel canonical construction for multiworker plans. Add fp64 tensor kernels and large differential coverage.
Resolve the prover streaming and one-hot cutover on main while preserving typed evaluation tables and multicore sumcheck kernels. Keep main's compact exact-prefix state and split profiling modules, and port the branch's live kernel bounds and compact SIMD dispatch.
Resolve the setup-prefix and late-reduction overlap while preserving the runtime sumcheck table kernels. Take main's canonical compact subfield opening APIs, which subsume the branch's high-half consistency operation.

Summary
This draft lands the canonical storage and design contract for a sumcheck prover that selects SIMD operations at runtime. It does not change the transcript or verifier.
It currently lands:
EvaluationTable<F, E>representation that works on every architecture.Vec<E>toEvaluationTable<F, E>and deletes the old dense traversal module.The normative design is
specs/sumcheck-kernel-architecture.md.Diff metadata
c9ca8e9dd5f19e234a7cd5e3c09fd4089b8454fd(main)ccb6b8a0a47782b1a00eb525c4c7a4ed6e40e7f5specs/sumcheck-kernel-architecture.mdMotivation
Before this PR, dense EOR stored extension field evaluations as
Vec<E>. That layout works for scalar arithmetic, but SIMD code must repeatedly gather or repack the coefficients. Ice Lake microbenchmarks show the cost. A fold over two persistent contiguous halves takes 56.619 microseconds for 16,384 output rows. Repacking the adjacent layout takes 372.140 microseconds. The scalar path takes 466.053 microseconds.In the pinned full fp32 profile, a D128 one hot proof with
nv=28takes 6.304 seconds. Sumchecks account for about 4.32 seconds, or 68 percent. The extension opening reduction accounts for 2.232 seconds. Stage 1 takes 1.400 seconds, and stage 2 takes 1.296 seconds.The previous packed sumcheck design tied packing to compile time types and storage that depended on the register width. It could not give scalar, ARM, AVX2, and AVX-512 code one durable representation. This PR archives that design and specifies runtime selection over one table layout.
Architecture
flowchart LR A[Protocol evaluations] --> B[EvaluationTable] C[Detected CPU features] --> D[Opaque kernel plan] B --> E[Scalar reference operations] B --> F[NEON, AVX2, or AVX-512 IFMA operations] D --> E D --> F E --> G[Unchanged round claims and transcript] F --> GChange overview
Vec<E>with interleaved coefficientsEvaluationTable<F, E>allocationCanonical evaluation storage
EvaluationTable<F, E>owns one flatBox<[F]>. It stores extension coefficients in contiguous slabs, with a private live length and fixed stride. Its constructors cover sparse values that preserve row order and dense multilinear values in bit reversed binding order. Folding can shrink the live length without reallocating or permuting later rounds.The type exposes live coefficient slices for arithmetic but keeps its allocation shape private. Tests cover extension degrees 1, 2, 4, and 8, constructor order, truncation, inactive storage, empty sparse inputs, and invalid dense lengths.
ExtField<F>now exposesfrom_base_fnandbase_coefficientas its primitive operations that do not allocate. The existing slice and vector operations derive from those primitives, which keeps coefficient conversion in one source of truth.fold_first_variable_scalaris the portable correctness oracle. It reads the two binding order halves, reuses the field's optimized scalar fold context, writes the first half, and truncates the live rows without allocation or permutation. Differential tests compare every round and the final multilinear evaluation with the existing logical order fold.Canonical product round
compute_product_round_scalarcomputes the constant and quadratic coefficients from two binding order tables. It usesw0 * f0for the constant term and(w1 - w0) * (f1 - f0)for the quadratic term. The existing previous claim formula still derives the linear coefficient.akita-sumchecknow ownsProductRoundAccumulatorand its two reduction policies. The delayed policy sums wide products and reduces once when the field declares that exact. The direct policy reduces every product before adding it. The scalar table operation and sparse EOR traversal use these same implementations. Dense EOR calls the table operations directly. The delayed accumulator rejects fields whose exactness flag is false.fold_and_compute_product_round_scalarfolds both tables and computes the next round in one pass. Each iteration reads the two children for the current variable and the two children for the next variable before it writes either folded output. Differential tests compare both mutated tables and both returned coefficients against separate scalar operations.Runtime CPU kernels and benchmark coverage
PackedFpExt4::from_coeff_fnand itspack_sliceimplementation remove temporary coefficient vectors. The benchmark suite now measures packing from the current adjacent layout, folding over persistent halves, root factor accumulation, and coefficient gathering separately.Portable x86 builds include the existing fp32 AVX2 and AVX-512 arithmetic without changing the compile time
HasPackingaliases for other code.akita-fieldowns the unsafe slice loops that require specific CPU features.akita-proverownsSumcheckKernelPlan, which detects one safe operation before entering a table loop. Its fields and operation choices are private, so a safe caller cannot run AVX2 or AVX-512 code on an unsupported host.The 512 bit operation is named and guarded as AVX-512 IFMA because its portable release assembly contains
vpmadd52luq. Detection and target functions require AVX-512F, AVX-512DQ, and AVX-512IFMA. A host without that exact set uses AVX2 for this operation.The x86 table loops validate all slice lengths once, then use raw pointers inside complete SIMD chunks. The safe prover operation supplies disjoint coefficient halves and handles short tables with the scalar oracle. Differential tests force every operation that the test host supports.
The runtime plan also selects the fp32 product round separately from the fold. Its packed loop computes the constant and quadratic products in every SIMD lane, sums those packed results across chunks, and reduces the lanes to two ordinary extension field values at the operation boundary. Portable assembly contains only YMM instructions in the AVX2 function. The AVX-512 IFMA function uses ZMM registers and
vpmadd52instructions.The fused fp32 operation folds both tables and computes the next product round in the same packed pass. It has its own private plan choice, so later measurements can select AVX2 for one operation and IFMA for another. Its loop reads all source chunks before writing either folded output half. Portable assembly confirms vector stores in both x86 functions.
The committed Criterion group contains separate rows for the detected runtime plan, AVX2, and AVX-512 IFMA. A preliminary Ice Lake binary built with global
target-cpu=nativefolded 16,384 output rows in 76.089 microseconds median. The matching scalar row took 466.02 microseconds, so this preliminary result is 6.12 times faster. This is not the acceptance result. The final comparison must use one portable binary and the explicit AVX2 and IFMA rows.These microbenchmarks are not a full proof performance claim. Dense EOR now uses the runtime operations in production. The exact portable Ice Lake end-to-end measurement is still pending.
The AArch64 path uses the same opaque whole-operation dispatch.
FpExt4<Fp32>selects NEON for fold, product, and fused fold plus next product round.FpExt2<Fp64>deliberately mixes operations: the initial product round stays scalar because it measured faster, while fold and fused rounds use NEON. Identity fields use direct degree-one slice kernels, which avoid reconstructing a field value through the generic extension interface.The shared packed traversal is independent of SIMD width. Thin target-feature wrappers supply the NEON, AVX2, or AVX-512 packed base field. Portable x86 builds compile fp64 AVX2 and AVX-512 candidates, but production x86 fp64 selection stays scalar until the pinned Ice Lake benchmark selects each operation independently.
Dense EOR cutover
ExtensionOpeningTables::Densenow owns coefficient-major witness and factor tables from construction through the final fold. Logical-order vectors are validated once, their unscaled input claim is computed with the existing canonical claim function, and they are converted once into binding order. Later rounds do not convert, repack, allocate, or retain a parallelVec<E>representation.ExtensionOpeningReductionTerm<F, E>andExtensionOpeningReductionProver<F, E>name the base and extension fields directly. Each prover detects oneSumcheckKernelPlanduring construction. Generic EOR code reaches the plan throughSumcheckTableOperations<F>. Its default methods are the canonical scalar operations.FpExt4<Fp32>andFpExt2<Fp64>override the dense operations with field-specific runtime choices. Identity fields use direct scalar slice operations. Other extension shapes keep the scalar defaults.The old 205-line dense EOR traversal module was deleted. Round zero calls the canonical product operation. Later rounds call the canonical fused fold and next product operation while another round remains. The final short fold uses the canonical fold operation. Sparse and cylindrical terms keep their current representation and behavior in this slice.
Dense field acceptance results
The matched benchmark uses 65,536 dense rows and one Rayon worker on Apple Silicon. Construction and cloning stay outside the measured proof interval.
The all-NEON fp64 path measured 0.89332 ms. Selecting the scalar initial product round and NEON fused rounds reduced it to 0.80775 ms. This is why the plan stores a choice per operation rather than one field-wide SIMD tier.
A generic coefficient-slice experiment was rejected. It regressed fp32 to 4.5623 ms, fp64 to 1.2935 ms, and fp128 to 1.0500 ms. The accepted code uses field-shaped kernels while keeping the table and protocol operation as the single public representation and call path.
Protocol and safety
Breaking changes
This is an intentional breaking API change.
ExtField<F>implementors must provide coefficient primitives that do not allocate. EOR terms and provers now name bothFandE, andinput_claim_from_termsreturns its infallible field value directly. Custom application fields implementSumcheckTableOperations<F>to use the scalar defaults or provide accepted optimized operations. No compatibility wrapper or duplicate table representation is added.The protocol, proof format, transcript, and verifier behavior remain unchanged.
Commit map
6d1bf5c89removes temporary FpExt4 packing allocations and adds layout microbenchmarks.d0b7cbba4specifies the canonical table, runtime selection, operation boundaries, rollout order, and performance gates.83f07e128adds extension coefficient primitives andEvaluationTable<F, E>.c92aad58eadds the scalar table fold without allocation and adds differential tests over every round.0433cb9b2adds opaque runtime selection, AVX2 and AVX-512 IFMA table folds, forced differential tests, and portable benchmark rows.d72c190beadds the scalar product round and moves the existing EOR accumulator policies to their canonical sumcheck owner.361e99068adds the fused scalar fold and next product round oracle.a854759c3adds runtime AVX2 and AVX-512 IFMA product rounds and matched benchmark rows.6552d7d7eadds runtime AVX2 and AVX-512 IFMA fused operations and matched benchmark rows.8b1e06f27moves dense EOR toEvaluationTable<F, E>, deletes the old dense traversal, and preserves validated input claims at the conversion boundary.93d477df8stores one runtime plan per EOR prover and dispatches production fp32 dense rounds through the selected operations.7bad074b4adds shared packed traversal, production NEON fp32 kernels, mixed scalar and NEON fp64 selection, direct identity-field kernels, and matched dense EOR benchmarks.3fae22a76propagates the required sumcheck operation bound through generic PCS tests and profile drivers.ccb6b8a0aseparates dispatch, fp32, fp64, portable scalar, and differential test ownership without changing the public API.Validation completed at
ccb6b8a0aThe following passed:
cargo machete,typos, and all documentation guardrails passed.akita-sumchecklibrary tests passed.akita-field,akita-sumcheck, andakita-prover.akita-pcsfield arithmetic benchmark.parallel,disk-persistencepasses with warnings denied.x86_64-unknown-linux-gnu.vpmadd52instructions inside the IFMA target functions.git diff --checkpassed.The
7bad074b4GitHub run exposed missing generic bounds in the end-to-end test and profile helpers before any benchmark ran.3fae22a76fixes those compile failures. Checks forccb6b8a0aare running.Specification status and remaining work
The specification is active. Dense EOR is now on the production table and runtime operation path. The remaining work is:
The Ice Lake acceptance target for the 16,384-row fp32 fold is at most 75 microseconds median, compared with the measured 466.053-microsecond scalar reference.
Reviewer map
Suggested review order:
specs/sumcheck-kernel-architecture.mdcrates/akita-field/src/ext/lift.rscrates/akita-sumcheck/src/evaluation_table.rscrates/akita-sumcheck/src/accum.rscrates/akita-prover/src/kernels/sumcheck.rscrates/akita-field/src/packed/runtime_common.rscrates/akita-field/src/packed/runtime_neon.rscrates/akita-field/src/packed/runtime_x86.rscrates/akita-field/src/packed/ext/mod.rscrates/akita-pcs/benches/extension_opening_reduction.rs