C-ref NTT: merge layer pairs, switch fqmul to twisted Montgomery

hanno-becker · hanno-becker · commit 5b598cc7d528 · 2026-05-15T14:23:43.000+01:00
Rewrite mlk_poly_ntt_c / mlk_poly_invntt_tomont_c to process two
layers at a time, with three 2-layer passes plus the leftover layer 7
as a single layer.

Introduces shared mlk_ct_butterfly and mlk_gs_butterfly helpers;
the inverse 2-layer block applies four GS butterflies and then
Barrett-reduces the additive outputs explicitly.

mlk_fqmul now takes a precomputed b_twisted = b * MLKEM_Q^{-1} mod 2^16
and uses a hi-mul / lo-mul-and-correct sequence in place of an inline
mlk_montgomery_reduce, dropping the QINV multiply. The mlk_zetas table
is regenerated as int16_t[128][2] of (zeta_mont, zeta_twisted) pairs.

Signed-off-by: Hanno Becker &lt;beckphan@amazon.co.uk&gt;
diff --git a/mlkem/src/poly.c b/mlkem/src/poly.c
diff --git a/mlkem/src/zetas.inc b/mlkem/src/zetas.inc
@@ -12,19 +12,39 @@
 
 /*
  * Table of zeta values used in the reference NTT and inverse NTT.
- * See autogen for details.
+ * Each row is (zeta_mont, zeta_twisted); see autogen for details.
  */
-static MLK_ALIGN const int16_t mlk_zetas[128] = {
-    -1044, -758,  -359,  -1517, 1493,  1422,  287,   202,  -171,  622,   1577,
-    182,   962,   -1202, -1474, 1468,  573,   -1325, 264,  383,   -829,  1458,
-    -1602, -130,  -681,  1017,  732,   608,   -1542, 411,  -205,  -1571, 1223,
-    652,   -552,  1015,  -1293, 1491,  -282,  -1544, 516,  -8,    -320,  -666,
-    -1618, -1162, 126,   1469,  -853,  -90,   -271,  830,  107,   -1421, -247,
-    -951,  -398,  961,   -1508, -725,  448,   -1065, 677,  -1275, -1103, 430,
-    555,   843,   -1251, 871,   1550,  105,   422,   587,  177,   -235,  -291,
-    -460,  1574,  1653,  -246,  778,   1159,  -147,  -777, 1483,  -602,  1119,
-    -1590, 644,   -872,  349,   418,   329,   -156,  -75,  817,   1097,  603,
-    610,   1322,  -1285, -1465, 384,   -1215, -136,  1218, -1335, -874,  220,
-    -1187, -1659, -1185, -1530, -1278, 794,   -1510, -854, -870,  478,   -108,
-    -308,  996,   991,   958,   -1460, 1522,  1628,
+static MLK_ALIGN const int16_t mlk_zetas[128][2] = {
+    {-1044, -20},    {-758, 31498},   {-359, 14745},   {-1517, 787},
+    {1493, 13525},   {1422, -12402},  {287, 28191},    {202, -16694},
+    {-171, -20907},  {622, 27758},    {1577, -3799},   {182, -15690},
+    {962, 10690},    {-1202, 1358},   {-1474, -11202}, {1468, 31164},
+    {573, -5827},    {-1325, 17363},  {264, -26360},   {383, -29057},
+    {-829, 5571},    {1458, -1102},   {-1602, 21438},  {-130, -26242},
+    {-681, -28073},  {1017, 24313},   {732, -10532},   {608, 8800},
+    {-1542, 18426},  {411, 8859},     {-205, 26675},   {-1571, -16163},
+    {1223, -5689},   {652, -6516},    {-552, 1496},    {1015, 30967},
+    {-1293, -23565}, {1491, 20179},   {-282, 20710},   {-1544, 25080},
+    {516, -12796},   {-8, 26616},     {-320, 16064},   {-666, -12442},
+    {-1618, 9134},   {-1162, -650},   {126, -25986},   {1469, 27837},
+    {-853, 19883},   {-90, -28250},   {-271, -15887},  {830, -8898},
+    {107, -28309},   {-1421, 9075},   {-247, -30199},  {-951, 18249},
+    {-398, 13426},   {961, 14017},    {-1508, -29156}, {-725, -12757},
+    {448, 16832},    {-1065, 4311},   {677, -24155},   {-1275, -17915},
+    {-1103, -335},   {430, 11182},    {555, -11477},   {843, 13387},
+    {-1251, -32227}, {871, -14233},   {1550, 20494},   {105, -21655},
+    {422, -27738},   {587, 13131},    {177, 945},      {-235, -4587},
+    {-291, -14883},  {-460, 23092},   {1574, 6182},    {1653, 5493},
+    {-246, 32010},   {778, -32502},   {1159, 10631},   {-147, 30317},
+    {-777, 29175},   {1483, -18741},  {-602, -28762},  {1119, 12639},
+    {-1590, -18486}, {644, 20100},    {-872, 17560},   {349, 18525},
+    {418, -14430},   {329, 19529},    {-156, -5276},   {-75, -12619},
+    {817, -31183},   {1097, 20297},   {603, 25435},    {610, 2146},
+    {1322, -7382},   {-1285, 15355},  {-1465, 24391},  {384, -32384},
+    {-1215, -20927}, {-136, -6280},   {1218, 10946},   {-1335, -14903},
+    {-874, 24214},   {220, -11044},   {-1187, 16989},  {-1659, 14469},
+    {-1185, 10335},  {-1530, -21498}, {-1278, -7934},  {794, -20198},
+    {-1510, -22502}, {-854, 23210},   {-870, 10906},   {478, -17442},
+    {-108, 31636},   {-308, -23860},  {996, 28644},    {991, -20257},
+    {958, 23998},    {-1460, 7756},   {1522, -17422},  {1628, 23132},
 };
diff --git a/proofs/cbmc/fqmul/fqmul_harness.c b/proofs/cbmc/fqmul/fqmul_harness.c
@@ -6,11 +6,11 @@
 #include "common.h"
 
 
-int16_t mlk_fqmul(int16_t a, int16_t b);
+int16_t mlk_fqmul(int16_t a, int16_t b, int16_t b_twisted);
 
 void harness(void)
 {
-  int16_t a, b, r;
+  int16_t a, b, b_twisted, r;
 
-  r = mlk_fqmul(a, b);
+  r = mlk_fqmul(a, b, b_twisted);
 }
diff --git a/proofs/cbmc/invntt_2_layers/Makefile b/proofs/cbmc/invntt_2_layers/Makefile
@@ -0,0 +1,34 @@
+# Copyright (c) The mlkem-native project authors
+# SPDX-License-Identifier: Apache-2.0 OR ISC OR MIT
+
+include ../Makefile_params.common
+
+HARNESS_ENTRY = harness
+HARNESS_FILE = invntt_2_layers_harness
+
+# This should be a unique identifier for this proof, and will appear on the
+# Litani dashboard. It can be human-readable and contain spaces if you wish.
+PROOF_UID = mlk_invntt_2_layers
+
+DEFINES +=
+INCLUDES +=
+
+REMOVE_FUNCTION_BODY +=
+UNWINDSET +=
+
+PROOF_SOURCES += $(PROOFDIR)/$(HARNESS_FILE).c
+PROJECT_SOURCES += $(SRCDIR)/mlkem/src/poly.c
+
+CHECK_FUNCTION_CONTRACTS=mlk_invntt_2_layers
+USE_FUNCTION_CONTRACTS=mlk_invntt_2_layers_block
+APPLY_LOOP_CONTRACTS=on
+USE_DYNAMIC_FRAMES=1
+
+EXTERNAL_SAT_SOLVER=
+CBMCFLAGS=--smt2
+
+FUNCTION_NAME = mlk_invntt_2_layers
+
+CBMC_OBJECT_BITS = 8
+
+include ../Makefile.common
diff --git a/proofs/cbmc/invntt_2_layers/invntt_2_layers_harness.c b/proofs/cbmc/invntt_2_layers/invntt_2_layers_harness.c
@@ -0,0 +1,15 @@
+// Copyright (c) The mlkem-native project authors
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: MIT-0
+
+#include "poly.h"
+
+
+void mlk_invntt_2_layers(int16_t *r, unsigned layer);
+
+void harness(void)
+{
+  int16_t *r;
+  unsigned layer;
+  mlk_invntt_2_layers(r, layer);
+}
diff --git a/proofs/cbmc/invntt_2_layers_block/Makefile b/proofs/cbmc/invntt_2_layers_block/Makefile
@@ -0,0 +1,33 @@
+# Copyright (c) The mlkem-native project authors
+# SPDX-License-Identifier: Apache-2.0 OR ISC OR MIT
+
+include ../Makefile_params.common
+
+HARNESS_ENTRY = harness
+HARNESS_FILE = invntt_2_layers_block_harness
+
+PROOF_UID = mlk_invntt_2_layers_block
+
+DEFINES +=
+INCLUDES +=
+
+REMOVE_FUNCTION_BODY +=
+UNWINDSET +=
+
+PROOF_SOURCES += $(PROOFDIR)/$(HARNESS_FILE).c
+PROJECT_SOURCES += $(SRCDIR)/mlkem/src/poly.c
+
+CHECK_FUNCTION_CONTRACTS=mlk_invntt_2_layers_block
+USE_FUNCTION_CONTRACTS = mlk_fqmul
+USE_FUNCTION_CONTRACTS += mlk_barrett_reduce
+APPLY_LOOP_CONTRACTS=on
+USE_DYNAMIC_FRAMES=1
+
+EXTERNAL_SAT_SOLVER=
+CBMCFLAGS=--smt2
+
+FUNCTION_NAME = mlk_invntt_2_layers_block
+
+CBMC_OBJECT_BITS = 9
+
+include ../Makefile.common
diff --git a/proofs/cbmc/invntt_2_layers_block/invntt_2_layers_block_harness.c b/proofs/cbmc/invntt_2_layers_block/invntt_2_layers_block_harness.c
@@ -0,0 +1,20 @@
+// Copyright (c) The mlkem-native project authors
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: MIT-0
+
+#include <stdint.h>
+#include "common.h"
+
+
+void mlk_invntt_2_layers_block(int16_t *r, unsigned start, unsigned len,
+                               int16_t z0, int16_t z0_twst, int16_t z1,
+                               int16_t z1_twst, int16_t z2, int16_t z2_twst);
+
+void harness(void)
+{
+  int16_t *r;
+  unsigned start, len;
+  int16_t z0, z0_twst, z1, z1_twst, z2, z2_twst;
+  mlk_invntt_2_layers_block(r, start, len, z0, z0_twst, z1, z1_twst, z2,
+                            z2_twst);
+}
diff --git a/proofs/cbmc/ntt_2_layers/Makefile b/proofs/cbmc/ntt_2_layers/Makefile
@@ -0,0 +1,42 @@
+# Copyright (c) The mlkem-native project authors
+# SPDX-License-Identifier: Apache-2.0 OR ISC OR MIT
+
+include ../Makefile_params.common
+
+HARNESS_ENTRY = harness
+HARNESS_FILE = ntt_2_layers_harness
+
+# This should be a unique identifier for this proof, and will appear on the
+# Litani dashboard. It can be human-readable and contain spaces if you wish.
+PROOF_UID = mlk_ntt_2_layers
+
+DEFINES +=
+INCLUDES +=
+
+REMOVE_FUNCTION_BODY +=
+UNWINDSET +=
+
+PROOF_SOURCES += $(PROOFDIR)/$(HARNESS_FILE).c
+PROJECT_SOURCES += $(SRCDIR)/mlkem/src/poly.c
+
+CHECK_FUNCTION_CONTRACTS=mlk_ntt_2_layers
+USE_FUNCTION_CONTRACTS=mlk_ntt_2_layers_block
+APPLY_LOOP_CONTRACTS=on
+USE_DYNAMIC_FRAMES=1
+
+# Disable any setting of EXTERNAL_SAT_SOLVER, and choose SMT backend instead
+EXTERNAL_SAT_SOLVER=
+CBMCFLAGS=--bitwuzla
+
+FUNCTION_NAME = mlk_ntt_2_layers
+
+# If this proof is found to consume huge amounts of RAM, you can set the
+# EXPENSIVE variable. With new enough versions of the proof tools, this will
+# restrict the number of EXPENSIVE CBMC jobs running at once. See the
+# documentation in Makefile.common under the "Job Pools" heading for details.
+# EXPENSIVE = true
+
+# This function is large enough to need...
+CBMC_OBJECT_BITS = 8
+
+include ../Makefile.common
diff --git a/proofs/cbmc/ntt_2_layers/ntt_2_layers_harness.c b/proofs/cbmc/ntt_2_layers/ntt_2_layers_harness.c
@@ -0,0 +1,15 @@
+// Copyright (c) The mlkem-native project authors
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: MIT-0
+
+#include "poly.h"
+
+
+void mlk_ntt_2_layers(int16_t *r, unsigned layer);
+
+void harness(void)
+{
+  int16_t *r;
+  unsigned layer;
+  mlk_ntt_2_layers(r, layer);
+}
diff --git a/proofs/cbmc/ntt_2_layers_block/Makefile b/proofs/cbmc/ntt_2_layers_block/Makefile
@@ -0,0 +1,32 @@
+# Copyright (c) The mlkem-native project authors
+# SPDX-License-Identifier: Apache-2.0 OR ISC OR MIT
+
+include ../Makefile_params.common
+
+HARNESS_ENTRY = harness
+HARNESS_FILE = ntt_2_layers_block_harness
+
+PROOF_UID = mlk_ntt_2_layers_block
+
+DEFINES +=
+INCLUDES +=
+
+REMOVE_FUNCTION_BODY +=
+UNWINDSET +=
+
+PROOF_SOURCES += $(PROOFDIR)/$(HARNESS_FILE).c
+PROJECT_SOURCES += $(SRCDIR)/mlkem/src/poly.c
+
+CHECK_FUNCTION_CONTRACTS=mlk_ntt_2_layers_block
+USE_FUNCTION_CONTRACTS=mlk_fqmul
+APPLY_LOOP_CONTRACTS=on
+USE_DYNAMIC_FRAMES=1
+
+EXTERNAL_SAT_SOLVER=
+CBMCFLAGS=--bitwuzla
+
+FUNCTION_NAME = mlk_ntt_2_layers_block
+
+CBMC_OBJECT_BITS = 9
+
+include ../Makefile.common
diff --git a/proofs/cbmc/ntt_2_layers_block/ntt_2_layers_block_harness.c b/proofs/cbmc/ntt_2_layers_block/ntt_2_layers_block_harness.c
@@ -0,0 +1,21 @@
+// Copyright (c) The mlkem-native project authors
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: MIT-0
+
+#include <stdint.h>
+#include "common.h"
+
+
+void mlk_ntt_2_layers_block(int16_t *r, unsigned start, unsigned len,
+                            int16_t z0, int16_t z0_twst, int16_t z1,
+                            int16_t z1_twst, int16_t z2, int16_t z2_twst,
+                            const int16_t bound);
+
+void harness(void)
+{
+  int16_t *r;
+  unsigned start, len;
+  int16_t z0, z0_twst, z1, z1_twst, z2, z2_twst, bound;
+  mlk_ntt_2_layers_block(r, start, len, z0, z0_twst, z1, z1_twst, z2, z2_twst,
+                         bound);
+}
diff --git a/proofs/cbmc/ntt_butterfly_block/ntt_butterfly_block_harness.c b/proofs/cbmc/ntt_butterfly_block/ntt_butterfly_block_harness.c
@@ -6,12 +6,12 @@
 #include "common.h"
 
 
-void mlk_ntt_butterfly_block(int16_t *r, int16_t root, unsigned start,
-                             unsigned len, unsigned bound);
+void mlk_ntt_butterfly_block(int16_t *r, int16_t root, int16_t root_twisted,
+                             unsigned start, unsigned len, unsigned bound);
 
 void harness(void)
 {
-  int16_t *r, root;
+  int16_t *r, root, root_twisted;
   unsigned start, stride, bound;
-  mlk_ntt_butterfly_block(r, root, start, stride, bound);
+  mlk_ntt_butterfly_block(r, root, root_twisted, start, stride, bound);
 }
diff --git a/proofs/cbmc/poly_invntt_tomont_c/Makefile b/proofs/cbmc/poly_invntt_tomont_c/Makefile
@@ -22,6 +22,7 @@ PROJECT_SOURCES += $(SRCDIR)/mlkem/src/poly.c
 CHECK_FUNCTION_CONTRACTS=mlk_poly_invntt_tomont_c
 USE_FUNCTION_CONTRACTS = mlk_fqmul
 USE_FUNCTION_CONTRACTS += mlk_invntt_layer
+USE_FUNCTION_CONTRACTS += mlk_invntt_2_layers
 APPLY_LOOP_CONTRACTS=on
 USE_DYNAMIC_FRAMES=1
 
diff --git a/proofs/cbmc/poly_ntt_c/Makefile b/proofs/cbmc/poly_ntt_c/Makefile
@@ -21,6 +21,7 @@ PROJECT_SOURCES += $(SRCDIR)/mlkem/src/poly.c
 
 CHECK_FUNCTION_CONTRACTS=mlk_poly_ntt_c
 USE_FUNCTION_CONTRACTS=mlk_ntt_layer
+USE_FUNCTION_CONTRACTS+=mlk_ntt_2_layers
 APPLY_LOOP_CONTRACTS=on
 USE_DYNAMIC_FRAMES=1
 
diff --git a/scripts/autogen b/scripts/autogen
@@ -714,17 +714,25 @@ def signed_reduce(a):
 
 def gen_c_zetas():
     """Generate source and header file for zeta values used in
-    the reference NTT and invNTT"""
+    the reference NTT and invNTT.
+
+    Each yielded entry is a pair (zeta_mont, zeta_twisted) where:
+      - zeta_mont    = signed canonical representative of zeta * 2^16 mod q
+      - zeta_twisted = signed canonical 16-bit representative of
+                       zeta_mont * q^{-1} mod 2^16
+    The twisted form lets the C-reference fqmul skip the extra QINV
+    multiply that mlk_montgomery_reduce would otherwise perform."""
 
     # The zeta values are the powers of the chosen root of unity (17),
-    # converted to Montgomery form.
+    # converted to Montgomery form (and to its twisted-Montgomery counterpart).
 
-    zeta = []
+    zeta_pairs = []
     for i in range(128):
-        zeta.append(signed_reduce(pow(root_of_unity, i, modulus) * montgomery_factor))
+        raw = pow(root_of_unity, i, modulus)
+        zeta_pairs.append(prepare_root_for_montmul(raw))
 
     # The source code stores the zeta table in bit reversed form
-    yield from (zeta[bitreverse(i, 7)] for i in range(128))
+    yield from (zeta_pairs[bitreverse(i, 7)] for i in range(128))
 
 
 def gen_c_zeta_file():
@@ -733,10 +741,11 @@ def gen_c_zeta_file():
         yield ""
         yield "/*"
         yield " * Table of zeta values used in the reference NTT and inverse NTT."
-        yield " * See autogen for details."
+        yield " * Each row is (zeta_mont, zeta_twisted); see autogen for details."
         yield " */"
-        yield "static MLK_ALIGN const int16_t mlk_zetas[128] = {"
-        yield from map(lambda t: str(t) + ",", gen_c_zetas())
+        yield "static MLK_ALIGN const int16_t mlk_zetas[128][2] = {"
+        for z_mont, z_twisted in gen_c_zetas():
+            yield f"{{ {z_mont}, {z_twisted} }},"
         yield "};"
         yield ""
 
diff --git a/scripts/check-magic b/scripts/check-magic

Original file line number	Diff line number	Diff line change
`@@ -6,11 +6,11 @@`
`6`	`6`	`#include "common.h"`
`7`	`7`
`8`	`8`
`9`		`-int16_t mlk_fqmul(int16_t a, int16_t b);`
	`9`	`+int16_t mlk_fqmul(int16_t a, int16_t b, int16_t b_twisted);`
`10`	`10`
`11`	`11`	`void harness(void)`
`12`	`12`	`{`
`13`		`- int16_t a, b, r;`
	`13`	`+ int16_t a, b, b_twisted, r;`
`14`	`14`
`15`		`- r = mlk_fqmul(a, b);`
	`15`	`+ r = mlk_fqmul(a, b, b_twisted);`
`16`	`16`	`}`