Skip to content

Commit 5c8a993

Browse files
authored
feat(avm): fast entity indexing without macros (#13737)
Possibly the final iteration of the work started in #11605 . The macro was blowing up when adding more columns. Instead, I found a way to precompute an array that lets us efficiently access entities via the column enum. BEFORE ``` Compile time: 2m 20s Memory (compilation): 5GB ``` AFTER ``` Compile time: 1m 50s Memory (compilation): 2.6GB ``` Runtime is unaffected.
1 parent def1287 commit 5c8a993

26 files changed

+1305
-1247
lines changed

barretenberg/cpp/scripts/analyze_vm_compile_time.sh

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,26 +1,26 @@
11
#!/usr/bin/env bash
22
# This script summarises the compilation time for the vm
33
# The summary json file is outputted to $BUILD_DIR/avm_compilation_summary.json
4-
# it takes in two params the preset(e.g. clang16, clang16-dbg) and a target (e.g. bb, vm)
4+
# it takes in two params the preset(e.g. clang16, clang16-dbg) and a target (e.g. bb, vm)
55
# it can be called like this => ./analyze_vm_compile_time.sh clang16 bb
66
set -eu
77
# So we can glob recursively
88
shopt -s globstar
99

10-
PRESET="${1:-wasm-threads}"
11-
TARGET="${2:-barretenberg.wasm}"
10+
export PRESET="${1:-clang16}"
11+
export TARGET="${2:-vm2}"
1212

1313
BUILD_DIR=build-$PRESET-compiler-profile
1414

1515
cd $(dirname $0)/..
1616

1717
# Run the analyse script if we dont already have the specific directory
18-
if [ ! -d $BUILD_DIR ]; then
19-
echo -e "\n$BUILD_DIR not found, running $(dirname $0)/analyze_compile_time.sh $PRESET $TARGET"
20-
./scripts/analyze_compile_time.sh $PRESET $TARGET
21-
else
18+
if [ ! -d $BUILD_DIR ]; then
19+
echo -e "\n$BUILD_DIR not found, running $(dirname $0)/analyze_compile_time.sh with $PRESET $TARGET"
20+
./scripts/analyze_compile_time.sh
21+
else
2222
echo -e "\n$BUILD_DIR found, using existing results"
23-
fi
23+
fi
2424

2525
# Run summary analysis
2626
cd build-$PRESET-compiler-profile

barretenberg/cpp/src/barretenberg/vm2/common/macros.hpp

Lines changed: 0 additions & 16 deletions
This file was deleted.

barretenberg/cpp/src/barretenberg/vm2/constraining/benchmark/relations_acc.bench.cpp

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -16,14 +16,9 @@ using namespace bb::avm2;
1616

1717
namespace {
1818

19-
// Getters are needed for interaction accumulation to work.
20-
struct FullRowWithGetters : public AvmFullRow<FF> {
21-
DEFINE_GETTERS(DEFAULT_GETTERS, AVM2_ALL_ENTITIES);
22-
};
23-
24-
FullRowWithGetters get_random_row()
19+
AvmFullRow<FF> get_random_row()
2520
{
26-
FullRowWithGetters row;
21+
AvmFullRow<FF> row;
2722
for (size_t i = 0; i < NUM_COLUMNS_WITH_SHIFTS; i++) {
2823
row.get_column(static_cast<ColumnAndShifts>(i)) = FF::random_element();
2924
}

barretenberg/cpp/src/barretenberg/vm2/constraining/flavor.hpp

Lines changed: 29 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@
1212
#include "barretenberg/transcript/transcript.hpp"
1313

1414
#include "barretenberg/vm2/common/aztec_constants.hpp"
15-
#include "barretenberg/vm2/common/macros.hpp"
1615
#include "barretenberg/vm2/constraining/flavor_settings.hpp"
1716

1817
#include "barretenberg/vm2/generated/columns.hpp"
@@ -21,21 +20,21 @@
2120
// Metaprogramming to concatenate tuple types.
2221
template <typename... input_t> using tuple_cat_t = decltype(std::tuple_cat(std::declval<input_t>()...));
2322

24-
// clang-format off
25-
// These getters are used to speedup logderivative inverses.
26-
// See https://github.com/AztecProtocol/aztec-packages/pull/11605/ for a full explanation.
27-
#define DEFAULT_GETTERS(ENTITY) \
28-
inline auto& _##ENTITY() { return ENTITY; } \
29-
inline auto& _##ENTITY() const { return ENTITY; }
30-
#define ROW_PROXY_GETTERS(ENTITY) \
31-
inline auto& _##ENTITY() { return pp.ENTITY[row_idx]; } \
32-
inline auto& _##ENTITY() const { return pp.ENTITY[row_idx]; }
33-
#define DEFINE_GETTERS(GETTER_MACRO, ENTITIES) \
34-
FOR_EACH(GETTER_MACRO, ENTITIES)
35-
// clang-format on
36-
3723
namespace bb::avm2 {
3824

25+
// This helper lets us access entities by column.
26+
// It is critical to achieve performance when calculating lookup inverses.
27+
// See https://github.com/AztecProtocol/aztec-packages/pull/11605 for more details.
28+
template <typename Entities> auto& get_entity_by_column(Entities& entities, ColumnAndShifts c)
29+
{
30+
// A statically constructed pointer to members of the class, indexed by column.
31+
// This should only be created once per Entities class.
32+
static std::array<typename Entities::DataType(Entities::*), NUM_COLUMNS_WITH_SHIFTS> col_ptrs = {
33+
AVM2_ALL_ENTITIES_E(&Entities::)
34+
};
35+
return (entities.*col_ptrs[static_cast<size_t>(c)]);
36+
}
37+
3938
class AvmFlavor {
4039
public:
4140
using Curve = AvmFlavorSettings::Curve;
@@ -139,32 +138,28 @@ class AvmFlavor {
139138
template <typename DataType> class PrecomputedEntities {
140139
public:
141140
DEFINE_FLAVOR_MEMBERS(DataType, AVM2_PRECOMPUTED_ENTITIES)
142-
DEFINE_GETTERS(DEFAULT_GETTERS, AVM2_PRECOMPUTED_ENTITIES)
143141
};
144142

145143
private:
146144
template <typename DataType> class WireEntities {
147145
public:
148146
DEFINE_FLAVOR_MEMBERS(DataType, AVM2_WIRE_ENTITIES)
149-
DEFINE_GETTERS(DEFAULT_GETTERS, AVM2_WIRE_ENTITIES)
150147
};
151148

152149
template <typename DataType> class DerivedWitnessEntities {
153150
public:
154151
DEFINE_FLAVOR_MEMBERS(DataType, AVM2_DERIVED_WITNESS_ENTITIES)
155-
DEFINE_GETTERS(DEFAULT_GETTERS, AVM2_DERIVED_WITNESS_ENTITIES)
156152
};
157153

158154
template <typename DataType> class ShiftedEntities {
159155
public:
160156
DEFINE_FLAVOR_MEMBERS(DataType, AVM2_SHIFTED_ENTITIES)
161-
DEFINE_GETTERS(DEFAULT_GETTERS, AVM2_SHIFTED_ENTITIES)
162157
};
163158

164159
template <typename DataType, typename PrecomputedAndWitnessEntitiesSuperset>
165-
static auto get_to_be_shifted([[maybe_unused]] PrecomputedAndWitnessEntitiesSuperset& entities)
160+
static auto get_to_be_shifted(PrecomputedAndWitnessEntitiesSuperset& entities)
166161
{
167-
return RefArray<DataType, NUM_SHIFTED_ENTITIES>{ AVM2_TO_BE_SHIFTED(entities) };
162+
return RefArray<DataType, NUM_SHIFTED_ENTITIES>{ AVM2_TO_BE_SHIFTED_E(entities.) };
168163
}
169164

170165
public:
@@ -178,11 +173,12 @@ class AvmFlavor {
178173
static const auto& get_derived_labels() { return DerivedWitnessEntities<DataType>::get_labels(); }
179174
};
180175

181-
template <typename DataType>
182-
class AllEntities : public PrecomputedEntities<DataType>,
183-
public WitnessEntities<DataType>,
184-
public ShiftedEntities<DataType> {
176+
template <typename DataType_>
177+
class AllEntities : public PrecomputedEntities<DataType_>,
178+
public WitnessEntities<DataType_>,
179+
public ShiftedEntities<DataType_> {
185180
public:
181+
using DataType = DataType_;
186182
DEFINE_COMPOUND_GET_ALL(PrecomputedEntities<DataType>, WitnessEntities<DataType>, ShiftedEntities<DataType>)
187183

188184
auto get_unshifted()
@@ -200,6 +196,10 @@ class AvmFlavor {
200196
auto get_to_be_shifted() { return AvmFlavor::get_to_be_shifted<DataType>(*this); }
201197
auto get_shifted() { return ShiftedEntities<DataType>::get_all(); }
202198
auto get_precomputed() { return PrecomputedEntities<DataType>::get_all(); }
199+
200+
// We need both const and non-const versions.
201+
DataType& get(ColumnAndShifts c) { return get_entity_by_column(*this, c); }
202+
const DataType& get(ColumnAndShifts c) const { return get_entity_by_column(*this, c); }
203203
};
204204

205205
class ProvingKey : public PrecomputedEntities<Polynomial>, public WitnessEntities<Polynomial> {
@@ -258,6 +258,7 @@ class AvmFlavor {
258258
std::vector<FF> to_field_elements() const;
259259
};
260260

261+
// Used by sumcheck.
261262
class AllValues : public AllEntities<FF> {
262263
public:
263264
using Base = AllEntities<FF>;
@@ -270,7 +271,6 @@ class AvmFlavor {
270271
using BaseDataType = const FF;
271272
using DataType = BaseDataType&;
272273
DEFINE_FLAVOR_MEMBERS(DataType, AVM2_ALL_ENTITIES)
273-
DEFINE_GETTERS(DEFAULT_GETTERS, AVM2_ALL_ENTITIES)
274274
};
275275

276276
template <typename Polynomials> class PolynomialEntitiesAtFixedRow {
@@ -279,7 +279,10 @@ class AvmFlavor {
279279
: row_idx(row_idx)
280280
, pp(pp)
281281
{}
282-
DEFINE_GETTERS(ROW_PROXY_GETTERS, AVM2_ALL_ENTITIES)
282+
283+
// We need both const and non-const versions.
284+
auto& get(ColumnAndShifts c) { return pp.get(c)[row_idx]; }
285+
const auto& get(ColumnAndShifts c) const { return pp.get(c)[row_idx]; }
283286

284287
private:
285288
const size_t row_idx;

barretenberg/cpp/src/barretenberg/vm2/constraining/full_row.hpp

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,11 @@ template <typename FF_> struct AvmFullRow {
2121
static_assert(sizeof(*this) == sizeof(FF) * static_cast<size_t>(ColumnAndShifts::SENTINEL_DO_NOT_USE));
2222
return reinterpret_cast<const FF*>(this)[static_cast<size_t>(col)];
2323
}
24+
25+
// These are the names used by AllEntities, etc.
26+
// TODO(fcarreiro): Clean up duplication.
27+
FF& get(ColumnAndShifts col) { return get_column(col); }
28+
const FF& get(ColumnAndShifts col) const { return get_column(col); }
2429
};
2530

2631
} // namespace bb::avm2

barretenberg/cpp/src/barretenberg/vm2/generated/columns.hpp

Lines changed: 16 additions & 10 deletions
Large diffs are not rendered by default.

0 commit comments

Comments
 (0)