Skip to content

Repository navigation

Accelerate small-prime CRT reconstruction and BN254 export #175

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
odraccir1 merged 1 commit into from
Oct 16, 2026
Merged

Accelerate small-prime CRT reconstruction and BN254 export #175

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
79 changes: 75 additions & 4 deletions libs/math/include/nil/crypto3/math/detail/bn254_small_prime_fft.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<boost::multiprecision::cpp_int_modular_backend<crt_bits>>;
static constexpr unsigned export_split = 384;
using export_low_type =
boost::multiprecision::number<boost::multiprecision::cpp_int_modular_backend<export_split>>;
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;
Expand Down Expand Up @@ -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;
}
Expand Down Expand Up @@ -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<limb_type>(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<bool UseFma>
void reconstruct_impl(std::span<value_type> output, std::span<const double> 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);
Expand All @@ -304,7 +351,29 @@ namespace nil::crypto3::math::detail {
const double modulus = static_cast<double>(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<double, 4> 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<std::uint64_t>(canonical[lane]));
}
}
}
}
#endif
for (; i < count; ++i) {
double digit;
if constexpr (UseFma) {
digit = arithmetic.multiply_fma(arithmetic.reduce_fma(source[i]), context.inverse);
Expand All @@ -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<std::uint64_t>(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) {
Expand All @@ -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;
Expand All @@ -347,6 +416,8 @@ namespace nil::crypto3::math::detail {
std::array<integer_type, correction_count> multiples_;
integer_type half_modulus_;
integer_type bound_;
std::size_t accumulator_limbs_;
export_type export_radix_;
};

/**
Expand Down
111 changes: 111 additions & 0 deletions libs/math/test/small_prime_fft_crt.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<cpp_int> 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<cpp_int> 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_type> 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<cpp_int> 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_type> 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<const coordinates_type>(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_type, 1> coordinates {};
for (int sign : {-1, 1}) {
coordinates[0][0] = sign * bound;
check_reconstruction(reconstructor, coordinates, 1, true);
const std::vector<cpp_int> 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<value_type, 1> 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);
Expand Down
Loading