ML-KEM AArch64: decapsulation's decryption with 16-bit coefficients

lean/VerifiedGarbage/Impl/MlKem/AArch64/Kem16.lean

@@ -2,10 +2,7 @@ module
 
 public import VerifiedGarbage.Impl.MlKem.AArch64.Hybrid
 public import VerifiedGarbage.Impl.MlKem.AArch64.FixedHashDecaps
-public import VerifiedGarbage.Impl.MlKem.AArch64.AddSubW
-public import VerifiedGarbage.Impl.MlKem.AArch64.Codec12W
-public import VerifiedGarbage.Impl.MlKem.AArch64.DotW
-public import VerifiedGarbage.Impl.MlKem.AArch64.Widen16
+public import VerifiedGarbage.Impl.MlKem.AArch64.Ops16
 public import VerifiedGarbage.Impl.MlKem.AArch64.Decaps
 public import VerifiedGarbage.Impl.MlKem.AArch64.NoiseWith
 import VerifiedGarbage.Impl.MlKem.AArch64.Encaps
@@ -21,7 +18,8 @@ the sampler and the CBD write halfwords (`w16`), the transforms are
 and the sums `AddSubW.addW`, and keys are encoded and decoded by
 `Codec12W`. Compression still reads 32-bit coefficients: each row of `u` and
 `v` is widened (`Widen16.widen`) into `PP` first, and `v` adds the decoded
-message there with `vg_mlkem_add`. Decapsulation's decryption stays 32-bit.
+message there with `vg_mlkem_add`. Decapsulation's decryption (`Optimized.decryptM`) is
+16-bit too.
 The pieces of key generation (`Kem16P.keyGen`): each row of `t̂` (`kgRow`)
 as soon as the matrix has sampled it, the pairs absorbing `ek` into the third
 state of the matrix's `vg_keccak_f1600_x3_sha3` calls as the rows encode it
@@ -35,25 +33,6 @@ open VG VG.AArch64 VG.Impl.MlKem.AArch64 VG.Impl.MlKem.AArch64.KEM
 open VG.Impl.MlKem.AArch64.Noise (noiseOff noiseAllWith)
 open VG.Impl.MlKem.AArch64.Ntt16.Readonly (nttAtWithW nttInvAtWithW)
 
-/-- `f ← f + g`, both stored with 16 bits in `scratch`. -/
-def addWAt (f g : Nat) : Prog isa :=
-  .seq (.block (ptrTo .x0 .x28 f ++ ptrTo .x1 .x28 g)) AddSubW.addW
-
-/-- The polynomial at `f`, stored with 16 bits, at `g` with 32. -/
-def widenAt (f g : Nat) : Prog isa :=
-  .seq (.block (ptrTo .x0 .x28 f ++ ptrTo .x1 .x28 g)) Widen16.widen
-
-/-- `ByteDecode₁₂` of the 384 bytes at `b + o` into `off`, with 16 bits. -/
-def dec12WAt (b : Reg) (o off : Nat) : Prog isa :=
-  .seq (.block (ptrTo .x0 b o ++ ptrTo .x1 .x28 off)) Codec12W.decW
-
-/-- `ByteEncode₁₂` of the polynomial at `off`, stored with 16 bits, into `b + o`. -/
-def enc12WAt (off : Nat) (b : Reg) (o : Nat) : Prog isa :=
-  .seq (.block (ptrTo .x0 .x28 off ++ ptrTo .x1 b o)) Codec12W.encW
-
-/-- The fused dot product into `out`, of terms in `scratch`. -/
-def dotAt (out : Nat) (terms : List (Nat × Nat)) : Prog isa := DotW.dotW .x28 .x28 .x28 out terms
-
 /-- Row `i` of `u`: `NTT⁻¹(Σⱼ Âⱼᵢ ∘ ŷⱼ) + e₁[i]`, widened, compressed and encoded. -/
 def uRow (L : KemLay) (ct : Reg) (co i : Nat) : Prog isa :=
   seqs [dotAt (TP L) ((List.range L.k).map fun j => (aOff L j i, yOff L j)),

lean/VerifiedGarbage/Impl/MlKem/AArch64/Narrow16.lean

@@ -0,0 +1,26 @@
+module
+public import VerifiedGarbage.TCB.AArch64.Isa
+
+
+/-!
+# ML-KEM on AArch64: a polynomial from 32-bit to 16-bit coefficients
+
+`narrow` copies the polynomial stored with 32-bit coefficients at `x0`, each less than `2¹⁶`,
+to `x1` with 16-bit ones (each eight coefficients loaded at once, their low halfwords taken by
+`uzp1`).
+-/
+
+@[expose] public section
+
+namespace VG.Impl.MlKem.AArch64.Narrow16
+
+open VG.AArch64
+
+def narrowBody : List Instr :=
+  [.ldrq .v0 .x0 0, .ldrq .v1 .x0 16, .vop (.perm .uzp1 .h8 .v2 .v0 .v1), .strq .v2 .x1 0,
+   .addImm .x .x0 .x0 32, .addImm .x .x1 .x1 16, .subImm .x .x10 .x10 1]
+
+def narrow : Prog isa :=
+  .seq (.block [.movz .x .x10 32 0]) (.loop (.block narrowBody) (.nonzero .x .x10))
+
+end VG.Impl.MlKem.AArch64.Narrow16

lean/VerifiedGarbage/Impl/MlKem/AArch64/Ntt16/Readonly.lean

@@ -9,17 +9,8 @@ namespace VG.Impl.MlKem.AArch64.Ntt16.Readonly
 open VG.AArch64
 open VG.Impl.MlKem.AArch64 VG.Impl.MlKem.AArch64.Ntt16
 
-def fwd : Prog isa :=
- .seq (.block (consts ++ [mov .x2 .x0,mov .x3 .x1,.movz .x .x5 4 0]))
- (.seq (.loop (.block pass1Body) (.nonzero .x .x5))
- (.seq (.block [mov .x2 .x1,.adrSym .x6 "VG_MLKEM_NTT16_ZETAS",mov .x4 .x0,.movz .x .x5 8 0])
- (.loop (.block pass2Body) (.nonzero .x .x5))))
-def inv : Prog isa :=
- .seq (.block (qConsts ++ [mov .x2 .x0,mov .x3 .x1,.adrSym .x6 "VG_MLKEM_INTT16_ZETAS",.movz .x .x5 8 0]))
- (.seq (.loop (.block invPass1Body) (.nonzero .x .x5))
- (.seq (.block (invConsts ++ [mov .x2 .x1,mov .x3 .x0,.movz .x .x5 4 0]))
- (.loop (.block invPass2Body) (.nonzero .x .x5))))
-/-- `inv` of a polynomial stored with 16-bit coefficients (`invPass1BodyW`, `invPass2BodyW`). -/
+/-- The inverse transform of a polynomial stored with 16-bit coefficients (`invPass1BodyW`,
+`invPass2BodyW`). -/
 def invW : Prog isa :=
  .seq (.block (qConsts ++ [mov .x2 .x0,mov .x3 .x1,.adrSym .x6 "VG_MLKEM_INTT16_ZETAS",.movz .x .x5 8 0]))
  (.seq (.loop (.block invPass1BodyW) (.nonzero .x .x5))

lean/VerifiedGarbage/Impl/MlKem/AArch64/Ntt16/ReadonlyCalls.lean

@@ -8,22 +8,11 @@ public import VerifiedGarbage.Impl.MlKem.AArch64.Ntt16.Fast
 namespace VG.Impl.MlKem.AArch64.Ntt16.Readonly
 open VG.AArch64 VG.Impl.MlKem.AArch64
 
-def fwdName : String := "vg_mlkem_ntt_readonly"
-def invName : String := "vg_mlkem_inv_ntt_readonly"
-def fwdConsts := [("VG_MLKEM_NTT16_ZETAS",words 16 tableHalves)]
-def invConstsTable := [("VG_MLKEM_INTT16_ZETAS",words 16 invTableHalves)]
-
-def nttAtWith (off ns : Nat) : Prog isa :=
-  .seq (.block (ptrTo .x0 .x28 off ++ ptrTo .x1 .x28 ns)) fwd
-def nttInvAtWith (off ns : Nat) : Prog isa :=
-  .seq (.block (ptrTo .x0 .x28 off ++ ptrTo .x1 .x28 ns)) inv
-/-- `nttAtWith` and `nttInvAtWith` of a polynomial stored with 16-bit coefficients: the
-forward transform a call of `vg_mlkem16_ntt`. -/
+/-- The transforms of the polynomial stored with 16-bit coefficients at `off` from `x28`, with
+working space at `ns`: the forward transform a call of `vg_mlkem16_ntt`. -/
 def nttAtWithW (off ns : Nat) : Prog isa :=
   .seq (.block (ptrTo .x0 .x28 off ++ ptrTo .x1 .x28 ns)) (.call "vg_mlkem16_ntt" Fast.fwdFast)
 def nttInvAtWithW (off ns : Nat) : Prog isa :=
   .seq (.block (ptrTo .x0 .x28 off ++ ptrTo .x1 .x28 ns)) invW
-def nttAt (off : Nat) : Prog isa := nttAtWith off KEM.NS
-def nttInvAt (off : Nat) : Prog isa := nttInvAtWith off KEM.NS
 
 end VG.Impl.MlKem.AArch64.Ntt16.Readonly

lean/VerifiedGarbage/Impl/MlKem/AArch64/Ops16.lean

@@ -0,0 +1,44 @@
+module
+
+import VerifiedGarbage.Impl.MlKem.AArch64.Kem
+public import VerifiedGarbage.Impl.MlKem.AArch64.AddSubW
+public import VerifiedGarbage.Impl.MlKem.AArch64.Codec12W
+public import VerifiedGarbage.Impl.MlKem.AArch64.DotW
+public import VerifiedGarbage.Impl.MlKem.AArch64.Widen16
+public import VerifiedGarbage.Impl.MlKem.AArch64.Narrow16
+public import VerifiedGarbage.Impl.MlKem.AArch64.Top
+
+/-! # AArch64 ML-KEM with 16-bit coefficients: the operations on `scratch`
+
+The sums, conversions, codecs and dot products of polynomials stored with 16 bits in the first
+half of each 1024-byte slot of `scratch`, which `Kem16` and the decryption of
+`Optimized.decryptM` use. -/
+
+@[expose] public section
+namespace VG.Impl.MlKem.AArch64.Kem16
+open VG VG.AArch64 VG.Impl.MlKem.AArch64 VG.Impl.MlKem.AArch64.KEM
+
+/-- `f ← f + g`, both stored with 16 bits in `scratch`. -/
+def addWAt (f g : Nat) : Prog isa :=
+  .seq (.block (ptrTo .x0 .x28 f ++ ptrTo .x1 .x28 g)) AddSubW.addW
+
+/-- The polynomial at `f`, stored with 16 bits, at `g` with 32. -/
+def widenAt (f g : Nat) : Prog isa :=
+  .seq (.block (ptrTo .x0 .x28 f ++ ptrTo .x1 .x28 g)) Widen16.widen
+
+/-- `ByteDecode₁₂` of the 384 bytes at `b + o` into `off`, with 16 bits. -/
+def dec12WAt (b : Reg) (o off : Nat) : Prog isa :=
+  .seq (.block (ptrTo .x0 b o ++ ptrTo .x1 .x28 off)) Codec12W.decW
+
+/-- `ByteEncode₁₂` of the polynomial at `off`, stored with 16 bits, into `b + o`. -/
+def enc12WAt (off : Nat) (b : Reg) (o : Nat) : Prog isa :=
+  .seq (.block (ptrTo .x0 .x28 off ++ ptrTo .x1 b o)) Codec12W.encW
+
+/-- The polynomial at `f`, stored with 32 bits, at `g` with 16. -/
+def narrowAt (f g : Nat) : Prog isa :=
+  .seq (.block (ptrTo .x0 .x28 f ++ ptrTo .x1 .x28 g)) Narrow16.narrow
+
+/-- The fused dot product into `out`, of terms in `scratch`. -/
+def dotAt (out : Nat) (terms : List (Nat × Nat)) : Prog isa := DotW.dotW .x28 .x28 .x28 out terms
+
+end VG.Impl.MlKem.AArch64.Kem16

lean/VerifiedGarbage/Impl/MlKem/AArch64/Optimized.lean

@@ -1,7 +1,7 @@
 module
 
-public import VerifiedGarbage.Impl.MlKem.AArch64.ReadonlyDecrypt
-public import VerifiedGarbage.Impl.MlKem.AArch64.Dot
+public import VerifiedGarbage.Impl.MlKem.AArch64.Ntt16.ReadonlyCalls
+public import VerifiedGarbage.Impl.MlKem.AArch64.Ops16
 import VerifiedGarbage.Impl.MlKem.AArch64.Noise
 
 /-! # Paired sampling and fused polynomial products for AArch64 ML-KEM
@@ -16,12 +16,20 @@ before re-encryption.
 @[expose] public section
 namespace VG.Impl.MlKem.AArch64.Optimized
 open VG VG.AArch64 VG.Impl.MlKem.AArch64
-open VG.Impl.MlKem.AArch64.Dot VG.Impl.MlKem.AArch64.Noise
+open VG.Impl.MlKem.AArch64.Noise
 
+/-- Row `i` of `u`: decoded and decompressed with 32 bits into `PP`, narrowed to 16 bits at
+`yOff i` and transformed by `vg_mlkem16_ntt`. -/
+def deU (L : KemLay) (i : Nat) : Prog isa :=
+ seqs [L.ddLAt .x26 (32*L.du*i) L.du (KEM.PP L),Kem16.narrowAt (KEM.PP L) (KEM.yOff L i),
+  Ntt16.Readonly.nttAtWithW (KEM.yOff L i) KEM.NS]
+/-- `m = Compress₁(v - NTT⁻¹(ŝᵀ ∘ û))`, with `ŝ` and `û` stored with 16 bits: the dot product
+and the inverse transform in 16 bits, widened into `PP` for the subtraction from `v`. -/
 def decryptM (L : KemLay) : Prog isa :=
- seqs (((List.range L.k).map (readonlyDeU L)) ++ ((List.range L.k).map fun j => dec12At .x25 (384*j) (KEM.aOff L 0 j)) ++
- [dot (KEM.TP L) ((List.range L.k).map fun j => (KEM.aOff L 0 j,KEM.yOff L j)),
- Ntt16.Readonly.nttInvAt (KEM.TP L),L.ddLAt .x26 (32*L.du*L.k) L.dv (KEM.EP L),subAt (KEM.EP L) (KEM.TP L),ceAt (KEM.EP L) 1 .x28 KEM.MB])
+ seqs (((List.range L.k).map (deU L)) ++ ((List.range L.k).map fun j => Kem16.dec12WAt .x25 (384*j) (KEM.aOff L 0 j)) ++
+ [Kem16.dotAt (KEM.TP L) ((List.range L.k).map fun j => (KEM.aOff L 0 j,KEM.yOff L j)),
+ Ntt16.Readonly.nttInvAtWithW (KEM.TP L) KEM.NS,Kem16.widenAt (KEM.TP L) (KEM.PP L),
+ L.ddLAt .x26 (32*L.du*L.k) L.dv (KEM.EP L),subAt (KEM.EP L) (KEM.PP L),ceAt (KEM.EP L) 1 .x28 KEM.MB])
 def deStart (callee : Impl.Sha3.AArch64.Callee) (L : KemLay) : Prog isa :=
  seqs [.block (L.kemPrologue 3 [0,1,2,3]),decryptM L,
  hashWith callee .x28 KEM.ST KEM.WK 72 6 [⟨.x28,KEM.MB,32⟩,⟨.x25,768*L.k+32,32⟩]

lean/VerifiedGarbage/Proof/MlKem/AArch64/Kem16/Decaps.lean

@@ -8,7 +8,6 @@ public import VerifiedGarbage.Proof.MlKem.AArch64.Optimized.Tables
 
 import VerifiedGarbage.Proof.MlKem.AArch64.Kem16.Encrypt
 import VerifiedGarbage.Proof.MlKem.AArch64.Optimized.KemVDecode
-import VerifiedGarbage.Proof.MlKem.AArch64.Optimized.KemReadonlyTransform
 import VerifiedGarbage.Proof.MlKem.AArch64.Optimized.KemURow
 import VerifiedGarbage.Proof.MlKem.AArch64.Noise.ReadonlyLayoutTables
 import VerifiedGarbage.Proof.MlKem.AArch64.Optimized.KemNoise