Implement a vectorized "quad Keccak".

Gains 7% on AMD EPYC 7B13 and 2% on Apple M1 Pro.

Bug: 503700354
Change-Id: I34d14d97a76f8892d1c54dc934b3ca6c6a6a6964
Reviewed-on: https://boringssl-review.googlesource.com/c/boringssl/+/93087
Reviewed-by: David Benjamin <davidben@google.com>
Commit-Queue: Rudolf Polzer <rpolzer@google.com>
diff --git a/crypto/fipsmodule/keccak/internal.h b/crypto/fipsmodule/keccak/internal.h
index 4d11b48..a69081e 100644
--- a/crypto/fipsmodule/keccak/internal.h
+++ b/crypto/fipsmodule/keccak/internal.h
@@ -17,6 +17,8 @@
 
 #include <openssl/base.h>
 
+#include "../../internal.h"
+
 
 BSSL_NAMESPACE_BEGIN
 
@@ -72,7 +74,19 @@
 
 #if defined(__has_attribute)
 #if __has_attribute(vector_size)
+
 #define HAVE_KECCAK_X2
+
+#if defined(OPENSSL_AARCH64)
+#define HAVE_KECCAK_X4
+#define KECCAK_X4_TARGET
+#endif
+
+#if defined(OPENSSL_X86_64)
+#define HAVE_KECCAK_X4
+#define KECCAK_X4_TARGET __attribute__((target("avx2")))
+#endif
+
 #endif  // vector_size
 #endif  // __has_attribute
 
@@ -91,6 +105,31 @@
     enum boringssl_keccak_config_t config);
 #endif
 
+#if defined(HAVE_KECCAK_X4)
+// BORINGSSL_have_keccak_x4 needs to return true before using any of the _x4
+// variants here, as they may require additional CPU features.
+inline bool BORINGSSL_have_keccak_x4() {
+#if defined(OPENSSL_X86_64)
+  return CRYPTO_is_AVX2_capable();
+#else
+  return true;
+#endif
+}
+
+// BORINGSSL_keccak_squeeze_x4 performs BORINGSSL_keccak_squeeze in parallel
+// with four same-length outputs. The contexts must be in equivalent state (i.e.
+// same config, same amount of bytes absorbed and squeezed).
+OPENSSL_EXPORT void BORINGSSL_keccak_squeeze_x4(
+    struct BORINGSSL_keccak_st ctx[4], uint8_t *outs[4], size_t out_len);
+
+// BORINGSSL_keccak_short_x4 performs BORINGSSL_keccak in parallel on four
+// same-length strings with same-length outputs. in_len must be less than 72
+// (or actually |rate_bytes|).
+OPENSSL_EXPORT void BORINGSSL_keccak_short_x4(
+    uint8_t *outs[4], size_t out_len, const uint8_t *ins[4], size_t in_len,
+    enum boringssl_keccak_config_t config);
+#endif
+
 BSSL_NAMESPACE_END
 
 #endif  // OPENSSL_HEADER_CRYPTO_FIPSMODULE_KECCAK_INTERNAL_H
diff --git a/crypto/fipsmodule/keccak/keccak.cc.inc b/crypto/fipsmodule/keccak/keccak.cc.inc
index eef7e65..83b8647 100644
--- a/crypto/fipsmodule/keccak/keccak.cc.inc
+++ b/crypto/fipsmodule/keccak/keccak.cc.inc
@@ -23,14 +23,29 @@
 
 using namespace bssl;
 
+// A rotate-left function is needed for every type to use Keccak with.
+//
+// As argument passing may differ for AVX2 and non-AVX2 functions taking vector
+// arguments by value, this function takes a pointer instead; however, thanks
+// to `OPENSSL_ATTR_ALWAYS_INLINE`, this should not actually change generated
+// code.
+static inline OPENSSL_ATTR_ALWAYS_INLINE void rotl_u64(uint64_t *value,
+                                                       int shift) {
+  *value = CRYPTO_rotl_u64(*value, shift);
+}
+
 // keccak_f implements the Keccak-1600 permutation as described at
 // https://keccak.team/keccak_specs_summary.html. Each lane is represented as a
 // 64-bit value and the 5×5 lanes are stored as an array in row-major order.
 //
 // To support vectorization, U64 shall either be uint64_t, or an uint64_t-based
 // vector type.
-template <typename U64, U64 (*rotl)(U64 value, int shift) = CRYPTO_rotl_u64>
-static void keccak_f(U64 state[25]) {
+//
+// `OPENSSL_ATTR_ALWAYS_INLINE` is used here to ensure that AVX2 optimizations
+// are used whenever the calling function is `KECCAK_X4_TARGET`, as target
+// attributes are inherited into callees when inlining.
+template <typename U64, void (*rotl)(U64 *value, int shift) = rotl_u64>
+static inline OPENSSL_ATTR_ALWAYS_INLINE void keccak_f(U64 state[25]) {
   static const int kNumRounds = 24;
   for (int round = 0; round < kNumRounds; round++) {
     // θ step
@@ -41,7 +56,9 @@
     }
 
     for (int x = 0; x < 5; x++) {
-      const U64 d = c[(x + 4) % 5] ^ rotl(c[(x + 1) % 5], 1);
+      U64 next_value = c[(x + 1) % 5];
+      rotl(&next_value, 1);
+      const U64 d = c[(x + 4) % 5] ^ next_value;
       for (int y = 0; y < 5; y++) {
         state[y * 5 + x] ^= d;
       }
@@ -63,11 +80,12 @@
     // at (0, 0), but the rotation for that element is zero, and it goes to (0,
     // 0), so we can ignore it.
     U64 prev_value = state[1];
-#define PI_RHO_STEP(index, rotation)              \
-  do {                                            \
-    const U64 value = rotl(prev_value, rotation); \
-    prev_value = state[index];                    \
-    state[index] = value;                         \
+#define PI_RHO_STEP(index, rotation) \
+  do {                               \
+    U64 value = prev_value;          \
+    rotl(&value, rotation);          \
+    prev_value = state[index];       \
+    state[index] = value;            \
   } while (0)
 
     PI_RHO_STEP(10, 1);
@@ -134,8 +152,8 @@
 #if defined(HAVE_KECCAK_X2)
 typedef uint64_t v2u64 __attribute__((vector_size(16)));
 
-static inline v2u64 v_rotl(v2u64 x, int shift) {
-  return (x << shift) | (x >> ((-shift) & 63));
+static inline void v2_rotl(v2u64 *x, int shift) {
+  *x = (*x << shift) | (*x >> ((-shift) & 63));
 }
 
 static void keccak_f_x2(uint64_t state0[25], uint64_t state1[25]) {
@@ -143,7 +161,7 @@
   for (int i = 0; i < 25; i++) {
     s[i] = (v2u64){state0[i], state1[i]};
   }
-  keccak_f<v2u64, v_rotl>(s);
+  keccak_f<v2u64, v2_rotl>(s);
   for (int i = 0; i < 25; i++) {
     state0[i] = s[i][0];
     state1[i] = s[i][1];
@@ -151,6 +169,31 @@
 }
 #endif
 
+#if defined(HAVE_KECCAK_X4)
+typedef uint64_t v4u64 __attribute__((vector_size(32)));
+
+KECCAK_X4_TARGET static inline void v4_rotl(v4u64 *x, int shift) {
+  *x = (*x << shift) | (*x >> ((-shift) & 63));
+}
+
+KECCAK_X4_TARGET static void keccak_f_x4(uint64_t state0[25],
+                                         uint64_t state1[25],
+                                         uint64_t state2[25],
+                                         uint64_t state3[25]) {
+  v4u64 s[25];
+  for (int i = 0; i < 25; i++) {
+    s[i] = (v4u64){state0[i], state1[i], state2[i], state3[i]};
+  }
+  keccak_f<v4u64, v4_rotl>(s);
+  for (int i = 0; i < 25; i++) {
+    state0[i] = s[i][0];
+    state1[i] = s[i][1];
+    state2[i] = s[i][2];
+    state3[i] = s[i][3];
+  }
+}
+#endif
+
 static void keccak_init(struct BORINGSSL_keccak_st *ctx,
                         enum boringssl_keccak_config_t config) {
   size_t required_out_len;
@@ -189,9 +232,7 @@
                             enum boringssl_keccak_config_t config) {
   struct BORINGSSL_keccak_st ctx;
   BORINGSSL_keccak_init(&ctx, config);
-  if (ctx.required_out_len != 0 && out_len != ctx.required_out_len) {
-    abort();
-  }
+  BSSL_CHECK(ctx.required_out_len == 0 || out_len == ctx.required_out_len);
   BORINGSSL_keccak_absorb(&ctx, in, in_len);
   BORINGSSL_keccak_squeeze(&ctx, out, out_len);
 }
@@ -203,14 +244,13 @@
   struct BORINGSSL_keccak_st ctx[2];
   for (size_t i = 0; i < 2; ++i) {
     BORINGSSL_keccak_init(&ctx[i], config);
-    if (ctx[i].required_out_len != 0 && out_len != ctx[i].required_out_len) {
-      abort();
-    }
+    BSSL_CHECK(ctx[i].required_out_len == 0 ||
+               out_len == ctx[i].required_out_len);
 
-    // NOTE: this implementation is only efficient if in_len < ctx->rate_bytes,
-    // as right now only keccak_f calls in BORINGSSL_keccak_squeeze and
-    // keccak_finalize are vectorized. So just fail in every other case for
-    // now.
+    // NOTE: this implementation is only efficient if in_len <
+    // ctx->rate_bytes, as right now only keccak_f calls in
+    // BORINGSSL_keccak_squeeze and keccak_finalize are vectorized. So just
+    // fail in every other case for now.
     BSSL_CHECK(in_len < ctx[i].rate_bytes);
 
     BORINGSSL_keccak_absorb(&ctx[i], ins[i], in_len);
@@ -219,6 +259,28 @@
 }
 #endif
 
+#if defined(HAVE_KECCAK_X4)
+KECCAK_X4_TARGET void bssl::BORINGSSL_keccak_short_x4(
+    uint8_t *outs[4], size_t out_len, const uint8_t *ins[4], size_t in_len,
+    enum boringssl_keccak_config_t config) {
+  struct BORINGSSL_keccak_st ctx[4];
+  for (size_t i = 0; i < 4; ++i) {
+    BORINGSSL_keccak_init(&ctx[i], config);
+    BSSL_CHECK(ctx[i].required_out_len == 0 ||
+               out_len == ctx[i].required_out_len);
+
+    // NOTE: this implementation is only efficient if in_len <
+    // ctx->rate_bytes, as right now only keccak_f calls in
+    // BORINGSSL_keccak_squeeze and keccak_finalize are vectorized. So just
+    // fail in every other case for now.
+    BSSL_CHECK(in_len < ctx[i].rate_bytes);
+
+    BORINGSSL_keccak_absorb(&ctx[i], ins[i], in_len);
+  }
+  BORINGSSL_keccak_squeeze_x4(ctx, outs, out_len);
+}
+#endif
+
 void bssl::BORINGSSL_keccak_init(struct BORINGSSL_keccak_st *ctx,
                                  enum boringssl_keccak_config_t config) {
   keccak_init(ctx, config);
@@ -226,14 +288,12 @@
 
 void bssl::BORINGSSL_keccak_absorb(struct BORINGSSL_keccak_st *ctx,
                                    const uint8_t *in, size_t in_len) {
-  if (ctx->phase == boringssl_keccak_phase_squeeze) {
-    // It's illegal to call absorb() again after calling squeeze().
-    abort();
-  }
+  // It's illegal to call absorb() again after calling squeeze().
+  BSSL_CHECK(ctx->phase != boringssl_keccak_phase_squeeze);
 
   const size_t rate_words = ctx->rate_bytes / 8;
-  // XOR the input. Accessing `ctx->state` as a `uint8_t*` is allowed by strict
-  // aliasing because we require `uint8_t` to be a character type.
+  // XOR the input. Accessing `ctx->state` as a `uint8_t*` is allowed by
+  // strict aliasing because we require `uint8_t` to be a character type.
   uint8_t *state_bytes = (uint8_t *)ctx->state;
 
   // Absorb partial block.
@@ -298,8 +358,8 @@
 #if defined(HAVE_KECCAK_X2)
 static void keccak_finalize_x2(struct BORINGSSL_keccak_st ctx[2]) {
   for (size_t i = 0; i < 2; ++i) {
-    // XOR the terminator. Accessing `ctx->state` as a `uint8_t*` is allowed by
-    // strict aliasing because we require `uint8_t` to be a character type.
+    // XOR the terminator. Accessing `ctx->state` as a `uint8_t*` is allowed
+    // by strict aliasing because we require `uint8_t` to be a character type.
     uint8_t *state_bytes = (uint8_t *)ctx[i].state;
     state_bytes[ctx[i].absorb_offset] ^= keccak_terminator(&ctx[i]);
     state_bytes[ctx[i].rate_bytes - 1] ^= 0x80;
@@ -308,15 +368,39 @@
 }
 #endif
 
+#if defined(HAVE_KECCAK_X4)
+static void keccak_finalize_x4(struct BORINGSSL_keccak_st ctx[4]) {
+  for (size_t i = 0; i < 4; ++i) {
+    // XOR the terminator. Accessing `ctx->state` as a `uint8_t*` is allowed
+    // by strict aliasing because we require `uint8_t` to be a character type.
+    uint8_t terminator;
+    switch (ctx[i].config) {
+      case boringssl_sha3_256:
+      case boringssl_sha3_512:
+        terminator = 0x06;
+        break;
+      case boringssl_shake128:
+      case boringssl_shake256:
+        terminator = 0x1f;
+        break;
+      default:
+        abort();
+    }
+    uint8_t *state_bytes = (uint8_t *)ctx[i].state;
+    state_bytes[ctx[i].absorb_offset] ^= terminator;
+    state_bytes[ctx[i].rate_bytes - 1] ^= 0x80;
+  }
+  keccak_f_x4(ctx[0].state, ctx[1].state, ctx[2].state, ctx[3].state);
+}
+#endif
+
 void bssl::BORINGSSL_keccak_squeeze(struct BORINGSSL_keccak_st *ctx,
                                     uint8_t *out, size_t out_len) {
-  if (ctx->required_out_len != 0 &&
-      (ctx->phase == boringssl_keccak_phase_squeeze ||
-       out_len != ctx->required_out_len)) {
-    // The SHA-3 variants must be squeezed in a single call, to confirm that the
-    // output length is correct.
-    abort();
-  }
+  // The SHA-3 variants must be squeezed in a single call, to confirm that the
+  // output length is correct.
+  BSSL_CHECK(ctx->required_out_len == 0 ||
+             (ctx->phase != boringssl_keccak_phase_squeeze &&
+              out_len == ctx->required_out_len));
 
   if (ctx->phase == boringssl_keccak_phase_absorb) {
     keccak_finalize(ctx);
@@ -348,17 +432,15 @@
 void bssl::BORINGSSL_keccak_squeeze_x2(struct BORINGSSL_keccak_st ctx[2],
                                        uint8_t *outs[2], size_t out_len) {
   for (size_t i = 0; i < 2; ++i) {
-    if (ctx[i].required_out_len != 0 &&
-        (ctx[i].phase == boringssl_keccak_phase_squeeze ||
-         out_len != ctx[i].required_out_len)) {
-      // The SHA-3 variants must be squeezed in a single call, to confirm that
-      // the output length is correct.
-      abort();
-    }
+    // The SHA-3 variants must be squeezed in a single call, to confirm that
+    // the output length is correct.
+    BSSL_CHECK(ctx[i].required_out_len == 0 ||
+               (ctx[i].phase != boringssl_keccak_phase_squeeze &&
+                out_len == ctx[i].required_out_len));
   }
 
-  // These fields are processed in parallel. Everything here uses ctx[0]; at the
-  // end changes are mirrored back to ctx[1] just in case.
+  // These fields are processed in parallel. Everything here uses ctx[0]; at
+  // the end changes are mirrored back to ctx[1] just in case.
 #define FOR_COMMON_FIELDS(MACRO) \
   MACRO(phase)                   \
   MACRO(config)                  \
@@ -405,3 +487,69 @@
 #undef FOR_COMMON_FIELDS
 }
 #endif  // HAVE_KECCAK_X2
+
+#if defined(HAVE_KECCAK_X4)
+KECCAK_X4_TARGET void bssl::BORINGSSL_keccak_squeeze_x4(
+    struct BORINGSSL_keccak_st ctx[4], uint8_t *outs[4], size_t out_len) {
+  for (size_t i = 0; i < 4; ++i) {
+    // The SHA-3 variants must be squeezed in a single call, to confirm that
+    // the output length is correct.
+    BSSL_CHECK(ctx[i].required_out_len == 0 ||
+               (ctx[i].phase != boringssl_keccak_phase_squeeze &&
+                out_len == ctx[i].required_out_len));
+  }
+
+  // These fields are processed in parallel. Everything here uses ctx[0]; at
+  // the end changes are mirrored back to ctx[1..3] just in case.
+#define FOR_COMMON_FIELDS(MACRO) \
+  MACRO(phase)                   \
+  MACRO(config)                  \
+  MACRO(absorb_offset)           \
+  MACRO(squeeze_offset)          \
+  MACRO(rate_bytes)
+
+#define MUST_BE_EQUAL(field)                \
+  BSSL_CHECK(ctx[0].field == ctx[1].field); \
+  BSSL_CHECK(ctx[0].field == ctx[2].field); \
+  BSSL_CHECK(ctx[0].field == ctx[3].field);
+  FOR_COMMON_FIELDS(MUST_BE_EQUAL)
+#undef MUST_BE_EQUAL
+
+  if (ctx->phase == boringssl_keccak_phase_absorb) {
+    keccak_finalize_x4(ctx);
+    ctx->phase = boringssl_keccak_phase_squeeze;
+  }
+
+  // Accessing `ctx->state` as a `uint8_t*` is allowed by strict aliasing
+  // because we require `uint8_t` to be a character type.
+  uint8_t *optr[4] = {outs[0], outs[1], outs[2], outs[3]};
+  while (out_len) {
+    if (ctx->squeeze_offset == ctx->rate_bytes) {
+      keccak_f_x4(ctx[0].state, ctx[1].state, ctx[2].state, ctx[3].state);
+      ctx->squeeze_offset = 0;
+    }
+
+    size_t remaining = ctx->rate_bytes - ctx->squeeze_offset;
+    size_t todo = out_len;
+    if (todo > remaining) {
+      todo = remaining;
+    }
+    for (size_t i = 0; i < 4; ++i) {
+      const uint8_t *state_bytes = (const uint8_t *)ctx[i].state;
+      OPENSSL_memcpy(optr[i], &state_bytes[ctx->squeeze_offset], todo);
+      optr[i] += todo;
+    }
+    out_len -= todo;
+    ctx->squeeze_offset += todo;
+  }
+
+#define COPY_FIELD_VALUE(field) \
+  ctx[1].field = ctx[0].field;  \
+  ctx[2].field = ctx[0].field;  \
+  ctx[3].field = ctx[0].field;
+  FOR_COMMON_FIELDS(COPY_FIELD_VALUE)
+#undef COPY_FIELD_VALUE
+
+#undef FOR_COMMON_FIELDS
+}
+#endif  // HAVE_KECCAK_X4
diff --git a/crypto/fipsmodule/keccak/keccak_test.cc b/crypto/fipsmodule/keccak/keccak_test.cc
index 2501765..b20d587 100644
--- a/crypto/fipsmodule/keccak/keccak_test.cc
+++ b/crypto/fipsmodule/keccak/keccak_test.cc
@@ -107,6 +107,70 @@
   }
 #endif
 
+#if defined(HAVE_KECCAK_X4) && !defined(BORINGSSL_SHARED_LIBRARY)
+  if (BORINGSSL_have_keccak_x4() && input.size() > 0 && input.size() < 72) {
+    uint8_t noise[72] = {static_cast<uint8_t>(input[0] ^ 0xFF), 0};
+    const uint8_t *input_first[4] = {input.data(), noise, noise, noise};
+    const uint8_t *input_last[4] = {noise, noise, noise, input.data()};
+
+    uint8_t sha3_256_digests_buf[4][32];
+    uint8_t *sha3_256_digests[4] = {sha3_256_digests_buf[0],
+                                    sha3_256_digests_buf[1],
+                                    sha3_256_digests_buf[2],
+                                    sha3_256_digests_buf[3]};
+    BORINGSSL_keccak_short_x4(sha3_256_digests, 32, input_first, input.size(),
+                              boringssl_sha3_256);
+    EXPECT_EQ(Bytes(sha3_256_expected), Bytes(sha3_256_digests[0], 32));
+    EXPECT_NE(Bytes(sha3_256_expected), Bytes(sha3_256_digests[3], 32));
+    BORINGSSL_keccak_short_x4(sha3_256_digests, 32, input_last, input.size(),
+                              boringssl_sha3_256);
+    EXPECT_NE(Bytes(sha3_256_expected), Bytes(sha3_256_digests[0], 32));
+    EXPECT_EQ(Bytes(sha3_256_expected), Bytes(sha3_256_digests[3], 32));
+
+    uint8_t sha3_512_digests_buf[4][64];
+    uint8_t *sha3_512_digests[4] = {sha3_512_digests_buf[0],
+                                    sha3_512_digests_buf[1],
+                                    sha3_512_digests_buf[2],
+                                    sha3_512_digests_buf[3]};
+    BORINGSSL_keccak_short_x4(sha3_512_digests, 64, input_first, input.size(),
+                              boringssl_sha3_512);
+    EXPECT_EQ(Bytes(sha3_512_expected), Bytes(sha3_512_digests[0], 64));
+    EXPECT_NE(Bytes(sha3_512_expected), Bytes(sha3_512_digests[3], 64));
+    BORINGSSL_keccak_short_x4(sha3_512_digests, 64, input_last, input.size(),
+                              boringssl_sha3_512);
+    EXPECT_NE(Bytes(sha3_512_expected), Bytes(sha3_512_digests[0], 64));
+    EXPECT_EQ(Bytes(sha3_512_expected), Bytes(sha3_512_digests[3], 64));
+
+    uint8_t shake128_outputs_buf[4][512];
+    uint8_t *shake128_outputs[4] = {shake128_outputs_buf[0],
+                                    shake128_outputs_buf[1],
+                                    shake128_outputs_buf[2],
+                                    shake128_outputs_buf[3]};
+    BORINGSSL_keccak_short_x4(shake128_outputs, 512, input_first, input.size(),
+                              boringssl_shake128);
+    EXPECT_EQ(Bytes(shake128_expected), Bytes(shake128_outputs[0], 512));
+    EXPECT_NE(Bytes(shake128_expected), Bytes(shake128_outputs[3], 512));
+    BORINGSSL_keccak_short_x4(shake128_outputs, 512, input_last, input.size(),
+                              boringssl_shake128);
+    EXPECT_NE(Bytes(shake128_expected), Bytes(shake128_outputs[0], 512));
+    EXPECT_EQ(Bytes(shake128_expected), Bytes(shake128_outputs[3], 512));
+
+    uint8_t shake256_outputs_buf[4][512];
+    uint8_t *shake256_outputs[4] = {shake256_outputs_buf[0],
+                                    shake256_outputs_buf[1],
+                                    shake256_outputs_buf[2],
+                                    shake256_outputs_buf[3]};
+    BORINGSSL_keccak_short_x4(shake256_outputs, 512, input_first, input.size(),
+                              boringssl_shake256);
+    EXPECT_EQ(Bytes(shake256_expected), Bytes(shake256_outputs[0], 512));
+    EXPECT_NE(Bytes(shake256_expected), Bytes(shake256_outputs[3], 512));
+    BORINGSSL_keccak_short_x4(shake256_outputs, 512, input_last, input.size(),
+                              boringssl_shake256);
+    EXPECT_NE(Bytes(shake256_expected), Bytes(shake256_outputs[0], 512));
+    EXPECT_EQ(Bytes(shake256_expected), Bytes(shake256_outputs[3], 512));
+  }
+#endif
+
   struct BORINGSSL_keccak_st ctx;
 
   // Single-pass absorb/squeeze.
diff --git a/crypto/fipsmodule/mlkem/mlkem.cc.inc b/crypto/fipsmodule/mlkem/mlkem.cc.inc
index 98e0055..f0ef633 100644
--- a/crypto/fipsmodule/mlkem/mlkem.cc.inc
+++ b/crypto/fipsmodule/mlkem/mlkem.cc.inc
@@ -62,6 +62,13 @@
 }
 #endif
 
+#if defined(HAVE_KECCAK_X4)
+inline void prf_x4(uint8_t *outs[4], size_t out_len,
+                   const uint8_t *ins[4] /* 33 each */) {
+  BORINGSSL_keccak_short_x4(outs, out_len, ins, 33, boringssl_shake256);
+}
+#endif
+
 // Section 4.1
 void hash_h(uint8_t out[32], const uint8_t *in, size_t len) {
   BORINGSSL_keccak(out, 32, in, len, boringssl_sha3_256);
@@ -445,6 +452,29 @@
 }
 #endif
 
+#if defined(HAVE_KECCAK_X4)
+inline void scalar_from_keccak_vartime_x4(scalar *out[4],
+                                          BORINGSSL_keccak_st keccak_ctx[]) {
+  for (int i = 0; i < 4; i++) {
+    assert(keccak_ctx[i].squeeze_offset == 0);
+    assert(keccak_ctx[i].rate_bytes == 168);
+  }
+  static_assert(168 % 3 == 0, "block and coefficient boundaries do not align");
+
+  int done[4] = {0, 0, 0, 0};
+  while (done[0] < DEGREE || done[1] < DEGREE || done[2] < DEGREE ||
+         done[3] < DEGREE) {
+    uint8_t block[4][168];
+    uint8_t *blocks[] = {block[0], block[1], block[2], block[3]};
+    BORINGSSL_keccak_squeeze_x4(keccak_ctx, blocks, sizeof(block[0]));
+    scalar_from_keccak_block_vartime(out[0], &done[0], block[0]);
+    scalar_from_keccak_block_vartime(out[1], &done[1], block[1]);
+    scalar_from_keccak_block_vartime(out[2], &done[2], block[2]);
+    scalar_from_keccak_block_vartime(out[3], &done[3], block[3]);
+  }
+}
+#endif
+
 static void scalar_centered_binomial_distribution_eta_2(
     scalar *out, const uint8_t entropy[128]) {
   for (int i = 0; i < DEGREE; i += 2) {
@@ -504,6 +534,29 @@
 }
 #endif
 
+#if defined(HAVE_KECCAK_X4)
+// scalar_centered_binomial_distribution_eta_2_with_prf_x4 computes four blocks
+// of scalar_centered_binomial_distribution_eta_2_with_prf.
+static void scalar_centered_binomial_distribution_eta_2_with_prf_x4(
+    scalar out[4], uint8_t input[33]) {
+  uint8_t entropy[4][128];
+  static_assert(sizeof(entropy[0]) == 2 * /*kEta=*/2 * DEGREE / 8);
+
+  uint8_t inputs[4][33];
+  for (int i = 0; i < 4; i++) {
+    OPENSSL_memcpy(inputs[i], input, 33);
+    inputs[i][32] += i;
+  }
+  uint8_t *entropies[] = {entropy[0], entropy[1], entropy[2], entropy[3]};
+  const uint8_t *input_ptrs[4] = {inputs[0], inputs[1], inputs[2], inputs[3]};
+  prf_x4(entropies, sizeof(entropy[0]), input_ptrs);
+
+  for (int k = 0; k < 4; k++) {
+    scalar_centered_binomial_distribution_eta_2(&out[k], entropy[k]);
+  }
+}
+#endif
+
 // Generates a secret vector by using
 // `scalar_centered_binomial_distribution_eta_2_with_prf`, using the given seed
 // appending and incrementing `counter` for entry of the vector.
@@ -514,6 +567,16 @@
   OPENSSL_memcpy(input, seed, 32);
 
   int i = 0;
+#if defined(HAVE_KECCAK_X4)
+  if (BORINGSSL_have_keccak_x4()) {
+    for (; i + 4 <= RANK; i += 4) {
+      input[32] = *counter;
+      *counter += 4;
+      scalar_centered_binomial_distribution_eta_2_with_prf_x4(&out->v[i],
+                                                              input);
+    }
+  }
+#endif
 #if defined(HAVE_KECCAK_X2)
   for (; i + 2 <= RANK; i += 2) {
     input[32] = *counter;
@@ -540,7 +603,7 @@
 
 #if defined(HAVE_KECCAK_X2)
 template <int RANK>
-void matrix_expand_step_x2(matrix<RANK> out[2], uint8_t input[34], int i) {
+void matrix_expand_step_x2(matrix<RANK> *out, uint8_t input[34], int i) {
   input[32] = i / RANK;
   input[33] = i % RANK;
   uint8_t input1[34];
@@ -558,6 +621,26 @@
 }
 #endif
 
+#if defined(HAVE_KECCAK_X4)
+template <int RANK>
+void matrix_expand_step_x4(matrix<RANK> *out, uint8_t input[34], int i) {
+  uint8_t inputs[4][34];
+  BORINGSSL_keccak_st keccak_ctx[4];
+  for (int k = 0; k < 4; k++) {
+    OPENSSL_memcpy(inputs[k], input, 32);
+    inputs[k][32] = (i + k) / RANK;
+    inputs[k][33] = (i + k) % RANK;
+    BORINGSSL_keccak_init(&keccak_ctx[k], boringssl_shake128);
+    BORINGSSL_keccak_absorb(&keccak_ctx[k], inputs[k], 34);
+  }
+  scalar *outs[] = {&out->v[i / RANK][i % RANK],
+                    &out->v[(i + 1) / RANK][(i + 1) % RANK],
+                    &out->v[(i + 2) / RANK][(i + 2) % RANK],
+                    &out->v[(i + 3) / RANK][(i + 3) % RANK]};
+  scalar_from_keccak_vartime_x4(outs, keccak_ctx);
+}
+#endif
+
 // Expands the matrix of a seed for key generation and for encaps-CPA.
 template <int RANK>
 void matrix_expand(matrix<RANK> *out, const uint8_t rho[32]) {
@@ -565,9 +648,16 @@
   OPENSSL_memcpy(input, rho, 32);
 
   int i = 0;
+#if defined(HAVE_KECCAK_X4)
+  if (BORINGSSL_have_keccak_x4()) {
+    for (; i + 4 <= RANK * RANK; i += 4) {
+      matrix_expand_step_x4<RANK>(out, input, i);
+    }
+  }
+#endif
 #if defined(HAVE_KECCAK_X2)
   for (; i + 2 <= RANK * RANK; i += 2) {
-    matrix_expand_step_x2(out, input, i);
+    matrix_expand_step_x2<RANK>(out, input, i);
   }
 #endif
 
diff --git a/crypto/internal.h b/crypto/internal.h
index dcc109a..dda1933 100644
--- a/crypto/internal.h
+++ b/crypto/internal.h
@@ -143,6 +143,15 @@
 #define OPENSSL_ATTR_CONST
 #endif
 
+// OPENSSL_ATTR_ALWAYS_INLINE tells the compiler to always attempt inlining the
+// function. This has the beneficial side effects that it inherits attributes
+// for code generation (e.g. the target attribute) from its callee.
+#if defined(__GNUC__) || defined(__clang__)
+#define OPENSSL_ATTR_ALWAYS_INLINE __attribute__((always_inline))
+#else
+#define OPENSSL_ATTR_ALWAYS_INLINE
+#endif
+
 #if defined(BORINGSSL_MALLOC_FAILURE_TESTING)
 // OPENSSL_reset_malloc_counter_for_testing, when malloc testing is enabled,
 // resets the internal malloc counter, to simulate further malloc failures. This