Generate diskann cache asynchronously.

knowhere/index/vector_index/IndexDiskANN.cpp

@@ -290,9 +290,9 @@ IndexDiskANN<T>::Prepare(const Config& config) {
         KNOWHERE_THROW_MSG("Failed to generate cache, num_nodes_to_cache is larger than 1/3 of the total data number.");
     }
     if (num_nodes_to_cache > 0) {
-        std::vector<uint32_t> node_list;
         LOG_KNOWHERE_INFO_ << "Caching " << num_nodes_to_cache << " sample nodes around medoid(s).";
         if (prep_conf.use_bfs_cache) {
+            std::vector<uint32_t> node_list;
             auto gen_cache_successful = TryDiskANNCall<bool>([&]() -> bool {
                 pq_flash_index_->cache_bfs_levels(num_nodes_to_cache, node_list);
                 return true;
@@ -302,29 +302,28 @@ IndexDiskANN<T>::Prepare(const Config& config) {
                 LOG_KNOWHERE_ERROR_ << "Failed to generate bfs cache for DiskANN.";
                 return false;
             }
-        } else {
-            auto gen_cache_successful = TryDiskANNCall<bool>([&]() -> bool {
-                pq_flash_index_->generate_cache_list_from_sample_queries(warmup_query_file, 15, 6, num_nodes_to_cache,
-                                                                         prep_conf.num_threads, node_list);
+            auto load_cache_successful = TryDiskANNCall<bool>([&]() -> bool {
+                pq_flash_index_->load_cache_list(node_list);
                 return true;
             });
 
-            if (!gen_cache_successful.has_value()) {
-                LOG_KNOWHERE_ERROR_ << "Failed to generate cache from sample queries for DiskANN.";
+            if (!load_cache_successful.has_value()) {
+                LOG_KNOWHERE_ERROR_ << "Failed to load cache for DiskANN.";
                 return false;
             }
-        }
-        auto load_cache_successful = TryDiskANNCall<bool>([&]() -> bool {
-            pq_flash_index_->load_cache_list(node_list);
-            return true;
-        });
-
-        if (!load_cache_successful.has_value()) {
-            LOG_KNOWHERE_ERROR_ << "Failed to load cache for DiskANN.";
-            return false;
+        } else {
+            pq_flash_index_->set_async_cache_flag(true);
+            pool_->push([&, cache_num = num_nodes_to_cache,
+                        sample_nodes_file = warmup_query_file]() {
+                try {
+                    pq_flash_index_->generate_cache_list_from_sample_queries(
+                        sample_nodes_file, 15, 6, cache_num);
+                } catch (const std::exception& e) {
+                    LOG_KNOWHERE_ERROR_ << "DiskANN Exception: " << e.what();
+                }
+            });
         }
     }
-
     // warmup
     if (prep_conf.warm_up) {
         LOG_KNOWHERE_DEBUG_ << "Warming up.";

thirdparty/DiskANN/include/pq_flash_index.h

@@ -22,6 +22,7 @@
 #include "percentile_stats.h"
 #include "pq_table.h"
 #include "utils.h"
+#include "semaphore.h"
 #include "windows_customizations.h"
 
 #define MAX_GRAPH_DEGREE 512
@@ -80,16 +81,15 @@ namespace diskann {
 
     DISKANN_DLLEXPORT void load_cache_list(std::vector<uint32_t> &node_list);
 
+    // asynchronously collect the access frequency of each node in the graph
 #ifdef EXEC_ENV_OLS
     DISKANN_DLLEXPORT void generate_cache_list_from_sample_queries(
-        MemoryMappedFiles &files, std::string sample_bin, _u64 l_search,
-        _u64 beamwidth, _u64 num_nodes_to_cache, uint32_t nthreads,
-        std::vector<uint32_t> &node_list);
+        MemoryMappedFiles files, std::string sample_bin, _u64 l_search,
+        _u64 beamwidth, _u64 num_nodes_to_cache);
 #else
     DISKANN_DLLEXPORT void generate_cache_list_from_sample_queries(
         std::string sample_bin, _u64 l_search, _u64 beamwidth,
-        _u64 num_nodes_to_cache, uint32_t num_threads,
-        std::vector<uint32_t> &node_list);
+        _u64 num_nodes_to_cache);
 #endif
 
     DISKANN_DLLEXPORT void cache_bfs_levels(_u64 num_nodes_to_cache,
@@ -128,6 +128,8 @@ namespace diskann {
 
     DISKANN_DLLEXPORT diskann::Metric get_metric() const noexcept;
 
+    DISKANN_DLLEXPORT void set_async_cache_flag(const bool flag);
+
    protected:
     DISKANN_DLLEXPORT void use_medoids_data_as_centroids();
     DISKANN_DLLEXPORT void setup_thread_data(_u64 nthreads);
@@ -195,6 +197,7 @@ namespace diskann {
 
     std::string                        disk_index_file;
     std::vector<std::pair<_u32, _u32>> node_visit_counter;
+    std::atomic<_u32>                  search_counter = 0;
 
     // PQ data
     // n_chunks = # of chunks ndims is split into
@@ -233,12 +236,14 @@ namespace diskann {
     // coord_cache
     T *                       coord_cache_buf = nullptr;
     tsl::robin_map<_u32, T *> coord_cache;
+    Semaphore                 semaph;
+    std::atomic<bool>         async_generate_cache = false;
 
     // thread-specific scratch
     ConcurrentQueue<ThreadData<T>> thread_data;
     _u64                           max_nthreads;
     bool                           load_flag = false;
-    bool                           count_visited_nodes = false;
+    std::atomic<bool>              count_visited_nodes = false;
     bool                           reorder_data_exists = false;
     _u64                           reoreder_data_offset = 0;
 

thirdparty/DiskANN/include/semaphore.h

@@ -0,0 +1,35 @@
+#pragma once
+#include <mutex>
+#include <condition_variable>
+
+namespace diskann {
+class Semaphore {
+public:
+    Semaphore(long count = 0) : count(count) {}
+    void Signal()
+    {
+        std::unique_lock<std::mutex> unique(mt);
+        ++count;
+        if (count <= 0) {
+            cond.notify_one();
+        }
+    }
+    void Wait()
+    {
+        std::unique_lock<std::mutex> unique(mt);
+        --count;
+        if (count < 0) {
+            cond.wait(unique);
+        }
+    }
+    bool IsWaitting() {
+        std::unique_lock<std::mutex> unique(mt);
+        return count < 0;
+    }
+
+private:
+    std::mutex mt;
+    std::condition_variable cond;
+    long count;
+};
+}  // namespace diskann

thirdparty/DiskANN/src/pq_flash_index.cpp

@@ -14,6 +14,7 @@
 #include <iterator>
 #include <random>
 #include <thread>
+#include <mutex>
 #include "distance.h"
 #include "exceptions.h"
 #include "parameters.h"
@@ -98,7 +99,7 @@ namespace diskann {
   template<typename T>
   PQFlashIndex<T>::PQFlashIndex(
       std::shared_ptr<AlignedFileReader> fileReader, diskann::Metric m)
-      : reader(fileReader), metric(m) {
+      : reader(fileReader), metric(m), semaph(0) {
     if (m == diskann::Metric::COSINE || m == diskann::Metric::INNER_PRODUCT) {
       if (std::is_floating_point<T>::value) {
         LOG(INFO) << "Cosine metric chosen for (normalized) float data."
@@ -119,6 +120,9 @@ namespace diskann {
 
   template<typename T>
   PQFlashIndex<T>::~PQFlashIndex() {
+    if (this->async_generate_cache) {
+      this->semaph.Wait();
+    }
 #ifndef EXEC_ENV_OLS
     if (data != nullptr) {
       delete[] data;
@@ -216,6 +220,8 @@ namespace diskann {
   template<typename T>
   void PQFlashIndex<T>::load_cache_list(std::vector<uint32_t> &node_list) {
     LOG(DEBUG) << "Loading the cache list into memory...";
+    assert(this->nhood_cache_buf == nullptr && "nhoodc_cache_buf is not null");
+    assert(this->coord_cache_buf == nullptr && "coord_cache_buf is not null");
     _u64 num_cached_nodes = node_list.size();
 
     // borrow thread data
@@ -293,20 +299,20 @@ namespace diskann {
 #ifdef EXEC_ENV_OLS
   template<typename T>
   void PQFlashIndex<T>::generate_cache_list_from_sample_queries(
-      MemoryMappedFiles &files, std::string sample_bin, _u64 l_search,
-      _u64 beamwidth, _u64 num_nodes_to_cache, uint32_t nthreads,
-      std::vector<uint32_t> &node_list) {
+      MemoryMappedFiles files, std::string sample_bin, _u64 l_search,
+      _u64 beamwidth, _u64 num_nodes_to_cache) {
 #else
   template<typename T>
   void PQFlashIndex<T>::generate_cache_list_from_sample_queries(
       std::string sample_bin, _u64 l_search, _u64 beamwidth,
-      _u64 num_nodes_to_cache, uint32_t nthreads,
-      std::vector<uint32_t> &node_list) {
+      _u64 num_nodes_to_cache) {
 #endif
     auto s = std::chrono::high_resolution_clock::now();
-    this->count_visited_nodes = true;
+    this->search_counter.store(0);
     this->node_visit_counter.clear();
     this->node_visit_counter.resize(this->num_points);
+    this->count_visited_nodes.store(true);
+
     for (_u32 i = 0; i < node_visit_counter.size(); i++) {
       this->node_visit_counter[i].first = i;
       this->node_visit_counter[i].second = 0;
@@ -332,32 +338,47 @@ namespace diskann {
       return;
     }
 
-    std::vector<int64_t> tmp_result_ids_64(sample_num, 0);
-    std::vector<float>    tmp_result_dists(sample_num, 0);
+    int64_t tmp_result_ids_64;
+    float tmp_result_dists;
 
-#pragma omp parallel for schedule(dynamic, 1) num_threads(nthreads)
-    for (_s64 i = 0; i < (int64_t) sample_num; i++) {
-      cached_beam_search(samples + (i * sample_aligned_dim), 1, l_search,
-                         tmp_result_ids_64.data() + (i * 1),
-                         tmp_result_dists.data() + (i * 1), beamwidth);
+    auto id = 0;
+    while (this->search_counter.load() < sample_num && id < sample_num &&
+          !this->semaph.IsWaitting()) {
+      cached_beam_search(samples + (id * sample_aligned_dim), 1, l_search,
+                        &tmp_result_ids_64, &tmp_result_dists, beamwidth);
+      id++;
     }
 
+    if (this->semaph.IsWaitting()) {
+      this->semaph.Signal();
+      return;
+    }
+
+    this->count_visited_nodes.store(false);
     std::sort(this->node_visit_counter.begin(), node_visit_counter.end(),
               [](std::pair<_u32, _u32> &left, std::pair<_u32, _u32> &right) {
                 return left.second > right.second;
               });
+
+    std::vector<uint32_t> node_list;
     node_list.clear();
     node_list.shrink_to_fit();
     node_list.reserve(num_nodes_to_cache);
     for (_u64 i = 0; i < num_nodes_to_cache; i++) {
       node_list.push_back(this->node_visit_counter[i].first);
     }
-    this->count_visited_nodes = false;
+    this->node_visit_counter.clear();
+    this->node_visit_counter.shrink_to_fit();
+    this->search_counter.store(0);
 
     diskann::aligned_free(samples);
+    this->load_cache_list(node_list);
     auto e = std::chrono::high_resolution_clock::now();
     std::chrono::duration<double> diff = e - s;
     LOG(INFO) << "Using sample queries to generate cache, cost: " << diff.count() << "s";
+
+    this->semaph.Signal();
+    return;
   }
 
   template<typename T>
@@ -1464,6 +1485,9 @@ namespace diskann {
     if (stats != nullptr) {
       stats->total_us = (double) query_timer.elapsed();
     }
+    if (this->count_visited_nodes) {
+      this->search_counter.fetch_add(1);
+    }
   }
 
   // range search returns results of all neighbors within distance of range.
@@ -1548,6 +1572,11 @@ namespace diskann {
   diskann::Metric PQFlashIndex<T>::get_metric() const noexcept {
     return metric;
   }
+
+  template<typename T>
+  void PQFlashIndex<T>::set_async_cache_flag(const bool flag) {
+    this->async_generate_cache.exchange(flag);
+  }
 
 #ifdef EXEC_ENV_OLS
   template<typename T>