Skip to content

8383635: Improve performance of ML-KEM's implKyberNttMultJava() method #760

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
stooke wants to merge 1 commit into
base: master
Choose a base branch
from
+84 −53
Open

8383635: Improve performance of ML-KEM's implKyberNttMultJava() method #760

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 18 additions & 0 deletions src/hotspot/cpu/aarch64/stubGenerator_aarch64.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5721,6 +5721,24 @@ class StubGenerator: public StubCodeGenerator {
// static int implKyberNttMult(
// short[] result, short[] ntta, short[] nttb, short[] zetas) {}
//
// The actual algorithm that is used here differs from the one in the Java
// implementation, it uses Montgomery multiplications instead of Barrett
// reduction, but the end result modulo MLKEM_Q is the same. This is the
// Java equivalent of this intrinsic implementation:
// static void implKyberNttMultJava(short[] result, short[] ntta, short[] nttb) {
// for (int m = 0; m < ML_KEM_N / 2; m++) {
// int a0 = ntta[2 * m];
// int a1 = ntta[2 * m + 1];
// int b0 = nttb[2 * m];
// int b1 = nttb[2 * m + 1];
// int r = montMul(a0, b0) +
// montMul(montMul(a1, b1), MONT_ZETAS_FOR_NTT_MULT[m]);
// result[2 * m] = (short) montMul(r, MONT_R_SQUARE_MOD_Q);
// result[2 * m + 1] = (short) montMul(
// (montMul(a0, b1) + montMul(a1, b0)), MONT_R_SQUARE_MOD_Q);
// }
// }
//
// result (short[256]) = c_rarg0
// ntta (short[256]) = c_rarg1
// nttb (short[256]) = c_rarg2
Expand Down
21 changes: 21 additions & 0 deletions src/hotspot/cpu/x86/stubGenerator_x86_64_kyber.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -614,6 +614,27 @@ address generate_kyberInverseNtt_avx512(StubGenerator *stubgen,
}

// Kyber multiply polynomials in the NTT domain.
// Implements
// static int implKyberNttMult(
// short[] result, short[] ntta, short[] nttb, short[] zetas) {}
//
// The actual algorithm that is used here differs from the one in the Java
// implementation, it uses Montgomery multiplications instead of Barrett
// reduction, but the end result modulo MLKEM_Q is the same. This is the
// Java equivalent of this intrinsic implementation:
// static void implKyberNttMultJava(short[] result, short[] ntta, short[] nttb) {
// for (int m = 0; m < ML_KEM_N / 2; m++) {
// int a0 = ntta[2 * m];
// int a1 = ntta[2 * m + 1];
// int b0 = nttb[2 * m];
// int b1 = nttb[2 * m + 1];
// int r = montMul(a0, b0) +
// montMul(montMul(a1, b1), MONT_ZETAS_FOR_NTT_MULT[m]);
// result[2 * m] = (short) montMul(r, MONT_R_SQUARE_MOD_Q);
// result[2 * m + 1] = (short) montMul(
// (montMul(a0, b1) + montMul(a1, b0)), MONT_R_SQUARE_MOD_Q);
// }
// }
//
// result (short[256]) = c_rarg0
// ntta (short[256]) = c_rarg1
Expand Down
98 changes: 45 additions & 53 deletions src/java.base/share/classes/com/sun/crypto/provider/ML_KEM.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,14 +46,17 @@ public final class ML_KEM {
private static final int XOF_PAD = 24;
private static final int MONT_R_BITS = 20;
private static final int MONT_Q = 3329;
private static final int MONT_R_SQUARE_MOD_Q = 152;
private static final int MONT_Q_INV_MOD_R = 586497;

// toMont((ML_KEM_N / 2)^-1 mod ML_KEM_Q) using R = 2^MONT_R_BITS
private static final int MONT_DIM_HALF_INVERSE = 1534;
private static final int BARRETT_MULTIPLIER = 20159;
private static final int BARRETT_ADDEND = 1665;
private static final int BARRETT_SHIFT = 26;
private static final int[] MONT_ZETAS_FOR_NTT = new int[]{

// The values from Appendix A of the FIPS 203 standard converted to the
// Montgomery domain, i.e. toMont(zeta^ (bitrev_7(i)) for i = 0..127
private static final int[] MONT_ZETAS_FOR_NTT = new int[] {
1188, 914, -969, 585, -551, 1263, -97, 593,
-35, -1400, -417, -1253, 742, -281, 185, -819,
-1226, 895, -530, 52, 25, 1000, 1249, -909,
Expand All @@ -72,7 +75,7 @@ public final class ML_KEM {
-1599, -709, -789, -1317, -57, 1049, -584
};

private static final short[] montZetasForVectorNttArr = new short[]{
private static final short[] montZetasForVectorNttArr = new short[] {
// level 0
-758, -758, -758, -758, -758, -758, -758, -758,
-758, -758, -758, -758, -758, -758, -758, -758,
Expand Down Expand Up @@ -193,26 +196,8 @@ public final class ML_KEM {
-108, -108, -308, -308, 996, 996, 991, 991,
958, 958, -1460, -1460, 1522, 1522, 1628, 1628
};
private static final int[] MONT_ZETAS_FOR_INVERSE_NTT = new int[]{
584, -1049, 57, 1317, 789, 709, 1599, -1601,
-990, 604, 348, 857, 612, 474, 1177, -1014,
-88, -982, -191, 668, 1386, 486, -1153, -534,
514, 137, 586, -1178, 227, 339, -907, 244,
1200, -833, 1394, -30, 1074, 636, -317, -1192,
-1259, -355, -425, -884, -977, 1430, 868, 607,
184, 1448, 702, 1327, 431, 497, 595, -94,
1649, -1497, -620, 42, -172, 1107, -222, 1003,
426, -845, 395, -510, 1613, 825, 1269, -290,
-1429, 623, -567, 1617, 36, 1007, 1440, 332,
-201, 1313, -1382, -744, 669, -1538, 128, -1598,
1401, 1183, -553, 714, 405, -1155, -445, 406,
-1496, -49, 82, 1369, 259, 1604, 373, 909,
-1249, -1000, -25, -52, 530, -895, 1226, 819,
-185, 281, -742, 1253, 417, 1400, 35, -593,
97, -1263, 551, -585, 969, -914, -1188
};

private static final short[] montZetasForVectorInverseNttArr = new short[]{
private static final short[] montZetasForVectorInverseNttArr = new short[] {
// level 0
-1628, -1628, -1522, -1522, 1460, 1460, -958, -958,
-991, -991, -996, -996, 308, 308, 108, 108,
Expand Down Expand Up @@ -334,25 +319,28 @@ public final class ML_KEM {
758, 758, 758, 758, 758, 758, 758, 758
};

private static final int[] MONT_ZETAS_FOR_NTT_MULT = new int[]{
-1003, 1003, 222, -222, -1107, 1107, 172, -172,
-42, 42, 620, -620, 1497, -1497, -1649, 1649,
94, -94, -595, 595, -497, 497, -431, 431,
-1327, 1327, -702, 702, -1448, 1448, -184, 184,
-607, 607, -868, 868, -1430, 1430, 977, -977,
884, -884, 425, -425, 355, -355, 1259, -1259,
1192, -1192, 317, -317, -636, 636, -1074, 1074,
30, -30, -1394, 1394, 833, -833, -1200, 1200,
-244, 244, 907, -907, -339, 339, -227, 227,
1178, -1178, -586, 586, -137, 137, -514, 514,
534, -534, 1153, -1153, -486, 486, -1386, 1386,
-668, 668, 191, -191, 982, -982, 88, -88,
1014, -1014, -1177, 1177, -474, 474, -612, 612,
-857, 857, -348, 348, -604, 604, 990, -990,
1601, -1601, -1599, 1599, -709, 709, -789, 789,
-1317, 1317, -57, 57, 1049, -1049, -584, 584
// modulo MLKEM_Q positive equivalents of the values listed for
// the MultiplyNTTs algorithm in the FIPS 203 standard
private static final int[] ZETAS_FOR_NTT_MULT = new int[] {
17, 3312, 2761, 568, 583, 2746, 2649, 680,
1637, 1692, 723, 2606, 2288, 1041, 1100, 2229,
1409, 1920, 2662, 667, 3281, 48, 233, 3096,
756, 2573, 2156, 1173, 3015, 314, 3050, 279,
1703, 1626, 1651, 1678, 2789, 540, 1789, 1540,
1847, 1482, 952, 2377, 1461, 1868, 2687, 642,
939, 2390, 2308, 1021, 2437, 892, 2388, 941,
733, 2596, 2337, 992, 268, 3061, 641, 2688,
1584, 1745, 2298, 1031, 2037, 1292, 3220, 109,
375, 2954, 2549, 780, 2090, 1239, 1645, 1684,
1063, 2266, 319, 3010, 2773, 556, 757, 2572,
2099, 1230, 561, 2768, 2466, 863, 2594, 735,
2804, 525, 1092, 2237, 403, 2926, 1026, 2303,
1143, 2186, 2150, 1179, 2775, 554, 886, 2443,
1722, 1607, 1212, 2117, 1874, 1455, 1029, 2300,
2110, 1219, 2935, 394, 885, 2444, 2154, 1175
};
private static final short[] montZetasForVectorNttMultArr = new short[]{

private static final short[] montZetasForVectorNttMultArr = new short[] {
-1103, 1103, 430, -430, 555, -555, 843, -843,
-1251, 1251, 871, -871, 1550, -1550, 105, -105,
422, -422, 587, -587, 177, -177, -235, 235,
Expand Down Expand Up @@ -1143,17 +1131,20 @@ static int implKyberNttMult(short[] result, short[] ntta, short[] nttb,
}

static void implKyberNttMultJava(short[] result, short[] ntta, short[] nttb) {
for (int m = 0; m < ML_KEM_N / 2; m++) {

int a0 = ntta[2 * m];
int a1 = ntta[2 * m + 1];
int b0 = nttb[2 * m];
int b1 = nttb[2 * m + 1];
int r = montMul(a0, b0) +
montMul(montMul(a1, b1), MONT_ZETAS_FOR_NTT_MULT[m]);
result[2 * m] = (short) montMul(r, MONT_R_SQUARE_MOD_Q);
result[2 * m + 1] = (short) montMul(
(montMul(a0, b1) + montMul(a1, b0)), MONT_R_SQUARE_MOD_Q);
for (int m = 0; m < ML_KEM_N; m += 2) {
int a0 = ntta[m];
int a1 = ntta[m + 1];
int b0 = nttb[m];
int b1 = nttb[m + 1];
long r = a1 * b1;
r -= ((r * BARRETT_MULTIPLIER) >> BARRETT_SHIFT) * ML_KEM_Q;
r *= ZETAS_FOR_NTT_MULT[m >> 1];
r += a0 * b0;
result[m] = (short) (r - (((r + BARRETT_ADDEND) *
BARRETT_MULTIPLIER) >> BARRETT_SHIFT) * ML_KEM_Q);
long r1 = a0 * b1 + a1 * b0;
result[m + 1] = (short) (r1 - (((r1 + BARRETT_ADDEND) *
BARRETT_MULTIPLIER) >> BARRETT_SHIFT) * ML_KEM_Q);
}
}

Expand Down Expand Up @@ -1521,9 +1512,10 @@ static int implKyberBarrettReduce(short[] coeffs) {
}

static void implKyberBarrettReduceJava(short[] poly) {
int tmp = 0;
for (int m = 0; m < ML_KEM_N; m++) {
int tmp = ((int) poly[m] * BARRETT_MULTIPLIER) >> BARRETT_SHIFT;
poly[m] = (short) (poly[m] - tmp * ML_KEM_Q);
tmp = poly[m];
poly[m] = (short) (tmp - ((tmp * BARRETT_MULTIPLIER) >> BARRETT_SHIFT) * ML_KEM_Q);
}
}

Expand Down