diff --git a/libs/math/include/nil/crypto3/math/detail/bn254_small_prime_fft.hpp b/libs/math/include/nil/crypto3/math/detail/bn254_small_prime_fft.hpp index 94f3032c4..1592c38a0 100644 --- a/libs/math/include/nil/crypto3/math/detail/bn254_small_prime_fft.hpp +++ b/libs/math/include/nil/crypto3/math/detail/bn254_small_prime_fft.hpp @@ -198,11 +198,17 @@ namespace nil::crypto3::math::detail { * This operation has no constant-time guarantee. */ class bn254_small_prime_fft_reconstructor { + using base_field_type = algebra::fields::alt_bn128_base_field<254>; static constexpr std::size_t block_size = 32; static constexpr std::size_t correction_count = std::bit_width(small_prime_fft_primes.size() - 1); // Each CRT term is < M, so their sum is < prime_count*M < 12*2^600 < 2^604. static constexpr unsigned crt_bits = 50 * small_prime_fft_primes.size() + correction_count; using integer_type = boost::multiprecision::number>; + static constexpr unsigned export_split = 384; + using export_low_type = + boost::multiprecision::number>; + using export_type = boost::multiprecision::uint512_modular_t; + static_assert(50 * small_prime_fft_primes.size() <= 600, "CRT export folding requires M < 2^600"); struct prime_context { small_prime_fft_arithmetic arithmetic; @@ -242,6 +248,10 @@ namespace nil::crypto3::math::detail { multiples_[0].backend().from_cpp_int(setup_backend(product.backend())); bound_.backend().from_cpp_int(setup_backend(coefficient_bound.backend())); half_modulus_ = multiples_[0] >> 1; + const boost::multiprecision::cpp_int sum_bound = params.prime_count * product - 1; + constexpr unsigned limb_bits = integer_type::backend_type::limb_bits; + accumulator_limbs_ = (boost::multiprecision::msb(sum_bound) + 1 + limb_bits - 1) / limb_bits; + export_radix_ = (export_type(1) << export_split) % export_type(base_field_type::modulus); for (std::size_t bit = 1; bit < correction_count; ++bit) { multiples_[bit] = multiples_[bit - 1] << 1; } @@ -290,10 +300,47 @@ namespace nil::crypto3::math::detail { } private: + void add_cofactor_product(integer_type &sum, const integer_type &cofactor, std::uint64_t digit) const { + using limb_type = boost::multiprecision::limb_type; + using wide_type = boost::multiprecision::double_limb_type; + constexpr unsigned limb_bits = integer_type::backend_type::limb_bits; + if constexpr (limb_bits >= 50) { + auto *destination = sum.backend().limbs(); + const auto *source = cofactor.backend().limbs(); + wide_type carry = 0; + // A digit is < 2^50. Each product plus the accumulator limb and carry fits + // in a double limb. The complete sum is < prime_count*M, so its final carry + // fits in the active width prepared by the constructor; higher limbs stay zero. + for (std::size_t limb = 0; limb < accumulator_limbs_; ++limb) { + const wide_type value = wide_type(source[limb]) * digit + destination[limb] + carry; + destination[limb] = static_cast(value); + carry = value >> limb_bits; + } + assert(carry == 0); + assert((destination[sum.backend().size() - 1] & ~integer_type::backend_type::upper_limb_mask) == 0); + } else { + // A 50-bit digit spans more than one limb on a 32-bit limb backend. + sum += cofactor * digit; + } + } + + base_field_type::value_type export_magnitude(const integer_type &magnitude) const { + // Centering and the coefficient-bound check have already established magnitude < M < 2^600. + // Write magnitude = low + high*2^384 and replace 2^384 by its cached residue modulo p. + // The fold is < 2^384 + 2^216*2^254 < 2^471, within the existing BN254 Barrett reducer's + // range. This avoids its general-division fallback while retaining the field's normal + // Montgomery conversion. The fixed-width low conversion intentionally discards high bits. + const export_type high(magnitude >> export_split); + export_type folded {export_low_type(magnitude)}; + if (high != 0) { + folded += high * export_radix_; + } + return base_field_type::value_type(folded); + } + template void reconstruct_impl(std::span output, std::span residues, std::size_t stride) const { - using base_field_type = algebra::fields::alt_bn128_base_field<254>; for (std::size_t coordinate = 0; coordinate < field_type::arity; ++coordinate) { for (std::size_t offset = 0; offset < output.size();) { const std::size_t count = std::min(block_size, output.size() - offset); @@ -304,7 +351,29 @@ namespace nil::crypto3::math::detail { const double modulus = static_cast(arithmetic.modulus()); const double *source = residues.data() + (prime * field_type::arity + coordinate) * stride + offset; - for (std::size_t i = 0; i < count; ++i) { + std::size_t i = 0; +#if defined(CRYPTO3_SMALL_PRIME_FFT_AVX2_SELECTED) + if constexpr (UseFma) { + if (count >= 4) { + const small_prime_fft_arithmetic::packed_fma_kernel packed(arithmetic); + const __m256d inverse = _mm256_set1_pd(context.inverse); + const __m256d q = _mm256_set1_pd(modulus); + for (; count - i >= 4; i += 4) { + __m256d digits = + packed.multiply_fma(packed.reduce_fma(_mm256_loadu_pd(source + i)), inverse); + const __m256d negative = _mm256_cmp_pd(digits, _mm256_setzero_pd(), _CMP_LT_OQ); + digits = _mm256_add_pd(digits, _mm256_and_pd(negative, q)); + std::array canonical; + _mm256_storeu_pd(canonical.data(), digits); + for (std::size_t lane = 0; lane < canonical.size(); ++lane) { + add_cofactor_product(sums[i + lane], context.cofactor, + static_cast(canonical[lane])); + } + } + } + } +#endif + for (; i < count; ++i) { double digit; if constexpr (UseFma) { digit = arithmetic.multiply_fma(arithmetic.reduce_fma(source[i]), context.inverse); @@ -315,7 +384,7 @@ namespace nil::crypto3::math::detail { // Both operands have magnitude < q, so the product remainder is in (-q,q). // Canonical digits give 0 <= digit*(M/q) < M and bound the complete sum. const auto canonical = static_cast(digit < 0 ? digit + modulus : digit); - sums[i] += context.cofactor * canonical; + add_cofactor_product(sums[i], context.cofactor, canonical); } } for (std::size_t i = 0; i < count; ++i) { @@ -335,7 +404,7 @@ namespace nil::crypto3::math::detail { throw std::overflow_error( "bn254_small_prime_fft_reconstructor: coefficient exceeds CRT bound"); } - const base_field_type::value_type value(magnitude); + const auto value = export_magnitude(magnitude); output[offset + i].coordinate(coordinate) = negative ? -value : value; } offset += count; @@ -347,6 +416,8 @@ namespace nil::crypto3::math::detail { std::array multiples_; integer_type half_modulus_; integer_type bound_; + std::size_t accumulator_limbs_; + export_type export_radix_; }; /** diff --git a/libs/math/test/small_prime_fft_crt.cpp b/libs/math/test/small_prime_fft_crt.cpp index d2b4f4d09..3faf7ccad 100644 --- a/libs/math/test/small_prime_fft_crt.cpp +++ b/libs/math/test/small_prime_fft_crt.cpp @@ -200,6 +200,117 @@ BOOST_AUTO_TEST_CASE(random_signed_integers_cover_all_prime_counts_and_partial_b } } +BOOST_AUTO_TEST_CASE(wide_export_preserves_split_boundaries_and_field_multiples) { + const cpp_int characteristic(fq::modulus.backend().to_cpp_int()); + for (std::size_t count = 8; count <= detail::small_prime_fft_primes.size(); ++count) { + const auto params = make_params(count); + const cpp_int bound = params.modulus / 2; + const reconstructor_type reconstructor(params, bound); + std::vector values {0, bound, -bound}; + for (unsigned bit : {254, 256, 384, 507, 508, 512}) { + const cpp_int power = cpp_int(1) << bit; + const cpp_int multiple = (power / characteristic) * characteristic; + for (const cpp_int &anchor : {power, multiple}) { + for (int delta : {-1, 0, 1}) { + const cpp_int value = anchor + delta; + if (value >= 0 && value <= bound) { + values.push_back(value); + values.push_back(-value); + } + } + } + } + // Exercise nonzero high and low parts together, including the largest high part. + const cpp_int max_high = bound >> 384; + const std::vector lows { + 0, 1, characteristic - 1, characteristic, characteristic + 1, (cpp_int(1) << 384) - 1}; + for (const cpp_int &high : {cpp_int(1), max_high}) { + for (const cpp_int &low : lows) { + const cpp_int value = (high << 384) + low; + if (value <= bound) { + values.push_back(value); + values.push_back(-value); + } + } + } + std::vector coordinates(values.size()); + for (std::size_t i = 0; i < coordinates.size(); ++i) { + for (std::size_t c = 0; c < 12; ++c) { + coordinates[i][c] = values[(i + c) % values.size()]; + } + } + check_reconstruction(reconstructor, coordinates, coordinates.size() + 3, true); + } +} + +BOOST_AUTO_TEST_CASE(cofactor_accumulation_preserves_carries_for_every_active_width) { + for (std::size_t count = 1; count <= detail::small_prime_fft_primes.size(); ++count) { + const auto params = make_params(count); + const cpp_int bound = params.modulus / 2; + const reconstructor_type reconstructor(params, bound); + std::vector values; + for (unsigned bit : {0, 1, 31, 32, 48, 49}) { + for (std::size_t parity = 0; parity < 2; ++parity) { + cpp_int sum = 0; + for (std::size_t prime = 0; prime < count; ++prime) { + const auto q = detail::small_prime_fft_primes[prime].modulus; + const auto digit = prime % 2 == parity ? q - 1 : (std::uint64_t(1) << bit) - 1; + sum += (params.modulus / q) * digit; + } + cpp_int value = canonical(sum, params.modulus); + if (value > bound) { + value -= params.modulus; + } + values.push_back(value); + } + } + std::vector coordinates(values.size()); + for (std::size_t i = 0; i < coordinates.size(); ++i) { + for (std::size_t c = 0; c < 12; ++c) { + coordinates[i][c] = values[(i + c) % values.size()]; + } + } + check_reconstruction(reconstructor, coordinates, 32, true); + } +} + +BOOST_AUTO_TEST_CASE(packed_lanes_and_block_tails_preserve_prefixes_and_rounding_modes) { + restore_rounding restore; + for (std::size_t count = 1; count <= detail::small_prime_fft_primes.size(); ++count) { + const auto params = make_params(count); + const cpp_int bound = params.modulus / 2; + const reconstructor_type reconstructor(params, bound); + const auto coordinates = random_coordinates(67, bound); + for (int mode : {FE_TONEAREST, FE_TOWARDZERO}) { + BOOST_REQUIRE_EQUAL(std::fesetround(mode), 0); + for (const std::size_t size : {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 31, 32, 33, 34, 35, 63, 64, 65, 66, 67}) { + check_reconstruction(reconstructor, std::span(coordinates).first(size), 71, + true); + BOOST_CHECK_EQUAL(std::fegetround(), mode); + } + } + } +} + +BOOST_AUTO_TEST_CASE(wide_bound_violations_are_rejected_before_field_reduction) { + const auto params = make_params(12); + const cpp_int characteristic(fq::modulus.backend().to_cpp_int()); + const cpp_int bound = (cpp_int(1) << 384) - 1; + const reconstructor_type reconstructor(params, bound); + std::array coordinates {}; + for (int sign : {-1, 1}) { + coordinates[0][0] = sign * bound; + check_reconstruction(reconstructor, coordinates, 1, true); + const std::vector invalid {bound + 1, (bound / characteristic + 1) * characteristic}; + for (const cpp_int &value : invalid) { + coordinates[0][0] = sign * value; + const auto residues = make_residues(coordinates, 12, 1, true); + std::array output {value_type::one()}; + BOOST_CHECK_THROW(reconstructor.reconstruct(output, residues, 1), std::overflow_error); + } + } +} + BOOST_AUTO_TEST_CASE(workload_plan_bound_is_recovered_with_both_signs) { const auto plan = detail::make_bn254_small_prime_fft_params(14939, 14939); BOOST_REQUIRE(plan);