ML-DSA on x86-64: polynomial arithmetic in AVX2 (sign −23%, verify −10% instructions)

[alex]

Stacked on #476, which makes key generation, signing and verification generic over the arithmetic (and #476 is stacked on #467). The base here is #476's branch, so the diff shows only this change. I'll retarget it as those merge.

What

This adds 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 (Impl/MlDsa/X86_64/Arith/Avx2.lean). Each works on eight coefficients at a time, four in each 128-bit lane.

• Per-lane reuse. In each lane, the VEX.256 form of the SSE2 code (toY) does what the SSE2 code does to an xmm register.
• NTT layers with len ≥ 8. Each iteration loads eight coefficients of each half of a block.
• len = 4. Two blocks at a time, regrouped with vperm2i128.
• len = 2 and len = 1. The SSE2 gatherings run in each lane. The zetas for each lane are arranged with vpshufd/vpblendd (yzetaS), or with vpermq 0xD8 (forward) or 0x27 + vpshufd 0xB1 (inverse).
• Multiplications. They keep the SSE2 code's MXCSR handling: withMxcsr saves MXCSR through the last 8 bytes of h, so there is no stack. The products are inside the MCDT prologue and epilogue (check_mcdt passes).
• vzeroupper runs before every return.

Because #476 made keygen, sign and verify generic, the emitter generates vg_mldsa{44,65,87}_{keygen,sign,verify}_avx2 with nothing listed by hand (Variants/MlDsaArith/X86_64/Avx2.lean). On the Rust side, mldsa_common::Backend (Scalar / AArch64 Sha3 / x86-64 Avx2) wraps the Keccak backend's choice. It selects Avx2 when the CPU has the generated _AVX2_FEATURES (avx, avx2). Every call site matches on it exhaustively.

ci.yml's x86-64 CPU-feature matrix now runs the mldsa tests as well. The p4p chip covers the SSE2 path, and hsw and native cover AVX2. Without this, no CI job took the SSE2 branch of Backend::select, and coverage failed on it.

Proofs (untrusted)

• YBase, YAddSub, YMul: lanes of 256-bit loads and stores; add/sub; products (the h-slot MXCSR design as in ML-DSA on x86-64: polynomial arithmetic in SSE2 (sign −19%, verify −11% instructions) #467, now with ymmHi kept by withMxcsrH_ok).
• YLay: layers with len ≥ 8 (ystep, yblock_ok, ylay_ok) and len = 4 (ystep4, ylay4_ok), each layer coefficient by coefficient (layF_get). Also the zeta loads (yzeta1_ok, yzetaS_ok).
• YLay21: one generic iteration for len = 2 and 1 (ystep21), the per-lane gather → butterfly → scatter (core2_ok, core1_ok), vpermq 0x27 (ypermq27_ok), and the eight-zeta loads (yzeta8_ok, yzeta8R_ok).
• YNtt: prologue, yscale, and nttY_correct / nttInvY_correct. Constant time is proved by taint_decide, then Verified follows.
• BackendAvx2: ArithImpl.avx2, with features := ["avx", "avx2"].

There are no changes to TCB or Spec. One small change to existing untrusted code: withMxcsrH_ok now also says the body's ymmHi is kept, and Mul.lean is adjusted to match.

Performance

Instructions in each verified function, ML-DSA-65 (callgrind, sign_deterministic, 21 calls each). Both columns are on the same CPU: SSE2 is VG_CPU_FEATURES=bmi1,bmi2,adx,sha, AVX2 is the default.

	SSE2	AVX2
keygen	39.0M	36.5M	−6.4%
sign	107.4M	83.1M	−22.6%
verify	32.5M	29.1M	−10.3%

The rest of keygen is mostly Keccak and sampling, which this PR doesn't change. docs/algorithms/ml-dsa-*.toml now say "SSE2 and AVX2 polynomial arithmetic", and the README tables are regenerated.

Checks run

• lake build and lake env lean --run Emit.lean --check
• check_lean_imports, check_lean_speed, check_vectors, check_arch_gates, check_variants, check_mcdt, algorithms_table
• cargo fmt --check
• cargo clippy --all-targets -D warnings for x86_64, i686 and aarch64, each with and without cpu-features-env
• cargo test --release with Wycheproof, both on the default AVX2 path and with VG_CPU_FEATURES=none --features cpu-features-env: all pass

🤖 Generated with Claude Code

https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr

[alex]

merge conflicts

[alex]

Merged main in 4563920. The conflicts were in Arith/Mxcsr.lean and Arith/Mul.lean (kept this PR's ymmHi additions) and in the ml-dsa-*.toml notes; the README tables are regenerated. lake build, the emitter (src/asm unchanged), all ci/check_*.py and cargo test with and without VG_CPU_FEATURES=none pass locally.

---

Generated by Claude Code