ML-KEM iNTT proof for x86 AVX2 implementation

This change adds the implementation of ML-KEM Inverse NTT
function and its HolLight proof of correctness. At this moment,
the implementation is the same as the one in mlkem-native
repository at https://github.com/pq-code-package/mlkem-native.

The proof was done in collaboration with @jargh.

Signed-off-by: dkostic <[email protected]>

Author: dkostic

benchmarks/benchmark.c

@@ -831,14 +831,14 @@ void call_mldsa_ntt(void) repeat(mldsa_ntt((int32_t*)b0,(const int32_t*)b1))
 void call_mldsa_poly_reduce(void) repeat(mldsa_poly_reduce((int32_t*)b0))
 
 void call_mlkem_ntt(void) repeat(mlkem_ntt_x86((int16_t*)b0,(int16_t*)b1))
+void call_mlkem_intt(void) repeat(mlkem_intt_x86((int16_t*)b0,(int16_t*)b1))
 
 void call_bignum_copy_row_from_table_8n__32_16(void) {}
 void call_bignum_copy_row_from_table_8n__32_32(void) {}
 void call_bignum_copy_row_from_table_16__32(void) {}
 void call_bignum_copy_row_from_table_32__32(void) {}
 
 void call_bignum_emontredc_8n_cdiff__32(void) {}
-void call_mlkem_intt(void) {}
 void call_mlkem_mulcache_compute(void) {}
 void call_mlkem_tobytes(void) {}
 void call_mlkem_tomont(void) {}
@@ -1528,7 +1528,7 @@ int main(int argc, char *argv[])
   timingtest(all,"mlkem_basemul_k2",call_mlkem_basemul_k2);
   timingtest(all,"mlkem_basemul_k3",call_mlkem_basemul_k3);
   timingtest(all,"mlkem_basemul_k4",call_mlkem_basemul_k4);
-  timingtest(arm,"mlkem_intt",call_mlkem_intt);
+  timingtest(all,"mlkem_intt",call_mlkem_intt);
   timingtest(arm,"mlkem_mulcache_compute",call_mlkem_mulcache_compute);
   timingtest(all,"mlkem_ntt",call_mlkem_ntt);
   timingtest(all,"mlkem_reduce",call_mlkem_reduce);

common/mlkem_mldsa.ml

@@ -52,6 +52,55 @@ let avx2_ntt_order = define
  `avx2_ntt_order i =
     bitreverse7(64 * (i DIV 64) + ((i MOD 64) DIV 16) + 4 * (i MOD 16))`;;
 
+let avx2_ntt_order' = define
+ `avx2_ntt_order' i =
+    let j = bitreverse7 i in
+    (64 * (j DIV 64) + 16 * (j MOD 4) + (j MOD 64) DIV 4)`;;
+
+let avx2_reorder = define
+ `avx2_reorder i =
+    let r = (i DIV 16) MOD 2
+    and q = 16 * (i DIV 32) + i MOD 16 in
+    2 * avx2_ntt_order q + r`;;
+
+let avx2_reorder' = define
+ `avx2_reorder' i =
+    let r = i MOD 2
+    and q = avx2_ntt_order'(i DIV 2) in
+    (q DIV 16) * 32 + r * 16 + q MOD 16`;;
+
+(* ------------------------------------------------------------------------- *)
+(* The simpler ones as used on ARM are actually involutions.                 *)
+(* ------------------------------------------------------------------------- *)
+
+let BITREVERSE7_INVOLUTION = prove
+ (`!n. n < 128 ==> bitreverse7(bitreverse7 n) = n`,
+  CONV_TAC EXPAND_CASES_CONV THEN REWRITE_TAC[bitreverse7] THEN
+  CONV_TAC(DEPTH_CONV WORD_NUM_RED_CONV));;
+
+let BITREVERSE_PAIRS_INVOLUTION = prove
+ (`!n. n < 256 ==> bitreverse_pairs(bitreverse_pairs n) = n`,
+  CONV_TAC EXPAND_CASES_CONV THEN
+  REWRITE_TAC[bitreverse_pairs; bitreverse7] THEN
+  CONV_TAC(DEPTH_CONV WORD_NUM_RED_CONV));;
+
+let AVX2_NTT_ORDER_INVOLUTION = prove
+ (`!n. n < 128 ==> avx2_ntt_order'(avx2_ntt_order n) = n /\
+                   avx2_ntt_order(avx2_ntt_order' n) = n`,
+  CONV_TAC EXPAND_CASES_CONV THEN
+  REWRITE_TAC[avx2_ntt_order; avx2_ntt_order'; bitreverse7] THEN
+  CONV_TAC(TOP_DEPTH_CONV let_CONV) THEN
+  CONV_TAC(DEPTH_CONV WORD_NUM_RED_CONV));;
+
+let AVX2_REORDER_INVOLUTION = prove
+ (`!n. n < 256 ==> avx2_reorder'(avx2_reorder n) = n /\
+                   avx2_reorder(avx2_reorder' n) = n`,
+  CONV_TAC EXPAND_CASES_CONV THEN
+  REWRITE_TAC[avx2_reorder; avx2_reorder';
+              avx2_ntt_order; avx2_ntt_order'; bitreverse7] THEN
+  CONV_TAC(TOP_DEPTH_CONV let_CONV) THEN
+  CONV_TAC(DEPTH_CONV WORD_NUM_RED_CONV));;
+
 (* ------------------------------------------------------------------------- *)
 (* AVX2-optimized ordering for ML-DSA NTT (swaps bit fields then reverses)   *)
 (* ------------------------------------------------------------------------- *)
@@ -107,6 +156,15 @@ let avx2_forward_ntt = define
                        &17 pow ((2 * avx2_ntt_order q + 1) * j))
     rem &3329`;;
 
+let avx2_inverse_ntt = define
+ `avx2_inverse_ntt f k =
+    (&512 * isum (0..127)
+                 (\j. f(avx2_ntt_order' j DIV 16 * 32 +
+                        k MOD 2 * 16 +
+                        avx2_ntt_order' j MOD 16) *
+                      &1175 pow ((2 * j + 1) * k DIV 2)))
+    rem &3329`;;
+
 let mldsa_forward_ntt = define
  `mldsa_forward_ntt f k =
     isum (0..255) (\j. f j * &1753 pow ((2 * mldsa_avx2_ntt_order k + 1) * j))
@@ -133,6 +191,26 @@ let INVERSE_NTT = prove
   CONV_TAC INT_REM_DOWN_CONV THEN REWRITE_TAC[INT_MUL_ASSOC] THEN
   ONCE_REWRITE_TAC[GSYM INT_MUL_REM] THEN CONV_TAC INT_REDUCE_CONV);;
 
+let AVX2_FORWARD_NTT = prove
+ (`avx2_forward_ntt = reorder avx2_reorder o pure_forward_ntt`,
+  REWRITE_TAC[FUN_EQ_THM; o_DEF; avx2_reorder; reorder] THEN
+  REWRITE_TAC[avx2_forward_ntt; pure_forward_ntt] THEN
+  MAP_EVERY X_GEN_TAC [`x:num->int`; `k:num`] THEN
+  CONV_TAC(ONCE_DEPTH_CONV let_CONV) THEN
+  SIMP_TAC[MOD_MULT_ADD; DIV_MULT_ADD; ARITH_EQ; MOD_MOD_REFL] THEN
+  REWRITE_TAC[ARITH_RULE `x MOD 2 DIV 2 = 0`; ADD_CLAUSES]);;
+
+let AVX2_INVERSE_NTT = prove
+ (`avx2_inverse_ntt = tomont_3329 o pure_inverse_ntt o reorder avx2_reorder'`,
+  REWRITE_TAC[FUN_EQ_THM; o_DEF; avx2_reorder'; reorder] THEN
+  REWRITE_TAC[avx2_inverse_ntt; pure_inverse_ntt; tomont_3329] THEN
+  REWRITE_TAC[ARITH_RULE `(2 * x + i MOD 2) DIV 2 = x`] THEN
+  REWRITE_TAC[MOD_MULT_ADD; MOD_MOD_REFL] THEN
+  MAP_EVERY X_GEN_TAC [`x:num->int`; `k:num`] THEN
+  CONV_TAC(ONCE_DEPTH_CONV let_CONV) THEN
+  CONV_TAC INT_REM_DOWN_CONV THEN REWRITE_TAC[INT_MUL_ASSOC] THEN
+  ONCE_REWRITE_TAC[GSYM INT_MUL_REM] THEN CONV_TAC INT_REDUCE_CONV);;
+
 let MLDSA_FORWARD_NTT = prove
  (`mldsa_forward_ntt f k =
    isum (0..255) (\j. f j * &1753 pow ((2 * mldsa_avx2_ntt_order k + 1) * j)) rem &8380417`,
@@ -198,6 +276,25 @@ let INVERSE_NTT_ALT = prove
   CONV_TAC INT_REM_DOWN_CONV THEN
   AP_THM_TAC THEN AP_TERM_TAC THEN CONV_TAC INT_ARITH);;
 
+let AVX2_INVERSE_NTT_ALT = prove
+ (`avx2_inverse_ntt f k =
+    isum (0..127)
+      (\j. f(avx2_ntt_order' j DIV 16 * 32 +
+             k MOD 2 * 16 +
+             avx2_ntt_order' j MOD 16) *
+           (&512 *
+            (&1175 pow ((2 * j + 1) * k DIV 2)) rem &3329)
+           rem &3329) rem &3329`,
+  REWRITE_TAC[avx2_inverse_ntt; GSYM ISUM_LMUL] THEN
+  MATCH_MP_TAC (REWRITE_RULE[] (ISPEC
+      `(\x y. x rem &3329 = y rem &3329)` ISUM_RELATED)) THEN
+  REWRITE_TAC[INT_REM_EQ; FINITE_NUMSEG; INT_CONG_ADD] THEN
+  X_GEN_TAC `i:num` THEN DISCH_TAC THEN
+  REWRITE_TAC[GSYM INT_OF_NUM_REM; GSYM INT_OF_NUM_CLAUSES;
+              GSYM INT_REM_EQ] THEN
+  CONV_TAC INT_REM_DOWN_CONV THEN
+  AP_THM_TAC THEN AP_TERM_TAC THEN CONV_TAC INT_ARITH);;
+
 let FORWARD_NTT_CONV =
   GEN_REWRITE_CONV I [FORWARD_NTT_ALT] THENC
   LAND_CONV EXPAND_ISUM_CONV THENC
@@ -212,6 +309,12 @@ let AVX2_NTT_ORDER_CLAUSES = end_itlist CONJ (map
   GEN_REWRITE_CONV I [BITREVERSE7_CLAUSES])
  (map (curry mk_comb `avx2_ntt_order` o mk_small_numeral) (0--127)));;
 
+let AVX2_NTT_ORDER_CLAUSES' = end_itlist CONJ (map
+ (GEN_REWRITE_CONV I [avx2_ntt_order'] THENC DEPTH_CONV WORD_NUM_RED_CONV THENC
+ DEPTH_CONV let_CONV THENC
+ GEN_REWRITE_CONV ONCE_DEPTH_CONV [BITREVERSE7_CLAUSES] THENC NUM_REDUCE_CONV)
+ (map (curry mk_comb `avx2_ntt_order'` o mk_small_numeral) (0--127)));;
+
 let AVX2_FORWARD_NTT_CONV =
   GEN_REWRITE_CONV I [AVX2_FORWARD_NTT_ALT] THENC
   NUM_REDUCE_CONV THENC ONCE_DEPTH_CONV let_CONV THENC
@@ -231,6 +334,16 @@ let INVERSE_NTT_CONV =
   GEN_REWRITE_CONV DEPTH_CONV [INT_OF_NUM_POW; INT_OF_NUM_REM] THENC
   ONCE_DEPTH_CONV EXP_MOD_CONV THENC INT_REDUCE_CONV;;
 
+let AVX2_INVERSE_NTT_CONV =
+  GEN_REWRITE_CONV I [AVX2_INVERSE_NTT_ALT] THENC
+  NUM_REDUCE_CONV THENC ONCE_DEPTH_CONV let_CONV THENC
+  LAND_CONV EXPAND_ISUM_CONV THENC
+  DEPTH_CONV NUM_RED_CONV THENC
+  GEN_REWRITE_CONV ONCE_DEPTH_CONV [AVX2_NTT_ORDER_CLAUSES'] THENC
+  DEPTH_CONV NUM_RED_CONV THENC
+  GEN_REWRITE_CONV DEPTH_CONV [INT_OF_NUM_POW; INT_OF_NUM_REM] THENC
+  ONCE_DEPTH_CONV EXP_MOD_CONV THENC INT_REDUCE_CONV;;
+
 (* ------------------------------------------------------------------------- *)
 (* Explicit computation rules to evaluate mod-8380417 powers less naively.   *)
 (* ------------------------------------------------------------------------- *)
@@ -672,7 +785,7 @@ let CONGBOUND_BARRED_X86 = prove
  (`!a a' l u.
         ((ival a == a') (mod &3329) /\ l <= ival a /\ ival a <= u)
         ==> (ival(barred_x86 a) == a') (mod &3329) /\
-            &0 <= ival(barred_x86 a) /\ ival(barred_x86 a) < &6658`,
+            &0 <= ival(barred_x86 a) /\ ival(barred_x86 a) <= &6657`,
   REPEAT GEN_TAC THEN STRIP_TAC THEN REWRITE_TAC[barred_x86] THEN
   REWRITE_TAC[WORD_BLAST
    `word_ishr (word_subword (x:int32) (16,16):int16) 10 =

include/s2n-bignum.h

@@ -1002,6 +1002,10 @@ extern void mlkem_basemul_k4(int16_t r[S2N_BIGNUM_STATIC 256],const int16_t a[S2
 // Input a[256] (signed 16-bit words), z_01234[80] (signed 16-bit words), z_56[384] (signed 16-bit words); output a[256] (signed 16-bit words)
 extern void mlkem_intt(int16_t a[S2N_BIGNUM_STATIC 256],const int16_t z_01234[S2N_BIGNUM_STATIC 80],const int16_t z_56[S2N_BIGNUM_STATIC 384]);
 
+// Inverse number-theoretic transform from ML-KEM
+// Input a[256] (signed 16-bit words), qdata[624]; output a[256] (signed 16-bit words)
+extern void mlkem_intt_x86(int16_t a[S2N_BIGNUM_STATIC 256],const int16_t qdata[S2N_BIGNUM_STATIC 624]);
+
 // Precompute the mulcache data for a polynomial in the NTT domain
 // Inputs a[256], z[128] and t[128] (signed 16-bit words); output x[128] (signed 16-bit words)
 extern void mlkem_mulcache_compute(int16_t x[S2N_BIGNUM_STATIC 128],const int16_t a[S2N_BIGNUM_STATIC 256],const int16_t z[S2N_BIGNUM_STATIC 128],const int16_t t[S2N_BIGNUM_STATIC 128]);

tests/test.c

@@ -12106,26 +12106,33 @@ uint64_t t, i;
 
 int test_mlkem_intt(void)
 {
-#ifdef __x86_64__
-  return 1;
-#else
   uint64_t t, i;
-  int16_t a[256], b[256], c[256];
+  int16_t a[256] __attribute__((aligned(32)));
+  int16_t b[256] __attribute__((aligned(32)));
+  int16_t c[256] __attribute__((aligned(32)));
   printf("Testing mlkem_intt with %d cases\n",tests);
 
   for (t = 0; t < tests; ++t)
    { for (i = 0; i < 256; ++i)
         a[i] = (int16_t) (random64()); // any int16_t inputs allowed
      for (i = 0; i < 256; ++i) b[i] = a[i];
+#ifdef __x86_64__
+     mlkem_poly_to_avx2_layout(b);
+     mlkem_intt_x86(b,mlkem_qdata);
+#else
      mlkem_intt(b,intt_zetas_layer01234,intt_zetas_layer56);
+#endif
+
      reference_bitreverse(c,a);
      reference_inverse_ntt(c,c);
      reference_tomont3329(c,c);
+
+
      for (i = 0; i < 256; ++i)
       { if (rem_3329(b[i]) != rem_3329(c[i]))
          { printf("Error in iNTT element i = %"PRIu64"; code[i] = 0x%04"PRIx16
                   " while reference[i] = 0x%04"PRIx16"\n",
-                  i,b[i],c[i]);
+                  i,rem_3329(b[i]),rem_3329(c[i]));
            return 1;
          }
       }
@@ -12140,7 +12147,6 @@ int test_mlkem_intt(void)
    }
   printf("All OK\n");
   return 0;
-#endif
 }
 
 int test_mlkem_mulcache_compute(void)
@@ -15622,6 +15628,7 @@ int main(int argc, char *argv[])
   functionaltest(all,"mlkem_basemul_k2",test_mlkem_basemul_k2);
   functionaltest(all,"mlkem_basemul_k3",test_mlkem_basemul_k3);
   functionaltest(all,"mlkem_basemul_k4",test_mlkem_basemul_k4);
+  functionaltest(all,"mlkem_intt",test_mlkem_intt);
   functionaltest(all,"mlkem_ntt",test_mlkem_ntt);
   functionaltest(all,"mlkem_reduce",test_mlkem_reduce);
   functionaltest(bmi,"p256_montjadd",test_p256_montjadd);
@@ -15682,7 +15689,6 @@ int main(int argc, char *argv[])
     functionaltest(all,"bignum_copy_row_from_table_16",test_bignum_copy_row_from_table_16);
     functionaltest(all,"bignum_copy_row_from_table_32",test_bignum_copy_row_from_table_32);
     functionaltest(all,"bignum_emontredc_8n_cdiff",test_bignum_emontredc_8n_cdiff);
-    functionaltest(arm,"mlkem_intt",test_mlkem_intt);
     functionaltest(arm,"mlkem_mulcache_compute",test_mlkem_mulcache_compute);
     functionaltest(arm,"mlkem_tobytes",test_mlkem_tobytes);
     functionaltest(arm,"mlkem_tomont",test_mlkem_tomont);

tools/collect-signatures.py

@@ -333,6 +333,7 @@ def stripPrefixes(s, prefixes):
   "mldsa_ntt",
   "mldsa_poly_reduce",
   "mlkem_ntt_x86",
+  "mlkem_intt_x86",
 ]
 
 for arch in ["arm","x86"]:

x86/Makefile

@@ -255,6 +255,7 @@ BIGNUM_OBJ = curve25519/bignum_add_p25519.o \
              mlkem/mlkem_basemul_k3.o \
              mlkem/mlkem_basemul_k4.o \
              mlkem/mlkem_ntt.o \
+             mlkem/mlkem_intt.o \
              mlkem/mlkem_reduce.o \
              p256/bignum_add_p256.o \
              p256/bignum_bigendian_4.o \