From cde884d7a992940e93007d206f9f5c16ff89b2e1 Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@gmail.com>
Date: Fri, 3 Nov 2023 21:47:48 -0400
Subject: [PATCH] Fixing threading

---
 cpp/bench/ann/src/common/benchmark.hpp                   | 9 ++++++---
 .../src/raft-ann-bench/run/conf/datasets.yaml            | 6 ++++++
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/cpp/bench/ann/src/common/benchmark.hpp b/cpp/bench/ann/src/common/benchmark.hpp
index 7a9f0b99f8..9ddb232f5c 100644
--- a/cpp/bench/ann/src/common/benchmark.hpp
+++ b/cpp/bench/ann/src/common/benchmark.hpp
@@ -198,7 +198,7 @@ void bench_search(::benchmark::State& state,
    * Make sure the first thread loads the algo and dataset
    */
   if (state.thread_index() == 0) {
-    std::lock_guard lk(init_mutex);
+    std::unique_lock lk(init_mutex);
     // algo is static to cache it between close search runs to save time on index loading
     static std::string index_file = "";
     if (index.file != index_file) {
@@ -249,9 +249,11 @@ void bench_search(::benchmark::State& state,
     query_set = dataset->query_set(current_algo_props->query_memory_type);
     cond_var.notify_all();
   } else {
-    // All other threads will wait for the first thread to initialize the algo.
     std::unique_lock lk(init_mutex);
-    cond_var.wait(lk, [] { return current_algo_props.get() != nullptr; });
+    // All other threads will wait for the first thread to initialize the algo.
+
+    cond_var.wait(
+      lk, [] { return current_algo_props.get() != nullptr && current_algo.get() != nullptr; });
     // gbench ensures that all threads are synchronized at the start of the benchmark loop.
     // We are accessing shared variables (like current_algo, current_algo_probs) before the
     // benchmark loop, therefore the synchronization here is necessary.
@@ -292,6 +294,7 @@ void bench_search(::benchmark::State& state,
 
       // advance to the next batch
       batch_offset = (batch_offset + n_queries) % query_set_size;
+
       queries_processed += n_queries;
     }
   }
diff --git a/python/raft-ann-bench/src/raft-ann-bench/run/conf/datasets.yaml b/python/raft-ann-bench/src/raft-ann-bench/run/conf/datasets.yaml
index 11a91f91fd..a2b948a464 100644
--- a/python/raft-ann-bench/src/raft-ann-bench/run/conf/datasets.yaml
+++ b/python/raft-ann-bench/src/raft-ann-bench/run/conf/datasets.yaml
@@ -39,36 +39,42 @@
   dims: 960
   base_file: gist-960-euclidean/base.fbin
   query_file: gist-960-euclidean/query.fbin
+  groundtruth_neighbors_file: gist-960-euclidean/groundtruth.neighbors.ibin
   distance: euclidean
 
 - name: glove-50-angular
   dims: 50
   base_file: glove-50-angular/base.fbin
   query_file: glove-50-angular/query.fbin
+  groundtruth_neighbors_file: glove-50-angular/groundtruth.neighbors.ibin
   distance: euclidean
 
 - name: glove-50-inner
   dims: 50
   base_file: glove-50-inner/base.fbin
   query_file: glove-50-inner/query.fbin
+  groundtruth_neighbors_file: glove-50-inner/groundtruth.neighbors.ibin
   distance: euclidean
 
 - name: glove-100-angular
   dims: 100
   base_file: glove-100-angular/base.fbin
   query_file: glove-100-angular/query.fbin
+  groundtruth_neighbors_file: glove-100-angular/groundtruth.neighbors.ibin
   distance: euclidean
 
 - name: glove-100-inner
   dims: 100
   base_file: glove-100-inner/base.fbin
   query_file: glove-100-inner/query.fbin
+  groundtruth_neighbors_file: glove-100-inner/groundtruth.neighbors.ibin
   distance: euclidean
 
 - name: lastfm-65-angular
   dims: 65
   base_file: lastfm-65-angular/base.fbin
   query_file: lastfm-65-angular/query.fbin
+  groundtruth_neighbors_file: lastfm-65-angular/groundtruth.neighbors.ibin
   distance: euclidean
 
 - name: mnist-784-euclidean