fix: bb bench bug (#16972)

Benching could cause a segfault (does not seem related to our random
segfaults)

Author: AztecBot

barretenberg/cpp/src/barretenberg/common/bb_bench.cpp

@@ -206,7 +206,7 @@ AggregateData GlobalBenchStatsContainer::aggregate() const
 
     // Each count has a unique [thread, key] combo.
     // We therefore treat each count as a thread's contribution to that key.
-    for (const TimeStatsEntry* entry : entries) {
+    for (const std::shared_ptr<TimeStatsEntry>& entry : entries) {
         // A map from parent key => AggregateEntry
         auto& entry_map = result[entry->key];
         // combine all entries with same parent key
@@ -238,7 +238,7 @@ GlobalBenchStatsContainer::~GlobalBenchStatsContainer()
     }
 }
 
-void GlobalBenchStatsContainer::add_entry(const char* key, TimeStatsEntry* entry)
+void GlobalBenchStatsContainer::add_entry(const char* key, const std::shared_ptr<TimeStatsEntry>& entry)
 {
     std::unique_lock<std::mutex> lock(mutex);
     entry->key = key;
@@ -248,7 +248,7 @@ void GlobalBenchStatsContainer::add_entry(const char* key, TimeStatsEntry* entry
 void GlobalBenchStatsContainer::print() const
 {
     std::cout << "GlobalBenchStatsContainer::print() START" << "\n";
-    for (const TimeStatsEntry* entry : entries) {
+    for (const std::shared_ptr<TimeStatsEntry>& entry : entries) {
         print_stats_recursive(entry->key, &entry->count, "");
     }
     std::cout << "GlobalBenchStatsContainer::print() END" << "\n";
@@ -560,7 +560,7 @@ void GlobalBenchStatsContainer::print_aggregate_counts_hierarchical(std::ostream
 void GlobalBenchStatsContainer::clear()
 {
     std::unique_lock<std::mutex> lock(mutex);
-    for (TimeStatsEntry* entry : entries) {
+    for (std::shared_ptr<TimeStatsEntry>& entry : entries) {
         entry->count = TimeStats();
     }
 }

barretenberg/cpp/src/barretenberg/common/bb_bench.hpp

@@ -8,6 +8,7 @@
 #include <ostream>
 #include <string_view>
 #include <tracy/Tracy.hpp>
+#include <unordered_map>
 #include <vector>
 
 /**
@@ -65,19 +66,19 @@ struct AggregateEntry {
 // AggregateData: Result of normalizing benchmark data
 // entries: Key -> ParentKey -> Entry
 // Empty string is used as key if the entry has no parent.
-using AggregateData = std::map<OperationKey, std::map<OperationKey, AggregateEntry>>;
+using AggregateData = std::unordered_map<OperationKey, std::map<OperationKey, AggregateEntry>>;
 
 // Contains all statically known op counts
 struct GlobalBenchStatsContainer {
   public:
     static inline thread_local TimeStatsEntry* parent = nullptr;
     ~GlobalBenchStatsContainer();
     std::mutex mutex;
-    std::vector<TimeStatsEntry*> entries;
+    std::vector<std::shared_ptr<TimeStatsEntry>> entries;
     void print() const;
     // NOTE: Should be called when other threads aren't active
     void clear();
-    void add_entry(const char* key, TimeStatsEntry* entry);
+    void add_entry(const char* key, const std::shared_ptr<TimeStatsEntry>& entry);
     void print_stats_recursive(const OperationKey& key, const TimeStats* stats, const std::string& indent) const;
     void print_aggregate_counts(std::ostream&, size_t) const;
     void print_aggregate_counts_hierarchical(std::ostream&) const;
@@ -151,16 +152,17 @@ struct TimeStatsEntry {
 template <OperationLabel Op> struct ThreadBenchStats {
   public:
     // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
-    static inline thread_local TimeStatsEntry stats;
+    static inline thread_local std::shared_ptr<TimeStatsEntry> stats;
 
     static void init_entry(TimeStatsEntry& entry);
     // returns null if use_bb_bench not enabled
-    static TimeStatsEntry* ensure_stats()
+    static std::shared_ptr<TimeStatsEntry> ensure_stats()
     {
-        if (bb::detail::use_bb_bench && BB_UNLIKELY(stats.key.empty())) {
-            GLOBAL_BENCH_STATS.add_entry(Op.value, &stats);
+        if (bb::detail::use_bb_bench && BB_UNLIKELY(stats == nullptr)) {
+            stats = std::make_shared<TimeStatsEntry>();
+            GLOBAL_BENCH_STATS.add_entry(Op.value, stats);
         }
-        return bb::detail::use_bb_bench ? &stats : nullptr;
+        return stats;
     }
 };
 
@@ -198,7 +200,7 @@ struct BenchReporter {
 #define BB_BENCH_TRACY() BB_BENCH_ONLY_NAME(__func__)
 #define BB_BENCH_TRACY_NAME(name) BB_BENCH_ONLY_NAME(name)
 #define BB_BENCH_ONLY_NAME(name)                                                                                       \
-    bb::detail::BenchReporter _bb_bench_reporter((bb::detail::ThreadBenchStats<name>::ensure_stats()))
+    bb::detail::BenchReporter _bb_bench_reporter((bb::detail::ThreadBenchStats<name>::ensure_stats().get()))
 #define BB_BENCH_ENABLE_NESTING()                                                                                      \
     if (_bb_bench_reporter.stats)                                                                                      \
     bb::detail::GlobalBenchStatsContainer::parent = _bb_bench_reporter.stats

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

@@ -28,7 +28,7 @@ class ThreadPool {
 
     void start_tasks(size_t num_iterations, const std::function<void(size_t)>& func)
     {
-        parent = bb::detail::GlobalBenchStatsContainer::parent;
+        parent.store(bb::detail::GlobalBenchStatsContainer::parent);
         {
             std::unique_lock<std::mutex> lock(tasks_mutex);
             task_ = func;
@@ -48,7 +48,7 @@ class ThreadPool {
     }
 
   private:
-    bb::detail::TimeStatsEntry* parent = nullptr;
+    std::atomic<bb::detail::TimeStatsEntry*> parent = nullptr;
     std::vector<std::thread> workers;
     std::mutex tasks_mutex;
     std::function<void(size_t)> task_;
@@ -119,7 +119,7 @@ void ThreadPool::worker_loop(size_t /*unused*/)
         }
         // Make sure nested stats accounting works under multithreading
         // Note: parent is a thread-local variable.
-        bb::detail::GlobalBenchStatsContainer::parent = parent;
+        bb::detail::GlobalBenchStatsContainer::parent = parent.load();
         do_iterations();
     }
     // info("worker exit ", worker_num);

barretenberg/cpp/src/barretenberg/polynomials/gate_separator.hpp

@@ -143,7 +143,7 @@ template <typename FF> struct GateSeparatorPolynomial {
     BB_PROFILE static std::vector<FF> compute_beta_products(const std::vector<FF>& betas,
                                                             const size_t log_num_monomials)
     {
-
+        BB_BENCH_NAME("GateSeparatorPolynomial::compute_beta_products");
         size_t pow_size = 1 << log_num_monomials;
         std::vector<FF> beta_products(pow_size);