[alex]

Part of #991.

What

vg_mlkem_compress_encode and vg_mlkem_decode_decompress on AArch64 (Impl/MlKem/AArch64/Compress.lean), the codecs of ML-KEM-768's ciphertexts and messages (d = 1, 4 and 10), move from a scalar loop over groups to NEON. They keep their Apis, contracts and register usage. Coefficients stay u32, four to a register.

Compression. Compress_d(x) of a lane is still the scalar code's ((x · M_d + 262080) >> 19) mod 2ᵈ, as MUL, ADD, USHR and AND (cmpV), so its proof reuses compress_eq.

• Neighbouring lanes are paired, x₂ₑ + 2ᵇ x₂ₑ₊₁ (UZP1, UZP2, SHL, ORR; pairV). This repeats until a lane holds what one or more bytes encode:
  • d = 4: one level, then UZP1 narrows words to bytes (16 bytes per 32 coefficients);
  • d = 1: three levels (4 bytes per 32 coefficients);
  • d = 10: one level, then USHR and SLI put four values in a doubleword (dpairV), and TBL picks each group's five bytes (20 bytes per 16 coefficients).

Decompression. Decompress_d(y) of a lane is still (q · y + 2ᵈ⁻¹) >> d (MUL, ADD, USHR; decV).

• d = 4: AND and USHR split the nibbles of 16 bytes, ZIP1 and ZIP2 interleave them, and ZIPs with zero widen them to words.
• d = 1: each byte is broadcast (DUP), and its bit i is masked in word i (AND). CMEQ with zero and BIC turn the word into 0 or 1665, which is Decompress₁ of the bit.
• d = 10: TBL puts each group of five bytes in a doubleword. USHR, TRN1, the ANDs and the ZIPs split it into its four values (dd10Half).

Dynamic instructions per polynomial (counted on the model, Prog run on random and edge inputs, all matching Spec.MlKem.compressEncode / decodeDecompress):

	scalar	NEON
compress_encode, d = 1	1,984	620
compress_encode, d = 4	2,560	525
compress_encode, d = 10	2,688	681
decode_decompress, d = 1	1,760	337
decode_decompress, d = 4	2,432	438
decode_decompress, d = 10	2,752	654

The code uses v0–v7 and v16–v23, none of the registers the procedure call standard preserves, so the KEM callers (which use CE.correct / DD.correct, whose statements are unchanged) need no change.

ML-KEM-1024's d = 5 and d = 11 codecs are separate functions (Impl/MlKem1024/AArch64/Compress.lean) and stay as they are: the 16-bit K-PKE cores (#2763) will replace them with NEON codecs of 16-bit coefficients.

No TCB/ or Spec/ changes.

Constant time

Both functions are proven constant time with the usual taint_decide (CE.ct, DD.ct): every address and branch depends only on the pointers and len, and the coefficients and bytes (secret in decapsulation: Compress₁ of m′, Decompress₁ of the message) only flow through vector arithmetic.

Every instruction used is on Arm's DIT list (DDI 0487, "About PSTATE.DIT", FEAT_DIT): ADD (vector), AND, BIC, CMEQ (register), DUP (element and general), INS (general), MOVI, MUL (vector), ORR, SHL, SLI, TBL, TRN1, UMOV, USHR, UZP1, UZP2, ZIP1, ZIP2, the loads and stores LDR/STR (Q, W), and the scalar MOVZ/MOVK and ADD/SUB (immediate). The only branches are CBZ/CBNZ on the public length and the loop counters. The same vector instructions already carry secrets in the existing NEON code (the 16-bit NTT, vg_mlkem_encode12/decode12, CBD), under the same assumption that PSTATE.DIT is 1 (README, "AArch64 needs PSTATE.DIT set").

Proof

Proof/MlKem/AArch64/CodecV/:

• Lane, Narrow, Dword: the lanes of each operation, as numbers (cmpV_ok, pairV_ok, decV_ok, dpairV_ok, narrow4, …).
• Blocks: loads of coefficients (ld4_ok) and four compressions (cmp4_ok).
• Ce1, Ce4, Ce10: each width's arithmetic; Ce1Step, Ce4Step, Ce10Step: an iteration, against CE.L_eq.
• Dd4, Dd10: the decompression of d = 4 and the halves of d = 10; Dd1: a byte of d = 1 and its iteration; Dd4Step, Dd10Step: the other iterations, against DD.G_eq.
• CeLoop/CeRun, DdLoop/DdStore: the loops' invariants and constants, once for every width.
• CeSpec, DdSpec: the byte and coefficient lemmas of the scalar proofs, kept.

CompressEncode.lean and DecodeDecompress.lean keep the contracts, correct, ct and _verified, now dispatching to the vector loops.

Performance

One scoped run on the N2 runners (run 39130771921): arches=aarch64, the four configurations bench_arches lists, modules=mlkem (ML-KEM-768 and ML-KEM-1024), against 0bac9044. Fastest of 8 copies of each binary; every spread is 0.3% or less, but one (mlkem768_encaps_from_bytes, head, neon: 2.4%).

Benchmark	default	sha3	neon	none
mlkem768_decaps	−4.5%	−4.5%	−3.5%	−3.5%
mlkem768_encaps	−3.1%	−3.1%	−2.1%	−2.1%
mlkem768_encaps_from_bytes	−3.0%	−3.0%	−2.0%	−2.0%
mlkem768_keygen	−0.1%	0.0%	0.0%	0.0%
mlkem1024_decaps	−0.4%	−0.5%	−0.3%	−0.3%
mlkem1024_encaps	−0.3%	−0.4%	−0.2%	−0.2%
mlkem1024_encaps_from_bytes	−0.3%	−0.3%	−0.2%	−0.2%
mlkem1024_keygen	0.0%	0.0%	+0.0%	+0.0%

(default: 21.43 → 20.48 µs for decaps, 16.29 → 15.79 µs for encaps.) Keygen uses none of these codecs, and ML-KEM-1024 only the message's (d = 1; its d = 5 and d = 11 are separate functions), hence its small change. Nothing is slower by 0.1% or more.

Tests

• lake build (full), Emit.lean (only src/asm/aarch64/mlkem.rs changes), check_lean_modules, check_lean_imports, check_lean_speed, check_registrations, check_variants, check_mcdt, check_arch_gates, check_vectors, cpu_tests lint and algorithms_table.py (no change) pass.
• ML-KEM tests under qemu, Wycheproof included: cargo test --release --target aarch64-unknown-linux-gnu -- mlkem, 19 tests.

🤖 Generated with Claude Code

https://claude.ai/code/session_01Ma2fA9rUTPXEEQjYyCMCRw

[alex]

What's the perf here?

[alex]

N2 run: https://github.com/pyca/verified-garbage/actions/runs/39130771921 (all four aarch64 configurations, best of 8 copies, against 0bac9044).

• ML-KEM-768 decaps: −4.5% (default/sha3) and −3.5% (neon/none); default goes 21.43 → 20.48 µs.
• ML-KEM-768 encaps: −3.1% / −2.1%.
• Keygen and ML-KEM-1024 are unchanged (0 to −0.5%). No row is slower.

The gain comes from compress/decompress: they went from ~1.8–2.7k scalar instructions per polynomial to ~0.3–0.7k NEON.

---

Generated by Claude Code