chore(avm): update stats

barretenberg/cpp/src/barretenberg/bb/main.cpp

@@ -968,7 +968,7 @@ void avm_prove(const std::filesystem::path& bytecode_path,
 
     // Prove execution and return vk
     auto const [verification_key, proof] =
-        avm_trace::Execution::prove(bytecode, calldata, public_inputs_vec, avm_hints);
+        AVM_TRACK_TIME_V("prove/all", avm_trace::Execution::prove(bytecode, calldata, public_inputs_vec, avm_hints));
 
     // TODO(ilyas): <#4887>: Currently we only need these two parts of the vk, look into pcs_verification key reqs
     std::vector<uint64_t> vk_vector = { verification_key.circuit_size, verification_key.num_public_inputs };
@@ -989,7 +989,8 @@ void avm_prove(const std::filesystem::path& bytecode_path,
 #ifdef AVM_TRACK_STATS
     info("------- STATS -------");
     const auto& stats = avm_trace::Stats::get();
-    info(stats.to_string());
+    const int levels = std::getenv("AVM_STATS_LEVELS") != nullptr ? std::stoi(std::getenv("AVM_STATS_LEVELS")) : 2;
+    info(stats.to_string(levels));
 #endif
 }
 
@@ -1013,7 +1014,7 @@ bool avm_verify(const std::filesystem::path& proof_path, const std::filesystem::
     auto num_public_inputs = from_buffer<size_t>(vk_bytes, sizeof(size_t));
     auto vk = AvmFlavor::VerificationKey(circuit_size, num_public_inputs);
 
-    const bool verified = avm_trace::Execution::verify(vk, proof);
+    const bool verified = AVM_TRACK_TIME_V("verify/all", avm_trace::Execution::verify(vk, proof));
     vinfo("verified: ", verified);
     return verified;
 }

barretenberg/cpp/src/barretenberg/vm/avm_trace/avm_execution.cpp

@@ -16,6 +16,7 @@
 #include "barretenberg/vm/generated/avm_circuit_builder.hpp"
 #include "barretenberg/vm/generated/avm_composer.hpp"
 #include "barretenberg/vm/generated/avm_flavor.hpp"
+#include "barretenberg/vm/generated/avm_verifier.hpp"
 
 #include <cassert>
 #include <cstddef>
@@ -132,9 +133,9 @@ std::tuple<AvmFlavor::VerificationKey, HonkProof> Execution::prove(std::vector<u
 
     AVM_TRACK_TIME("prove/check_circuit", circuit_builder.check_circuit());
 
-    auto composer = AvmComposer();
-    auto prover = composer.create_prover(circuit_builder);
-    auto verifier = composer.create_verifier(circuit_builder);
+    auto composer = AVM_TRACK_TIME_V("prove/create_composer", AvmComposer());
+    auto prover = AVM_TRACK_TIME_V("prove/create_prover", composer.create_prover(circuit_builder));
+    auto verifier = AVM_TRACK_TIME_V("prove/create_verifier", composer.create_verifier(circuit_builder));
 
     vinfo("------- PROVING EXECUTION -------");
     // Proof structure: public_inputs | calldata_size | calldata | returndata_size | returndata | raw proof

barretenberg/cpp/src/barretenberg/vm/avm_trace/stats.cpp

@@ -24,22 +24,25 @@ void Stats::increment(const std::string& key, uint64_t value)
     stats[key] += value;
 }
 
-void Stats::time(const std::string& key, std::function<void()> f)
+void Stats::time(const std::string& key, const std::function<void()>& f)
 {
     auto start = std::chrono::system_clock::now();
     f();
     auto elapsed = std::chrono::system_clock::now() - start;
-    increment(key, static_cast<uint64_t>(std::chrono::duration_cast<std::chrono::milliseconds>(elapsed).count()));
+    increment(key + "_ms",
+              static_cast<uint64_t>(std::chrono::duration_cast<std::chrono::milliseconds>(elapsed).count()));
 }
 
-std::string Stats::to_string() const
+std::string Stats::to_string(int levels) const
 {
     std::lock_guard lock(stats_mutex);
 
     std::vector<std::string> result;
     result.reserve(stats.size());
     for (const auto& [key, value] : stats) {
-        result.push_back(key + ": " + std::to_string(value));
+        if (std::count(key.begin(), key.end(), '/') < levels) {
+            result.push_back(key + ": " + std::to_string(value));
+        }
     }
     std::sort(result.begin(), result.end());
     std::string joined;
@@ -49,17 +52,4 @@ std::string Stats::to_string() const
     return joined;
 }
 
-std::string Stats::aggregate_to_string(const std::string& key_prefix) const
-{
-    std::lock_guard lock(stats_mutex);
-
-    uint64_t result = 0;
-    for (const auto& [key, value] : stats) {
-        if (key.starts_with(key_prefix)) {
-            result += value;
-        }
-    }
-    return key_prefix + ": " + std::to_string(result);
-}
-
 } // namespace bb::avm_trace

barretenberg/cpp/src/barretenberg/vm/avm_trace/stats.hpp

@@ -13,9 +13,13 @@
 #endif
 
 #ifdef AVM_TRACK_STATS
+// For tracking time spent in a block of code.
 #define AVM_TRACK_TIME(key, body) ::bb::avm_trace::Stats::get().time(key, [&]() { body; });
+// For tracking time spent in a block of code and returning a value.
+#define AVM_TRACK_TIME_V(key, body) ::bb::avm_trace::Stats::get().template time_r(key, [&]() { return body; });
 #else
 #define AVM_TRACK_TIME(key, body) body
+#define AVM_TRACK_TIME_V(key, type, body) body
 #endif
 
 namespace bb::avm_trace {
@@ -25,9 +29,23 @@ class Stats {
     static Stats& get();
     void reset();
     void increment(const std::string& key, uint64_t value);
-    void time(const std::string& key, std::function<void()> f);
-    std::string to_string() const;
-    std::string aggregate_to_string(const std::string& key_prefix) const;
+    void time(const std::string& key, const std::function<void()>& f);
+
+    template <typename F> auto time_r(const std::string& key, F&& f)
+    {
+        auto start = std::chrono::system_clock::now();
+        auto result = f();
+        auto elapsed = std::chrono::system_clock::now() - start;
+        increment(key + "_ms",
+                  static_cast<uint64_t>(std::chrono::duration_cast<std::chrono::milliseconds>(elapsed).count()));
+        return result;
+    }
+
+    // Returns a string representation of the stats.
+    // E.g., if show_levels = 2, it will show the top 2 levels of the stats.
+    // That is, prove/logderiv_ms will be shown but
+    // prove/logderiv/relation_ms will not be shown.
+    std::string to_string(int show_levels = 2) const;
 
   private:
     Stats() = default;

barretenberg/cpp/src/barretenberg/vm/generated/avm_prover.cpp

@@ -79,7 +79,7 @@ void AvmProver::execute_log_derivative_inverse_round()
     bb::constexpr_for<0, std::tuple_size_v<Flavor::LookupRelations>, 1>([&]<size_t relation_idx>() {
         using Relation = std::tuple_element_t<relation_idx, Flavor::LookupRelations>;
         tasks.push_back([&]() {
-            AVM_TRACK_TIME(Relation::NAME + std::string("_ms"),
+            AVM_TRACK_TIME(std::string("prove/execute_log_derivative_inverse_round/") + Relation::NAME,
                            (compute_logderivative_inverse<Flavor, Relation>(
                                prover_polynomials, relation_parameters, key->circuit_size)));
         });
@@ -150,22 +150,22 @@ HonkProof AvmProver::construct_proof()
     execute_preamble_round();
 
     // Compute wire commitments
-    AVM_TRACK_TIME("prove/execute_wire_commitments_round_ms", execute_wire_commitments_round());
+    AVM_TRACK_TIME("prove/execute_wire_commitments_round", execute_wire_commitments_round());
 
     // Compute sorted list accumulator
-    AVM_TRACK_TIME("prove/execute_log_derivative_inverse_round_ms", execute_log_derivative_inverse_round());
+    AVM_TRACK_TIME("prove/execute_log_derivative_inverse_round", execute_log_derivative_inverse_round());
 
     // Compute commitments to logderivative inverse polynomials
-    AVM_TRACK_TIME("prove/execute_log_derivative_inverse_commitments_round_ms",
+    AVM_TRACK_TIME("prove/execute_log_derivative_inverse_commitments_round",
                    execute_log_derivative_inverse_commitments_round());
 
     // Fiat-Shamir: alpha
     // Run sumcheck subprotocol.
-    AVM_TRACK_TIME("prove/execute_relation_check_rounds_ms", execute_relation_check_rounds());
+    AVM_TRACK_TIME("prove/execute_relation_check_rounds", execute_relation_check_rounds());
 
     // Fiat-Shamir: rho, y, x, z
     // Execute Zeromorph multilinear PCS
-    AVM_TRACK_TIME("prove/execute_pcs_rounds_ms", execute_pcs_rounds());
+    AVM_TRACK_TIME("prove/execute_pcs_rounds", execute_pcs_rounds());
 
     return export_proof();
 }

bb-pilcom/bb-pil-backend/templates/prover.cpp.hbs

@@ -80,7 +80,7 @@ void AvmProver::execute_log_derivative_inverse_round()
     bb::constexpr_for<0, std::tuple_size_v<Flavor::LookupRelations>, 1>([&]<size_t relation_idx>() {
         using Relation = std::tuple_element_t<relation_idx, Flavor::LookupRelations>;
         tasks.push_back([&]() {
-            AVM_TRACK_TIME(Relation::NAME + std::string("_ms"),
+            AVM_TRACK_TIME(std::string("prove/execute_log_derivative_inverse_round/") + Relation::NAME,
                            (compute_logderivative_inverse<Flavor, Relation>(
                                prover_polynomials, relation_parameters, key->circuit_size)));
         });
@@ -151,21 +151,21 @@ HonkProof {{name}}Prover::construct_proof()
     execute_preamble_round();
 
     // Compute wire commitments
-    AVM_TRACK_TIME("prove/execute_wire_commitments_round_ms", execute_wire_commitments_round());
+    AVM_TRACK_TIME("prove/execute_wire_commitments_round", execute_wire_commitments_round());
 
     // Compute sorted list accumulator
-    AVM_TRACK_TIME("prove/execute_log_derivative_inverse_round_ms", execute_log_derivative_inverse_round());
+    AVM_TRACK_TIME("prove/execute_log_derivative_inverse_round", execute_log_derivative_inverse_round());
 
     // Compute commitments to logderivative inverse polynomials
-    AVM_TRACK_TIME("prove/execute_log_derivative_inverse_commitments_round_ms", execute_log_derivative_inverse_commitments_round());
+    AVM_TRACK_TIME("prove/execute_log_derivative_inverse_commitments_round", execute_log_derivative_inverse_commitments_round());
 
     // Fiat-Shamir: alpha
     // Run sumcheck subprotocol.
-    AVM_TRACK_TIME("prove/execute_relation_check_rounds_ms", execute_relation_check_rounds());
+    AVM_TRACK_TIME("prove/execute_relation_check_rounds", execute_relation_check_rounds());
 
     // Fiat-Shamir: rho, y, x, z
     // Execute Zeromorph multilinear PCS
-    AVM_TRACK_TIME("prove/execute_pcs_rounds_ms", execute_pcs_rounds());
+    AVM_TRACK_TIME("prove/execute_pcs_rounds", execute_pcs_rounds());
 
     return export_proof();
 }