ML-DSA on x86-64: key generation, signing and verification generic over the polynomial arithmetic - #476
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
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 #467 (which is stacked on #464). Its base is #467's branch, so the diff here shows only this change.
What
This is a refactor that prepares for AVX2 ML-DSA arithmetic. The generated code is unchanged: the emitter rewrites
src/asm/identically.Today x86-64 key generation, signing and verification call a fixed set of arithmetic functions (
vg_mldsa_ntt,vg_mldsa_inv_ntt,vg_mldsa_multiply_ntt,vg_mldsa_multiply_add_ntt,vg_mldsa_add,vg_mldsa_sub). This PR makes them generic over an implementation of those six functions, following CLAUDE.md's "An optimized implementation reaches everything built on it":MlDsaArithon x86-64. Its only variant is the existing SSE2 code (Variants/MlDsaArith/X86_64/Sse2.lean).Artifacts/MlDsa{KeyGen,Sign,Verify}/X86_64.leantoGeneric/MlDsaArith/X86_64/. They are emitted once per variant, named with its suffix and requiring its CPU features.Impl
Arith.Backendholds the six functions' code and a name suffixsfx.Primsgets ansfxfield (default""), and the arithmetic calls are named"vg_mldsa_ntt" ++ sfxand so on.primsWith Bswaps in backendB's arithmetic.Proof (untrusted)
ArithImpl(Proof/MlDsa/X86_64/Arith/Backend.lean) is a backend plus, for each function,FnOk: verified against its contract with no stack, never writesrsp, depth ≤ 2,ctlOk,spSafe.prims_okWithbuilds each caller'sPrimsOkfrom it.ctlOkandspSafeof signing and key generation were checked bydecide +kernelon the whole program, which needs a concrete backend.Proof/MlDsa/X86_64/Arith/Same.leanreplaces that:ctlC,Code.allInstrs q) gives the same result on the code with every arithmetic function empty (Same).same_tacproves this from the structure ofsign/keyGenfor any parameter set.SameC/SameQ). Itscalllemmas now allow the two call names to differ.There are no changes to TCB or Spec.
Checks run
lake buildlake env lean --run Emit.lean(no diff insrc/asm/)check_lean_imports,check_lean_speed,check_vectors,check_arch_gates,check_variants,check_mcdtalgorithms_table.pyNext
The AVX2 variant (NTT/NTT⁻¹, multiply(-add), add/sub on 8 lanes) plus the Rust backend dispatch will follow in a separate PR. I prototyped it in the Lean model: it matches the spec, and with it ML-DSA-65 executes −27% instructions in signing and −20% in verification relative to #467.
🤖 Generated with Claude Code
https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
Generated by Claude Code