From 5175ec455e1d6ae28f2ea0f74cb5727524ba9776 Mon Sep 17 00:00:00 2001 From: Riccardo Abbate Date: Mon, 19 Oct 2026 15:32:24 -0400 Subject: [PATCH] direct scalar/linear products and shared transforms and fused matrix products --- .../math/detail/bn254_small_prime_fft.hpp | 69 ++-- .../math/polynomial/arithmetic/half_gcd.hpp | 25 ++ .../bn254_small_prime_fft_backend.hpp | 173 +++++++++- .../backends/polynomial_backend.hpp | 14 + libs/math/test/polynomial_gcd.cpp | 84 +++++ libs/math/test/small_prime_fft_backend.cpp | 313 +++++++++++++++++- 6 files changed, 643 insertions(+), 35 deletions(-) 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 1592c38a0..22ae7501b 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 @@ -783,7 +783,7 @@ namespace nil::crypto3::math::detail { * arithmetic to the target modulus and cannot represent these small-prime intermediates. * * @pre Input cells are integer-valued doubles of magnitude < 2^53. scale is an integer with - * |scale| < q/2, normally the FFT plan's centered 1/N. The two buffers are disjoint or identical. + * |scale| < q/2, normally the FFT plan's centered 1/N. Buffers are disjoint or identical. * The thread's rounding mode stays unchanged during the call. Output cells are in (-q,q). * No buffers are allocated and no constant-time guarantee is provided. */ @@ -910,9 +910,9 @@ namespace nil::crypto3::math::detail { void multiply(std::span left, std::span right, std::size_t stride, double scale) const { check_buffers(left, right, stride, scale); if (arithmetic_.fma_supported()) { - apply(left, right, scale); + apply(left, left, right, scale); } else { - apply(left, right, scale); + apply(left, left, right, scale); } } @@ -920,9 +920,24 @@ namespace nil::crypto3::math::detail { void square(std::span values, std::size_t stride, double scale) const { check_buffers(values, values, stride, scale); if (arithmetic_.fma_supported()) { - apply(values, values, scale); + apply(values, values, values, scale); } else { - apply(values, values, scale); + apply(values, values, values, scale); + } + } + + /** Compute scale*(a*b+c*d). Every buffer has 12*stride cells; output may equal any input. */ + void multiply_sum(std::span output, std::span a, std::span b, + std::span c, std::span d, std::size_t stride, + double scale) const { + check_buffers(output, a, stride, scale); + check_buffers(output, b, stride, scale); + check_buffers(output, c, stride, scale); + check_buffers(output, d, stride, scale); + if (arithmetic_.fma_supported()) { + apply(output, a, b, scale, c, d); + } else { + apply(output, a, b, scale, c, d); } } @@ -936,32 +951,36 @@ namespace nil::crypto3::math::detail { std::abs(scale) < static_cast(arithmetic_.modulus()) / 2); } - template - void apply(std::span left, std::span right, double scale) const { - // Derive the indexing bound from the buffer extent checked by both public entry points. - const std::size_t stride = left.size() / 12; + template + void apply(std::span output, std::span left, std::span right, double scale, + std::span extra_left = {}, std::span extra_right = {}) const { + // Derive the indexing bound from the buffer extent checked by the public entry points. + const std::size_t stride = output.size() / 12; std::size_t offset = 0; #if defined(CRYPTO3_SMALL_PRIME_FFT_AVX2_SELECTED) if constexpr (UseFma) { offset = stride - stride % 4; if (offset != 0) { const small_prime_fft_arithmetic::packed_fma_kernel packed(arithmetic_); - apply_range(packed, left, right, 0, offset, _mm256_set1_pd(scale)); + apply_range(packed, output, left, right, extra_left, extra_right, 0, offset, + _mm256_set1_pd(scale)); } } #endif - apply_range(arithmetic_, left, right, offset, stride, scale); + apply_range(arithmetic_, output, left, right, extra_left, extra_right, offset, stride, + scale); } - template - static void apply_range(const Arithmetic &arithmetic, std::span left, std::span right, - std::size_t begin, std::size_t end, Value scale) { + template + static void apply_range(const Arithmetic &arithmetic, std::span output, std::span left, + std::span right, std::span extra_left, + std::span extra_right, std::size_t begin, std::size_t end, Value scale) { constexpr std::size_t width = sizeof(Value) / sizeof(double); - const std::size_t stride = left.size() / 12; + const std::size_t stride = output.size() / 12; const kernel operations {arithmetic}; using tower_type = typename kernel::tower_type; for (std::size_t i = begin; i < end; i += width) { - tower_type a, b; + tower_type a, b, c, d; for (std::size_t w = 0; w < 2; ++w) { for (std::size_t v = 0; v < 3; ++v) { for (std::size_t u = 0; u < 2; ++u) { @@ -971,19 +990,31 @@ namespace nil::crypto3::math::detail { if constexpr (!Square) { b[w][v][u] = operations.reduce(right[offset]); } + if constexpr (Sum) { + c[w][v][u] = operations.reduce(extra_left[offset]); + d[w][v][u] = operations.reduce(extra_right[offset]); + } #if defined(CRYPTO3_SMALL_PRIME_FFT_AVX2_SELECTED) } else { a[w][v][u] = operations.reduce(_mm256_loadu_pd(left.data() + offset)); if constexpr (!Square) { b[w][v][u] = operations.reduce(_mm256_loadu_pd(right.data() + offset)); } + if constexpr (Sum) { + c[w][v][u] = operations.reduce(_mm256_loadu_pd(extra_left.data() + offset)); + d[w][v][u] = operations.reduce(_mm256_loadu_pd(extra_right.data() + offset)); + } #endif } } } } const tower_type result = [&]() { - if constexpr (Square) { + if constexpr (Sum) { + // Each product coordinate is in (-q,q); adding two is exact below 2*q. + // Reduce that sum before the usual inverse-transform normalization. + return operations.add(operations.multiply(a, b), operations.multiply(c, d)); + } else if constexpr (Square) { return operations.square(a); } else { return operations.multiply(a, b); @@ -996,10 +1027,10 @@ namespace nil::crypto3::math::detail { const std::size_t offset = (6 * w + 2 * v + u) * stride + i; const Value value = operations.multiply(result[w][v][u], scale); if constexpr (width == 1) { - left[offset] = value; + output[offset] = value; #if defined(CRYPTO3_SMALL_PRIME_FFT_AVX2_SELECTED) } else { - _mm256_storeu_pd(left.data() + offset, value); + _mm256_storeu_pd(output.data() + offset, value); #endif } } diff --git a/libs/math/include/nil/crypto3/math/polynomial/arithmetic/half_gcd.hpp b/libs/math/include/nil/crypto3/math/polynomial/arithmetic/half_gcd.hpp index 650c693c4..b057b1465 100644 --- a/libs/math/include/nil/crypto3/math/polynomial/arithmetic/half_gcd.hpp +++ b/libs/math/include/nil/crypto3/math/polynomial/arithmetic/half_gcd.hpp @@ -25,6 +25,7 @@ #ifndef CRYPTO3_MATH_POLYNOMIAL_HALF_GCD_HPP #define CRYPTO3_MATH_POLYNOMIAL_HALF_GCD_HPP +#include #include #include #include @@ -114,6 +115,16 @@ namespace nil::crypto3::math::detail { polynomial_arithmetic::polynomial_context &arithmetic_context) { using polynomial_type = typename Backend::polynomial_type; + const std::array outputs {&first_output, &second_output}; + const std::array entries {&matrix.entry_00, &matrix.entry_01, &matrix.entry_10, + &matrix.entry_11}; + const std::array pair {&first, &second}; + if constexpr (requires { arithmetic_context.template try_multiply_matrix<1>(outputs, entries, pair); }) { + if (arithmetic_context.template try_multiply_matrix<1>(outputs, entries, pair)) { + return; + } + } + polynomial_type product_00; polynomial_type product_01; polynomial_type product_10; @@ -172,6 +183,20 @@ namespace nil::crypto3::math::detail { polynomial_arithmetic::polynomial_context &arithmetic_context) { using polynomial_type = typename Backend::polynomial_type; + const std::array outputs {&output.entry_00, &output.entry_01, &output.entry_10, + &output.entry_11}; + const std::array left_entries {&left.entry_00, &left.entry_01, &left.entry_10, + &left.entry_11}; + const std::array right_entries {&right.entry_00, &right.entry_01, &right.entry_10, + &right.entry_11}; + if constexpr (requires { + arithmetic_context.template try_multiply_matrix<2>(outputs, left_entries, right_entries); + }) { + if (arithmetic_context.template try_multiply_matrix<2>(outputs, left_entries, right_entries)) { + return; + } + } + half_gcd_matrix result; polynomial_type first_product; polynomial_type second_product; diff --git a/libs/math/include/nil/crypto3/math/polynomial/backends/bn254_small_prime_fft_backend.hpp b/libs/math/include/nil/crypto3/math/polynomial/backends/bn254_small_prime_fft_backend.hpp index 848b248bb..8843a6a1e 100644 --- a/libs/math/include/nil/crypto3/math/polynomial/backends/bn254_small_prime_fft_backend.hpp +++ b/libs/math/include/nil/crypto3/math/polynomial/backends/bn254_small_prime_fft_backend.hpp @@ -26,6 +26,7 @@ #define CRYPTO3_MATH_BN254_SMALL_PRIME_FFT_BACKEND_HPP #include +#include #include #include #include @@ -37,6 +38,7 @@ #include #include +#include #include #include @@ -46,7 +48,8 @@ namespace nil::crypto3::math::polynomial_arithmetic { * Exact BN254 Fp12 polynomial products through small-prime transforms and signed CRT. * Configure a power-of-two maximum transform length. All smaller power-of-two plans are * cached at construction; ordinary products use the smallest length containing their full linear - * product. multiply_low first discards input coefficients that cannot affect the output. Prepared + * product. Scalar and linear operands use direct coefficient arithmetic. multiply_low first + * discards input coefficients that cannot affect the output. Prepared * cyclic remainder reconstruction can use a shorter transform once the exact quotient is known. * * One catalogue prefix and coefficient bound cover every supported product, allowing all @@ -141,11 +144,14 @@ namespace nil::crypto3::math::polynomial_arithmetic { * max_spectrum_bytes limits one spectrum at the configured maximum length. A second * spectrum, cached twiddle tables and coefficient output storage are additional allocations. * Invalid/unsupported lengths or an insufficient byte limit throw before allocating plans. + * max_matrix_spectrum_bytes limits the additional shared input spectra used by matrix + * products. Zero disables that optional path; exceeding its budget declines the operation. */ - explicit bn254_small_prime_fft_backend( - std::size_t max_transform_size, - std::size_t max_spectrum_bytes = std::numeric_limits::max()) : - bn254_small_prime_fft_backend(checked_params(max_transform_size, max_spectrum_bytes)) { + explicit bn254_small_prime_fft_backend(std::size_t max_transform_size, + std::size_t max_spectrum_bytes = std::numeric_limits::max(), + std::size_t max_matrix_spectrum_bytes = 256 * 1024 * 1024) : + bn254_small_prime_fft_backend(checked_params(max_transform_size, max_spectrum_bytes), + max_matrix_spectrum_bytes) { } std::size_t max_transform_size() const { @@ -173,6 +179,135 @@ namespace nil::crypto3::math::polynomial_arithmetic { } } + /** + * Multiply a row-major 2-by-2 matrix by a 2-by-Columns matrix, sharing input transforms and + * reconstructing each sum of two products once. Columns is one for a polynomial pair. + * Outputs must be distinct and all pointers nonnull; outputs may alias any input. + * Return false without changing outputs when the shape, CRT bound or workspace budget does + * not permit fusion. Small or heavily unbalanced products retain the ordinary dispatch. + */ + template + requires(Columns == 1 || Columns == 2) + bool try_multiply_matrix(const std::array &output, + const std::array &left, + const std::array &right) { + if (plans_.empty() || max_matrix_spectrum_bytes_ == 0) { + return false; + } + for (std::size_t i = 0; i < output.size(); ++i) { + if (output[i] == nullptr || + std::find(output.begin(), output.begin() + i, output[i]) != output.begin() + i) { + return false; + } + } + std::array operands; + std::copy(left.begin(), left.end(), operands.begin()); + std::copy(right.begin(), right.end(), operands.begin() + 4); + for (const auto *operand : operands) { + if (operand == nullptr || operand->size() <= 2 || math::is_zero(*operand)) { + return false; + } + } + + std::array result_sizes {}; + std::size_t max_product_size = 0; + std::size_t min_product_size = max_transform_size_; + const std::size_t term_limit = std::max(std::size_t(1), max_transform_size_ / 2); + for (std::size_t row = 0; row < 2; ++row) { + for (std::size_t column = 0; column < Columns; ++column) { + std::size_t terms = 0; + for (std::size_t inner = 0; inner < 2; ++inner) { + const std::size_t a = left[2 * row + inner]->size(); + const std::size_t b = right[Columns * inner + column]->size(); + if (a > max_transform_size_ || b > max_transform_size_ - (a - 1)) { + return false; + } + // Each output coefficient sums at most min(a,b) tower products from this + // convolution. Their combined bound must fit the existing signed CRT plan: + // 102 * (min(a,b) + min(c,d)) * (p-1)^2 <= its configured coefficient bound. + const std::size_t product_terms = std::min(a, b); + if (product_terms > term_limit - terms) { + return false; + } + terms += product_terms; + const std::size_t product_size = a + b - 1; + auto &result_size = result_sizes[row * Columns + column]; + result_size = std::max(result_size, product_size); + max_product_size = std::max(max_product_size, product_size); + min_product_size = std::min(min_product_size, product_size); + } + } + } + const std::size_t size = std::bit_ceil(max_product_size); + if (std::bit_ceil(min_product_size) < size / 2) { + return false; + } + + // Equal input objects share storage as well as their conversion and forward transforms. + std::array unique_operands {}; + std::array spectrum_indices; + std::size_t unique_count = 0; + for (std::size_t i = 0; i < operands.size(); ++i) { + auto end = unique_operands.begin() + unique_count; + auto found = std::find(unique_operands.begin(), end, operands[i]); + spectrum_indices[i] = found - unique_operands.begin(); + if (found == end) { + unique_operands[unique_count++] = operands[i]; + } + } + const std::size_t cells = field_type::arity * prime_count() * size; + if (cells > max_matrix_spectrum_bytes_ / sizeof(double) / unique_count) { + return false; + } + const std::size_t total_cells = unique_count * cells; + if (matrix_spectra_.capacity() < total_cells) { + // Avoid geometric vector growth beyond the configured spectrum budget. + std::vector(total_cells).swap(matrix_spectra_); + } else { + matrix_spectra_.resize(total_cells); + } + for (std::size_t i = 0; i < unique_count; ++i) { + const auto &operand = *unique_operands[i]; + prepare_spectrum(std::span(matrix_spectra_).subspan(i * cells, cells), + std::span(operand.data(), operand.size()), size); + } + + left_spectrum_.resize(cells); + std::array results; + const std::size_t plan_offset = std::countr_zero(size) * prime_count(); + for (std::size_t row = 0; row < 2; ++row) { + for (std::size_t column = 0; column < Columns; ++column) { + for (std::size_t prime = 0; prime < prime_count(); ++prime) { + const auto &plan = plans_[plan_offset + prime]; + const std::size_t count = field_type::arity * size; + const std::size_t offset = prime * count; + const auto spectrum = [&](std::size_t operand) { + return std::span(matrix_spectra_) + .subspan(spectrum_indices[operand] * cells + offset, count); + }; + auto destination = std::span(left_spectrum_).subspan(offset, count); + pointwise_[prime].multiply_sum(destination, spectrum(2 * row), spectrum(4 + column), + spectrum(2 * row + 1), spectrum(4 + Columns + column), size, + plan.size_inverse()); + for (std::size_t coordinate = 0; coordinate < field_type::arity; ++coordinate) { + plan.inverse_fft(destination.subspan(coordinate * size, size), false); + } + } + const std::size_t index = row * Columns + column; + auto &result = results[index]; + result.resize(result_sizes[index]); + reconstructor_.reconstruct(std::span(result.data(), result.size()), left_spectrum_, + size); + math::condense(result); + } + } + // All inputs have been consumed and all allocations/reconstruction have succeeded. + for (std::size_t i = 0; i < output.size(); ++i) { + output[i]->get_storage().swap(results[i].get_storage()); + } + return true; + } + /** * Prepare the fixed right operand for multiply_low, selecting the same prefixes and smallest * transform as an ordinary product. Positive precision with no variable coefficients returns @@ -323,15 +458,20 @@ namespace nil::crypto3::math::polynomial_arithmetic { std::vector prepare_spectrum(std::span coefficients, std::size_t size) const { // Allocate just this spectrum, so a small preparation never retains larger scratch capacity. std::vector spectrum(field_type::arity * prime_count() * size); + prepare_spectrum(spectrum, coefficients, size); + return spectrum; + } + + void prepare_spectrum(std::span spectrum, std::span coefficients, + std::size_t size) const { converter_.convert(spectrum, coefficients, size); const std::size_t plan_offset = std::countr_zero(size) * prime_count(); for (std::size_t prime = 0; prime < prime_count(); ++prime) { for (std::size_t coordinate = 0; coordinate < field_type::arity; ++coordinate) { plans_[plan_offset + prime].fft( - std::span(spectrum).subspan((prime * field_type::arity + coordinate) * size, size)); + spectrum.subspan((prime * field_type::arity + coordinate) * size, size)); } } - return spectrum; } static params_type checked_params(std::size_t max_size, std::size_t max_spectrum_bytes) { @@ -349,9 +489,9 @@ namespace nil::crypto3::math::polynomial_arithmetic { return std::move(*params); } - explicit bn254_small_prime_fft_backend(const params_type ¶ms) : - max_transform_size_(params.transform_size), converter_(params.crt.prime_count), - reconstructor_(params.crt, params.coefficient_bound) { + explicit bn254_small_prime_fft_backend(const params_type ¶ms, std::size_t max_matrix_spectrum_bytes) : + max_transform_size_(params.transform_size), max_matrix_spectrum_bytes_(max_matrix_spectrum_bytes), + converter_(params.crt.prime_count), reconstructor_(params.crt, params.coefficient_bound) { const unsigned depth = std::countr_zero(max_transform_size_); plans_.reserve((depth + 1) * prime_count()); pointwise_.reserve(prime_count()); @@ -383,8 +523,17 @@ namespace nil::crypto3::math::polynomial_arithmetic { throw std::invalid_argument("bn254_small_prime_fft_backend: product exceeds configured transform size"); } const std::size_t product_size = left_size + right_size - 1; - const std::size_t size = std::bit_ceil(product_size); const std::size_t result_size = std::min(product_size, coefficient_count); + // Euclidean quotients are commonly linear. Their matrix updates need only one or two + // coefficient products per output, without conversion, transforms or reconstruction. + // Keep the capacity check above: direct arithmetic has the same supported shapes. + if (std::min(left_size, right_size) <= 2) { + const bool left_is_shorter = left_size <= right_size; + schoolbook_backend {}.multiply_low(output, left_is_shorter ? left : right, + left_is_shorter ? right : left, result_size); + return; + } + const std::size_t size = std::bit_ceil(product_size); const std::size_t channels = field_type::arity * prime_count(); // Constructor checks guarantee that even the largest spectrum's byte count fits size_t. left_spectrum_.resize(channels * size); @@ -448,12 +597,14 @@ namespace nil::crypto3::math::polynomial_arithmetic { } std::size_t max_transform_size_; + std::size_t max_matrix_spectrum_bytes_; math::detail::bn254_small_prime_fft_converter converter_; math::detail::bn254_small_prime_fft_reconstructor reconstructor_; std::vector plans_; std::vector pointwise_; std::vector left_spectrum_; std::vector right_spectrum_; + std::vector matrix_spectra_; polynomial_type result_; }; diff --git a/libs/math/include/nil/crypto3/math/polynomial/backends/polynomial_backend.hpp b/libs/math/include/nil/crypto3/math/polynomial/backends/polynomial_backend.hpp index f628ca09d..2884eb897 100644 --- a/libs/math/include/nil/crypto3/math/polynomial/backends/polynomial_backend.hpp +++ b/libs/math/include/nil/crypto3/math/polynomial/backends/polynomial_backend.hpp @@ -25,6 +25,7 @@ #ifndef CRYPTO3_MATH_POLYNOMIAL_BACKEND_HPP #define CRYPTO3_MATH_POLYNOMIAL_BACKEND_HPP +#include #include #include #include @@ -133,6 +134,19 @@ namespace nil::crypto3::math::polynomial_arithmetic { backend_.multiply_low(output, left, right, coefficient_count); } + // Optional row-major 2-by-2 times 2-by-Columns product. A declined operation leaves every + // output unchanged, so algorithms can use their ordinary multiplication/addition fallback. + template + bool try_multiply_matrix(const std::array &output, + const std::array &left, + const std::array &right) + requires requires(backend_type &backend) { + { backend.template try_multiply_matrix(output, left, right) } -> std::same_as; + } + { + return backend_.template try_multiply_matrix(output, left, right); + } + // Optional capabilities: backends without prepared low products still satisfy PolynomialBackend. // Forward the backend's owned preparation without adding another context or exposing its scratch storage. auto prepare_low_product(const polynomial_type &fixed, std::size_t variable_coefficient_count, diff --git a/libs/math/test/polynomial_gcd.cpp b/libs/math/test/polynomial_gcd.cpp index 37b895a91..3ec06c235 100644 --- a/libs/math/test/polynomial_gcd.cpp +++ b/libs/math/test/polynomial_gcd.cpp @@ -24,8 +24,10 @@ #define BOOST_TEST_MODULE polynomial_gcd_test +#include #include #include +#include #include #include @@ -36,6 +38,7 @@ #include #include +#include #include #include @@ -49,6 +52,35 @@ namespace { using fq12_field_type = fields::fp12_2over3over2>; using fq12_value_type = fq12_field_type::value_type; + struct matrix_product_counts { + std::array accepted {}; + std::size_t declined = 0; + }; + + struct counting_matrix_backend : polynomial_arithmetic::bn254_small_prime_fft_backend { + using base_type = polynomial_arithmetic::bn254_small_prime_fft_backend; + + counting_matrix_backend(std::size_t max_size, matrix_product_counts &counts, std::size_t byte_limit) : + base_type(max_size, std::numeric_limits::max(), byte_limit), counts(counts) { + } + + template + requires(Columns == 1 || Columns == 2) + bool try_multiply_matrix(const std::array &output, + const std::array &left, + const std::array &right) { + const bool accepted = base_type::try_multiply_matrix(output, left, right); + if (accepted) { + ++counts.accepted[Columns - 1]; + } else { + ++counts.declined; + } + return accepted; + } + + matrix_product_counts &counts; + }; + fq12_value_type fq12_value(std::size_t first_coordinate) { fq12_value_type value = fq12_value_type::zero(); for (std::size_t i = 0; i < fq12_field_type::arity; ++i) { @@ -330,4 +362,56 @@ BOOST_AUTO_TEST_CASE(gcd_is_alias_safe_when_half_gcd_is_forced) { BOOST_CHECK(right_alias == expected); } +BOOST_AUTO_TEST_CASE(small_prime_half_gcd_uses_fused_matrices_and_preserves_the_fallback) { + using reference_backend = polynomial_arithmetic::schoolbook_backend; + using polynomial_type = reference_backend::polynomial_type; + polynomial_arithmetic::polynomial_context euclidean_context; + polynomial_arithmetic::polynomial_context_options options; + options.gcd_half_gcd_cutoff = 1; + options.half_gcd_basecase_cutoff = 8; + matrix_product_counts fused_counts, fallback_counts; + polynomial_arithmetic::polynomial_context fused_context( + counting_matrix_backend(1024, fused_counts, 32 * 1024 * 1024), options); + polynomial_arithmetic::polynomial_context fallback_context( + counting_matrix_backend(1024, fallback_counts, 0), options); + + for (const std::size_t common_size : {1, 5}) { + const polynomial_type common = fq12_polynomial(common_size, 0x123456789abcdefULL); + const polynomial_type left = product(reference_backend {}, common, fq12_polynomial(129, 0x31415926)); + const polynomial_type right = product(reference_backend {}, common, fq12_polynomial(128, 0xfedcba98)); + polynomial_type expected, actual; + math::gcd(expected, left, right, euclidean_context); + math::gcd(actual, left, right, fused_context); + BOOST_CHECK(actual == expected); + math::gcd(actual, left, right, fallback_context); + BOOST_CHECK(actual == expected); + + auto left_alias = left; + auto right_alias = right; + math::gcd(left_alias, left_alias, right, fused_context); + math::gcd(right_alias, left, right_alias, fused_context); + BOOST_CHECK(left_alias == expected); + BOOST_CHECK(right_alias == expected); + + math::detail::half_gcd_matrix matrix; + polynomial_type reduced_first, reduced_second; + math::detail::half_gcd_reduce(matrix, reduced_first, reduced_second, left, right, + options.half_gcd_basecase_cutoff, fused_context); + BOOST_CHECK(math::is_zero(reduced_second) || reduced_second.size() <= left.size() / 2); + polynomial_type expected_first, expected_second; + math::addition(expected_first, product(reference_backend {}, matrix.entry_00, left), + product(reference_backend {}, matrix.entry_01, right)); + math::addition(expected_second, product(reference_backend {}, matrix.entry_10, left), + product(reference_backend {}, matrix.entry_11, right)); + BOOST_CHECK(reduced_first == expected_first); + BOOST_CHECK(reduced_second == expected_second); + } + BOOST_CHECK_GT(fused_counts.accepted[0], 0); + BOOST_CHECK_GT(fused_counts.accepted[1], 0); + BOOST_CHECK_GT(fused_counts.declined, 0); + BOOST_CHECK_EQUAL(fallback_counts.accepted[0], 0); + BOOST_CHECK_EQUAL(fallback_counts.accepted[1], 0); + BOOST_CHECK_GT(fallback_counts.declined, 0); +} + BOOST_AUTO_TEST_SUITE_END() diff --git a/libs/math/test/small_prime_fft_backend.cpp b/libs/math/test/small_prime_fft_backend.cpp index 2958b6fe6..dfe37431a 100644 --- a/libs/math/test/small_prime_fft_backend.cpp +++ b/libs/math/test/small_prime_fft_backend.cpp @@ -63,6 +63,16 @@ namespace { static_assert(math::detail::SupportsPreparedLowProduct); static_assert(math::detail::SupportsPreparedCyclicRemainder); + template + concept SupportsMatrixProduct = + requires(Context &context, std::array output, + std::array left, std::array right) { + { context.template try_multiply_matrix(output, left, right) } -> std::same_as; + }; + static_assert(SupportsMatrixProduct && SupportsMatrixProduct); + static_assert(!SupportsMatrixProduct); + static_assert(!SupportsMatrixProduct, 2>); + struct prepared_product_counts { std::size_t preparations = 0; std::size_t prepared_products = 0; @@ -198,6 +208,37 @@ namespace { return result; } + template + std::array + expected_matrix_product(const std::array &left, + const std::array &right) { + std::array result; + for (std::size_t row = 0; row < 2; ++row) { + for (std::size_t column = 0; column < Columns; ++column) { + math::addition(result[row * Columns + column], expected_product(left[2 * row], right[column]), + expected_product(left[2 * row + 1], right[Columns + column])); + } + } + return result; + } + + template + bool try_matrix_product(backend_type &backend, std::array &output, + const std::array &left, + const std::array &right) { + std::array outputs; + std::array left_entries; + std::array right_entries; + for (std::size_t i = 0; i < left.size(); ++i) { + left_entries[i] = &left[i]; + } + for (std::size_t i = 0; i < output.size(); ++i) { + outputs[i] = &output[i]; + right_entries[i] = &right[i]; + } + return backend.try_multiply_matrix(outputs, left_entries, right_entries); + } + polynomial_type expected_low(polynomial_type product, std::size_t count) { if (count == 0) { return {value_type::zero()}; @@ -296,30 +337,292 @@ BOOST_AUTO_TEST_CASE(pointwise_invalid_dimensions_throw_before_writes) { BOOST_CHECK_THROW(pointwise.multiply(left, right, std::numeric_limits::max(), 1.0), std::invalid_argument); BOOST_CHECK_THROW(pointwise.square(left, std::numeric_limits::max(), 1.0), std::invalid_argument); + BOOST_CHECK_THROW(pointwise.multiply_sum(left, right, right, right, right, 1, 1.0), std::invalid_argument); + BOOST_CHECK_THROW( + pointwise.multiply_sum(left, right, right, right, std::span(right).first(23), 2, 1.0), + std::invalid_argument); BOOST_CHECK_EQUAL_COLLECTIONS(left.begin(), left.end(), before.begin(), before.end()); BOOST_CHECK_NO_THROW(pointwise.multiply({}, {}, 0, 1.0)); BOOST_CHECK_NO_THROW(pointwise.square({}, 0, 1.0)); + BOOST_CHECK_NO_THROW(pointwise.multiply_sum({}, {}, {}, {}, {}, 0, 1.0)); +} + +BOOST_AUTO_TEST_CASE(pointwise_sums_match_expanded_integer_tower_and_preserve_aliases) { + restore_rounding restore; + BOOST_REQUIRE_EQUAL(std::fesetround(FE_TONEAREST), 0); + std::mt19937_64 random(0x46465421); + for (const auto &prime : detail::small_prime_fft_primes) { + const auto q = static_cast(prime.modulus); + const detail::bn254_small_prime_fft_pointwise pointwise(prime.modulus); + for (const std::size_t stride : {0, 1, 2, 3, 4, 5, 6, 7, 8, 9}) { + std::array, 4> inputs; + for (auto &input : inputs) { + input.resize(12 * stride); + } + std::vector expected(stride); + for (std::size_t i = 0; i < stride; ++i) { + std::array values; + for (std::size_t operand = 0; operand < inputs.size(); ++operand) { + for (std::size_t c = 0; c < 12; ++c) { + std::int64_t value = static_cast(random() % (6 * q - 1)) - (3 * q - 1); + if ((c + i + operand) % 7 == 0) { + value = (std::int64_t(1) << 53) - 1; + if (operand % 2) { + value = -value; + } + } + values[operand][c] = value; + inputs[operand][c * stride + i] = static_cast(value); + } + } + expected[i] = tower_product(values[0], values[1]); + const auto second = tower_product(values[2], values[3]); + for (std::size_t c = 0; c < 12; ++c) { + expected[i][c] += second[c]; + } + } + for (const double scale : + {0.0, 1.0, -1.0, static_cast((q - 1) / 2), -static_cast((q - 1) / 2)}) { + std::vector storage(12 * stride + 2, -19.0); + auto output = std::span(storage).subspan(1, 12 * stride); + pointwise.multiply_sum(output, inputs[0], inputs[1], inputs[2], inputs[3], stride, scale); + check_pointwise(output, expected, prime.modulus, scale); + BOOST_CHECK_EQUAL(storage.front(), -19.0); + BOOST_CHECK_EQUAL(storage.back(), -19.0); + for (std::size_t alias = 0; alias < inputs.size(); ++alias) { + auto aliased = inputs; + pointwise.multiply_sum(aliased[alias], aliased[0], aliased[1], aliased[2], aliased[3], stride, + scale); + check_pointwise(aliased[alias], expected, prime.modulus, scale); + } + } + } + } +} + +BOOST_AUTO_TEST_CASE(fused_matrix_products_match_schoolbook_and_preserve_aliases) { + backend_type backend(256); + std::mt19937_64 random(0x46465422); + for (const std::size_t size : {3, 17, 64, 8}) { + std::array left, right, actual; + for (std::size_t i = 0; i < left.size(); ++i) { + left[i] = random_polynomial(size, random); + right[i] = random_polynomial(size, random); + } + const auto expected = expected_matrix_product<2>(left, right); + BOOST_REQUIRE(try_matrix_product<2>(backend, actual, left, right)); + for (std::size_t i = 0; i < actual.size(); ++i) { + check_polynomial(actual[i], expected[i]); + } + auto left_alias = left; + BOOST_REQUIRE(try_matrix_product<2>(backend, left_alias, left_alias, right)); + auto right_alias = right; + BOOST_REQUIRE(try_matrix_product<2>(backend, right_alias, left, right_alias)); + for (std::size_t i = 0; i < actual.size(); ++i) { + check_polynomial(left_alias[i], expected[i]); + check_polynomial(right_alias[i], expected[i]); + } + std::array pair {right[0], right[1]}; + const auto pair_expected = expected_matrix_product<1>(left, pair); + BOOST_REQUIRE(try_matrix_product<1>(backend, pair, left, pair)); + check_polynomial(pair[0], pair_expected[0]); + check_polynomial(pair[1], pair_expected[1]); + } + // Entries differing by one coefficient can straddle a power-of-two transform boundary. + std::array left, right, actual; + for (std::size_t i = 0; i < left.size(); ++i) { + left[i] = random_polynomial(16 + i % 2, random); + right[i] = random_polynomial(16 + (i / 2), random); + } + const auto expected = expected_matrix_product<2>(left, right); + BOOST_REQUIRE(try_matrix_product<2>(backend, actual, left, right)); + for (std::size_t i = 0; i < actual.size(); ++i) { + check_polynomial(actual[i], expected[i]); + } +} + +BOOST_AUTO_TEST_CASE(fused_matrix_products_check_the_sum_bound_and_workspace_before_writes) { + std::mt19937_64 random(0x46465423); + backend_type backend(64); + std::array left, right, output; + for (std::size_t i = 0; i < left.size(); ++i) { + left[i] = random_polynomial(16, random); + right[i] = random_polynomial(16, random); + } + const auto expected = expected_matrix_product<2>(left, right); + BOOST_REQUIRE(try_matrix_product<2>(backend, output, left, right)); + for (std::size_t i = 0; i < output.size(); ++i) { + check_polynomial(output[i], expected[i]); + } + const auto before = output; + const auto check_unchanged = [&]() { + for (std::size_t i = 0; i < output.size(); ++i) { + check_polynomial(output[i], before[i]); + } + }; + const std::size_t bytes = 8 * 12 * backend.prime_count() * 32 * sizeof(double); + backend_type exact_budget(64, std::numeric_limits::max(), bytes); + BOOST_REQUIRE(try_matrix_product<2>(exact_budget, output, left, right)); + check_unchanged(); + for (const std::size_t budget : {std::size_t(0), bytes - 1}) { + backend_type limited(64, std::numeric_limits::max(), budget); + BOOST_CHECK(!try_matrix_product<2>(limited, output, left, right)); + check_unchanged(); + } + // Both individual products fit, but 17+17 terms exceed the plan's 32-term CRT bound. + for (std::size_t i = 0; i < left.size(); ++i) { + left[i] = random_polynomial(17, random); + right[i] = random_polynomial(17, random); + } + BOOST_CHECK(!try_matrix_product<2>(backend, output, left, right)); + check_unchanged(); + for (const std::size_t size : {1, 2, 65}) { + left[0] = random_polynomial(size, random); + BOOST_CHECK(!try_matrix_product<2>(backend, output, left, right)); + check_unchanged(); + } + left[0] = polynomial_type(3, value_type::zero()); + BOOST_CHECK(!try_matrix_product<2>(backend, output, left, right)); + check_unchanged(); + + std::array outputs {&output[0], &output[1], &output[2], &output[3]}; + std::array inputs {&right[0], &right[1], &right[2], &right[3]}; + outputs[3] = outputs[0]; + BOOST_CHECK(!backend.try_multiply_matrix<2>(outputs, inputs, inputs)); + check_unchanged(); + outputs[3] = nullptr; + BOOST_CHECK(!backend.try_multiply_matrix<2>(outputs, inputs, inputs)); + check_unchanged(); + outputs[3] = &output[3]; + inputs[3] = nullptr; + BOOST_CHECK(!backend.try_multiply_matrix<2>(outputs, inputs, inputs)); + check_unchanged(); + + backend_type large_backend(512); + for (std::size_t i = 0; i < left.size(); ++i) { + left[i] = random_polynomial(i == 0 ? 3 : 40, random); + right[i] = random_polynomial(i == 0 ? 3 : 40, random); + } + BOOST_CHECK(!try_matrix_product<2>(large_backend, output, left, right)); + check_unchanged(); +} + +BOOST_AUTO_TEST_CASE(fused_matrix_products_share_identical_inputs_and_survive_backend_moves) { + std::mt19937_64 random(0x46465424); + polynomial_type input = random_polynomial(16, random); + const auto product = expected_product(input, input); + polynomial_type expected; + math::addition(expected, product, product); + backend_type ordinary(64); + const std::size_t one_spectrum_bytes = 12 * ordinary.prime_count() * 32 * sizeof(double); + backend_type backend(64, std::numeric_limits::max(), one_spectrum_bytes); + std::array output; + const std::array outputs {&input, &output[1], &output[2], &output[3]}; + const std::array inputs {&input, &input, &input, &input}; + BOOST_REQUIRE(backend.try_multiply_matrix<2>(outputs, inputs, inputs)); + check_polynomial(input, expected); + for (std::size_t i = 1; i < output.size(); ++i) { + check_polynomial(output[i], expected); + } + input = random_polynomial(16, random); + auto moved = std::move(backend); + BOOST_CHECK(!backend.try_multiply_matrix<2>(outputs, inputs, inputs)); + const auto square = expected_product(input, input); + math::addition(expected, square, square); + BOOST_REQUIRE(moved.try_multiply_matrix<2>(outputs, inputs, inputs)); + check_polynomial(input, expected); + for (std::size_t i = 1; i < output.size(); ++i) { + check_polynomial(output[i], expected); + } +} + +BOOST_AUTO_TEST_CASE(fused_matrix_products_preserve_cancellation_and_rounding_modes) { + restore_rounding restore; + std::mt19937_64 random(0x46465425); + backend_type backend(64); + std::array left, right, output; + left[0] = left[1] = random_polynomial(15, random); + left[2] = left[3] = random_polynomial(16, random); + right[0] = random_polynomial(17, random); + right[1] = random_polynomial(16, random); + const polynomial_type zero {value_type::zero()}; + math::subtraction(right[2], zero, right[0]); + math::subtraction(right[3], zero, right[1]); + for (const int mode : {FE_TONEAREST, FE_DOWNWARD, FE_UPWARD, FE_TOWARDZERO}) { + BOOST_REQUIRE_EQUAL(std::fesetround(mode), 0); + BOOST_REQUIRE(try_matrix_product<2>(backend, output, left, right)); + BOOST_CHECK_EQUAL(std::fegetround(), mode); + for (const auto &result : output) { + check_polynomial(result, zero); + } + } } BOOST_AUTO_TEST_CASE(all_tower_basis_products_cross_the_conversion_and_crt_pipeline) { - backend_type backend(1); + backend_type backend(8); polynomial_type actual; + // Degree-two monomials keep these checks on the transform path when scalar and linear + // operands use direct arithmetic. Their product has its only nonzero coefficient at X^4. for (std::size_t i = 0; i < 12; ++i) { coordinates_type a {}; a[i] = 1; - const polynomial_type left {field_value(a)}; + const polynomial_type left {value_type::zero(), value_type::zero(), field_value(a)}; for (std::size_t j = 0; j < 12; ++j) { coordinates_type b {}; b[j] = 1; - const polynomial_type right {field_value(b)}; + const polynomial_type right {value_type::zero(), value_type::zero(), field_value(b)}; backend.multiply(actual, left, right); - check_polynomial(actual, polynomial_type {field_value(tower_product(a, b))}); + polynomial_type expected(5, value_type::zero()); + expected[4] = field_value(tower_product(a, b)); + check_polynomial(actual, expected); } backend.square(actual, left); - check_polynomial(actual, polynomial_type {field_value(tower_product(a, a))}); + polynomial_type expected(5, value_type::zero()); + expected[4] = field_value(tower_product(a, a)); + check_polynomial(actual, expected); } } +BOOST_AUTO_TEST_CASE(short_operands_preserve_low_products_aliases_and_capacity) { + backend_type backend(512); + std::mt19937_64 random(0x46465440); + for (const std::size_t short_size : {1, 2, 3}) { + const auto short_operand = random_polynomial(short_size, random); + for (const std::size_t long_size : {17, 64, 255}) { + const auto long_operand = random_polynomial(long_size, random); + const auto full = expected_product(short_operand, long_operand); + for (const std::size_t count : + {std::size_t(0), std::size_t(1), std::size_t(2), std::size_t(3), std::size_t(17), full.size()}) { + const auto expected = expected_low(full, count); + auto left_alias = short_operand; + auto right_alias = long_operand; + backend.multiply_low(left_alias, left_alias, long_operand, count); + backend.multiply_low(right_alias, short_operand, right_alias, count); + check_polynomial(left_alias, expected); + check_polynomial(right_alias, expected); + backend.multiply_low(right_alias, long_operand, short_operand, count); + check_polynomial(right_alias, expected); + } + } + } + + backend_type limited(32); + const auto linear = random_polynomial(2, random); + const auto at_limit = random_polynomial(31, random); + polynomial_type result; + limited.multiply(result, linear, at_limit); + check_polynomial(result, expected_product(linear, at_limit)); + const auto before = result; + BOOST_CHECK_THROW(limited.multiply(result, linear, random_polynomial(32, random)), std::invalid_argument); + BOOST_CHECK_THROW( + limited.multiply_low(result, random_polynomial(33, random), polynomial_type {value_type::one()}, 33), + std::invalid_argument); + check_polynomial(result, before); + backend_type moved(std::move(limited)); + BOOST_CHECK_THROW(limited.multiply(result, linear, at_limit), std::invalid_argument); + check_polynomial(result, before); +} + BOOST_AUTO_TEST_CASE(full_width_products_match_native_fp12_schoolbook) { backend_type backend(512); std::mt19937_64 random(0x46465411);