ML-DSA on x86-64: sample  four entries at a time in signing (sign −11 to −20% instructions with AVX2) - #536
Merged
Merged
Conversation
vg_mldsa_ntt and vg_mldsa_inv_ntt now compute on four coefficients at a
time in SSE2 registers, as ML-KEM's x86-64 NTT does: a Montgomery
multiplication with pmuludq (the even doublewords, then the odd ones moved
down by pshufd), conditional additions of q with psrad masks, and for the
layers with len 2 and 1 the coefficients of two or four blocks gathered with
punpck{l,h}qdq (and pshufd) and interleaved back. The zetas are a table in
Montgomery form that the prologue stores in scratch; the multiplications run
inside ML-KEM's withMxcsr, so Intel's MCDT mitigation holds.
Proofs: the lanes' arithmetic (VArith), the butterflies on registers
(VLanes), the loads and stores of four coefficients and the zetas (VMem),
the layers (VLay, VLay21), and the functions (Ntt, NttInv). Signing and
verification now check that their code loads MXCSR only to restore it
(ctlOk, through a compositional ctlC for verification) rather than never,
since their primitives now do.
ML-DSA-65 on this machine: sign 1.53 ms -> 0.81 ms, verify 281 us ->
181 us, keygen 276 us -> 248 us.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
vg_mldsa_multiply_ntt, vg_mldsa_multiply_add_ntt, vg_mldsa_add and vg_mldsa_sub now compute on four coefficients at a time in SSE2 registers, with the vector helpers of the NTT (Vec.lean). The products are two Montgomery multiplications: f·g·2⁻³², then by 2⁶⁴ mod q, which is f·g in [0, 2q), reduced with vcsub (for multiply_add, h is then added and the sum reduced). They use pmuludq, so they run inside ML-KEM's withMxcsr. These functions have no working space and use no stack, so MXCSR goes through the last 8 bytes of h, addressed through r8 = h. The last four coefficients of h are loaded into xmm6 first. The loop stores the first 252 coefficients. The last four are computed from registers before MXCSR is loaded back, and stored after it. Their callers' proofs are unchanged: the functions still need no stack and never write rsp. add and sub use paddd/psubd and vcsub/vcadd. Proofs: the lanes of a product (mul_lane, mulAdd_lane), the loop (Mul.step, Mul.loop_ok), the last block (Mul.last), the function (Mul.fn_ok), and withMxcsr through the end of a polynomial (withMxcsrH_ok). AddSub is reproven the same way. ML-DSA-65 instructions executed (callgrind, PR #464 -> this): sign -19%, verify -11%, keygen -2%. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
…into claude/fervent-einstein-ukl7t7-mul
…er the arithmetic
Key generation, signing and verification on x86-64 now call the
polynomial arithmetic of an implementation given as a variant of the new
interface MlDsaArith (Variants/MlDsaArith/X86_64/). They are registered
in Generic/MlDsaArith/X86_64/ and emitted once for each implementation,
so a faster arithmetic (e.g. AVX2) reaches them without editing them. The
only variant is the existing SSE2 code (Sse2), and the generated code is
unchanged.
- Impl: Arith.Backend holds the six functions' code and a name suffix.
Each Prims gets a sfx field, and the arithmetic calls are named
"vg_mldsa_ntt" ++ sfx and so on. primsWith B swaps in B's arithmetic.
- Proof: ArithImpl is a backend with FnOk for each function: verified
without stack, no rsp writes, depth <= 2, ctlOk, spSafe. prims_okWith
builds each caller's PrimsOk from it.
- ctlOk and spSafe of signing and key generation are no longer evaluated
on the whole code with a concrete backend. Same/same_tac
(Proof/MlDsa/X86_64/Arith/Same.lean) shows that a check composing over
the code's structure (ctlC, Code.allInstrs) gives the same result as
on the code with every arithmetic function empty, which the kernel
evaluates once per parameter set. Verification already did this by
hand; its call lemmas now allow the names to differ.
- The registration files Artifacts/MlDsa{KeyGen,Sign,Verify}/X86_64.lean
move to Generic/MlDsaArith/X86_64/.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
…de/fervent-einstein-ukl7t7-generic Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
Add an AVX2 variant of MlDsaArith on x86-64: vg_mldsa_ntt_avx2, vg_mldsa_inv_ntt_avx2, vg_mldsa_multiply_ntt_avx2, vg_mldsa_multiply_add_ntt_avx2, vg_mldsa_add_avx2 and vg_mldsa_sub_avx2 compute on eight coefficients at a time, four in each 128-bit lane of an AVX2 register. In each lane the VEX.256 form of the SSE2 code does what the SSE2 code does to an xmm register, so the proofs lift the SSE2 lemmas to each lane (ML-KEM's ylanes). The NTT layers with len >= 8 load eight coefficients of each half of a block; len = 4 regroups two blocks with vperm2i128; len = 2 and 1 run the SSE2 gatherings in each lane, with the zetas of each lane arranged by vpshufd/vpblendd or vpermq. Key generation, signing and verification are generic over MlDsaArith, so the emitter generates their _avx2 instances; the Rust API selects them on CPUs with AVX and AVX2 (mldsa_common::Backend). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
The contract of a function that samples four elements of the matrix A at once, as ML-KEM's vg_mlkem_sample_ntt4 does: for each of the four 34-byte seeds, RejNTTPoly (FIPS 204 Algorithm 30) of it, reduced, or 0 if the loop does not finish within Appendix C's least bound for one of them. Its scratch is that of vg_mldsa_rej_ntt_poly (256 u64s), so callers can pass the same working space. No implementation yet: an x86-64 one, with four SHAKE128 instances in AVX2 registers, follows in its own PR. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
Four interleaved Keccak states, the second buffer of the 4-way permutation and its table of round constants already take 2368 bytes, before the squeezed output; ML-KEM's vg_mlkem_sample_ntt4 has 1024 u64s for the same. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
ML-DSA now chooses between its SSE2 and AVX2 polynomial arithmetic by CPU feature, so test it end to end where the choice differs: under SDE's Pentium 4 (the SSE2 code, and the line that chooses it) and Haswell. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
…rvent-einstein-ukl7t7-rej4
vg_mldsa_rej_ntt_poly4 runs RejNTTPoly on four seeds. Its AVX2 variant runs the four SHAKE128 instances at once in the 64-bit elements of ymm registers, with the absorb and squeeze code of vg_mlkem_sample_ntt4_avx2, and runs vg_mldsa_rej_ntt_poly's loop over the same 1008 bytes of each seed's output (squeezed in two rounds of three blocks). The baseline calls vg_mldsa_rej_ntt_poly on each seed. Both are proven against rejNTT4Contract with 24 bytes of stack, and are a new function of the MlDsaArith backend. Verification now samples each row of  four entries at a time from SB4 (four copies of ρ with the row's indices): the entries from 0, then for ℓ = 7 the last four (sampling entry 3 again), or for ℓ = 5 the last one alone. Its rej4 working space is the 8 KiB after the last row of Â, which the scratch size already allows. Its stack grows from 24 to 32 bytes, as its baseline callee calls three deep. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
Key generation calls vg_mldsa_rej_ntt_poly4 (four SHAKE128 instances at once with AVX2, four calls of vg_mldsa_rej_ntt_poly without) for each group of four consecutive entries of Â, from four seeds at oSA4 that each hold ρ and the indices of their entry, with its 8 KiB of working space after the polynomials; the last kℓ mod 4 entries are sampled one at a time as before. Each call's result is ANDed into r15 and masks the four polynomials it sampled, as for one entry. Its calls now use 24 bytes of stack, within the 32 its contract already allows. Instructions per key generation (callgrind, AVX2): ML-DSA-44 23.5M → 16.8M (−29%), ML-DSA-65 38.4M → 26.7M (−31%), ML-DSA-87 64.4M → 40.8M (−37%); the baseline is unchanged (+0.06%). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
Signing copies ρ to four seeds at RS4 and, for each group of four consecutive entries of Â, sets their indices and calls vg_mldsa_rej_ntt_poly4 (four SHAKE128 instances at once with AVX2), with 8 KiB of working space after the polynomials, ANDing its result into r15; the last kℓ mod 4 entries are sampled one at a time as before. Signing branches on r15, so the proofs need, as of vg_mldsa_rej_ntt_poly, that the result depends only on the seeds and that it is 1 only if RejNTTPoly finishes within maxBounds: both implementations return whether each seed has 256 coefficients in 1008 bytes of output (Rej4Ok.ret, from their proofs' contract r4K). Its calls now use 32 bytes of stack (24 for vg_mldsa_rej_ntt_poly4 and the return address), as key generation's and verification's. Instructions for 20 deterministic signatures (callgrind, AVX2): ML-DSA-44 60.4M → 53.6M (−11%), ML-DSA-65 85.6M → 73.8M (−14%), ML-DSA-87 117.1M → 93.5M (−20%); the baseline is unchanged (+0.02%). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
vg_mldsa_high_bits_avx2 and vg_mldsa_low_bits_avx2 compute on eight coefficients at a time: in each 128-bit lane, the VEX.256 form of SSE2 code on four doublewords (hbX, lbX). Decompose is the reference implementation's, in 32 bits, with the multiplications by M and by 2γ₂ sums of shifts (no multiplication instructions, so nothing MCDT affects), r₁ = f mod m as f ANDed with the sign of f - m, and r₀ plus q if negative. They branch once on the public γ₂. They join the polynomial arithmetic's backend, as variants of vg_mldsa_high_bits and vg_mldsa_low_bits, so signing with AVX2 calls them. The proofs: what the SSE2 code computes in a doubleword is r₁ and r₀ of Decompose (YLane), each lane does it (YBlock, ylanes), and the loop stores all 256 (YBits). Instructions for 20 deterministic signatures (callgrind, AVX2): ML-DSA-44 60.7M → 58.2M (−4.1%), ML-DSA-65 86.2M → 82.2M (−4.6%), ML-DSA-87 118.3M → 113.8M (−3.8%); the baseline is unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
…t-einstein-ukl7t7-ybits2 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
…t-einstein-ukl7t7-rej4 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
…t-einstein-ukl7t7-rej4kg Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
…ent-einstein-ukl7t7-rej4sg Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
Also fixes the conflict markers the previous merge left in ci.yml: the x86-64 CPU-feature matrix runs both poly1305's and mldsa's tests. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
…t-einstein-ukl7t7-ybits2 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
…t-einstein-ukl7t7-rej4 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
…t-einstein-ukl7t7-rej4kg Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
…ent-einstein-ukl7t7-rej4sg Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
…t-einstein-ukl7t7-rej4kg Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
…ent-einstein-ukl7t7-rej4sg Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
…t-einstein-ukl7t7-ybits2 Both add functions to the polynomial arithmetic's Backend: the merged Backend has highBits, lowBits and rej4 (and BackendOk their proofs). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
Main has #534's vg_mldsa_rej_ntt_poly4, which this branch already merged with its own Backend fields (highBits, lowBits and rej4). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
…ent-einstein-ukl7t7-rej4sg Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
…tructions) (#519) * ML-DSA on x86-64: the norm check and MakeHint with AVX2 vg_mldsa_norm_lt_avx2 and vg_mldsa_make_hint_avx2 compute on eight coefficients at a time, in the VEX.256 form of SSE2 code on the four doublewords of each 128-bit lane, as HighBits and LowBits do. norm_lt clamps the public bound to q (which changes no result), so that a - b and (q - b) - a fit in 32 bits and one of them is negative exactly when a coefficient is out of bounds; it ANDs their ORs into an accumulator, spreads its sign bits and returns (vpmovmskb + 1) >> 32. There is no branch on data. make_hint computes r₁ of r and of (r + z) mod q with hbX, the hint as (r₁ ^ r₁' + 63) >> 6, and counts the eight hints of an iteration from the byte mask of the hints shifted to bit 7 (vpmovmskb), summing its nibbles with three shifts and additions. Both join the polynomial arithmetic's backend, as variants of vg_mldsa_norm_lt and vg_mldsa_make_hint, so signing (both) and verification (norm_lt) with AVX2 call them. The proofs: YNorm (the differences, the mask and the result) and YHint (the hint in a doubleword, the count, the loop, and the count is hintOnes). Instructions (callgrind, AVX2), 20 deterministic signatures: ML-DSA-44 58.2M → 53.5M (−8.0%), ML-DSA-65 82.2M → 75.5M (−8.2%), ML-DSA-87 113.8M → 106.1M (−6.8%); 20 verifications: 23.5M → 23.0M (−2.0%), 35.7M → 35.4M (−0.8%), 72.0M → 70.2M (−2.4%). The baseline is unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr * Regenerate src/asm after merging claude/fervent-einstein-ukl7t7-ybits2 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr --------- Co-authored-by: Claude <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
The Backend has highBits, lowBits, normLt, makeHint and rej4, and sign_same takes all of them. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
github-merge-queue
Bot
removed this pull request from the merge queue due to a conflict with the base branch
Oct 2, 2026
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
Regenerates the ML-DSA asm, which conflicted with the 8-byte copy. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
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.
Stacked on #535; the diff shows only this change.
What
Signing copies ρ to four seeds at
RS4. Then, for each group of four consecutive entries of Â, it:vg_mldsa_rej_ntt_poly4(four SHAKE128 instances at once with AVX2), with 8 KiB of working space after the polynomials;r15.The last kℓ mod 4 entries are sampled one at a time, as before.
Signing branches on
r15. So the proofs need ofvg_mldsa_rej_ntt_poly4what they already need ofvg_mldsa_rej_ntt_poly:rej4Ret);RejNTTPolyfinishes withinmaxBoundson each seed (rej4Max).Both implementations return whether each seed gets 256 coefficients from 1008 bytes of output. This follows from their proofs' contract,
r4K(Rej4Ok.ret). Signing's calls now use 32 bytes of stack (24 forvg_mldsa_rej_ntt_poly4plus the return address), the same as key generation and verification.Proofs (untrusted)
Sign/PhaseA4(new): the seeds of a group (slot_ok), the call (call4_ok,and4_ok), andExpandA(expandA_ok).PhaseACT: its leakage.PrimsB:rej4Call_ok/_tr.Prims,Inst: therej4callee withrej4Retandrej4Max.Verified:sign_samewithrej4.There are no TCB or Spec changes.
Performance
Instructions for 20 deterministic signatures (callgrind, AVX2):
The baseline is unchanged (+0.02%).
Checks run
lake buildand the emittercheck_lean_imports,check_lean_speed,check_vectors,check_arch_gates,check_variants,check_mcdt,algorithms_tablecargo fmt --checkandcargo clippy --all-targets -D warningscargo test --releasewith Wycheproof, both on the default path and withVG_CPU_FEATURES=none --features cpu-features-env🤖 Generated with Claude Code
https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
Generated by Claude Code