From 63ea71167f4663beff38396356ae208bde30a37a Mon Sep 17 00:00:00 2001 From: Riccardo Abbate Date: Wed, 21 Oct 2026 22:28:36 -0400 Subject: [PATCH] optimize fft --- .../math/detail/bn254_small_prime_fft.hpp | 776 +++++++++++++++++- .../detail/small_prime_fft_arithmetic.hpp | 201 +++++ .../bn254_small_prime_fft_backend.hpp | 4 +- libs/math/test/small_prime_fft_arithmetic.cpp | 70 +- libs/math/test/small_prime_fft_transform.cpp | 118 ++- 5 files changed, 1105 insertions(+), 64 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 03f0bc146..d17eb5397 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 @@ -37,6 +37,8 @@ #include #include +#include + #include #include #include @@ -45,6 +47,13 @@ #include #include +#ifndef CRYPTO3_SMALL_PRIME_FFT_RADIX +#define CRYPTO3_SMALL_PRIME_FFT_RADIX 4 +#endif +static_assert(CRYPTO3_SMALL_PRIME_FFT_RADIX == 4 || CRYPTO3_SMALL_PRIME_FFT_RADIX == 8 || + CRYPTO3_SMALL_PRIME_FFT_RADIX == 16, + "small-prime FFT stage radix must be 4, 8 or 16"); + namespace nil::crypto3::math::detail { /** @@ -578,11 +587,13 @@ namespace nil::crypto3::math::detail { /** * In-place, power-of-two transforms for one prime from the catalogue. The forward transform * maps natural-order coefficients to evaluations at omega^k in bit-reversed index order. - * The inverse consumes that same order, so convolution needs no permutation passes. + * Both directions use this order, so convolution needs no permutation passes. * * 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. + * AVX2/FMA builds fuse 8-, 16- and 32-point leaves; AVX-512 adds paired 32-/64-point + * leaves and eight-lane stages. CRYPTO3_SMALL_PRIME_FFT_RADIX selects 4 (default), 8 + * or 16 for large SIMD stages; the scalar/integer fallback keeps the radix-4 schedule. * 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 @@ -663,6 +674,21 @@ namespace nil::crypto3::math::detail { } } + /** + * Unscaled inverse for pointwise products whose integer representatives are in (-q,q). + * The caller has already folded 1/N into its pointwise scale. This stricter input bound + * permits omitting the first leaf reductions; the output still lies in (-2*q,2*q). + * Dimension errors throw before writes, as in inverse_fft(). + */ + void inverse_fft_pointwise(std::span values) const { + check_buffer<1>(values); + if (arithmetic_.fma_supported()) { + inverse_impl(values, false); + } else { + inverse_impl(values, false); + } + } + private: static const small_prime_fft_prime &checked_prime(std::size_t prime_index, std::size_t size) { if (prime_index >= small_prime_fft_primes.size() || !std::has_single_bit(size) || @@ -672,12 +698,13 @@ namespace nil::crypto3::math::detail { return small_prime_fft_primes[prime_index]; } + template void check_buffer(std::span values) const { if (values.size() != size_ || twiddles_.empty()) { throw std::invalid_argument("small_prime_fft_plan: invalid buffer dimensions or moved-from plan"); } #ifndef NDEBUG - const double limit = 3 * static_cast(modulus()); + const double limit = Bound * static_cast(modulus()); for (double value : values) { assert(std::isfinite(value) && std::trunc(value) == value && std::abs(value) < limit); } @@ -689,7 +716,7 @@ namespace nil::crypto3::math::detail { return value > q / 2 ? value - q : (value < -q / 2 ? value + q : value); } - double inverse_twiddle(std::size_t index) const { + BOOST_FORCEINLINE double inverse_twiddle(std::size_t index) const { if (index == 0) { return 1.0; } @@ -700,7 +727,7 @@ namespace nil::crypto3::math::detail { // Keep arithmetic selection outside the butterfly loops, with compile-time dispatch inside them. template - static Value reduce(const Arithmetic &arithmetic, Value value) { + BOOST_FORCEINLINE static Value reduce(const Arithmetic &arithmetic, Value value) { if constexpr (UseFma) { return arithmetic.reduce_fma(value); } else { @@ -709,7 +736,7 @@ namespace nil::crypto3::math::detail { } template - static Value multiply(const Arithmetic &arithmetic, Value left, Value right) { + BOOST_FORCEINLINE static Value multiply(const Arithmetic &arithmetic, Value left, Value right) { if constexpr (UseFma) { return arithmetic.multiply_fma(left, right); } else { @@ -717,22 +744,47 @@ namespace nil::crypto3::math::detail { } } + template + BOOST_FORCEINLINE static void multiply_pair(const Arithmetic &arithmetic, Value &a, Value &b, Value wa, + Value wb) { +#if defined(CRYPTO3_SMALL_PRIME_FFT_AVX2_SELECTED) + if constexpr (UseFma && requires { arithmetic.multiply_fma_pair(a, b, wa, wb); }) { + arithmetic.multiply_fma_pair(a, b, wa, wb); + } else +#endif + { + a = multiply(arithmetic, a, wa); + b = multiply(arithmetic, b, wb); + } + } + template - static void forward_butterfly(const Arithmetic &arithmetic, Value &v0, Value &v1, Value &v2, Value &v3, - Value w2, Value w, Value iw) { + BOOST_FORCEINLINE 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); + Value a2 = v2; + Value a3 = v3; + if constexpr (UnitTwiddle) { + a2 = reduce(arithmetic, a2); + a3 = reduce(arithmetic, a3); + } else { + multiply_pair(arithmetic, a2, a3, w2, 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); + Value odd_sum = a1 + a3; + Value odd_difference = a1 - a3; + if constexpr (UnitTwiddle) { + odd_sum = reduce(arithmetic, odd_sum); + odd_difference = multiply(arithmetic, odd_difference, iw); + } else { + multiply_pair(arithmetic, odd_sum, odd_difference, w, 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. @@ -742,12 +794,475 @@ namespace nil::crypto3::math::detail { v3 = even_difference - odd_difference; } +#if defined(CRYPTO3_SMALL_PRIME_FFT_AVX2_SELECTED) + BOOST_FORCEINLINE static void transpose4(__m256d &v0, __m256d &v1, __m256d &v2, __m256d &v3) { + const __m256d a = _mm256_unpacklo_pd(v0, v1); + const __m256d b = _mm256_unpackhi_pd(v0, v1); + const __m256d c = _mm256_unpacklo_pd(v2, v3); + const __m256d d = _mm256_unpackhi_pd(v2, v3); + v0 = _mm256_permute2f128_pd(a, c, 0x20); + v1 = _mm256_permute2f128_pd(b, d, 0x20); + v2 = _mm256_permute2f128_pd(a, c, 0x31); + v3 = _mm256_permute2f128_pd(b, d, 0x31); + } + + // Aligned table indices, with unaligned loads: the owning vector needs no special allocator. + template + BOOST_FORCEINLINE __m256d twiddle4(std::size_t index) const { + if constexpr (!Inverse) { + return _mm256_loadu_pd(twiddles_.data() + index); + } else { + if (index == 0) { + return _mm256_setr_pd(1.0, inverse_twiddle(1), inverse_twiddle(2), inverse_twiddle(3)); + } + // Four consecutive entries stay in one [b,2*b) segment. Reflect, reverse, negate. + const std::size_t base = std::bit_floor(index); + const __m256d reflected = _mm256_loadu_pd(twiddles_.data() + 3 * base - index - 4); + return _mm256_sub_pd(_mm256_setzero_pd(), _mm256_permute4x64_pd(reflected, 0x1b)); + } + } + + template + BOOST_FORCEINLINE void leaf_twiddles(std::size_t node, __m256d &w2, __m256d &w, __m256d &iw) const { + w2 = twiddle4(4 * node); + const __m256d low = twiddle4(8 * node); + const __m256d high = twiddle4(8 * node + 4); + w = _mm256_permute4x64_pd(_mm256_unpacklo_pd(low, high), 0xd8); + iw = _mm256_permute4x64_pd(_mm256_unpackhi_pd(low, high), 0xd8); + } + + template + BOOST_FORCEINLINE void forward16(const small_prime_fft_arithmetic::packed_fma_kernel &packed, __m256d &v0, + __m256d &v1, __m256d &v2, __m256d &v3, std::size_t node) const { + forward_butterfly(packed, v0, v1, v2, v3, _mm256_set1_pd(twiddles_[node]), + _mm256_set1_pd(twiddles_[2 * node]), + _mm256_set1_pd(twiddles_[2 * node + 1])); + // The four child transforms now run across SIMD lanes, without scalar leaves or reloads. + transpose4(v0, v1, v2, v3); + __m256d w2, w, iw; + leaf_twiddles(node, w2, w, iw); + forward_butterfly(packed, v0, v1, v2, v3, w2, w, iw); + transpose4(v0, v1, v2, v3); + } + + template + BOOST_FORCEINLINE void inverse16(const small_prime_fft_arithmetic::packed_fma_kernel &packed, __m256d &v0, + __m256d &v1, __m256d &v2, __m256d &v3, std::size_t node) const { + transpose4(v0, v1, v2, v3); + __m256d w2, w, iw; + leaf_twiddles(node, w2, w, iw); + inverse_butterfly(packed, v0, v1, v2, v3, w2, w, iw, !BoundedInput); + transpose4(v0, v1, v2, v3); + inverse_butterfly(packed, v0, v1, v2, v3, _mm256_set1_pd(inverse_twiddle(node)), + _mm256_set1_pd(inverse_twiddle(2 * node)), + _mm256_set1_pd(inverse_twiddle(2 * node + 1)), false); + } + + template + void base16(double *values, std::size_t node) const { + const small_prime_fft_arithmetic::packed_fma_kernel packed(arithmetic_); + __m256d v0 = _mm256_loadu_pd(values); + __m256d v1 = _mm256_loadu_pd(values + 4); + __m256d v2 = _mm256_loadu_pd(values + 8); + __m256d v3 = _mm256_loadu_pd(values + 12); + if constexpr (Inverse) { + inverse16(packed, v0, v1, v2, v3, node); + } else { + forward16(packed, v0, v1, v2, v3, node); + } + _mm256_storeu_pd(values, v0); + _mm256_storeu_pd(values + 4, v1); + _mm256_storeu_pd(values + 8, v2); + _mm256_storeu_pd(values + 12, v3); + } + + template + BOOST_FORCEINLINE static void radix2_packed(const Arithmetic &packed, Value &left, Value &right, Value weight) { + if constexpr (Inverse) { + // Both subtransforms already return magnitudes < 2*q. No input reductions are needed. + const Value sum = left + right; + const Value difference = left - right; + left = packed.reduce_fma(sum); + right = UnitTwiddle ? packed.reduce_fma(difference) : packed.multiply_fma(difference, weight); + } else { + const Value a = packed.reduce_fma(left); + const Value b = UnitTwiddle ? packed.reduce_fma(right) : packed.multiply_fma(right, weight); + left = a + b; + right = a - b; + } + } + + template + BOOST_FORCEINLINE static void stage4(const Arithmetic &arithmetic, Value &v0, Value &v1, Value &v2, Value &v3, + const double *weights) { + const Value w2 = Arithmetic::broadcast(weights[0]); + const Value w = Arithmetic::broadcast(weights[1]); + const Value iw = Arithmetic::broadcast(weights[2]); + if constexpr (Inverse) { + inverse_butterfly(arithmetic, v0, v1, v2, v3, w2, w, iw, false); + } else { + forward_butterfly(arithmetic, v0, v1, v2, v3, w2, w, iw); + } + } + + template + BOOST_FORCEINLINE void stage_weights(double *weights, std::size_t node) const { + weights[0] = Inverse ? inverse_twiddle(node) : twiddles_[node]; + weights[1] = Inverse ? inverse_twiddle(2 * node) : twiddles_[2 * node]; + weights[2] = Inverse ? inverse_twiddle(2 * node + 1) : twiddles_[2 * node + 1]; + } + + template + BOOST_FORCEINLINE void stage4(const Arithmetic &arithmetic, Value &v0, Value &v1, Value &v2, Value &v3, + std::size_t node) const { + double weights[3]; + stage_weights(weights, node); + stage4(arithmetic, v0, v1, v2, v3, weights); + } + + template + BOOST_FORCEINLINE static void fused_outer(const Arithmetic &arithmetic, Value *v, const double *weights) { + if constexpr (Radix == 8) { + const Value weight = Arithmetic::broadcast(weights[0]); + radix2_packed(arithmetic, v[0], v[4], weight); + radix2_packed(arithmetic, v[1], v[5], weight); + radix2_packed(arithmetic, v[2], v[6], weight); + radix2_packed(arithmetic, v[3], v[7], weight); + } else { + stage4(arithmetic, v[0], v[4], v[8], v[12], weights); + stage4(arithmetic, v[1], v[5], v[9], v[13], weights); + stage4(arithmetic, v[2], v[6], v[10], v[14], weights); + stage4(arithmetic, v[3], v[7], v[11], v[15], weights); + } + } + + // Fuse three/four binary layers without intermediate array passes. Every constituent + // butterfly retains the forward <3*q and inverse <2*q bounds; no wider lazy range is used. + // UnitTwiddle is valid only at node zero, hence only the first child inherits it. + template + BOOST_FORCEINLINE static void fused_stage(const Arithmetic &arithmetic, Value *v, const double *weights) { + static_assert(Radix == 8 || Radix == 16); + if constexpr (!Inverse) { + fused_outer(arithmetic, v, weights); + } + constexpr std::size_t outer_weights = Radix == 8 ? 1 : 3; + stage4(arithmetic, v[0], v[1], v[2], v[3], weights + outer_weights); + stage4(arithmetic, v[4], v[5], v[6], v[7], weights + outer_weights + 3); + if constexpr (Radix == 16) { + stage4(arithmetic, v[8], v[9], v[10], v[11], weights + outer_weights + 6); + stage4(arithmetic, v[12], v[13], v[14], v[15], weights + outer_weights + 9); + } + if constexpr (Inverse) { + fused_outer(arithmetic, v, weights); + } + } + + template + void large_radix(double *values, std::size_t stride, std::size_t node) const { +#if defined(CRYPTO3_SMALL_PRIME_FFT_AVX512_SELECTED) + using arithmetic_type = small_prime_fft_arithmetic::wide_fma_kernel; +#else + using arithmetic_type = small_prime_fft_arithmetic::packed_fma_kernel; +#endif + const arithmetic_type arithmetic(arithmetic_); + // Dispatch requires stride >=16, a power of two divisible by the SIMD width. + assert(stride >= 16 && stride % arithmetic_type::lanes == 0); + // The node is fixed throughout this coefficient loop. Resolve reflections and + // table lookups once, retaining scalar weights to avoid forcing Radix-1 live vectors. + double weights[Radix - 1]; + if constexpr (Radix == 8) { + weights[0] = Inverse ? inverse_twiddle(node) : twiddles_[node]; + } else { + stage_weights(weights, node); + } + constexpr std::size_t outer_weights = Radix == 8 ? 1 : 3; +#pragma GCC unroll 4 + for (std::size_t child = 0; child < Radix / 4; ++child) { + stage_weights(weights + outer_weights + 3 * child, Radix / 4 * node + child); + } + for (std::size_t i = 0; i < stride; i += arithmetic_type::lanes) { + arithmetic_type::value_type v[Radix]; +#pragma GCC unroll 16 + for (std::size_t part = 0; part < Radix; ++part) { + v[part] = arithmetic_type::load(values + part * stride + i); + } + fused_stage(arithmetic, v, weights); +#pragma GCC unroll 16 + for (std::size_t part = 0; part < Radix; ++part) { + arithmetic_type::store(values + part * stride + i, v[part]); + } + } + } + + template + void base32(double *values, std::size_t node) const { + const small_prime_fft_arithmetic::packed_fma_kernel packed(arithmetic_); + __m256d v0 = _mm256_loadu_pd(values); + __m256d v1 = _mm256_loadu_pd(values + 4); + __m256d v2 = _mm256_loadu_pd(values + 8); + __m256d v3 = _mm256_loadu_pd(values + 12); + __m256d v4 = _mm256_loadu_pd(values + 16); + __m256d v5 = _mm256_loadu_pd(values + 20); + __m256d v6 = _mm256_loadu_pd(values + 24); + __m256d v7 = _mm256_loadu_pd(values + 28); + const __m256d weight = _mm256_set1_pd(Inverse ? inverse_twiddle(node) : twiddles_[node]); + // One radix-2 layer and two register-resident 16-point transforms cover all five layers. + if constexpr (Inverse) { + inverse16(packed, v0, v1, v2, v3, 2 * node); + inverse16(packed, v4, v5, v6, v7, 2 * node + 1); + } + radix2_packed(packed, v0, v4, weight); + radix2_packed(packed, v1, v5, weight); + radix2_packed(packed, v2, v6, weight); + radix2_packed(packed, v3, v7, weight); + if constexpr (!Inverse) { + forward16(packed, v0, v1, v2, v3, 2 * node); + forward16(packed, v4, v5, v6, v7, 2 * node + 1); + } + _mm256_storeu_pd(values, v0); + _mm256_storeu_pd(values + 4, v1); + _mm256_storeu_pd(values + 8, v2); + _mm256_storeu_pd(values + 12, v3); + _mm256_storeu_pd(values + 16, v4); + _mm256_storeu_pd(values + 20, v5); + _mm256_storeu_pd(values + 24, v6); + _mm256_storeu_pd(values + 28, v7); + } + + template + void base8(double *values, std::size_t node) const { + const small_prime_fft_arithmetic::packed_fma_kernel packed(arithmetic_); + const __m256d low = _mm256_loadu_pd(values); + const __m256d high = _mm256_loadu_pd(values + 4); + __m256d v0, v1, v2, v3; + if constexpr (Inverse) { + __m256d a = _mm256_permute4x64_pd(_mm256_unpacklo_pd(low, high), 0xd8); + __m256d b = _mm256_permute4x64_pd(_mm256_unpackhi_pd(low, high), 0xd8); + if constexpr (!BoundedInput) { + a = packed.reduce_fma(a); + b = packed.reduce_fma(b); + } + const __m256d sum = packed.reduce_fma(a + b); + const __m256d difference = packed.multiply_fma(a - b, twiddle4(4 * node)); + const __m256d even = _mm256_unpacklo_pd(sum, difference); + const __m256d odd = _mm256_unpackhi_pd(sum, difference); + v0 = _mm256_permute2f128_pd(even, even, 0x00); + v1 = _mm256_permute2f128_pd(odd, odd, 0x00); + v2 = _mm256_permute2f128_pd(even, even, 0x11); + v3 = _mm256_permute2f128_pd(odd, odd, 0x11); + inverse_butterfly(packed, v0, v1, v2, v3, _mm256_set1_pd(inverse_twiddle(node)), + _mm256_set1_pd(inverse_twiddle(2 * node)), + _mm256_set1_pd(inverse_twiddle(2 * node + 1)), false); + _mm256_storeu_pd(values, _mm256_permute2f128_pd(v0, v1, 0x20)); + _mm256_storeu_pd(values + 4, _mm256_permute2f128_pd(v2, v3, 0x20)); + } else { + // Two independent columns occupy the lower lanes; the duplicated upper lanes are discarded. + v0 = _mm256_permute2f128_pd(low, low, 0x00); + v1 = _mm256_permute2f128_pd(low, low, 0x11); + v2 = _mm256_permute2f128_pd(high, high, 0x00); + v3 = _mm256_permute2f128_pd(high, high, 0x11); + forward_butterfly(packed, v0, v1, v2, v3, _mm256_set1_pd(twiddles_[node]), + _mm256_set1_pd(twiddles_[2 * node]), + _mm256_set1_pd(twiddles_[2 * node + 1])); + const __m256d a = packed.reduce_fma( + _mm256_permute2f128_pd(_mm256_unpacklo_pd(v0, v1), _mm256_unpacklo_pd(v2, v3), 0x20)); + const __m256d b = packed.multiply_fma( + _mm256_permute2f128_pd(_mm256_unpackhi_pd(v0, v1), _mm256_unpackhi_pd(v2, v3), 0x20), + twiddle4(4 * node)); + const __m256d even = _mm256_unpacklo_pd(a + b, a - b); + const __m256d odd = _mm256_unpackhi_pd(a + b, a - b); + _mm256_storeu_pd(values, _mm256_permute2f128_pd(even, odd, 0x20)); + _mm256_storeu_pd(values + 4, _mm256_permute2f128_pd(even, odd, 0x31)); + } + } + + template + void basecase(double *values, std::size_t length, std::size_t node) const { + if (length == 16) { + base16(values, node); + } else if (length == 32) { + base32(values, node); + } else { + base8(values, node); + } + } + +#if defined(CRYPTO3_SMALL_PRIME_FFT_AVX512_SELECTED) + BOOST_FORCEINLINE static __m512d join4(__m256d low, __m256d high) { + return _mm512_insertf64x4(_mm512_castpd256_pd512(low), high, 1); + } + + // Two independent 4x4 transposes, one in each 256-bit half. The stored spectrum + // remains bit-reversed; these permutations only change the lanes during a leaf. + BOOST_FORCEINLINE static void transpose4(__m512d &v0, __m512d &v1, __m512d &v2, __m512d &v3) { + const __m512d a = _mm512_unpacklo_pd(v0, v1); + const __m512d b = _mm512_unpackhi_pd(v0, v1); + const __m512d c = _mm512_unpacklo_pd(v2, v3); + const __m512d d = _mm512_unpackhi_pd(v2, v3); + const __m512i low = _mm512_setr_epi64(0, 1, 8, 9, 4, 5, 12, 13); + const __m512i high = _mm512_setr_epi64(2, 3, 10, 11, 6, 7, 14, 15); + v0 = _mm512_permutex2var_pd(a, low, c); + v1 = _mm512_permutex2var_pd(b, low, d); + v2 = _mm512_permutex2var_pd(a, high, c); + v3 = _mm512_permutex2var_pd(b, high, d); + } + + template + BOOST_FORCEINLINE __m512d paired_twiddle(std::size_t low, std::size_t high) const { + return join4(_mm256_set1_pd(Inverse ? inverse_twiddle(low) : twiddles_[low]), + _mm256_set1_pd(Inverse ? inverse_twiddle(high) : twiddles_[high])); + } + + // Keep the complete leaf in one function: outlining this step forces live vectors + // through memory across the call. Recursive transform boundaries remain uninlined. + template + BOOST_FORCEINLINE void paired_outer(const small_prime_fft_arithmetic::wide_fma_kernel &arithmetic, __m512d &v0, + __m512d &v1, __m512d &v2, __m512d &v3, std::size_t node) const { + const __m512d w2 = paired_twiddle(node, node + 1); + const __m512d w = paired_twiddle(2 * node, 2 * node + 2); + const __m512d iw = paired_twiddle(2 * node + 1, 2 * node + 3); + if constexpr (Inverse) { + inverse_butterfly(arithmetic, v0, v1, v2, v3, w2, w, iw, false); + } else { + forward_butterfly(arithmetic, v0, v1, v2, v3, w2, w, iw); + } + } + + template + BOOST_FORCEINLINE void base16_pair(const small_prime_fft_arithmetic::wide_fma_kernel &arithmetic, __m512d &v0, + __m512d &v1, __m512d &v2, __m512d &v3, std::size_t node) const { + // Group matching rows of two adjacent 16-point transforms in the two vector halves. + const __m512d a = _mm512_shuffle_f64x2(v0, v2, 0x44); + const __m512d b = _mm512_shuffle_f64x2(v0, v2, 0xee); + const __m512d c = _mm512_shuffle_f64x2(v1, v3, 0x44); + const __m512d d = _mm512_shuffle_f64x2(v1, v3, 0xee); + v0 = a; + v1 = b; + v2 = c; + v3 = d; + if constexpr (!Inverse) { + paired_outer(arithmetic, v0, v1, v2, v3, node); + } + transpose4(v0, v1, v2, v3); + __m256d lo_w2, lo_w, lo_iw, hi_w2, hi_w, hi_iw; + leaf_twiddles(node, lo_w2, lo_w, lo_iw); + leaf_twiddles(node + 1, hi_w2, hi_w, hi_iw); + if constexpr (Inverse) { + inverse_butterfly(arithmetic, v0, v1, v2, v3, join4(lo_w2, hi_w2), join4(lo_w, hi_w), + join4(lo_iw, hi_iw), !BoundedInput); + } else { + forward_butterfly(arithmetic, v0, v1, v2, v3, join4(lo_w2, hi_w2), join4(lo_w, hi_w), + join4(lo_iw, hi_iw)); + } + transpose4(v0, v1, v2, v3); + if constexpr (Inverse) { + paired_outer(arithmetic, v0, v1, v2, v3, node); + } + // Restore contiguous bit-reversed order in the same four input/output vectors. + const __m512d out0 = _mm512_shuffle_f64x2(v0, v1, 0x44); + const __m512d out1 = _mm512_shuffle_f64x2(v2, v3, 0x44); + const __m512d out2 = _mm512_shuffle_f64x2(v0, v1, 0xee); + const __m512d out3 = _mm512_shuffle_f64x2(v2, v3, 0xee); + v0 = out0; + v1 = out1; + v2 = out2; + v3 = out3; + } + + template + void base32_wide(double *values, std::size_t node) const { + const small_prime_fft_arithmetic::wide_fma_kernel arithmetic(arithmetic_); + __m512d v0 = _mm512_loadu_pd(values); + __m512d v1 = _mm512_loadu_pd(values + 8); + __m512d v2 = _mm512_loadu_pd(values + 16); + __m512d v3 = _mm512_loadu_pd(values + 24); + const __m512d weight = _mm512_set1_pd(Inverse ? inverse_twiddle(node) : twiddles_[node]); + if constexpr (Inverse) { + base16_pair(arithmetic, v0, v1, v2, v3, 2 * node); + } + radix2_packed(arithmetic, v0, v2, weight); + radix2_packed(arithmetic, v1, v3, weight); + if constexpr (!Inverse) { + base16_pair(arithmetic, v0, v1, v2, v3, 2 * node); + } + _mm512_storeu_pd(values, v0); + _mm512_storeu_pd(values + 8, v1); + _mm512_storeu_pd(values + 16, v2); + _mm512_storeu_pd(values + 24, v3); + } + + template + void base64_wide(double *values, std::size_t node) const { + const small_prime_fft_arithmetic::wide_fma_kernel arithmetic(arithmetic_); + __m512d v[8]; +#pragma GCC unroll 8 + for (std::size_t i = 0; i < 8; ++i) { + v[i] = _mm512_loadu_pd(values + 8 * i); + } + if constexpr (!Inverse) { + stage4(arithmetic, v[0], v[2], v[4], v[6], node); + stage4(arithmetic, v[1], v[3], v[5], v[7], node); + } + base16_pair(arithmetic, v[0], v[1], v[2], v[3], 4 * node); + base16_pair(arithmetic, v[4], v[5], v[6], v[7], 4 * node + 2); + if constexpr (Inverse) { + stage4(arithmetic, v[0], v[2], v[4], v[6], node); + stage4(arithmetic, v[1], v[3], v[5], v[7], node); + } +#pragma GCC unroll 8 + for (std::size_t i = 0; i < 8; ++i) { + _mm512_storeu_pd(values + 8 * i, v[i]); + } + } + + template + void wide_basecase(double *values, std::size_t length, std::size_t node) const { + if (length == 32) { + if (node == 0) { + base32_wide(values, node); + } else { + base32_wide(values, node); + } + } else { + if (node == 0) { + base64_wide(values, node); + } else { + base64_wide(values, node); + } + } + } +#endif +#endif + template - void forward_radix4(double *values, std::size_t quarter, std::size_t node) const { + void forward_radix4(double *values, std::size_t quarter, std::size_t node, std::size_t first, + std::size_t last) const { const double w2 = twiddles_[node]; const double w = twiddles_[2 * node]; const double iw = twiddles_[2 * node + 1]; - std::size_t i = 0; + std::size_t i = first; +#if defined(CRYPTO3_SMALL_PRIME_FFT_AVX512_SELECTED) + if constexpr (UseFma) { + if (quarter >= 8) { + const small_prime_fft_arithmetic::wide_fma_kernel packed(arithmetic_); + const __m512d packed_w2 = _mm512_set1_pd(w2); + const __m512d packed_w = _mm512_set1_pd(w); + const __m512d packed_iw = _mm512_set1_pd(iw); +#pragma GCC unroll 2 + for (; last - i >= 8; i += 8) { + __m512d v0 = _mm512_loadu_pd(values + i); + __m512d v1 = _mm512_loadu_pd(values + i + quarter); + __m512d v2 = _mm512_loadu_pd(values + i + 2 * quarter); + __m512d v3 = _mm512_loadu_pd(values + i + 3 * quarter); + forward_butterfly(packed, v0, v1, v2, v3, packed_w2, packed_w, packed_iw); + _mm512_storeu_pd(values + i, v0); + _mm512_storeu_pd(values + i + quarter, v1); + _mm512_storeu_pd(values + i + 2 * quarter, v2); + _mm512_storeu_pd(values + i + 3 * quarter, v3); + } + } + } +#endif #if defined(CRYPTO3_SMALL_PRIME_FFT_AVX2_SELECTED) if constexpr (UseFma) { if (quarter >= 4) { @@ -755,7 +1270,8 @@ namespace nil::crypto3::math::detail { 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) { +#pragma GCC unroll 2 + for (; last - 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); @@ -769,14 +1285,98 @@ namespace nil::crypto3::math::detail { } } #endif - for (; i < quarter; ++i) { + for (; i < last; ++i) { forward_butterfly(arithmetic_, values[i], values[i + quarter], values[i + 2 * quarter], values[i + 3 * quarter], w2, w, iw); } } +#if defined(CRYPTO3_SMALL_PRIME_FFT_AVX2_SELECTED) + template + void four_layer_block(double *values, std::size_t length, std::size_t node) const { + const std::size_t quarter = length / 4; + const std::size_t sixteenth = length / 16; + // Sixteen strips of 64 doubles occupy 8 KiB. Complete both radix-4 passes on + // each strip group while it is hot, before recursing into its sixteen children. + // This only changes visitation order; every butterfly retains the same inputs and bounds. + constexpr std::size_t strip_size = 64; + const auto pass = [&](double *base, std::size_t stride, std::size_t j, std::size_t first, + std::size_t last) { + if constexpr (Inverse) { + if (j == 0) { + inverse_radix4(base, stride, j, first, last); + } else { + inverse_radix4(base, stride, j, first, last); + } + } else { + if (j == 0) { + forward_radix4(base, stride, j, first, last); + } else { + forward_radix4(base, stride, j, first, last); + } + } + }; + for (std::size_t first = 0; first < sixteenth; first += strip_size) { + const std::size_t last = std::min(first + strip_size, sixteenth); + if constexpr (Inverse) { + for (std::size_t child = 0; child < 4; ++child) { + pass(values + child * quarter, sixteenth, 4 * node + child, first, last); + } + } + for (std::size_t part = 0; part < 4; ++part) { + pass(values, quarter, node, first + part * sixteenth, last + part * sixteenth); + } + if constexpr (!Inverse) { + for (std::size_t child = 0; child < 4; ++child) { + pass(values + child * quarter, sixteenth, 4 * node + child, first, last); + } + } + } + } +#endif + template void forward_recursive(double *values, std::size_t length, std::size_t node) const { +#if defined(CRYPTO3_SMALL_PRIME_FFT_AVX2_SELECTED) + if constexpr (UseFma) { +#if defined(CRYPTO3_SMALL_PRIME_FFT_AVX512_SELECTED) + if (length == 32 || length == 64) { + wide_basecase(values, length, node); + return; + } +#endif + if (length >= 8 && length <= 32) { + if (node == 0) { + basecase(values, length, node); + } else { + basecase(values, length, node); + } + return; + } +#if CRYPTO3_SMALL_PRIME_FFT_RADIX != 4 + constexpr std::size_t radix = CRYPTO3_SMALL_PRIME_FFT_RADIX; + if (length >= 16 * radix) { + const auto stride = length / radix; + if (node == 0) { + large_radix(values, stride, node); + } else { + large_radix(values, stride, node); + } + for (std::size_t child = 0; child < radix; ++child) { + forward_recursive(values + child * stride, stride, radix * node + child); + } + return; + } +#endif + if (length >= 4096) { + four_layer_block(values, length, node); + for (std::size_t child = 0; child < 16; ++child) { + forward_recursive(values + child * (length / 16), length / 16, 16 * node + child); + } + return; + } + } +#endif if (length == 1) { return; } @@ -790,9 +1390,9 @@ namespace nil::crypto3::math::detail { } const std::size_t quarter = length / 4; if (node == 0) { - forward_radix4(values, quarter, node); + forward_radix4(values, quarter, node, 0, quarter); } else { - forward_radix4(values, quarter, node); + forward_radix4(values, quarter, node, 0, quarter); } if (quarter > 1) { for (std::size_t child = 0; child < 4; ++child) { @@ -802,8 +1402,8 @@ namespace nil::crypto3::math::detail { } template - static void inverse_butterfly(const Arithmetic &arithmetic, Value &v0, Value &v1, Value &v2, Value &v3, - Value w2, Value w, Value iw, bool leaf) { + BOOST_FORCEINLINE 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) @@ -816,26 +1416,59 @@ namespace nil::crypto3::math::detail { 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); + Value left = b0 - b1; + Value right = b2 - b3; + if constexpr (UnitTwiddle) { + left = reduce(arithmetic, left); + right = multiply(arithmetic, right, iw); + } else { + multiply_pair(arithmetic, left, right, w, 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); + v2 = even_sum - odd_sum; + v3 = left - right; + if constexpr (UnitTwiddle) { + v2 = reduce(arithmetic, v2); + v3 = reduce(arithmetic, v3); + } else { + multiply_pair(arithmetic, v2, v3, w2, w2); + } } - template - void inverse_radix4(double *values, std::size_t quarter, std::size_t node) const { + template + void inverse_radix4(double *values, std::size_t quarter, std::size_t node, std::size_t first, + std::size_t last) const { const double w2 = inverse_twiddle(node); const double w = inverse_twiddle(2 * node); const double iw = inverse_twiddle(2 * node + 1); - std::size_t i = 0; + std::size_t i = first; +#if defined(CRYPTO3_SMALL_PRIME_FFT_AVX512_SELECTED) + if constexpr (UseFma) { + if (quarter >= 8) { + const small_prime_fft_arithmetic::wide_fma_kernel packed(arithmetic_); + const __m512d packed_w2 = _mm512_set1_pd(w2); + const __m512d packed_w = _mm512_set1_pd(w); + const __m512d packed_iw = _mm512_set1_pd(iw); +#pragma GCC unroll 2 + for (; last - i >= 8; i += 8) { + __m512d v0 = _mm512_loadu_pd(values + i); + __m512d v1 = _mm512_loadu_pd(values + i + quarter); + __m512d v2 = _mm512_loadu_pd(values + i + 2 * quarter); + __m512d v3 = _mm512_loadu_pd(values + i + 3 * quarter); + inverse_butterfly(packed, v0, v1, v2, v3, packed_w2, packed_w, packed_iw, + false); + _mm512_storeu_pd(values + i, v0); + _mm512_storeu_pd(values + i + quarter, v1); + _mm512_storeu_pd(values + i + 2 * quarter, v2); + _mm512_storeu_pd(values + i + 3 * quarter, v3); + } + } + } +#endif #if defined(CRYPTO3_SMALL_PRIME_FFT_AVX2_SELECTED) if constexpr (UseFma) { if (quarter >= 4) { @@ -843,7 +1476,8 @@ namespace nil::crypto3::math::detail { 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) { +#pragma GCC unroll 2 + for (; last - 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); @@ -858,22 +1492,65 @@ namespace nil::crypto3::math::detail { } } #endif - for (; i < quarter; ++i) { + for (; i < last; ++i) { inverse_butterfly(arithmetic_, values[i], values[i + quarter], values[i + 2 * quarter], values[i + 3 * quarter], w2, w, iw, - quarter == 1); + quarter == 1 && !BoundedInput); } } - template + template void inverse_recursive(double *values, std::size_t length, std::size_t node) const { +#if defined(CRYPTO3_SMALL_PRIME_FFT_AVX2_SELECTED) + if constexpr (UseFma) { +#if defined(CRYPTO3_SMALL_PRIME_FFT_AVX512_SELECTED) + if (length == 32 || length == 64) { + wide_basecase(values, length, node); + return; + } +#endif + if (length >= 8 && length <= 32) { + if (node == 0) { + basecase(values, length, node); + } else { + basecase(values, length, node); + } + return; + } +#if CRYPTO3_SMALL_PRIME_FFT_RADIX != 4 + constexpr std::size_t radix = CRYPTO3_SMALL_PRIME_FFT_RADIX; + if (length >= 16 * radix) { + const auto stride = length / radix; + for (std::size_t child = 0; child < radix; ++child) { + inverse_recursive(values + child * stride, stride, radix * node + child); + } + if (node == 0) { + large_radix(values, stride, node); + } else { + large_radix(values, stride, node); + } + return; + } +#endif + if (length >= 4096) { + for (std::size_t child = 0; child < 16; ++child) { + inverse_recursive(values + child * (length / 16), length / 16, + 16 * node + child); + } + four_layer_block(values, length, node); + return; + } + } +#endif if (length == 1) { - values[0] = reduce(arithmetic_, values[0]); + if constexpr (!BoundedInput) { + values[0] = reduce(arithmetic_, values[0]); + } return; } if (length == 2) { - const double left = reduce(arithmetic_, values[0]); - const double right = reduce(arithmetic_, values[1]); + const double left = BoundedInput ? values[0] : reduce(arithmetic_, values[0]); + const double right = BoundedInput ? values[1] : 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)); @@ -882,22 +1559,35 @@ namespace nil::crypto3::math::detail { const std::size_t quarter = length / 4; if (quarter > 1) { for (std::size_t child = 0; child < 4; ++child) { - inverse_recursive(values + child * quarter, quarter, 4 * node + child); + inverse_recursive(values + child * quarter, quarter, 4 * node + child); } } if (node == 0) { - inverse_radix4(values, quarter, node); + inverse_radix4(values, quarter, node, 0, quarter); } else { - inverse_radix4(values, quarter, node); + inverse_radix4(values, quarter, node, 0, quarter); } } - template + template void inverse_impl(std::span values, bool normalize) const { - inverse_recursive(values.data(), size_, 0); + inverse_recursive(values.data(), size_, 0); if (normalize) { const double q = static_cast(modulus()); std::size_t i = 0; +#if defined(CRYPTO3_SMALL_PRIME_FFT_AVX512_SELECTED) + if constexpr (UseFma) { + const small_prime_fft_arithmetic::wide_fma_kernel packed(arithmetic_); + const __m512d scale = _mm512_set1_pd(size_inverse_); + const __m512d packed_q = _mm512_set1_pd(q); + for (; values.size() - i >= 8; i += 8) { + __m512d value = packed.multiply_fma(_mm512_loadu_pd(values.data() + i), scale); + const auto negative = _mm512_cmp_pd_mask(value, _mm512_setzero_pd(), _CMP_LT_OQ); + value = _mm512_mask_add_pd(value, negative, value, packed_q); + _mm512_storeu_pd(values.data() + i, value); + } + } +#endif #if defined(CRYPTO3_SMALL_PRIME_FFT_AVX2_SELECTED) if constexpr (UseFma) { if (values.size() >= 4) { 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 190232db7..72a530d14 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 @@ -39,6 +39,12 @@ #if defined(__x86_64__) && defined(__AVX2__) && defined(__FMA__) && (defined(__GNUC__) || defined(__clang__)) #include #define CRYPTO3_SMALL_PRIME_FFT_AVX2_SELECTED +#if defined(__AVX512F__) && defined(__AVX512DQ__) +#define CRYPTO3_SMALL_PRIME_FFT_AVX512_AVAILABLE +#if !defined(CRYPTO3_SMALL_PRIME_FFT_DISABLE_AVX512) +#define CRYPTO3_SMALL_PRIME_FFT_AVX512_SELECTED +#endif +#endif #endif namespace nil::crypto3::math::detail { @@ -164,6 +170,21 @@ namespace nil::crypto3::math::detail { */ class packed_fma_kernel { public: + using value_type = __m256d; + static constexpr std::size_t lanes = 4; + + static value_type broadcast(double value) { + return _mm256_set1_pd(value); + } + + static value_type load(const double *values) { + return _mm256_loadu_pd(values); + } + + static void store(double *values, value_type value) { + _mm256_storeu_pd(values, value); + } + 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_)) { @@ -227,11 +248,191 @@ namespace nil::crypto3::math::detail { return high; } + /** + * Two independent compensated products. The distinct left vectors are overwritten; + * right operands are copied before either write. Interleave the independent product + * and quotient work while sharing one quotient scratch register. Bounds match multiply_fma. + */ + void multiply_fma_pair(__m256d &left0, __m256d &left1, __m256d right0, __m256d right1) const { + assert(&left0 != &left1); + assert(arithmetic_.fma_supported()); +#ifndef NDEBUG + double a[8], b[8]; + _mm256_storeu_pd(a, left0); + _mm256_storeu_pd(a + 4, left1); + _mm256_storeu_pd(b, right0); + _mm256_storeu_pd(b + 4, right1); + for (std::size_t i = 0; i < 8; ++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 high0, high1, quotient; + // The early-clobber left operands hold their exact low products after vfmsub213pd. + // Keep right operands distinct from both outputs until both residuals are computed. + // clang-format off + asm( + "vmulpd %[right0], %[left0], %[high0]\n" + "vmulpd %[right1], %[left1], %[high1]\n" + "vmulpd %[reciprocal], %[high0], %[quotient]\n" + "vfmsub213pd %[high0], %[right0], %[left0]\n" + "vfmsub213pd %[high1], %[right1], %[left1]\n" + "vroundpd $8, %[quotient], %[quotient]\n" + "vfnmadd231pd %[modulus], %[quotient], %[high0]\n" + "vmulpd %[reciprocal], %[high1], %[quotient]\n" + "vroundpd $8, %[quotient], %[quotient]\n" + "vfnmadd231pd %[modulus], %[quotient], %[high1]\n" + "vaddpd %[high0], %[left0], %[left0]\n" + "vaddpd %[high1], %[left1], %[left1]\n" + : [left0] "+&x"(left0), [left1] "+&x"(left1), + [high0] "=&x"(high0), [high1] "=&x"(high1), [quotient] "=&x"(quotient) + : [right0] "x"(right0), [right1] "x"(right1), + [modulus] "x"(modulus_), [reciprocal] "x"(reciprocal_) + ); + // clang-format on + } + private: const small_prime_fft_arithmetic &arithmetic_; const __m256d modulus_; const __m256d reciprocal_; }; + +#if defined(CRYPTO3_SMALL_PRIME_FFT_AVX512_AVAILABLE) + // Eight independent lanes use the same reciprocal certificate and exact product bounds. + // AVX-512 register constraints permit all 32 ZMM registers; no runtime CPU dispatch is + // needed because the compiler target already requires AVX-512F/DQ. + class wide_fma_kernel { + public: + using value_type = __m512d; + static constexpr std::size_t lanes = 8; + + static value_type broadcast(double value) { + return _mm512_set1_pd(value); + } + + static value_type load(const double *values) { + return _mm512_loadu_pd(values); + } + + static void store(double *values, value_type value) { + _mm512_storeu_pd(values, value); + } + + explicit wide_fma_kernel(const small_prime_fft_arithmetic &arithmetic) : + arithmetic_(arithmetic), modulus_(_mm512_set1_pd(static_cast(arithmetic.modulus_))), + reciprocal_(_mm512_set1_pd(arithmetic.reciprocal_)) { + assert(arithmetic.fma_supported()); + } + + /** Eight independent reductions with the same preconditions and bounds as the scalar overload. */ + __m512d reduce_fma(__m512d value) const { + assert(arithmetic_.fma_supported()); +#ifndef NDEBUG + double lanes[8]; + _mm512_storeu_pd(lanes, value); + for (double lane : lanes) { + assert(valid_integer(lane)); + } +#endif + __m512d 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" + "vrndscalepd $8, %[quotient], %[quotient]\n" + "vfnmadd231pd %[modulus], %[quotient], %[value]\n" + : [value] "+v"(value), [quotient] "=&v"(quotient) + : [modulus] "v"(modulus_), [reciprocal] "v"(reciprocal_) + ); + // clang-format on + return value; + } + + /** Eight independent compensated products; the scalar product bounds apply to every lane. */ + __m512d multiply_fma(__m512d left, __m512d right) const { + assert(arithmetic_.fma_supported()); +#ifndef NDEBUG + double a[8], b[8]; + _mm512_storeu_pd(a, left); + _mm512_storeu_pd(b, right); + for (std::size_t i = 0; i < 8; ++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 + __m512d 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" + "vrndscalepd $8, %[quotient], %[quotient]\n" + "vfnmadd231pd %[modulus], %[quotient], %[high]\n" + "vaddpd %[low], %[high], %[high]\n" + : [high] "=&v"(high), [low] "=&v"(low), [quotient] "=&v"(quotient) + : [left] "v"(left), [right] "v"(right), + [modulus] "v"(modulus_), [reciprocal] "v"(reciprocal_) + ); + // clang-format on + return high; + } + + /** + * Two independent compensated products. The distinct left vectors are overwritten; + * right operands are copied before either write. Interleave the independent product + * and quotient work with separate scratch registers. Bounds match multiply_fma. + */ + void multiply_fma_pair(__m512d &left0, __m512d &left1, __m512d right0, __m512d right1) const { + assert(&left0 != &left1); + assert(arithmetic_.fma_supported()); +#ifndef NDEBUG + double a[16], b[16]; + _mm512_storeu_pd(a, left0); + _mm512_storeu_pd(a + 8, left1); + _mm512_storeu_pd(b, right0); + _mm512_storeu_pd(b + 8, right1); + for (std::size_t i = 0; i < 16; ++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 + __m512d high0, high1, quotient0, quotient1; + // The early-clobber left operands hold their exact low products after vfmsub213pd. + // Keep right operands distinct from both outputs until both residuals are computed. + // clang-format off + asm( + "vmulpd %[right0], %[left0], %[high0]\n" + "vmulpd %[right1], %[left1], %[high1]\n" + "vmulpd %[reciprocal], %[high0], %[quotient0]\n" + "vmulpd %[reciprocal], %[high1], %[quotient1]\n" + "vfmsub213pd %[high0], %[right0], %[left0]\n" + "vfmsub213pd %[high1], %[right1], %[left1]\n" + "vrndscalepd $8, %[quotient0], %[quotient0]\n" + "vrndscalepd $8, %[quotient1], %[quotient1]\n" + "vfnmadd231pd %[modulus], %[quotient0], %[high0]\n" + "vfnmadd231pd %[modulus], %[quotient1], %[high1]\n" + "vaddpd %[high0], %[left0], %[left0]\n" + "vaddpd %[high1], %[left1], %[left1]\n" + : [left0] "+&v"(left0), [left1] "+&v"(left1), + [high0] "=&v"(high0), [high1] "=&v"(high1), + [quotient0] "=&v"(quotient0), [quotient1] "=&v"(quotient1) + : [right0] "v"(right0), [right1] "v"(right1), + [modulus] "v"(modulus_), [reciprocal] "v"(reciprocal_) + ); + // clang-format on + } + + private: + const small_prime_fft_arithmetic &arithmetic_; + const __m512d modulus_; + const __m512d reciprocal_; + }; +#endif #endif /** Exact fallback for reduce_fma, with the same input/output bounds and any rounding mode. */ 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 57fa79ad2..6876af248 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 @@ -290,7 +290,7 @@ namespace nil::crypto3::math::polynomial_arithmetic { spectrum(2 * row + 1), spectrum(4 + Columns + column), size, normalizations_[plan_offset + prime]); for (std::size_t coordinate = 0; coordinate < field_type::arity; ++coordinate) { - plan.inverse_fft(destination.subspan(coordinate * size, size), false); + plan.inverse_fft_pointwise(destination.subspan(coordinate * size, size)); } } const std::size_t index = row * Columns + column; @@ -580,7 +580,7 @@ namespace nil::crypto3::math::polynomial_arithmetic { pointwise_[prime].multiply(a, b, size, normalizations_[plan_offset + prime]); } for (std::size_t coordinate = 0; coordinate < field_type::arity; ++coordinate) { - plan.inverse_fft(a.subspan(coordinate * size, size), false); + plan.inverse_fft_pointwise(a.subspan(coordinate * size, size)); } } result_.resize(result_size); diff --git a/libs/math/test/small_prime_fft_arithmetic.cpp b/libs/math/test/small_prime_fft_arithmetic.cpp index e65354d71..0b32f3864 100644 --- a/libs/math/test/small_prime_fft_arithmetic.cpp +++ b/libs/math/test/small_prime_fft_arithmetic.cpp @@ -247,7 +247,9 @@ 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) { +template +void check_packed_products() { + constexpr auto lanes = Packed::lanes; restore_rounding restore; BOOST_REQUIRE_EQUAL(std::fesetround(FE_TONEAREST), 0); std::mt19937_64 random(0x41565832); @@ -256,7 +258,7 @@ BOOST_AUTO_TEST_CASE(packed_products_match_scalar_and_exact_integer_oracles) { if (!arithmetic.fma_supported()) { continue; } - const arithmetic_type::packed_fma_kernel packed(arithmetic); + const Packed packed(arithmetic); const auto q = static_cast(prime.modulus); const std::array, 12> boundaries {{{0, 4 * q - 1}, {q - 1, 1 - q}, @@ -271,25 +273,37 @@ BOOST_AUTO_TEST_CASE(packed_products_match_scalar_and_exact_integer_oracles) { {1, -1}, {-q, q}}}; for (std::size_t batch = 0; batch < 1003; ++batch) { - alignas(32) std::array left {}, right {}, output; + alignas(64) 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] : + for (std::size_t lane = 0; lane < lanes; ++lane) { + const auto a = batch < 3 ? boundaries[(lanes * batch + lane) % boundaries.size()][0] : static_cast(random() % (8 * q - 1)) - (4 * q - 1); - const auto b = batch < 3 ? boundaries[4 * batch + lane][1] : + const auto b = batch < 3 ? boundaries[(lanes * batch + lane) % boundaries.size()][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) { + Packed::store(output.data() + 1, + packed.multiply_fma(Packed::load(left.data() + 1), Packed::load(right.data() + 1))); + std::array paired, reversed, reversed_left, reversed_right; + for (std::size_t lane = 0; lane < lanes; ++lane) { + reversed_left[lane] = left[lanes - lane]; + reversed_right[lane] = right[lanes - lane]; + } + auto a = Packed::load(left.data() + 1); + auto b = Packed::load(reversed_left.data()); + packed.multiply_fma_pair(a, b, Packed::load(right.data() + 1), Packed::load(reversed_right.data())); + Packed::store(paired.data(), a); + Packed::store(reversed.data(), b); + for (std::size_t lane = 1; lane <= lanes; ++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(paired[lane - 1], output[lane]); + BOOST_CHECK_EQUAL(reversed[lanes - lane], output[lane]); } } BOOST_CHECK_EQUAL(output.front(), -19.0); @@ -299,7 +313,9 @@ BOOST_AUTO_TEST_CASE(packed_products_match_scalar_and_exact_integer_oracles) { BOOST_CHECK_EQUAL(std::fegetround(), FE_TONEAREST); } -BOOST_AUTO_TEST_CASE(packed_reductions_preserve_lane_order_and_large_integer_bounds) { +template +void check_packed_reductions() { + constexpr auto lanes = Packed::lanes; restore_rounding restore; BOOST_REQUIRE_EQUAL(std::fesetround(FE_TONEAREST), 0); std::mt19937_64 random(0x41565833); @@ -309,19 +325,22 @@ BOOST_AUTO_TEST_CASE(packed_reductions_preserve_lane_order_and_large_integer_bou if (!arithmetic.fma_supported()) { continue; } - const arithmetic_type::packed_fma_kernel packed(arithmetic); + const Packed packed(arithmetic); const auto q = static_cast(prime.modulus); for (std::size_t batch = 0; batch < 1001; ++batch) { - alignas(32) std::array input {}, output; + alignas(64) std::array input {}, output; output.fill(-23.0); - for (std::size_t lane = 1; lane <= 4; ++lane) { + for (std::size_t lane = 1; lane <= lanes; ++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}; + input[1] = double(limit - 1); + input[2] = double(1 - limit); + input[3] = double(q / 2); + input[4] = double(-q / 2 - 1); } - _mm256_storeu_pd(output.data() + 1, packed.reduce_fma(_mm256_loadu_pd(input.data() + 1))); - for (std::size_t lane = 1; lane <= 4; ++lane) { + Packed::store(output.data() + 1, packed.reduce_fma(Packed::load(input.data() + 1))); + for (std::size_t lane = 1; lane <= lanes; ++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)); @@ -334,6 +353,25 @@ BOOST_AUTO_TEST_CASE(packed_reductions_preserve_lane_order_and_large_integer_bou } BOOST_CHECK_EQUAL(std::fegetround(), FE_TONEAREST); } + +BOOST_AUTO_TEST_CASE(packed_products_match_scalar_and_exact_integer_oracles) { + check_packed_products(); +} + +BOOST_AUTO_TEST_CASE(packed_reductions_preserve_lane_order_and_large_integer_bounds) { + check_packed_reductions(); +} + +#if defined(CRYPTO3_SMALL_PRIME_FFT_AVX512_AVAILABLE) +BOOST_AUTO_TEST_CASE(wide_products_match_scalar_and_exact_integer_oracles) { + check_packed_products(); +} + +BOOST_AUTO_TEST_CASE(wide_reductions_preserve_lane_order_and_large_integer_bounds) { + check_packed_reductions(); +} +#endif + #endif BOOST_AUTO_TEST_SUITE_END() diff --git a/libs/math/test/small_prime_fft_transform.cpp b/libs/math/test/small_prime_fft_transform.cpp index 9d53195e8..39d60aa1d 100644 --- a/libs/math/test/small_prime_fft_transform.cpp +++ b/libs/math/test/small_prime_fft_transform.cpp @@ -136,7 +136,7 @@ BOOST_AUTO_TEST_SUITE(small_prime_fft_transform_tests) BOOST_AUTO_TEST_CASE(forward_matches_integer_dft_for_every_prime) { std::mt19937_64 random(0x46465401); for (std::size_t prime = 0; prime < detail::small_prime_fft_primes.size(); ++prime) { - for (std::size_t size = 1; size <= 64; size *= 2) { + for (std::size_t size = 1; size <= 256; size *= 2) { BOOST_TEST_CONTEXT("prime=" << prime << " N=" << size) { const plan_type plan(prime, size); const auto q = detail::small_prime_fft_primes[prime].modulus; @@ -160,7 +160,7 @@ BOOST_AUTO_TEST_CASE(forward_matches_integer_dft_for_every_prime) { BOOST_AUTO_TEST_CASE(inverse_matches_independent_spectra) { std::mt19937_64 random(0x46465402); for (std::size_t prime = 0; prime < detail::small_prime_fft_primes.size(); ++prime) { - for (std::size_t size = 1; size <= 64; size *= 2) { + for (std::size_t size = 1; size <= 256; size *= 2) { BOOST_TEST_CONTEXT("prime=" << prime << " N=" << size) { const plan_type plan(prime, size); const auto q = plan.modulus(); @@ -334,7 +334,7 @@ BOOST_AUTO_TEST_CASE(rounding_modes_preserve_exactness_and_environment) { BOOST_AUTO_TEST_CASE(plan_reuse_and_in_place_buffer_boundaries) { std::mt19937_64 random(0x46465406); - for (const std::size_t size : {1, 2, 4, 8, 16, 32, 64}) { + for (const std::size_t size : {1, 2, 4, 8, 16, 32, 64, 128, 256, 512}) { const plan_type original(0, size); auto copy = original; const plan_type moved = std::move(copy); @@ -365,6 +365,86 @@ BOOST_AUTO_TEST_CASE(plan_reuse_and_in_place_buffer_boundaries) { } } +BOOST_AUTO_TEST_CASE(pointwise_inverse_preserves_bounds_and_matches_general_inverse) { + restore_rounding restore; + std::mt19937_64 random(0x46465407); + for (std::size_t prime = 0; prime < detail::small_prime_fft_primes.size(); ++prime) { + for (const std::size_t size : {1, 2, 4, 8, 16, 32, 64, 128, 512, 4096, 32768}) { + const plan_type plan(prime, size); + const auto q = static_cast(plan.modulus()); + std::vector input(size); + for (std::size_t i = 0; i < size; ++i) { + input[i] = i % 3 == 0 ? double(q - 1) : + i % 3 == 1 ? double(1 - q) : + double(static_cast(random() % (2 * q - 1)) - (q - 1)); + } + for (int mode : {FE_TONEAREST, FE_DOWNWARD}) { + BOOST_REQUIRE_EQUAL(std::fesetround(mode), 0); + auto general = input; + auto bounded = input; + plan.inverse_fft(general, false); + plan.inverse_fft_pointwise(bounded); + for (std::size_t i = 0; i < size; ++i) { + BOOST_CHECK_EQUAL(std::trunc(bounded[i]), bounded[i]); + BOOST_CHECK(std::abs(bounded[i]) < 2 * double(q)); + BOOST_CHECK_EQUAL(canonical(as_integer(bounded[i]), cpp_int(q)), + canonical(as_integer(general[i]), cpp_int(q))); + } + BOOST_CHECK_EQUAL(std::fegetround(), mode); + } + } + } +} + +BOOST_AUTO_TEST_CASE(bit_reversed_order_is_exact_across_rounding_mode_changes_and_unaligned_buffers) { + restore_rounding restore; + std::mt19937_64 random(0x46465408); + for (std::size_t prime = 0; prime < detail::small_prime_fft_primes.size(); ++prime) { + for (const std::size_t size : {1, 2, 4, 8, 16, 32, 64, 128, 256, 1024, 4096}) { + BOOST_REQUIRE_EQUAL(std::fesetround(FE_TONEAREST), 0); + const plan_type plan(prime, size); + const plan_type ordinary(prime, size); + const detail::small_prime_fft_arithmetic arithmetic(plan.modulus()); + const auto input = lazy_input(size, plan.modulus(), random); + auto reference = input; + ordinary.fft(reference); + const auto direct = size <= 64 ? dft(input, prime) : std::vector(); + // Switching modes between calls preserves ordering across FMA and integer arithmetic. + for (int forward_mode : {FE_TONEAREST, FE_DOWNWARD}) { + const int inverse_mode = forward_mode == FE_TONEAREST ? FE_DOWNWARD : FE_TONEAREST; + std::vector storage(size + 5, -29.0); + auto values = std::span(storage).subspan(1, size); + std::copy(input.begin(), input.end(), values.begin()); + BOOST_REQUIRE_EQUAL(std::fesetround(forward_mode), 0); + plan.fft(values); + for (std::size_t i = 0; i < size; ++i) { + const cpp_int expected = direct.empty() ? + canonical(as_integer(reference[i]), cpp_int(plan.modulus())) : + direct[detail::bitreverse(i, std::countr_zero(size))]; + check_residue(values[i], expected, plan.modulus(), 3); + } + auto pointwise = std::vector(values.begin(), values.end()); + for (double &value : pointwise) { + value = arithmetic.multiply_integer(value, plan.size_inverse()); + } + BOOST_REQUIRE_EQUAL(std::fesetround(inverse_mode), 0); + plan.inverse_fft(values); + plan.inverse_fft_pointwise(pointwise); + for (std::size_t i = 0; i < size; ++i) { + const auto expected = canonical(as_integer(input[i]), cpp_int(plan.modulus())); + check_residue(values[i], expected, plan.modulus(), 1, true); + check_residue(pointwise[i], expected, plan.modulus(), 2); + } + BOOST_CHECK_EQUAL(storage.front(), -29.0); + for (std::size_t i = size + 1; i < storage.size(); ++i) { + BOOST_CHECK_EQUAL(storage[i], -29.0); + } + BOOST_CHECK_EQUAL(std::fegetround(), inverse_mode); + } + } + } +} + BOOST_AUTO_TEST_CASE(invalid_parameters_and_dimensions_throw_before_writes) { BOOST_CHECK_THROW(plan_type(detail::small_prime_fft_primes.size(), 1), std::invalid_argument); BOOST_CHECK_THROW(plan_type(std::numeric_limits::max(), 1), std::invalid_argument); @@ -385,8 +465,40 @@ BOOST_AUTO_TEST_CASE(invalid_parameters_and_dimensions_throw_before_writes) { BOOST_CHECK_THROW(plan.fft(values), std::invalid_argument); BOOST_CHECK_THROW(plan.inverse_fft(values), std::invalid_argument); BOOST_CHECK_THROW(plan.inverse_fft(values, false), std::invalid_argument); + BOOST_CHECK_THROW(plan.inverse_fft_pointwise(values), std::invalid_argument); BOOST_CHECK_EQUAL_COLLECTIONS(values.begin(), values.end(), before.begin(), before.end()); } } +BOOST_AUTO_TEST_CASE(blocked_transforms_match_direct_evaluations_at_large_even_and_odd_depths) { + std::mt19937_64 random(0x46465409); + for (const std::size_t prime : {0, 10}) { + for (const std::size_t size : {2048, 4096, 8192, 16384, 32768, 65536}) { + { + const plan_type plan(prime, size); + const auto input = lazy_input(size, plan.modulus(), random); + std::vector storage(size + 2, -31.0); + auto values = std::span(storage).subspan(1, size); + std::copy(input.begin(), input.end(), values.begin()); + plan.fft(values); + const cpp_int q = plan.modulus(); + for (const std::size_t k : {std::size_t(0), size / 7, size / 3, size / 2, size - 1}) { + const cpp_int step = boost::multiprecision::powm(cpp_int(plan.omega()), cpp_int(k), q); + cpp_int expected = 0; + for (std::size_t i = size; i-- > 0;) { + expected = canonical(expected * step + as_integer(input[i]), q); + } + check_residue(values[detail::bitreverse(k, std::countr_zero(size))], expected, plan.modulus(), 3); + } + plan.inverse_fft(values); + for (std::size_t i = 0; i < size; ++i) { + check_residue(values[i], canonical(as_integer(input[i]), q), plan.modulus(), 1, true); + } + BOOST_CHECK_EQUAL(storage.front(), -31.0); + BOOST_CHECK_EQUAL(storage.back(), -31.0); + } + } + } +} + BOOST_AUTO_TEST_SUITE_END()