First shot at ML-DSA-44.

Makefile

@@ -1,5 +1,5 @@
-PARAMETER_SET = 65
-IMPLEMENTATION_TYPE = ref
+PARAMETER_SET ?= 65
+IMPLEMENTATION_TYPE ?= ref
 
 # --------------------------------------------------------------------
 JASMINC ?= jasminc
@@ -22,11 +22,17 @@ $(OUTPUT_FILE_NAME).s: $(IMPLEMENTATION_SOURCES)
 # --------------------------------------------------------------------
 .PHONY: test
 test: $(OUTPUT_FILE_NAME).so
-	python3 -m pytest tests/
+	python3 -m pytest \
+		--parameter-set=$(PARAMETER_SET) \
+		--implementation-type=$(IMPLEMENTATION_TYPE) \
+		tests/
 
 .PHONY: nist-drbg-kat-test
 nist-drbg-kat-test: $(OUTPUT_FILE_NAME).so
-	python3 -m pytest tests/test_nist_drbg_kats.py
+	python3 -m pytest \
+		--parameter-set=$(PARAMETER_SET) \
+		--implementation-type=$(IMPLEMENTATION_TYPE) \
+		tests/test_nist_drbg_kats.py
 
 $(OUTPUT_FILE_NAME).so: $(OUTPUT_FILE_NAME).s
 	$(CC) $^ -fPIC -shared -o $@

ref/ml_dsa_44/api.h

@@ -0,0 +1,33 @@
+#ifndef ML_DSA_65_API_H
+#define ML_DSA_65_API_H
+
+#define VERIFICATION_KEY_SIZE 1312
+#define SIGNING_KEY_SIZE 2560
+#define SIGNATURE_SIZE 2420
+
+#include <stdint.h>
+#include <stddef.h>
+
+void ml_dsa_44_keygen(
+  uint8_t verification_key[VERIFICATION_KEY_SIZE],
+  uint8_t signing_key[SIGNING_KEY_SIZE],
+  const uint8_t randomness[32]
+);
+
+void ml_dsa_44_sign(
+  uint8_t signature[SIGNATURE_SIZE],
+  const uint8_t signing_key[SIGNING_KEY_SIZE],
+  const uint8_t *message,
+  const size_t message_size,
+  const uint8_t randomness[32]
+);
+
+int ml_dsa_44_verify(
+  uint8_t verification_key[VERIFICATION_KEY_SIZE],
+  const uint8_t *message,
+  const size_t message_size,
+  const uint8_t signature[SIGNATURE_SIZE]
+);
+
+#endif /* ML_DSA_65_API_H */
+

ref/ml_dsa_44/arithmetic/modular.jinc

@@ -0,0 +1,63 @@
+require "../constants.jinc"
+
+param int INVERSE_OF_MODULUS_MOD_MONTGOMERY_R = 58_728_449;
+
+inline
+fn montgomery_reduce(reg u64 a) -> reg u32
+{
+    reg u32 a32;
+    reg u64 v64;
+    reg u32 t;
+
+    // t = (int64_t)(int32_t)a*QINV;
+    a32 = (32s) a;
+    v64 = (64s) a32;
+    v64 *= INVERSE_OF_MODULUS_MOD_MONTGOMERY_R;
+    t = (32s) v64;
+
+    // t = (a - (int64_t)t*Q) >> 32;
+    v64 = (64s) t;
+    v64 *= -MODULUS;
+    v64 += a;
+    v64 >>= 32;
+    t = (32s) v64;
+
+    return t;
+}
+
+
+namespace coefficient {
+    inline
+    fn reduce32(reg u32 coefficient) -> reg u32
+    {
+        reg u32 t quotient;
+
+        // t = (coefficient + (1 << 22)) >> 23;
+        t = coefficient;
+        t += (1 << 22);
+        t = t >>s 23;
+
+        // t = coefficient - t*Q;
+        quotient = t;
+        quotient *= MODULUS;
+        t = coefficient;
+        t -= quotient;
+
+        return t;
+    }
+
+    inline
+    fn conditionally_add_modulus(reg u32 coefficient) -> reg u32 {
+        reg u32 add_by;
+        reg u32 result;
+
+        add_by = coefficient;
+        add_by >>s= 31;
+        add_by &= MODULUS;
+
+        result = coefficient;
+        result += add_by;
+
+        return result;
+    }
+}

ref/ml_dsa_44/arithmetic/ntt.jinc

@@ -0,0 +1,172 @@
+require "modular.jinc"
+require "../constants.jinc"
+
+u32[256] ZETAS_TIMES_MONTGOMERY_R = {
+    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, 0
+};
+
+namespace polynomial {
+    inline
+    fn ntt_at_layer(
+        inline int LAYER,
+        reg u64 zeta_i,
+        reg ptr u32[COEFFICIENTS_IN_POLYNOMIAL] polynomial
+    ) -> reg u64, reg ptr u32[COEFFICIENTS_IN_POLYNOMIAL]
+    {
+        inline int STEP = 1 << LAYER;
+        inline int ROUNDS = 128 >> LAYER;
+
+        reg u64 round offset bound;
+
+        reg u64 t64 zeta;
+        reg u32 t left;
+
+        reg ptr u32[COEFFICIENTS_IN_POLYNOMIAL] zetas = ZETAS_TIMES_MONTGOMERY_R;
+
+        round = 0;
+        while (round < ROUNDS) {
+            zeta_i += 1;
+
+            offset = round;
+            offset *= STEP;
+            offset *= 2;
+
+            bound = offset;
+            bound += STEP;
+
+            while (offset < bound) {
+                // TODO: Loading these values as a 64u will result in the
+                // wrong answer, maybe this should be reflected in the types?
+                t64 = (64s)polynomial[offset + STEP];
+                zeta = (64s)zetas[zeta_i];
+
+                t64 *= zeta;
+                t = montgomery_reduce(t64);
+
+                left = polynomial[offset];
+                left -= t;
+                polynomial[offset + STEP] = left;
+
+                left = polynomial[offset];
+                left += t;
+                polynomial[offset] = left;
+
+                offset += 1;
+            }
+
+            round += 1;
+        }
+
+        return zeta_i, polynomial;
+    }
+    fn ntt(reg ptr u32[COEFFICIENTS_IN_POLYNOMIAL] polynomial) -> reg ptr u32[COEFFICIENTS_IN_POLYNOMIAL]
+    {
+        reg u64 zeta_i = -1;
+
+        inline int LAYER;
+        for LAYER = 7 downto -1 {
+            zeta_i, polynomial = ntt_at_layer(LAYER, zeta_i, polynomial);
+        }
+
+        return polynomial;
+    }
+
+    inline
+    fn invert_ntt_at_layer(
+        inline int LAYER,
+        reg u64 zeta_i,
+        reg ptr u32[COEFFICIENTS_IN_POLYNOMIAL] polynomial
+    ) -> reg u64, reg ptr u32[COEFFICIENTS_IN_POLYNOMIAL]
+    {
+        inline int STEP = 1 << LAYER;
+        inline int ROUNDS = 128 >> LAYER;
+
+        reg u64 round offset bound;
+
+        reg u64 t64 zeta;
+        reg u32 t left right;
+
+        reg ptr u32[COEFFICIENTS_IN_POLYNOMIAL] zetas = ZETAS_TIMES_MONTGOMERY_R;
+
+        round = 0;
+        while (round < ROUNDS) {
+            zeta_i -= 1;
+
+            offset = round;
+            offset *= STEP;
+            offset *= 2;
+
+            bound = offset;
+            bound += STEP;
+
+            while (offset < bound) {
+                t = polynomial[offset + STEP];
+                t -= polynomial[offset];
+
+                left = polynomial[offset];
+                right = polynomial[offset + STEP];
+                left += right;
+                polynomial[offset] = left;
+
+                zeta = (64s)zetas[zeta_i];
+                t64 = (64s)t;
+
+                t64 *= zeta;
+                t = montgomery_reduce(t64);
+                polynomial[offset + STEP] = t;
+
+                offset += 1;
+            }
+
+            round += 1;
+        }
+
+        return zeta_i, polynomial;
+    }
+
+    fn invert_ntt_montgomery(
+        reg ptr u32[COEFFICIENTS_IN_POLYNOMIAL] polynomial
+    ) -> reg ptr u32[COEFFICIENTS_IN_POLYNOMIAL]
+    {
+        reg u64 zeta_i coefficient64;
+        reg u32 coefficient;
+
+        reg u64 j;
+
+        inline int LAYER;
+
+        zeta_i = 255;
+        for LAYER = 0 to 8 {
+            zeta_i, polynomial = invert_ntt_at_layer(LAYER, zeta_i, polynomial);
+        }
+
+        j = 0;
+        while(j < COEFFICIENTS_IN_POLYNOMIAL) {
+            coefficient64 = (64s)polynomial[j];
+            coefficient64 *= 41_978;
+
+            coefficient = montgomery_reduce(coefficient64);
+            polynomial[j] = coefficient;
+
+            j += 1;
+        }
+
+        return polynomial;
+    }
+}