On Thu, 22 Oct 2026 02:24:10 GMT, Volodymyr Paprotski <[email protected]> wrote:
>> This provides AVX2 intrinsics for MLKEM operations. I also refactored AVX512 >> for code sharing. (A lot of this is similar to the MLDSA refactoring I did, >> and general shape follows that design) >> >> Performance: >> - KEMBench.MLKEM.(en|de)capsulate on AVX2 got 50-80% faster (depending on >> keysize and operation) >> - On AVX512, the intrinsics themselves were ~10-30% faster, but it doesnt >> really show on the top-level KEMBench.MLKEM.(en|de)capsulate (other parts of >> the benchmark take over, eg SHA3 which we optimized previously) >> >> --------- >> - [X] I confirm that I make this contribution in accordance with the >> [OpenJDK Interim AI Policy](https://openjdk.org/legal/ai). > > Volodymyr Paprotski has updated the pull request with a new target base due > to a merge or a rebase. The pull request now contains eight commits: > > - re-enable test > - Merge remote-tracking branch 'origin/master' into avx2-mlkem > - whitespace > - reduce fuzz test repetition > - Merge remote-tracking branch 'origin/master' into avx2-mlkem > - prepare for rebase > - comments from Anas > - avx2 mlkem drop src/hotspot/cpu/x86/stubGenerator_x86_64_kyber.cpp line 243: > 241: const XMMRegister* scratch1 = scratch == input1 ? output : scratch; > 242: const XMMRegister* scratch2 = scratch == input1 ? scratch : output; > 243: bool input2ScratchOverlap = input2[0] == scratch1[1]; Is this input2ScratchOverlap check sufficient? src/hotspot/cpu/x86/stubGenerator_x86_64_kyber.cpp line 280: > 278: // NOTE: size 0 and 1 are used for final and initial shuffles > respectively of > 279: // kyberNtt and kyberInverseNtt. For size 0 and 1, input1[] > 280: // and input2[] are modified in-place (and output2 is used as a > temporary) For the AVX_256bit 16/32 path also input2 is modified in place. src/hotspot/cpu/x86/stubGenerator_x86_64_kyber.cpp line 1201: > 1199: // condensed (byte[168]) = c_rarg0 > 1200: // condensedOffs (int) = c_rarg1 > 1201: // parsed (short[112]) = c_rarg2 Should this be byte[192] and short[128]? The comment above MLKEM.twelve2Sixteen says so. src/hotspot/cpu/x86/stubGenerator_x86_64_kyber.cpp line 1483: > 1481: void StubGenerator::generate_kyber_stubs() { > 1482: int vector_len = Assembler::AVX_256bit; > 1483: if (VM_Version::supports_evex() && VM_Version::supports_avx512bw()) { Should this be VM_Version::supports_avx512vlbw()? The 12To16 stub at 1361–1366 uses AVX_256bit vector length. test/jdk/sun/security/provider/pqc/ML_KEM_Intrinsic_Test.java line 38: > 36: * @key randomness > 37: * @modules java.base/com.sun.crypto.provider:+open > 38: * @run main/othervm -XX:UseAVX=2 ML_KEM_Intrinsic_Test I think we should add AVX2 to the requires directive, as we are using UseAVX=2 on run command line: requires os.simpleArch == "x64" & vm.cpu.features ~= ".*avx2.*" ------------- PR Review Comment: https://git.openjdk.org/jdk/pull/33444#discussion_r4519236194 PR Review Comment: https://git.openjdk.org/jdk/pull/33444#discussion_r4519473362 PR Review Comment: https://git.openjdk.org/jdk/pull/33444#discussion_r4519392024 PR Review Comment: https://git.openjdk.org/jdk/pull/33444#discussion_r4518656297 PR Review Comment: https://git.openjdk.org/jdk/pull/33444#discussion_r4519607215
