Early exit if backend is idle and there is work to do. (#1503)

* Early exit gathering if all collisions and backend is idle.

* Fix logic.

* Even more aggressive.

* Don't enable if threads=1.

* Parameterise the behavior.

Author: Tilps

src/mcts/params.cc

@@ -297,6 +297,16 @@ const OptionId SearchParams::kMinimumWorkPerTaskForProcessingId{
     "minimum-per-task-processing", "MinimumPerTaskProcessing",
     "Processing work won't be split into chunks smaller than this (unless its "
     "more than half of MinimumProcessingWork)."};
+const OptionId SearchParams::kIdlingMinimumWorkId{
+    "idling-minimum-work", "IdlingMinimumWork",
+    "Only early exit gathering due to 'idle' backend if more than this many "
+    "nodes will be sent to the backend."};
+const OptionId SearchParams::kThreadIdlingThresholdId{
+    "thread-idling-threshold", "ThreadIdlingThreshold",
+    "If there are more than this number of search threads that are not "
+    "actively in the process of either sending data to the backend or waiting "
+    "for data from the backend, assume that the backend is idle."};
+
 void SearchParams::Populate(OptionsParser* options) {
   // Here the uci optimized defaults" are set.
   // Many of them are overridden with training specific values in tournament.cc.
@@ -369,6 +379,8 @@ void SearchParams::Populate(OptionsParser* options) {
   options->Add<IntOption>(kMinimumRemainingWorkSizeForPickingId, 0, 100000) =
       20;
   options->Add<IntOption>(kMinimumWorkPerTaskForProcessingId, 1, 100000) = 8;
+  options->Add<IntOption>(kIdlingMinimumWorkId, 0, 10000) = 0;
+  options->Add<IntOption>(kThreadIdlingThresholdId, 0, 128) = 1;
 
   options->HideOption(kNoiseEpsilonId);
   options->HideOption(kNoiseAlphaId);
@@ -449,7 +461,9 @@ SearchParams::SearchParams(const OptionsDict& options)
       kMinimumRemainingWorkSizeForPicking(
           options.Get<int>(kMinimumRemainingWorkSizeForPickingId)),
       kMinimumWorkPerTaskForProcessing(
-          options.Get<int>(kMinimumWorkPerTaskForProcessingId)) {
+          options.Get<int>(kMinimumWorkPerTaskForProcessingId)),
+      kIdlingMinimumWork(options.Get<int>(kIdlingMinimumWorkId)),
+      kThreadIdlingThreshold(options.Get<int>(kThreadIdlingThresholdId)) {
   if (std::max(std::abs(kDrawScoreSidetomove), std::abs(kDrawScoreOpponent)) +
           std::max(std::abs(kDrawScoreWhite), std::abs(kDrawScoreBlack)) >
       1.0f) {

src/mcts/params.h

@@ -128,6 +128,8 @@ class SearchParams {
   int GetMinimumWorkPerTaskForProcessing() const {
     return kMinimumWorkPerTaskForProcessing;
   }
+  int GetIdlingMinimumWork() const { return kIdlingMinimumWork; }
+  int GetThreadIdlingThreshold() const { return kThreadIdlingThreshold; }
 
   // Search parameter IDs.
   static const OptionId kMiniBatchSizeId;
@@ -187,6 +189,8 @@ class SearchParams {
   static const OptionId kMinimumWorkSizeForPickingId;
   static const OptionId kMinimumRemainingWorkSizeForPickingId;
   static const OptionId kMinimumWorkPerTaskForProcessingId;
+  static const OptionId kIdlingMinimumWorkId;
+  static const OptionId kThreadIdlingThresholdId;
 
  private:
   const OptionsDict& options_;
@@ -239,6 +243,8 @@ class SearchParams {
   const int kMinimumWorkSizeForPicking;
   const int kMinimumRemainingWorkSizeForPicking;
   const int kMinimumWorkPerTaskForProcessing;
+  const int kIdlingMinimumWork;
+  const int kThreadIdlingThreshold;
 };
 
 }  // namespace lczero

src/mcts/search.cc

@@ -795,6 +795,7 @@ EdgeAndNode Search::GetBestRootChildWithTemperature(float temperature) const {
 }
 
 void Search::StartThreads(size_t how_many) {
+  thread_count_.store(how_many, std::memory_order_release);
   Mutex::Lock lock(threads_mutex_);
   // First thread is a watchdog thread.
   if (threads_.size() == 0) {
@@ -1073,6 +1074,7 @@ void SearchWorker::ExecuteOneIteration() {
   } else {
     GatherMinibatch();
   }
+  search_->backend_waiting_counter_.fetch_add(1, std::memory_order_relaxed);
 
   // 2b. Collect collisions.
   CollectCollisions();
@@ -1086,6 +1088,7 @@ void SearchWorker::ExecuteOneIteration() {
 
   // 4. Run NN computation.
   RunNNComputation();
+  search_->backend_waiting_counter_.fetch_add(-1, std::memory_order_relaxed);
 
   // 5. Retrieve NN computations (and terminal values) into nodes.
   FetchMinibatchResults();
@@ -1210,6 +1213,8 @@ void SearchWorker::GatherMinibatch2() {
   // Number of nodes processed out of order.
   number_out_of_order_ = 0;
 
+  int thread_count = search_->thread_count_.load(std::memory_order_acquire);
+
   // Gather nodes to process in the current batch.
   // If we had too many nodes out of order, also interrupt the iteration so
   // that search can exit.
@@ -1218,6 +1223,20 @@ void SearchWorker::GatherMinibatch2() {
     // If there's something to process without touching slow neural net, do it.
     if (minibatch_size > 0 && computation_->GetCacheMisses() == 0) return;
 
+    // If there is backend work to be done, and the backend is idle - exit
+    // immediately.
+    // Only do this fancy work if there are multiple threads as otherwise we
+    // early exit from every batch since there is never another search thread to
+    // be keeping the backend busy. Which would mean that threads=1 has a
+    // massive nps drop.
+    if (thread_count > 1 && minibatch_size > 0 &&
+        computation_->GetCacheMisses() > params_.GetIdlingMinimumWork() &&
+        thread_count - search_->backend_waiting_counter_.load(
+                           std::memory_order_relaxed) >
+            params_.GetThreadIdlingThreshold()) {
+      return;
+    }
+
     int new_start = static_cast<int>(minibatch_.size());
 
     PickNodesToExtend(

src/mcts/search.h

@@ -193,6 +193,8 @@ class Search {
       GUARDED_BY(counters_mutex_);
 
   std::atomic<int> pending_searchers_{0};
+  std::atomic<int> backend_waiting_counter_{0};
+  std::atomic<int> thread_count_{0};
 
   std::vector<std::pair<Node*, int>> shared_collisions_
       GUARDED_BY(nodes_mutex_);