Implement a vectorized "double Keccak".

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

Bug: 503700354
Change-Id: If3faab601d5d30a41f54cb7306d04a586a6a6964
Reviewed-on: https://boringssl-review.googlesource.com/c/boringssl/+/92829
Commit-Queue: David Benjamin <davidben@google.com>
Reviewed-by: David Benjamin <davidben@google.com>
Auto-Submit: Rudolf Polzer <rpolzer@google.com>
Presubmit-BoringSSL-Verified: boringssl-scoped@luci-project-accounts.iam.gserviceaccount.com <boringssl-scoped@luci-project-accounts.iam.gserviceaccount.com>
diff --git a/crypto/fipsmodule/keccak/internal.h b/crypto/fipsmodule/keccak/internal.h
index c7c3f94..ff10c44 100644
--- a/crypto/fipsmodule/keccak/internal.h
+++ b/crypto/fipsmodule/keccak/internal.h
@@ -70,6 +70,27 @@
 OPENSSL_EXPORT void BORINGSSL_keccak_squeeze(struct BORINGSSL_keccak_st *ctx,
                                              uint8_t *out, size_t out_len);
 
+#if defined(__has_attribute)
+#if __has_attribute(vector_size)
+#define HAVE_KECCAK_X2
+#endif  // vector_size
+#endif  // __has_attribute
+
+#if defined(HAVE_KECCAK_X2)
+// BORINGSSL_keccak_squeeze_x2 performs BORINGSSL_keccak_squeeze in parallel
+// with two 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_x2(
+    struct BORINGSSL_keccak_st ctx[2], uint8_t *outs[2], size_t out_len);
+
+// BORINGSSL_keccak_short_x2 performs BORINGSSL_keccak in parallel on two
+// same-length strings with same-length outputs. |in_len| must be less than 72
+// (or actually |rate_bytes|).
+OPENSSL_EXPORT void BORINGSSL_keccak_short_x2(
+    uint8_t *outs[2], size_t out_len, const uint8_t *ins[2], 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 32be718..5008cd4 100644
--- a/crypto/fipsmodule/keccak/keccak.cc.inc
+++ b/crypto/fipsmodule/keccak/keccak.cc.inc
@@ -26,18 +26,22 @@
 // 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.
-static void keccak_f(uint64_t state[25]) {
+//
+// 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]) {
   static const int kNumRounds = 24;
   for (int round = 0; round < kNumRounds; round++) {
     // θ step
-    uint64_t c[5];
+    U64 c[5];
     for (int x = 0; x < 5; x++) {
       c[x] = state[x] ^ state[x + 5] ^ state[x + 10] ^ state[x + 15] ^
              state[x + 20];
     }
 
     for (int x = 0; x < 5; x++) {
-      const uint64_t d = c[(x + 4) % 5] ^ CRYPTO_rotl_u64(c[(x + 1) % 5], 1);
+      const U64 d = c[(x + 4) % 5] ^ rotl(c[(x + 1) % 5], 1);
       for (int y = 0; y < 5; y++) {
         state[y * 5 + x] ^= d;
       }
@@ -58,12 +62,12 @@
     // and the sequence will repeat. All that remains is to handle the element
     // at (0, 0), but the rotation for that element is zero, and it goes to (0,
     // 0), so we can ignore it.
-    uint64_t prev_value = state[1];
-#define PI_RHO_STEP(index, rotation)                              \
-  do {                                                            \
-    const uint64_t value = CRYPTO_rotl_u64(prev_value, rotation); \
-    prev_value = state[index];                                    \
-    state[index] = value;                                         \
+    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;                         \
   } while (0)
 
     PI_RHO_STEP(10, 1);
@@ -96,8 +100,8 @@
     // χ step
     for (int y = 0; y < 5; y++) {
       const int row_index = 5 * y;
-      const uint64_t orig_x0 = state[row_index];
-      const uint64_t orig_x1 = state[row_index + 1];
+      const U64 orig_x0 = state[row_index];
+      const U64 orig_x1 = state[row_index + 1];
       state[row_index] ^= ~orig_x1 & state[row_index + 2];
       state[row_index + 1] ^= ~state[row_index + 2] & state[row_index + 3];
       state[row_index + 2] ^= ~state[row_index + 3] & state[row_index + 4];
@@ -127,6 +131,26 @@
   }
 }
 
+#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 void keccak_f_x2(uint64_t state0[25], uint64_t state1[25]) {
+  v2u64 s[25];
+  for (int i = 0; i < 25; i++) {
+    s[i] = (v2u64){state0[i], state1[i]};
+  }
+  keccak_f<v2u64, v_rotl>(s);
+  for (int i = 0; i < 25; i++) {
+    state0[i] = s[i][0];
+    state1[i] = s[i][1];
+  }
+}
+#endif
+
 static void keccak_init(struct BORINGSSL_keccak_st *ctx,
                         enum boringssl_keccak_config_t config) {
   size_t required_out_len;
@@ -164,7 +188,7 @@
                             size_t in_len,
                             enum boringssl_keccak_config_t config) {
   struct BORINGSSL_keccak_st ctx;
-  keccak_init(&ctx, config);
+  BORINGSSL_keccak_init(&ctx, config);
   if (ctx.required_out_len != 0 && out_len != ctx.required_out_len) {
     abort();
   }
@@ -172,6 +196,29 @@
   BORINGSSL_keccak_squeeze(&ctx, out, out_len);
 }
 
+#if defined(HAVE_KECCAK_X2)
+void bssl::BORINGSSL_keccak_short_x2(uint8_t *outs[2], size_t out_len,
+                                     const uint8_t *ins[2], size_t in_len,
+                                     enum boringssl_keccak_config_t config) {
+  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();
+    }
+
+    // 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_x2(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,29 +273,41 @@
   ctx->absorb_offset = in_len;
 }
 
-static void keccak_finalize(struct BORINGSSL_keccak_st *ctx) {
-  uint8_t terminator;
+static uint8_t keccak_terminator(struct BORINGSSL_keccak_st *ctx) {
   switch (ctx->config) {
     case boringssl_sha3_256:
     case boringssl_sha3_512:
-      terminator = 0x06;
-      break;
+      return 0x06;
     case boringssl_shake128:
     case boringssl_shake256:
-      terminator = 0x1f;
-      break;
+      return 0x1f;
     default:
       abort();
   }
+}
 
+static void keccak_finalize(struct BORINGSSL_keccak_st *ctx) {
   // 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->state;
-  state_bytes[ctx->absorb_offset] ^= terminator;
+  state_bytes[ctx->absorb_offset] ^= keccak_terminator(ctx);
   state_bytes[ctx->rate_bytes - 1] ^= 0x80;
   keccak_f(ctx->state);
 }
 
+#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.
+    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;
+  }
+  keccak_f_x2(ctx[0].state, ctx[1].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 &&
@@ -284,3 +343,65 @@
     ctx->squeeze_offset += todo;
   }
 }
+
+#if defined(HAVE_KECCAK_X2)
+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();
+    }
+  }
+
+  // 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)                  \
+  MACRO(absorb_offset)           \
+  MACRO(squeeze_offset)          \
+  MACRO(rate_bytes)
+
+#define MUST_BE_EQUAL(field) BSSL_CHECK(ctx[0].field == ctx[1].field);
+  FOR_COMMON_FIELDS(MUST_BE_EQUAL)
+#undef MUST_BE_EQUAL
+
+  if (ctx->phase == boringssl_keccak_phase_absorb) {
+    keccak_finalize_x2(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[2] = {outs[0], outs[1]};
+  while (out_len) {
+    if (ctx->squeeze_offset == ctx->rate_bytes) {
+      keccak_f_x2(ctx[0].state, ctx[1].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 < 2; ++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;
+  FOR_COMMON_FIELDS(COPY_FIELD_VALUE)
+#undef COPY_FIELD_VALUE
+
+#undef FOR_COMMON_FIELDS
+}
+#endif  // HAVE_KECCAK_X2
diff --git a/crypto/fipsmodule/keccak/keccak_test.cc b/crypto/fipsmodule/keccak/keccak_test.cc
index d1fcd51..2501765 100644
--- a/crypto/fipsmodule/keccak/keccak_test.cc
+++ b/crypto/fipsmodule/keccak/keccak_test.cc
@@ -55,6 +55,58 @@
   EXPECT_EQ(Bytes(shake128_expected), Bytes(shake128_output));
   EXPECT_EQ(Bytes(shake256_expected), Bytes(shake256_output));
 
+#if defined(HAVE_KECCAK_X2)
+  if (input.size() > 0 && input.size() < 72) {
+    uint8_t noise[72] = {static_cast<uint8_t>(~input[0]), 0};
+    const uint8_t *input_first[2] = {input.data(), noise};
+    const uint8_t *input_last[2] = {noise, input.data()};
+
+    uint8_t sha3_256_digest1[32];
+    uint8_t *sha3_256_digests[2] = {sha3_256_digest, sha3_256_digest1};
+    BORINGSSL_keccak_short_x2(sha3_256_digests, sizeof(sha3_256_digest),
+                              input_first, input.size(), boringssl_sha3_256);
+    EXPECT_EQ(Bytes(sha3_256_expected), Bytes(sha3_256_digest));
+    EXPECT_NE(Bytes(sha3_256_expected), Bytes(sha3_256_digest1));
+    BORINGSSL_keccak_short_x2(sha3_256_digests, sizeof(sha3_256_digest),
+                              input_last, input.size(), boringssl_sha3_256);
+    EXPECT_NE(Bytes(sha3_256_expected), Bytes(sha3_256_digest));
+    EXPECT_EQ(Bytes(sha3_256_expected), Bytes(sha3_256_digest1));
+
+    uint8_t sha3_512_digest1[64];
+    uint8_t *sha3_512_digests[2] = {sha3_512_digest, sha3_512_digest1};
+    BORINGSSL_keccak_short_x2(sha3_512_digests, sizeof(sha3_512_digest),
+                              input_first, input.size(), boringssl_sha3_512);
+    EXPECT_EQ(Bytes(sha3_512_expected), Bytes(sha3_512_digest));
+    EXPECT_NE(Bytes(sha3_512_expected), Bytes(sha3_512_digest1));
+    BORINGSSL_keccak_short_x2(sha3_512_digests, sizeof(sha3_512_digest),
+                              input_last, input.size(), boringssl_sha3_512);
+    EXPECT_NE(Bytes(sha3_512_expected), Bytes(sha3_512_digest));
+    EXPECT_EQ(Bytes(sha3_512_expected), Bytes(sha3_512_digest1));
+
+    uint8_t shake128_output1[512];
+    uint8_t *shake128_outputs[2] = {shake128_output, shake128_output1};
+    BORINGSSL_keccak_short_x2(shake128_outputs, sizeof(shake128_output),
+                              input_first, input.size(), boringssl_shake128);
+    EXPECT_EQ(Bytes(shake128_expected), Bytes(shake128_output));
+    EXPECT_NE(Bytes(shake128_expected), Bytes(shake128_output1));
+    BORINGSSL_keccak_short_x2(shake128_outputs, sizeof(shake128_output),
+                              input_last, input.size(), boringssl_shake128);
+    EXPECT_NE(Bytes(shake128_expected), Bytes(shake128_output));
+    EXPECT_EQ(Bytes(shake128_expected), Bytes(shake128_output1));
+
+    uint8_t shake256_output1[512];
+    uint8_t *shake256_outputs[2] = {shake256_output, shake256_output1};
+    BORINGSSL_keccak_short_x2(shake256_outputs, sizeof(shake256_output),
+                              input_first, input.size(), boringssl_shake256);
+    EXPECT_EQ(Bytes(shake256_expected), Bytes(shake256_output));
+    EXPECT_NE(Bytes(shake256_expected), Bytes(shake256_output1));
+    BORINGSSL_keccak_short_x2(shake256_outputs, sizeof(shake256_output),
+                              input_last, input.size(), boringssl_shake256);
+    EXPECT_NE(Bytes(shake256_expected), Bytes(shake256_output));
+    EXPECT_EQ(Bytes(shake256_expected), Bytes(shake256_output1));
+  }
+#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 ecb071e..540b5b3 100644
--- a/crypto/fipsmodule/mlkem/mlkem.cc.inc
+++ b/crypto/fipsmodule/mlkem/mlkem.cc.inc
@@ -55,6 +55,13 @@
   BORINGSSL_keccak(out, out_len, in, 33, boringssl_shake256);
 }
 
+#if defined(HAVE_KECCAK_X2)
+inline void prf_x2(uint8_t *outs[2], size_t out_len,
+                   const uint8_t *ins[2] /* 33 each */) {
+  BORINGSSL_keccak_short_x2(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);
@@ -387,7 +394,21 @@
   }
 }
 
-// Algorithm 6 from the spec. Rejection samples a Keccak stream to get
+inline void scalar_from_keccak_block_vartime(scalar *out, int *done,
+                                             const uint8_t block[168]) {
+  for (size_t i = 0; i < 168 && *done < DEGREE; i += 3) {
+    uint16_t d1 = block[i] + 256 * (block[i + 1] % 16);
+    uint16_t d2 = block[i + 1] / 16 + 16 * block[i + 2];
+    if (d1 < kPrime) {
+      out->c[(*done)++] = d1;
+    }
+    if (d2 < kPrime && *done < DEGREE) {
+      out->c[(*done)++] = d2;
+    }
+  }
+}
+
+// Algorithm 7 from the spec. Rejection samples a Keccak stream to get
 // uniformly distributed elements. This is used for matrix expansion and only
 // operates on public inputs.
 inline void scalar_from_keccak_vartime(scalar *out,
@@ -400,31 +421,32 @@
   while (done < DEGREE) {
     uint8_t block[168];
     BORINGSSL_keccak_squeeze(keccak_ctx, block, sizeof(block));
-    for (size_t i = 0; i < sizeof(block) && done < DEGREE; i += 3) {
-      uint16_t d1 = block[i] + 256 * (block[i + 1] % 16);
-      uint16_t d2 = block[i + 1] / 16 + 16 * block[i + 2];
-      if (d1 < kPrime) {
-        out->c[done++] = d1;
-      }
-      if (d2 < kPrime && done < DEGREE) {
-        out->c[done++] = d2;
-      }
-    }
+    scalar_from_keccak_block_vartime(out, &done, block);
   }
 }
 
-// Algorithm 7 from the spec, with eta fixed to two and the PRF call
-// included. Creates binominally distributed elements by sampling 2*|eta| bits,
-// and setting the coefficient to the count of the first bits minus the count of
-// the second bits, resulting in a centered binomial distribution. Since eta is
-// two this gives -2/2 with a probability of 1/16, -1/1 with probability 1/4,
-// and 0 with probability 3/8.
-void scalar_centered_binomial_distribution_eta_2_with_prf(
-    scalar *out, const uint8_t input[33]) {
-  uint8_t entropy[128];
-  static_assert(sizeof(entropy) == 2 * /*kEta=*/2 * DEGREE / 8);
-  prf(entropy, sizeof(entropy), input);
+#if defined(HAVE_KECCAK_X2)
+inline void scalar_from_keccak_vartime_x2(scalar *out[2],
+                                          BORINGSSL_keccak_st keccak_ctx[2]) {
+  assert(keccak_ctx[0].squeeze_offset == 0);
+  assert(keccak_ctx[0].rate_bytes == 168);
+  assert(keccak_ctx[1].squeeze_offset == 0);
+  assert(keccak_ctx[1].rate_bytes == 168);
+  static_assert(168 % 3 == 0, "block and coefficient boundaries do not align");
 
+  int done[2] = {0, 0};
+  while (done[0] < DEGREE || done[1] < DEGREE) {
+    uint8_t block[2][168];
+    uint8_t *blocks[] = {block[0], block[1]};
+    BORINGSSL_keccak_squeeze_x2(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]);
+  }
+}
+#endif
+
+static void scalar_centered_binomial_distribution_eta_2(
+    scalar *out, const uint8_t entropy[128]) {
   for (int i = 0; i < DEGREE; i += 2) {
     uint8_t byte = entropy[i / 2];
 
@@ -446,6 +468,42 @@
   }
 }
 
+// Algorithm 8 from the spec, with eta fixed to two and the PRF call
+// included. Creates binominally distributed elements by sampling 2*|eta| bits,
+// and setting the coefficient to the count of the first bits minus the count of
+// the second bits, resulting in a centered binomial distribution. Since eta is
+// two this gives -2/2 with a probability of 1/16, -1/1 with probability 1/4,
+// and 0 with probability 3/8.
+void scalar_centered_binomial_distribution_eta_2_with_prf(
+    scalar *out, const uint8_t input[33]) {
+  uint8_t entropy[128];
+  static_assert(sizeof(entropy) == 2 * /*kEta=*/2 * DEGREE / 8);
+  prf(entropy, sizeof(entropy), input);
+  scalar_centered_binomial_distribution_eta_2(out, entropy);
+}
+
+#if defined(HAVE_KECCAK_X2)
+// scalar_centered_binomial_distribution_eta_2_with_prf_x2 computes two blocks
+// of scalar_centered_binomial_distribution_eta_2_with_prf.
+static void scalar_centered_binomial_distribution_eta_2_with_prf_x2(
+    scalar out[2], uint8_t input[33]) {
+  uint8_t entropy[2][128];
+  static_assert(sizeof(entropy[0]) == 2 * /*kEta=*/2 * DEGREE / 8);
+
+  uint8_t input1[33];
+  OPENSSL_memcpy(input1, input, 33);
+  ++input1[32];  // Counter.
+
+  uint8_t *entropies[] = {entropy[0], entropy[1]};
+  const uint8_t *inputs[2] = {input, input1};
+  prf_x2(entropies, sizeof(entropy[0]), inputs);
+
+  for (int k = 0; k < 2; 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.
@@ -454,26 +512,67 @@
                                   const uint8_t seed[32]) {
   uint8_t input[33];
   OPENSSL_memcpy(input, seed, 32);
-  for (int i = 0; i < RANK; i++) {
+
+  int i = 0;
+#if defined(HAVE_KECCAK_X2)
+  for (; i + 2 <= RANK; i += 2) {
+    input[32] = *counter;
+    *counter += 2;
+    scalar_centered_binomial_distribution_eta_2_with_prf_x2(&out->v[i], input);
+  }
+#endif
+
+  for (; i < RANK; i++) {
     input[32] = (*counter)++;
     scalar_centered_binomial_distribution_eta_2_with_prf(&out->v[i], input);
   }
 }
 
+template <int RANK>
+void matrix_expand_step(scalar *out, uint8_t input[34], int i) {
+  input[32] = i / RANK;
+  input[33] = i % RANK;
+  BORINGSSL_keccak_st keccak_ctx;
+  BORINGSSL_keccak_init(&keccak_ctx, boringssl_shake128);
+  BORINGSSL_keccak_absorb(&keccak_ctx, input, 34);
+  scalar_from_keccak_vartime(out, &keccak_ctx);
+}
+
+#if defined(HAVE_KECCAK_X2)
+template <int RANK>
+void matrix_expand_step_x2(matrix<RANK> out[2], uint8_t input[34], int i) {
+  input[32] = i / RANK;
+  input[33] = i % RANK;
+  uint8_t input1[34];
+  OPENSSL_memcpy(input1, input, 32);
+  input1[32] = (i + 1) / RANK;
+  input1[33] = (i + 1) % RANK;
+  BORINGSSL_keccak_st keccak_ctx[2];
+  BORINGSSL_keccak_init(&keccak_ctx[0], boringssl_shake128);
+  BORINGSSL_keccak_init(&keccak_ctx[1], boringssl_shake128);
+  BORINGSSL_keccak_absorb(&keccak_ctx[0], input, 34);
+  BORINGSSL_keccak_absorb(&keccak_ctx[1], input1, 34);
+  scalar *outs[] = {&out->v[i / RANK][i % RANK],
+                    &out->v[(i + 1) / RANK][(i + 1) % RANK]};
+  scalar_from_keccak_vartime_x2(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]) {
   uint8_t input[34];
   OPENSSL_memcpy(input, rho, 32);
-  for (int i = 0; i < RANK; i++) {
-    for (int j = 0; j < RANK; j++) {
-      input[32] = i;
-      input[33] = j;
-      BORINGSSL_keccak_st keccak_ctx;
-      BORINGSSL_keccak_init(&keccak_ctx, boringssl_shake128);
-      BORINGSSL_keccak_absorb(&keccak_ctx, input, sizeof(input));
-      scalar_from_keccak_vartime(&out->v[i][j], &keccak_ctx);
-    }
+
+  int i = 0;
+#if defined(HAVE_KECCAK_X2)
+  for (; i + 2 <= RANK * RANK; i += 2) {
+    matrix_expand_step_x2(out, input, i);
+  }
+#endif
+
+  for (; i < RANK * RANK; i++) {
+    matrix_expand_step<RANK>(&out->v[i / RANK][i % RANK], input, i);
   }
 }