ML-DSA: add Keccak x4 support.

This is basically a port of 8daf854545e60454740b3c2286de391bcf8b1685
from ML-KEM.

Only small improvements on aarch64, but massive improvements on x86-64.

aarch64 Apple M1 Pro:

Did 23153 BM_SpeedMLDSAKeyGen/ml_dsa_44/threads:1 operations (33067.7 ops/sec) [+2.5%]
Did 5817 BM_SpeedMLDSASign/ml_dsa_44/threads:1 operations (8109.4 ops/sec) [-0.2%]
Did 253549 BM_SpeedMLDSAParsePubKey/ml_dsa_44/threads:1 operations (357425.4 ops/sec) [-1.1%]
Did 26062 BM_SpeedMLDSAVerify/ml_dsa_44/threads:1 operations (38089.8 ops/sec) [+0.8%]
Did 26751 BM_SpeedMLDSAVerifyBadSignature/ml_dsa_44/threads:1 operations (38209.4 ops/sec) [+1.1%]
Did 11327 BM_SpeedMLDSAKeyGen/ml_dsa_65/threads:1 operations (16189.1 ops/sec) [+0.9%]
Did 3650 BM_SpeedMLDSASign/ml_dsa_65/threads:1 operations (5250.0 ops/sec) [+1.8%]
Did 185026 BM_SpeedMLDSAParsePubKey/ml_dsa_65/threads:1 operations (264028.7 ops/sec) [-0.5%]
Did 17160 BM_SpeedMLDSAVerify/ml_dsa_65/threads:1 operations (24510.6 ops/sec) [+4.2%]
Did 17186 BM_SpeedMLDSAVerifyBadSignature/ml_dsa_65/threads:1 operations (24555.3 ops/sec) [+1.7%]
Did 8932 BM_SpeedMLDSAKeyGen/ml_dsa_87/threads:1 operations (12802.3 ops/sec) [+1.7%]
Did 3165 BM_SpeedMLDSASign/ml_dsa_87/threads:1 operations (4381.6 ops/sec) [-0.1%]
Did 145645 BM_SpeedMLDSAParsePubKey/ml_dsa_87/threads:1 operations (208208.8 ops/sec) [+1.1%]
Did 10307 BM_SpeedMLDSAVerify/ml_dsa_87/threads:1 operations (14741.2 ops/sec) [+1.7%]
Did 10302 BM_SpeedMLDSAVerifyBadSignature/ml_dsa_87/threads:1 operations (14715.9 ops/sec) [+3.9%]

x86-64 AMD EPYC 7B13:

Did 13796 BM_SpeedMLDSAKeyGen/ml_dsa_44/threads:1 operations (19726.2 ops/sec) [+19.3%]
Did 3675 BM_SpeedMLDSASign/ml_dsa_44/threads:1 operations (5120.6 ops/sec) [+1.4%]
Did 155308 BM_SpeedMLDSAParsePubKey/ml_dsa_44/threads:1 operations (221784.4 ops/sec) [+1.1%]
Did 16906 BM_SpeedMLDSAVerify/ml_dsa_44/threads:1 operations (24154.9 ops/sec) [+21.3%]
Did 16865 BM_SpeedMLDSAVerifyBadSignature/ml_dsa_44/threads:1 operations (24128.6 ops/sec) [+21.4%]
Did 7246 BM_SpeedMLDSAKeyGen/ml_dsa_65/threads:1 operations (10344.9 ops/sec) [+17.9%]
Did 2369 BM_SpeedMLDSASign/ml_dsa_65/threads:1 operations (3461.7 ops/sec) [+11.5%]
Did 103849 BM_SpeedMLDSAParsePubKey/ml_dsa_65/threads:1 operations (148677.9 ops/sec) [+0.2%]
Did 10579 BM_SpeedMLDSAVerify/ml_dsa_65/threads:1 operations (15057.3 ops/sec) [+22.9%]
Did 10711 BM_SpeedMLDSAVerifyBadSignature/ml_dsa_65/threads:1 operations (15178.8 ops/sec) [+23.9%]
Did 5502 BM_SpeedMLDSAKeyGen/ml_dsa_87/threads:1 operations (7879.5 ops/sec) [+24.9%]
Did 2044 BM_SpeedMLDSASign/ml_dsa_87/threads:1 operations (3009.9 ops/sec) [+13.4%]
Did 77857 BM_SpeedMLDSAParsePubKey/ml_dsa_87/threads:1 operations (111054.2 ops/sec) [+0.7%]
Did 6569 BM_SpeedMLDSAVerify/ml_dsa_87/threads:1 operations (9391.4 ops/sec) [+28.1%]
Did 6445 BM_SpeedMLDSAVerifyBadSignature/ml_dsa_87/threads:1 operations (9275.8 ops/sec) [+26.7%]

Bug: 503700354
Change-Id: I09c1ef686314a05586ecefe1382239556a6a6964
Reviewed-on: https://boringssl-review.googlesource.com/c/boringssl/+/102127
Commit-Queue: David Benjamin <davidben@google.com>
Presubmit-BoringSSL-Verified: boringssl-scoped@luci-project-accounts.iam.gserviceaccount.com <boringssl-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: David Benjamin <davidben@google.com>
Auto-Submit: Rudolf Polzer <rpolzer@google.com>
diff --git a/crypto/fipsmodule/mldsa/mldsa.cc.inc b/crypto/fipsmodule/mldsa/mldsa.cc.inc
index 9d6a4dc..9879738 100644
--- a/crypto/fipsmodule/mldsa/mldsa.cc.inc
+++ b/crypto/fipsmodule/mldsa/mldsa.cc.inc
@@ -1243,6 +1243,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] < kDegree || done[1] < kDegree || done[2] < kDegree ||
+         done[3] < kDegree) {
+    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
+
 template <int ETA>
 inline bool coefficient_from_nibble(uint32_t nibble, uint32_t *result);
 
@@ -1325,6 +1348,29 @@
 }
 #endif
 
+#if defined(HAVE_KECCAK_X4)
+template <int ETA>
+inline void scalar_uniform_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 == 136);
+  }
+
+  int done[4] = {0, 0, 0, 0};
+  while (done[0] < kDegree || done[1] < kDegree || done[2] < kDegree ||
+         done[3] < kDegree) {
+    uint8_t block[4][136];
+    uint8_t *blocks[] = {block[0], block[1], block[2], block[3]};
+    BORINGSSL_keccak_squeeze_x4(keccak_ctx, blocks, sizeof(block[0]));
+    scalar_uniform_block<ETA>(out[0], &done[0], block[0]);
+    scalar_uniform_block<ETA>(out[1], &done[1], block[1]);
+    scalar_uniform_block<ETA>(out[2], &done[2], block[2]);
+    scalar_uniform_block<ETA>(out[3], &done[3], block[3]);
+  }
+}
+#endif
+
 // FIPS 204, Algorithm 34 (`ExpandMask`), but just a single step.
 template <int K>
 inline void scalar_sample_mask(scalar *out,
@@ -1421,12 +1467,41 @@
 }
 #endif
 
+#if defined(HAVE_KECCAK_X4)
+template <int K, int L>
+inline void matrix_expand_step_x4(matrix<K, L> *out,
+                                  const uint8_t rho[kRhoBytes], int index) {
+  BORINGSSL_keccak_st keccak_ctx[4];
+  for (int k = 0; k < 4; k++) {
+    BORINGSSL_keccak_init(&keccak_ctx[k], boringssl_shake128);
+    BORINGSSL_keccak_absorb(&keccak_ctx[k], rho, kRhoBytes);
+    uint8_t index_buf[2];
+    index_buf[0] = static_cast<uint8_t>((index + k) % L);
+    index_buf[1] = static_cast<uint8_t>((index + k) / L);
+    BORINGSSL_keccak_absorb(&keccak_ctx[k], index_buf, 2);
+  }
+
+  scalar *outs[] = {&out->v[index / L][index % L],
+                    &out->v[(index + 1) / L][(index + 1) % L],
+                    &out->v[(index + 2) / L][(index + 2) % L],
+                    &out->v[(index + 3) / L][(index + 3) % L]};
+  scalar_from_keccak_vartime_x4(outs, keccak_ctx);
+}
+#endif
+
 template <int K, int L>
 inline void matrix_expand(matrix<K, L> *out, const uint8_t rho[kRhoBytes]) {
   static_assert(K <= 0x100, "K must fit in 8 bits");
   static_assert(L <= 0x100, "L must fit in 8 bits");
 
   int index = 0;
+#if defined(HAVE_KECCAK_X4)
+  if (BORINGSSL_have_keccak_x4()) {
+    for (; index + 4 <= K * L; index += 4) {
+      matrix_expand_step_x4<K, L>(out, rho, index);
+    }
+  }
+#endif
 #if defined(HAVE_KECCAK_X2)
   for (; index + 2 <= K * L; index += 2) {
     matrix_expand_step_x2(out, rho, index);
@@ -1475,6 +1550,27 @@
 }
 #endif
 
+#if defined(HAVE_KECCAK_X4)
+template <int K>
+inline void vector_expand_short_step_x4(scalar *out0, scalar *out1,
+                                        scalar *out2, scalar *out3,
+                                        const uint8_t sigma[kSigmaBytes],
+                                        int index) {
+  BORINGSSL_keccak_st keccak_ctx[4];
+  for (int k = 0; k < 4; k++) {
+    BORINGSSL_keccak_init(&keccak_ctx[k], boringssl_shake256);
+    BORINGSSL_keccak_absorb(&keccak_ctx[k], sigma, kSigmaBytes);
+    uint8_t index_buf[2];
+    index_buf[0] = static_cast<uint8_t>(index + k);
+    index_buf[1] = 0;
+    BORINGSSL_keccak_absorb(&keccak_ctx[k], index_buf, 2);
+  }
+
+  scalar *outs[] = {out0, out1, out2, out3};
+  scalar_uniform_x4<eta<K>()>(outs, keccak_ctx);
+}
+#endif
+
 // FIPS 204, Algorithm 33 (`ExpandS`).
 template <int K, int L>
 inline void vector_expand_short(vector<L> *s1, vector<K> *s2,
@@ -1492,6 +1588,15 @@
   }
 
   int index = 0;
+#if defined(HAVE_KECCAK_X4)
+  if (BORINGSSL_have_keccak_x4()) {
+    for (; index + 4 <= K + L; index += 4) {
+      vector_expand_short_step_x4<K>(all_scalars[index], all_scalars[index + 1],
+                                     all_scalars[index + 2],
+                                     all_scalars[index + 3], sigma, index);
+    }
+  }
+#endif
 #if defined(HAVE_KECCAK_X2)
   for (; index + 2 <= K + L; index += 2) {
     vector_expand_short_step_x2<K>(all_scalars[index], all_scalars[index + 1],