Leave intermediates partially reduced in ML-DSA NTT Credit to Filippo Valsorda for doing something similar in https://go-review.googlesource.com/c/go/+/822040/ This is quite effective on aarch64 (Apple M1 Pro): Benchmark Time CPU Time Old Time New CPU Old CPU New -------------------------------------------------------------------------------------------------------------------------------------------------- BM_SpeedMLDSAKeyGen/ml_dsa_44/threads:1 -0.1039 -0.1039 33944 30417 33944 30416 BM_SpeedMLDSASign/ml_dsa_44/threads:1 -0.1277 -0.1280 140310 122385 140308 122343 BM_SpeedMLDSAParsePubKey/ml_dsa_44/threads:1 +0.0080 +0.0049 2417 2436 2417 2429 BM_SpeedMLDSAVerify/ml_dsa_44/threads:1 -0.0991 -0.0990 29312 26406 29308 26406 BM_SpeedMLDSAVerifyBadSignature/ml_dsa_44/threads:1 -0.0993 -0.0993 29266 26360 29266 26360 BM_SpeedMLDSAKeyGen/ml_dsa_65/threads:1 -0.0786 -0.0784 67418 62122 67405 62121 BM_SpeedMLDSASign/ml_dsa_65/threads:1 -0.1127 -0.1127 221116 196194 221066 196151 BM_SpeedMLDSAParsePubKey/ml_dsa_65/threads:1 +0.0040 +0.0034 3581 3595 3580 3593 BM_SpeedMLDSAVerify/ml_dsa_65/threads:1 -0.0940 -0.0937 45508 41230 45492 41229 BM_SpeedMLDSAVerifyBadSignature/ml_dsa_65/threads:1 -0.0920 -0.0913 45578 41383 45535 41378 BM_SpeedMLDSAKeyGen/ml_dsa_87/threads:1 -0.0817 -0.0818 86174 79132 86174 79126 BM_SpeedMLDSASign/ml_dsa_87/threads:1 -0.1084 -0.1084 251821 224521 251813 224520 BM_SpeedMLDSAParsePubKey/ml_dsa_87/threads:1 +0.0004 -0.0004 4793 4795 4786 4785 BM_SpeedMLDSAVerify/ml_dsa_87/threads:1 -0.0695 -0.0696 74345 69180 74344 69170 BM_SpeedMLDSAVerifyBadSignature/ml_dsa_87/threads:1 -0.0717 -0.0717 74405 69072 74405 69071 Less visible on x86_64 (AMD Ryzen Threadripper PRO 7945WX 12-Cores) Benchmark Time CPU Time Old Time New CPU Old CPU New -------------------------------------------------------------------------------------------------------------------------------------------------- BM_SpeedMLDSAKeyGen/ml_dsa_44/threads:1 -0.0184 -0.0184 38295 37591 38293 37589 BM_SpeedMLDSASign/ml_dsa_44/threads:1 -0.0041 -0.0041 124470 123958 124458 123949 BM_SpeedMLDSAParsePubKey/ml_dsa_44/threads:1 -0.0043 -0.0043 2740 2728 2740 2728 BM_SpeedMLDSAVerify/ml_dsa_44/threads:1 -0.0216 -0.0216 31537 30855 31534 30854 BM_SpeedMLDSAVerifyBadSignature/ml_dsa_44/threads:1 -0.0272 -0.0271 31575 30718 31572 30716 BM_SpeedMLDSAKeyGen/ml_dsa_65/threads:1 -0.0191 -0.0191 72158 70780 72153 70775 BM_SpeedMLDSASign/ml_dsa_65/threads:1 -0.0304 -0.0305 198536 192510 198519 192470 BM_SpeedMLDSAParsePubKey/ml_dsa_65/threads:1 -0.0004 -0.0003 4084 4083 4084 4082 BM_SpeedMLDSAVerify/ml_dsa_65/threads:1 -0.0168 -0.0168 49948 49111 49944 49107 BM_SpeedMLDSAVerifyBadSignature/ml_dsa_65/threads:1 -0.0197 -0.0197 50288 49299 50286 49297 BM_SpeedMLDSAKeyGen/ml_dsa_87/threads:1 -0.0161 -0.0161 101991 100347 101979 100335 BM_SpeedMLDSASign/ml_dsa_87/threads:1 -0.0132 -0.0132 240216 237045 240183 237024 BM_SpeedMLDSAParsePubKey/ml_dsa_87/threads:1 -0.0052 -0.0051 5468 5439 5467 5439 BM_SpeedMLDSAVerify/ml_dsa_87/threads:1 -0.0275 -0.0276 84974 82635 84969 82625 BM_SpeedMLDSAVerifyBadSignature/ml_dsa_87/threads:1 -0.0258 -0.0257 84841 82652 84829 82649 Bug: 503700354 Change-Id: I14b83d24c3585a42959241907bab2005d1e4f5aa Reviewed-on: https://boringssl-review.googlesource.com/c/boringssl/+/101807 Reviewed-by: Lily Chen <chlily@google.com> Commit-Queue: David Benjamin <davidben@google.com>
diff --git a/crypto/fipsmodule/mldsa/mldsa.cc.inc b/crypto/fipsmodule/mldsa/mldsa.cc.inc index d5d1e96..9d6a4dc 100644 --- a/crypto/fipsmodule/mldsa/mldsa.cc.inc +++ b/crypto/fipsmodule/mldsa/mldsa.cc.inc
@@ -308,6 +308,22 @@ return (mask & x) | (~mask & subtracted); } +// reduce_from_u32 returns `x mod kPrime`. `x` can be any `uint32_t`. +inline uint32_t reduce_from_u32(uint32_t x) { + // kPrime is 2^23 - 2^13 + 1. We first split `x` into the upper 9 bits and the + // lower 23 bits: + uint32_t hi = x >> 23, lo = x & ((uint32_t{1} << 23) - 1); + // 2^23 = 2^13 - 1 mod kPrime + // x = hi * 2^23 + lo + // = hi * 2^13 - hi + lo mod kPrime + // + // This value does not overflow and is below 2 * kPrime. We can `reduce_once`. + static constexpr uint32_t kMaxHi = UINT32_MAX >> 23; + static constexpr uint32_t kMaxLo = (uint32_t{1} << 23) - 1; + static_assert((kMaxHi << 13) - kMaxHi + kMaxLo < 2 * kPrime); + return reduce_once((hi << 13) - hi + lo); +} + // Returns the absolute value in constant time, interpreting the high bit as a // sign bit. inline uint32_t abs_signed(uint32_t x) { @@ -358,13 +374,20 @@ } } -inline uint32_t reduce_montgomery(uint64_t x) { +// reduce_montgomery_partial returns a number equivalent to x * R^-1 mod kPrime. +// The output is less than 2 * kPrime. `x` must be at most kPrime * 2^32. +inline uint32_t reduce_montgomery_partial(uint64_t x) { declassify_assert(x <= ((uint64_t)kPrime << 32)); uint64_t a = (uint32_t)x * kPrimeNegInverse; uint64_t b = x + a * kPrime; declassify_assert((b & 0xffffffff) == 0); - uint32_t c = b >> 32; - return reduce_once(c); + return b >> 32; +} + +// reduce_montgomery returns x * R^-1 mod kPrime, fully reduced. `x` must be at +// most kPrime * 2^32. +inline uint32_t reduce_montgomery(uint64_t x) { + return reduce_once(reduce_montgomery_partial(x)); } // Montgomery-multiply two scalars in the number theoretically transformed @@ -382,23 +405,28 @@ // FIPS 204, Algorithm 41 (`NTT`). 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 = kNTTRootsMontgomery[step + i]; \ - for (int j = k; j < k + offset; j++) { \ - uint32_t even = s->c[j]; \ - /* `reduce_montgomery` works on values up to kPrime*R and R \ - * 2*kPrime. `step_root` < kPrime because it's static data. \ - * `s->c[...]` is < kPrime by the invariants of that struct. */ \ - uint32_t odd = reduce_montgomery((uint64_t)step_root * \ - (uint64_t)s->c[j + offset]); \ - s->c[j] = reduce_once(odd + even); \ - s->c[j + offset] = mod_sub(even, odd); \ - } \ - 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 eight iterations, each element + // will be less than 17*kPrime. This fits in `uint32_t`, so we never overflow. +#define ITER(step, offset) \ + { \ + int k = 0; \ + for (int i = 0; i < step; i++) { \ + const uint32_t step_root = kNTTRootsMontgomery[step + i]; \ + for (int j = k; j < k + offset; j++) { \ + uint32_t even = s->c[j]; \ + /* `reduce_montgomery_partial` works on values up to kPrime*R. \ + * `step_root < kPrime` because it's static data. `s->c[...] < R` \ + * because it's 32-bit. */ \ + uint32_t odd = reduce_montgomery_partial(uint64_t{step_root} * \ + uint64_t{s->c[j + offset]}); \ + s->c[j] = even + odd; \ + s->c[j + offset] = even + 2 * kPrime - odd; \ + } \ + k += 2 * offset; \ + } \ } ITER(1, 128) @@ -411,6 +439,11 @@ ITER(128, 1) static_assert(kDegree == 256); #undef ITER + + // Now `s` is correct but the elements are not fully reduced. Reduce them. + for (int i = 0; i < kDegree; i++) { + s->c[i] = reduce_from_u32(s->c[i]); + } } // In place inverse number theoretic transform of a given scalar, but with each @@ -421,27 +454,31 @@ // FIPS 204, Algorithm 42 (`NTT^-1`). inline void scalar_inverse_ntt_montgomery(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 = \ - kPrime - kNTTRootsMontgomery[step + (step - 1 - i)]; \ - for (int j = k; j < k + offset; j++) { \ - uint32_t even = s->c[j]; \ - uint32_t odd = s->c[j + offset]; \ - s->c[j] = reduce_once(odd + even); \ - /* `reduce_montgomery` works on values up to kPrime*R and R > \ - * 2*kPrime. kPrime + even < 2*kPrime because `even` < kPrime, by the \ - * invariants of that structure. Thus kPrime + even - odd < 2*kPrime \ - * because odd >= 0, because it's unsigned and less than kPrime. \ - * Lastly step_root < kPrime, because `kNTTRootsMontgomery` is static \ - * data. */ \ - s->c[j + offset] = reduce_montgomery((uint64_t)step_root * \ - (uint64_t)(kPrime + even - odd)); \ - } \ - k += 2 * offset; \ - } \ + // + // Intermediate values of `s` are not fully reduced. `ITER` writes + // to every element of `s` exactly once. Each write either adds two values + // from the previous iteration, or is at most 2 * kPrime. After eight + // iterations, each element will be less than 2^8 * kPrime = 256 * kPrime. + // This fits in `uint32_t`, so we never overflow. +#define ITER(step, offset) \ + { \ + int k = 0; \ + for (int i = 0; i < step; i++) { \ + const uint32_t step_root = \ + kPrime - kNTTRootsMontgomery[step + (step - 1 - i)]; \ + for (int j = k; j < k + offset; j++) { \ + uint32_t even = s->c[j]; \ + uint32_t odd = s->c[j + offset]; \ + s->c[j] = odd + even; \ + /* `reduce_montgomery_partial` works on values up to kPrime*R. \ + * `step_root < kPrime` because it's static data. The other term is \ + * less than R because it's 32-bit, and does not overflow because \ + * `odd` is at most 128 * kPrime per the above. */ \ + s->c[j + offset] = reduce_montgomery_partial( \ + uint64_t{step_root} * uint64_t{128 * kPrime + even - odd}); \ + } \ + k += 2 * offset; \ + } \ } ITER(128, 1) @@ -456,7 +493,8 @@ #undef ITER // `kInverseDegreeDoubleMontgomery` has been double-converted to Montgomery - // form, so the result is an extra R term in the final result. + // form, so the result is an extra R term in the final result. This also fully + // reduces the partially-reduced values above. for (int i = 0; i < kDegree; i++) { s->c[i] = reduce_montgomery(uint64_t{s->c[i]} * uint64_t{kInverseDegreeDoubleMontgomery});