Skip to content

Repository navigation

Optimize BN254 small-prime pointwise products #183

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

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

Optimize BN254 small-prime pointwise products #183

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

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
219 changes: 143 additions & 76 deletions libs/math/include/nil/crypto3/math/detail/bn254_small_prime_fft.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1823,16 +1823,16 @@ 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.
* 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.
// The packed kernels use vector coordinates in the existing tower arrays.
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wignored-attributes"
#endif
Expand Down Expand Up @@ -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<bool Reduce = true>
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<bool Reduce = true>
Value subtract(Value a, Value b, Value c) const {
if constexpr (Reduce) {
return reduce(a - b - c);
} else {
return a - b - c;
}
}

template<typename T, std::size_t N>
template<bool Reduce = true, typename T, std::size_t N>
std::array<T, N> add(const std::array<T, N> &a, const std::array<T, N> &b) const {
std::array<T, N> result;
for (std::size_t i = 0; i < N; ++i) {
result[i] = add(a[i], b[i]);
result[i] = add<Reduce>(a[i], b[i]);
}
return result;
}

template<typename T, std::size_t N>
std::array<T, N> subtract(const std::array<T, N> &a, const std::array<T, N> &b) const {
template<bool Reduce = true, typename T, std::size_t N>
std::array<T, N> subtract(const std::array<T, N> &a, const std::array<T, N> &b,
const std::array<T, N> &c) const {
std::array<T, N> result;
for (std::size_t i = 0; i < N; ++i) {
result[i] = subtract(a[i], b[i]);
result[i] = subtract<Reduce>(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)};
}
Expand All @@ -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]};
}
Expand All @@ -1920,29 +1957,69 @@ 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<false>(multiply(add(a[0], a[1]), add(b[0], b[1])), diagonal0, diagonal1);
result[1] = add_multiply_xi(cross01, diagonal2);
}
{
const auto cross02 =
subtract<false>(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<false>(multiply(add(a[0], a[1]), add(b[0], b[1])), diagonal0, diagonal1);
return {add<false>(diagonal0, multiply_v(diagonal1)), cross};
}

tower_type square(const tower_type &a) const {
// The quadratic-extension square needs two cubic products instead of three:
// 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<false>(combined, cross, multiply_v(cross)), add<false>(cross, cross)};
}

template<bool Square>
BOOST_FORCEINLINE void load(tower_type &a, tower_type &b, std::span<const double> left,
std::span<const double> 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
}
}
}
}
}
};

Expand Down Expand Up @@ -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<true, Square, Sum>(packed, output, left, right, extra_left, extra_right, 0, offset,
_mm256_set1_pd(scale));
const wide_kernel wide(arithmetic_);
apply_range<true, Square, Sum>(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<true, Square, Sum>(packed, output, left, right, extra_left, extra_right, offset,
packed_end, packed_kernel::broadcast(scale));
offset = packed_end;
}
}
#endif
Expand All @@ -2024,59 +2112,38 @@ namespace nil::crypto3::math::detail {
const kernel<UseFma, Value, Arithmetic> operations {arithmetic};
using tower_type = typename kernel<UseFma, Value, Arithmetic>::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<Square>(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<false>(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
}
}
}
}
Expand Down
Loading
Loading