feat(avm): parallelize polynomial alloc and set (#8520)

This PR * Parallelizes allocation and initialization of prover polys * Removes duplicated initialization * Removes copy (replaces with move) Improvements in proving time for a 2^20 trace * Creation of prover time goes down 90% * Overall proving time goes down ~40% 2^20 before ``` >>> prove/create_prover_ms: 53011 >>> prove/all_ms: 95457 prove/create_composer_ms: 0 prove/create_verifier_ms: 322 prove/execute_log_derivative_inverse_commitments_round_ms: 2579 prove/execute_log_derivative_inverse_round_ms: 2294 prove/execute_pcs_rounds_ms: 1915 prove/execute_relation_check_rounds_ms: 12066 prove/execute_wire_commitments_round_ms: 1138 prove/gen_trace_ms: 18690 ``` 2^20 after ``` >>> prove/create_prover_ms: 5797 >>> prove/all_ms: 50688 circuit_builder/init_polys_to_be_shifted_ms: 660 circuit_builder/init_polys_unshifted_ms: 301 circuit_builder/set_polys_shifted_ms: 0 circuit_builder/set_polys_unshifted_ms: 4395 composer/create_prover:commitment_key_ms: 152 composer/create_prover:construct_prover_ms: 55 composer/create_prover:proving_key_ms: 177 composer/create_prover:witness_ms: 5410 prove/create_composer_ms: 0 prove/create_verifier_ms: 326 prove/execute_log_derivative_inverse_commitments_round_ms: 2523 prove/execute_log_derivative_inverse_round_ms: 2185 prove/execute_pcs_rounds_ms: 1766 prove/execute_relation_check_rounds_ms: 11966 prove/execute_wire_commitments_round_ms: 1148 prove/gen_trace_ms: 20765 ```
AztecProtocol · Sep 12, 2024 · 7e73531 · 7e73531
1 parent c458a79
commit 7e73531
Show file tree

Hide file tree

Showing 9 changed files with 735 additions and 718 deletions.
diff --git a/barretenberg/cpp/src/barretenberg/vm/avm/generated/circuit_builder.cpp b/barretenberg/cpp/src/barretenberg/vm/avm/generated/circuit_builder.cpp
diff --git a/barretenberg/cpp/src/barretenberg/vm/avm/generated/composer.cpp b/barretenberg/cpp/src/barretenberg/vm/avm/generated/composer.cpp
@@ -1,5 +1,6 @@
 // AUTOGENERATED FILE
 #include "barretenberg/vm/avm/generated/composer.hpp"
+#include "barretenberg/vm/stats.hpp"
 
 namespace bb {
 
@@ -14,20 +15,20 @@ void AvmComposer::compute_witness(CircuitConstructor& circuit)
 
     for (auto [key_poly, prover_poly] : zip_view(proving_key->get_all(), polynomials.get_unshifted())) {
         ASSERT(flavor_get_label(*proving_key, key_poly) == flavor_get_label(polynomials, prover_poly));
-        key_poly = prover_poly;
+        key_poly = std::move(prover_poly);
     }
 
     computed_witness = true;
 }
 
 AvmProver AvmComposer::create_prover(CircuitConstructor& circuit_constructor)
 {
-    compute_proving_key(circuit_constructor);
-    compute_witness(circuit_constructor);
-    compute_commitment_key(circuit_constructor.get_circuit_subgroup_size());
-
-    AvmProver output_state(proving_key, proving_key->commitment_key);
-
+    AVM_TRACK_TIME("composer/create_prover:proving_key", compute_proving_key(circuit_constructor));
+    AVM_TRACK_TIME("composer/create_prover:witness", compute_witness(circuit_constructor));
+    AVM_TRACK_TIME("composer/create_prover:commitment_key",
+                   compute_commitment_key(circuit_constructor.get_circuit_subgroup_size()));
+    auto output_state = AVM_TRACK_TIME_V("composer/create_prover:construct_prover",
+                                         AvmProver(proving_key, proving_key->commitment_key));
     return output_state;
 }
 
@@ -44,10 +45,8 @@ std::shared_ptr<Flavor::ProvingKey> AvmComposer::compute_proving_key(CircuitCons
     }
 
     // Initialize proving_key
-    {
-        const size_t subgroup_size = circuit_constructor.get_circuit_subgroup_size();
-        proving_key = std::make_shared<Flavor::ProvingKey>(subgroup_size, 0);
-    }
+    const size_t subgroup_size = circuit_constructor.get_circuit_subgroup_size();
+    proving_key = std::make_shared<Flavor::ProvingKey>(subgroup_size, 0);
 
     return proving_key;
 }

diff --git a/barretenberg/cpp/src/barretenberg/vm/avm/generated/flavor.cpp b/barretenberg/cpp/src/barretenberg/vm/avm/generated/flavor.cpp
@@ -2273,14 +2273,8 @@ AvmFlavor::ProvingKey::ProvingKey(const size_t circuit_size, const size_t num_pu
     this->log_circuit_size = numeric::get_msb(circuit_size);
     this->num_public_inputs = num_public_inputs;
 
-    // Allocate memory for precomputed polynomials
-    for (auto& poly : PrecomputedEntities<Polynomial>::get_all()) {
-        poly = Polynomial(circuit_size);
-    }
-    // Allocate memory for witness polynomials
-    for (auto& poly : WitnessEntities<Polynomial>::get_all()) {
-        poly = Polynomial(circuit_size);
-    }
+    // The proving key's polynomials are not allocated here because they are later overwritten
+    // AvmComposer::compute_witness(). We should probably refactor this flow.
 };
 
 /**

diff --git a/barretenberg/cpp/src/barretenberg/vm/avm/generated/prover.cpp b/barretenberg/cpp/src/barretenberg/vm/avm/generated/prover.cpp
@@ -27,8 +27,8 @@ using FF = Flavor::FF;
  * @tparam settings Settings class.
  */
 AvmProver::AvmProver(std::shared_ptr<Flavor::ProvingKey> input_key, std::shared_ptr<PCSCommitmentKey> commitment_key)
-    : key(input_key)
-    , commitment_key(commitment_key)
+    : key(std::move(input_key))
+    , commitment_key(std::move(commitment_key))
 {
     for (auto [prover_poly, key_poly] : zip_view(prover_polynomials.get_unshifted(), key->get_all())) {
         ASSERT(bb::flavor_get_label(prover_polynomials, prover_poly) == bb::flavor_get_label(*key, key_poly));

diff --git a/bb-pilcom/bb-pil-backend/src/lookup_builder.rs b/bb-pilcom/bb-pil-backend/src/lookup_builder.rs
@@ -109,7 +109,6 @@ pub fn get_counts_from_lookups(lookups: &[Lookup]) -> Vec<String> {
 
 fn create_lookup_settings_data(lookup: &Lookup) -> Json {
     let columns_per_set = lookup.left.cols.len();
-    let counts_poly_name = lookup.counts_poly.to_owned();
 
     // NOTE: https://github.com/AztecProtocol/aztec-packages/issues/3879
     // Settings are not flexible enough to combine inverses

diff --git a/bb-pilcom/bb-pil-backend/templates/circuit_builder.cpp.hbs b/bb-pilcom/bb-pil-backend/templates/circuit_builder.cpp.hbs
@@ -10,6 +10,7 @@
 #include "barretenberg/relations/generic_permutation/generic_permutation_relation.hpp"
 #include "barretenberg/relations/generic_lookup/generic_lookup_relation.hpp"
 #include "barretenberg/honk/proof_system/logderivative_library.hpp"
+#include "barretenberg/vm/stats.hpp"
 
 namespace bb {
 
@@ -18,26 +19,40 @@ namespace bb {
     ProverPolynomials polys;
 
     // Allocate mem for each column
-
-    for (auto& poly : polys.get_to_be_shifted()) {
-        poly = Polynomial::shiftable(num_rows);
-    }
+    AVM_TRACK_TIME("circuit_builder/init_polys_to_be_shifted", ({
+                       for (auto& poly : polys.get_to_be_shifted()) {
+                           poly = Polynomial{ /*memory size*/ num_rows - 1,
+                                              /*largest possible index*/ num_rows,
+                                              /*make shiftable with offset*/ 1 };
+                       }
+                   }));
     // catch-all with fully formed polynomials
-    for (auto& poly : polys.get_unshifted()) {
-        if (poly.is_empty()) {
-            poly = Polynomial{ num_rows };
-        }
-    }
-
-    for (size_t i = 0; i < rows.size(); i++) {
+    AVM_TRACK_TIME("circuit_builder/init_polys_unshifted", ({
+                       auto unshifted = polys.get_unshifted();
+                       bb::parallel_for(unshifted.size(), [&](size_t i) {
+                           auto& poly = unshifted[i];
+                           if (poly.is_empty()) {
+                               // Not set above
+                               poly = Polynomial{ /*memory size*/ num_rows, /*largest possible index*/ num_rows };
+                           }
+                       });
+                   }));
+
+    AVM_TRACK_TIME(
+        "circuit_builder/set_polys_unshifted", ({
+            bb::parallel_for(rows.size(), [&](size_t i) {
         {{#each all_cols_without_inverses as |poly|}}
         polys.{{poly}}.set_if_valid_index(i, rows[i].{{poly}});
         {{/each}}
-    }
+        });
+    }));
 
-    for (auto [shifted, to_be_shifted] : zip_view(polys.get_shifted(), polys.get_to_be_shifted())) {
-        shifted = to_be_shifted.shifted();
-    }
+    AVM_TRACK_TIME(
+        "circuit_builder/set_polys_shifted", ({
+        for (auto [shifted, to_be_shifted] : zip_view(polys.get_shifted(), polys.get_to_be_shifted())) {
+            shifted = to_be_shifted.shifted();
+        }
+    }));
 
     return polys;
 }

diff --git a/bb-pilcom/bb-pil-backend/templates/composer.cpp.hbs b/bb-pilcom/bb-pil-backend/templates/composer.cpp.hbs
@@ -1,5 +1,6 @@
 // AUTOGENERATED FILE
 #include "barretenberg/vm/{{snakeCase name}}/generated/composer.hpp"
+#include "barretenberg/vm/stats.hpp"
 
 namespace bb {
 
@@ -14,20 +15,20 @@ void {{name}}Composer::compute_witness(CircuitConstructor& circuit)
 
     for (auto [key_poly, prover_poly] : zip_view(proving_key->get_all(), polynomials.get_unshifted())) {
         ASSERT(flavor_get_label(*proving_key, key_poly) == flavor_get_label(polynomials, prover_poly));
-        key_poly = prover_poly;
+        key_poly = std::move(prover_poly);
     }
 
     computed_witness = true;
 }
 
 {{name}}Prover {{name}}Composer::create_prover(CircuitConstructor& circuit_constructor)
 {
-    compute_proving_key(circuit_constructor);
-    compute_witness(circuit_constructor);
-    compute_commitment_key(circuit_constructor.get_circuit_subgroup_size());
-
-    {{name}}Prover output_state(proving_key, proving_key->commitment_key);
-
+    AVM_TRACK_TIME("composer/create_prover:proving_key", compute_proving_key(circuit_constructor));
+    AVM_TRACK_TIME("composer/create_prover:witness", compute_witness(circuit_constructor));
+    AVM_TRACK_TIME("composer/create_prover:commitment_key",
+                   compute_commitment_key(circuit_constructor.get_circuit_subgroup_size()));
+    auto output_state = AVM_TRACK_TIME_V("composer/create_prover:construct_prover",
+                                         AvmProver(proving_key, proving_key->commitment_key));
     return output_state;
 }
 

diff --git a/bb-pilcom/bb-pil-backend/templates/flavor.cpp.hbs b/bb-pilcom/bb-pil-backend/templates/flavor.cpp.hbs
@@ -114,14 +114,8 @@ AvmFlavor::ProvingKey::ProvingKey(const size_t circuit_size, const size_t num_pu
     this->log_circuit_size = numeric::get_msb(circuit_size);
     this->num_public_inputs = num_public_inputs;
 
-    // Allocate memory for precomputed polynomials
-    for (auto& poly : PrecomputedEntities<Polynomial>::get_all()) {
-        poly = Polynomial(circuit_size);
-    }
-    // Allocate memory for witness polynomials
-    for (auto& poly : WitnessEntities<Polynomial>::get_all()) {
-        poly = Polynomial(circuit_size);
-    }
+    // The proving key's polynomials are not allocated here because they are later overwritten
+    // AvmComposer::compute_witness(). We should probably refactor this flow.
 };
 
 /**

diff --git a/bb-pilcom/bb-pil-backend/templates/prover.cpp.hbs b/bb-pilcom/bb-pil-backend/templates/prover.cpp.hbs
@@ -28,8 +28,8 @@ using FF = Flavor::FF;
  */
 {{name}}Prover::{{name}}Prover(std::shared_ptr<Flavor::ProvingKey> input_key,
                                     std::shared_ptr<PCSCommitmentKey> commitment_key)
-    : key(input_key)
-    , commitment_key(commitment_key)
+    : key(std::move(input_key))
+    , commitment_key(std::move(commitment_key))
 {
     for (auto [prover_poly, key_poly] : zip_view(prover_polynomials.get_unshifted(), key->get_all())) {
         ASSERT(bb::flavor_get_label(prover_polynomials, prover_poly) ==