ML-DSA on x86-64: RejNTTPoly four at a time, in verification (verify −28% instructions with AVX2)

[alex]

Based on main, now that #491 (AVX2 polynomial arithmetic) has merged. It implements vg_mldsa_rej_ntt_poly4, whose spec and contract were reviewed in #492. #535, #536, #537 and #544 are stacked on it.

What

vg_mldsa_rej_ntt_poly4 runs RejNTTPoly on four 34-byte seeds and writes four polynomials (Impl/MlDsa/X86_64/Sample/RejNtt4.lean). There are two implementations:

• _avx2. It runs the four SHAKE128 instances at once, one in each 64-bit element of the ymm registers, using the absorb and squeeze code of vg_mlkem_sample_ntt4_avx2. It squeezes three blocks of each instance twice, then runs vg_mldsa_rej_ntt_poly's loop over the same 1008 bytes of each seed's output.
• Baseline. It calls vg_mldsa_rej_ntt_poly on each seed, saving its caller's callee-saved registers in scratch.

Both are proven against rejNTT4Contract with 24 bytes of stack. They are a new function of the MlDsaArith backend, so the generic callers get the _avx2 variant with nothing listed by hand.

Verification. Each row of  is now sampled four entries at a time:

• the seeds come from SB4, four copies of ρ with the row's indices;
• the first call covers entries 0–3;
• for ℓ = 7, a second call covers entries 3–6, so entry 3 is sampled again;
• for ℓ = 5, the last entry is sampled alone.

The rej4 working space is the 8 KiB after the last row of Â, which the scratch size already allows. Verification's stack grows from 24 to 32 bytes, because the baseline callee calls three deep.

Proofs (untrusted)

• Rej4Top, Rej4Sq, Rej4Parse: correctness of the AVX2 variant.
• Rej4Scalar: correctness of the baseline (four calls).
• Rej4CT: constant time; the result depends only on the seeds.
• Rej4Verified: Verified against the shared contract.
• Verify/StageA, CallSample, CTSample, Flag: the new sampling of each row, and its leakage.
• The single-poly samplers' proofs only gain the scratch frame facts the four-way callers need.

There are no TCB or Spec changes.

Performance

Instructions for key generation, one signature and 20 verifications (callgrind; base is #491's head):

	AVX2	baseline (VG_CPU_FEATURES=none)
ML-DSA-44	23.7M → 16.9M (−28%)	+0.06%
ML-DSA-65	35.7M → 25.6M (−28%)	+0.04%
ML-DSA-87	72.8M → 52.0M (−29%)	83.5M → 89.6M (+7.3%)

The baseline ML-DSA-87 slowdown is the duplicated entry 3 of each row: without AVX2 it costs a whole extra vg_mldsa_rej_ntt_poly call. #544, at the top of this stack, samples verification's  as one run of kℓ entries, as key generation and signing do. That removes the duplicate (baseline ML-DSA-87 −9.3% per verification) and saves another ~8% with AVX2.

Checks run

• lake build and the emitter
• check_lean_imports, check_lean_speed, check_vectors, check_arch_gates, check_variants, check_mcdt, algorithms_table
• cargo fmt --check and cargo clippy --all-targets -D warnings
• cargo test --release with Wycheproof, both on the default path and with VG_CPU_FEATURES=none --features cpu-features-env

🤖 Generated with Claude Code

https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr

[alex]

merge conflicts

[alex]

Merged main in 8cc896c. The conflicts came from #491's squash: its files were add/add against this branch's copies, which already had #491's final content plus this PR's changes. So the only new content is #529's Rust changes, and the README tables are regenerated. #535, #536, #537 and #544 are updated on top.

---

Generated by Claude Code