Commit 809b8180c1 for openssl.org
commit 809b8180c1b0a041fda20f9bd0f18f3a29802d4f
Author: Paul Dale <paul.dale@oracle.com>
Date: Wed Sep 30 11:22:59 2026 +1000
ml_kem: add AVX2 base multiplication and accumulation
Use AVX2 NTT-domain multiplication and accumulation for ML-KEM inner
products, matrix products, and the transposed matrix product used in
key generation.
Rename the x86_64 perlasm source to reflect its wider scope.
Assisted-by: ChatGPT:gpt-6-Sol
Reviewed-by: Neil Horman <nhorman@openssl.org>
Reviewed-by: Mounir Idrassi <mounir.idrassi@idrix.fr>
Merge-date: Wed Oct 21 16:38:25 2026
Merged-from: https://github.com/openssl/openssl/pull/33023
diff --git a/crypto/ml_kem/asm/mlkem_ntt-x86_64.pl b/crypto/ml_kem/asm/ml_kem_avx2.pl
similarity index 66%
rename from crypto/ml_kem/asm/mlkem_ntt-x86_64.pl
rename to crypto/ml_kem/asm/ml_kem_avx2.pl
index 9c649bb48f..60d126df27 100644
--- a/crypto/ml_kem/asm/mlkem_ntt-x86_64.pl
+++ b/crypto/ml_kem/asm/ml_kem_avx2.pl
@@ -68,6 +68,8 @@ my @roots = map { powmod(17, bitreverse7($_), 3329) } 0..127;
my @shoup = map { int($_ * 65536 / 3329) } @roots;
my @invroots = (0, map { 3329 - $roots[128 - $_] } 1..127);
my @invshoup = map { int($_ * 65536 / 3329) } @invroots;
+my @modroots = map { powmod(17, 2 * bitreverse7($_) + 1, 3329) } 0..127;
+my @modshoup = map { int($_ * 65536 / 3329) } @modroots;
my $code = "";
# Windows preserves XMM6..XMM15, so clear only the volatile vector registers.
# Unix has no nonvolatile vector registers and can clear all of them at once.
@@ -80,15 +82,15 @@ if ($avx2) {
$code .= <<___;
.text
.extern OPENSSL_ia32cap_P
-.globl mlkem_ntt_avx2_capable
-.type mlkem_ntt_avx2_capable,\@abi-omnipotent
+.globl mlkem_avx2_capable
+.type mlkem_avx2_capable,\@abi-omnipotent
.align 16
-mlkem_ntt_avx2_capable:
+mlkem_avx2_capable:
mov OPENSSL_ia32cap_P+8(%rip), %eax
and \$32, %eax
ret
-.Lmlkem_ntt_avx2_capable_end:
-.size mlkem_ntt_avx2_capable, .-mlkem_ntt_avx2_capable
+.Lmlkem_avx2_capable_end:
+.size mlkem_avx2_capable, .-mlkem_avx2_capable
.globl mlkem_ntt_avx2
.type mlkem_ntt_avx2,\@abi-omnipotent
@@ -343,6 +345,122 @@ ___
ret
.Lmlkem_inverse_ntt_avx2_end:
.size mlkem_inverse_ntt_avx2, .-mlkem_inverse_ntt_avx2
+___
+
+ # Each AVX2 vector holds eight quadratic NTT elements. Shoup-scale the
+ # odd left coefficients by their known roots, then VPMADDWD computes the
+ # two 32-bit coefficients. Four products still fit a signed 32-bit lane.
+ # Montgomery reduction followed by multiplication by R mod q gives the
+ # canonical, unscaled result required at the C function boundary.
+ # The transpose-add form strides over the matrix and adds the existing
+ # error polynomial in the output, without copying a matrix column.
+ for my $variant (['mlkem_basemul_acc_avx2', 0],
+ ['mlkem_basemul_acc_transpose_add_avx2', 1]) {
+ my ($name, $transpose_add) = @$variant;
+ $code .= ".globl $name\n.type $name,\@abi-omnipotent\n.align 32\n$name:\n";
+ if ($win64) {
+ $code .= " mov %rcx, %r10\n mov %rdx, %r11\n";
+ } else {
+ $code .= " mov %rdi, %r10\n mov %rsi, %r11\n"
+ . " mov %rdx, %r8\n";
+ $code .= " mov %ecx, %r9d\n";
+ }
+ $code .= " mov %r9d, %r9d\n shl \$9, %r9\n"
+ if $transpose_add;
+ my $rank_load = $transpose_add
+ ? " mov %r9, %rcx\n shr \$9, %rcx"
+ : " mov %r9d, %ecx";
+ my $lhs_advance = $transpose_add
+ ? " add %r9, %r11" : " add \$512, %r11";
+ my $rewind = $transpose_add
+ ? " mov %r9, %rcx\n shr \$9, %rcx\n"
+ . " imulq %r9, %rcx\n sub %rcx, %r11\n"
+ . " sub %r9, %r8"
+ : " mov %r9d, %ecx\n shl \$9, %rcx\n"
+ . " sub %rcx, %r11\n sub %rcx, %r8";
+ my $body = <<'___';
+ lea .Lmlkem_modroots(%rip), %rdx
+ xor %eax, %eax
+.L@NAME@_block:
+ vpxor %ymm0, %ymm0, %ymm0
+ vpxor %ymm1, %ymm1, %ymm1
+@RANK_LOAD@
+.L@NAME@_rank:
+ vmovdqu (%r11,%rax), %ymm2
+ vmovdqu (%r8,%rax), %ymm3
+ vmovdqu (%rdx,%rax,2), %ymm4
+ vpmullw %ymm2, %ymm4, %ymm4
+ vmovdqu 32(%rdx,%rax,2), %ymm5
+ vpmulhuw %ymm2, %ymm5, %ymm5
+ vpmullw .Lmlkem_vec_q(%rip), %ymm5, %ymm5
+ vpsubw %ymm5, %ymm4, %ymm4
+ vpmaddwd %ymm3, %ymm4, %ymm5
+ vpaddd %ymm5, %ymm0, %ymm0
+ vpshufb .Lmlkem_swap_pairs(%rip), %ymm3, %ymm5
+ vpmaddwd %ymm5, %ymm2, %ymm5
+ vpaddd %ymm5, %ymm1, %ymm1
+@LHS_ADVANCE@
+ add $512, %r8
+ dec %ecx
+ jnz .L@NAME@_rank
+@REWIND@
+
+ vpmulld .Lmlkem_vec_qinv(%rip), %ymm0, %ymm4
+ vpand .Lmlkem_vec_mask16(%rip), %ymm4, %ymm4
+ vpmulld .Lmlkem_vec_q32(%rip), %ymm4, %ymm4
+ vpaddd %ymm4, %ymm0, %ymm0
+ vpsrld $16, %ymm0, %ymm0
+ vpsubd .Lmlkem_vec_q32(%rip), %ymm0, %ymm4
+ vpminud %ymm4, %ymm0, %ymm0
+
+ vpmulld .Lmlkem_vec_qinv(%rip), %ymm1, %ymm4
+ vpand .Lmlkem_vec_mask16(%rip), %ymm4, %ymm4
+ vpmulld .Lmlkem_vec_q32(%rip), %ymm4, %ymm4
+ vpaddd %ymm4, %ymm1, %ymm1
+ vpsrld $16, %ymm1, %ymm1
+ vpsubd .Lmlkem_vec_q32(%rip), %ymm1, %ymm4
+ vpminud %ymm4, %ymm1, %ymm1
+
+ vpackusdw %ymm1, %ymm0, %ymm2
+ vpshufb .Lmlkem_interleave(%rip), %ymm2, %ymm2
+ vpmulhuw .Lmlkem_vec_r_shoup(%rip), %ymm2, %ymm3
+ vpmullw .Lmlkem_vec_r(%rip), %ymm2, %ymm2
+ vpmullw .Lmlkem_vec_q(%rip), %ymm3, %ymm3
+ vpsubw %ymm3, %ymm2, %ymm2
+ vpsubw .Lmlkem_vec_q(%rip), %ymm2, %ymm3
+ vpminuw %ymm3, %ymm2, %ymm2
+___
+ $body =~ s/\@NAME\@/$name/g;
+ $body =~ s/\@RANK_LOAD\@/$rank_load/;
+ $body =~ s/\@LHS_ADVANCE\@/$lhs_advance/;
+ $body =~ s/\@REWIND\@/$rewind/;
+ $code .= $body;
+ if ($transpose_add) {
+ $code .= <<'___';
+ vmovdqu (%r10,%rax), %ymm3
+ vpaddw %ymm3, %ymm2, %ymm2
+ vpsubw .Lmlkem_vec_q(%rip), %ymm2, %ymm3
+ vpminuw %ymm3, %ymm2, %ymm2
+___
+ }
+ $code .= <<'___';
+ vmovdqu %ymm2, (%r10,%rax)
+ add $32, %rax
+ cmp $512, %rax
+___
+ $code .= " jb .L" . $name . "_block\n";
+ $code .= <<'___';
+ xor %r8, %r8
+ xor %r10, %r10
+ xor %r11, %r11
+ xor %rdx, %rdx
+___
+ $code .= " xor %rdi, %rdi\n xor %rsi, %rsi\n" unless $win64;
+ $code .= $clear_vectors;
+ $code .= " ret\n.L" . $name . "_end:\n.size $name, .-$name\n";
+ }
+
+ $code .= <<'___';
.section .rodata
.align 32
@@ -411,16 +529,51 @@ ___
}
}
}
+ my $packed16 = sub {
+ my ($value) = @_;
+ return $value + ($value << 16);
+ };
+ if (!$win64) {
+ $code .= ".align 32\n.Lmlkem_vec_q:\n .long "
+ . join(", ", ($packed16->(3329)) x 8) . "\n";
+ }
+ $code .= ".align 32\n.Lmlkem_vec_r:\n .long "
+ . join(", ", ($packed16->(65536 % 3329)) x 8) . "\n";
+ $code .= ".Lmlkem_vec_r_shoup:\n .long "
+ . join(", ", ($packed16->(int((65536 % 3329) * 65536 / 3329))) x 8) . "\n";
+ for my $constant (['q32', 3329], ['qinv', 3327],
+ ['mask16', 65535]) {
+ $code .= ".Lmlkem_vec_" . $constant->[0] . ":\n .long "
+ . join(", ", ($constant->[1]) x 8) . "\n";
+ }
+ my @swap = (2,3,0,1,6,7,4,5,10,11,8,9,14,15,12,13);
+ my @interleave = (0,1,8,9,2,3,10,11,4,5,12,13,6,7,14,15);
+ $code .= ".Lmlkem_swap_pairs:\n .byte "
+ . join(", ", (@swap, @swap)) . "\n";
+ $code .= ".Lmlkem_interleave:\n .byte "
+ . join(", ", (@interleave, @interleave)) . "\n";
+ $code .= ".align 32\n.Lmlkem_modroots:\n";
+ for my $block (0..15) {
+ for my $kind (0, 1) {
+ my @words = map { ($_ & 1)
+ ? ($kind ? $modshoup[8 * $block + ($_ >> 1)]
+ : $modroots[8 * $block + ($_ >> 1)])
+ : ($kind ? 0 : 1) } 0..15;
+ my @packed = map { $words[2 * $_]
+ + ($words[2 * $_ + 1] << 16) } 0..7;
+ $code .= " .long " . join(", ", @packed) . "\n";
+ }
+ }
} else {
$code .= <<'___';
.text
-.globl mlkem_ntt_avx2_capable
-.type mlkem_ntt_avx2_capable,@abi-omnipotent
-mlkem_ntt_avx2_capable:
+.globl mlkem_avx2_capable
+.type mlkem_avx2_capable,@abi-omnipotent
+mlkem_avx2_capable:
xor %eax, %eax
ret
-.Lmlkem_ntt_avx2_capable_end:
-.size mlkem_ntt_avx2_capable, .-mlkem_ntt_avx2_capable
+.Lmlkem_avx2_capable_end:
+.size mlkem_avx2_capable, .-mlkem_avx2_capable
.globl mlkem_ntt_avx2
.type mlkem_ntt_avx2,@abi-omnipotent
mlkem_ntt_avx2:
@@ -433,33 +586,57 @@ mlkem_inverse_ntt_avx2:
.byte 0x0f,0x0b
.Lmlkem_inverse_ntt_avx2_end:
.size mlkem_inverse_ntt_avx2, .-mlkem_inverse_ntt_avx2
+.globl mlkem_basemul_acc_avx2
+.type mlkem_basemul_acc_avx2,@abi-omnipotent
+mlkem_basemul_acc_avx2:
+ .byte 0x0f,0x0b
+.Lmlkem_basemul_acc_avx2_end:
+.size mlkem_basemul_acc_avx2, .-mlkem_basemul_acc_avx2
+.globl mlkem_basemul_acc_transpose_add_avx2
+.type mlkem_basemul_acc_transpose_add_avx2,@abi-omnipotent
+mlkem_basemul_acc_transpose_add_avx2:
+ .byte 0x0f,0x0b
+.Lmlkem_basemul_acc_transpose_add_avx2_end:
+.size mlkem_basemul_acc_transpose_add_avx2, .-mlkem_basemul_acc_transpose_add_avx2
___
}
if ($win64) {
- # All three routines are leaf functions. They do not change RSP or any
- # nonvolatile register, so a zero-code UNWIND_INFO describes each one.
+ # Win64 uses only volatile GPRs and YMM0..YMM5. The preserved halves of
+ # XMM6..XMM15 and all nonvolatile GPRs stay unchanged, so no register
+ # save/restore is needed. These leaf functions leave RSP unchanged;
+ # zero-code UNWIND_INFO describes each one.
$code .= <<'___';
.section .pdata
.align 4
- .rva mlkem_ntt_avx2_capable
- .rva .Lmlkem_ntt_avx2_capable_end
- .rva .Lmlkem_ntt_avx2_capable_unwind
+ .rva mlkem_avx2_capable
+ .rva .Lmlkem_avx2_capable_end
+ .rva .Lmlkem_avx2_capable_unwind
.rva mlkem_ntt_avx2
.rva .Lmlkem_ntt_avx2_end
.rva .Lmlkem_ntt_avx2_unwind
.rva mlkem_inverse_ntt_avx2
.rva .Lmlkem_inverse_ntt_avx2_end
.rva .Lmlkem_inverse_ntt_avx2_unwind
+ .rva mlkem_basemul_acc_avx2
+ .rva .Lmlkem_basemul_acc_avx2_end
+ .rva .Lmlkem_basemul_acc_avx2_unwind
+ .rva mlkem_basemul_acc_transpose_add_avx2
+ .rva .Lmlkem_basemul_acc_transpose_add_avx2_end
+ .rva .Lmlkem_basemul_acc_transpose_add_avx2_unwind
.section .xdata
.align 4
-.Lmlkem_ntt_avx2_capable_unwind:
+.Lmlkem_avx2_capable_unwind:
.byte 1,0,0,0
.Lmlkem_ntt_avx2_unwind:
.byte 1,0,0,0
.Lmlkem_inverse_ntt_avx2_unwind:
.byte 1,0,0,0
+.Lmlkem_basemul_acc_avx2_unwind:
+ .byte 1,0,0,0
+.Lmlkem_basemul_acc_transpose_add_avx2_unwind:
+ .byte 1,0,0,0
___
}
diff --git a/crypto/ml_kem/build.info b/crypto/ml_kem/build.info
index 0b6ab49e83..4a8a4458ca 100644
--- a/crypto/ml_kem/build.info
+++ b/crypto/ml_kem/build.info
@@ -2,8 +2,8 @@ LIBS = ../../libcrypto
$MLKEMASM=
IF[{- !$disabled{asm} -}]
- $MLKEMASM_x86_64=mlkem_ntt-x86_64.s
- $MLKEMDEF_x86_64=MLKEM_NTT_X86_64_ASM
+ $MLKEMASM_x86_64=ml_kem_avx2.s
+ $MLKEMDEF_x86_64=MLKEM_AVX2_X86_64_ASM
$MLKEMDEF_ppc64=MLKEM_NTT_PPC_ASM
IF[{- $target{sys_id} ne "AIX" && $target{sys_id} ne "MACOSX" -}]
@@ -39,4 +39,4 @@ ENDIF
DEFINE[../../libcrypto]=$MLKEMDEF
DEFINE[../../providers/libfips.a]=$MLKEMDEF
-GENERATE[mlkem_ntt-x86_64.s]=asm/mlkem_ntt-x86_64.pl
+GENERATE[ml_kem_avx2.s]=asm/ml_kem_avx2.pl
diff --git a/crypto/ml_kem/ml_kem.c b/crypto/ml_kem/ml_kem.c
index ba10e11615..5fbf8b41a3 100644
--- a/crypto/ml_kem/ml_kem.c
+++ b/crypto/ml_kem/ml_kem.c
@@ -488,8 +488,8 @@ static ml_kem_scalar_inverse_ntt_fn scalar_inverse_ntt = ossl_ml_kem_scalar_inve
*/
static ml_kem_scalar_inverse_ntt_demontgomerize_fn scalar_inverse_ntt_demontgomerize = ossl_ml_kem_scalar_inverse_ntt_generic;
-#if defined(MLKEM_NTT_X86_64_ASM)
-int mlkem_ntt_avx2_capable(void);
+#if defined(MLKEM_AVX2_X86_64_ASM)
+int mlkem_avx2_capable(void);
void mlkem_ntt_avx2(uint16_t *c);
void mlkem_inverse_ntt_avx2(uint16_t *c);
@@ -565,12 +565,48 @@ static void scalar_inverse_ntt_ppc(scalar *s)
typedef void (*ml_kem_scalar_mult_add_fn)(scalar *out, const scalar *lhs, const scalar *rhs);
typedef void (*ml_kem_inner_product_montgomery_fn)(scalar *out, const scalar *lhs, const scalar *rhs, int rank);
typedef void (*ml_kem_matrix_mult_intt_fn)(scalar *out, const scalar *m, const scalar *a, int rank);
+typedef void (*ml_kem_matrix_mult_transpose_add_fn)(scalar *out, const scalar *m, const scalar *a, int rank);
/* Forward declarations */
static void scalar_mult_generic(scalar *out, const scalar *lhs, const scalar *rhs);
static void scalar_mult_add_generic(scalar *out, const scalar *lhs, const scalar *rhs);
static void inner_product_generic(scalar *out, const scalar *lhs, const scalar *rhs, int rank);
static void matrix_mult_intt_generic(scalar *out, const scalar *m, const scalar *a, int rank);
+static void matrix_mult_transpose_add_generic(scalar *out, const scalar *m, const scalar *a, int rank);
+
+#if defined(MLKEM_AVX2_X86_64_ASM)
+void mlkem_basemul_acc_avx2(uint16_t *out, const uint16_t *lhs,
+ const uint16_t *rhs, int rank);
+void mlkem_basemul_acc_transpose_add_avx2(uint16_t *out, const uint16_t *lhs,
+ const uint16_t *rhs, int rank);
+
+static void inner_product_avx2(scalar *out, const scalar *lhs,
+ const scalar *rhs, int rank)
+{
+ mlkem_basemul_acc_avx2(out->c, lhs->c, rhs->c, rank);
+}
+
+static void matrix_mult_intt_avx2(scalar *out, const scalar *m,
+ const scalar *a, int rank)
+{
+ int i;
+
+ for (i = 0; i < rank; i++, out++, m += rank) {
+ mlkem_basemul_acc_avx2(out->c, m->c, a->c, rank);
+ scalar_inverse_ntt(out);
+ }
+}
+
+static void matrix_mult_transpose_add_avx2(scalar *out, const scalar *m,
+ const scalar *a, int rank)
+{
+ int i;
+
+ for (i = 0; i < rank; i++)
+ mlkem_basemul_acc_transpose_add_avx2(out[i].c, m[i].c, a[0].c, rank);
+}
+
+#endif
/* Function pointers for dispatch */
static ml_kem_scalar_mult_add_fn scalar_mult_add = scalar_mult_add_generic;
@@ -582,14 +618,18 @@ static ml_kem_scalar_mult_add_fn scalar_mult_add = scalar_mult_add_generic;
*/
static ml_kem_inner_product_montgomery_fn inner_product_montgomery = inner_product_generic;
static ml_kem_matrix_mult_intt_fn matrix_mult_intt = matrix_mult_intt_generic;
+static ml_kem_matrix_mult_transpose_add_fn matrix_mult_transpose_add = matrix_mult_transpose_add_generic;
static void ml_kem_ntt_init(void)
{
-#if defined(MLKEM_NTT_X86_64_ASM)
- if (mlkem_ntt_avx2_capable()) {
+#if defined(MLKEM_AVX2_X86_64_ASM)
+ if (mlkem_avx2_capable()) {
scalar_ntt = scalar_ntt_avx2;
scalar_inverse_ntt = scalar_inverse_ntt_avx2;
scalar_inverse_ntt_demontgomerize = scalar_inverse_ntt_avx2;
+ inner_product_montgomery = inner_product_avx2;
+ matrix_mult_intt = matrix_mult_intt_avx2;
+ matrix_mult_transpose_add = matrix_mult_transpose_add_avx2;
}
#endif
@@ -1157,7 +1197,7 @@ static void matrix_mult_intt_generic(scalar *out, const scalar *m, const scalar
/* Here, the output vector must not overlap with the inputs */
static void
-matrix_mult_transpose_add(scalar *out, const scalar *m, const scalar *a, int rank)
+matrix_mult_transpose_add_generic(scalar *out, const scalar *m, const scalar *a, int rank)
{
const scalar *mc = m, *mr, *ar;
int i, j;