Cache per-level FFT twiddle tables and batch evaluator slice math - #290
Closed
lamb356 wants to merge 4 commits into
Closed
Cache per-level FFT twiddle tables and batch evaluator slice math#290lamb356 wants to merge 4 commits into
lamb356 wants to merge 4 commits into
Conversation
Adds an ifma feature to zakura-pasta-curves with 8-way AVX-512 IFMA kernels for batched Montgomery multiplication, squaring, scaling, and a deferred-reduction dot product over Fp/Fq, with runtime CPU detection and scalar fallbacks. halo2 gains a TypeId-dispatched batch module used for iFFT divisor scaling and deferred inner products. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Proving-key twiddles now retain contiguous per-recursion-level twiddle tables so the butterfly combine step multiplies the right half against a contiguous table with the batched slice multiplier instead of a strided walk. The polynomial evaluator uses the batched elementwise multiply/square helpers for product, square, and Horner nodes, and processes one chunk per thread. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Contributor
|
Thank you! Extracted the FFT part of this into github.com//pull/324 , and added you as a commit co-author to it (appears in both PR and main's squash commit!) TYSM! |
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
2 participants
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.
Summary
Caches per-level FFT twiddle tables in the proving key and batches the evaluator's hot slice math through the Pasta batched field ops.
poly/domain.rs:ProvingKeyTwiddlesnow also retainsbase_inverse_tables/extended_forward_tables: per-recursion-level contiguous twiddle tables (butterfly_twiddle_tables), built once per proving key.BUTTERFLY_BATCH_MIN_HALF = 32multiply the right half against the cached contiguous table viabatch::mul_sliceinstead of strided per-elementtwiddles[(i+1)*chunk]reads; smaller nodes keep the scalar loop. Composes with the paired-chunk interleaving from Interleave paired field FFT chunks #284 (the paired helper carries the cached tables too).poly/evaluator.rs:EVALUATOR_CHUNKS_PER_THREAD = 1) — larger contiguous slices make the batched ops effective.assign_rotated_chunk+mul_rotated_chunk/batch::sqr_slice, dispatching to the IFMA slice kernels when available and to the scalar loop otherwise.The batched slice ops come from the IFMA PR (piece 2, #286); this branch includes that commit so it builds standalone — only the top commits are new here; it rebases cleanly once piece 2 merges. Without AVX-512 IFMA the batch dispatch falls back to scalar loops, and the twiddle caching + scheduling changes still help on their own.
Several ms end-to-end at 8 threads in the full stack. 16-core numbers welcome:
RAYON_NUM_THREADS=16 ORCHARD_K11_PROVER_THREADS=16 cargo bench -p zakura-orchard --features circuit,orbits --bench orchard_k11_proverPart of the split of #273 (piece 6 of 6, depends on piece 2, #286).
Testing
cargo test --release -p zakura-halo2-proofs(137 tests pass, FFT round-trip and proof tests included).cargo check --locked --all-features --all-targets, MSRV (1.91) check,cargo clippy(no new warnings vs main), andcargo fmt --all -- --checkclean.