ML-DSA on x86-64: polynomial arithmetic in SSE2 (sign −19%, verify −11% instructions)

README.md

@@ -827,7 +827,7 @@ yours to keep:
 
 <td>✅</td>
 
-<td>✅ SSE2 NTT</td>
+<td>✅ SSE2 polynomial arithmetic</td>
 
 <td>✅ SHA extensions</td>
 
@@ -843,7 +843,7 @@ yours to keep:
 
 <td>✅</td>
 
-<td>✅ SSE2 NTT</td>
+<td>✅ SSE2 polynomial arithmetic</td>
 
 <td>✅ SHA extensions</td>
 
@@ -859,7 +859,7 @@ yours to keep:
 
 <td>✅</td>
 
-<td>✅ SSE2 NTT</td>
+<td>✅ SSE2 polynomial arithmetic</td>
 
 <td>✅ SHA extensions</td>
 

docs/algorithms/ml-dsa-44.toml

@@ -3,4 +3,4 @@ family = "Signatures"
 specs = ["MlDsa"]
 modules = ["src/mldsa44.rs"]
 asm = ["mldsa44", "mldsa"]
-optimized = { x86_64 = "SSE2 NTT" }
+optimized = { x86_64 = "SSE2 polynomial arithmetic" }

docs/algorithms/ml-dsa-65.toml

@@ -3,4 +3,4 @@ family = "Signatures"
 specs = ["MlDsa"]
 modules = ["src/mldsa65.rs"]
 asm = ["mldsa65", "mldsa"]
-optimized = { x86_64 = "SSE2 NTT" }
+optimized = { x86_64 = "SSE2 polynomial arithmetic" }

docs/algorithms/ml-dsa-87.toml

@@ -3,4 +3,4 @@ family = "Signatures"
 specs = ["MlDsa"]
 modules = ["src/mldsa87.rs"]
 asm = ["mldsa87", "mldsa"]
-optimized = { x86_64 = "SSE2 NTT" }
+optimized = { x86_64 = "SSE2 polynomial arithmetic" }

lean/VerifiedGarbage/Artifacts/MlDsaArith/X86_64.lean

@@ -46,27 +46,37 @@ def artifacts : List Artifact := [
   { Spec.MlDsa.mulApi with
     target := X86_64.target
     doc := Spec.MlDsa.mulApi.doc
+      (notes := ["The function computes on four coefficients at a time in SSE2 registers. It sets MXCSR to \
+        `0x1FBF` around its multiplications (Intel's mitigation of MXCSR-configuration-dependent timing), \
+        through the last 8 bytes of `h`, which it stores last, and loads the caller's MXCSR back before \
+        returning."])
     code := Impl.MlDsa.X86_64.Arith.mul
     contract := Spec.MlDsa.mulContract X86_64.abi
     verified := Proof.MlDsa.X86_64.Arith.mul_verified
     spSafe := Code.all_of_allInstrs (by lit_decide) },
   { Spec.MlDsa.mulAddApi with
     target := X86_64.target
     doc := Spec.MlDsa.mulAddApi.doc
+      (notes := ["The function computes on four coefficients at a time in SSE2 registers. It sets MXCSR to \
+        `0x1FBF` around its multiplications (Intel's mitigation of MXCSR-configuration-dependent timing), \
+        through the last 8 bytes of `h`, which it stores last, and loads the caller's MXCSR back before \
+        returning."])
     code := Impl.MlDsa.X86_64.Arith.mulAdd
     contract := Spec.MlDsa.mulAddContract X86_64.abi
     verified := Proof.MlDsa.X86_64.Arith.mulAdd_verified
     spSafe := Code.all_of_allInstrs (by lit_decide) },
   { Spec.MlDsa.addApi with
     target := X86_64.target
     doc := Spec.MlDsa.addApi.doc
+      (notes := ["The function computes on four coefficients at a time in SSE2 registers."])
     code := Impl.MlDsa.X86_64.Arith.add
     contract := Spec.MlDsa.addContract X86_64.abi
     verified := Proof.MlDsa.X86_64.Arith.add_verified
     spSafe := Code.all_of_allInstrs (by lit_decide) },
   { Spec.MlDsa.subApi with
     target := X86_64.target
     doc := Spec.MlDsa.subApi.doc
+      (notes := ["The function computes on four coefficients at a time in SSE2 registers."])
     code := Impl.MlDsa.X86_64.Arith.sub
     contract := Spec.MlDsa.subContract X86_64.abi
     verified := Proof.MlDsa.X86_64.Arith.sub_verified

lean/VerifiedGarbage/Impl/MlDsa/X86_64/Arith/AddSub.lean

@@ -1,33 +1,39 @@
-import VerifiedGarbage.Impl.MlDsa.X86_64.Arith.Common
+import VerifiedGarbage.Impl.MlDsa.X86_64.Arith.Vec
 
 /-!
 # ML-DSA on x86-64: `vg_mldsa_add` and `vg_mldsa_sub`
 
-`add(f = rdi, g = rsi)` and `sub(f = rdi, g = rsi)` run over the 256
-coefficients with `rdi` and `rsi` pointing at coefficient `i` of `f` and
-`g`, and `rcx = 256 - i` counting down: `f[i] + g[i]` (for `sub`,
-`f[i] + q - g[i]`), less than `2q`, is reduced with `csubQ` and stored to
-`f[i]`. Every address and branch depends only on the pointers.
+`add(f = rdi, g = rsi)` and `sub(f = rdi, g = rsi)` compute on four
+coefficients at a time, as doublewords of SSE registers (see `Vec.lean`),
+with `rdi` and `rsi` pointing at coefficient `4i` of `f` and `g`, and
+`rcx = 64 - i` counting down: `f + g`, less than `2q`, is reduced with
+`vcsub` (for `sub`, `f - g`, in `(-q, q)`, with `vcadd`) and stored to `f`,
+with `q` in the doublewords of `xmm15`. Every address and branch depends
+only on the pointers.
 -/
 
 namespace VG.Impl.MlDsa.X86_64.Arith
 
 open VG.X86_64
+open VG.Impl.MlKem.X86_64 (xb rcxLoop)
 
-/-- Advance the two pointers and count down. -/
-def step2 : List Instr :=
-  [.alu .add .rdi (.imm 4), .alu .add .rsi (.imm 4), .alu .sub .rcx (.imm 1)]
+/-- `q` in the doublewords of `xmm15`, through `rax`. -/
+def qPro : List Instr := [.mov32 .rax (.imm 8380417), .xop (.movq .xmm15 .rax), .xop (.pshufd .xmm15 .xmm15 0)]
+
+/-- Store `xmm0` to `f` and advance the two pointers. -/
+def accTail : List Instr :=
+  [.movdquStore (at_ .rdi 0) .xmm0, .alu .add .rdi (.imm 16), .alu .add .rsi (.imm 16)]
 
 def addBody : List Instr :=
-  [.mov32 .rax (.mem (at_ .rdi 0)), .alu32 .add .rax (.mem (at_ .rsi 0))] ++ csubQ .rax .rdx ++
-    [.store32 (at_ .rdi 0) .rax] ++ step2
+  [.movdquLoad .xmm0 (at_ .rdi 0), .movdquLoad .xmm1 (at_ .rsi 0), xb .paddd .xmm0 .xmm1] ++
+    vcsub .xmm0 .xmm2 ++ accTail
 
 def subBody : List Instr :=
-  [.mov32 .rax (.mem (at_ .rdi 0)), .alu32 .add .rax (.imm qImm), .alu32 .sub .rax (.mem (at_ .rsi 0))] ++
-    csubQ .rax .rdx ++ [.store32 (at_ .rdi 0) .rax] ++ step2
+  [.movdquLoad .xmm0 (at_ .rdi 0), .movdquLoad .xmm1 (at_ .rsi 0), xb .psubd .xmm0 .xmm1] ++
+    vcadd .xmm0 .xmm2 ++ accTail
 
-def add : Prog isa := .seq (.block [.mov32 .rcx (.imm 256)]) (.loop (.block addBody) .ne)
+def add : Prog isa := .seq (.block qPro) (rcxLoop 64 addBody)
 
-def sub : Prog isa := .seq (.block [.mov32 .rcx (.imm 256)]) (.loop (.block subBody) .ne)
+def sub : Prog isa := .seq (.block qPro) (rcxLoop 64 subBody)
 
 end VG.Impl.MlDsa.X86_64.Arith

lean/VerifiedGarbage/Impl/MlDsa/X86_64/Arith/Mul.lean

@@ -1,41 +1,63 @@
-import VerifiedGarbage.Impl.MlDsa.X86_64.Arith.Common
+import VerifiedGarbage.Impl.MlDsa.X86_64.Arith.Vec
 
 /-!
 # ML-DSA on x86-64: `vg_mldsa_multiply_ntt` and `vg_mldsa_multiply_add_ntt`
 
 `multiplyNTT(h = rdi, f = rsi, g = rdx)` and `multiplyAddNTT(h = rdi,
-f = rsi, g = rdx)` run over the 256 coefficients with `rdi`, `rsi` and `r8`
-(`g`, as `mul` writes `rdx`) pointing at coefficient `i` of `h`, `f` and
-`g`, and `rcx = 256 - i` counting down: the product `f[i] · g[i]` (by
-`mul`, less than `q²`; for `multiplyAddNTT`, plus `h[i]`) is reduced with
-`reduce` and stored to `h[i]`. Every address and branch depends only on the
-pointers.
+f = rsi, g = rdx)` compute on four coefficients at a time, as doublewords of
+SSE registers (see `Vec.lean`), with `rdi`, `rsi` and `rdx` pointing at
+coefficient `4i` of `h`, `f` and `g`, and `rcx = 63 - i` counting down:
+`vmont` of `f` by `g` is `f · g · 2⁻³²`, and `vmont` of that by
+`2⁶⁴ mod q = 2365951` (in `xmm11`) is `f · g`, less than `2q`, which `vcsub`
+reduces (for `multiplyAddNTT`, `h` is then added and the sum reduced), and
+it is stored to `h`.
+
+The functions have no working space and use no stack, so `withMxcsr` saves
+MXCSR through the last 8 bytes of `h` (`r8 = h`): the last four coefficients
+of `h` are loaded to `xmm6` first, the loop stores the first 252, the last
+four are computed from registers before MXCSR is loaded back, and stored
+after it. Every address and branch depends only on the pointers.
 -/
 
 namespace VG.Impl.MlDsa.X86_64.Arith
 
 open VG.X86_64
+open VG.Impl.MlKem.X86_64 (xb xmov withMxcsr rcxLoop)
 
-/-- Advance the three pointers and count down. -/
-def step3 : List Instr :=
-  [.alu .add .rdi (.imm 4), .alu .add .rsi (.imm 4), .alu .add .r8 (.imm 4), .alu .sub .rcx (.imm 1)]
+/-- The constants and `2⁶⁴ mod q` in the doublewords of `xmm11`. -/
+def mulPro : List Instr :=
+  vconsts ++ [.mov32 .rax (.imm 2365951), .xop (.movq .xmm11 .rax), .xop (.pshufd .xmm11 .xmm11 0)]
 
-/-- `f[i] · g[i]`, in `rax`. -/
-def mulHead : List Instr :=
-  [.mov32 .rax (.mem (at_ .rsi 0)), .mov32 .r9 (.mem (at_ .r8 0)), .mul .r9]
+/-- The coefficients of `f`, `g` and `h` to `xmm3`, `xmm13` and `xmm5`. -/
+def mulLoads : List Instr :=
+  [.movdquLoad .xmm3 (at_ .rsi 0), .movdquLoad .xmm13 (at_ .rdx 0), .movdquLoad .xmm5 (at_ .rdi 0)]
 
-/-- `f[i] · g[i] + h[i]`, in `rax`. -/
-def mulAddHead : List Instr :=
-  mulHead ++ [.mov32 .r9 (.mem (at_ .rdi 0)), .alu .add .rax (.reg .r9)]
+/-- `f · g` for four coefficients, in `[0, q)`, in `xmm3`. -/
+def mulCore : List Instr :=
+  [.xop (.pshufd .xmm12 .xmm13 0xF5)] ++ vmont .xmm3 .xmm13 .xmm12 .xmm2 .xmm4 ++
+    vmont .xmm3 .xmm11 .xmm11 .xmm2 .xmm4 ++ vcsub .xmm3 .xmm2
 
-def mulBody : List Instr := mulHead ++ reduce ++ [.store32 (at_ .rdi 0) .r10] ++ step3
+/-- `h + f · g` for four coefficients, in `[0, q)`, in `xmm3`. -/
+def mulAddCore : List Instr := mulCore ++ [xb .paddd .xmm3 .xmm5] ++ vcsub .xmm3 .xmm2
 
-def mulAddBody : List Instr := mulAddHead ++ reduce ++ [.store32 (at_ .rdi 0) .r10] ++ step3
+/-- Store `xmm3` to `h` and advance the three pointers. -/
+def mulTail : List Instr :=
+  [.movdquStore (at_ .rdi 0) .xmm3, .alu .add .rdi (.imm 16), .alu .add .rsi (.imm 16),
+    .alu .add .rdx (.imm 16)]
 
-def mul : Prog isa :=
-  .seq (.block [.mov .r8 (.reg .rdx)]) (.seq (.block [.mov32 .rcx (.imm 256)]) (.loop (.block mulBody) .ne))
+/-- The last four coefficients, with those of `h` in `xmm6`, to `xmm3`. -/
+def mulLast (core : List Instr) : List Instr :=
+  [.movdquLoad .xmm3 (at_ .rsi 0), .movdquLoad .xmm13 (at_ .rdx 0), xmov .xmm5 .xmm6] ++ core
 
-def mulAdd : Prog isa :=
-  .seq (.block [.mov .r8 (.reg .rdx)]) (.seq (.block [.mov32 .rcx (.imm 256)]) (.loop (.block mulAddBody) .ne))
+/-- A function of `h`, `f` and `g` four coefficients at a time by `core`. -/
+def mulFn (core : List Instr) : Prog isa :=
+  .seq (.block [.mov .r8 (.reg .rdi), .movdquLoad .xmm6 (at_ .rdi 1008)])
+    (.seq (withMxcsr .r8 1016
+        (.seq (.block mulPro) (.seq (rcxLoop 63 (mulLoads ++ core ++ mulTail)) (.block (mulLast core)))))
+      (.block [.movdquStore (at_ .rdi 0) .xmm3]))
+
+def mul : Prog isa := mulFn mulCore
+
+def mulAdd : Prog isa := mulFn mulAddCore
 
 end VG.Impl.MlDsa.X86_64.Arith