Leave intermediates partially reduced in ML-KEM NTT This ports a074f282d026a0ebbed7c9efef5a0cf63f72338d's change to ML-DSA. Changes relative to the ML-DSA changes: - Barrett reduction needed to be adjusted from 24 to 32 bits so it can successfully reduce any uint32_t to 0..2*kPrime range. - NTT using partial reduction uses intermediates of up to 15*kPrime, which fits in uint16_t. - INTT using partial reduction uses intermediates of up to 128*kPrime, which required storing intermediates in an uint32_t array. This wasn't a problem in ML-DSA as it already stores vectors of uint32_t to begin with. Saves about 4.9% to 8.1% for TLS ML-KEM operations (most savings being in Encap and Decap, namely 11.0% to 14.7%) on aarch64 Apple M1 Pro: Did 27910 BM_SpeedMLKEM768KeyGenDecap/threads:1 operations (39989.3 ops/sec) [+4.9%] Did 48780 BM_SpeedMLKEM768ParseEncap/threads:1 operations (69751.4 ops/sec) [+7.6%] Did 19862 BM_SpeedMLKEM1024KeyGenDecap/threads:1 operations (28153.9 ops/sec) [+5.1%] Did 33890 BM_SpeedMLKEM1024ParseEncap/threads:1 operations (48063.6 ops/sec) [+8.1%] Did 46596 BM_SpeedMLKEM768KeyGenOnly/threads:1 operations (66675.0 ops/sec) [+2.9%] Did 45743 BM_SpeedMLKEM768PrivateKeyFromSeedOnly/threads:1 operations (65436.1 ops/sec) [+1.9%] Did 68136 BM_SpeedMLKEM768DecapOnly/threads:1 operations (96674.0 ops/sec) [+11.0%] Did 85036 BM_SpeedMLKEM768ParseOnly/threads:1 operations (120954.2 ops/sec) [+2.7%] Did 105067 BM_SpeedMLKEM768EncapOnly/threads:1 operations (150126.4 ops/sec) [+14.4%] Did 30992 BM_SpeedMLKEM1024KeyGenOnly/threads:1 operations (44467.0 ops/sec) [+3.2%] Did 30577 BM_SpeedMLKEM1024PrivateKeyFromSeedOnly/threads:1 operations (43805.5 ops/sec) [+2.9%] Did 50684 BM_SpeedMLKEM1024DecapOnly/threads:1 operations (73954.1 ops/sec) [+14.7%] Did 54923 BM_SpeedMLKEM1024ParseOnly/threads:1 operations (77869.7 ops/sec) [+2.6%] Did 82865 BM_SpeedMLKEM1024EncapOnly/threads:1 operations (118470.5 ops/sec) [+14.1%] Minor improvements (2.9% to 4.5% but only for parse/encap) on AMD EPYC 7B13: Did 12436 BM_SpeedMLKEM768KeyGenDecap/threads:1 operations (17751.4 ops/sec) [+0.0%] Did 22079 BM_SpeedMLKEM768ParseEncap/threads:1 operations (31509.8 ops/sec) [+2.9%] Did 8851 BM_SpeedMLKEM1024KeyGenDecap/threads:1 operations (12692.4 ops/sec) [-0.2%] Did 15821 BM_SpeedMLKEM1024ParseEncap/threads:1 operations (22667.0 ops/sec) [+4.5%] Did 23994 BM_SpeedMLKEM768KeyGenOnly/threads:1 operations (34374.9 ops/sec) [-1.9%] Did 23451 BM_SpeedMLKEM768PrivateKeyFromSeedOnly/threads:1 operations (33470.7 ops/sec) [-1.7%] Did 24538 BM_SpeedMLKEM768DecapOnly/threads:1 operations (35050.9 ops/sec) [+1.9%] Did 50112 BM_SpeedMLKEM768ParseOnly/threads:1 operations (71421.1 ops/sec) [+1.1%] Did 36402 BM_SpeedMLKEM768EncapOnly/threads:1 operations (52087.7 ops/sec) [+4.9%] Did 16656 BM_SpeedMLKEM1024KeyGenOnly/threads:1 operations (23778.1 ops/sec) [-0.7%] Did 16274 BM_SpeedMLKEM1024PrivateKeyFromSeedOnly/threads:1 operations (23304.1 ops/sec) [-0.5%] Did 18824 BM_SpeedMLKEM1024DecapOnly/threads:1 operations (26908.8 ops/sec) [+3.0%] Did 33258 BM_SpeedMLKEM1024ParseOnly/threads:1 operations (47471.2 ops/sec) [+1.0%] Did 28276 BM_SpeedMLKEM1024EncapOnly/threads:1 operations (40380.6 ops/sec) [+7.4%] Bug: 503700354 Change-Id: Idd4b6e8986253288e09e40b8ffbf3adf6a6a6964 Reviewed-on: https://boringssl-review.googlesource.com/c/boringssl/+/101967 Auto-Submit: Rudolf Polzer <rpolzer@google.com> 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 f0ef633..13382e3 100644 --- a/crypto/fipsmodule/mlkem/mlkem.cc.inc +++ b/crypto/fipsmodule/mlkem/mlkem.cc.inc
@@ -91,8 +91,8 @@ // Constants that are common across all sizes. #define DEGREE 256 -const size_t kBarrettMultiplier = 5039; -const unsigned kBarrettShift = 24; +const size_t kBarrettMultiplier = 1290167; +const unsigned kBarrettShift = 32; const uint16_t kPrime = 3329; const int kLog2Prime = 12; const uint16_t kHalfPrime = (/*kPrime=*/3329 - 1) / 2; @@ -223,16 +223,19 @@ return (mask & x) | (~mask & subtracted); } -// constant time reduce x mod kPrime using Barrett reduction. x must be less -// than kPrime + 2×kPrime². -inline uint16_t reduce(uint32_t x) { - declassify_assert(x < kPrime + 2u * kPrime * kPrime); +// reduce_partial returns a number equivalent to x mod kPrime. The output is +// less than 2 * kPrime. `x` can be any `uint32_t`. +inline uint32_t reduce_partial(uint32_t x) { uint64_t product = (uint64_t)x * kBarrettMultiplier; uint32_t quotient = (uint32_t)(product >> kBarrettShift); uint32_t remainder = x - quotient * kPrime; - return reduce_once(remainder); + return remainder; } +// constant time reduce x mod kPrime using Barrett reduction. `x` can be any +// `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> @@ -248,20 +251,30 @@ // elements being consecutive entries in `s->c`. inline void scalar_ntt(scalar *s) { // Manually unrolled loop to maximize vectorization. -#define ITER(step, offset) \ - { \ - int k = 0; \ - for (int i = 0; i < step; i++) { \ - const uint32_t step_root = kNTTRoots[i + step]; \ - for (int j = k; j < k + offset; j++) { \ - uint16_t odd = reduce(step_root * s->c[j + offset]); \ - uint16_t even = s->c[j]; \ - s->c[j] = reduce_once(odd + even); \ - s->c[j + offset] = reduce_once(even - odd + kPrime); \ - } \ - k += 2 * offset; \ - } \ + // + // Intermediate values of `s` are not fully reduced. `ITER` writes to every + // element of `s` exactly once. Each write adds less than 2*kPrime to some + // value from the previous iteration. After seven iterations, each element + // will be less than 15*kPrime. This fits in `uint16_t`, so we never overflow. +#define ITER(step, offset) \ + { \ + int k = 0; \ + for (int i = 0; i < step; i++) { \ + const uint32_t step_root = kNTTRoots[step + i]; \ + for (int j = k; j < k + offset; j++) { \ + uint32_t even = s->c[j]; \ + /* `reduce_partial` works on values up to 2^32. \ + * `step_root < kPrime` because it's static data. `s->c[...] < 2^16` \ + * because it's 16-bit. */ \ + uint32_t odd = \ + reduce_partial(uint32_t{step_root} * uint32_t{s->c[j + offset]}); \ + s->c[j] = even + odd; \ + s->c[j + offset] = even + 2 * kPrime - odd; \ + } \ + k += 2 * offset; \ + } \ } + // for (int step = 1; step < DEGREE / 2; step <<= 1) ITER(1, 128) ITER(2, 64) @@ -272,6 +285,11 @@ ITER(64, 2) static_assert(DEGREE == 256); #undef ITER + + // Now `s` is correct but the elements are not fully reduced. Reduce them. + for (int i = 0; i < DEGREE; i++) { + s->c[i] = reduce(s->c[i]); + } } template <int RANK> @@ -286,35 +304,52 @@ // number theoretic transform, this leaves off the first step of the normal iFFT // to account for the fact that 3329 does not have a 512th root of unity, using // the precomputed 128 roots of unity stored in `kInverseNTTRoots`. +// +// FIPS 203, Algorithm 10 (`NTT^-1`). void scalar_inverse_ntt(scalar *s) { // Manually unrolled loop to maximize vectorization. -#define ITER(step, offset) \ - { \ - int k = 0; \ - for (int i = 0; i < step; i++) { \ - uint32_t step_root = kInverseNTTRoots[i + step]; \ - for (int j = k; j < k + offset; j++) { \ - uint16_t odd = s->c[j + offset]; \ - uint16_t even = s->c[j]; \ - s->c[j] = reduce_once(odd + even); \ - s->c[j + offset] = reduce(step_root * (even - odd + kPrime)); \ - } \ - k += 2 * offset; \ - } \ + // + // Intermediate values are not fully reduced. `ITER` writes to every element + // exactly once. Each write either adds two values from the previous + // iteration, or is at most 2 * kPrime. After seven iterations, each element + // will be less than 2^7 * kPrime = 128 * kPrime. This fits in `uint32_t`, so + // we never overflow. + uint32_t c[DEGREE]; +#define ITER(step, offset, src) \ + { \ + int k = 0; \ + for (int i = 0; i < step; i++) { \ + const uint32_t step_root = kInverseNTTRoots[step + i]; \ + for (int j = k; j < k + offset; j++) { \ + uint32_t even = (src)[j]; \ + uint32_t odd = (src)[j + offset]; \ + c[j] = odd + even; \ + /* `reduce_partial` works on values up to 2^32. \ + * `step_root < kPrime` because it's static data. The other term is \ + * less than 2^32, and does not overflow because `odd` is at most \ + * 128 * kPrime per the above. */ \ + c[j + offset] = reduce_partial(uint32_t{step_root} * \ + uint32_t{128 * kPrime + even - odd}); \ + } \ + k += 2 * offset; \ + } \ } + // for (int offset = 2; offset < DEGREE; offset <<= 1) - ITER(64, 2) - ITER(32, 4) - ITER(16, 8) - ITER(8, 16) - ITER(4, 32) - ITER(2, 64) - ITER(1, 128) + ITER(64, 2, s->c) + ITER(32, 4, c) + ITER(16, 8, c) + ITER(8, 16, c) + ITER(4, 32, c) + ITER(2, 64, c) + ITER(1, 128, c) static_assert(DEGREE == 256); #undef ITER + // The final multiplication by `kInverseDegree` also fully reduces the + // partially-reduced values above. for (int i = 0; i < DEGREE; i++) { - s->c[i] = reduce(s->c[i] * kInverseDegree); + s->c[i] = reduce(c[i] * kInverseDegree); } }