8347608: Optimize Java implementation of ML-KEM

Reviewed-by: weijun
This commit is contained in:
Ben Perez 2025-05-14 19:38:34 +00:00
parent e91088a9e8
commit ecabea6cd3

View File

@ -703,7 +703,7 @@ public final class ML_KEM {
private K_PKE_CipherText kPkeEncrypt(
K_PKE_EncryptionKey publicKey, byte[] message, byte[] sigma) {
short[][] zeroes = new short[mlKem_k][ML_KEM_N];
short[][] zeroes = shortMatrixAlloc(mlKem_k, ML_KEM_N);
byte[] pkBytes = publicKey.keyBytes;
byte[] rho = Arrays.copyOfRange(pkBytes,
pkBytes.length - 32, pkBytes.length);
@ -792,7 +792,7 @@ public final class ML_KEM {
System.arraycopy(rho, 0, seedBuf, 0, rho.length);
seedBuf[rhoLen + 2] = 0x1F;
seedBuf[XOF_BLOCK_LEN - 1] = (byte)0x80;
byte[][] xofBufArr = new byte[nrPar][XOF_BLOCK_LEN + XOF_PAD];
byte[][] xofBufArr = byteMatrixAlloc(nrPar, XOF_BLOCK_LEN + XOF_PAD);
int[] iIndex = new int[nrPar];
int[] jIndex = new int[nrPar];
@ -1550,4 +1550,22 @@ public final class ML_KEM {
return (aHigh - ((m * MONT_Q) >> MONT_R_BITS)); // subtract signed high product
}
// For multidimensional array initialization, manually allocating each entry is
// faster than doing the entire initialization in one go
static short[][] shortMatrixAlloc(int first, int second) {
short[][] res = new short[first][];
for (int i = 0; i < first; i++) {
res[i] = new short[second];
}
return res;
}
static byte[][] byteMatrixAlloc(int first, int second) {
byte[][] res = new byte[first][];
for (int i = 0; i < first; i++) {
res[i] = new byte[second];
}
return res;
}
}