[fft] AVX2 Shoup butterflies with runtime dispatch - #65
Open
devin-ai-integration[bot] wants to merge 3 commits into
Open
[fft] AVX2 Shoup butterflies with runtime dispatch#65devin-ai-integration[bot] wants to merge 3 commits into
devin-ai-integration[bot] wants to merge 3 commits into
Conversation
Co-Authored-By: Andrew He <he.andrew.mail@gmail.com>
Contributor
Author
🤖 Devin AI EngineerI'll be helping with this pull request! Here's what you should know: ✅ I will automatically:
Note: I can only respond to comments from users who have write access to this repository. ⚙️ Control Options:
|
GCC Code Coverage Report📂 Overall coverage
|
Co-Authored-By: Andrew He <he.andrew.mail@gmail.com>
Co-Authored-By: Andrew He <he.andrew.mail@gmail.com>
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
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
Vectorizes the hot NTT loops for 32-bit prime moduli using Shoup multiplication (precomputed quotients
w' = floor(w·2^32 / MOD)stored alongsidert/inv_rt), avoiding Montgomery form entirely: data stays in the normal residue representation, lazily in[0, 2·MOD)between passes and normalized back to[0, MOD)on a transform's final pass — outputs are bit-identical to the scalar path.New AVX2 kernels in
fft_core(gated byntt_simd<num>, true formodnum<MOD>withMOD < 2^30on x86):forward_simd/inverse_simd: butterflies atk >= 8are 8-wide (mul_shoup= mulhi32 + 2 mullo,lazy_reduce=min_epu32(x, x - lim)); thek < 8tail stays scalar Shoup.scale_simd: constant-multiply, used byfft_engine::finish's 1/n scale for assign ops (finish_linearregains an assign fast-path so the cut fixup doesn't force the scalarcut_oproute).twiddle_simd: theb[i] = coeffs[i] * rt[n+i]pass inextend.The kernels carry
[[gnu::target("avx2")]]and are selected at runtime via__builtin_cpu_supports("avx2"), so no special compile flags are needed and non-x86 / non-AVX2 / other value types (goldilocks, cnum, dynamic modnum) keep the existing scalar loops untouched.Measured (g++ -O2, no arch flags, N=2^18): isolated butterflies ~0.58 → ~0.20 ns/weighted-op (~2.8×); end-to-end multipoint eval 325 → ~233 ms (−29%), interpolation 455 → ~337 ms (−26%). Remaining scalar hot loops are the variable×variable pointwise
mul/mul2passes, which Shoup can't cover (they'd need Montgomery-domain transforms).Link to Devin session: https://app.devin.ai/sessions/44c622b53d88445c9431863199c4fe7b
Requested by: @ecnerwala