Speed up C-reference NTT/invNTT with twisted zetas + 2-layer merging

Replace the single-layer C-reference forward and inverse NTT in
`mldsa/src/poly.c` with one that merges two layers each.

Also, store each twiddle alongside its precomputed twist, letting
`mld_fqmul(a, b, b_twisted)` drop the multiply with MLDSA_Q^{-1}
that was previously hidden inside `mld_montgomery_reduce`.

Mirrors pq-code-package/mlkem-native/#463 and pq-code/package/mlkem-native/#683

Signed-off-by: Hanno Becker <beckphan@amazon.co.uk>

Author: hanno-becker

mldsa/src/poly.c

@@ -12,44 +12,263 @@
 
 /*
  * 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 const int32_t mld_zetas[MLDSA_N] = {
-    0,        25847,    -2608894, -518909,  237124,   -777960,  -876248,
-    466468,   1826347,  2353451,  -359251,  -2091905, 3119733,  -2884855,
-    3111497,  2680103,  2725464,  1024112,  -1079900, 3585928,  -549488,
-    -1119584, 2619752,  -2108549, -2118186, -3859737, -1399561, -3277672,
-    1757237,  -19422,   4010497,  280005,   2706023,  95776,    3077325,
-    3530437,  -1661693, -3592148, -2537516, 3915439,  -3861115, -3043716,
-    3574422,  -2867647, 3539968,  -300467,  2348700,  -539299,  -1699267,
-    -1643818, 3505694,  -3821735, 3507263,  -2140649, -1600420, 3699596,
-    811944,   531354,   954230,   3881043,  3900724,  -2556880, 2071892,
-    -2797779, -3930395, -1528703, -3677745, -3041255, -1452451, 3475950,
-    2176455,  -1585221, -1257611, 1939314,  -4083598, -1000202, -3190144,
-    -3157330, -3632928, 126922,   3412210,  -983419,  2147896,  2715295,
-    -2967645, -3693493, -411027,  -2477047, -671102,  -1228525, -22981,
-    -1308169, -381987,  1349076,  1852771,  -1430430, -3343383, 264944,
-    508951,   3097992,  44288,    -1100098, 904516,   3958618,  -3724342,
-    -8578,    1653064,  -3249728, 2389356,  -210977,  759969,   -1316856,
-    189548,   -3553272, 3159746,  -1851402, -2409325, -177440,  1315589,
-    1341330,  1285669,  -1584928, -812732,  -1439742, -3019102, -3881060,
-    -3628969, 3839961,  2091667,  3407706,  2316500,  3817976,  -3342478,
-    2244091,  -2446433, -3562462, 266997,   2434439,  -1235728, 3513181,
-    -3520352, -3759364, -1197226, -3193378, 900702,   1859098,  909542,
-    819034,   495491,   -1613174, -43260,   -522500,  -655327,  -3122442,
-    2031748,  3207046,  -3556995, -525098,  -768622,  -3595838, 342297,
-    286988,   -2437823, 4108315,  3437287,  -3342277, 1735879,  203044,
-    2842341,  2691481,  -2590150, 1265009,  4055324,  1247620,  2486353,
-    1595974,  -3767016, 1250494,  2635921,  -3548272, -2994039, 1869119,
-    1903435,  -1050970, -1333058, 1237275,  -3318210, -1430225, -451100,
-    1312455,  3306115,  -1962642, -1279661, 1917081,  -2546312, -1374803,
-    1500165,  777191,   2235880,  3406031,  -542412,  -2831860, -1671176,
-    -1846953, -2584293, -3724270, 594136,   -3776993, -2013608, 2432395,
-    2454455,  -164721,  1957272,  3369112,  185531,   -1207385, -3183426,
-    162844,   1616392,  3014001,  810149,   1652634,  -3694233, -1799107,
-    -3038916, 3523897,  3866901,  269760,   2213111,  -975884,  1717735,
-    472078,   -426683,  1723600,  -1803090, 1910376,  -1667432, -1104333,
-    -260646,  -3833893, -2939036, -2235985, -420899,  -2286327, 183443,
-    -976891,  1612842,  -3545687, -554416,  3919660,  -48306,   -1362209,
-    3937738,  1400424,  -846154,  1976782,
+static const int32_t mld_zetas[MLDSA_N][2] = {
+    {0, 0},
+    {25847, 1830765815},
+    {-2608894, -1929875198},
+    {-518909, -1927777021},
+    {237124, 1640767044},
+    {-777960, 1477910808},
+    {-876248, 1612161320},
+    {466468, 1640734244},
+    {1826347, 308362795},
+    {2353451, -1815525077},
+    {-359251, -1374673747},
+    {-2091905, -1091570561},
+    {3119733, -1929495947},
+    {-2884855, 515185417},
+    {3111497, -285697463},
+    {2680103, 625853735},
+    {2725464, 1727305304},
+    {1024112, 2082316400},
+    {-1079900, -1364982364},
+    {3585928, 858240904},
+    {-549488, 1806278032},
+    {-1119584, 222489248},
+    {2619752, -346752664},
+    {-2108549, 684667771},
+    {-2118186, 1654287830},
+    {-3859737, -878576921},
+    {-1399561, -1257667337},
+    {-3277672, -748618600},
+    {1757237, 329347125},
+    {-19422, 1837364258},
+    {4010497, -1443016191},
+    {280005, -1170414139},
+    {2706023, -1846138265},
+    {95776, -1631226336},
+    {3077325, -1404529459},
+    {3530437, 1838055109},
+    {-1661693, 1594295555},
+    {-3592148, -1076973524},
+    {-2537516, -1898723372},
+    {3915439, -594436433},
+    {-3861115, -202001019},
+    {-3043716, -475984260},
+    {3574422, -561427818},
+    {-2867647, 1797021249},
+    {3539968, -1061813248},
+    {-300467, 2059733581},
+    {2348700, -1661512036},
+    {-539299, -1104976547},
+    {-1699267, -1750224323},
+    {-1643818, -901666090},
+    {3505694, 418987550},
+    {-3821735, 1831915353},
+    {3507263, -1925356481},
+    {-2140649, 992097815},
+    {-1600420, 879957084},
+    {3699596, 2024403852},
+    {811944, 1484874664},
+    {531354, -1636082790},
+    {954230, -285388938},
+    {3881043, -1983539117},
+    {3900724, -1495136972},
+    {-2556880, -950076368},
+    {2071892, -1714807468},
+    {-2797779, -952438995},
+    {-3930395, -1574918427},
+    {-1528703, -654783359},
+    {-3677745, 1350681039},
+    {-3041255, -1974159335},
+    {-1452451, -2143979939},
+    {3475950, 1651689966},
+    {2176455, 1599739335},
+    {-1585221, 140455867},
+    {-1257611, -1285853323},
+    {1939314, -1039411342},
+    {-4083598, -993005454},
+    {-1000202, 1955560694},
+    {-3190144, -1440787840},
+    {-3157330, 1529189038},
+    {-3632928, 568627424},
+    {126922, -2131021878},
+    {3412210, -783134478},
+    {-983419, -247357819},
+    {2147896, -588790216},
+    {2715295, 1518161567},
+    {-2967645, 289871779},
+    {-3693493, -86965173},
+    {-411027, -1262003603},
+    {-2477047, 1708872713},
+    {-671102, 2135294594},
+    {-1228525, 1787797779},
+    {-22981, -1018755525},
+    {-1308169, 1638590967},
+    {-381987, -889861155},
+    {1349076, -120646188},
+    {1852771, 1665705315},
+    {-1430430, -1669960606},
+    {-3343383, 1321868265},
+    {264944, -916321552},
+    {508951, 1225434135},
+    {3097992, 1155548552},
+    {44288, -1784632064},
+    {-1100098, 2143745726},
+    {904516, 666258756},
+    {3958618, 1210558298},
+    {-3724342, 675310538},
+    {-8578, -1261461890},
+    {1653064, -1555941048},
+    {-3249728, -318346816},
+    {2389356, -1999506068},
+    {-210977, 628664287},
+    {759969, -1499481951},
+    {-1316856, -1729304568},
+    {189548, -695180180},
+    {-3553272, 1422575624},
+    {3159746, -1375177022},
+    {-1851402, 1424130038},
+    {-2409325, 1777179795},
+    {-177440, -1185330464},
+    {1315589, 334803717},
+    {1341330, 235321234},
+    {1285669, -178766299},
+    {-1584928, 168022240},
+    {-812732, -518252220},
+    {-1439742, 1206536194},
+    {-3019102, 1957047970},
+    {-3881060, 985155484},
+    {-3628969, 1146323031},
+    {3839961, -894060583},
+    {2091667, -898413},
+    {3407706, 991903578},
+    {2316500, 1363007700},
+    {3817976, 746144248},
+    {-3342478, -1363460238},
+    {2244091, 912367099},
+    {-2446433, 30313375},
+    {-3562462, -1420958686},
+    {266997, -605900043},
+    {2434439, -44694137},
+    {-1235728, -326425360},
+    {3513181, 2032221021},
+    {-3520352, 2027833504},
+    {-3759364, 1176904444},
+    {-1197226, 1683520342},
+    {-3193378, 1904936414},
+    {900702, 14253662},
+    {1859098, -421552614},
+    {909542, -517299994},
+    {819034, 1257750362},
+    {495491, 1014493059},
+    {-1613174, -818371958},
+    {-43260, 2027935492},
+    {-522500, 1926727420},
+    {-655327, 863641633},
+    {-3122442, 1747917558},
+    {2031748, -1372618620},
+    {3207046, 1931587462},
+    {-3556995, 1819892093},
+    {-525098, -325927722},
+    {-768622, 128353682},
+    {-3595838, 1258381762},
+    {342297, 2124962073},
+    {286988, 908452108},
+    {-2437823, -1123881663},
+    {4108315, 885133339},
+    {3437287, -1223601433},
+    {-3342277, 1851023419},
+    {1735879, 137583815},
+    {203044, 1629985060},
+    {2842341, -1920467227},
+    {2691481, -1176751719},
+    {-2590150, -635454918},
+    {1265009, 1967222129},
+    {4055324, -1637785316},
+    {1247620, -1354528380},
+    {2486353, -642772911},
+    {1595974, 6363718},
+    {-3767016, -1536588520},
+    {1250494, -72690498},
+    {2635921, 45766801},
+    {-3548272, -1287922800},
+    {-2994039, 694382729},
+    {1869119, -314284737},
+    {1903435, 671509323},
+    {-1050970, 1136965286},
+    {-1333058, 235104446},
+    {1237275, 985022747},
+    {-3318210, -2070602178},
+    {-1430225, 1779436847},
+    {-451100, -1045062172},
+    {1312455, 963438279},
+    {3306115, 419615363},
+    {-1962642, 1116720494},
+    {-1279661, 831969619},
+    {1917081, -1078959975},
+    {-2546312, 1216882040},
+    {-1374803, 1042326957},
+    {1500165, -300448763},
+    {777191, 604552167},
+    {2235880, -270590488},
+    {3406031, 1405999311},
+    {-542412, 756955444},
+    {-2831860, -1021949428},
+    {-1671176, -1276805128},
+    {-1846953, 713994583},
+    {-2584293, -260312805},
+    {-3724270, 608791570},
+    {594136, 371462360},
+    {-3776993, 940195359},
+    {-2013608, 1554794072},
+    {2432395, 173440395},
+    {2454455, -1357098057},
+    {-164721, -1542497137},
+    {1957272, 1339088280},
+    {3369112, -2126092136},
+    {185531, -384158533},
+    {-1207385, 2061661095},
+    {-3183426, -2040058690},
+    {162844, -1316619236},
+    {1616392, 827959816},
+    {3014001, -883155599},
+    {810149, -853476187},
+    {1652634, -1039370342},
+    {-3694233, -596344473},
+    {-1799107, 1726753853},
+    {-3038916, -2047270596},
+    {3523897, 6087993},
+    {3866901, 702390549},
+    {269760, -1547952704},
+    {2213111, -1723816713},
+    {-975884, -110126092},
+    {1717735, -279505433},
+    {472078, 394851342},
+    {-426683, -1591599803},
+    {1723600, 565464272},
+    {-1803090, -260424530},
+    {1910376, 283780712},
+    {-1667432, -440824168},
+    {-1104333, -1758099917},
+    {-260646, -71875110},
+    {-3833893, 776003547},
+    {-2939036, 1119856484},
+    {-2235985, -1600929361},
+    {-420899, -1208667171},
+    {-2286327, 1123958025},
+    {183443, 1544891539},
+    {-976891, 879867909},
+    {1612842, -1499603926},
+    {-3545687, 201262505},
+    {-554416, 155290192},
+    {3919660, -1809756372},
+    {-48306, 2036925262},
+    {-1362209, 1934038751},
+    {3937738, -973777462},
+    {1400424, 400711272},
+    {-846154, -540420426},
+    {1976782, 374860238},
 };

mldsa/src/zetas.inc

@@ -20,7 +20,7 @@ PROOF_SOURCES += $(PROOFDIR)/$(HARNESS_FILE).c
 PROJECT_SOURCES += $(SRCDIR)/mldsa/src/poly.c
 
 CHECK_FUNCTION_CONTRACTS=mld_fqmul
-USE_FUNCTION_CONTRACTS=mld_montgomery_reduce
+USE_FUNCTION_CONTRACTS=
 APPLY_LOOP_CONTRACTS=on
 USE_DYNAMIC_FRAMES=1
 

proofs/cbmc/fqmul/Makefile

@@ -3,9 +3,9 @@
 
 #include "poly.h"
 
-int32_t mld_fqmul(int32_t a, int32_t b);
+int32_t mld_fqmul(int32_t a, int32_t b, int32_t b_twisted);
 void harness(void)
 {
-  int32_t a, b, r;
-  r = mld_fqmul(a, b);
+  int32_t a, b, b_twisted, r;
+  r = mld_fqmul(a, b, b_twisted);
 }

proofs/cbmc/fqmul/fqmul_harness.c

@@ -20,7 +20,7 @@ PROOF_SOURCES += $(PROOFDIR)/$(HARNESS_FILE).c
 PROJECT_SOURCES += $(SRCDIR)/mldsa/src/poly.c
 
 CHECK_FUNCTION_CONTRACTS=mld_fqscale
-USE_FUNCTION_CONTRACTS=mld_montgomery_reduce
+USE_FUNCTION_CONTRACTS=mld_fqmul
 APPLY_LOOP_CONTRACTS=on
 USE_DYNAMIC_FRAMES=1
 

proofs/cbmc/fqscale/Makefile

@@ -4,11 +4,11 @@
 include ../Makefile_params.common
 
 HARNESS_ENTRY = harness
-HARNESS_FILE = ntt_butterfly_block_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 = mld_ntt_butterfly_block
+PROOF_UID = mld_invntt_2_layers
 
 DEFINES +=
 INCLUDES +=
@@ -19,16 +19,16 @@ UNWINDSET +=
 PROOF_SOURCES += $(PROOFDIR)/$(HARNESS_FILE).c
 PROJECT_SOURCES += $(SRCDIR)/mldsa/src/poly.c
 
-CHECK_FUNCTION_CONTRACTS=mld_ntt_butterfly_block
-USE_FUNCTION_CONTRACTS=mld_fqmul
+CHECK_FUNCTION_CONTRACTS=mld_invntt_2_layers
+USE_FUNCTION_CONTRACTS=mld_invntt_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 = mld_ntt_butterfly_block
+FUNCTION_NAME = mld_invntt_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

proofs/cbmc/ntt_butterfly_block/Makefile proofs/cbmc/invntt_2_layers/Makefileproofs/cbmc/ntt_butterfly_block/Makefile renamed to proofs/cbmc/invntt_2_layers/Makefile

@@ -4,11 +4,11 @@
 #include <stdint.h>
 #include "params.h"
 
-void mld_invntt_layer(int32_t r[MLDSA_N], unsigned layer);
+void mld_invntt_2_layers(int32_t r[MLDSA_N], unsigned layer);
 
 void harness(void)
 {
   int32_t *r;
   unsigned layer;
-  mld_invntt_layer(r, layer);
+  mld_invntt_2_layers(r, layer);
 }