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 33a919c63..27256ca4e 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 @@ -1823,8 +1823,8 @@ 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 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. + * identity with scalar, four-lane or eight-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 * |scale| < q/2, normally the FFT plan's centered 1/N. Buffers are disjoint or identical. @@ -1832,7 +1832,7 @@ namespace nil::crypto3::math::detail { * 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. +// The packed kernels use vector coordinates in the existing tower arrays. #pragma GCC diagnostic push #pragma GCC diagnostic ignored "-Wignored-attributes" #endif @@ -1862,42 +1862,69 @@ namespace nil::crypto3::math::detail { } } + BOOST_FORCEINLINE void multiply_pair(Value &a, Value &b, Value wa, Value wb) const { +#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(a, wa); + b = multiply(b, wb); + } + } + + template Value add(Value a, Value b) const { - return reduce(a + b); + if constexpr (Reduce) { + return reduce(a + b); + } else { + return a + b; + } } - Value subtract(Value a, Value b) const { - return reduce(a - b); + // Three reduced coordinates give an exact integer of magnitude < 3*q. + // Callers can defer this reduction to the next bounded sum or final scaling. + template + Value subtract(Value a, Value b, Value c) const { + if constexpr (Reduce) { + return reduce(a - b - c); + } else { + return a - b - c; + } } - template + template std::array add(const std::array &a, const std::array &b) const { std::array result; for (std::size_t i = 0; i < N; ++i) { - result[i] = add(a[i], b[i]); + result[i] = add(a[i], b[i]); } return result; } - template - std::array subtract(const std::array &a, const std::array &b) const { + template + std::array subtract(const std::array &a, const std::array &b, + const std::array &c) const { std::array result; for (std::size_t i = 0; i < N; ++i) { - result[i] = subtract(a[i], b[i]); + result[i] = subtract(a[i], b[i], c[i]); } return result; } - // All tower inputs and returned coordinates are in (-q,q). Pair sums are < 2*q; + // Quadratic inputs and returned coordinates are in (-q,q). Pair sums are < 2*q; // the Karatsuba product is < 4*q^2 and its remainder < 1.5*q. Correcting by the // two diagonal products gives a sum < 3.5*q, hence exact in binary64 before reduction. - quadratic_type multiply(const quadratic_type &a, const quadratic_type &b) const { + BOOST_FORCEINLINE quadratic_type multiply(const quadratic_type &a, const quadratic_type &b) const { #if defined(__clang__) #pragma clang fp reassociate(off) #pragma clang fp contract(off) #endif - const Value diagonal0 = multiply(a[0], b[0]); - const Value diagonal1 = multiply(a[1], b[1]); + // Keep the two independent products in registers and interleave their FMA chains. + Value diagonal0 = a[0]; + Value diagonal1 = a[1]; + multiply_pair(diagonal0, diagonal1, b[0], b[1]); const Value cross = multiply(a[0] + a[1], b[0] + b[1]); return {reduce(diagonal0 - diagonal1), reduce(cross - diagonal0 - diagonal1)}; } @@ -1912,6 +1939,16 @@ namespace nil::crypto3::math::detail { return {reduce(reduce(8 * a[0]) + a[0] - a[1]), reduce(reduce(8 * a[1]) + a[1] + a[0])}; } + quadratic_type add_multiply_xi(const quadratic_type &a, const quadratic_type &b) const { +#if defined(__clang__) +#pragma clang fp reassociate(off) +#pragma clang fp contract(off) +#endif + // a + (9+u)*b, with |a| < 3*q and |b| < q. After reducing 8*b, + // each four-term sum is < 6*q < 2^53, so only its final reduction is needed. + return {reduce(reduce(8 * b[0]) + b[0] - b[1] + a[0]), reduce(reduce(8 * b[1]) + b[1] + b[0] + a[1])}; + } + cubic_type multiply_v(const cubic_type &a) const { return {multiply_xi(a[2]), a[0], a[1]}; } @@ -1920,21 +1957,35 @@ namespace nil::crypto3::math::detail { const auto diagonal0 = multiply(a[0], b[0]); const auto diagonal1 = multiply(a[1], b[1]); const auto diagonal2 = multiply(a[2], b[2]); - const auto cross12 = - subtract(subtract(multiply(add(a[1], a[2]), add(b[1], b[2])), diagonal1), diagonal2); - const auto cross01 = - subtract(subtract(multiply(add(a[0], a[1]), add(b[0], b[1])), diagonal0), diagonal1); - const auto cross02 = - subtract(subtract(multiply(add(a[0], a[2]), add(b[0], b[2])), diagonal0), diagonal2); - return {add(diagonal0, multiply_xi(cross12)), add(cross01, multiply_xi(diagonal2)), - add(cross02, diagonal1)}; + // cross12 is reduced before multiplication by 8 inside multiply_xi. + // The other cross terms stay < 3*q until these final sums; the cubic result is < q. + // Finish each output pair before forming the next cross term to limit live temporaries. + cubic_type result; + { + const auto cross12 = subtract(multiply(add(a[1], a[2]), add(b[1], b[2])), diagonal1, diagonal2); + result[0] = add_multiply_xi(diagonal0, cross12); + } + { + const auto cross01 = + subtract(multiply(add(a[0], a[1]), add(b[0], b[1])), diagonal0, diagonal1); + result[1] = add_multiply_xi(cross01, diagonal2); + } + { + const auto cross02 = + subtract(multiply(add(a[0], a[2]), add(b[0], b[2])), diagonal0, diagonal2); + result[2] = add(cross02, diagonal1); + } + return result; } + // Both tower operations return coordinates of magnitude < 3*q. The final + // |scale| < q/2 gives a product < 1.5*q^2, whose FMA remainder is already in (-q,q). + // Keeping these sums unreduced is also safe when two tower products are added below. tower_type multiply(const tower_type &a, const tower_type &b) const { const auto diagonal0 = multiply(a[0], b[0]); const auto diagonal1 = multiply(a[1], b[1]); - const auto cross = subtract(subtract(multiply(add(a[0], a[1]), add(b[0], b[1])), diagonal0), diagonal1); - return {add(diagonal0, multiply_v(diagonal1)), cross}; + const auto cross = subtract(multiply(add(a[0], a[1]), add(b[0], b[1])), diagonal0, diagonal1); + return {add(diagonal0, multiply_v(diagonal1)), cross}; } tower_type square(const tower_type &a) const { @@ -1942,7 +1993,33 @@ namespace nil::crypto3::math::detail { // c0 = (a0+a1)*(a0+v*a1) - a0*a1 - v*a0*a1; c1 = 2*a0*a1. const auto cross = multiply(a[0], a[1]); const auto combined = multiply(add(a[0], a[1]), add(a[0], multiply_v(a[1]))); - return {subtract(subtract(combined, cross), multiply_v(cross)), add(cross, cross)}; + return {subtract(combined, cross, multiply_v(cross)), add(cross, cross)}; + } + + template + BOOST_FORCEINLINE void load(tower_type &a, tower_type &b, std::span left, + std::span right, std::size_t stride, std::size_t index) const { + constexpr std::size_t width = sizeof(Value) / sizeof(double); + 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 + index; + if constexpr (width == 1) { + a[w][v][u] = reduce(left[offset]); + if constexpr (!Square) { + b[w][v][u] = reduce(right[offset]); + } +#if defined(CRYPTO3_SMALL_PRIME_FFT_AVX2_SELECTED) + } else { + a[w][v][u] = reduce(Arithmetic::load(left.data() + offset)); + if constexpr (!Square) { + b[w][v][u] = reduce(Arithmetic::load(right.data() + offset)); + } +#endif + } + } + } + } } }; @@ -2003,11 +2080,22 @@ namespace nil::crypto3::math::detail { std::size_t offset = 0; #if defined(CRYPTO3_SMALL_PRIME_FFT_AVX2_SELECTED) if constexpr (UseFma) { - offset = stride - stride % 4; +#if defined(CRYPTO3_SMALL_PRIME_FFT_AVX512_SELECTED) + using wide_kernel = small_prime_fft_arithmetic::wide_fma_kernel; + offset = stride - stride % wide_kernel::lanes; if (offset != 0) { - const small_prime_fft_arithmetic::packed_fma_kernel packed(arithmetic_); - apply_range(packed, output, left, right, extra_left, extra_right, 0, offset, - _mm256_set1_pd(scale)); + const wide_kernel wide(arithmetic_); + apply_range(wide, output, left, right, extra_left, extra_right, 0, offset, + wide_kernel::broadcast(scale)); + } +#endif + using packed_kernel = small_prime_fft_arithmetic::packed_fma_kernel; + const std::size_t packed_end = stride - stride % packed_kernel::lanes; + if (offset < packed_end) { + const packed_kernel packed(arithmetic_); + apply_range(packed, output, left, right, extra_left, extra_right, offset, + packed_end, packed_kernel::broadcast(scale)); + offset = packed_end; } } #endif @@ -2024,59 +2112,38 @@ namespace nil::crypto3::math::detail { 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, 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) { - const std::size_t offset = (6 * w + 2 * v + u) * stride + i; - if constexpr (width == 1) { - a[w][v][u] = operations.reduce(left[offset]); - 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 = [&]() { + tower_type a, b; + operations.template load(a, b, left, right, stride, i); + tower_type result; + if constexpr (Square) { + result = operations.square(a); + } else { + result = operations.multiply(a, b); if constexpr (Sum) { - // Each product coordinate is in (-q,q); adding two is exact below 2*q. + // Reuse the input slots after the first product. Load both remaining + // inputs before storing anything so output may still alias any input. + operations.template load(a, b, extra_left, extra_right, stride, i); + // Each product coordinate has magnitude < 3*q; their sum is exact below 6*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); + result = operations.add(result, operations.multiply(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) { - // Fold inverse-transform normalization into this existing write pass. - 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) { - output[offset] = value; + // Fold inverse-transform normalization into this existing write pass, + // pairing the independent scales of the two quadratic coordinates. + const std::size_t offset = (6 * w + 2 * v) * stride + i; + Value first = result[w][v][0]; + Value second = result[w][v][1]; + operations.multiply_pair(first, second, scale, scale); + if constexpr (width == 1) { + output[offset] = first; + output[offset + stride] = second; #if defined(CRYPTO3_SMALL_PRIME_FFT_AVX2_SELECTED) - } else { - _mm256_storeu_pd(output.data() + offset, value); + } else { + Arithmetic::store(output.data() + offset, first); + Arithmetic::store(output.data() + offset + stride, second); #endif - } } } } diff --git a/libs/math/test/small_prime_fft_backend.cpp b/libs/math/test/small_prime_fft_backend.cpp index dfe37431a..294d4188e 100644 --- a/libs/math/test/small_prime_fft_backend.cpp +++ b/libs/math/test/small_prime_fft_backend.cpp @@ -280,8 +280,8 @@ BOOST_AUTO_TEST_CASE(pointwise_products_and_squares_match_expanded_integer_tower 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); - // 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}) { + // Cover both vector widths, every eight-lane remainder, and multiple wide blocks. + for (const std::size_t stride : {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 31, 32, 33}) { std::vector left(12 * stride), right(12 * stride); std::vector products(stride), squares(stride); for (std::size_t i = 0; i < stride; ++i) { @@ -354,7 +354,7 @@ BOOST_AUTO_TEST_CASE(pointwise_sums_match_expanded_integer_tower_and_preserve_al 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}) { + for (const std::size_t stride : {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 31, 32, 33}) { std::array, 4> inputs; for (auto &input : inputs) { input.resize(12 * stride); @@ -400,6 +400,59 @@ BOOST_AUTO_TEST_CASE(pointwise_sums_match_expanded_integer_tower_and_preserve_al } } +BOOST_AUTO_TEST_CASE(pointwise_boundary_products_and_cancellation_preserve_rounding_modes) { + restore_rounding restore; + BOOST_REQUIRE_EQUAL(std::fesetround(FE_TONEAREST), 0); + // Multiple wide blocks and a scalar tail, with all tower coordinates active. + constexpr std::size_t stride = 65; + 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); + const std::array edges {0, + 1, + -1, + q - 1, + 1 - q, + (q - 1) / 2, + -(q - 1) / 2, + 3 * q - 1, + 1 - 3 * q, + (std::int64_t(1) << 53) - 1, + 1 - (std::int64_t(1) << 53)}; + std::vector left(12 * stride), right(12 * stride), negative_left(12 * stride); + std::vector products(stride), squares(stride), zero(stride); + for (std::size_t i = 0; i < stride; ++i) { + coordinates_type a {}, b {}; + for (std::size_t c = 0; c < 12; ++c) { + const auto x = edges[(i + 3 * c) % edges.size()]; + const auto y = edges[(i / edges.size() + 7 * c) % edges.size()]; + a[c] = x; + b[c] = y; + left[c * stride + i] = static_cast(x); + right[c * stride + i] = static_cast(y); + negative_left[c * stride + i] = static_cast(-x); + } + products[i] = tower_product(a, b); + squares[i] = tower_product(a, a); + } + for (const int mode : {FE_TONEAREST, FE_DOWNWARD, FE_UPWARD, FE_TOWARDZERO}) { + BOOST_REQUIRE_EQUAL(std::fesetround(mode), 0); + // The largest allowed scales exercise the bound on unreduced tower outputs. + for (const double scale : {static_cast((q - 1) / 2), -static_cast((q - 1) / 2)}) { + auto actual = left; + pointwise.multiply(actual, right, stride, scale); + check_pointwise(actual, products, prime.modulus, scale); + actual = left; + pointwise.square(actual, stride, scale); + check_pointwise(actual, squares, prime.modulus, scale); + pointwise.multiply_sum(actual, left, right, negative_left, right, stride, scale); + check_pointwise(actual, zero, prime.modulus, scale); + BOOST_CHECK_EQUAL(std::fegetround(), mode); + } + } + } +} + BOOST_AUTO_TEST_CASE(fused_matrix_products_match_schoolbook_and_preserve_aliases) { backend_type backend(256); std::mt19937_64 random(0x46465422);