add beam_width, rm dbg statements

rapidsai · Oct 15, 2024 · df54939 · df54939
1 parent a13bf1a
commit df54939
Show file tree

Hide file tree

Showing 6 changed files with 18 additions and 15 deletions.
diff --git a/cpp/bench/ann/src/common/benchmark.hpp b/cpp/bench/ann/src/common/benchmark.hpp
@@ -241,8 +241,6 @@ void bench_search(::benchmark::State& state,
     return;
   }
 
-  std::cout << "state.threads" << state.threads() << std::endl;
-
   /**
    * Make sure the first thread loads the algo and dataset
    */

diff --git a/cpp/bench/ann/src/diskann/diskann_benchmark.cpp b/cpp/bench/ann/src/diskann/diskann_benchmark.cpp
@@ -68,6 +68,9 @@ void parse_search_param(const nlohmann::json& conf,
   if (conf.contains("num_nodes_to_cache")) {
     param.num_nodes_to_cache = conf.at("num_nodes_to_cache");
   }
+  if (conf.contains("beam_width")) {
+    param.beam_width = conf.at("beam_width");
+  }
 }
 
 template <typename T, template <typename> class Algo>

diff --git a/cpp/bench/ann/src/diskann/diskann_wrapper.h b/cpp/bench/ann/src/diskann/diskann_wrapper.h
@@ -210,13 +210,14 @@ class diskann_ssd : public algo<T> {
     uint32_t L_search;
     uint32_t num_threads        = omp_get_num_procs();
     uint32_t num_nodes_to_cache = 10000;
+    int beam_width              = 2;
     // Mode metric_objective;
   };
 
   diskann_ssd(Metric metric, int dim, const build_param& param);
 
   void build_from_bin(std::string dataset_path, std::string path_to_index, size_t nrow) override;
-  void build(const T* dataset, size_t nrow) override{
+  void build(const T* dataset, size_t nrow) override {
     // do nothing. will not be used.
   };
 
@@ -244,14 +245,14 @@ class diskann_ssd : public algo<T> {
  private:
   std::string index_build_params_str;
   std::shared_ptr<diskann::PQFlashIndex<T, uint32_t>> p_flash_index_;
+  int beam_width_;
+  uint32_t num_nodes_to_cache_;
 
-  uint32_t num_nodes_to_cache_ = 10000;
-  uint32_t num_search_threads_;
-
-  // mem index params
-  uint32_t max_points_;
+  // in-memory index params
   uint32_t build_pq_bytes_ = 0;
+  uint32_t max_points_;
   int num_threads_;
+  int num_search_threads_;
   uint32_t L_search_;
   Mode bench_mode_;
   std::shared_ptr<fixed_thread_pool> thread_pool_;
@@ -299,6 +300,7 @@ void diskann_ssd<T>::set_search_param(const search_param_base& param_)
   L_search_           = param.L_search;
   num_search_threads_ = param.num_threads;
   num_nodes_to_cache_ = param.num_nodes_to_cache;
+  beam_width_         = param.beam_width;
 
   // only latency mode supported with thread pool
   bench_mode_ = Mode::kLatency;
@@ -315,10 +317,10 @@ void diskann_ssd<T>::search(
     // diskann ssd index can only handle a single vector at a time.
     p_flash_index_->cached_beam_search(queries + (i * this->dim_),
                                        static_cast<size_t>(k),
-                                       this->L_search_,
+                                       L_search_,
                                        reinterpret_cast<uint64_t*>(neighbors + i * k),
                                        distances + i * k,
-                                       1,
+                                       beam_width_,
                                        false,
                                        nullptr);
   };

diff --git a/python/cuvs_bench/cuvs_bench/config/algos/cuvs_vamana.yaml b/python/cuvs_bench/cuvs_bench/config/algos/cuvs_vamana.yaml
@@ -3,8 +3,8 @@ groups:
   base:
     build:
       graph_degree: [64, 96]
-      visited_size: [100, 200, 300]
+      visited_size: [128, 256, 512]
       alpha: [1.2]
     search:
       L_search: [10, 20, 30, 40, 50, 100, 200, 300]
-      num_threads: [32]
+      num_threads: [32]
diff --git a/python/cuvs_bench/cuvs_bench/config/algos/diskann_memory.yaml b/python/cuvs_bench/cuvs_bench/config/algos/diskann_memory.yaml
@@ -5,7 +5,7 @@ groups:
   base:
     build:
       R: [64, 96]
-      L_build: [100, 200, 300]
+      L_build: [128, 256, 384]
       alpha: [1.2]
       num_threads: [32]
     search:

diff --git a/python/cuvs_bench/cuvs_bench/config/algos/diskann_ssd.yaml b/python/cuvs_bench/cuvs_bench/config/algos/diskann_ssd.yaml
@@ -5,9 +5,9 @@ groups:
   base:
     build:
       R: [64, 96]
-      L_build: [100, 200, 300]
+      L_build: [128, 256, 384]
       QD: [192]
       num_threads: [32]
     search:
       L_search: [10, 20, 30, 40, 50, 100, 200, 300]
-      num_threads: [32]
+      num_threads: [32]