feat(avm): lock-free sharded TraceContainer (#24300)

## Summary

Port of AztecProtocol/aztec-packages-private#312 (by @fcarreiro) to
public `next`.

Replaces the AVM `TraceContainer`'s per-column `std::shared_mutex` +
`unordered_flat_map` storage with a lock-free, sharded design. Each
column is a fixed-size table of atomic pointers to dense, fixed-size row
shards; the hot `get`/`set` path is a single atomic load with no lock,
and shards are created at most once via compare-and-swap. The public API
is unchanged.

Motivation: under parallel tracegen (`parallel_for` over jobs), libc++'s
`std::shared_mutex::lock_shared()` is backed by a real mutex, so the
per-column lock serialized concurrent writers. Profiling showed lock
machinery dominating tracegen and poor scaling at 16 threads.

## Changes

- Column storage is now a fixed-size
`std::array<std::atomic<ColumnInterval*>, NUM_SHARDS>`, indexed by `row
/ INTERVAL_SIZE`, held inline and freed in the destructor.
- Each shard stores rows in a dense, value-initialized `std::array<FF,
INTERVAL_SIZE>` (zero == absent). Distinct rows are distinct elements,
so concurrent writers to different rows of a column never race or
serialize — including when two chunks meet at a shard boundary.
- Shard creation is lock-free via `compare_exchange_strong`. The
per-column lock is gone; the bulk `reserve_column` / `clear_column` /
`invert_column` paths are documented as not thread-safe (they run after
the parallel fill phase joins).
- In-shard offset uses `row % INTERVAL_SIZE` (correct for any
`INTERVAL_SIZE`; strength-reduces to the same instruction as a bitmask
for the power-of-two value).
- Documented the thread-safety contract: concurrent `get`/`set` is safe
as long as no two threads touch the same `(column, row)` with at least
one writer, and reads are barrier-separated from writes.

With `INTERVAL_SIZE = 2^11` and `MAX_AVM_TRACE_SIZE = 2^21`, `NUM_SHARDS
= 1024` and `INTERVAL_SIZE` divides the max trace size evenly.

## Performance & memory (from the source PR)

Measured with `HARDWARE_CONCURRENCY=16`, RelWithAssert `bb-avm`.

**Bulk test** (`avm_bulk.test.ts`), source branch vs `next`:

| | `next` (shared_mutex + hash-map) | this change (lock-free dense) |
delta |
|---|---:|---:|---:|
| Tracegen (all) | ~3.0 s | ~0.78 s | **≈ 4× faster** |
| Peak `bb-avm` process memory | ~4.1 GiB | ~3.2–3.4 GiB | **~0.8–0.9
GiB lower (~20%)** |

**Opcode spammer** (`avm_opcode_spam.test.ts`, proving skipped, 70
opcode cases), aggregate tracegen: ~878.6 s → ~230.2 s (**≈ 3.8×
faster**).

The lower memory comes from the trace being top-dense: a dense array
stores one `FF` per cell, versus the hash map's per-entry key + bucket +
load-factor overhead.

## Port notes

- The two changed files
(`barretenberg/cpp/src/barretenberg/vm2/tracegen/trace_container.{cpp,hpp}`)
were byte-identical between public `next` and the source PR's base, so
the change applies as-is (+186/−54, identical to the source PR).
- Verified locally: the changed translation unit compiles against `next`
(`cmake --preset default -DAVM=ON`, built
`vm2_objects/.../trace_container.cpp.o`).

Source: AztecProtocol/aztec-packages-private#312 · Part of Linear
AVM-200.

---
*Created by
[claudebox](https://claudebox.work/v2/sessions/7ad6d36a577e36fd) ·
group: `slackbot`*

Author: fcarreiro

barretenberg/cpp/src/barretenberg/vm2/tracegen/trace_container.cpp

@@ -1,8 +1,9 @@
 #include "barretenberg/vm2/tracegen/trace_container.hpp"
 
 #include <algorithm>
-#include <mutex>
+#include <ranges>
 
+#include "barretenberg/common/assert.hpp"
 #include "barretenberg/common/log.hpp"
 #include "barretenberg/common/ref_vector.hpp"
 #include "barretenberg/vm2/common/field.hpp"
@@ -23,9 +24,15 @@ TraceContainer::TraceContainer()
 const FF& TraceContainer::get(Column col, uint32_t row) const
 {
     auto& column_data = (*trace)[static_cast<size_t>(col)];
-    std::shared_lock lock(column_data.mutex);
-    const auto it = column_data.rows.find(row);
-    return it == column_data.rows.end() ? zero : it->second;
+    const size_t shard_idx = row / INTERVAL_SIZE;
+    if (shard_idx >= NUM_SHARDS) {
+        return zero;
+    }
+    ColumnInterval* shard = column_data.slots[shard_idx].load(std::memory_order_acquire);
+    if (shard == nullptr) {
+        return zero;
+    }
+    return shard->rows[row % INTERVAL_SIZE];
 }
 
 const FF& TraceContainer::get_column_or_shift(ColumnAndShifts col, uint32_t row) const
@@ -36,19 +43,42 @@ const FF& TraceContainer::get_column_or_shift(ColumnAndShifts col, uint32_t row)
     return get(static_cast<Column>(col), row);
 }
 
+TraceContainer::ColumnInterval& TraceContainer::get_or_create_shard(SparseColumn& column_data, size_t shard_idx)
+{
+    ColumnInterval* shard = column_data.slots[shard_idx].load(std::memory_order_acquire);
+    if (shard != nullptr) {
+        return *shard;
+    }
+    // Slow path: this slot has never been written. Construct a shard and install it with a single CAS.
+    // Creators of different shards target different atomics and run fully in parallel with no lock. If we
+    // lose the race for this slot (only possible when two chunks share a shard at a boundary), we discard
+    // our spare copy and use the winner's.
+    auto fresh = std::make_unique<ColumnInterval>();
+    ColumnInterval* expected = nullptr;
+    if (column_data.slots[shard_idx].compare_exchange_strong(
+            expected, fresh.get(), std::memory_order_acq_rel, std::memory_order_acquire)) {
+        return *fresh.release();
+    }
+    return *expected; // CAS failure loaded the winning pointer into `expected` (acquire).
+}
+
 void TraceContainer::set(Column col, uint32_t row, const FF& value)
 {
     auto& column_data = (*trace)[static_cast<size_t>(col)];
-    std::unique_lock lock(column_data.mutex);
+    const size_t shard_idx = row / INTERVAL_SIZE;
+    BB_ASSERT_LT(shard_idx, NUM_SHARDS, "row exceeds the maximum trace size");
+    const uint32_t offset = row % INTERVAL_SIZE;
+
     if (!value.is_zero()) {
-        column_data.rows.insert_or_assign(row, value);
-        column_data.max_row_number = std::max(column_data.max_row_number, static_cast<int64_t>(row));
+        // Lock-free: a single atomic load finds the shard (created on first write), then we write our
+        // own dense cell directly. Different rows are distinct array elements, so concurrent writers of
+        // this column (or even of the same shard, at a chunk boundary) never race and never serialize.
+        get_or_create_shard(column_data, shard_idx).rows[offset] = value;
     } else {
-        auto num_erased = column_data.rows.erase(row);
-        if (column_data.max_row_number == row && num_erased > 0) {
-            // This shouldn't happen often. We delay recalculation of the max row number
-            // until someone actually needs it.
-            column_data.row_number_dirty = true;
+        // Zero value: clear if present. We never create a shard (clearing an absent row is a no-op).
+        ColumnInterval* shard = column_data.slots[shard_idx].load(std::memory_order_acquire);
+        if (shard != nullptr) {
+            shard->rows[offset] = FF::zero();
         }
     }
 }
@@ -62,24 +92,39 @@ void TraceContainer::set(uint32_t row, std::span<const std::pair<Column, FF>> va
 
 void TraceContainer::reserve_column(Column col, size_t size)
 {
+    if (size == 0) {
+        return;
+    }
     auto& column_data = (*trace)[static_cast<size_t>(col)];
-    std::unique_lock lock(column_data.mutex);
-    column_data.rows.reserve(size);
+    const size_t num_shards = std::min((size + INTERVAL_SIZE - 1) / INTERVAL_SIZE, NUM_SHARDS);
+    // Each shard's dense row array is full size on creation, so reserving just materializes the shards up
+    // front (e.g. for precomputed columns). Lock-free: get_or_create_shard installs each via CAS.
+    for (size_t k = 0; k < num_shards; ++k) {
+        get_or_create_shard(column_data, k);
+    }
 }
 
 uint32_t TraceContainer::get_column_rows(Column col) const
 {
+    // The number of rows is (highest non-zero absolute row + 1). We find it by scanning shards from the
+    // top down and, within the first non-empty shard, scanning its rows from the top. Shards are
+    // top-dense, so this terminates almost immediately. This is only called after the parallel fill
+    // phase, so no lock is needed. Lower shards cannot hold a higher row, so the first hit is the answer.
     auto& column_data = (*trace)[static_cast<size_t>(col)];
-    std::unique_lock lock(column_data.mutex);
-    if (column_data.row_number_dirty) {
-        // Trigger recalculation of max row number.
-        auto keys = std::views::keys(column_data.rows);
-        const auto it = std::ranges::max_element(keys);
-        // We use -1 to indicate that the column is empty.
-        column_data.max_row_number = it == keys.end() ? -1 : static_cast<int64_t>(*it);
-        column_data.row_number_dirty = false;
-    }
-    return static_cast<uint32_t>(column_data.max_row_number + 1);
+    for (size_t k = NUM_SHARDS; k-- > 0;) {
+        ColumnInterval* shard_ptr = column_data.slots[k].load(std::memory_order_acquire);
+        if (shard_ptr == nullptr) {
+            continue;
+        }
+        const auto& rows = shard_ptr->rows;
+        const uint32_t base = static_cast<uint32_t>(k) * INTERVAL_SIZE;
+        for (uint32_t off = INTERVAL_SIZE; off-- > 0;) {
+            if (!rows[off].is_zero()) {
+                return base + off + 1;
+            }
+        }
+    }
+    return 0;
 }
 
 uint32_t TraceContainer::get_num_witness_rows() const
@@ -103,9 +148,18 @@ uint32_t TraceContainer::get_num_rows() const
 void TraceContainer::visit_column(Column col, const std::function<void(uint32_t, const FF&)>& visitor) const
 {
     auto& column_data = (*trace)[static_cast<size_t>(col)];
-    std::shared_lock lock(column_data.mutex);
-    for (const auto& [row, value] : column_data.rows) {
-        visitor(row, value);
+    for (size_t k = 0; k < NUM_SHARDS; ++k) {
+        ColumnInterval* shard_ptr = column_data.slots[k].load(std::memory_order_acquire);
+        if (shard_ptr == nullptr) {
+            continue;
+        }
+        auto& shard = *shard_ptr;
+        const uint32_t base = static_cast<uint32_t>(k) * INTERVAL_SIZE;
+        for (uint32_t off = 0; off < INTERVAL_SIZE; ++off) {
+            if (!shard.rows[off].is_zero()) {
+                visitor(base + off, shard.rows[off]);
+            }
+        }
     }
 }
 
@@ -120,20 +174,28 @@ void TraceContainer::invert_column(Column col)
 {
     RefVector<FF> ff_vector;
     auto& column_data = (*trace)[static_cast<size_t>(col)];
-    std::unique_lock lock(column_data.mutex);
-    for (auto& [row, value] : column_data.rows) {
-        ff_vector.push_back(value);
+    for (size_t k = 0; k < NUM_SHARDS; ++k) {
+        ColumnInterval* shard_ptr = column_data.slots[k].load(std::memory_order_acquire);
+        if (shard_ptr == nullptr) {
+            continue;
+        }
+        auto& shard = *shard_ptr;
+        for (auto& value : shard.rows) {
+            if (!value.is_zero()) {
+                ff_vector.push_back(value);
+            }
+        }
     }
     FF::batch_invert<RefVector<FF>>(ff_vector);
 }
 
 void TraceContainer::clear_column(Column col)
 {
+    // Lock-free: exchange hands each non-null shard pointer to exactly one caller, which frees it.
     auto& column_data = (*trace)[static_cast<size_t>(col)];
-    std::unique_lock lock(column_data.mutex);
-    column_data.rows.clear();
-    column_data.max_row_number = -1;
-    column_data.row_number_dirty = false;
+    for (size_t k = 0; k < NUM_SHARDS; ++k) {
+        delete column_data.slots[k].exchange(nullptr, std::memory_order_acq_rel);
+    }
 }
 
 } // namespace bb::avm2::tracegen

barretenberg/cpp/src/barretenberg/vm2/tracegen/trace_container.hpp

@@ -2,15 +2,17 @@
 
 #include <algorithm>
 #include <array>
+#include <atomic>
 #include <cstddef>
+#include <cstdint>
 #include <functional>
 #include <memory>
-#include <shared_mutex>
 #include <span>
-#include <unordered_map>
 #include <utility>
+#include <vector>
 
 #include "barretenberg/common/tuple.hpp"
+#include "barretenberg/vm2/common/constants.hpp"
 #include "barretenberg/vm2/common/field.hpp"
 #include "barretenberg/vm2/common/map.hpp"
 #include "barretenberg/vm2/constraining/flavor_settings.hpp"
@@ -19,10 +21,48 @@
 
 namespace bb::avm2::tracegen {
 
-// This container is thread-safe.
-// Contention can only happen when concurrently accessing the same column.
+// Thread-safety contract.
+//
+// This container supports concurrent get/set during the parallel tracegen fill phase, *as long as the
+// caller never lets two threads touch the same cell (column, row) with at least one of them writing*.
+// Under that contract the following are all race-free and never serialize against each other:
+//   - writes to different columns (independent storage);
+//   - writes to different rows of the same column, even when those rows fall in the same shard (e.g.
+//     two chunks that meet at a shard boundary) — distinct rows are distinct objects;
+//   - concurrent first writes that race to create the same shard — the shard is installed with a single
+//     compare-and-swap, so it is created at most once (the loser frees its spare copy).
+//
+// What is NOT protected, and what the caller must therefore guarantee never happens concurrently:
+//   - two writes to the same (column, row), or a read and a write of the same (column, row). There is
+//     no per-cell synchronization, so this is a data race (undefined behavior). Tracegen satisfies the
+//     contract because it never writes a given cell from two threads, and every read of a column
+//     happens in a later, barrier-separated phase.
+//   - get/set concurrent with reserve_column, clear_column, invert_column, or the destructor on the
+//     same column. Those operations mutate or free the shard table and assume the parallel fill phase
+//     for that column has already joined.
+//
+// Design. Each column is partitioned into fixed-size row intervals ("shards") of INTERVAL_SIZE rows.
+// The per-column shard table is a fixed-size array of atomic pointers indexed by shard (row /
+// INTERVAL_SIZE), so the hot get/set path finds its shard with a single lock-free atomic load; that
+// path takes no per-column or per-shard lock, so concurrent writers of a column never serialize. A
+// shard is created at most once with a lock-free compare-and-swap: a writer that finds a null slot
+// constructs a shard and atomically installs it only if the slot is still null, so creators of
+// different shards never serialize. Each shard holds its rows in a dense, fixed-size array, so once a
+// shard exists, a write is a plain store to a fixed address.
 class TraceContainer {
   public:
+    // Number of rows owned by a single shard, and the unit of lazy allocation. Writes to different rows
+    // never serialize regardless of shard, so this does not affect write parallelism; it only sets how
+    // finely columns are allocated and the (rare) granularity at which two chunks can race to create the
+    // same shard. Chunked tracegen sizes its chunks as a multiple of this value (and ideally aligns to
+    // it) so that concurrent chunks touch disjoint shards. Smaller => less memory wasted in
+    // sparsely-filled regions, at the cost of a larger shard table (more atomic slots) and more shard
+    // allocations per column.
+    static constexpr uint32_t INTERVAL_SIZE = 1u << 11;
+    // Number of shards in a column's (fixed-size) shard table. The trace never exceeds the circuit size,
+    // so this bounds the shard index. INTERVAL_SIZE divides MAX_AVM_TRACE_SIZE evenly.
+    static constexpr size_t NUM_SHARDS = MAX_AVM_TRACE_SIZE / INTERVAL_SIZE;
+
     TraceContainer();
 
     const FF& get(Column col, uint32_t row) const;
@@ -42,7 +82,7 @@ class TraceContainer {
     // Reserve column size. Useful for precomputed columns.
     void reserve_column(Column col, size_t size);
 
-    // Visits non-zero values in a column.
+    // Visits non-zero values in a column. The visit order is unspecified.
     void visit_column(Column col, const std::function<void(uint32_t, const FF&)>& visitor) const;
     // Returns the number of rows in a column. That is, the maximum non-zero row index + 1.
     uint32_t get_column_rows(Column col) const;
@@ -53,30 +93,60 @@ class TraceContainer {
     // Number of columns (without shifts).
     static constexpr size_t num_columns() { return NUM_COLUMNS_WITHOUT_SHIFTS; }
 
-    // Batch inverts a set of columns.
+    // Batch inverts a set of columns. Not thread-safe.
     void invert_columns(std::span<const Column> cols);
 
-    // Free column memory.
+    // Free column memory. Not thread-safe.
     void clear_column(Column col);
 
   private:
-    // We use a mutex per column to allow for concurrent writes.
-    // Observe that therefore concurrent write access to different columns is cheap.
+    // A contiguous range of rows [k*INTERVAL_SIZE, (k+1)*INTERVAL_SIZE) within a column.
+    // Rows are stored densely: rows[row % INTERVAL_SIZE] holds the value for that absolute row,
+    // with zero meaning "absent" (an unset cell reads as zero anyway). Shards are top-dense — once a
+    // shard exists it tends to fill — so a flat array beats a hash map: no hashing/rehashing, sequential
+    // cache-friendly writes, stable element addresses, and less memory than the map's key + load-factor
+    // overhead.
+    //
+    // Concurrency: there is no per-shard lock. The array is fixed-size and never reallocates, so
+    // concurrent writes to different rows of this shard touch distinct objects and are race-free by the
+    // C++ memory model, including when two chunks meet at a shard boundary. Same-cell concurrency is the
+    // caller's responsibility (see the thread-safety contract on TraceContainer). The row count is derived
+    // lazily by scanning (see get_column_rows), avoiding any per-write bookkeeping.
+    //
+    // NOTE: the {} initializer on rows is required. FF's default constructor is trivial, so it zeroes
+    // only under value-initialization; a plain `std::array<FF, INTERVAL_SIZE> rows;` would leave garbage,
+    // breaking the "unset cell reads as zero" invariant.
+    struct ColumnInterval {
+        std::array<FF, INTERVAL_SIZE> rows{};
+    };
+    // A column is a fixed-size table of lazily-created shards (index = row / INTERVAL_SIZE), held inline
+    // as atomic pointers. The hot get/set path reads a slot with a single relaxed/acquire atomic load and
+    // never takes a lock, so concurrent writers of the same column do not serialize. A shard is created
+    // at most once: a writer that finds a null slot constructs one and installs it with a compare-and-swap,
+    // keeping it only if the slot was still null. Slots are never cleared concurrently with get/set
+    // (clear_column and the destructor run after the parallel fill phase).
     struct SparseColumn {
-        std::shared_mutex mutex;
-        int64_t max_row_number = -1; // We use -1 to indicate that the column is empty.
-        bool row_number_dirty;       // Needs recalculation.
-        // Future memory optimization notes: we can do the same trick as in Operand.
-        // That is, store a variant with a unique_ptr. However, we should benchmark this.
-        // (see serialization.hpp).
-        unordered_flat_map<uint32_t, FF> rows;
+        SparseColumn() = default;
+        ~SparseColumn()
+        {
+            for (size_t k = 0; k < NUM_SHARDS; ++k) {
+                delete slots[k].load(std::memory_order_relaxed);
+            }
+        }
+        SparseColumn(const SparseColumn&) = delete;
+        SparseColumn& operator=(const SparseColumn&) = delete;
+
+        // The {} zero-initializes every slot to nullptr (std::atomic value-initializes its value in C++20).
+        std::array<std::atomic<ColumnInterval*>, NUM_SHARDS> slots{};
     };
-    // We store the trace as a sparse matrix.
-    // We use a unique_ptr to allocate the array in the heap vs the stack.
-    // Even if the _content_ of each unordered_map is always heap-allocated, if we have 3k columns
-    // we could unnecessarily put strain on the stack with sizeof(unordered_map) * 3k bytes.
+    // We store the trace as a sparse matrix. Each SparseColumn holds its shard table inline, so
+    // sizeof(SparseColumn) is ~NUM_SHARDS atomic pointers; with ~3k columns the whole array is tens of
+    // MB, so we heap-allocate it via unique_ptr rather than placing it on the stack.
     std::unique_ptr<std::array<SparseColumn, NUM_COLUMNS_WITHOUT_SHIFTS>> trace;
 
+    // Returns the shard owning shard_idx, creating it (once, via a lock-free compare-and-swap) if absent.
+    static ColumnInterval& get_or_create_shard(SparseColumn& column_data, size_t shard_idx);
+    // Not thread-safe.
     void invert_column(Column col);
 };