Merge e288bb8 into df415ad

sam-panzer · web-flow · commit de38336323d5 · 2018-01-31T19:10:44.000Z
diff --git a/include/benchmark/benchmark.h b/include/benchmark/benchmark.h
@@ -172,6 +172,7 @@ BENCHMARK(BM_test)->Unit(benchmark::kMillisecond);
 
 #include <stdint.h>
 
+#include <algorithm>
 #include <cassert>
 #include <cstddef>
 #include <iosfwd>
@@ -429,16 +430,18 @@ class State {
   // Returns true if the benchmark should continue through another iteration.
   // NOTE: A benchmark may not return from the test until KeepRunning() has
   // returned false.
-  bool KeepRunning() {
-    if (BENCHMARK_BUILTIN_EXPECT(!started_, false)) {
-      StartKeepRunning();
-    }
-    bool const res = (--total_iterations_ != 0);
-    if (BENCHMARK_BUILTIN_EXPECT(!res, false)) {
-      FinishKeepRunning();
-    }
-    return res;
-  }
+  bool KeepRunning();
+
+  // Returns true iff the benchmark should run n more iterations.
+  // NOTE: A benchmark must not return from the test until KeepRunningBatch()
+  // has returned false.
+  // NOTE: KeepRunningBatch() may overshoot by up to 'n' iterations.
+  //
+  // Intended usage:
+  //   while (state.KeepRunningBatch(1000)) {
+  //     // process 1000 elements
+  //   }
+  bool KeepRunningBatch(int n);
 
   // REQUIRES: timer is running and 'SkipWithError(...)' has not been called
   //           by the current thread.
@@ -565,12 +568,15 @@ class State {
   int range_y() const { return range(1); }
 
   BENCHMARK_ALWAYS_INLINE
-  size_t iterations() const { return (max_iterations - total_iterations_) + 1; }
+  size_t iterations() const { return (completed_iterations_ - total_iterations_); }
 
  private:
   bool started_;
   bool finished_;
+  // When total_iterations_ is 0, KeepRunning() and friends will return false.
   size_t total_iterations_;
+  // May be larger than max_iterations.
+  size_t completed_iterations_;
 
   std::vector<int> range_;
 
@@ -581,6 +587,11 @@ class State {
 
   bool error_occurred_;
 
+  // When using KeepRunningBatch(), batch_leftover_ holds the number of
+  // iterations beyond max_iters that were run. Used to track
+  // completed_iterations_ accurately.
+  size_t batch_leftover_;
+
  public:
   // Container for user-defined counters.
   UserCounters counters;
@@ -597,12 +608,43 @@ class State {
 
  private:
   void StartKeepRunning();
-  void FinishKeepRunning();
+  void FinishKeepRunning(size_t n);
   internal::ThreadTimer* timer_;
   internal::ThreadManager* manager_;
   BENCHMARK_DISALLOW_COPY_AND_ASSIGN(State);
 };
 
+BENCHMARK_ALWAYS_INLINE
+inline bool State::KeepRunningBatch(int n) {
+  if (BENCHMARK_BUILTIN_EXPECT(!started_, false)) {
+    StartKeepRunning();
+  }
+  // leftover only has a sensible value when n > total_iterations.
+  size_t leftover = n - total_iterations_;
+  // total_iterations_ is a size_t. Avoid wrapping around past 0.
+  bool const res = (total_iterations_ > 0);
+  total_iterations_ -= std::min<size_t>(n, total_iterations_);
+  if (BENCHMARK_BUILTIN_EXPECT(!res, false)) {
+    FinishKeepRunning(batch_leftover_);
+  }
+  // If this batch has 0 < n < total_iterations iterations, (i.e. it's the last
+  // batch), track how many extra iterations past max_iters we ran.
+  batch_leftover_ = leftover;
+  return res;
+}
+
+BENCHMARK_ALWAYS_INLINE
+inline bool State::KeepRunning() {
+    if (BENCHMARK_BUILTIN_EXPECT(!started_, false)) {
+      StartKeepRunning();
+    }
+    bool const res = (total_iterations_-- != 0);
+    if (BENCHMARK_BUILTIN_EXPECT(!res, false)) {
+      FinishKeepRunning(0);
+    }
+    return res;
+}
+
 struct State::StateIterator {
   struct BENCHMARK_UNUSED Value {};
   typedef std::forward_iterator_tag iterator_category;
@@ -633,7 +675,7 @@ struct State::StateIterator {
   BENCHMARK_ALWAYS_INLINE
   bool operator!=(StateIterator const&) const {
     if (BENCHMARK_BUILTIN_EXPECT(cached_ != 0, true)) return true;
-    parent_->FinishKeepRunning();
+    parent_->FinishKeepRunning(0);
     return false;
   }
 
diff --git a/src/benchmark.cc b/src/benchmark.cc
@@ -268,7 +268,7 @@ void RunInThread(const benchmark::internal::Benchmark::Instance* b,
   internal::ThreadTimer timer;
   State st(iters, b->arg, thread_id, b->threads, &timer, manager);
   b->benchmark->Run(st);
-  CHECK(st.iterations() == st.max_iterations)
+  CHECK(st.iterations() >= st.max_iterations)
       << "Benchmark returned before State::KeepRunning() returned false!";
   {
     MutexLock l(manager->GetBenchmarkMutex());
@@ -399,12 +399,13 @@ State::State(size_t max_iters, const std::vector<int>& ranges, int thread_i,
              internal::ThreadManager* manager)
     : started_(false),
       finished_(false),
-      total_iterations_(max_iters + 1),
+      total_iterations_(max_iters),
       range_(ranges),
       bytes_processed_(0),
       items_processed_(0),
       complexity_n_(0),
       error_occurred_(false),
+      batch_leftover_(0),
       counters(),
       thread_index(thread_i),
       threads(n_threads),
@@ -437,7 +438,7 @@ void State::SkipWithError(const char* msg) {
       manager_->results.has_error_ = true;
     }
   }
-  total_iterations_ = 1;
+  total_iterations_ = 0;
   if (timer_->running()) timer_->StopTimer();
 }
 
@@ -453,17 +454,20 @@ void State::SetLabel(const char* label) {
 void State::StartKeepRunning() {
   CHECK(!started_ && !finished_);
   started_ = true;
+  completed_iterations_ = max_iterations;
   manager_->StartStopBarrier();
   if (!error_occurred_) ResumeTiming();
 }
 
-void State::FinishKeepRunning() {
+void State::FinishKeepRunning(size_t n) {
   CHECK(started_ && (!finished_ || error_occurred_));
   if (!error_occurred_) {
     PauseTiming();
   }
-  // Total iterations has now wrapped around zero. Fix this.
-  total_iterations_ = 1;
+  // Batch size overshot total_iterations_. Account for it.
+  completed_iterations_ += n;
+  // Total iterations has now wrapped around past 0. Fix this.
+  total_iterations_ = 0;
   finished_ = true;
   manager_->StartStopBarrier();
 }
diff --git a/test/basic_test.cc b/test/basic_test.cc
@@ -102,10 +102,20 @@ void BM_KeepRunning(benchmark::State& state) {
   while (state.KeepRunning()) {
     ++iter_count;
   }
-  assert(iter_count == state.max_iterations);
+  assert(iter_count == state.iterations());
 }
 BENCHMARK(BM_KeepRunning);
 
+void BM_KeepRunningBatch(benchmark::State& state) {
+  const size_t batch_size = 100;
+  size_t iter_count = 0;
+  while (state.KeepRunningBatch(batch_size)) {
+    iter_count += batch_size;
+  }
+  assert(state.iterations() == iter_count);
+}
+BENCHMARK(BM_KeepRunningBatch);
+
 void BM_RangedFor(benchmark::State& state) {
   size_t iter_count = 0;
   for (auto _ : state) {
