refactor(bb): batch commit wire polys

barretenberg/bbup/bbup

@@ -1,6 +1,6 @@
 #!/usr/bin/env bash
 
-set -e
+set -ex
 
 # Colors and symbols
 RED='\033[0;31m'

barretenberg/cpp/src/barretenberg/commitment_schemes/commitment_key.hpp

@@ -14,6 +14,7 @@
  */
 
 #include "barretenberg/common/bb_bench.hpp"
+#include "barretenberg/common/ref_span.hpp"
 #include "barretenberg/constants.hpp"
 #include "barretenberg/ecc/batched_affine_addition/batched_affine_addition.hpp"
 #include "barretenberg/ecc/scalar_multiplication/scalar_multiplication.hpp"
@@ -25,6 +26,8 @@
 #include "barretenberg/srs/global_crs.hpp"
 
 #include <cstddef>
+#include <cstdlib>
+#include <limits>
 #include <memory>
 #include <string_view>
 
@@ -103,6 +106,116 @@ template <class Curve> class CommitmentKey {
         Commitment point(r);
         return point;
     };
+    /**
+     * @brief Batch commitment to multiple polynomials
+     * @details Uses batch_multi_scalar_mul for more efficient processing when committing to multiple polynomials
+     *
+     * @param polynomials vector of polynomial spans to commit to
+     * @return std::vector<Commitment> vector of commitments, one for each polynomial
+     */
+    std::vector<Commitment> batch_commit(RefSpan<Polynomial<Fr>> polynomials,
+                                         size_t max_batch_size = std::numeric_limits<size_t>::max()) const
+    {
+        BB_BENCH_NAME("CommitmentKey::batch_commit");
+        max_batch_size = 1;
+        // std::vector<Commitment> commitments;
+        // commitments.reserve(polynomials.size());
+        // for (const auto& polynomial : polynomials) {
+        //     commitments.emplace_back(commit(polynomial));
+        // }
+        // return commitments;
+        // BB_BENCH_NAME("CommitmentKey::batch_commit");
+        // Check environment variable for max batch size
+        // static const size_t max_batch_size = []() -> size_t {
+        //     const char* env_val = std::getenv("BB_BATCH_MSM_MAX");
+        //     if (env_val != nullptr) {
+        //         return static_cast<size_t>(std::stoull(env_val));
+        //     }
+        //     return static_cast<size_t>(4); // Default to 4 polynomials
+        // }();
+
+        std::span<const G1> point_table = srs->get_monomial_points();
+
+        // We can only commit max_batch_size at a time
+        // This is to prevent excessive memory usage in the pippenger algorithm
+
+        // First batch, create the commitments vector
+        std::vector<Commitment> commitments;
+
+        for (size_t i = 0; i < polynomials.size();) {
+            // Note: have to be careful how we compute this to not overlow e.g. max_batch_size + 1 would
+            size_t batch_size = std::min(max_batch_size, polynomials.size() - i);
+            size_t batch_end = i + batch_size;
+
+            // Prepare spans for batch MSM
+            std::vector<std::span<const G1>> points_spans;
+            // Note, we need to const_cast unfortunately as pippenger takes non-const spans
+            // as it converts back and forth from montgomery form
+            std::vector<std::span<Fr>> scalar_spans;
+
+            for (auto& polynomial : polynomials.subspan(i, batch_end - i)) {
+                std::span<const G1> point_table = srs->get_monomial_points().subspan(polynomial.start_index());
+                size_t consumed_srs = polynomial.start_index() + polynomial.size();
+                if (consumed_srs > srs->get_monomial_size()) {
+                    throw_or_abort(format("Attempting to commit to a polynomial that needs ",
+                                          consumed_srs,
+                                          " points with an SRS of size ",
+                                          srs->get_monomial_size()));
+                }
+                scalar_spans.emplace_back(polynomial.coeffs());
+                points_spans.emplace_back(point_table);
+            }
+
+            // Perform batch MSM
+            auto results = scalar_multiplication::MSM<Curve>::batch_multi_scalar_mul(points_spans, scalar_spans, false);
+            for (const auto& result : results) {
+                commitments.emplace_back(result);
+            }
+            i += batch_size;
+        }
+        return commitments;
+    };
+
+    // helper builder struct for constructing a batch to commit at once
+    struct CommitBatch {
+        CommitmentKey* key;
+        RefVector<Polynomial<Fr>> wires;
+        RefVector<const std::string> labels;
+        void mask_and_send_to_verifier(auto transcript, size_t max_batch_size = std::numeric_limits<size_t>::max())
+        {
+            send_to_verifier(transcript, max_batch_size, true);
+        }
+        void send_to_verifier(auto transcript, bool is_zk)
+        {
+            send_to_verifier(transcript, std::numeric_limits<size_t>::max(), is_zk);
+        }
+        void send_to_verifier(auto transcript,
+                              size_t max_batch_size = std::numeric_limits<size_t>::max(),
+                              bool is_zk = false)
+        {
+            if (is_zk) {
+                for (auto& wire : wires) {
+                    wire.mask();
+                }
+            }
+            (void)max_batch_size;
+            for (auto [label, wire] : zip_view(labels, wires)) {
+                transcript->send_to_verifier(label, key->commit(wire));
+            }
+            // std::vector<Commitment> commitments = key->batch_commit(wires, max_batch_size);
+            // for (size_t i = 0; i < commitments.size(); ++i) {
+            //     transcript->send_to_verifier(labels[i], commitments[i]);
+            // }
+        }
+
+        void add_to_batch(Polynomial<Fr>& poly, const std::string& label)
+        {
+            wires.push_back(poly);
+            labels.push_back(label);
+        }
+    };
+
+    CommitBatch start_batch() { return CommitBatch{ this, {}, {} }; }
 
     /**
      * @brief Efficiently commit to a polynomial whose nonzero elements are arranged in discrete blocks
@@ -253,17 +366,14 @@ template <class Curve> class CommitmentKey {
         return result;
     }
 
-    enum class CommitType { Default, Structured, Sparse, StructuredNonZeroComplement };
+    enum class CommitType { Default, StructuredNonZeroComplement };
 
     Commitment commit_with_type(PolynomialSpan<const Fr> poly,
                                 CommitType type,
                                 const std::vector<std::pair<size_t, size_t>>& active_ranges = {},
                                 size_t final_active_wire_idx = 0)
     {
         switch (type) {
-        case CommitType::Structured:
-        case CommitType::Sparse:
-            return commit(poly);
         case CommitType::StructuredNonZeroComplement:
             return commit_structured_with_nonzero_complement(poly, active_ranges, final_active_wire_idx);
         case CommitType::Default:

barretenberg/cpp/src/barretenberg/common/parallel_for_mutex_pool.cpp

@@ -41,7 +41,7 @@ class ThreadPool {
         do_iterations();
 
         {
-            BB_BENCH_NAME("spinning main thread");
+            // BB_BENCH_NAME("spinning main thread");
             std::unique_lock<std::mutex> lock(tasks_mutex);
             complete_condition_.wait(lock, [this] { return complete_ == num_iterations_; });
         }
@@ -72,7 +72,7 @@ class ThreadPool {
                 }
                 iteration = iteration_++;
             }
-            BB_BENCH_NAME("do_iterations()");
+            // BB_BENCH_NAME("do_iterations()");
             task_(iteration);
             {
                 std::unique_lock<std::mutex> lock(tasks_mutex);

barretenberg/cpp/src/barretenberg/common/ref_span.hpp

@@ -27,7 +27,7 @@ template <typename T> class RefSpan {
     {}
 
     // Constructor from an array of pointers and size
-    RefSpan(T** ptr_array, std::size_t size)
+    RefSpan(T* const* ptr_array, std::size_t size)
         : storage(ptr_array)
         , array_size(size)
     {}
@@ -57,6 +57,18 @@ template <typename T> class RefSpan {
     // Get size of the RefSpan
     constexpr std::size_t size() const { return array_size; }
 
+    RefSpan subspan(std::size_t offset, std::size_t count)
+    {
+        // NOTE: like std::span, assumes the caller ensures offset and count are within bounds.
+        return RefSpan(storage + offset, count);
+    }
+
+    RefSpan subspan(std::size_t offset)
+    {
+        // NOTE: like std::span, assumes the caller ensures offset and count are within bounds.
+        return RefSpan(storage + offset, array_size - offset);
+    }
+
     // Iterator implementation
     class iterator {
       public:

barretenberg/cpp/src/barretenberg/common/ref_vector.hpp

@@ -97,7 +97,7 @@ template <typename T> class RefVector {
 
     std::size_t size() const { return storage.size(); }
 
-    void push_back(T& element) { storage.push_back(element); }
+    void push_back(T& element) { storage.push_back(&element); }
     iterator begin() const { return iterator(this, 0); }
     iterator end() const { return iterator(this, storage.size()); }
 

barretenberg/cpp/src/barretenberg/ecc/scalar_multiplication/scalar_multiplication.cpp

@@ -115,7 +115,7 @@ void MSM<Curve>::transform_scalar_and_get_nonzero_scalar_indices(std::span<typen
  */
 template <typename Curve>
 std::vector<typename MSM<Curve>::ThreadWorkUnits> MSM<Curve>::get_work_units(
-    std::vector<std::span<ScalarField>>& scalars, std::vector<std::vector<uint32_t>>& msm_scalar_indices) noexcept
+    std::span<std::span<ScalarField>> scalars, std::vector<std::vector<uint32_t>>& msm_scalar_indices) noexcept
 {
 
     const size_t num_msms = scalars.size();
@@ -599,9 +599,10 @@ void MSM<Curve>::consume_point_schedule(std::span<const uint64_t> point_schedule
         }
 
         // We do some branchless programming here to minimize instruction pipeline flushes
-        // TODO(@zac-williamson, cc @ludamad) check these ternary operators are not branching!
-        // We are iterating through our points and can come across the following scenarios:
-        // 1: The next 2 points in `point_schedule` belong to the *same* bucket
+        // TODO(@zac-williamson, cc @ludamad) check these ternary operators are not branching! -> (ludamad: they don't,
+        // but its not clear that the conditional move is fundamentally less expensive)
+        // We are iterating through our points and
+        // can come across the following scenarios: 1: The next 2 points in `point_schedule` belong to the *same* bucket
         //    (happy path - can put both points into affine_addition_scratch_space)
         // 2: The next 2 points have different bucket destinations AND point_schedule[point_it].bucket contains a point
         //    (happyish path - we can put points[lhs_schedule] and buckets[lhs_bucket] into
@@ -761,8 +762,8 @@ void MSM<Curve>::consume_point_schedule(std::span<const uint64_t> point_schedule
  */
 template <typename Curve>
 std::vector<typename Curve::AffineElement> MSM<Curve>::batch_multi_scalar_mul(
-    std::vector<std::span<const typename Curve::AffineElement>>& points,
-    std::vector<std::span<ScalarField>>& scalars,
+    std::span<std::span<const typename Curve::AffineElement>> points,
+    std::span<std::span<ScalarField>> scalars,
     bool handle_edge_cases) noexcept
 {
     BB_ASSERT_EQ(points.size(), scalars.size());

barretenberg/cpp/src/barretenberg/ecc/scalar_multiplication/scalar_multiplication.hpp

@@ -94,7 +94,7 @@ template <typename Curve> class MSM {
     static void transform_scalar_and_get_nonzero_scalar_indices(std::span<typename Curve::ScalarField> scalars,
                                                                 std::vector<uint32_t>& consolidated_indices) noexcept;
 
-    static std::vector<ThreadWorkUnits> get_work_units(std::vector<std::span<ScalarField>>& scalars,
+    static std::vector<ThreadWorkUnits> get_work_units(std::span<std::span<ScalarField>> scalars,
                                                        std::vector<std::vector<uint32_t>>& msm_scalar_indices) noexcept;
     static uint32_t get_scalar_slice(const ScalarField& scalar, size_t round, size_t normal_slice_size) noexcept;
     static size_t get_optimal_log_num_buckets(const size_t num_points) noexcept;
@@ -122,8 +122,8 @@ template <typename Curve> class MSM {
                                        size_t num_input_points_processed,
                                        size_t num_queued_affine_points) noexcept;
 
-    static std::vector<AffineElement> batch_multi_scalar_mul(std::vector<std::span<const AffineElement>>& points,
-                                                             std::vector<std::span<ScalarField>>& scalars,
+    static std::vector<AffineElement> batch_multi_scalar_mul(std::span<std::span<const AffineElement>> points,
+                                                             std::span<std::span<ScalarField>> scalars,
                                                              bool handle_edge_cases = true) noexcept;
     static AffineElement msm(std::span<const AffineElement> points,
                              PolynomialSpan<const ScalarField> _scalars,