diff --git a/barretenberg/cpp/src/barretenberg/vm2/tracegen/lib/lookup_builder.hpp b/barretenberg/cpp/src/barretenberg/vm2/tracegen/lib/lookup_builder.hpp index 2d372e03082d..0fe56077cbb1 100644 --- a/barretenberg/cpp/src/barretenberg/vm2/tracegen/lib/lookup_builder.hpp +++ b/barretenberg/cpp/src/barretenberg/vm2/tracegen/lib/lookup_builder.hpp @@ -51,10 +51,10 @@ template class IndexedLookupTraceBuilder : public Int } trace.set(LookupSettings::COUNTS, dst_row, trace.get(LookupSettings::COUNTS, dst_row) + 1); - // Set the fine grained inner selector if it's not already one. - if (LookupSettings::DST_SELECTOR != this->outer_dst_selector && - trace.get(LookupSettings::DST_SELECTOR, dst_row) != 1) { - trace.set(LookupSettings::DST_SELECTOR, dst_row, 1); + if (LookupSettings::DST_SELECTOR != this->outer_dst_selector) { + // This step might write to the same cell from multiple threads, so we use atomic limbs to avoid UB. + // Since we are always writing a 1, the end result will be 1 even under concurrency. + trace.set(LookupSettings::DST_SELECTOR, dst_row, 1, /*use_atomic_limbs=*/true); } }); } @@ -198,10 +198,10 @@ template class LookupIntoDynamicTableSequential : publ if (dst_selector == 1 && src_values == trace.get_multiple(LookupSettings::DST_COLUMNS, dst_row)) { trace.set(LookupSettings::COUNTS, dst_row, trace.get(LookupSettings::COUNTS, dst_row) + 1); - // Set the fine grained inner selector if it's not already one. - if (LookupSettings::DST_SELECTOR != this->outer_dst_selector && - trace.get(LookupSettings::DST_SELECTOR, dst_row) != 1) { - trace.set(LookupSettings::DST_SELECTOR, dst_row, 1); + if (LookupSettings::DST_SELECTOR != this->outer_dst_selector) { + // This step might write to the same cell from multiple threads, so we use atomic limbs to avoid + // UB. Since we are always writing a 1, the end result will be 1 even under concurrency. + trace.set(LookupSettings::DST_SELECTOR, dst_row, 1, /*use_atomic_limbs=*/true); } found = true; diff --git a/barretenberg/cpp/src/barretenberg/vm2/tracegen/trace_container.cpp b/barretenberg/cpp/src/barretenberg/vm2/tracegen/trace_container.cpp index 9d6b8799928d..18937979df87 100644 --- a/barretenberg/cpp/src/barretenberg/vm2/tracegen/trace_container.cpp +++ b/barretenberg/cpp/src/barretenberg/vm2/tracegen/trace_container.cpp @@ -4,6 +4,7 @@ #include #include "barretenberg/common/assert.hpp" +#include "barretenberg/common/compiler_hints.hpp" #include "barretenberg/common/log.hpp" #include "barretenberg/common/ref_vector.hpp" #include "barretenberg/vm2/common/field.hpp" @@ -13,7 +14,19 @@ namespace bb::avm2::tracegen { namespace { // We need a zero value to return (a reference to) when a value is not found. -static const FF zero = FF::zero(); +const FF zero = FF::zero(); + +// Writes each 64-bit limb of the field with a relaxed atomic store. Each limb is naturally aligned (FF is +// alignas(32), 4x uint64_t), so this lowers to 4 plain `movq` stores on x86-64 — no lock, no libatomic call +// (unlike a whole-field std::atomic_ref, whose 32 bytes exceed the lock-free width). Lets set() make a +// same-cell concurrent write data-race-free when every writer stores the same value. +inline void store_per_limb(FF& cell, const FF& value) +{ + static_assert(sizeof(FF) == 4 * sizeof(uint64_t)); + for (size_t i = 0; i < 4; ++i) { + std::atomic_ref(cell.data[i]).store(value.data[i], std::memory_order_relaxed); + } +} } // namespace @@ -62,7 +75,7 @@ TraceContainer::ColumnInterval& TraceContainer::get_or_create_shard(SparseColumn return *expected; // CAS failure loaded the winning pointer into `expected` (acquire). } -void TraceContainer::set(Column col, uint32_t row, const FF& value) +void TraceContainer::set(Column col, uint32_t row, const FF& value, bool use_atomic_limbs) { auto& column_data = (*trace)[static_cast(col)]; const size_t shard_idx = row / INTERVAL_SIZE; @@ -73,12 +86,22 @@ void TraceContainer::set(Column col, uint32_t row, const FF& value) // 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; + auto& cell = get_or_create_shard(column_data, shard_idx).rows[offset]; + if (BB_UNLIKELY(use_atomic_limbs)) { + store_per_limb(cell, value); + } else { + cell = value; + } } else { - // Zero value: clear if present. We never create a shard (clearing an absent row is a no-op). + // Zero value: clear if present. We never create a shard, so sparse (mostly-zero) columns are not + // materialized (an unset cell already reads as zero). ColumnInterval* shard = column_data.slots[shard_idx].load(std::memory_order_acquire); if (shard != nullptr) { - shard->rows[offset] = FF::zero(); + if (BB_UNLIKELY(use_atomic_limbs)) { + store_per_limb(shard->rows[offset], zero); + } else { + shard->rows[offset] = FF::zero(); + } } } } diff --git a/barretenberg/cpp/src/barretenberg/vm2/tracegen/trace_container.hpp b/barretenberg/cpp/src/barretenberg/vm2/tracegen/trace_container.hpp index 22f307c9dca0..87e067aec5fc 100644 --- a/barretenberg/cpp/src/barretenberg/vm2/tracegen/trace_container.hpp +++ b/barretenberg/cpp/src/barretenberg/vm2/tracegen/trace_container.hpp @@ -76,7 +76,11 @@ class TraceContainer { // Extended version of get that works with shifted columns. More expensive. const FF& get_column_or_shift(ColumnAndShifts col, uint32_t row) const; - void set(Column col, uint32_t row, const FF& value); + // Sets the value of a cell. Thread-safe if the same cell is not written to from multiple threads. + // If writing to the same cell from multiple threads, use_atomic_limbs=true to use atomic limbs. + // This makes the write slower, but it will not be UB. However, it is also not thread-safe. + // Use only if you know what you are doing. + void set(Column col, uint32_t row, const FF& value, bool use_atomic_limbs = false); // Bulk setting for a given row. void set(uint32_t row, std::span> values); // Reserve column size. Useful for precomputed columns.