Introduce polynomial_sum for fast batched addition

NilFoundation · May 17, 2024 · 6cb1ae3 · 6cb1ae3
1 parent bd9fc5e
commit 6cb1ae3
Show file tree

Hide file tree

Showing 4 changed files with 92 additions and 11 deletions.
diff --git a/include/nil/crypto3/math/polynomial/polynomial.hpp b/include/nil/crypto3/math/polynomial/polynomial.hpp
@@ -96,6 +96,7 @@ namespace nil {
 
                 polynomial(const polynomial& x) : val(x.val) {
                 }
+
                 polynomial(const polynomial& x, const allocator_type& a) : val(x.val, a) {
                 }
 
@@ -117,13 +118,8 @@ namespace nil {
                     this->operator[](power) = value;
                 }
 
-                explicit polynomial(const container_type &c) : val(c) {
-                    if (val.empty()) {
-                        val.push_back(FieldValueType::zero());
-                    }
-                }
-
-                explicit polynomial(container_type &&c) : val(c) {
+                template<typename T>
+                explicit polynomial(T&& c) : val(std::forward<T>(c)) {
                     if (val.empty()) {
                         val.push_back(FieldValueType::zero());
                     }
@@ -179,6 +175,10 @@ namespace nil {
                     return this->val.__alloc();
                 }
 
+                container_type& get_storage() {
+                    return val;
+                }
+
                 iterator begin() BOOST_NOEXCEPT {
                     return val.begin();
                 }
@@ -383,7 +383,7 @@ namespace nil {
                  * Returns true if polynomial is a one polynomial.
                  */
                 bool is_one() const {
-                    return (*this->begin() == FieldValueType(1)) && 
+                    return (*this->begin() == FieldValueType(1)) &&
                         std::all_of(++this->begin(), this->end(),
                             [](FieldValueType i) { return i == FieldValueType::zero(); });
                 }

diff --git a/include/nil/crypto3/math/polynomial/polynomial_dfs.hpp b/include/nil/crypto3/math/polynomial/polynomial_dfs.hpp
@@ -35,9 +35,10 @@
 #include <iterator>
 #include <unordered_map>
 
-#include <nil/crypto3/math/polynomial/basic_operations.hpp>
-#include <nil/crypto3/math/domains/evaluation_domain.hpp>
 #include <nil/crypto3/math/algorithms/make_evaluation_domain.hpp>
+#include <nil/crypto3/math/domains/evaluation_domain.hpp>
+#include <nil/crypto3/math/polynomial/basic_operations.hpp>
+#include <nil/crypto3/math/polynomial/polynomial.hpp>
 
 namespace nil {
     namespace crypto3 {
@@ -164,6 +165,10 @@ namespace nil {
                     return this->val.__alloc();
                 }
 
+                container_type& get_storage() {
+                    return val;
+                }
+
                 iterator begin() BOOST_NOEXCEPT {
                     return val.begin();
                 }
@@ -761,6 +766,41 @@ namespace nil {
                 return os;
             }
 
+            template<typename FieldType>
+            static inline polynomial_dfs<typename FieldType::value_type> polynomial_sum(
+                    std::vector<math::polynomial_dfs<typename FieldType::value_type>> addends) {
+                using FieldValueType = typename FieldType::value_type;
+                std::size_t max_size = 0;
+                std::unordered_map<std::size_t, polynomial_dfs<FieldValueType>> size_to_part_sum;
+                for (auto& addend : addends) {
+                    max_size = std::max(max_size, addend.size());
+                    auto it = size_to_part_sum.find(addend.size());
+                    if (it == size_to_part_sum.end()) {
+                        size_to_part_sum[addend.size()] = std::move(addend);
+                    } else {
+                        it->second += addend;
+                        // Free the memory we are not going to use anymore.
+                        addend = math::polynomial_dfs<FieldValueType>();
+                    }
+                }
+
+                auto coef_result = polynomial<FieldValueType>(max_size, FieldValueType::zero());
+                for (auto& [_, partial_sum] : size_to_part_sum) {
+                    FieldValueType omega = unity_root<FieldType>(partial_sum.size());
+                    detail::basic_radix2_fft<FieldType>(partial_sum, omega.inversed());
+                    const FieldValueType sconst = FieldValueType(partial_sum.size()).inversed();
+                    for (auto& el : partial_sum) {
+                        el *= sconst;
+                    }
+                    coef_result += polynomial<FieldValueType>(std::move(partial_sum.get_storage()));
+                }
+
+                polynomial_dfs<FieldValueType> dfs_result;
+                dfs_result.from_coefficients(coef_result.get_storage());
+
+                return dfs_result;
+            }
+
             template<typename FieldType>
             static inline polynomial_dfs<typename FieldType::value_type> polynomial_product(
                     std::vector<math::polynomial_dfs<typename FieldType::value_type>> multipliers) {

diff --git a/test/benchmarks/CMakeLists.txt b/test/benchmarks/CMakeLists.txt
@@ -6,7 +6,10 @@
 # http://www.boost.org/LICENSE_1_0.txt
 #---------------------------------------------------------------------------#
 
-find_package(Boost REQUIRED COMPONENTS unit_test_framework timer)
+find_package(Boost REQUIRED COMPONENTS
+    timer
+    unit_test_framework
+)
 
 cm_test_link_libraries(${CMAKE_WORKSPACE_NAME}_${CURRENT_PROJECT_NAME}
 

diff --git a/test/benchmarks/polynomial_dfs_benchmark.cpp b/test/benchmarks/polynomial_dfs_benchmark.cpp
@@ -174,4 +174,42 @@ BENCHMARK_AUTO_TEST_CASE(polynomial_product_test, 20) {
     STOP_TIMER("polynomial_product")
 }
 
+BENCHMARK_AUTO_TEST_CASE(polynomial_sum_real_test, 20) {
+    using Field = nil::crypto3::algebra::fields::bls12_fr<381>;
+
+    std::vector<polynomial_dfs<typename Field::value_type>> random_polynomials;
+    random_polynomials.reserve(8);
+    Field::value_type a = alg_rnd_engine();
+    std::vector<std::size_t> sizes = {23, 15, 21, 16, 22, 17, 18};
+    for (auto size : sizes) {
+        random_polynomials.emplace_back(
+            generate_random_polynomial<Field>(
+                1u << size,
+                alg_rnd_engine
+            )
+        );
+    }
+    auto random_polynomials_copy = random_polynomials;
+
+    START_TIMER("polynomial_naive")
+    std::size_t max_size = 0;
+    for (const auto& polynomial : random_polynomials_copy) {
+        max_size = std::max(max_size, polynomial.size());
+    }
+    auto max_domain = make_evaluation_domain<FieldType>(max_size);
+    for (auto& polynomial : random_polynomials_copy) {
+        polynomial.resize(max_size, nullptr, max_domain);
+    }
+    polynomial_dfs<typename Field::value_type> naive_res(0, max_size);
+    for (auto& polynomial : random_polynomials_copy) {
+        naive_res += std::move(polynomial);
+    }
+    STOP_TIMER("polynomial_naive")
+
+    START_TIMER("polynomial_sum")
+    const auto res = polynomial_sum<Field>(std::move(random_polynomials));
+    STOP_TIMER("polynomial_sum")
+    BOOST_CHECK_EQUAL(naive_res, res);
+}
+
 BOOST_AUTO_TEST_SUITE_END()