Skip to content

ML-DSA on x86-64: copy 8 bytes at a time in signing (sign −4 to −7% instructions) - #523

Merged
alex merged 39 commits into
mainfrom
claude/fervent-einstein-ukl7t7-copy8
Oct 2, 2026
Merged

alex merged 39 commits into
mainfrom
claude/fervent-einstein-ukl7t7-copy8

Conversation

@alex

@alex alex commented Oct 1, 2026

Copy link
Copy Markdown
Member

Stacked on #519, so that the generated src/asm/x86_64/mldsa{44,65,87}.rs don't conflict. Only the last commit is new.

It changes implementation and proof only; nothing in Spec/ or TCB/.

Why

A callgrind profile of signing with AVX2 found that 98% of the instructions in the body of vg_mldsa65_sign_avx2 (outside its callees) are in byte-by-byte copy loops. Each copied byte costs 6 instructions. The largest copy is y to ŷ, 1024 bytes on every iteration of the rejection loop. In total that is about 8% of a signature.

What

  • The copy. Every copy in signing has a length that is a multiple of 8: 1024 bytes, 32 bytes, and cLen (32, 48 or 64). Signing's copy (Impl/MlDsa/X86_64/Sign/Frag.lean) now moves one quadword per iteration, with n / 8 iterations.
  • The layout check. copyChk now also requires n % 8 = 0. It is decided per parameter set, as before.
  • The proofs.
    • copy_ok (in Sign/Blocks.lean; until now signing used ML-KEM's byte copy proof) and copy_tr (in Sign/BlockTr.lean) prove the quadword loop. Their invariants are the same as before.
    • copy_byte gives each byte of an 8-byte store of an 8-byte load, from writeW_byte, byte_readW and Offset.sub_toNat'.
    • inRegions_off gives the 8-byte subrange of a region that each iteration accesses.
  • Key generation and verification have copies of their own, which this PR doesn't change.

Measurements

Instructions for 20 deterministic signatures (callgrind), compared with #519:

AVX2 baseline ISA
ML-DSA-44 53.5M → 49.7M (−7.2%) 78.3M → 74.5M (−5.0%)
ML-DSA-65 75.5M → 70.4M (−6.7%) 110.9M → 105.8M (−4.6%)
ML-DSA-87 106.1M → 100.2M (−5.6%) 148.9M → 143.0M (−4.0%)

Checks

Run locally:

  • lake build and Emit.lean
  • the six ci/check_*.py scripts
  • cargo fmt and cargo clippy
  • cargo test with Wycheproof, with default features and with VG_CPU_FEATURES=none --features cpu-features-env

The README algorithm table doesn't change, because the description of the optimizations is unchanged.

🤖 Generated with Claude Code

https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr


Generated by Claude Code

claude added 16 commits October 1, 2026 13:33
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
…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
…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
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
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
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
Signing's copies (y to ŷ, 1024 bytes in every iteration, and the
32-byte seed and c̃) moved one byte per iteration: 6 instructions a byte,
about 8% of the instructions of a signature with AVX2. All of them are
of a multiple of 8 bytes, so `copy` now moves a quadword per iteration,
and `copyChk` checks that the length is one.

copy_ok and copy_tr are the byte copy's proofs with the quadword
loop: each iteration writes the 8 bytes it read (copy_byte, by
writeW_byte and byte_readW).

Instructions for 20 deterministic signatures (callgrind):
AVX2: ML-DSA-44 53.5M → 49.7M (−7.2%), ML-DSA-65 75.5M → 70.4M (−6.7%),
ML-DSA-87 106.1M → 100.2M (−5.6%); baseline ISA: 78.3M → 74.5M (−5.0%),
110.9M → 105.8M (−4.6%), 148.9M → 143.0M (−4.0%).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
claude added 13 commits October 1, 2026 22:40
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
…ent-einstein-ukl7t7-ynorm

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
…nt-einstein-ukl7t7-copy8

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
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
…ent-einstein-ukl7t7-ynorm

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
…nt-einstein-ukl7t7-copy8

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
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-ynorm

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
…nt-einstein-ukl7t7-copy8

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
claude and others added 7 commits October 2, 2026 01:11
…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
…ent-einstein-ukl7t7-ynorm

The Backend has highBits, lowBits, normLt, makeHint and rej4.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
…nt-einstein-ukl7t7-copy8

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>
Base automatically changed from claude/fervent-einstein-ukl7t7-ynorm to claude/fervent-einstein-ukl7t7-ybits2 October 2, 2026 02:58
claude added 2 commits October 2, 2026 03:22
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
…ent-einstein-ukl7t7-copy8

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
Base automatically changed from claude/fervent-einstein-ukl7t7-ybits2 to main October 2, 2026 03:53
@alex
alex added this pull request to the merge queue Oct 2, 2026
Merged via the queue into main with commit b97c94a Oct 2, 2026
44 checks passed
@alex
alex deleted the claude/fervent-einstein-ukl7t7-copy8 branch October 2, 2026 04:02
alex pushed a commit that referenced this pull request Oct 2, 2026
Main has #513, #519 and #523, which this branch already has.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Development

Successfully merging this pull request may close these issues.

2 participants