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 c50ac286c..94f3032c4 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 @@ -356,6 +356,7 @@ namespace nil::crypto3::math::detail { * * A reusable table of max(1,N/2) centered twiddles serves both directions. Radix-4 recursion * fuses pairs of layers, with radix-2 leaves for odd depths and special cases for unit twiddles. + * AVX2/FMA builds process four independent butterflies at a time; short layers stay scalar. * Transform calls allocate no buffers; recursion uses O(log N) stack space. * * @pre Inputs are integer-valued doubles in (-3*q,3*q). The calling thread's rounding mode @@ -472,51 +473,79 @@ namespace nil::crypto3::math::detail { } // Keep arithmetic selection outside the butterfly loops, with compile-time dispatch inside them. - template - double reduce(double value) const { + template + static Value reduce(const Arithmetic &arithmetic, Value value) { if constexpr (UseFma) { - return arithmetic_.reduce_fma(value); + return arithmetic.reduce_fma(value); } else { - return arithmetic_.reduce_integer(value); + return arithmetic.reduce_integer(value); } } - template - double multiply(double left, double right) const { + template + static Value multiply(const Arithmetic &arithmetic, Value left, Value right) { if constexpr (UseFma) { - return arithmetic_.multiply_fma(left, right); + return arithmetic.multiply_fma(left, right); } else { - return arithmetic_.multiply_integer(left, right); + return arithmetic.multiply_integer(left, right); } } - template - void forward_radix4(double *values, std::size_t quarter, std::size_t node) const { + template + static void forward_butterfly(const Arithmetic &arithmetic, Value &v0, Value &v1, Value &v2, Value &v3, + Value w2, Value w, Value iw) { #if defined(__clang__) #pragma clang fp reassociate(off) #pragma clang fp contract(off) #endif + const Value a0 = reduce(arithmetic, v0); + const Value a1 = v1; + const Value a2 = UnitTwiddle ? reduce(arithmetic, v2) : multiply(arithmetic, v2, w2); + const Value a3 = UnitTwiddle ? reduce(arithmetic, v3) : multiply(arithmetic, v3, w2); + const Value even_sum = a0 + a2; + const Value even_difference = a0 - a2; + const Value odd_sum = + UnitTwiddle ? reduce(arithmetic, a1 + a3) : multiply(arithmetic, a1 + a3, w); + const Value odd_difference = multiply(arithmetic, a1 - a3, iw); + // Input magnitudes < 3*q and weights < q/2 imply |a2|,|a3| < q and + // |a1 +/- a3| < 4*q. Every product is < 2*q^2, so the odd terms are < q. + // Thus all outputs are < 3*q, with every integer sum exact below 2^53. + v0 = even_sum + odd_sum; + v1 = even_sum - odd_sum; + v2 = even_difference + odd_difference; + v3 = even_difference - odd_difference; + } + + template + void forward_radix4(double *values, std::size_t quarter, std::size_t node) const { const double w2 = twiddles_[node]; const double w = twiddles_[2 * node]; const double iw = twiddles_[2 * node + 1]; - for (std::size_t i = 0; i < quarter; ++i) { - const double a0 = reduce(values[i]); - const double a1 = values[i + quarter]; - const double a2 = UnitTwiddle ? reduce(values[i + 2 * quarter]) : - multiply(values[i + 2 * quarter], w2); - const double a3 = UnitTwiddle ? reduce(values[i + 3 * quarter]) : - multiply(values[i + 3 * quarter], w2); - const double even_sum = a0 + a2; - const double even_difference = a0 - a2; - const double odd_sum = UnitTwiddle ? reduce(a1 + a3) : multiply(a1 + a3, w); - const double odd_difference = multiply(a1 - a3, iw); - // Input magnitudes < 3*q and weights < q/2 imply |a2|,|a3| < q and - // |a1 +/- a3| < 4*q. Every product is < 2*q^2, so the odd terms are < q. - // Thus all outputs are < 3*q, with every integer sum exact below 2^53. - values[i] = even_sum + odd_sum; - values[i + quarter] = even_sum - odd_sum; - values[i + 2 * quarter] = even_difference + odd_difference; - values[i + 3 * quarter] = even_difference - odd_difference; + std::size_t i = 0; +#if defined(CRYPTO3_SMALL_PRIME_FFT_AVX2_SELECTED) + if constexpr (UseFma) { + if (quarter >= 4) { + const small_prime_fft_arithmetic::packed_fma_kernel packed(arithmetic_); + const __m256d packed_w2 = _mm256_set1_pd(w2); + const __m256d packed_w = _mm256_set1_pd(w); + const __m256d packed_iw = _mm256_set1_pd(iw); + for (; quarter - i >= 4; i += 4) { + __m256d v0 = _mm256_loadu_pd(values + i); + __m256d v1 = _mm256_loadu_pd(values + i + quarter); + __m256d v2 = _mm256_loadu_pd(values + i + 2 * quarter); + __m256d v3 = _mm256_loadu_pd(values + i + 3 * quarter); + forward_butterfly(packed, v0, v1, v2, v3, packed_w2, packed_w, packed_iw); + _mm256_storeu_pd(values + i, v0); + _mm256_storeu_pd(values + i + quarter, v1); + _mm256_storeu_pd(values + i + 2 * quarter, v2); + _mm256_storeu_pd(values + i + 3 * quarter, v3); + } + } + } +#endif + for (; i < quarter; ++i) { + forward_butterfly(arithmetic_, values[i], values[i + quarter], + values[i + 2 * quarter], values[i + 3 * quarter], w2, w, iw); } } @@ -526,9 +555,9 @@ namespace nil::crypto3::math::detail { return; } if (length == 2) { - const double left = reduce(values[0]); - const double right = - node == 0 ? reduce(values[1]) : multiply(values[1], twiddles_[node]); + const double left = reduce(arithmetic_, values[0]); + const double right = node == 0 ? reduce(arithmetic_, values[1]) : + multiply(arithmetic_, values[1], twiddles_[node]); values[0] = left + right; values[1] = left - right; return; @@ -546,50 +575,82 @@ namespace nil::crypto3::math::detail { } } - template - void inverse_radix4(double *values, std::size_t quarter, std::size_t node) const { + template + static void inverse_butterfly(const Arithmetic &arithmetic, Value &v0, Value &v1, Value &v2, Value &v3, + Value w2, Value w, Value iw, bool leaf) { #if defined(__clang__) #pragma clang fp reassociate(off) #pragma clang fp contract(off) #endif + // Leaf inputs may be forward residues of magnitude < 3*q. Reduce them before + // adding four entries; non-leaf inputs already satisfy the inverse bound < 2*q. + const Value b0 = leaf ? reduce(arithmetic, v0) : v0; + const Value b1 = leaf ? reduce(arithmetic, v1) : v1; + const Value b2 = leaf ? reduce(arithmetic, v2) : v2; + const Value b3 = leaf ? reduce(arithmetic, v3) : v3; + const Value even_sum = b0 + b1; + const Value odd_sum = b2 + b3; + const Value left = + UnitTwiddle ? reduce(arithmetic, b0 - b1) : multiply(arithmetic, b0 - b1, w); + const Value right = multiply(arithmetic, b2 - b3, iw); + // Four inputs < 2*q give sums < 8*q < 2^53. Times a centered weight this is + // < 4*q^2, so the a2 remainder is < 1.5*q. Left and right are each < q; + // their sum is < 2*q and their weighted difference is < q. The invariant closes. + v0 = reduce(arithmetic, even_sum + odd_sum); + v1 = left + right; + v2 = UnitTwiddle ? reduce(arithmetic, even_sum - odd_sum) : + multiply(arithmetic, even_sum - odd_sum, w2); + v3 = + UnitTwiddle ? reduce(arithmetic, left - right) : multiply(arithmetic, left - right, w2); + } + + template + void inverse_radix4(double *values, std::size_t quarter, std::size_t node) const { const double w2 = inverse_twiddle(node); const double w = inverse_twiddle(2 * node); const double iw = inverse_twiddle(2 * node + 1); - for (std::size_t i = 0; i < quarter; ++i) { - // Leaf inputs may be forward residues of magnitude < 3*q. Reduce them before - // adding four entries; non-leaf inputs already satisfy the inverse bound < 2*q. - const double b0 = quarter == 1 ? reduce(values[i]) : values[i]; - const double b1 = quarter == 1 ? reduce(values[i + quarter]) : values[i + quarter]; - const double b2 = quarter == 1 ? reduce(values[i + 2 * quarter]) : values[i + 2 * quarter]; - const double b3 = quarter == 1 ? reduce(values[i + 3 * quarter]) : values[i + 3 * quarter]; - const double even_sum = b0 + b1; - const double odd_sum = b2 + b3; - const double left = UnitTwiddle ? reduce(b0 - b1) : multiply(b0 - b1, w); - const double right = multiply(b2 - b3, iw); - // Four inputs < 2*q give sums < 8*q < 2^53. Times a centered weight this is - // < 4*q^2, so the a2 remainder is < 1.5*q. Left and right are each < q; - // their sum is < 2*q and their weighted difference is < q. The invariant closes. - values[i] = reduce(even_sum + odd_sum); - values[i + quarter] = left + right; - values[i + 2 * quarter] = - UnitTwiddle ? reduce(even_sum - odd_sum) : multiply(even_sum - odd_sum, w2); - values[i + 3 * quarter] = - UnitTwiddle ? reduce(left - right) : multiply(left - right, w2); + std::size_t i = 0; +#if defined(CRYPTO3_SMALL_PRIME_FFT_AVX2_SELECTED) + if constexpr (UseFma) { + if (quarter >= 4) { + const small_prime_fft_arithmetic::packed_fma_kernel packed(arithmetic_); + const __m256d packed_w2 = _mm256_set1_pd(w2); + const __m256d packed_w = _mm256_set1_pd(w); + const __m256d packed_iw = _mm256_set1_pd(iw); + for (; quarter - i >= 4; i += 4) { + __m256d v0 = _mm256_loadu_pd(values + i); + __m256d v1 = _mm256_loadu_pd(values + i + quarter); + __m256d v2 = _mm256_loadu_pd(values + i + 2 * quarter); + __m256d v3 = _mm256_loadu_pd(values + i + 3 * quarter); + inverse_butterfly(packed, v0, v1, v2, v3, packed_w2, packed_w, packed_iw, + false); + _mm256_storeu_pd(values + i, v0); + _mm256_storeu_pd(values + i + quarter, v1); + _mm256_storeu_pd(values + i + 2 * quarter, v2); + _mm256_storeu_pd(values + i + 3 * quarter, v3); + } + } + } +#endif + for (; i < quarter; ++i) { + inverse_butterfly(arithmetic_, values[i], values[i + quarter], + values[i + 2 * quarter], values[i + 3 * quarter], w2, w, iw, + quarter == 1); } } template void inverse_recursive(double *values, std::size_t length, std::size_t node) const { if (length == 1) { - values[0] = reduce(values[0]); + values[0] = reduce(arithmetic_, values[0]); return; } if (length == 2) { - const double left = reduce(values[0]); - const double right = reduce(values[1]); - values[0] = reduce(left + right); - values[1] = - node == 0 ? reduce(left - right) : multiply(left - right, inverse_twiddle(node)); + const double left = reduce(arithmetic_, values[0]); + const double right = reduce(arithmetic_, values[1]); + values[0] = reduce(arithmetic_, left + right); + values[1] = node == 0 ? reduce(arithmetic_, left - right) : + multiply(arithmetic_, left - right, inverse_twiddle(node)); return; } const std::size_t quarter = length / 4; @@ -610,8 +671,25 @@ namespace nil::crypto3::math::detail { inverse_recursive(values.data(), size_, 0); if (normalize) { const double q = static_cast(modulus()); - for (double &value : values) { - value = multiply(value, size_inverse_); + std::size_t i = 0; +#if defined(CRYPTO3_SMALL_PRIME_FFT_AVX2_SELECTED) + if constexpr (UseFma) { + if (values.size() >= 4) { + const small_prime_fft_arithmetic::packed_fma_kernel packed(arithmetic_); + const __m256d scale = _mm256_set1_pd(size_inverse_); + const __m256d packed_q = _mm256_set1_pd(q); + for (; values.size() - i >= 4; i += 4) { + __m256d value = packed.multiply_fma(_mm256_loadu_pd(values.data() + i), scale); + const __m256d negative = _mm256_cmp_pd(value, _mm256_setzero_pd(), _CMP_LT_OQ); + value = _mm256_add_pd(value, _mm256_and_pd(negative, packed_q)); + _mm256_storeu_pd(values.data() + i, value); + } + } + } +#endif + for (; i < values.size(); ++i) { + double &value = values[i]; + value = multiply(arithmetic_, value, size_inverse_); value = value < 0 ? value + q : value; } } @@ -630,7 +708,7 @@ namespace nil::crypto3::math::detail { * Coordinate channels have the same order as the converter, with one prime per call. * * Reuse the field library's Karatsuba tower identities and quadratic-extension squaring - * identity with the scalar residue kernels. The existing field elements bind their base + * identity with scalar or four-lane residue kernels. The existing field elements bind their base * 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 @@ -638,16 +716,22 @@ namespace nil::crypto3::math::detail { * 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. */ +#if defined(CRYPTO3_SMALL_PRIME_FFT_AVX2_SELECTED) +// The packed kernel uses __m256d as the coordinate type of the existing tower arrays. +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wignored-attributes" +#endif class bn254_small_prime_fft_pointwise { - using quadratic_type = std::array; - using cubic_type = std::array; - using tower_type = std::array; - template + template struct kernel { - const small_prime_fft_arithmetic &arithmetic; + using quadratic_type = std::array; + using cubic_type = std::array; + using tower_type = std::array; + + const Arithmetic &arithmetic; - double reduce(double value) const { + Value reduce(Value value) const { if constexpr (UseFma) { return arithmetic.reduce_fma(value); } else { @@ -655,7 +739,7 @@ namespace nil::crypto3::math::detail { } } - double multiply(double a, double b) const { + Value multiply(Value a, Value b) const { if constexpr (UseFma) { return arithmetic.multiply_fma(a, b); } else { @@ -663,11 +747,11 @@ namespace nil::crypto3::math::detail { } } - double add(double a, double b) const { + Value add(Value a, Value b) const { return reduce(a + b); } - double subtract(double a, double b) const { + Value subtract(Value a, Value b) const { return reduce(a - b); } @@ -697,9 +781,9 @@ namespace nil::crypto3::math::detail { #pragma clang fp reassociate(off) #pragma clang fp contract(off) #endif - const double diagonal0 = multiply(a[0], b[0]); - const double diagonal1 = multiply(a[1], b[1]); - const double cross = multiply(a[0] + a[1], b[0] + b[1]); + const Value diagonal0 = multiply(a[0], b[0]); + const Value diagonal1 = multiply(a[1], b[1]); + const Value cross = multiply(a[0] + a[1], b[0] + b[1]); return {reduce(diagonal0 - diagonal1), reduce(cross - diagonal0 - diagonal1)}; } @@ -785,16 +869,44 @@ namespace nil::crypto3::math::detail { 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; - const kernel operations {arithmetic_}; - for (std::size_t i = 0; i < stride; ++i) { + 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)); + } + } +#endif + apply_range(arithmetic_, left, 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) { + constexpr std::size_t width = sizeof(Value) / sizeof(double); + const std::size_t stride = left.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; 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) { const std::size_t offset = (6 * w + 2 * v + u) * stride + i; - a[w][v][u] = operations.reduce(left[offset]); - if constexpr (!Square) { - b[w][v][u] = operations.reduce(right[offset]); + if constexpr (width == 1) { + a[w][v][u] = operations.reduce(left[offset]); + if constexpr (!Square) { + b[w][v][u] = operations.reduce(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)); + } +#endif } } } @@ -810,7 +922,15 @@ namespace nil::crypto3::math::detail { for (std::size_t v = 0; v < 3; ++v) { for (std::size_t u = 0; u < 2; ++u) { // Fold inverse-transform normalization into this existing write pass. - left[(6 * w + 2 * v + u) * stride + i] = operations.multiply(result[w][v][u], scale); + 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; +#if defined(CRYPTO3_SMALL_PRIME_FFT_AVX2_SELECTED) + } else { + _mm256_storeu_pd(left.data() + offset, value); +#endif + } } } } @@ -819,6 +939,9 @@ namespace nil::crypto3::math::detail { small_prime_fft_arithmetic arithmetic_; }; +#if defined(CRYPTO3_SMALL_PRIME_FFT_AVX2_SELECTED) +#pragma GCC diagnostic pop +#endif } // namespace nil::crypto3::math::detail diff --git a/libs/math/include/nil/crypto3/math/detail/small_prime_fft_arithmetic.hpp b/libs/math/include/nil/crypto3/math/detail/small_prime_fft_arithmetic.hpp index 1b1ca1082..190232db7 100644 --- a/libs/math/include/nil/crypto3/math/detail/small_prime_fft_arithmetic.hpp +++ b/libs/math/include/nil/crypto3/math/detail/small_prime_fft_arithmetic.hpp @@ -29,16 +29,22 @@ #include #include #include +#include #include #include #include #include +#if defined(__x86_64__) && defined(__AVX2__) && defined(__FMA__) && (defined(__GNUC__) || defined(__clang__)) +#include +#define CRYPTO3_SMALL_PRIME_FFT_AVX2_SELECTED +#endif + namespace nil::crypto3::math::detail { /** - * Scalar residue arithmetic for an odd 50-bit modulus q. Values are integer-valued binary64 + * Residue arithmetic for an odd 50-bit modulus q. Values are integer-valued binary64 * numbers, not approximate coefficients. Primality is a caller concern (the prime catalogue * is separately certified); these arithmetic operations also work for odd composite moduli. * @@ -145,6 +151,89 @@ namespace nil::crypto3::math::detail { return std::fma(-quotient, static_cast(modulus_), high) + low; } +#if defined(CRYPTO3_SMALL_PRIME_FFT_AVX2_SELECTED) + /** + * Snapshot the two packed constants once before a batch of FMA operations. Keeping them + * here avoids repeating integer conversion and broadcasts at each arithmetic operation. + * The source arithmetic must remain unchanged and outlive this kernel. The scalar FMA + * preconditions still apply. + * + * Under the required fixed FE_TONEAREST mode, the assembly computes register outputs + * from its inputs. It is non-volatile so the compiler can reuse the packed constants + * across operations; the instructions within each block still stay together. + */ + class packed_fma_kernel { + public: + explicit packed_fma_kernel(const small_prime_fft_arithmetic &arithmetic) : + arithmetic_(arithmetic), modulus_(_mm256_set1_pd(static_cast(arithmetic.modulus_))), + reciprocal_(_mm256_set1_pd(arithmetic.reciprocal_)) { + assert(arithmetic.fma_supported()); + } + + /** Four independent reductions with the same preconditions and bounds as the scalar overload. */ + __m256d reduce_fma(__m256d value) const { + assert(arithmetic_.fma_supported()); +#ifndef NDEBUG + double lanes[4]; + _mm256_storeu_pd(lanes, value); + for (double lane : lanes) { + assert(valid_integer(lane)); + } +#endif + __m256d quotient; + // Round to nearest, ties to even, suppressing the inexact exception just like nearbyint. + // All operands are registers; the early-clobber scratch cannot overlap a live input. + // clang-format off + asm( + "vmulpd %[reciprocal], %[value], %[quotient]\n" + "vroundpd $8, %[quotient], %[quotient]\n" + "vfnmadd231pd %[modulus], %[quotient], %[value]\n" + : [value] "+x"(value), [quotient] "=&x"(quotient) + : [modulus] "x"(modulus_), [reciprocal] "x"(reciprocal_) + ); + // clang-format on + return value; + } + + /** Four independent compensated products; the scalar product bounds apply to every lane. */ + __m256d multiply_fma(__m256d left, __m256d right) const { + assert(arithmetic_.fma_supported()); +#ifndef NDEBUG + double a[4], b[4]; + _mm256_storeu_pd(a, left); + _mm256_storeu_pd(b, right); + for (std::size_t i = 0; i < 4; ++i) { + assert(valid_integer(a[i]) && valid_integer(b[i])); + assert(product_magnitude(a[i], b[i]) < 4 * wide_type(arithmetic_.modulus_) * arithmetic_.modulus_); + } +#endif + __m256d high, low, quotient; + // Keep the rounded high product and its exact FMA residual separate. The quotient + // estimate and low product are independent once high is available. + // clang-format off + asm( + "vmulpd %[right], %[left], %[high]\n" + "vmulpd %[reciprocal], %[high], %[quotient]\n" + "vmovapd %[high], %[low]\n" + "vfmsub231pd %[right], %[left], %[low]\n" + "vroundpd $8, %[quotient], %[quotient]\n" + "vfnmadd231pd %[modulus], %[quotient], %[high]\n" + "vaddpd %[low], %[high], %[high]\n" + : [high] "=&x"(high), [low] "=&x"(low), [quotient] "=&x"(quotient) + : [left] "x"(left), [right] "x"(right), + [modulus] "x"(modulus_), [reciprocal] "x"(reciprocal_) + ); + // clang-format on + return high; + } + + private: + const small_prime_fft_arithmetic &arithmetic_; + const __m256d modulus_; + const __m256d reciprocal_; + }; +#endif + /** Exact fallback for reduce_fma, with the same input/output bounds and any rounding mode. */ double reduce_integer(double value) const { assert(valid_integer(value)); diff --git a/libs/math/test/small_prime_fft_arithmetic.cpp b/libs/math/test/small_prime_fft_arithmetic.cpp index c3a2a8436..e65354d71 100644 --- a/libs/math/test/small_prime_fft_arithmetic.cpp +++ b/libs/math/test/small_prime_fft_arithmetic.cpp @@ -246,4 +246,94 @@ BOOST_AUTO_TEST_CASE(invalid_moduli_are_rejected_and_valid_endpoints_work) { } } +#if defined(CRYPTO3_SMALL_PRIME_FFT_AVX2_SELECTED) +BOOST_AUTO_TEST_CASE(packed_products_match_scalar_and_exact_integer_oracles) { + restore_rounding restore; + BOOST_REQUIRE_EQUAL(std::fesetround(FE_TONEAREST), 0); + std::mt19937_64 random(0x41565832); + for (const auto &prime : detail::small_prime_fft_primes) { + const arithmetic_type arithmetic(prime.modulus); + if (!arithmetic.fma_supported()) { + continue; + } + const arithmetic_type::packed_fma_kernel packed(arithmetic); + const auto q = static_cast(prime.modulus); + const std::array, 12> boundaries {{{0, 4 * q - 1}, + {q - 1, 1 - q}, + {2 * q - 1, 2 * q + 1}, + {1 - 2 * q, 2 * q + 1}, + {4 * q - 1, q - 1}, + {1 - 4 * q, 1 - q}, + {3 * q - 1, q / 2}, + {q / 2 + 1, 1 - 3 * q}, + {(std::int64_t(1) << 53) - 1, 1}, + {1 - (std::int64_t(1) << 53), -1}, + {1, -1}, + {-q, q}}}; + for (std::size_t batch = 0; batch < 1003; ++batch) { + alignas(32) std::array left {}, right {}, output; + output.fill(-19.0); + for (std::size_t lane = 0; lane < 4; ++lane) { + const auto a = batch < 3 ? boundaries[4 * batch + lane][0] : + static_cast(random() % (8 * q - 1)) - (4 * q - 1); + const auto b = batch < 3 ? boundaries[4 * batch + lane][1] : + static_cast(random() % (2 * q - 1)) - (q - 1); + left[lane + 1] = static_cast(a); + right[lane + 1] = static_cast(b); + } + _mm256_storeu_pd(output.data() + 1, + packed.multiply_fma(_mm256_loadu_pd(left.data() + 1), _mm256_loadu_pd(right.data() + 1))); + for (std::size_t lane = 1; lane <= 4; ++lane) { + const cpp_int product = + cpp_int(static_cast(left[lane])) * static_cast(right[lane]); + const double bound = boost::multiprecision::abs(product) < 2 * cpp_int(q) * q ? double(q) : 1.5 * q; + BOOST_TEST_CONTEXT("q=" << q << " batch=" << batch << " lane=" << lane) { + check_residue(output[lane], positive_residue(product, prime.modulus), prime.modulus, bound); + BOOST_CHECK_EQUAL(output[lane], arithmetic.multiply_fma(left[lane], right[lane])); + } + } + BOOST_CHECK_EQUAL(output.front(), -19.0); + BOOST_CHECK_EQUAL(output.back(), -19.0); + } + } + BOOST_CHECK_EQUAL(std::fegetround(), FE_TONEAREST); +} + +BOOST_AUTO_TEST_CASE(packed_reductions_preserve_lane_order_and_large_integer_bounds) { + restore_rounding restore; + BOOST_REQUIRE_EQUAL(std::fesetround(FE_TONEAREST), 0); + std::mt19937_64 random(0x41565833); + constexpr std::int64_t limit = std::int64_t(1) << 53; + for (const auto &prime : detail::small_prime_fft_primes) { + const arithmetic_type arithmetic(prime.modulus); + if (!arithmetic.fma_supported()) { + continue; + } + const arithmetic_type::packed_fma_kernel packed(arithmetic); + const auto q = static_cast(prime.modulus); + for (std::size_t batch = 0; batch < 1001; ++batch) { + alignas(32) std::array input {}, output; + output.fill(-23.0); + for (std::size_t lane = 1; lane <= 4; ++lane) { + input[lane] = static_cast(static_cast(random() % (2 * limit - 1)) - (limit - 1)); + } + if (batch == 0) { + input = {0.0, double(limit - 1), double(1 - limit), double(q / 2), double(-q / 2 - 1), 0.0}; + } + _mm256_storeu_pd(output.data() + 1, packed.reduce_fma(_mm256_loadu_pd(input.data() + 1))); + for (std::size_t lane = 1; lane <= 4; ++lane) { + BOOST_TEST_CONTEXT("q=" << q << " batch=" << batch << " lane=" << lane) { + check_residue(output[lane], positive_residue(static_cast(input[lane]), prime.modulus), + prime.modulus, double(q)); + BOOST_CHECK_EQUAL(output[lane], arithmetic.reduce_fma(input[lane])); + } + } + BOOST_CHECK_EQUAL(output.front(), -23.0); + BOOST_CHECK_EQUAL(output.back(), -23.0); + } + } + BOOST_CHECK_EQUAL(std::fegetround(), FE_TONEAREST); +} +#endif + 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 d686fcf7a..2958b6fe6 100644 --- a/libs/math/test/small_prime_fft_backend.cpp +++ b/libs/math/test/small_prime_fft_backend.cpp @@ -233,42 +233,54 @@ namespace { BOOST_AUTO_TEST_SUITE(small_prime_fft_backend_tests) BOOST_AUTO_TEST_CASE(pointwise_products_and_squares_match_expanded_integer_tower) { - constexpr std::size_t stride = 9; + restore_rounding restore; + BOOST_REQUIRE_EQUAL(std::fesetround(FE_TONEAREST), 0); std::mt19937_64 random(0x46465410); 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); - std::vector left(12 * stride), right(12 * stride); - std::vector products(stride), squares(stride); - for (std::size_t i = 0; i < stride; ++i) { - coordinates_type a {}, b {}; - for (std::size_t c = 0; c < 12; ++c) { - const std::int64_t x = - c % 3 == 0 ? 3 * q - 1 : static_cast(random() % (6 * q - 1)) - (3 * q - 1); - const std::int64_t y = - c % 3 == 1 ? 1 - 3 * q : static_cast(random() % (6 * q - 1)) - (3 * q - 1); - a[c] = x; - b[c] = y; - left[c * stride + i] = static_cast(x); - right[c * stride + i] = static_cast(y); + // Cover empty/scalar buffers, complete four-lane blocks, and every possible scalar tail. + for (const std::size_t stride : {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 16, 17}) { + std::vector left(12 * stride), right(12 * stride); + std::vector products(stride), squares(stride); + for (std::size_t i = 0; i < stride; ++i) { + coordinates_type a {}, b {}; + for (std::size_t c = 0; c < 12; ++c) { + std::int64_t x; + if (i % 4 == 0 && c % 2 == 0) { + x = (std::int64_t(1) << 53) - 1; + } else { + x = c % 3 == 0 ? 3 * q - 1 : static_cast(random() % (6 * q - 1)) - (3 * q - 1); + } + std::int64_t y; + if (i % 4 == 1 && c % 2 == 1) { + y = 1 - (std::int64_t(1) << 53); + } else { + y = c % 3 == 1 ? 1 - 3 * q : static_cast(random() % (6 * q - 1)) - (3 * q - 1); + } + a[c] = x; + b[c] = y; + left[c * stride + i] = static_cast(x); + right[c * stride + i] = static_cast(y); + } + products[i] = tower_product(a, b); + squares[i] = tower_product(a, a); + } + for (double scale : {0.0, 1.0, -1.0, static_cast((q - 1) / 2), -static_cast((q - 1) / 2)}) { + std::vector storage(left.size() + 2, -19.0); + auto actual = std::span(storage).subspan(1, left.size()); + std::copy(left.begin(), left.end(), actual.begin()); + pointwise.multiply(actual, right, stride, scale); + check_pointwise(actual, products, prime.modulus, scale); + std::copy(left.begin(), left.end(), actual.begin()); + pointwise.square(actual, stride, scale); + check_pointwise(actual, squares, prime.modulus, scale); + std::copy(left.begin(), left.end(), actual.begin()); + pointwise.multiply(actual, actual, stride, scale); + check_pointwise(actual, squares, prime.modulus, scale); + BOOST_CHECK_EQUAL(storage.front(), -19.0); + BOOST_CHECK_EQUAL(storage.back(), -19.0); } - products[i] = tower_product(a, b); - squares[i] = tower_product(a, a); - } - for (double scale : {0.0, 1.0, -1.0, static_cast((q - 1) / 2), -static_cast((q - 1) / 2)}) { - std::vector storage(left.size() + 2, -19.0); - auto actual = std::span(storage).subspan(1, left.size()); - std::copy(left.begin(), left.end(), actual.begin()); - pointwise.multiply(actual, right, stride, scale); - check_pointwise(actual, products, prime.modulus, scale); - std::copy(left.begin(), left.end(), actual.begin()); - pointwise.square(actual, stride, scale); - check_pointwise(actual, squares, prime.modulus, scale); - std::copy(left.begin(), left.end(), actual.begin()); - pointwise.multiply(actual, actual, stride, scale); - check_pointwise(actual, squares, prime.modulus, scale); - BOOST_CHECK_EQUAL(storage.front(), -19.0); - BOOST_CHECK_EQUAL(storage.back(), -19.0); } } } diff --git a/libs/math/test/small_prime_fft_transform.cpp b/libs/math/test/small_prime_fft_transform.cpp index bbc540605..9d53195e8 100644 --- a/libs/math/test/small_prime_fft_transform.cpp +++ b/libs/math/test/small_prime_fft_transform.cpp @@ -333,26 +333,33 @@ BOOST_AUTO_TEST_CASE(rounding_modes_preserve_exactness_and_environment) { } BOOST_AUTO_TEST_CASE(plan_reuse_and_in_place_buffer_boundaries) { - constexpr std::size_t size = 32; - const plan_type original(0, size); - auto copy = original; - const plan_type moved = std::move(copy); std::mt19937_64 random(0x46465406); - for (const plan_type *plan : {&original, &moved}) { - for (unsigned trial = 0; trial < 3; ++trial) { - const auto input = lazy_input(size, plan->modulus(), random); - std::vector storage(size + 2, -17.0); - auto values = std::span(storage).subspan(1, size); - std::copy(input.begin(), input.end(), values.begin()); - const double *address = values.data(); - plan->fft(values); - plan->inverse_fft(values); - BOOST_CHECK(values.data() == address); - BOOST_CHECK_EQUAL(storage.front(), -17.0); - BOOST_CHECK_EQUAL(storage.back(), -17.0); - for (std::size_t i = 0; i < size; ++i) { - check_residue(values[i], canonical(as_integer(input[i]), cpp_int(plan->modulus())), plan->modulus(), 1, - true); + for (const std::size_t size : {1, 2, 4, 8, 16, 32, 64}) { + const plan_type original(0, size); + auto copy = original; + const plan_type moved = std::move(copy); + for (const plan_type *plan : {&original, &moved}) { + for (unsigned trial = 0; trial < 3; ++trial) { + const auto input = lazy_input(size, plan->modulus(), random); + for (std::size_t padding = 1; padding <= 4; ++padding) { + std::vector storage(size + 8, -17.0); + auto values = std::span(storage).subspan(padding, size); + std::copy(input.begin(), input.end(), values.begin()); + const double *address = values.data(); + plan->fft(values); + plan->inverse_fft(values); + BOOST_CHECK(values.data() == address); + for (std::size_t i = 0; i < padding; ++i) { + BOOST_CHECK_EQUAL(storage[i], -17.0); + } + for (std::size_t i = padding + size; i < storage.size(); ++i) { + BOOST_CHECK_EQUAL(storage[i], -17.0); + } + for (std::size_t i = 0; i < size; ++i) { + check_residue(values[i], canonical(as_integer(input[i]), cpp_int(plan->modulus())), + plan->modulus(), 1, true); + } + } } } }