chore(avm): basic stat collection (#7283)

Colelcts and prints stats for `avm_prove` when BB is compiled with in `RelWithAssert` mode or `AVM_TRACK_STATS` is defined.

Has no impact otherwise.

```
------- STATS -------
compute_logderivative_inverse/incl_main_tag_err_ms: 0
compute_logderivative_inverse/incl_mem_tag_err_ms: 0
compute_logderivative_inverse/kernel_output_lookup_ms: 0
compute_logderivative_inverse/lookup_byte_lengths_ms: 0
compute_logderivative_inverse/lookup_byte_operations_ms: 0
compute_logderivative_inverse/lookup_div_u16_0_ms: 0
compute_logderivative_inverse/lookup_div_u16_1_ms: 0
compute_logderivative_inverse/lookup_div_u16_2_ms: 0
compute_logderivative_inverse/lookup_div_u16_3_ms: 0
compute_logderivative_inverse/lookup_div_u16_4_ms: 0
compute_logderivative_inverse/lookup_div_u16_5_ms: 0
compute_logderivative_inverse/lookup_div_u16_6_ms: 0
compute_logderivative_inverse/lookup_div_u16_7_ms: 0
compute_logderivative_inverse/lookup_into_kernel_ms: 0
compute_logderivative_inverse/lookup_mem_rng_chk_hi_ms: 0
compute_logderivative_inverse/lookup_mem_rng_chk_lo_ms: 0
compute_logderivative_inverse/lookup_mem_rng_chk_mid_ms: 0
compute_logderivative_inverse/lookup_opcode_gas_ms: 0
compute_logderivative_inverse/lookup_pow_2_0_ms: 0
compute_logderivative_inverse/lookup_pow_2_1_ms: 0
compute_logderivative_inverse/lookup_u16_0_ms: 0
compute_logderivative_inverse/lookup_u16_10_ms: 0
compute_logderivative_inverse/lookup_u16_11_ms: 0
compute_logderivative_inverse/lookup_u16_12_ms: 0
compute_logderivative_inverse/lookup_u16_13_ms: 0
compute_logderivative_inverse/lookup_u16_14_ms: 0
compute_logderivative_inverse/lookup_u16_1_ms: 0
compute_logderivative_inverse/lookup_u16_2_ms: 0
compute_logderivative_inverse/lookup_u16_3_ms: 0
compute_logderivative_inverse/lookup_u16_4_ms: 0
compute_logderivative_inverse/lookup_u16_5_ms: 0
compute_logderivative_inverse/lookup_u16_6_ms: 0
compute_logderivative_inverse/lookup_u16_7_ms: 0
compute_logderivative_inverse/lookup_u16_8_ms: 0
compute_logderivative_inverse/lookup_u16_9_ms: 0
compute_logderivative_inverse/lookup_u8_0_ms: 0
compute_logderivative_inverse/lookup_u8_1_ms: 0
compute_logderivative_inverse/perm_main_alu_ms: 1
compute_logderivative_inverse/perm_main_bin_ms: 0
compute_logderivative_inverse/perm_main_conv_ms: 0
compute_logderivative_inverse/perm_main_mem_a_ms: 0
compute_logderivative_inverse/perm_main_mem_b_ms: 0
compute_logderivative_inverse/perm_main_mem_c_ms: 0
compute_logderivative_inverse/perm_main_mem_d_ms: 0
compute_logderivative_inverse/perm_main_mem_ind_addr_a_ms: 0
compute_logderivative_inverse/perm_main_mem_ind_addr_b_ms: 0
compute_logderivative_inverse/perm_main_mem_ind_addr_c_ms: 0
compute_logderivative_inverse/perm_main_mem_ind_addr_d_ms: 0
compute_logderivative_inverse/perm_main_pedersen_ms: 0
compute_logderivative_inverse/perm_main_pos2_perm_ms: 0
compute_logderivative_inverse/range_check_da_gas_hi_ms: 0
compute_logderivative_inverse/range_check_da_gas_lo_ms: 0
compute_logderivative_inverse/range_check_l2_gas_hi_ms: 0
compute_logderivative_inverse/range_check_l2_gas_lo_ms: 0
proving/log_derivative_inverse_round_ms: 82
proving/pcs_rounds_ms: 158
proving/relation_check_rounds_ms: 66
proving/wire_commitments_round_ms: 193
```

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

@@ -712,6 +712,12 @@ void avm_prove(const std::filesystem::path& bytecode_path,
     vinfo("vk written to: ", vk_path);
     write_file(vk_fields_path, { vk_json.begin(), vk_json.end() });
     vinfo("vk as fields written to: ", vk_fields_path);
+
+#ifdef AVM_TRACK_STATS
+    info("------- STATS -------");
+    const auto& stats = avm_trace::Stats::get();
+    info(stats.to_string());
+#endif
 }
 
 /**

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

@@ -3,6 +3,7 @@
 #include "barretenberg/common/throw_or_abort.hpp"
 #include "barretenberg/serialize/msgpack.hpp"
 #include "barretenberg/vm/avm_trace/constants.hpp"
+#include "barretenberg/vm/avm_trace/stats.hpp"
 #include "barretenberg/vm/generated/avm_flavor.hpp"
 
 #include <array>

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

@@ -0,0 +1,60 @@
+#include "barretenberg/vm/avm_trace/stats.hpp"
+
+#include <chrono>
+#include <cstdint>
+#include <string>
+#include <vector>
+
+namespace bb::avm_trace {
+
+Stats& Stats::get()
+{
+    static Stats stats;
+    return stats;
+}
+
+void Stats::reset()
+{
+    stats.clear();
+}
+
+void Stats::increment(const std::string& key, uint64_t value)
+{
+    stats[key] += value;
+}
+
+void Stats::time(const std::string& key, 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()));
+}
+
+std::string Stats::to_string() const
+{
+    std::vector<std::string> result;
+    result.reserve(stats.size());
+    for (const auto& [key, value] : stats) {
+        result.push_back(key + ": " + std::to_string(value));
+    }
+    std::sort(result.begin(), result.end());
+    std::string joined;
+    for (auto& s : result) {
+        joined += std::move(s) + "\n";
+    }
+    return joined;
+}
+
+std::string Stats::aggregate_to_string(const std::string& key_prefix) const
+{
+    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

@@ -0,0 +1,37 @@
+#pragma once
+
+#include <cstdint>
+#include <functional>
+#include <string>
+#include <unordered_map>
+
+// To enable stats tracking, compile in RelWithAssert mode.
+// cmake --preset $PRESET -DCMAKE_BUILD_TYPE=RelWithAssert
+#ifndef NDEBUG
+#define AVM_TRACK_STATS
+#endif
+
+#ifdef AVM_TRACK_STATS
+#define AVM_TRACK_TIME(key, body) ::bb::avm_trace::Stats::get().time(key, [&]() { body; });
+#else
+#define AVM_TRACK_TIME(key, body) body
+#endif
+
+namespace bb::avm_trace {
+
+class Stats {
+  public:
+    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;
+
+  private:
+    Stats() = default;
+
+    std::unordered_map<std::string, uint64_t> stats;
+};
+
+} // namespace bb::avm_trace

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

@@ -9,6 +9,7 @@
 #include "barretenberg/polynomials/polynomial.hpp"
 #include "barretenberg/relations/permutation_relation.hpp"
 #include "barretenberg/sumcheck/sumcheck.hpp"
+#include "barretenberg/vm/avm_trace/stats.hpp"
 
 namespace bb {
 
@@ -240,18 +241,18 @@ HonkProof AvmProver::construct_proof()
     execute_preamble_round();
 
     // Compute wire commitments
-    execute_wire_commitments_round();
+    AVM_TRACK_TIME("proving/wire_commitments_round_ms", execute_wire_commitments_round());
 
     // Compute sorted list accumulator and commitment
-    execute_log_derivative_inverse_round();
+    AVM_TRACK_TIME("proving/log_derivative_inverse_round_ms", execute_log_derivative_inverse_round());
 
     // Fiat-Shamir: alpha
     // Run sumcheck subprotocol.
-    execute_relation_check_rounds();
+    AVM_TRACK_TIME("proving/relation_check_rounds_ms", execute_relation_check_rounds());
 
     // Fiat-Shamir: rho, y, x, z
     // Execute Zeromorph multilinear PCS
-    execute_pcs_rounds();
+    AVM_TRACK_TIME("proving/pcs_rounds_ms", execute_pcs_rounds());
 
     return export_proof();
 }

bb-pilcom/bb-pil-backend/src/flavor_builder.rs

@@ -128,6 +128,7 @@ fn flavor_includes(name: &str, relation_file_names: &[String], lookups: &[String
 
 #include \"barretenberg/relations/generic_permutation/generic_permutation_relation.hpp\"
 
+#include \"barretenberg/vm/avm_trace/stats.hpp\"
 #include \"barretenberg/flavor/flavor_macros.hpp\"
 #include \"barretenberg/transcript/transcript.hpp\"
 #include \"barretenberg/polynomials/evaluation_domain.hpp\"
@@ -568,7 +569,7 @@ fn create_compute_logderivative_inverses(flavor_name: &str, lookups: &[String])
     }
 
     let compute_inverse_transformation = |lookup_name: &String| {
-        format!("bb::compute_logderivative_inverse<{flavor_name}Flavor, {lookup_name}_relation<FF>>(prover_polynomials, relation_parameters, this->circuit_size);")
+        format!("AVM_TRACK_TIME(\"compute_logderivative_inverse/{lookup_name}_ms\", (bb::compute_logderivative_inverse<{flavor_name}Flavor, {lookup_name}_relation<FF>>(prover_polynomials, relation_parameters, this->circuit_size)));")
     };
 
     let compute_inverses = map_with_newline(lookups, compute_inverse_transformation);

bb-pilcom/bb-pil-backend/src/prover_builder.rs

@@ -82,16 +82,7 @@ impl ProverBuilder for BBFiles {
         let include_str = includes_cpp(&snake_case(name));
 
         let polynomial_commitment_phase = create_commitments_phase();
-
-        let (call_log_derivative_phase, log_derivative_inverse_phase): (String, String) =
-            if lookup_names.is_empty() {
-                ("".to_owned(), "".to_owned())
-            } else {
-                (
-                    "execute_log_derivative_inverse_round();".to_owned(),
-                    create_log_derivative_inverse_round(lookup_names),
-                )
-            };
+        let log_derivative_inverse_phase = create_log_derivative_inverse_round(lookup_names);
 
         let prover_cpp = format!("
     {include_str}
@@ -206,18 +197,18 @@ impl ProverBuilder for BBFiles {
         execute_preamble_round();
     
         // Compute wire commitments
-        execute_wire_commitments_round();
+        AVM_TRACK_TIME(\"proving/wire_commitments_round_ms\", execute_wire_commitments_round());
     
         // Compute sorted list accumulator and commitment
-        {call_log_derivative_phase}
+        AVM_TRACK_TIME(\"proving/log_derivative_inverse_round_ms\", execute_log_derivative_inverse_round());
     
         // Fiat-Shamir: alpha
         // Run sumcheck subprotocol.
-        execute_relation_check_rounds();
+        AVM_TRACK_TIME(\"proving/relation_check_rounds_ms\", execute_relation_check_rounds());
     
         // Fiat-Shamir: rho, y, x, z
         // Execute Zeromorph multilinear PCS
-        execute_pcs_rounds();
+        AVM_TRACK_TIME(\"proving/pcs_rounds_ms\", execute_pcs_rounds());
     
         return export_proof();
     }}
@@ -264,6 +255,7 @@ fn includes_cpp(name: &str) -> String {
     #include \"barretenberg/polynomials/polynomial.hpp\"
     #include \"barretenberg/relations/permutation_relation.hpp\"
     #include \"barretenberg/sumcheck/sumcheck.hpp\"
+    #include \"barretenberg/vm/avm_trace/stats.hpp\"
     "
     )
 }
@@ -296,6 +288,10 @@ fn create_commitments_phase() -> String {
 }
 
 fn create_log_derivative_inverse_round(lookup_operations: &[String]) -> String {
+    if lookup_operations.is_empty() {
+        return "".to_owned();
+    }
+
     let all_commit_operations = map_with_newline(lookup_operations, commitment_transform);
     let send_to_verifier_operations =
         map_with_newline(lookup_operations, send_to_verifier_transform);