Add Optimized and HOL Light verified AVX2 Keccak x4 (#3020)

### Issues:
Import AVX2 Optimized and HOL Light verified 4x Keccak permutation
awslabs/s2n-bignum#354

NOTE:: Once awslabs/s2n-bignum#354 is merged,
the assembly files would be imported directly with the importer script.

### Description of changes: 
This PR introduces an optimized AVX2 implementation of the
Keccak-f[1600] x4 permutation, formally verified using HOL Light. This
batched Keccak implementation processes four independent Keccak
permutations in parallel using AVX2 SIMD instructions, significantly
accelerating the core hash operations underlying ML-KEM (FIPS 203) and
ML-DSA (FIPS 204).

The 4-way parallel Keccak permutation is a critical building block for
lattice-based cryptographic schemes, as it is heavily used in:

- ML-KEM: Matrix/vector sampling, seed expansion, and hash operations
during keygen, encapsulation, and decapsulation
- ML-DSA: Key generation, signing (rejection sampling), and verification

<h4>Performance Results</h4>
<p>The optimization delivers substantial throughput improvements across
all tested EC2 instance types:</p>
<p><strong>Average Speedups by Algorithm Family:</strong></p>

Algorithm | c7i | c7a | c6i | c6a
-- | -- | -- | -- | --
ML-KEM-512 | +29.2% | +38.7% | +41.4% | +37.6%
ML-KEM-768 | +29.3% | +37.6% | +37.4% | +37.4%
ML-KEM-1024 | +34.8% | +46.7% | +51.0% | +48.4%
MLDSA44 | +16.6% | +17.9% | +23.0% | +21.1%
MLDSA65 | +19.6% | +18.4% | +23.9% | +20.8%
MLDSA87 | +28.5% | +28.0% | +34.7% | +31.9%

<p><strong>Notable highlights:</strong></p>
<ul>
<li>Peak speedup of <strong>+59.0%</strong> for ML-KEM-1024 keygen on
c6i</li>
<li>ML-KEM-1024 benefits the most across all platforms (up to
<strong>+52.7%</strong> on c7a), as larger parameter sets invoke more
Keccak calls</li>
<li>MLDSA signing shows modest gains (+2–13%) since its runtime is 
dominated by rejection sampling rather than Keccak permutation 
throughput</li>
<li>Improvements are consistent across both Intel and AMD platforms and 
across both current (Gen 7) and previous generation (Gen 6)
instances</li>
</ul>

### Call-outs:

- Reviewers should pay attention to the integration points where the new
Keccak x4 is wired into the ML-KEM and ML-DSA call paths

### Testing:
- All existing ML-KEM and ML-DSA KAT (Known Answer Tests) pass,
confirming functional correctness
`./crypto/crypto_test `
- Performance benchmarked using ./tool/bssl speed on four EC2 instance
types (c7i, c7a, c6i, c6a) to validate throughput improvements
`tool/bssl speed -filter "ML-KEM"`
`tool/bssl speed -filter "MLDSA"`

### More Performance Data

<h4>EC2 c7i</h4>

Algorithm | Operation | Original (ops/sec) | New (ops/sec) | Speedup
-- | -- | -- | -- | --
ML-KEM-512 | keygen | 102,039.0 | 137,883.4 | +35.1%
ML-KEM-512 | encaps | 92,432.1 | 118,961.3 | +28.7%
ML-KEM-512 | decaps | 77,155.5 | 95,523.5 | +23.8%
ML-KEM-768 | keygen | 65,240.3 | 86,148.5 | +32.1%
ML-KEM-768 | encaps | 60,583.8 | 79,416.7 | +31.1%
ML-KEM-768 | decaps | 51,275.5 | 64,007.9 | +24.8%
ML-KEM-1024 | keygen | 43,752.6 | 62,079.1 | +41.9%
ML-KEM-1024 | encaps | 40,528.9 | 54,745.9 | +35.1%
ML-KEM-1024 | decaps | 35,182.9 | 44,833.4 | +27.4%
MLDSA44 | keygen | 19,594.8 | 23,784.0 | +21.4%
MLDSA44 | signing | 4,776.1 | 5,105.4 | +6.9%
MLDSA44 | verify | 18,485.0 | 22,439.7 | +21.4%
MLDSA65 | keygen | 10,078.2 | 12,485.9 | +23.9%
MLDSA65 | signing | 3,030.3 | 3,263.0 | +7.7%
MLDSA65 | verify | 11,629.3 | 14,807.7 | +27.3%
MLDSA87 | keygen | 7,177.4 | 9,908.2 | +38.0%
MLDSA87 | signing | 2,534.6 | 2,776.0 | +9.5%
MLDSA87 | verify | 7,049.3 | 9,737.1 | +38.1%
 
<h4>EC2 c7a</h4>

Algorithm | Operation | Original (ops/sec) | New (ops/sec) | Speedup
-- | -- | -- | -- | --
ML-KEM-512 | keygen | 94,563.7 | 137,392.5 | +45.3%
ML-KEM-512 | encaps | 85,020.5 | 118,473.0 | +39.3%
ML-KEM-512 | decaps | 71,284.2 | 93,645.4 | +31.4%
ML-KEM-768 | keygen | 56,037.5 | 79,772.7 | +42.4%
ML-KEM-768 | encaps | 52,744.2 | 73,353.1 | +39.1%
ML-KEM-768 | decaps | 44,832.4 | 58,874.9 | +31.3%
ML-KEM-1024 | keygen | 37,007.2 | 56,511.5 | +52.7%
ML-KEM-1024 | encaps | 34,843.2 | 51,659.7 | +48.3%
ML-KEM-1024 | decaps | 30,052.9 | 41,833.0 | +39.2%
MLDSA44 | keygen | 17,087.6 | 21,781.5 | +27.5%
MLDSA44 | signing | 3,833.2 | 3,941.9 | +2.8%
MLDSA44 | verify | 15,055.9 | 18,594.2 | +23.5%
MLDSA65 | keygen | 9,295.8 | 11,665.2 | +25.5%
MLDSA65 | signing | 2,418.1 | 2,468.3 | +2.1%
MLDSA65 | verify | 9,658.5 | 12,321.2 | +27.6%
MLDSA87 | keygen | 6,458.0 | 9,079.5 | +40.6%
MLDSA87 | signing | 2,021.0 | 2,147.5 | +6.3%
MLDSA87 | verify | 6,094.7 | 8,355.6 | +37.1%

<h4>EC2 c6i</h4>

Algorithm | Operation | Original (ops/sec) | New (ops/sec) | Speedup
-- | -- | -- | -- | --
ML-KEM-512 | keygen | 74,243.5 | 110,577.2 | +48.9%
ML-KEM-512 | encaps | 66,855.4 | 94,949.4 | +42.0%
ML-KEM-512 | decaps | 56,641.7 | 75,435.7 | +33.2%
ML-KEM-768 | keygen | 44,861.7 | 63,758.2 | +42.1%
ML-KEM-768 | encaps | 42,938.1 | 59,149.7 | +37.7%
ML-KEM-768 | decaps | 35,953.7 | 47,646.6 | +32.5%
ML-KEM-1024 | keygen | 30,333.7 | 48,230.3 | +59.0%
ML-KEM-1024 | encaps | 28,685.5 | 43,577.8 | +51.9%
ML-KEM-1024 | decaps | 23,941.1 | 33,996.0 | +42.0%
MLDSA44 | keygen | 14,708.0 | 19,526.7 | +32.8%
MLDSA44 | signing | 3,488.7 | 3,693.8 | +5.9%
MLDSA44 | verify | 13,581.8 | 17,702.9 | +30.3%
MLDSA65 | keygen | 7,868.8 | 10,223.9 | +29.9%
MLDSA65 | signing | 2,153.5 | 2,309.9 | +7.3%
MLDSA65 | verify | 8,542.6 | 11,490.7 | +34.5%
MLDSA87 | keygen | 5,428.8 | 8,082.0 | +48.9%
MLDSA87 | signing | 1,819.2 | 1,973.5 | +8.5%
MLDSA87 | verify | 5,258.6 | 7,708.2 | +46.6%

<h4>EC2 c6a</h4>

Algorithm | Operation | Original (ops/sec) | New (ops/sec) | Speedup
-- | -- | -- | -- | --
ML-KEM-512 | keygen | 94,817.9 | 138,020.5 | +45.6%
ML-KEM-512 | encaps | 87,240.8 | 120,129.4 | +37.7%
ML-KEM-512 | decaps | 72,457.8 | 93,790.5 | +29.4%
ML-KEM-768 | keygen | 60,954.7 | 87,137.1 | +42.9%
ML-KEM-768 | encaps | 57,065.8 | 79,197.9 | +38.8%
ML-KEM-768 | decaps | 47,498.6 | 62,029.9 | +30.6%
ML-KEM-1024 | keygen | 41,115.4 | 64,320.3 | +56.4%
ML-KEM-1024 | encaps | 38,250.1 | 57,234.5 | +49.6%
ML-KEM-1024 | decaps | 32,493.6 | 45,190.5 | +39.1%
MLDSA44 | keygen | 16,540.2 | 21,594.4 | +30.6%
MLDSA44 | signing | 3,670.8 | 3,916.2 | +6.7%
MLDSA44 | verify | 14,447.5 | 18,216.4 | +26.1%
MLDSA65 | keygen | 8,977.5 | 11,496.0 | +28.0%
MLDSA65 | signing | 2,346.1 | 2,422.6 | +3.3%
MLDSA65 | verify | 9,377.8 | 12,294.7 | +31.1%
MLDSA87 | keygen | 6,224.4 | 8,939.6 | +43.6%
MLDSA87 | signing | 1,961.1 | 2,208.0 | +12.6%
MLDSA87 | verify | 5,862.3 | 8,183.7 | +39.6%



By submitting this pull request, I confirm that my contribution is made
under the terms of the Apache 2.0 license and the ISC license.

Author: manastasova

crypto/fipsmodule/CMakeLists.txt

@@ -250,6 +250,9 @@ if((((ARCH STREQUAL "x86_64") AND NOT MY_ASSEMBLER_IS_TOO_OLD_FOR_512AVX) OR
     ${S2N_BIGNUM_DIR}/curve25519/edwards25519_scalarmulbase_alt.S
     ${S2N_BIGNUM_DIR}/curve25519/edwards25519_scalarmuldouble.S
     ${S2N_BIGNUM_DIR}/curve25519/edwards25519_scalarmuldouble_alt.S
+
+    ${S2N_BIGNUM_DIR}/sha3/sha3_keccak_f1600.S
+    ${S2N_BIGNUM_DIR}/sha3/sha3_keccak4_f1600_alt.S
   )
 
   if(ARCH STREQUAL "x86_64")
@@ -264,7 +267,6 @@ if((((ARCH STREQUAL "x86_64") AND NOT MY_ASSEMBLER_IS_TOO_OLD_FOR_512AVX) OR
                 ${S2N_BIGNUM_DIR}/curve25519/curve25519_x25519_alt.S
                 ${S2N_BIGNUM_DIR}/curve25519/curve25519_x25519base.S
                 ${S2N_BIGNUM_DIR}/curve25519/curve25519_x25519base_alt.S
-                ${S2N_BIGNUM_DIR}/sha3/sha3_keccak_f1600.S
     )
   elseif(ARCH STREQUAL "aarch64")
     # byte-level interface for aarch64 s2n-bignum x25519 are in
@@ -303,14 +305,8 @@ if((((ARCH STREQUAL "x86_64") AND NOT MY_ASSEMBLER_IS_TOO_OLD_FOR_512AVX) OR
     # (gcc supports it since gcc8, clang supports it since clang7)
     check_compiler("neon_sha3_check.c" MY_ASSEMBLER_SUPPORTS_NEON_SHA3_EXTENSION "-march=armv8.4-a+sha3")
 
-    list(APPEND BCM_ASM_SOURCES
-                # Scalar Keccak-x1 assembly from s2n-bignum/mlkem-native
-                ${S2N_BIGNUM_DIR}/sha3/sha3_keccak_f1600.S
-
-                # Batched Keccak-x4 assembly from s2n-bignum
-                # Scalar version for Neoverse N1
-                ${S2N_BIGNUM_DIR}/sha3/sha3_keccak4_f1600_alt.S
-    )
+    # Note: Scalar Keccak-x1 and Batched Neoverse N1 Keccak-x4 assembly from s2n-bignum
+    # are included in the common s2n-bignum block above
 
     if(MY_ASSEMBLER_SUPPORTS_NEON_SHA3_EXTENSION)
         list(APPEND BCM_ASM_SOURCES

crypto/fipsmodule/sha/keccak1600.c

@@ -33,6 +33,18 @@ static const uint64_t iotas[] = {
     0x8000000080008008ULL
 };
 
+#if defined(OPENSSL_X86_64)
+static const uint64_t keccak_rho8[4] = {
+    0x0605040302010007ULL, 0x0E0D0C0B0A09080FULL,
+    0x0605040302010007ULL, 0x0E0D0C0B0A09080FULL
+};
+
+static const uint64_t keccak_rho56[4] = {
+    0x0007060504030201ULL, 0x080F0E0D0C0B0A09ULL,
+    0x0007060504030201ULL, 0x080F0E0D0C0B0A09ULL
+};
+#endif
+
 #if !defined(KECCAK1600_ASM)
 
 static const uint8_t rhotates[KECCAK1600_ROWS][KECCAK1600_ROWS] = {
@@ -315,7 +327,7 @@ void Keccak1600_Squeeze(uint64_t A[KECCAK1600_ROWS][KECCAK1600_ROWS], uint8_t *o
 // Scalar implementation from OpenSSL provided by keccak1600-armv8.pl
 extern void KeccakF1600_hw(uint64_t state[25]);
 
-#if defined(OPENSSL_AARCH64)
+#if defined(OPENSSL_AARCH64) || defined(OPENSSL_X86_64)
 static void keccak_log_dispatch(size_t id) {
 #if BORINGSSL_DISPATCH_TEST
     BORINGSSL_function_hit[id] = 1;
@@ -366,6 +378,7 @@ void KeccakF1600(uint64_t A[KECCAK1600_ROWS][KECCAK1600_ROWS]) {
     KeccakF1600_hw((uint64_t *) A);
 
 #elif defined(OPENSSL_X86_64)
+    keccak_log_dispatch(9); // kFlag_sha3_keccak_f1600
     sha3_keccak_f1600((uint64_t *)A, iotas);
 #endif
 }
@@ -443,6 +456,14 @@ static void Keccak1600_x4(uint64_t A[4][KECCAK1600_ROWS][KECCAK1600_ROWS]) {
 #endif
 #endif
 
+#if defined(KECCAK1600_S2N_BIGNUM_ASM) && defined(OPENSSL_X86_64)
+    if (CRYPTO_is_AVX2_capable()) {
+        keccak_log_dispatch(10); // kFlag_sha3_keccak4_f1600_alt
+        sha3_keccak4_f1600_alt((uint64_t *)A, iotas, keccak_rho8, keccak_rho56);
+        return;
+    }
+#endif
+
     // Fallback: 4x individual KeccakF1600 calls (each with their own dispatch)
     KeccakF1600(A[0]);
     KeccakF1600(A[1]);

crypto/impl_dispatch_test.cc

@@ -31,6 +31,7 @@ class ImplDispatchTest : public ::testing::Test {
     aes_hw_ = CRYPTO_is_AESNI_capable();
     avx_movbe_ = CRYPTO_is_AVX_capable() && CRYPTO_is_MOVBE_capable();
     aes_vpaes_ = CRYPTO_is_SSSE3_capable();
+    is_avx2_ = CRYPTO_is_AVX2_capable();
     ifma_avx512 = CRYPTO_is_AVX512IFMA_capable();
     sha_ext_ =
     // sha_ext_ isn't enabled on 32-bit x86 architectures.
@@ -85,13 +86,13 @@ class ImplDispatchTest : public ::testing::Test {
         true;
 #else
         false;
+#endif
 #endif
     have_s2n_bignum_asm_ =
 #if defined(KECCAK1600_S2N_BIGNUM_ASM)
         true;
 #else
         false;
-#endif
 #endif
   }
 
@@ -122,10 +123,12 @@ class ImplDispatchTest : public ::testing::Test {
   bool aes_hw_ = false;
   bool aes_vpaes_ = false;
   bool sha_ext_ = false;
+  bool have_s2n_bignum_asm_ = false;
 #if defined(OPENSSL_X86) || defined(OPENSSL_X86_64)
   bool vaes_vpclmulqdq_ = false;
   bool avx_movbe_ = false;
   bool is_x86_64_ = false;
+  bool is_avx2_ = false;
   bool is_assembler_too_old = false;
   bool is_assembler_too_old_avx512 = false;
   bool ifma_avx512 = false;
@@ -138,7 +141,6 @@ class ImplDispatchTest : public ::testing::Test {
   bool neoverse_v1_ = false;
   bool neoverse_v2_ = false;
   bool assembler_has_neon_sha3_extension_ = false;
-  bool have_s2n_bignum_asm_ = false;
 #endif
 
 };
@@ -156,6 +158,8 @@ constexpr size_t kFlag_sha256_hw = 6;
 constexpr size_t kFlag_aesni_gcm_encrypt = 2;
 constexpr size_t kFlag_aes_gcm_encrypt_avx512 = 7;
 constexpr size_t kFlag_RSAZ_mod_exp_avx512_x2 = 8;
+constexpr size_t kFlag_sha3_keccak_f1600 = 9;
+constexpr size_t kFlag_sha3_keccak4_f1600_alt = 10;
 #else // AARCH64
 constexpr size_t kFlag_aes_gcm_enc_kernel = 2;
 constexpr size_t kFlag_aesv8_gcm_8x_enc_128 = 7;
@@ -297,8 +301,24 @@ TEST_F(ImplDispatchTest, SHA3_512) {
         SHA3_512(in, 32, out);
       });
 }
+#endif // OPENSSL_AARCH64
 
 TEST_F(ImplDispatchTest, SHAKE256_Batched) {
+#if defined(OPENSSL_X86_64) || defined(OPENSSL_X86)
+  // Assembly dispatch logic for Keccak-x4 on x86:
+  // - For platforms with AVX2 support, we use batched Keccak assembly from s2n-bignum
+  //   (`sha3_keccak4_f1600_alt()`).
+  // - Otherwise, fall back to scalar Keccak implementation from s2n-bignum,
+  //   (`sha3_keccak_f1600()`).
+  AssertFunctionsHit(
+      {
+          {kFlag_sha3_keccak4_f1600_alt,
+           have_s2n_bignum_asm_ &&
+           is_avx2_ },
+           {kFlag_sha3_keccak_f1600,
+           have_s2n_bignum_asm_ && is_x86_64_ && !is_avx2_ },
+      },
+#else // AARCH64
   // Assembly dispatch logic for Keccak-x4 on AArch64:
   // - For Neoverse N1, we use scalar batched hybrid Keccak assembly from s2n-bignum
   //   (`sha3_keccak4_f1600_alt()`) leveraging Neon and scalar assembly with
@@ -342,14 +362,13 @@ TEST_F(ImplDispatchTest, SHAKE256_Batched) {
 	       !(assembler_has_neon_sha3_extension_ && sha3_ext_)
 	   ) },
       },
+#endif
       [] {
         const uint8_t in[32] = {0};
         uint8_t out0[32], out1[32], out2[32], out3[32];
         SHAKE256_x4(in, in, in, in, 32, out0, out1, out2, out3, 32);
       });
 }
-#endif // OPENSSL_AARCH64
-
 
 #if defined(OPENSSL_X86) || defined(OPENSSL_X86_64)
 static bssl::UniquePtr<BIGNUM> GetBIGNUM(FileTest *t, const char *attr);

crypto/internal.h

@@ -1288,6 +1288,8 @@ OPENSSL_INLINE int boringssl_fips_break_test(const char *test) {
 //   6: sha256_block_data_order_shaext
 //   7: aes_gcm_encrypt_avx512
 //   8: RSAZ_mod_exp_avx512_x2
+//   9: sha3_keccak_f1600
+//  10: sha3_keccak4_f1600_alt
 // On AARCH64:
 //   0: aes_hw_ctr32_encrypt_blocks
 //   1: aes_hw_encrypt

third_party/s2n-bignum/META.yml

@@ -1,5 +1,5 @@
 name: s2n-bignum-imported
 source: awslabs/s2n-bignum.git
-commit: 2b5350cf955a32d2f7aced172f7bd28dd85a8587
+commit: 4b5f8214e85b6b239077c278825b7fa9c2ab9cf5
 target: main
-imported-at: 2025-09-05T04:04:05+0000
+imported-at: 2026-03-16T15:51:12+0000

third_party/s2n-bignum/s2n-bignum-imported/README.md

@@ -11,9 +11,11 @@ tuned for highest performance both by hand and using automatic optimization
 techniques such as the [SLOTHY](https://github.com/slothy-optimizer/slothy)
 superoptimizer, and each function is accompanied by a machine-checked formal
 proof in [HOL-Light](https://hol-light.github.io/) that its mathematical
-result is correct, based on a formal model of the underlying machine. Each
-function is moreover written in a constant-time style to avoid timing
-side-channels.
+result is correct, based on a formal model of the underlying machine.
+Moreover, each function is written in a constant-time style to avoid timing
+side-channels. For a detailed analysis of the formal verification process, the
+assumptions made, and the correspondence of formal models with reality, please
+refer to the [s2n-bignum soundness review](doc/s2n_bignum_soundness.md).
 
 For the SHA-3 and ML-KEM code currently part of s2n-bignum, some of the
 comments in the main part of this README do not apply exactly. See the section

third_party/s2n-bignum/s2n-bignum-imported/arm/Makefile

@@ -146,6 +146,7 @@ BIGNUM_OBJ = curve25519/bignum_add_p25519.o \
              curve25519/bignum_invsqrt_p25519_alt.o \
              curve25519/bignum_madd_n25519.o \
              curve25519/bignum_madd_n25519_alt.o \
+             curve25519/bignum_mod_m25519.o \
              curve25519/bignum_mod_m25519_4.o \
              curve25519/bignum_mod_n25519.o \
              curve25519/bignum_mod_n25519_4.o \
@@ -358,6 +359,7 @@ BIGNUM_OBJ = curve25519/bignum_add_p25519.o \
              secp256k1/bignum_triple_p256k1.o \
              sha3/sha3_keccak_f1600.o \
              sha3/sha3_keccak_f1600_alt.o \
+             sha3/sha3_keccak_f1600_alt2.o \
              sha3/sha3_keccak2_f1600.o \
              sha3/sha3_keccak2_f1600_alt.o \
              sha3/sha3_keccak4_f1600.o \
@@ -408,7 +410,9 @@ OBJ = $(POINT_OBJ) $(BIGNUM_OBJ)
 
 TUTORIAL_PROOFS = $(wildcard tutorial/*.ml)
 
-TUTORIAL_OBJ = $(TUTORIAL_PROOFS:.ml=.o) tutorial/rel_loop2.o \
+TUTORIAL_OBJ = $(filter-out tutorial/safety.o, $(TUTORIAL_PROOFS:.ml=.o)) \
+    curve25519/bignum_mod_n25519.o p256/bignum_mux_4.o \
+    tutorial/rel_loop2.o \
     tutorial/rel_simp2.o tutorial/rel_veceq2.o tutorial/rel_equivtac2.o \
     tutorial/rel_reordertac2.o tutorial/rodata_local.o
 
@@ -517,6 +521,11 @@ sm2/sm2_montjscalarmul_alt.native: sm2/sm2_montjadd_alt.native sm2/sm2_montjdoub
 
 # Tutorial
 
+# safety tutorial does not have safety.o and uses existing ones
+tutorial/safety.native: tutorial/safety.ml proofs/bignum_mux_4.ml \
+                        curve25519/bignum_mod_n25519.o ; \
+  ../tools/build-proof.sh "$<" "$(HOLLIGHT)" "$@"
+# other tutorials have their .o files
 .SECONDEXPANSION:
 tutorial/%.native: tutorial/%.ml tutorial/%.o ; ../tools/build-proof.sh "$<" "$(HOLLIGHT)" "$@"
 # Additional dependencies on .o files

third_party/s2n-bignum/s2n-bignum-imported/arm/allowed_asm

@@ -52,6 +52,10 @@
 : lsl$
 : lsr$
 : madd$
+: mla$
+: mla.2s$
+: mla.4s$
+: mla.8h$
 : mls$
 : mls.2s$
 : mls.8h$

third_party/s2n-bignum/s2n-bignum-imported/arm/curve25519/Makefile

@@ -29,6 +29,7 @@ OBJ = bignum_add_p25519.o \
       bignum_invsqrt_p25519_alt.o \
       bignum_madd_n25519.o \
       bignum_madd_n25519_alt.o \
+      bignum_mod_m25519.o \
       bignum_mod_m25519_4.o \
       bignum_mod_n25519.o \
       bignum_mod_n25519_4.o \