Use lazy reduction in ML-KEM matrix-vector products In the NTT domain, a product of two ML-KEM polynomials consists of 128 products in rings of the form GF(3329)[X]/(X^2 - zeta), each computing a0*b0 + a1*b1*zeta and a0*b1 + a1*b0. scalar_mult reduced a1*b1 before multiplying it by zeta, and then reduced both results, for three Barrett reductions per pair of coefficients. The matrix-vector and inner product functions then added each product to the running sum with another conditional subtraction. These products are only ever summed. Since the inputs are fully reduced, each term is below 2 * 3329^2, and a sum of up to four terms fits in 32 bits. Instead, accumulate the unreduced terms in a uint32_t array and reduce each coefficient once, at the end. A static_assert checks the bound. Also, b1*zeta only depends on the right-hand side, which is the same for every row of a matrix-vector product. Compute it once per vector element and pass it to the products. When encrypting, the matrix-vector product and the inner product share the cache. Both techniques are described for Kyber's matrix-vector product in "Faster Kyber and Dilithium on the Cortex-M4" (https://eprint.iacr.org/2022/112), as "better accumulation" and "asymmetric multiplication". The latter comes from "Neon NTT" (https://eprint.iacr.org/2021/986). For one ML-KEM-768 matrix-vector product, this replaces 3,456 Barrett reductions and 2,304 conditional subtractions with 1,152 Barrett reductions. The output is unchanged, as are the self-test known-answer values, and the code remains constant-time. scalar_zero and vector_zero are no longer used, so remove them. Median of 15 runs on an Intel Xeon (Skylake-SP) at 2.0 GHz, in microseconds: ML-KEM-768 ML-KEM-1024 Before After Speedup Before After Speedup keygen 35.1 32.5 1.08x 49.2 43.3 1.14x encap 24.4 20.8 1.17x 30.8 24.5 1.26x decap 37.0 33.1 1.12x 48.6 41.4 1.17x Change-Id: Ib8ff5c08d24bf4c8c02218ad270ba448e96fc5f8 Reviewed-on: https://boringssl-review.googlesource.com/c/boringssl/+/104068 Reviewed-by: David Benjamin <davidben@google.com> 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>
diff --git a/crypto/fipsmodule/mlkem/mlkem.cc.inc b/crypto/fipsmodule/mlkem/mlkem.cc.inc index 8071a42..ab640ef 100644 --- a/crypto/fipsmodule/mlkem/mlkem.cc.inc +++ b/crypto/fipsmodule/mlkem/mlkem.cc.inc
@@ -236,13 +236,6 @@ // `uint32_t`. inline uint16_t reduce(uint32_t x) { return reduce_once(reduce_partial(x)); } -inline void scalar_zero(scalar *out) { OPENSSL_memset(out, 0, sizeof(*out)); } - -template <int RANK> -inline void vector_zero(vector<RANK> *out) { - OPENSSL_memset(out->v, 0, sizeof(scalar) * RANK); -} - // In place number theoretic transform of a given scalar. // Note that MLKEM's kPrime 3329 does not have a 512th root of unity, so this // transform leaves off the last iteration of the usual FFT code, with the 128 @@ -376,19 +369,75 @@ // 3329 does not have a 512th root of unity, this means we have to interpret // the 2*ith and (2*i+1)th entries of the scalar as elements of GF(3329)[X]/(X^2 // - 17^(2*bitreverse(i)+1)) The value of 17^(2*bitreverse(i)+1) mod 3329 is -// stored in the precomputed `kModRoots` table. Note that our Barrett transform -// only allows us to multiply two reduced numbers together, so we need some -// intermediate reduction steps, even if an uint64_t could hold 3 multiplied -// numbers. -inline void scalar_mult(scalar *out, const scalar *lhs, const scalar *rhs) { +// stored in the precomputed `kModRoots` table. +// +// Writing the two elements as a0 + a1*X and b0 + b1*X, and the root as zeta, +// their product is (a0*b0 + a1*(b1*zeta)) + (a0*b1 + a1*b0)*X. Products are +// only ever needed as terms of a sum, in matrix-vector and inner products, so +// they are computed lazily: +// +// - b1*zeta mod kPrime only depends on the right-hand side, which is shared by +// every row of a matrix-vector product. It is computed once per scalar and +// stored in a `scalar_mulcache`. This is the "asymmetric multiplication" of +// https://eprint.iacr.org/2021/986. +// - The terms are summed in a `scalar_accumulator` without any modular +// reduction, and the sum is reduced once, at the end. + +// scalar_mulcache holds, for a scalar b, b[2*i+1] * kModRoots[i] mod kPrime. +struct scalar_mulcache { + // On every function entry and exit, 0 <= c < kPrime. + uint16_t c[DEGREE / 2]; +}; + +template <int RANK> +struct vector_mulcache { + scalar_mulcache v[RANK]; +}; + +inline void scalar_mulcache_compute(scalar_mulcache *out, const scalar *b) { for (int i = 0; i < DEGREE / 2; i++) { - uint32_t real_real = (uint32_t)lhs->c[2 * i] * rhs->c[2 * i]; - uint32_t img_img = (uint32_t)lhs->c[2 * i + 1] * rhs->c[2 * i + 1]; - uint32_t real_img = (uint32_t)lhs->c[2 * i] * rhs->c[2 * i + 1]; - uint32_t img_real = (uint32_t)lhs->c[2 * i + 1] * rhs->c[2 * i]; - out->c[2 * i] = - reduce(real_real + (uint32_t)reduce(img_img) * kModRoots[i]); - out->c[2 * i + 1] = reduce(img_real + real_img); + out->c[i] = reduce(uint32_t{b->c[2 * i + 1]} * kModRoots[i]); + } +} + +template <int RANK> +inline void vector_mulcache_compute(vector_mulcache<RANK> *out, + const vector<RANK> *b) { + for (int i = 0; i < RANK; i++) { + scalar_mulcache_compute(&out->v[i], &b->v[i]); + } +} + +// scalar_accumulator holds a sum of products of scalars, not reduced mod +// kPrime. +struct scalar_accumulator { + uint32_t c[DEGREE]; +}; + +// Scalars are fully reduced, so each product adds at most 2 * (kPrime - 1)^2 to +// each element of a `scalar_accumulator`. Sums have at most RANK1024 terms, so +// the elements never overflow. +static_assert(uint64_t{RANK1024} * 2 * (kPrime - 1) * (kPrime - 1) <= + UINT32_MAX); + +// scalar_mult_accumulate adds `lhs * rhs` to `acc`. `rhs_cache` must have been +// computed from `rhs`. +inline void scalar_mult_accumulate(scalar_accumulator *acc, const scalar *lhs, + const scalar *rhs, + const scalar_mulcache *rhs_cache) { + for (int i = 0; i < DEGREE / 2; i++) { + const uint32_t lhs_real = lhs->c[2 * i]; + const uint32_t lhs_img = lhs->c[2 * i + 1]; + acc->c[2 * i] += lhs_real * rhs->c[2 * i] + lhs_img * rhs_cache->c[i]; + acc->c[2 * i + 1] += lhs_real * rhs->c[2 * i + 1] + lhs_img * rhs->c[2 * i]; + } +} + +// scalar_from_accumulator sets `out` to `acc` mod kPrime. +inline void scalar_from_accumulator(scalar *out, + const scalar_accumulator *acc) { + for (int i = 0; i < DEGREE; i++) { + out->c[i] = reduce(acc->c[i]); } } @@ -401,39 +450,39 @@ template <int RANK> inline void matrix_mult(vector<RANK> *out, const matrix<RANK> *m, - const vector<RANK> *a) { - vector_zero(out); + const vector<RANK> *a, + const vector_mulcache<RANK> *a_cache) { for (int i = 0; i < RANK; i++) { + scalar_accumulator acc = {}; for (int j = 0; j < RANK; j++) { - scalar product; - scalar_mult(&product, &m->v[i][j], &a->v[j]); - scalar_add(&out->v[i], &product); + scalar_mult_accumulate(&acc, &m->v[i][j], &a->v[j], &a_cache->v[j]); } + scalar_from_accumulator(&out->v[i], &acc); } } template <int RANK> inline void matrix_mult_transpose(vector<RANK> *out, const matrix<RANK> *m, - const vector<RANK> *a) { - vector_zero(out); + const vector<RANK> *a, + const vector_mulcache<RANK> *a_cache) { for (int i = 0; i < RANK; i++) { + scalar_accumulator acc = {}; for (int j = 0; j < RANK; j++) { - scalar product; - scalar_mult(&product, &m->v[j][i], &a->v[j]); - scalar_add(&out->v[i], &product); + scalar_mult_accumulate(&acc, &m->v[j][i], &a->v[j], &a_cache->v[j]); } + scalar_from_accumulator(&out->v[i], &acc); } } template <int RANK> inline void scalar_inner_product(scalar *out, const vector<RANK> *lhs, - const vector<RANK> *rhs) { - scalar_zero(out); + const vector<RANK> *rhs, + const vector_mulcache<RANK> *rhs_cache) { + scalar_accumulator acc = {}; for (int i = 0; i < RANK; i++) { - scalar product; - scalar_mult(&product, &lhs->v[i], &rhs->v[i]); - scalar_add(out, &product); + scalar_mult_accumulate(&acc, &lhs->v[i], &rhs->v[i], &rhs_cache->v[i]); } + scalar_from_accumulator(out, &acc); } inline void scalar_from_keccak_block_vartime(scalar *out, int *done, @@ -1088,11 +1137,13 @@ vector_decode<du>(&u, ciphertext); vector_decompress(&u, du); vector_ntt(&u); + vector_mulcache<RANK> u_cache; + vector_mulcache_compute(&u_cache, &u); scalar v; scalar_decode<dv>(&v, ciphertext + compressed_vector_size(RANK)); scalar_decompress(&v, dv); scalar mask; - scalar_inner_product(&mask, &priv->s, &u); + scalar_inner_product(&mask, &priv->s, &u, &u_cache); scalar_inverse_ntt(&mask); scalar_sub(&v, &mask); scalar_compress(&v, 1); @@ -1141,7 +1192,9 @@ vector<RANK> error; vector_generate_secret_eta_2(&error, &counter, sigma); vector_ntt(&error); - matrix_mult_transpose(&priv->pub.t, &priv->pub.m, &priv->s); + vector_mulcache<RANK> s_cache; + vector_mulcache_compute(&s_cache, &priv->s); + matrix_mult_transpose(&priv->pub.t, &priv->pub.m, &priv->s, &s_cache); vector_add(&priv->pub.t, &error); // t is part of the public key and thus is public. CONSTTIME_DECLASSIFY(&priv->pub.t, sizeof(priv->pub.t)); @@ -1180,6 +1233,8 @@ mlkem::vector<RANK> secret; vector_generate_secret_eta_2(&secret, &counter, randomness); vector_ntt(&secret); + mlkem::vector_mulcache<RANK> secret_cache; + vector_mulcache_compute(&secret_cache, &secret); mlkem::vector<RANK> error; vector_generate_secret_eta_2(&error, &counter, randomness); uint8_t input[33]; @@ -1188,11 +1243,11 @@ mlkem::scalar scalar_error; scalar_centered_binomial_distribution_eta_2_with_prf(&scalar_error, input); mlkem::vector<RANK> u; - matrix_mult(&u, &pub->m, &secret); + matrix_mult(&u, &pub->m, &secret, &secret_cache); vector_inverse_ntt(&u); vector_add(&u, &error); mlkem::scalar v; - scalar_inner_product(&v, &pub->t, &secret); + scalar_inner_product(&v, &pub->t, &secret, &secret_cache); scalar_inverse_ntt(&v); scalar_add(&v, &scalar_error); mlkem::scalar expanded_message;