ML-DSA on x86-64: ExpandMask four polynomials at a time in signing (sign −5 to −11% instructions with AVX2)

README.md

@@ -924,7 +924,7 @@ yours to keep:
 
 <td>✅</td>
 
-<td>✅ AVX2; SSE2 and AVX2 polynomial arithmetic, rounding and norm check; matrix sampled with four SHAKE128 instances at once</td>
+<td>✅ AVX2; SSE2 and AVX2 polynomial arithmetic, rounding and norm check; matrix and masks sampled with four SHAKE128 or SHAKE256 instances at once</td>
 
 <td>✅ SHA extensions</td>
 
@@ -940,7 +940,7 @@ yours to keep:
 
 <td>✅</td>
 
-<td>✅ AVX2; SSE2 and AVX2 polynomial arithmetic, rounding and norm check; matrix sampled with four SHAKE128 instances at once</td>
+<td>✅ AVX2; SSE2 and AVX2 polynomial arithmetic, rounding and norm check; matrix and masks sampled with four SHAKE128 or SHAKE256 instances at once</td>
 
 <td>✅ SHA extensions</td>
 
@@ -956,7 +956,7 @@ yours to keep:
 
 <td>✅</td>
 
-<td>✅ AVX2; SSE2 and AVX2 polynomial arithmetic, rounding and norm check; matrix sampled with four SHAKE128 instances at once</td>
+<td>✅ AVX2; SSE2 and AVX2 polynomial arithmetic, rounding and norm check; matrix and masks sampled with four SHAKE128 or SHAKE256 instances at once</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 and AVX2 polynomial arithmetic, rounding and norm check; matrix sampled with four SHAKE128 instances at once" }
+optimized = { x86_64 = "SSE2 and AVX2 polynomial arithmetic, rounding and norm check; matrix and masks sampled with four SHAKE128 or SHAKE256 instances at once" }

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 and AVX2 polynomial arithmetic, rounding and norm check; matrix sampled with four SHAKE128 instances at once" }
+optimized = { x86_64 = "SSE2 and AVX2 polynomial arithmetic, rounding and norm check; matrix and masks sampled with four SHAKE128 or SHAKE256 instances at once" }

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 and AVX2 polynomial arithmetic, rounding and norm check; matrix sampled with four SHAKE128 instances at once" }
+optimized = { x86_64 = "SSE2 and AVX2 polynomial arithmetic, rounding and norm check; matrix and masks sampled with four SHAKE128 or SHAKE256 instances at once" }

lean/VerifiedGarbage/Artifacts/MlDsaSample/X86_64.lean

@@ -5,6 +5,7 @@ import VerifiedGarbage.Proof.MlDsa.X86_64.Sample.RejBoundedCT
 import VerifiedGarbage.Proof.MlDsa.X86_64.Sample.ExpandMask
 import VerifiedGarbage.Proof.MlDsa.X86_64.Sample.BallCT
 import VerifiedGarbage.Proof.MlDsa.X86_64.Sample.Rej4Verified
+import VerifiedGarbage.Proof.MlDsa.X86_64.Sample.M4Verified
 
 /-! # ML-DSA (FIPS 204) on x86-64: the sampling primitives -/
 
@@ -63,6 +64,28 @@ def artifacts : List Artifact := [
     stack := 16
     verified := Proof.MlDsa.X86_64.Sample.expandMask_verified
     spSafe := Code.all_of_allInstrs (by lit_decide) },
+  { Spec.MlDsa.expandMask4Api with
+    target := X86_64.target
+    doc := Spec.MlDsa.expandMask4Api.doc (notes := ["It calls `vg_mldsa_expand_mask_poly` on each seed, and \
+      saves its caller's callee-saved registers in `scratch`."])
+    code := Impl.MlDsa.X86_64.Sample.Mask4.expandMask4
+    contract := Spec.MlDsa.expandMask4Contract X86_64.abi 24
+    stack := 24
+    verified := Proof.MlDsa.X86_64.Mask4.expandMask4_verified
+    spSafe := Code.all_of_allInstrs (by decide +kernel) },
+  { Spec.MlDsa.expandMask4Api with
+    name := Spec.MlDsa.expandMask4Api.name ++ "_avx2"
+    target := X86_64.target
+    doc := Spec.MlDsa.expandMask4Api.doc (notes := ["It runs the four instances of SHAKE256 at once, in the \
+      four 64-bit elements of AVX2 registers (as `vg_mldsa_rej_ntt_poly4_avx2` does with SHAKE128), \
+      squeezing five blocks of each, and unpacks the 576 or 640 bytes of each seed's output as \
+      `vg_mldsa_expand_mask_poly` does. It saves its caller's callee-saved registers in `scratch`."])
+    code := Impl.MlDsa.X86_64.Sample.Mask4.expandMask4Avx2
+    contract := Spec.MlDsa.expandMask4Contract X86_64.abi 24
+    stack := 24
+    verified := Proof.MlDsa.X86_64.Mask4.expandMask4Avx2_verified
+    spSafe := Code.all_of_allInstrs (by decide +kernel)
+    features := ["avx", "avx2"] },
   { Spec.MlDsa.sampleInBallApi with
     target := X86_64.target
     doc := Spec.MlDsa.sampleInBallApi.doc (notes := ["It squeezes 272 bytes of SHAKE256 output (2 blocks) and \

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

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

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

@@ -3,6 +3,7 @@ import VerifiedGarbage.Impl.MlDsa.X86_64.Arith.Mul
 import VerifiedGarbage.Impl.MlDsa.X86_64.Arith.AddSub
 import VerifiedGarbage.Impl.MlDsa.X86_64.Round.Round
 import VerifiedGarbage.Impl.MlDsa.X86_64.Sample.RejNtt4
+import VerifiedGarbage.Impl.MlDsa.X86_64.Sample.ExpandMask4
 
 /-!
 # ML-DSA on x86-64: implementations of the polynomial arithmetic
@@ -11,8 +12,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`, `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.
+`vg_mldsa_norm_lt`, `vg_mldsa_make_hint` and `vg_mldsa_use_hint`, and of `vg_mldsa_rej_ntt_poly4` and
+`vg_mldsa_expand_mask_poly4` (which sample four polynomials 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
@@ -37,19 +38,21 @@ structure Backend where
   makeHint : Prog isa
   useHint : Prog isa
   rej4 : Prog isa
+  expandMask4 : 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, Round.useHint, Sample.Rej4.rejNTT4, ""⟩
+    Round.lowBits, Round.normLt, Round.makeHint, Round.useHint, Sample.Rej4.rejNTT4,
+    Sample.Mask4.expandMask4, ""⟩
 
 /-- 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 [], .block [], .block [], ""⟩
 
 end VG.Impl.MlDsa.X86_64.Arith

lean/VerifiedGarbage/Impl/MlDsa/X86_64/Sample/ExpandMask4.lean

@@ -0,0 +1,113 @@
+import VerifiedGarbage.Impl.MlDsa.X86_64.Sample.ExpandMask
+import VerifiedGarbage.Impl.MlKem.X86_64.Sample4
+
+/-!
+# ML-DSA on x86-64: `vg_mldsa_expand_mask_poly4` and `vg_mldsa_expand_mask_poly4_avx2`
+
+`expandMask4(seeds = rdi, gamma1 = esi, a = rdx, scratch = rcx)` runs
+`ExpandMask`'s sampling of a polynomial (`vg_mldsa_expand_mask_poly`) on four
+66-byte seeds, to the four polynomials from `a`. Both keep `scratch` in
+`rbx`, `seeds` in `r12`, `a` in `r13` and `γ₁` in `r14`, and save their
+caller's values and `rbp`'s in `scratch[5088..5128)`. The baseline
+implementation (`expandMask4`) calls `vg_mldsa_expand_mask_poly` on each
+seed, with its scratch space from byte 6144 of `scratch`.
+
+The one for AVX2 (`expandMask4Avx2`) runs the four SHAKE256 instances at
+once in the four 64-bit elements of `ymm` registers, as
+`vg_mlkem_sample_ntt4_avx2` runs four of SHAKE128 (`Impl/MlKem/X86_64/Sample4.lean`,
+whose layout of the first 2368 bytes of `scratch` it shares: the four
+states, the second buffer of the permutation and the table of the round
+constants). Each seed is 66 bytes, so the padded message is one block of
+136 bytes: the code zeroes the states, writes the seed's bytes, SHAKE's
+suffix `0x1f` (at byte 66) and the last bit of the padding (`0x80`, at byte
+135) into each, and permutes them; it then copies the first 136 bytes of
+each state to its output (from byte `2368 + 680 k`), and permutes them
+again, five times in all (680 bytes, which hold the 640 that `γ₁ = 2¹⁹`
+needs and the 576 of `γ₁ = 2¹⁷`). It then unpacks the output of each seed
+with the loop of `vg_mldsa_expand_mask_poly` (`emBody`, for `c = 18` or
+`20`, chosen by a branch on the public `γ₁`).
+
+There is no branch on data, and every address depends only on the pointers
+and `γ₁`: it is constant time.
+-/
+
+namespace VG.Impl.MlDsa.X86_64.Sample.Mask4
+
+open VG.X86_64
+open VG.Impl.MlKem.X86_64 (at_)
+open VG.Impl.MlKem.X86_64.Sample4 (zero4 permArgs oRc oBuf)
+open VG.Impl.Sha3.X86_64.X4 (permute4 rcTable)
+
+/-- Where the callee-saved registers are saved. -/
+def oSave : Nat := 5088
+
+/-- The scratch space of `vg_mldsa_expand_mask_poly`, in the baseline implementation. -/
+def oScalar : Nat := 6144
+
+/-- The registers saved, at `scratch + oSave + 8 k`. -/
+def saved : List Reg := [.rbx, .rbp, .r12, .r13, .r14]
+
+/-- Save the callee-saved registers, and keep the pointers and `γ₁`. -/
+def pro : List Instr :=
+  (List.range 5).map (fun k => .store (at_ .rcx (oSave + 8 * k)) (saved.getD k .rbx)) ++
+    [.mov .rbx (.reg .rcx), .mov .r12 (.reg .rdi), .mov .r13 (.reg .rdx), .mov .r14 (.reg .rsi)]
+
+/-- Restore the callee-saved registers (`rbx` last). -/
+def epi : List Instr :=
+  ((List.range 4).map fun k => .mov (saved.getD (4 - k) .rbx) (.mem (at_ .rbx (oSave + 8 * (4 - k))))) ++
+    [.mov .rbx (.mem (at_ .rbx oSave))]
+
+/-- Bytes 0 to 65 of state `k`: the 66 bytes of seed `k`, as eight lanes and
+two bytes. -/
+def seedLanes (k : Nat) : List Instr :=
+  (List.range 8).flatMap (fun i =>
+    [.mov .rax (.mem (at_ .r12 (66 * k + 8 * i))), .store (at_ .rbx (32 * i + 8 * k)) .rax]) ++
+  [.movzx8 .rax (at_ .r12 (66 * k + 64)), .store8 (at_ .rbx (256 + 8 * k)) .rax,
+    .movzx8 .rax (at_ .r12 (66 * k + 65)), .store8 (at_ .rbx (256 + 8 * k + 1)) .rax]
+
+/-- The padded blocks of the four seeds, XORed into the zero states: the
+seeds, SHAKE's suffix at byte 66 (byte 2 of lane 8) and `0x80` at byte 135
+(byte 7 of lane 16). -/
+def absorb4 : List Instr :=
+  zero4 ++ (List.range 4).flatMap seedLanes ++
+    .mov32 .rax (.imm 0x1f) :: (List.range 4).flatMap (fun k => [.store8 (at_ .rbx (256 + 8 * k + 2)) .rax]) ++
+    .mov32 .rax (.imm 0x80) :: (List.range 4).flatMap fun k => [.store8 (at_ .rbx (512 + 8 * k + 7)) .rax]
+
+/-- The first 136 bytes of each state to block `b` of its output. -/
+def extract (b : Nat) : List Instr :=
+  (List.range 4).flatMap fun k => (List.range 17).flatMap fun i =>
+    [.mov .rax (.mem (at_ .rbx (32 * i + 8 * k))), .store (at_ .rbx (oBuf + 680 * k + 136 * b + 8 * i)) .rax]
+
+/-- Permute the states and squeeze block `b`. -/
+def squeeze4 (b : Nat) : Prog isa :=
+  .seq (.block permArgs) (.seq permute4 (.block (extract b)))
+
+/-- The coefficients of polynomial `k` from the output of seed `k`, `c` bits each. -/
+def unpack (c k : Nat) : Prog isa :=
+  .seq (.block [.mov .rsi (.reg .rbx), .alu .add .rsi (.imm (BitVec.ofNat 32 (oBuf + 680 * k))),
+      .mov .rdi (.reg .r13), .alu .add .rdi (.imm (BitVec.ofNat 32 (1024 * k))), .mov32 .rcx (.imm 64)])
+    (.loop (.block (emBody c)) .ne)
+
+def unpack4 (c : Nat) : Prog isa := .seq (unpack c 0) (.seq (unpack c 1) (.seq (unpack c 2) (unpack c 3)))
+
+/-- `vg_mldsa_expand_mask_poly4_avx2`. `vzeroupper` clears the upper halves
+of the vector registers after the last permutation. -/
+def expandMask4Avx2 : Prog isa :=
+  .seq (.block (pro ++ rcTable .rbx (oRc / 32) ++ absorb4))
+    (.seq (squeeze4 0) (.seq (squeeze4 1) (.seq (squeeze4 2) (.seq (squeeze4 3) (.seq (squeeze4 4)
+      (.seq (.block [.vop .vzeroupper, .alu32 .cmp .r14 (.imm 0x20000)])
+        (.seq (.ite .e (unpack4 18) (unpack4 20)) (.block epi))))))))
+
+/-- `vg_mldsa_expand_mask_poly` on seed `k`. -/
+def callK (k : Nat) : Prog isa :=
+  .seq (.block [.mov .rdi (.reg .r12), .alu .add .rdi (.imm (BitVec.ofNat 32 (66 * k))), .mov .rsi (.reg .r14),
+      .mov .rdx (.reg .r13), .alu .add .rdx (.imm (BitVec.ofNat 32 (1024 * k))), .mov .rcx (.reg .rbx),
+      .alu .add .rcx (.imm (BitVec.ofNat 32 oScalar))])
+    (.call "vg_mldsa_expand_mask_poly" expandMask)
+
+/-- `vg_mldsa_expand_mask_poly4`: `vg_mldsa_expand_mask_poly` on each seed,
+with the prologue and epilogue of `vg_mldsa_expand_mask_poly4_avx2`. -/
+def expandMask4 : Prog isa :=
+  .seq (.block pro) (.seq (callK 0) (.seq (callK 1) (.seq (callK 2) (.seq (callK 3) (.block epi)))))
+
+end VG.Impl.MlDsa.X86_64.Sample.Mask4

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

@@ -46,6 +46,8 @@ structure Prims where
   hintBitPack : Prog isa
   /-- `vg_mldsa_rej_ntt_poly4` -/
   rej4 : Prog isa
+  /-- `vg_mldsa_expand_mask_poly4` -/
+  expandMask4 : Prog isa
   /-- What the names of the polynomial arithmetic's functions end with (`Arith.Backend`). -/
   sfx : String := ""
 
@@ -57,7 +59,8 @@ the iterations left (`CNT`), the counter `κ` (`KAP`) and the number of 1s
 of the hint (`ONES`) at 888, 896 and 904 (8 bytes each); the seed of
 `RejNTTPoly` (`RS`, 34 bytes) at 912; the seed of `ExpandMask` (`MS`,
 `ρ″` and two bytes) at 960; `c̃` (`CT`, up to 64 bytes) at 1040; the four
-seeds of `vg_mldsa_rej_ntt_poly4` (`RS4`, 136 bytes) at 1152; the
+seeds of `vg_mldsa_rej_ntt_poly4` (`RS4`, 136 bytes) at 1152; the four
+seeds of `vg_mldsa_expand_mask_poly4` (`MS4`, 264 bytes) at 1296; the
 encoding of `w₁` (`W1`, up to 1024 bytes) at 2048; the working space of
 the primitives (`PS`, 2048 bytes) at 3072; and polynomials of 1024 bytes
 from 5120 (`P i`). -/
@@ -70,6 +73,7 @@ def oRS : Nat := 912
 def oMS : Nat := 960
 def oCT : Nat := 1040
 def oRS4 : Nat := 1152
+def oMS4 : Nat := 1296
 def oW1 : Nat := 2048
 def oPS : Nat := 3072
 /-- Polynomial `i`. -/
@@ -194,6 +198,11 @@ def rej4At (a w : Ptr) : Prog isa :=
   .seq (callP ("vg_mldsa_rej_ntt_poly4" ++ P.sfx) P.rej4 [.ptr (sc oRS4), .ptr a, .ptr w])
     (.block [.alu32 .and .r15 (.reg .rax)])
 
+/-- Four polynomials of `ExpandMask` from the four seeds at `MS4` to the four polynomials from `a`,
+with the working space `w`. -/
+def mask4At (gamma1 : Nat) (a w : Ptr) : Prog isa :=
+  callP ("vg_mldsa_expand_mask_poly4" ++ P.sfx) P.expandMask4 [.ptr (sc oMS4), .imm gamma1, .ptr a, .ptr w]
+
 /-- A polynomial of `ExpandMask` from the seed at `MS` to `a`. -/
 def maskAt (gamma1 : Nat) (a : Ptr) : Prog isa :=
   callP "vg_mldsa_expand_mask_poly" P.expandMask [.ptr (sc oMS), .imm gamma1, .ptr a, .ptr (sc oPS)]

lean/VerifiedGarbage/Impl/MlDsa/X86_64/Sign/Sign.lean

@@ -21,7 +21,9 @@ keeps `scratch` in `rbx`, `sk` in `rbp`, `mu` in `r12`, `rnd` in `r13` and
 2. `ŝ₁`, `ŝ₂` and `t̂₀`: the `NTT` of the `BitUnpack` of their pieces of
    `sk`; and `ρ″ = H(K ‖ rnd ‖ μ, 64)` to `MS`.
 3. The rejection sampling loop, at most 814 iterations (`minBounds.sign`),
-   with the counter `κ` at `KAP`: the commitment (`y ← ExpandMask(ρ″, κ)`,
+   with the counter `κ` at `KAP`: the commitment (`y ← ExpandMask(ρ″, κ)`:
+   four polynomials at a time (`vg_mldsa_expand_mask_poly4`, from the seeds
+   `ρ″ ‖ κ + r` at `MS4`), then the last `ℓ mod 4` one at a time;
    `ŷ = NTT(y)`, `w = NTT⁻¹(Â ŷ)`, and `c̃ = H(μ ‖ w1Encode(HighBits(w)), λ/4)`
    at `CT`), then `c = SampleInBall(c̃)` and, if it succeeded, every validity
    check of the iteration, combined without a branch into `r15`: the norms
@@ -144,15 +146,29 @@ def decode : Prog isa :=
 
 /-! ### An iteration -/
 
+/-- The two bytes of `κ + r` to `scratch + o`. -/
+def setKap (o r : Nat) : List Instr :=
+  [.mov .rax (.mem (at_ .rbx oKAP)), .alu .add .rax (.imm (BitVec.ofNat 32 r)), .store8 (at_ .rbx o) .rax,
+    .shift .shr .rax 8, .store8 (at_ .rbx (o + 1)) .rax]
+
 /-- The two bytes of `κ + r` to `MS + 64`. -/
-def setKappa (r : Nat) : List Instr :=
-  [.mov .rax (.mem (at_ .rbx oKAP)), .alu .add .rax (.imm (BitVec.ofNat 32 r)), .store8 (at_ .rbx (oMS + 64)) .rax,
-    .shift .shr .rax 8, .store8 (at_ .rbx (oMS + 65)) .rax]
+def setKappa (r : Nat) : List Instr := setKap (oMS + 64) r
 
 /-- `y[r]` and `ŷ[r] = NTT(y[r])`. -/
 def maskR (r : Nat) : Prog isa :=
   .seq (.block (setKappa r)) (.seq (maskAt P p.γ₁ (yP p r)) (.seq (copy (yhP p r) (yP p r) 1024) (nttAt P (yhP p r))))
 
+/-- `ŷ[r] = NTT(y[r])`. -/
+def yhR (r : Nat) : Prog isa := .seq (copy (yhP p r) (yP p r) 1024) (nttAt P (yhP p r))
+
+/-- `ρ″ ‖ κ + 4g + k` to seed `k` of `MS4`. -/
+def cpM4 (g k : Nat) : Prog isa :=
+  .seq (copy (sc (oMS4 + 66 * k)) (sc oMS) 64) (.block (setKap (oMS4 + 66 * k + 64) (4 * g + k)))
+
+/-- `y[4g], …, y[4g + 3]`, four at a time, and their `ŷ`. -/
+def mask4 (g : Nat) : Prog isa :=
+  .seq (seqR (cpM4 g) 0 4) (.seq (mask4At P p.γ₁ (yP p (4 * g)) (r4P p)) (seqR (yhR P p) (4 * g) 4))
+
 /-- `w[i] = NTT⁻¹(∑_j Â[i, j] ŷ[j])`. -/
 def rowW (i : Nat) : Prog isa :=
   .seq (mulAt P (wP p i) (aP p i 0) (yhP p 0))
@@ -164,8 +180,9 @@ def w1R (i : Nat) : Prog isa :=
 
 /-- `y`, `ŷ`, `w`, `w₁` and `c̃ = H(μ ‖ w1Encode(w₁), λ/4)` to `CT`. -/
 def commit : Prog isa :=
-  .seq (seqR (maskR P p) 0 p.ℓ) (.seq (seqR (rowW P p) 0 p.k) (.seq (seqR (w1R P p) 0 p.k)
-    (shakeAt [((.r12, 0), 64), (sc oW1, p.k * w1Len p)] (sc oCT) (cLen p))))
+  .seq (seqR (mask4 P p) 0 (p.ℓ / 4)) (.seq (seqR (maskR P p) (4 * (p.ℓ / 4)) (p.ℓ % 4))
+    (.seq (seqR (rowW P p) 0 p.k) (.seq (seqR (w1R P p) 0 p.k)
+      (shakeAt [((.r12, 0), 64), (sc oW1, p.k * w1Len p)] (sc oCT) (cLen p)))))
 
 /-- `z[r] = y[r] + NTT⁻¹(ĉ ŝ₁[r])` (in `y[r]`), and its norm. -/
 def zR (r : Nat) : Prog isa :=

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

@@ -26,8 +26,9 @@ of `vg_mldsa_rej_ntt_poly4` (`SB4`, 136 bytes) at 2560; the working
 space of the primitives at 4096 (2048 bytes); and polynomials of 1024 bytes
 from 8192 (`P j`): the hint `h` (polynomials 0 to 7, of which the first
 `k`), `z` (8 to 14), `c` (15), two temporaries (16, 17), `w′` (18), `w′₁`
-(19), and `Â[r, s]` (`20 + 8r + s`), then the working space of
-`vg_mldsa_rej_ntt_poly4` (8 KiB, after the last row of `Â`).
+(19), and `Â[r, s]` (`20 + ℓr + s`, entry `ℓr + s` of `Â` in order), then the
+working space of `vg_mldsa_rej_ntt_poly4` (8 KiB, from polynomial `20 + 8k`, after
+`Â` for every `ℓ` ≤ 8).
 -/
 
 namespace VG.Impl.MlDsa.X86_64.Verify
@@ -62,7 +63,8 @@ abbrev pT : Ptr := pS 16
 abbrev pT2 : Ptr := pS 17
 abbrev pW : Ptr := pS 18
 abbrev pW1 : Ptr := pS 19
-abbrev pA (r s : Nat) : Ptr := pS (20 + 8 * r + s)
+/-- `Â[r, s]`, for rows of `l` entries. -/
+abbrev pA (l r s : Nat) : Ptr := pS (20 + l * r + s)
 
 /-! ## Moves -/