ML-DSA on x86-64: UseHint with AVX2 (verify −2% instructions)

lean/VerifiedGarbage/Artifacts/MlDsaRound/X86_64.lean

@@ -6,6 +6,7 @@ import VerifiedGarbage.Proof.MlDsa.X86_64.Round.UseHint
 import VerifiedGarbage.Proof.MlDsa.X86_64.Round.MakeHint
 import VerifiedGarbage.Proof.MlDsa.X86_64.Round.YBits
 import VerifiedGarbage.Proof.MlDsa.X86_64.Round.YHint
+import VerifiedGarbage.Proof.MlDsa.X86_64.Round.YUse
 
 /-!
 # ML-DSA (FIPS 204) on x86-64: rounding and hints
@@ -105,6 +106,16 @@ def artifacts : List Artifact := [
     code := Impl.MlDsa.X86_64.Round.useHint
     contract := Spec.MlDsa.useHintContract X86_64.abi
     verified := useHint_verified
-    spSafe := Code.all_of_allInstrs (by decide +kernel) }]
+    spSafe := Code.all_of_allInstrs (by decide +kernel) },
+  { Spec.MlDsa.useHintApi with
+    name := Spec.MlDsa.useHintApi.name ++ "_avx2"
+    target := X86_64.target
+    doc := Spec.MlDsa.useHintApi.doc (notes := ["The function computes on eight coefficients at a time in AVX2 \
+      registers, multiplying by shifts and additions; it needs AVX and AVX2."])
+    code := Impl.MlDsa.X86_64.Round.useHintAvx2
+    contract := Spec.MlDsa.useHintContract X86_64.abi
+    verified := useHintY_verified
+    spSafe := Code.all_of_allInstrs (by decide +kernel)
+    features := ["avx", "avx2"] }]
 
 end VG.Artifacts.MlDsaRound.X86_64

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

@@ -201,6 +201,7 @@ def subAvx2 : Prog isa :=
 /-- The AVX2 code. -/
 def Backend.avx2 : Backend :=
   ⟨nttAvx2, nttInvAvx2, mulAvx2, mulAddAvx2, addAvx2, subAvx2, Round.highBitsAvx2,
-    Round.lowBitsAvx2, Round.normLtAvx2, Round.makeHintAvx2, Sample.Rej4.rejNTT4Avx2, "_avx2"⟩
+    Round.lowBitsAvx2, Round.normLtAvx2, Round.makeHintAvx2, Round.useHintAvx2, Sample.Rej4.rejNTT4Avx2,
+    "_avx2"⟩
 
 end VG.Impl.MlDsa.X86_64.Arith

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

@@ -11,8 +11,8 @@ Key generation, signing and verification call the polynomial arithmetic of
 one implementation, a `Backend`: the code of `vg_mldsa_ntt`,
 `vg_mldsa_inv_ntt`, `vg_mldsa_multiply_ntt`, `vg_mldsa_multiply_add_ntt`,
 `vg_mldsa_add` and `vg_mldsa_sub`, of `vg_mldsa_high_bits`, `vg_mldsa_low_bits`,
-`vg_mldsa_norm_lt` and `vg_mldsa_make_hint`, and of `vg_mldsa_rej_ntt_poly4` (which samples four
-entries of the matrix `Â` at once), whose names end with `sfx` (e.g.
+`vg_mldsa_norm_lt`, `vg_mldsa_make_hint` and `vg_mldsa_use_hint`, and of `vg_mldsa_rej_ntt_poly4`
+(which samples four entries of the matrix `Â` at once), whose names end with `sfx` (e.g.
 `_avx2`; nothing for the SSE2 code, `sse2`). Each is a variant of the
 interface `MlDsaArith` on x86-64 (`Variants/MlDsaArith/X86_64/`), and the
 functions that call them are emitted once for each
@@ -35,20 +35,21 @@ structure Backend where
   lowBits : Prog isa
   normLt : Prog isa
   makeHint : Prog isa
+  useHint : Prog isa
   rej4 : Prog isa
   /-- What the names of its functions, and of those calling them, end with. -/
   sfx : String
 
 /-- The SSE2 code. -/
 def Backend.sse2 : Backend :=
   ⟨Arith.ntt, Arith.nttInv, Arith.mul, Arith.mulAdd, Arith.add, Arith.sub, Round.highBits,
-    Round.lowBits, Round.normLt, Round.makeHint, Sample.Rej4.rejNTT4, ""⟩
+    Round.lowBits, Round.normLt, Round.makeHint, Round.useHint, Sample.Rej4.rejNTT4, ""⟩
 
 /-- Every function empty, which the proofs that the functions calling a
 backend never write `rsp` (and load MXCSR only to restore it) evaluate in
 its place. -/
 def Backend.empty : Backend :=
   ⟨.block [], .block [], .block [], .block [], .block [], .block [], .block [], .block [], .block [],
-    .block [], .block [], ""⟩
+    .block [], .block [], .block [], ""⟩
 
 end VG.Impl.MlDsa.X86_64.Arith

lean/VerifiedGarbage/Impl/MlDsa/X86_64/Round/Avx2.lean

@@ -6,12 +6,13 @@ import VerifiedGarbage.Impl.MlKem.X86_64.Avx
 # ML-DSA on x86-64: rounding with AVX2
 
 `vg_mldsa_high_bits_avx2`, `vg_mldsa_low_bits_avx2`,
-`vg_mldsa_norm_lt_avx2` and `vg_mldsa_make_hint_avx2` are
-`vg_mldsa_high_bits`, `vg_mldsa_low_bits`, `vg_mldsa_norm_lt` and
-`vg_mldsa_make_hint` (`Round.lean`) on eight
+`vg_mldsa_norm_lt_avx2`, `vg_mldsa_make_hint_avx2` and
+`vg_mldsa_use_hint_avx2` are `vg_mldsa_high_bits`, `vg_mldsa_low_bits`,
+`vg_mldsa_norm_lt`, `vg_mldsa_make_hint` and `vg_mldsa_use_hint`
+(`Round.lean`) on eight
 coefficients at a time, in the doublewords of `ymm` registers: in each
 128-bit lane, the VEX.256 form (`toY`) of SSE2 code on the four doublewords
-of an `xmm` register (`hbX`, `lbX`, `nlX`, `mhX`), with `rdi` and `r10` at the eight
+of an `xmm` register (`hbX`, `lbX`, `nlX`, `mhX`, `uhX`), with `rdi` and `r10` at the eight
 coefficients of `r` and `out` and `rcx` counting down the 32 vectors.
 
 `Decompose` is the reference implementation's, as in `Round.lean`:
@@ -41,11 +42,14 @@ def mulX (sh : List Nat) : List Instr :=
   xmov .xmm1 .xmm0 :: sh.flatMap fun k =>
     [xmov .xmm2 .xmm1, .xop (.shift .pslld .xmm2 (BitVec.ofNat 8 k)), xb .paddd .xmm0 .xmm2]
 
+/-- `xmm0 ← f` of `xmm0`, with `127` and `2^(S-1)` in `xmm8` and `xmm9`. -/
+def hfX (g : Nat) : List Instr :=
+  [xb .paddd .xmm0 .xmm8, .xop (.shift .psrld .xmm0 7)] ++ mulX (dSh g) ++
+    [xb .paddd .xmm0 .xmm9, .xop (.shift .psrld .xmm0 (BitVec.ofNat 8 (dShift g)))]
+
 /-- `xmm0 ← r₁` of `xmm0`, with `127`, `2^(S-1)` and `m` in `xmm8`, `xmm9` and `xmm10`. -/
 def hbX (g : Nat) : List Instr :=
-  [xb .paddd .xmm0 .xmm8, .xop (.shift .psrld .xmm0 7)] ++ mulX (dSh g) ++
-    [xb .paddd .xmm0 .xmm9, .xop (.shift .psrld .xmm0 (BitVec.ofNat 8 (dShift g))), xmov .xmm1 .xmm0,
-      xb .psubd .xmm1 .xmm10, .xop (.shift .psrad .xmm1 31), xb .pand .xmm0 .xmm1]
+  hfX g ++ [xmov .xmm1 .xmm0, xb .psubd .xmm1 .xmm10, .xop (.shift .psrad .xmm1 31), xb .pand .xmm0 .xmm1]
 
 /-- `xmm1 ← xmm0 · 2γ₂`, through `xmm2`. -/
 def mul2X (g : Nat) : List Instr :=
@@ -155,4 +159,36 @@ def makeHintAvx2 : Prog isa :=
   .seq (.block (gammaCmp .rdx ++ [.mov .r10 (.reg .rcx), .mov32 .r9 (.imm 0)]))
     (.seq (.ite .e (mhY g32) (mhY g88)) (.block [.mov .rax (.reg .r9), .vop .vzeroupper]))
 
+/-! ## `vg_mldsa_use_hint_avx2`
+
+In each lane, `f` of `r` (`hfX`, which `hbX` reduces modulo `m`), the sign
+`P` of `f · 2γ₂ - r` (all ones exactly when `r₀ > 0`), and the sign `H` of
+`h | -h` (all ones exactly when the hint is not 0); then `δ = ¬(P + P) ∧ H`
+is `1` if `P` is set, `-1` if not, and `0` if the hint is 0, and the result
+is `(f + δ + m) mod m`, by two conditional subtractions of `m` (`msub`), as
+in `vg_mldsa_use_hint` (`Round.lean`). -/
+
+/-- `x ← x - m`, plus `m` if negative, with `m` in `xmm10`, through `t`. -/
+def msub (x t : XReg) : List Instr :=
+  [xb .psubd x .xmm10, xmov t x, .xop (.shift .psrad t 31), xb .pand t .xmm10, xb .paddd x t]
+
+/-- `xmm4 ← UseHint` of the hint in `xmm5` and `r` in `xmm0`. -/
+def uhX (g : Nat) : List Instr :=
+  xmov .xmm3 .xmm0 :: hfX g ++ xmov .xmm4 .xmm0 :: mul2X g ++
+    [xb .psubd .xmm1 .xmm3, .xop (.shift .psrad .xmm1 31), xb .paddd .xmm1 .xmm1, xb .pxor .xmm2 .xmm2,
+      xb .psubd .xmm2 .xmm5, xb .por .xmm2 .xmm5, .xop (.shift .psrad .xmm2 31), xb .pandn .xmm1 .xmm2,
+      xb .paddd .xmm4 .xmm1, xb .paddd .xmm4 .xmm10] ++ msub .xmm4 .xmm1 ++ msub .xmm4 .xmm1
+
+/-- Eight hints (at `rdi`) and coefficients of `r` (at `rsi`) to `out` (at `r10`). -/
+def uhBodyY (g : Nat) : List Instr :=
+  [.vmovdquLoad .l256 .xmm0 (at_ .rsi 0), .vmovdquLoad .l256 .xmm5 (at_ .rdi 0)] ++ toY (uhX g) ++
+    [.vmovdquStore .l256 (at_ .r10 0) .xmm4, .alu .add .rdi (.imm 32), .alu .add .rsi (.imm 32),
+      .alu .add .r10 (.imm 32)]
+
+def uhY (g : Nat) : Prog isa := .seq (.block (yC g)) (rcxLoop 32 (uhBodyY g))
+
+def useHintAvx2 : Prog isa :=
+  .seq (.block (gammaCmp .rdx ++ [.mov .r10 (.reg .rcx)]))
+    (.seq (.ite .e (uhY g32) (uhY g88)) (.block [.vop .vzeroupper]))
+
 end VG.Impl.MlDsa.X86_64.Round

lean/VerifiedGarbage/Impl/MlDsa/X86_64/Verify/Frag.lean

@@ -171,7 +171,7 @@ def ballAt (ct : Ptr) (len tau : Nat) (c : Ptr) : Prog isa :=
     [(.rdi, .ptr ct), (.rsi, .imm len), (.rdx, .imm tau), (.rcx, .ptr c), (.r8, .ptr (sc oSS))]
 
 def useHintAt (h r : Ptr) (g2 : Nat) (out : Ptr) : Prog isa :=
-  callAt "vg_mldsa_use_hint" P.useHint [(.rdi, .ptr h), (.rsi, .ptr r), (.rdx, .imm g2), (.rcx, .ptr out)]
+  callAt ("vg_mldsa_use_hint" ++ P.sfx) P.useHint [(.rdi, .ptr h), (.rsi, .ptr r), (.rdx, .imm g2), (.rcx, .ptr out)]
 
 def sbpAt (f : Ptr) (b : Nat) (out : Ptr) (len : Nat) : Prog isa :=
   callAt "vg_mldsa_simple_bit_pack" P.simpleBitPack

lean/VerifiedGarbage/Proof/MlDsa/X86_64/Arith/Backend.lean

@@ -60,6 +60,7 @@ structure BackendOk (B : Backend) : Prop where
   lowBits : FnOk (fun S => Spec.MlDsa.lowBitsContract X86_64.abi S) B.lowBits
   normLt : FnOk (fun S => Spec.MlDsa.normLtContract X86_64.abi S) B.normLt
   makeHint : FnOk (fun S => Spec.MlDsa.makeHintContract X86_64.abi S) B.makeHint
+  useHint : FnOk (fun S => Spec.MlDsa.useHintContract X86_64.abi S) B.useHint
   rej4 : Rej4Ok B.rej4
 
 /-- An implementation of the polynomial arithmetic on x86-64. -/
@@ -99,6 +100,8 @@ def ArithImpl.sse2 : ArithImpl where
         (by decide +kernel)
       makeHint := FnOk.of Round.makeHint_verified (by decide +kernel) (by decide +kernel) (by decide +kernel)
         (by decide +kernel)
+      useHint := FnOk.of Round.useHint_verified (by decide +kernel) (by decide +kernel) (by decide +kernel)
+        (by decide +kernel)
       rej4 := ⟨Rej4.rejNTT4_verified, Proof.MlKem.X86_64.nosp_of (by decide +kernel), by decide +kernel,
         by decide +kernel, Code.all_of_allInstrs (by decide +kernel), fun _ _ _ => Rej4.rejNTT4_ret⟩ }
   features := []

lean/VerifiedGarbage/Proof/MlDsa/X86_64/Arith/BackendAvx2.lean

@@ -3,6 +3,7 @@ import VerifiedGarbage.Proof.MlDsa.X86_64.Arith.YNtt
 import VerifiedGarbage.Proof.MlDsa.X86_64.Arith.YMul
 import VerifiedGarbage.Proof.MlDsa.X86_64.Arith.YAddSub
 import VerifiedGarbage.Proof.MlDsa.X86_64.Round.YHint
+import VerifiedGarbage.Proof.MlDsa.X86_64.Round.YUse
 
 /-!
 # ML-DSA on x86-64: the polynomial arithmetic with AVX2, as an `ArithImpl`
@@ -40,6 +41,8 @@ def ArithImpl.avx2 : ArithImpl where
         (by decide +kernel)
       makeHint := FnOk.of Round.makeHintY_verified (by decide +kernel) (by decide +kernel) (by decide +kernel)
         (by decide +kernel)
+      useHint := FnOk.of Round.useHintY_verified (by decide +kernel) (by decide +kernel) (by decide +kernel)
+        (by decide +kernel)
       rej4 := ⟨Rej4.rejNTT4Avx2_verified, Proof.MlKem.X86_64.nosp_of (by decide +kernel), by decide +kernel,
         by decide +kernel, Code.all_of_allInstrs (by decide +kernel), fun _ _ _ => Rej4.rejNTT4Avx2_ret⟩ }
   features := ["avx", "avx2"]

lean/VerifiedGarbage/Proof/MlDsa/X86_64/Round/YBlock.lean

@@ -39,7 +39,7 @@ theorem hbX_ok {g : Nat} (hg : g = g32 ∨ g = g88) (s : State) (hc : HbC g s) :
     WP isa (.block (hbX g)) s fun s' => (∀ e < 4, dword (s'.xmm .xmm0) e = hbL g (dword (s.xmm .xmm0) e)) ∧
       XOnly [.xmm0, .xmm1, .xmm2] s s' := by
   rcases hg with rfl | rfl <;>
-  · simp only [hbX, mulX, dSh_32, dSh_88, dShift_32, dShift_88, xmov, xb, List.cons_append, List.nil_append,
+  · simp only [hbX, hfX, mulX, dSh_32, dSh_88, dShift_32, dShift_88, xmov, xb, List.cons_append, List.nil_append,
       List.flatMap_cons, List.flatMap_nil, List.append_nil]
     vrun [VG.X86_64.eval_movdqa]
     refine ⟨fun e he => ?_, by xonly⟩
@@ -57,7 +57,7 @@ theorem lbX_ok {g : Nat} (hg : g = g32 ∨ g = g88) (s : State) (hc : HbC g s) (
     WP isa (.block (lbX g)) s fun s' => (∀ e < 4, dword (s'.xmm .xmm3) e = lbL g (dword (s.xmm .xmm0) e)) ∧
       XOnly [.xmm0, .xmm1, .xmm2, .xmm3] s s' := by
   rcases hg with rfl | rfl <;>
-  · simp only [lbX, hbX, mulX, mul2X_32, mul2X_88, dSh_32, dSh_88, dShift_32, dShift_88,
+  · simp only [lbX, hbX, hfX, mulX, mul2X_32, mul2X_88, dSh_32, dSh_88, dShift_32, dShift_88,
       VG.Impl.MlDsa.X86_64.Arith.vcadd, xmov, xb, List.cons_append, List.nil_append, List.flatMap_cons,
       List.flatMap_nil, List.append_nil]
     vrun [VG.X86_64.eval_movdqa]