Avoids fetching the same position many times in a batch

src/engine.cc

@@ -169,7 +169,7 @@ void Engine::UpdateBackendConfig() {
       backend_->UpdateConfiguration(options_) == Backend::NEED_RESTART) {
     backend_name_ = backend_name;
     backend_ = CreateMemCache(BackendManager::Get()->CreateFromParams(options_),
-                              options_);
+                              options_, search_->GetMaxOutOfOrderFactor());
     search_->SetBackend(backend_.get());
   } else {
     backend_->SetCacheSize(

src/neural/backend.h

@@ -79,6 +79,8 @@ class BackendComputation {
   enum AddInputResult {
     ENQUEUED_FOR_EVAL = 0,    // Will be computed during ComputeBlocking();
     FETCHED_IMMEDIATELY = 1,  // Was in cache, the result is already populated.
+    FETCHED_DELAYED = 2,      // Was already queued for evaluation but results
+                              // aren't available yet.
   };
   virtual AddInputResult AddInput(
       const EvalPosition& pos,    // Input position.
@@ -133,4 +135,4 @@ class BackendFactory {
   virtual std::unique_ptr<Backend> Create(const OptionsDict&) = 0;
 };
 
-}  // namespace lczero
+}  // namespace lczero

src/neural/memcache.cc

@@ -28,9 +28,9 @@
 #include "neural/memcache.h"
 
 #include "neural/shared_params.h"
+#include "utils/atomic.h"
 #include "utils/atomic_vector.h"
 #include "utils/cache.h"
-#include "utils/smallarray.h"
 
 namespace lczero {
 namespace {
@@ -43,6 +43,19 @@ uint64_t ComputeEvalPositionHash(const EvalPosition& pos) {
 }
 
 struct CachedValue {
+  // State transitions happen atomically using release and aquire sematics for
+  // dependant reads and writes. The state progresses in order. Each transition
+  // must happen only once which requires compare and exchange. Secondary
+  // readers must wait for READY state to read the cached value.
+  enum State {
+    NOT_QUEUED, // Initial state before NN submision.
+    NO_WAITERS, // One thread has taken this position to be evaluated. None is
+                // yet waiting for the result.
+    WAITERS,    // Another thread is waiting for results. Setting READY state
+                // must be folled by notify_all to wake up waiters.
+    READY,      // The value is ready. Waiters can read the value and proceed.
+  };
+  WaitableAtomic<State> state = NOT_QUEUED;
   float q;
   float d;
   float m;
@@ -54,15 +67,18 @@ void CachedValueToEvalResult(const CachedValue& cv, const EvalResultPtr& ptr) {
   if (ptr.d) *ptr.d = cv.d;
   if (ptr.q) *ptr.q = cv.q;
   if (ptr.m) *ptr.m = cv.m;
+  assert(cv.num_moves >= ptr.p.size());
   std::copy(cv.p.get(), cv.p.get() + ptr.p.size(), ptr.p.begin());
 }
 
 class MemCache : public CachingBackend {
  public:
-  MemCache(std::unique_ptr<Backend> wrapped, const OptionsDict& options)
+  MemCache(std::unique_ptr<Backend> wrapped, const OptionsDict& options, float max_out_of_order_evals_factor)
       : wrapped_backend_(std::move(wrapped)),
         cache_(options.Get<int>(SharedBackendParams::kNNCacheSizeId)),
-        max_batch_size_(wrapped_backend_->GetAttributes().maximum_batch_size) {}
+        max_batch_size_(
+            wrapped_backend_->GetAttributes().maximum_batch_size *
+            (1.0f + max_out_of_order_evals_factor)) {}
 
   BackendAttributes GetAttributes() const override {
     return wrapped_backend_->GetAttributes();
@@ -107,51 +123,121 @@ class MemCacheComputation : public BackendComputation {
 
  private:
   size_t UsedBatchSize() const override {
-    return wrapped_computation_->UsedBatchSize();
+    return entries_.size();
   }
   virtual AddInputResult AddInput(const EvalPosition& pos,
                                   EvalResultPtr result) override {
     assert(pos.legal_moves.size() == result.p.size() || result.p.empty());
     const uint64_t hash = ComputeEvalPositionHash(pos);
-    {
-      HashKeyedCacheLock<CachedValue> lock(&memcache_->cache_, hash);
-      // Sometimes search queries NN without passing the legal moves. It is
-      // still cached in this case, but in subsequent queries we only return it
-      // if legal moves are not passed again. Otherwise check the size to guard
-      // against hash collisions.
-      if (lock.holds_value() &&
-          (pos.legal_moves.empty() ||
-           (lock->p && lock->num_moves == pos.legal_moves.size()))) {
+    bool to_be_queued = false;
+    auto value = std::make_unique<CachedValue>();
+    EvalResultPtr result_ptr;
+    value->num_moves = pos.legal_moves.size();
+    value->p.reset(pos.legal_moves.empty()
+        ? nullptr
+        : new float[pos.legal_moves.size()]);
+    memcache_->cache_.Insert(hash, std::move(value));
+
+    HashKeyedCacheLock<CachedValue> lock(&memcache_->cache_, hash);
+    // Sometimes search queries NN without passing the legal moves. It is
+    // still cached in this case, but in subsequent queries we only return it
+    // if legal moves are not passed again. Otherwise check the size to guard
+    // against hash collisions.
+    if (lock.holds_value() && (pos.legal_moves.empty() ||
+                               (lock->p && lock->num_moves == pos.legal_moves.size()))) {
+      value.reset();
+      auto state = lock->state.load(std::memory_order_acquire);
+      while (state == CachedValue::NOT_QUEUED) {
+        if (lock->state.compare_exchange_weak(state, CachedValue::NO_WAITERS,
+                                              std::memory_order_acq_rel)) {
+          to_be_queued = true;
+          break;
+        }
+      }
+      if (state == CachedValue::READY) {
         CachedValueToEvalResult(**lock, result);
         return AddInputResult::FETCHED_IMMEDIATELY;
       }
+      result_ptr = EvalResultPtr{
+        &lock->q, &lock->d, &lock->m,
+          lock->p ? std::span<float>{lock->p.get(), pos.legal_moves.size()}
+        : std::span<float>{}};
+    } else {
+      // No space, hash collision, or value was removed after insert.
+      lock = HashKeyedCacheLock<CachedValue>();  // release the lock
+      if (!value) {
+        value = std::make_unique<CachedValue>();
+        value->num_moves = pos.legal_moves.size();
+        value->p.reset(pos.legal_moves.empty()
+            ? nullptr
+            : new float[pos.legal_moves.size()]);
+      }
+      to_be_queued = true;
+      result_ptr = EvalResultPtr{
+          &value->q, &value->d, &value->m,
+          value->p ? std::span<float>{value->p.get(), pos.legal_moves.size()}
+                   : std::span<float>{}};
     }
-    size_t entry_idx = entries_.emplace_back(
-        Entry{hash, std::make_unique<CachedValue>(), result});
-    auto& value = entries_[entry_idx].value;
-    value->p.reset(pos.legal_moves.empty() ? nullptr
-                                           : new float[pos.legal_moves.size()]);
-    value->num_moves = pos.legal_moves.size();
-    return wrapped_computation_->AddInput(
-        pos, EvalResultPtr{&value->q, &value->d, &value->m,
-                           value->p ? std::span<float>{value->p.get(),
-                                                       pos.legal_moves.size()}
-                                    : std::span<float>{}});
+    entries_.emplace_back(
+        Entry{std::move(lock), std::move(value), result, to_be_queued});
+    if (!to_be_queued) {
+      // Another thread is already computing the value, we'll fetch it in
+      // ComputeBlocking.
+      return AddInputResult::FETCHED_DELAYED;
+    }
+    return wrapped_computation_->AddInput(pos, result_ptr);
   }
 
   virtual void ComputeBlocking() override {
-    if (wrapped_computation_->UsedBatchSize() == 0) return;
-    wrapped_computation_->ComputeBlocking();
+    if (wrapped_computation_->UsedBatchSize() != 0) {
+      wrapped_computation_->ComputeBlocking();
+    }
+    // Process results from our branch.
+    for (auto& entry : entries_) {
+      if (entry.queued_for_eval) {
+        if (entry.value) {
+          // There is no cache entry.
+          CachedValueToEvalResult(*entry.value, entry.result_ptr);
+        } else {
+          // There is a cache entry.
+          auto& lock = entry.lock;
+          assert(lock.holds_value());
+          CachedValueToEvalResult(**lock, entry.result_ptr);
+          auto state = lock->state.exchange(CachedValue::READY,
+                                            std::memory_order_release);
+          // Wake up waiters if there are any,
+          if (state == CachedValue::WAITERS) {
+            lock->state.notify_all();
+          }
+        }
+      }
+    }
+    // Process results from other batches which we got through cache before
+    // results were ready.
     for (auto& entry : entries_) {
-      CachedValueToEvalResult(*entry.value, entry.result_ptr);
-      memcache_->cache_.Insert(entry.key, std::move(entry.value));
+      if (!entry.queued_for_eval) {
+        auto& lock = entry.lock;
+        assert(lock.holds_value());
+        auto state = lock->state.load(std::memory_order_acquire);
+        // Make sure writing side knows about waiters
+        if (state == CachedValue::NO_WAITERS) {
+          lock->state.compare_exchange_strong(state, CachedValue::WAITERS,
+                                              std::memory_order_acquire);
+        }
+        // Wait until the value is ready.
+        lock->state.wait(CachedValue::WAITERS, std::memory_order_acquire);
+        assert(lock->state.load(std::memory_order_acquire) ==
+               CachedValue::READY);
+        CachedValueToEvalResult(**lock, entry.result_ptr);
+      }
     }
   }
 
   struct Entry {
-    uint64_t key;
+    HashKeyedCacheLock<CachedValue> lock;
     std::unique_ptr<CachedValue> value;
     EvalResultPtr result_ptr;
+    bool queued_for_eval = false;
   };
 
   std::unique_ptr<BackendComputation> wrapped_computation_;
@@ -168,6 +254,7 @@ std::optional<EvalResult> MemCache::GetCachedEvaluation(
   const uint64_t hash = ComputeEvalPositionHash(pos);
   HashKeyedCacheLock<CachedValue> lock(&cache_, hash);
   if (!lock.holds_value() ||
+      lock->state.load(std::memory_order_acquire) != CachedValue::READY ||
       (!pos.legal_moves.empty() &&
        !(lock->p && lock->num_moves == pos.legal_moves.size()))) {
     return std::nullopt;
@@ -186,9 +273,11 @@ std::optional<EvalResult> MemCache::GetCachedEvaluation(
 
 }  // namespace
 
-std::unique_ptr<CachingBackend> CreateMemCache(std::unique_ptr<Backend> wrapped,
-                                               const OptionsDict& options) {
-  return std::make_unique<MemCache>(std::move(wrapped), options);
+std::unique_ptr<CachingBackend> CreateMemCache(
+    std::unique_ptr<Backend> wrapped, const OptionsDict& options,
+    float max_out_of_order_evals_factor) {
+  return std::make_unique<MemCache>(std::move(wrapped), options,
+                                    max_out_of_order_evals_factor);
 }
 
 }  // namespace lczero

src/neural/memcache.h

@@ -42,6 +42,7 @@ class CachingBackend : public Backend {
 // are found, and forwards the request to the wrapped backend otherwise (and
 // caches the result).
 std::unique_ptr<CachingBackend> CreateMemCache(std::unique_ptr<Backend> parent,
-                                               const OptionsDict& options);
+                                               const OptionsDict& options,
+                                               const float out_of_order_eval_factor);
 
-}  // namespace lczero
+}  // namespace lczero

src/search/classic/search.cc

@@ -1411,6 +1411,9 @@ void SearchWorker::GatherMinibatch() {
           minibatch_.erase(minibatch_.begin() + i);
           --minibatch_size;
           ++number_out_of_order_;
+        } else if (minibatch_[i].is_delayed_cache_hit) {
+          --minibatch_size;
+          ++number_out_of_order_;
         }
       }
     }
@@ -1466,13 +1469,22 @@ void SearchWorker::ProcessPickedTask(int start_idx, int end_idx,
                        std::back_inserter(legal_moves),
                        [](const auto& edge) { return edge.GetMove(); });
         picked_node.eval->p.resize(legal_moves.size());
-        picked_node.is_cache_hit = computation_->AddInput(
-                                       EvalPosition{
-                                           .pos = history.GetPositions(),
-                                           .legal_moves = legal_moves,
-                                       },
-                                       picked_node.eval->AsPtr()) ==
-                                   BackendComputation::FETCHED_IMMEDIATELY;
+        auto cache_result = computation_->AddInput(
+            EvalPosition{
+                .pos = history.GetPositions(),
+                .legal_moves = legal_moves,
+            },
+            picked_node.eval->AsPtr());
+        switch (cache_result) {
+          case BackendComputation::ENQUEUED_FOR_EVAL:
+            break;
+          case BackendComputation::FETCHED_IMMEDIATELY:
+            picked_node.is_cache_hit = true;
+            break;
+          case BackendComputation::FETCHED_DELAYED:
+            picked_node.is_delayed_cache_hit = true;
+            break;
+        }
       }
     }
     if (params_.GetOutOfOrderEval() && picked_node.CanEvalOutOfOrder()) {
@@ -2292,7 +2304,8 @@ void SearchWorker::DoBackupUpdateSingleNode(
     }
   }
   search_->total_playouts_ += node_to_process.multivisit;
-  if (node_to_process.nn_queried && !node_to_process.is_cache_hit) {
+  if (node_to_process.nn_queried && !node_to_process.is_cache_hit &&
+      !node_to_process.is_delayed_cache_hit) {
     search_->network_evaluations_++;
   }
   search_->cum_depth_ += node_to_process.depth * node_to_process.multivisit;

src/search/classic/search.h

@@ -321,6 +321,7 @@ class SearchWorker {
     uint16_t depth;
     bool nn_queried = false;
     bool is_cache_hit = false;
+    bool is_delayed_cache_hit = false;
     bool is_collision = false;
     // Only populated for visits,
     std::vector<Move> moves_to_visit;

src/search/classic/wrapper.cc

@@ -72,6 +72,10 @@ class ClassicSearch : public SearchBase {
     if (search_) search_->Abort();
   }
 
+  float GetMaxOutOfOrderFactor() const override {
+    return options_->Get<float>(BaseSearchParams::kMaxOutOfOrderEvalsFactorId);
+  }
+
   const OptionsDict* options_;
   std::unique_ptr<TimeManager> time_manager_;
   std::unique_ptr<Search> search_;

src/search/dag_classic/search.cc

@@ -1488,6 +1488,9 @@ void SearchWorker::GatherMinibatch() {
           minibatch_.erase(minibatch_.begin() + i);
           --minibatch_size;
           ++number_out_of_order_;
+        } else if (minibatch_[i].is_delayed_cache_hit) {
+          --minibatch_size;
+          ++number_out_of_order_;
         }
       }
     }
@@ -2105,13 +2108,22 @@ void SearchWorker::ExtendNode(NodeToProcess& picked_node) {
     picked_node.tt_low_node = std::make_shared<LowNode>(legal_moves);
     picked_node.nn_queried = true;
     picked_node.eval->p.resize(legal_moves.size());
-    picked_node.is_cache_hit = computation_->AddInput(
-                                   EvalPosition{
-                                       .pos = history.GetPositions(),
-                                       .legal_moves = legal_moves,
-                                   },
-                                   picked_node.eval->AsPtr()) ==
-                               BackendComputation::FETCHED_IMMEDIATELY;
+    auto cache_result = computation_->AddInput(
+        EvalPosition{
+            .pos = history.GetPositions(),
+            .legal_moves = legal_moves,
+        },
+        picked_node.eval->AsPtr());
+    switch (cache_result) {
+      case BackendComputation::ENQUEUED_FOR_EVAL:
+        break;
+      case BackendComputation::FETCHED_IMMEDIATELY:
+        picked_node.is_cache_hit = true;
+        break;
+      case BackendComputation::FETCHED_DELAYED:
+        picked_node.is_delayed_cache_hit = true;
+        break;
+    }
   }
 }
 
@@ -2380,7 +2392,8 @@ void SearchWorker::DoBackupUpdateSingleNode(
     nm = pm;
   }
   search_->total_playouts_ += node_to_process.multivisit;
-  if (node_to_process.nn_queried && !node_to_process.is_cache_hit) {
+  if (node_to_process.nn_queried && !node_to_process.is_cache_hit &&
+      !node_to_process.is_delayed_cache_hit) {
     search_->network_evaluations_++;
   }
   search_->cum_depth_ +=

src/search/dag_classic/search.h

@@ -332,6 +332,7 @@ class SearchWorker {
     bool nn_queried = false;
     bool is_tt_hit = false;
     bool is_cache_hit = false;
+    bool is_delayed_cache_hit = false;
     bool is_collision = false;
 
     // Details that are filled in as we go.

src/search/dag_classic/wrapper.cc

@@ -74,6 +74,11 @@ class DagClassicSearch : public SearchBase {
     if (search_) search_->Abort();
   }
 
+  float GetMaxOutOfOrderFactor() const override {
+    return options_->Get<float>(
+        classic::BaseSearchParams::kMaxOutOfOrderEvalsFactorId);
+  }
+
   const OptionsDict* options_;
   std::unique_ptr<classic::TimeManager> time_manager_;
   std::unique_ptr<Search> search_;