Reduce overhead in TCMalloc benchmark

PiperOrigin-RevId: 783348068
Change-Id: If7c9612666ff31ac4f3dd6dc53752fb0410b5fc9

Author: andreas-abel

fleetbench/tcmalloc/empirical.cc

@@ -86,7 +86,9 @@ EmpiricalData::EmpiricalData(size_t seed, const absl::Span<const Entry> weights,
       total_bytes_allocated_(0),
       birth_sampler_(BirthRateDistribution(weights)),
       total_birth_rate_(0),
-      death_sampler_(weights.size()) {
+      death_sampler_(weights.size()),
+      num_allocated_recorded_(0),
+      bytes_allocated_recorded_(0) {
   // First, compute average live count for each size in a heap of size
   // <total_mem>.
   double total = 0;
@@ -168,9 +170,12 @@ void EmpiricalData::DoDeath(const size_t i) {
 }
 
 void EmpiricalData::RecordBirth(const size_t i) {
+  birth_or_death_.push_back(true);
   birth_or_death_sizes_.push_back(i);
   SizeState& s = state_[i];
   death_sampler_.AdjustWeight(i, s.death_rate);
+  num_allocated_recorded_++;
+  bytes_allocated_recorded_ += s.size;
   // We only care about keeping the number of objects correct when building the
   // trace.  When we replay we will actually push the allocated address but
   // when building the trace we can just push nullptr to keep the length of live
@@ -181,14 +186,13 @@ void EmpiricalData::RecordBirth(const size_t i) {
 void* EmpiricalData::ReplayBirth(const size_t i) {
   SizeState& s = state_[i];
   const size_t size = s.size;
-  total_num_allocated_++;
-  total_bytes_allocated_ += size;
   void* p = alloc_(size);
   s.objs.push_back(p);
   return p;
 }
 
 void EmpiricalData::RecordDeath(const size_t i) {
+  birth_or_death_.push_back(false);
   SizeState& s = state_[i];
   CHECK(!s.objs.empty());
   birth_or_death_sizes_.push_back(i);
@@ -215,7 +219,6 @@ void EmpiricalData::RecordNext() {
   const double Both = B + T;
   absl::uniform_real_distribution<double> which(0, Both);
   bool do_birth = which(rng_) < B;
-  birth_or_death_.push_back(do_birth);
 
   if (do_birth) {
     size_t i = birth_sampler_(rng_);
@@ -226,18 +229,25 @@ void EmpiricalData::RecordNext() {
   }
 }
 
-void EmpiricalData::ReplayNext() {
-  bool do_birth = birth_or_death_[birth_or_death_index_];
-  if (do_birth) {
-    void* allocated = ReplayBirth(birth_or_death_sizes_[birth_or_death_index_]);
-    TouchAllocated(allocated);
-  } else {
-    ReplayDeath(birth_or_death_sizes_[birth_or_death_index_],
-                death_objects_[death_object_index_]);
-    __builtin_prefetch(death_object_pointers_[death_object_index_], 1, 3);
-    death_object_index_++;
+void EmpiricalData::ReplayTrace() {
+  for (birth_or_death_index_ = 0, death_object_index_ = 0;
+       birth_or_death_index_ < birth_or_death_.size();
+       ++birth_or_death_index_) {
+    bool do_birth = birth_or_death_[birth_or_death_index_];
+    if (do_birth) {
+      void* allocated =
+          ReplayBirth(birth_or_death_sizes_[birth_or_death_index_]);
+      TouchAllocated(allocated);
+    } else {
+      ReplayDeath(birth_or_death_sizes_[birth_or_death_index_],
+                  death_objects_[death_object_index_]);
+      __builtin_prefetch(death_object_pointers_[death_object_index_], /*rw=*/1,
+                         /*locality*/ 3);
+      ++death_object_index_;
+    }
   }
-  birth_or_death_index_++;
+  total_num_allocated_ += num_allocated_recorded_;
+  total_bytes_allocated_ += bytes_allocated_recorded_;
 }
 
 void EmpiricalData::SnapshotLiveObjects() {
@@ -301,7 +311,7 @@ void EmpiricalData::BuildDeathObjectPointers() {
               death_object_pointers_.end());
 }
 
-void EmpiricalData::RepairToSnapshotState() {
+void EmpiricalData::RecordRepairToSnapshotState() {
   // Compared to the number of live objects when the snapshot was taken each
   // size state either
   // 1) Contains the same number of live objects as when the snapshot was taken,
@@ -312,29 +322,14 @@ void EmpiricalData::RepairToSnapshotState() {
   //    number of true deallocations.
   for (int i = 0; i < state_.size(); i++) {
     while (state_[i].objs.size() < snapshot_state_[i].objs.size()) {
-      DoBirth(i);
+      RecordBirth(i);
     }
     while (state_[i].objs.size() > snapshot_state_[i].objs.size()) {
-      DoDeath(i);
+      RecordDeath(i);
     }
   }
 }
 
-void EmpiricalData::RestartTraceIfNecessary() {
-  if (birth_or_death_index_ == birth_or_death_.size()) {
-    // As the snapshotted lists of live objects will contain addresses which
-    // have already been freed we can't just call RestoreSnapshot().  Instead
-    // let's do the necessary allocations / deallocations to end up with the
-    // identical number of live objects we had when initially building the
-    // trace.
-    RepairToSnapshotState();
-    // After the above call we can safely run through the recorded trace
-    // again.
-    birth_or_death_index_ = 0;
-    death_object_index_ = 0;
-  }
-}
-
 std::vector<EmpiricalData::Entry> GetEmpiricalDataEntries(
     absl::string_view file) {
   std::vector<EmpiricalData::Entry> distribution;

fleetbench/tcmalloc/empirical.h

@@ -203,10 +203,11 @@ class EmpiricalData {
   // logic of this function please see the comments within Next().
   void RecordNext();
 
-  // Replays the next alloc or dealloc we recorded when building the trace.
-  // Also updates the indices into the recorded birth / death trace.
-  // incremented.
-  void ReplayNext();
+  // Replays the recorded trace.
+  void ReplayTrace();
+
+  // Returns the number of allocs and deallocs in the recorded trace.
+  size_t TraceLength() const { return birth_or_death_.size(); }
 
   std::default_random_engine* const rng() { return &rng_; }
 
@@ -222,23 +223,17 @@ class EmpiricalData {
   // starting to replay the trace.
   void RestoreSnapshot();
 
-  // Restores the *lengths* of the number of live objects within each size class
-  // to what it was after the warmup allocations were complete.  This is
-  // accomplished by either allocating or deallocating objects until the same
-  // number of objects are live within each size class as were live after the
-  // warmup allocations were complete.  This is safe to call repeatedly.
-  void RepairToSnapshotState();
+  // Records a sequence of allocations and deallocations that restores the
+  // *lengths* of the number of live objects within each size class to what it
+  // was after the warmup allocations were complete. This function should only
+  // be called once (at the end of the recording phase).
+  void RecordRepairToSnapshotState();
 
   // Computes addresses to prefetch when executing in record and replay mode.
   // This is necessary to minimize the impact of indexing into SizeState.objs
   // when freeing an object.
   void BuildDeathObjectPointers();
 
-  // Tests whether we have reached the end of the birth / death trace.  If so
-  // performs the actions necessary so that we can start replaying allocs /
-  // deallocs from the beginning of the trace again.
-  void RestartTraceIfNecessary();
-
  private:
   std::default_random_engine rng_;
 
@@ -276,6 +271,8 @@ class EmpiricalData {
   std::vector<void**> death_object_pointers_;
   uint32_t birth_or_death_index_ = 0;
   uint32_t death_object_index_ = 0;
+  size_t num_allocated_recorded_;
+  size_t bytes_allocated_recorded_;
 };
 
 std::vector<EmpiricalData::Entry> GetEmpiricalDataEntries(

fleetbench/tcmalloc/empirical_driver.cc

@@ -44,17 +44,21 @@ namespace {
 
 void* alloc(size_t s) { return ::operator new(s); }
 
-static constexpr int64_t kBatch = 100;
 // When non-zero, empirical driver will simulate tick of ReleaseMemoryToOS
 // iteration, given number of bytes allocated.
 static constexpr int64_t kSimulatedBytesPerSec = 0;
 // The total number of allocs / deallocs to precalculate for later replay.
 // Memory required to store replay buffers scales with the number of threads.
-static constexpr size_t kRecordAndReplayBufferSize = 1'000'000;
+// The actual number of allocs / deallocs can be slightly larger than this
+// value, as at the end of the replay, additional operations are performed to
+// bring the number of live objects back to where it was at the start of the
+// replay.
+static constexpr size_t kRecordAndReplayBufferSizeTarget = 1'000'000;
 // Number of bytes to try to release from the page heap per second.
 static constexpr int64_t kEmpiricalMallocReleaseBytesPerSec = 0;
-// Number of iterations to warm up the benchmark before the main benchmark loop.
-static constexpr size_t kNumWarmUpIterations = 500000;
+// Number of replays of the entire trace to warm up the benchmark before the
+// main benchmark loop.
+static constexpr size_t kNumWarmUpReplays = 50;
 
 class SimThread {
  public:
@@ -80,11 +84,10 @@ class SimThread {
   }
 
   void RecordBirthsAndDeaths(EmpiricalData* load) {
-    // Round number of births / deaths to record down to a multiple of kBatch.
-    const int buffer_size = (kRecordAndReplayBufferSize / kBatch) * kBatch;
-    for (int i = 0; i < buffer_size; ++i) {
+    for (int i = 0; i < kRecordAndReplayBufferSizeTarget; ++i) {
       load->RecordNext();
     }
+    load->RecordRepairToSnapshotState();
 
     load->RestoreSnapshot();
     load->BuildDeathObjectPointers();
@@ -106,10 +109,8 @@ class SimThread {
 
   void ReplayTrace() {
     DCHECK(done_recording_);
-    for (int i = 0; i < kBatch; i++) {
-      load_.ReplayNext();
-    }
-    load_.RestartTraceIfNecessary();
+    load_.ReplayTrace();
+
     auto allocated = load_.total_bytes_allocated();
     load_bytes_allocated_.store(allocated, std::memory_order_relaxed);
     auto total_num_allocated = load_.total_num_allocated();
@@ -122,6 +123,8 @@ class SimThread {
     }
   }
 
+  size_t TraceLength() { return load_.TraceLength(); }
+
  private:
   size_t n_;
   size_t transient_;
@@ -217,14 +220,14 @@ static void BM_TCMalloc_Empirical_Driver(benchmark::State& state) {
   // We do not use the MinWarmUpTime feature of the benchmark framework here,
   // as that feature calls Teardown and Setup after the warm-up phase, which
   // resets the state that we want to establish with the warm-up.
-  for (int i = 0; i < kNumWarmUpIterations; i++) {
+  for (int i = 0; i < kNumWarmUpReplays; i++) {
     sim_threads[thread_idx]->ReplayTrace();
   }
 
   size_t bytes_warm_up = sim_threads[thread_idx]->total_bytes_allocated();
   size_t allocations_warm_up = sim_threads[thread_idx]->load_allocations();
 
-  for (auto _ : state) {
+  while (state.KeepRunningBatch(sim_threads[thread_idx]->TraceLength())) {
     sim_threads[thread_idx]->ReplayTrace();
   }
 
@@ -315,7 +318,7 @@ void RegisterBenchmarks() {
           ->Teardown(BM_TCMalloc_Empirical_Driver_Teardown)
           ->ThreadRange(1, 1)
           ->UseRealTime()
-          ->Iterations(100000);
+          ->Iterations(10'000'000);
       min_threads = 2;
     }
 

fleetbench/tcmalloc/empirical_test.cc

@@ -56,15 +56,14 @@ TEST(EmpiricalRecordAndReplay, Basic) {
   for (int j = 0; j < kBufferSize; ++j) {
     data.RecordNext();
   }
+  data.RecordRepairToSnapshotState();
 
   data.RestoreSnapshot();
   data.BuildDeathObjectPointers();
 
   // We need one warmup iteration so we can compute the delta allocations and
   // bytes we should see from each time through the trace.
-  for (int j = 0; j < kBufferSize; ++j) {
-    data.ReplayNext();
-  }
+  data.ReplayTrace();
 
   size_t delta_allocations = data.total_num_allocated() - total_allocations;
   size_t delta_bytes_allocated =
@@ -76,16 +75,10 @@ TEST(EmpiricalRecordAndReplay, Basic) {
     EXPECT_EQ(delta_bytes_allocated,
               data.total_bytes_allocated() - total_bytes_allocated);
 
-    // Restart the trace before updating total_* so we don't capture the
-    // "repair" operations.
-    data.RestartTraceIfNecessary();
-
     total_allocations = data.total_num_allocated();
     total_bytes_allocated = data.total_bytes_allocated();
 
-    for (int j = 0; j < kBufferSize; ++j) {
-      data.ReplayNext();
-    }
+    data.ReplayTrace();
   }
 }