Store ML-DSA's s1, s2, and t0 in NTT form

ML-DSA sign only uses those in NTT form, so this avoids doing a
conversion on every signature operation. The result is that signing gets
faster, but keygen gets slower, because we shift some of those
conversions to keygen time.

Keygen itself needs s1 in NTT form anyway, so keygen gains two NTTs
while signing loses three, so even single-use keys (1x keygen + 1x sign)
get faster.

This complicates the ridiculous semi-expanded key format, which encodes
the three vectors in non-NTT form, but since we only implement it for
ACVP testing, it doesn't matter if that needs extra operations.

Benchmarks on a AMD Ryzen Threadripper PRO 7945WX 12-Cores below. Note
the percentages are misleading because the denominator for keygen was
lower.

Benchmark                                                 Time             CPU      Time Old      Time New       CPU Old       CPU New
--------------------------------------------------------------------------------------------------------------------------------------
BM_SpeedMLDSAKeyGen/ml_dsa_44/threads:1                +0.1576         +0.1576         33140         38362         33137         38358
BM_SpeedMLDSASign/ml_dsa_44/threads:1                  -0.0852         -0.0852        136175        124570        136165        124559
BM_SpeedMLDSAKeyGen/ml_dsa_65/threads:1                +0.1320         +0.1319         63872         72300         63862         72289
BM_SpeedMLDSASign/ml_dsa_65/threads:1                  -0.0384         -0.0383        210115        202048        210076        202027
BM_SpeedMLDSAKeyGen/ml_dsa_87/threads:1                +0.1130         +0.1131         90237        100437         90230        100431
BM_SpeedMLDSASign/ml_dsa_87/threads:1                  -0.0645         -0.0645        250679        234521        250656        234486

Change-Id: Ibdcf5591dc74c61fc694988828b2fa254f8112f2
Reviewed-on: https://boringssl-review.googlesource.com/c/boringssl/+/96810
Commit-Queue: David Benjamin <davidben@google.com>
Reviewed-by: Adam Langley <agl@google.com>
diff --git a/crypto/fipsmodule/mldsa/mldsa.cc.inc b/crypto/fipsmodule/mldsa/mldsa.cc.inc
index 504844f..d5d1e96 100644
--- a/crypto/fipsmodule/mldsa/mldsa.cc.inc
+++ b/crypto/fipsmodule/mldsa/mldsa.cc.inc
@@ -1513,6 +1513,24 @@
   }
 }
 
+// `vector_encode_signed_ntt` behaves like `vector_encode_signed` but takes its
+// input in NTT form. This is only used by the testing-only `skEncode` and
+// `skDecode` functions and is not performance-sensitive.
+template <int bits, uint32_t max, int X>
+inline void vector_encode_signed_ntt(uint8_t *out, const vector<X> *a_ntt) {
+  vector<X> a = *a_ntt;
+  vector_inverse_ntt_montgomery(&a);
+  // `vector_inverse_ntt_montgomery` adds an extra factor of R, which we must
+  // cancel with a Montgomery reduction.
+  for (int i = 0; i < X; i++) {
+    for (int j = 0; j < kDegree; j++) {
+      a.v[i].c[j] = reduce_montgomery(a.v[i].c[j]);
+    }
+  }
+
+  vector_encode_signed<bits, max>(out, &a);
+}
+
 // FIPS 204, Algorithm 19 (`BitUnpack`).
 template <int bits, uint32_t max, int X>
 inline int vector_decode_signed(vector<X> *out, const uint8_t *in) {
@@ -1591,9 +1609,9 @@
 struct private_key {
   public_key<K> pub;
   uint8_t k[kKBytes];
-  vector<L> s1;
-  vector<K> s2;
-  vector<K> t0;
+  vector<L> s1_ntt;
+  vector<K> s2_ntt;
+  vector<K> t0_ntt;
 };
 
 template <int K, int L>
@@ -1657,21 +1675,21 @@
   if (!CBB_add_space(out, &vectorl_output, scalar_bytes * L)) {
     return 0;
   }
-  vector_encode_signed<plus_minus_eta_bitlen<K>(), eta<K>()>(vectorl_output,
-                                                             &priv->s1);
+  vector_encode_signed_ntt<plus_minus_eta_bitlen<K>(), eta<K>()>(vectorl_output,
+                                                                 &priv->s1_ntt);
 
   uint8_t *s2_output;
   if (!CBB_add_space(out, &s2_output, scalar_bytes * K)) {
     return 0;
   }
-  vector_encode_signed<plus_minus_eta_bitlen<K>(), eta<K>()>(s2_output,
-                                                             &priv->s2);
+  vector_encode_signed_ntt<plus_minus_eta_bitlen<K>(), eta<K>()>(s2_output,
+                                                                 &priv->s2_ntt);
 
   uint8_t *t0_output;
   if (!CBB_add_space(out, &t0_output, 416 * K)) {
     return 0;
   }
-  vector_encode_signed<13, (1 << 12)>(t0_output, &priv->t0);
+  vector_encode_signed_ntt<13, (1 << 12)>(t0_output, &priv->t0_ntt);
 
   return 1;
 }
@@ -1688,15 +1706,16 @@
       !CBS_get_bytes(in, &public_key_hash, kTrBytes) ||
       !CBS_get_bytes(in, &s1_bytes, scalar_bytes * L) ||
       !vector_decode_signed<plus_minus_eta_bitlen<K>(), eta<K>()>(
-          &priv->s1, CBS_data(&s1_bytes)) ||
+          &priv->s1_ntt, CBS_data(&s1_bytes)) ||
       !CBS_get_bytes(in, &s2_bytes, scalar_bytes * K) ||
       !vector_decode_signed<plus_minus_eta_bitlen<K>(), eta<K>()>(
-          &priv->s2, CBS_data(&s2_bytes)) ||
+          &priv->s2_ntt, CBS_data(&s2_bytes)) ||
       !CBS_get_bytes(in, &t0_bytes, 416 * K)) {
     return 0;
   }
 
-  // Compute `t1`, which is not in the `skDecode` input.
+  // Compute `t1`, which is not in the `skDecode` input. This also computes
+  // `t0_ntt` and converts `s1_ntt` and `s2_ntt` to NTT form.
   uint8_t unused[public_key_bytes<K>()];
   if (!mldsa_finish_keygen(unused, priv)) {
     return 0;
@@ -1705,7 +1724,7 @@
   // As a side effect of computing `t1`, we also compute `t0` and
   // `public_key_hash`. Check they match the received bytes.
   uint8_t t0_computed[416 * K];
-  vector_encode_signed<13, (1 << 12)>(t0_computed, &priv->t0);
+  vector_encode_signed_ntt<13, (1 << 12)>(t0_computed, &priv->t0_ntt);
   if (!CBS_mem_equal(&public_key_hash, priv->pub.public_key_hash,
                      sizeof(priv->pub.public_key_hash)) ||
       !CBS_mem_equal(&t0_bytes, t0_computed, sizeof(t0_computed))) {
@@ -1758,6 +1777,9 @@
 
 // FIPS 204, Algorithm 6 (`ML-DSA.KeyGen_internal`), steps 3 and 5–11.
 // Returns 1 on success and 0 on failure.
+//
+// On input, `priv->s1_ntt` and `priv->s2_ntt` are not yet in NTT form. They are
+// converted to NTT form by this function.
 template <int K, int L>
 inline int mldsa_finish_keygen(
     uint8_t out_encoded_public_key[public_key_bytes<K>()],
@@ -1767,7 +1789,6 @@
   struct Values {
     enum { kAllowUniquePtr = true };
     matrix<K, L> a_ntt;
-    vector<L> s1_ntt;
     vector<K> t;
   };
   auto values = MakeUnique<Values>();
@@ -1778,16 +1799,17 @@
   // Step 3.
   matrix_expand(&values->a_ntt, priv->pub.rho);
 
-  // Step 5.
-  OPENSSL_memcpy(&values->s1_ntt, &priv->s1, sizeof(values->s1_ntt));
-  vector_ntt(&values->s1_ntt);
-
-  matrix_mult_montgomery(&values->t, &values->a_ntt, &values->s1_ntt);
+  // Step 5. Note that, on input, `s1_ntt` and `s2_ntt` are not yet in NTT form,
+  // but this function is responsible for converting them to NTT.
+  vector_ntt(&priv->s1_ntt);
+  matrix_mult_montgomery(&values->t, &values->a_ntt, &priv->s1_ntt);
   vector_inverse_ntt_montgomery(&values->t);
-  vector_add(&values->t, &values->t, &priv->s2);
+  vector_add(&values->t, &values->t, &priv->s2_ntt);
+  vector_ntt(&priv->s2_ntt);
 
-  // Step 6-7.
-  vector_power2_round(&priv->pub.t1, &priv->t0, &values->t);
+  // Step 6-7. We store t0 in NTT form.
+  vector_power2_round(&priv->pub.t1, &priv->t0_ntt, &values->t);
+  vector_ntt(&priv->t0_ntt);
   // t1 is public.
   CONSTTIME_DECLASSIFY(&priv->pub.t1, sizeof(priv->pub.t1));
 
@@ -1829,8 +1851,9 @@
   CONSTTIME_DECLASSIFY(rho, kRhoBytes);
   OPENSSL_memcpy(priv->pub.rho, rho, sizeof(priv->pub.rho));
   OPENSSL_memcpy(priv->k, k, sizeof(priv->k));
-  // Step 4. This is independent of A (step 3) and can be done first.
-  vector_expand_short(&priv->s1, &priv->s2, sigma);
+  // Steps 4. This is independent of A (step 3) and can be done first.
+  // `mldsa_finish_keygen` will convert `s1_ntt` and `s2_ntt` into NTT from.
+  vector_expand_short(&priv->s1_ntt, &priv->s2_ntt, sigma);
   // Steps 3 and 5-11.
   return mldsa_finish_keygen(out_encoded_public_key, priv);
 }
@@ -1865,9 +1888,6 @@
   struct Values {
     enum { kAllowUniquePtr = true };
     signature<K, L> sign;
-    vector<L> s1_ntt;
-    vector<K> s2_ntt;
-    vector<K> t0_ntt;
     matrix<K, L> a_ntt;
     vector<L> y;
     vector<K> w;
@@ -1879,15 +1899,6 @@
   if (values == nullptr) {
     return 0;
   }
-  OPENSSL_memcpy(&values->s1_ntt, &priv->s1, sizeof(values->s1_ntt));
-  vector_ntt(&values->s1_ntt);
-
-  OPENSSL_memcpy(&values->s2_ntt, &priv->s2, sizeof(values->s2_ntt));
-  vector_ntt(&values->s2_ntt);
-
-  OPENSSL_memcpy(&values->t0_ntt, &priv->t0, sizeof(values->t0_ntt));
-  vector_ntt(&values->t0_ntt);
-
   matrix_expand(&values->a_ntt, priv->pub.rho);
 
   // kappa must not exceed 2**16/L = 13107. But the probability of it
@@ -1917,9 +1928,9 @@
                                   sizeof(values->sign.c_tilde), tau<K>());
     scalar_ntt(&c_ntt);
 
-    vector_mult_scalar_montgomery(&values->cs1, &values->s1_ntt, &c_ntt);
+    vector_mult_scalar_montgomery(&values->cs1, &priv->s1_ntt, &c_ntt);
     vector_inverse_ntt_montgomery(&values->cs1);
-    vector_mult_scalar_montgomery(&values->cs2, &values->s2_ntt, &c_ntt);
+    vector_mult_scalar_montgomery(&values->cs2, &priv->s2_ntt, &c_ntt);
     vector_inverse_ntt_montgomery(&values->cs2);
 
     vector_add(&values->sign.z, &values->y, &values->cs1);
@@ -1951,7 +1962,7 @@
     }
 
     vector<K> *ct0 = &values->w1;
-    vector_mult_scalar_montgomery(ct0, &values->t0_ntt, &c_ntt);
+    vector_mult_scalar_montgomery(ct0, &priv->t0_ntt, &c_ntt);
     vector_inverse_ntt_montgomery(ct0);
     vector_make_hint(&values->sign.h, ct0, &values->cs2, &values->w);