Fixes and comments from Anas

vpaprotsk · vpaprotsk · commit 2ff3b82326a0 · 2025-10-06T21:23:40.000Z
diff --git a/src/hotspot/cpu/x86/stubGenerator_x86_64_dilithium.cpp b/src/hotspot/cpu/x86/stubGenerator_x86_64_dilithium.cpp
@@ -200,10 +200,10 @@ static auto whole_shuffle(Register scratch, KRegister mergeMask1, KRegister merg
             __ vmovdqu(output2[i], input2[i], vector_len);
           }
           for (int i = 0; i < regCnt; i++) {
-            __ evmovshdup(output2[i], k2, input1[i], true, vector_len);
+            __ evmovshdup(output2[i], mergeMask2, input1[i], true, vector_len);
           }
           for (int i = 0; i < regCnt; i++) {
-            __ evmovsldup(input1[i], k1, input2[i], true, vector_len);
+            __ evmovsldup(input1[i], mergeMask1, input2[i], true, vector_len);
           }
           break;
         // Special cases
@@ -390,7 +390,7 @@ static void storeXmms(Register destination, int offset, const XMMRegister xmmReg
 // static int implDilithiumAlmostNtt(int[] coeffs, int zetas[]) {}
 //
 // coeffs (int[256]) = c_rarg0
-// zetas (int[256]) = c_rarg1
+// zetas (int[128*8]) = c_rarg1
 //
 static address generate_dilithiumAlmostNtt_avx(StubGenerator *stubgen,
                             int vector_len, MacroAssembler *_masm) {
@@ -647,7 +647,7 @@ static address generate_dilithiumAlmostNtt_avx(StubGenerator *stubgen,
 // static int implDilithiumAlmostInverseNtt(int[] coeffs, int[] zetas) {}
 //
 // coeffs (int[256]) = c_rarg0
-// zetas (int[256]) = c_rarg1
+// zetas (int[128*8]) = c_rarg1
 static address generate_dilithiumAlmostInverseNtt_avx(StubGenerator *stubgen,
                                          int vector_len,MacroAssembler *_masm) {
   __ align(CodeEntryAlignment);
@@ -1017,12 +1017,13 @@ static address generate_dilithiumMontMulByConstant_avx(StubGenerator *stubgen,
     __ evpbroadcastd(constant, rConstant, Assembler::AVX_512bit); // constant multiplier
 
     __ mov64(scratch, 0b0101010101010101); //dw-mask
-    __ kmovwl(k2, scratch);
+    __ kmovwl(mergeMask, scratch);
   }
 
   // Total payload is 256*int32s.
   // - memStep is number of bytes one montMul64 processes.
   // - loopCnt is number of iterations it will take to process entire payload.
+  // - (two memSteps per loop)
   int memStep = 4 * 64;
   int loopCnt = 2;
   if (vector_len == Assembler::AVX_256bit) {
@@ -1321,15 +1322,15 @@ void StubGenerator::generate_dilithium_stubs() {
   }
   // Generate Dilithium intrinsics code
   if (UseDilithiumIntrinsics) {
-      StubRoutines::_dilithiumAlmostNtt =
+    StubRoutines::_dilithiumAlmostNtt =
         generate_dilithiumAlmostNtt_avx(this, vector_len, _masm);
-      StubRoutines::_dilithiumAlmostInverseNtt =
+    StubRoutines::_dilithiumAlmostInverseNtt =
         generate_dilithiumAlmostInverseNtt_avx(this, vector_len, _masm);
-      StubRoutines::_dilithiumNttMult =
+    StubRoutines::_dilithiumNttMult =
         generate_dilithiumNttMult_avx(this, vector_len, _masm);
-      StubRoutines::_dilithiumMontMulByConstant =
+    StubRoutines::_dilithiumMontMulByConstant =
         generate_dilithiumMontMulByConstant_avx(this, vector_len, _masm);
-      StubRoutines::_dilithiumDecomposePoly =
+    StubRoutines::_dilithiumDecomposePoly =
         generate_dilithiumDecomposePoly_avx(this, vector_len, _masm);
   }
 }
diff --git a/test/jdk/sun/security/provider/acvp/ML_DSA_Intrinsic_Test.java b/test/jdk/sun/security/provider/acvp/ML_DSA_Intrinsic_Test.java
@@ -31,6 +31,12 @@
 import java.lang.reflect.Constructor;
 import java.util.HexFormat;
 
+/*
+ * @test
+ * @library /test/lib
+ * @modules java.base/sun.security.provider:open
+ * @run main ML_DSA_Intrinsic_Test
+ */
 public class ML_DSA_Intrinsic_Test {
     public static void main(String[] args) throws Exception {
         MethodHandles.Lookup lookup = MethodHandles.lookup();
@@ -129,7 +135,7 @@ public static void testMult(int[] prod1, int[] prod2, int[] coeffs1, int[] coeff
         mult.invoke(prod1, coeffs1, coeffs2);
         multJava.invoke(prod2, coeffs1, coeffs2);
         
-        if (!Arrays.equals(prod1, parseHex(formatOf(prod2)))) {
+        if (!Arrays.equals(prod1, prod2)) {
                 throw new RuntimeException("[Seed "+seed+"@"+i+"] Result mult mismatch: " + formatOf(prod1) + " != " + formatOf(prod2));
         }
     }
@@ -141,7 +147,9 @@ public static void testMultConst(int[] prod1, int[] prod2,
         for (int j = 0; j<ML_DSA_N; j++) {
             prod1[j] = prod2[j] = rnd.nextInt();
         }
-        int c = rnd.nextInt();
+        // Per Algorithm 3 in https://eprint.iacr.org/2018/039.pdf, one of the inputs is bound, which prevents overflows
+        int dilithium_q = 8380417;
+        int c = rnd.nextInt(dilithium_q); 
 
         multConst.invoke(prod1, c);
         multConstJava.invoke(prod2, c);
@@ -189,9 +197,6 @@ public static void testAlmostNtt(int[] coeffs1, int[] coeffs2,
         almostNttJava.invoke(coeffs2);
 
         if (!Arrays.equals(coeffs1, coeffs2)) {
-            if (false) {
-                implDilithiumAlmostNttJava(coeffs2);
-            }
             throw new RuntimeException("[Seed "+seed+"@"+i+"] Result AlmostNtt mismatch: " + formatOf(coeffs1) + " != " + formatOf(coeffs2));
         }
     }
@@ -209,9 +214,6 @@ public static void testInverseNtt(int[] coeffs1, int[] coeffs2,
         inverseNttJava.invoke(coeffs2);
 
         if (!Arrays.equals(coeffs1, coeffs2)) {
-            if (true) {
-                implDilithiumAlmostInverseNttJava(coeffs3);
-            }
             throw new RuntimeException("[Seed "+seed+"@"+i+"] Result InverseNtt mismatch: " + formatOf(coeffs1) + " != " + formatOf(coeffs2));
         }
     }
@@ -225,18 +227,8 @@ private static CharSequence formatOf(int[] arr) {
         return b.toString();
     }
 
-    private static int[] parseHex(CharSequence string) {
-        assert(string.length()%8==0);
-
-        int[] r = new int[string.length()/8];
-        HexFormat hex = HexFormat.of();
-        for (int i = 0, j = 0; j<string.length(); i++, j+=8) {
-            r[i] = hex.fromHexDigits(string, j, j+8);
-        }
-        return r;
-    }
-
-private static final int[] MONT_ZETAS_FOR_VECTOR_INVERSE_NTT = new int[]{
+    // Copied constants from sun.security.provider.ML_DSA
+    private static final int[] MONT_ZETAS_FOR_VECTOR_INVERSE_NTT = new int[]{
             -1976782, 846154, -1400424, -3937738, 1362209, 48306, -3919660, 554416,
             3545687, -1612842, 976891, -183443, 2286327, 420899, 2235985, 2939036,
             3833893, 260646, 1104333, 1667432, -1910376, 1803090, -1723600, 426683,
@@ -511,107 +503,5 @@ private static int[] parseHex(CharSequence string) {
             -2939036, -2235985, -420899, -2286327, 183443, -976891, 1612842, -3545687,
             -554416, 3919660, -48306, -1362209, 3937738, 1400424, -846154, 1976782
     };
-
-
-    static void implDilithiumAlmostNttJava(int[] coeffs) {
-        HexFormat hex = HexFormat.of();
-        int dimension = 256;
-        int m = 0;
-        int testLevel = 2;
-        for (int l = dimension / 2; l >= testLevel; l /= 2) {
-            for (int s = 0; s < dimension; s += 2 * l) {
-                for (int j = s; j < s + l; j++) {
-                    StringBuilder bld = new StringBuilder();
-                    bld.append("l = " + l + ", m = " + m + ", j = " + j+": ");
-                    bld.append(hex.toHexDigits(coeffs[j]) + ", " + hex.toHexDigits(coeffs[j + l]) + " * " + hex.toHexDigits(MONT_ZETAS_FOR_NTT[m]));
-                    int tmp = montMul(MONT_ZETAS_FOR_NTT[m], coeffs[j + l]);
-                    coeffs[j + l] = coeffs[j] - tmp;
-                    coeffs[j] = coeffs[j] + tmp;
-                    bld.append(" -> " + hex.toHexDigits(coeffs[j]) + ", " + hex.toHexDigits(coeffs[j + l]) + " (tmp = " + hex.toHexDigits(tmp) + ")");
-                    if (l == testLevel) {
-                        System.out.println(bld.toString());
-                    }
-                }
-                m++;
-            }
-        }
-    }
-    static void implDilithiumAlmostInverseNttJava(int[] coeffs) {
-        HexFormat hex = HexFormat.of();
-        int dimension = 256;
-        int m = MONT_ZETAS_FOR_NTT.length - 1;
-        int testLevel = 1;
-        for (int l = 1; l < dimension; l *= 2) {
-            for (int s = 0; s < dimension; s += 2 * l) {
-                for (int j = s; j < s + l; j++) {
-                    StringBuilder bld = new StringBuilder();
-                    bld.append("l = " + l + ", m = " + m + ", j = " + j+": ");
-                    int tmp = coeffs[j];
-                    bld.append(hex.toHexDigits(coeffs[j]) + ", " + hex.toHexDigits(coeffs[j + l]) + " -> "  + hex.toHexDigits(tmp - coeffs[j + l]) + " * " + hex.toHexDigits(-MONT_ZETAS_FOR_NTT[m]));
-                    coeffs[j] = (tmp + coeffs[j + l]);
-                    coeffs[j + l] = montMul(tmp - coeffs[j + l], -MONT_ZETAS_FOR_NTT[m]);
-                    bld.append(" -> " + hex.toHexDigits(coeffs[j]) + ", " + hex.toHexDigits(coeffs[j + l]));
-                    if (l == testLevel) {
-                        System.out.println(bld.toString());
-                    }
-                }
-                m--;
-            }
-        }
-    }
-    private static int montMul(int b, int c) {
-        long a = (long) b * (long) c;
-        int aHigh = (int) (a >> 32);
-        int aLow = (int) a;
-        int m = 58728449 * aLow; // signed low product
-
-        // subtract signed high product
-        return (aHigh - (int) (((long)m * 8380417) >> 32));
-    }
-    // Zeta values for NTT with montgomery factor precomputed
-    private static final int[] MONT_ZETAS_FOR_NTT = new int[]{
-        25847, //0
-        -2608894, -518909, //1 
-        237124, -777960, -876248, 466468, //2
-        1826347,
-        2353451, -359251, -2091905, 3119733, -2884855, 3111497, 2680103, //3
-        2725464,
-        1024112, -1079900, 3585928, -549488, -1119584, 2619752, -2108549, -2118186,
-        -3859737, -1399561, -3277672, 1757237, -19422, 4010497, 280005,  // 4
-        
-        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, //5
-        
-        -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, //6
-        
-        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
-    };
 }
 // java --add-opens java.base/sun.security.provider=ALL-UNNAMED  -XX:+UseDilithiumIntrinsics test/jdk/sun/security/provider/acvp/ML_DSA_Intrinsic_Test.java
diff --git a/test/micro/org/openjdk/bench/javax/crypto/full/MLDSABench.java b/test/micro/org/openjdk/bench/javax/crypto/full/MLDSABench.java
@@ -48,9 +48,7 @@
 @OutputTimeUnit(TimeUnit.MILLISECONDS)
 @Fork(value = 1, jvmArgs = {"--add-opens", "java.base/sun.security.provider=ALL-UNNAMED"})
 public class MLDSABench extends CryptoBase {
-
     public static final int SET_SIZE = 128;
-    private static final int ML_DSA_N = 256;
 
     private int[][] coeffs1;
     private int[][] coeffs2;
@@ -109,7 +107,7 @@ public void setup() throws Exception {
     }
 
     @Benchmark
-    public void mult1() throws Exception, Throwable {
+    public void mult() throws Exception, Throwable {
         mult.invoke(prod1[index], coeffs1[index], coeffs2[index]);
         index = (index + 1) % SET_SIZE;
     }
@@ -143,6 +141,8 @@ public void multInverseNtt() throws Exception, Throwable {
         index = (index + 1) % SET_SIZE;
     }
 
+    // Copied constants from sun.security.provider.ML_DSA
+    private static final int ML_DSA_N = 256;
     private static final int[] MONT_ZETAS_FOR_VECTOR_INVERSE_NTT = new int[]{
             -1976782, 846154, -1400424, -3937738, 1362209, 48306, -3919660, 554416,
             3545687, -1612842, 976891, -183443, 2286327, 420899, 2235985, 2939036,
