[FEA] New algos and updates/corrections to Faiss cuvs-bench

cpp/bench/ann/CMakeLists.txt

@@ -21,11 +21,15 @@ option(CUVS_ANN_BENCH_USE_FAISS_GPU_FLAT "Include faiss' brute-force knn algorit
 option(CUVS_ANN_BENCH_USE_FAISS_GPU_IVF_FLAT "Include faiss' ivf flat algorithm in benchmark" ON)
 option(CUVS_ANN_BENCH_USE_FAISS_GPU_IVF_PQ "Include faiss' ivf pq algorithm in benchmark" ON)
 option(CUVS_ANN_BENCH_USE_FAISS_GPU_CAGRA "Include faiss' cagra algorithm in benchmark" ON)
+option(CUVS_ANN_BENCH_USE_FAISS_GPU_CAGRA_HNSW
+       "Include faiss' cagra algorithm for build and hnsw for search in benchmark" ON
+)
 option(CUVS_ANN_BENCH_USE_FAISS_CPU_FLAT "Include faiss' cpu brute-force algorithm in benchmark" ON)
 option(CUVS_ANN_BENCH_USE_FAISS_CPU_IVF_FLAT "Include faiss' cpu ivf flat algorithm in benchmark"
        ON
 )
 option(CUVS_ANN_BENCH_USE_FAISS_CPU_IVF_PQ "Include faiss' cpu ivf pq algorithm in benchmark" ON)
+option(CUVS_ANN_BENCH_USE_FAISS_CPU_HNSW_FLAT "Include faiss' hnsw algorithm in benchmark" ON)
 option(CUVS_ANN_BENCH_USE_CUVS_IVF_FLAT "Include cuVS ivf flat algorithm in benchmark" ON)
 option(CUVS_ANN_BENCH_USE_CUVS_IVF_PQ "Include cuVS ivf pq algorithm in benchmark" ON)
 option(CUVS_ANN_BENCH_USE_CUVS_CAGRA "Include cuVS CAGRA in benchmark" ON)
@@ -276,6 +280,12 @@ if(CUVS_ANN_BENCH_USE_FAISS_CPU_IVF_PQ)
   )
 endif()
 
+if(CUVS_ANN_BENCH_USE_FAISS_CPU_HNSW_FLAT)
+  ConfigureAnnBench(
+    NAME FAISS_CPU_HNSW_FLAT PATH src/faiss/faiss_cpu_benchmark.cpp LINKS ${CUVS_FAISS_TARGETS}
+  )
+endif()
+
 if(CUVS_ANN_BENCH_USE_FAISS_GPU_IVF_FLAT AND CUVS_FAISS_ENABLE_GPU)
   ConfigureAnnBench(
     NAME FAISS_GPU_IVF_FLAT PATH src/faiss/faiss_gpu_benchmark.cu LINKS ${CUVS_FAISS_TARGETS}
@@ -304,6 +314,13 @@ if(CUVS_ANN_BENCH_USE_FAISS_GPU_CAGRA AND CUVS_FAISS_ENABLE_GPU)
   )
 endif()
 
+if(CUVS_ANN_BENCH_USE_FAISS_GPU_CAGRA_HNSW AND CUVS_FAISS_ENABLE_GPU)
+  ConfigureAnnBench(
+    NAME FAISS_GPU_CAGRA_HNSW PATH src/faiss/faiss_gpu_benchmark.cu LINKS ${CUVS_FAISS_TARGETS}
+    raft::raft
+  )
+endif()
+
 if(CUVS_ANN_BENCH_USE_GGNN)
   include(cmake/thirdparty/get_glog)
   ConfigureAnnBench(
@@ -332,8 +349,9 @@ if(CUVS_ANN_BENCH_SINGLE_EXE)
   target_include_directories(ANN_BENCH PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
 
   target_link_libraries(
-    ANN_BENCH PRIVATE raft::raft nlohmann_json::nlohmann_json benchmark::benchmark dl
-                      $<$<TARGET_EXISTS:CUDA::nvtx3>:CUDA::nvtx3>
+    ANN_BENCH
+    PRIVATE raft::raft nlohmann_json::nlohmann_json benchmark::benchmark dl fmt::fmt-header-only
+            spdlog::spdlog_header_only $<$<TARGET_EXISTS:CUDA::nvtx3>:CUDA::nvtx3>
   )
   set_target_properties(
     ANN_BENCH

cpp/bench/ann/src/faiss/faiss_cpu_benchmark.cpp

@@ -68,6 +68,18 @@ void parse_build_param(const nlohmann::json& conf,
   param.quantizer_type = conf.at("quantizer_type");
 }
 
+template <typename T>
+void parse_build_param(const nlohmann::json& conf,
+                       typename cuvs::bench::faiss_cpu_hnsw_flat<T>::build_param& param)
+{
+  param.M = conf.at("M");
+  if (conf.contains("efConstruction")) {
+    param.efConstruction = conf.at("efConstruction");
+  } else {
+    param.efConstruction = 40;
+  }
+}
+
 template <typename T>
 void parse_search_param(const nlohmann::json& conf,
                         typename cuvs::bench::faiss_cpu<T>::search_param& param)
@@ -77,6 +89,19 @@ void parse_search_param(const nlohmann::json& conf,
   if (conf.contains("numThreads")) { param.num_threads = conf.at("numThreads"); }
 }
 
+template <typename T>
+void parse_search_param(const nlohmann::json& conf,
+                        typename cuvs::bench::faiss_cpu_hnsw_flat<T>::search_param& param)
+{
+  faiss::SearchParametersHNSW p;
+  if (conf.contains("efSearch")) { p.efSearch = conf.at("efSearch"); }
+  if (conf.contains("check_relative_distance")) {
+    p.check_relative_distance = conf.at("check_relative_distance");
+  }
+  if (conf.contains("bounded_queue")) { p.bounded_queue = conf.at("bounded_queue"); }
+  param.p = p;
+}
+
 template <typename T, template <typename> class Algo>
 auto make_algo(cuvs::bench::Metric metric, int dim, const nlohmann::json& conf)
   -> std::unique_ptr<cuvs::bench::algo<T>>
@@ -104,6 +129,8 @@ auto create_algo(const std::string& algo_name,
       a = make_algo<T, cuvs::bench::faiss_cpu_ivfsq>(metric, dim, conf);
     } else if (algo_name == "faiss_cpu_flat") {
       a = std::make_unique<cuvs::bench::faiss_cpu_flat<T>>(metric, dim);
+    } else if (algo_name == "faiss_cpu_hnsw_flat") {
+      a = make_algo<T, cuvs::bench::faiss_cpu_hnsw_flat>(metric, dim, conf);
     }
   }
 
@@ -126,6 +153,10 @@ auto create_search_param(const std::string& algo_name, const nlohmann::json& con
   } else if (algo_name == "faiss_cpu_flat") {
     auto param = std::make_unique<typename cuvs::bench::faiss_cpu<T>::search_param>();
     return param;
+  } else if (algo_name == "faiss_cpu_hnsw_flat") {
+    auto param = std::make_unique<typename cuvs::bench::faiss_cpu_hnsw_flat<T>::search_param>();
+    parse_search_param<T>(conf, *param);
+    return param;
   }
   // else
   throw std::runtime_error("invalid algo: '" + algo_name + "'");

cpp/bench/ann/src/faiss/faiss_cpu_wrapper.h

@@ -20,6 +20,7 @@
 #include "../common/util.hpp"
 
 #include <faiss/IndexFlat.h>
+#include <faiss/IndexHNSW.h>
 #include <faiss/IndexIVFFlat.h>
 #include <faiss/IndexIVFPQ.h>
 #include <faiss/IndexRefine.h>
@@ -88,11 +89,11 @@ class faiss_cpu : public algo<T> {
 
   // TODO(snanditale): if the number of results is less than k, the remaining elements of
   // 'neighbors' will be filled with (size_t)-1
-  void search(const T* queries,
-              int batch_size,
-              int k,
-              algo_base::index_type* neighbors,
-              float* distances) const final;
+  virtual void search(const T* queries,
+                      int batch_size,
+                      int k,
+                      algo_base::index_type* neighbors,
+                      float* distances) const;
 
   [[nodiscard]] auto get_preference() const -> algo_property override
   {
@@ -146,6 +147,7 @@ void faiss_cpu<T>::build(const T* dataset, size_t nrow)
     index_ivf->cp.max_points_per_centroid = max_ppc;
     index_ivf->cp.min_points_per_centroid = min_ppc;
   }
+  faiss::IndexHNSWFlat* hnsw_index = dynamic_cast<faiss::IndexHNSWFlat*>(index_.get());
   index_->train(nrow, dataset);  // faiss::IndexFlat::train() will do nothing
   assert(index_->is_trained);
   index_->add(nrow, dataset);
@@ -176,12 +178,7 @@ void faiss_cpu<T>::search(
   static_assert(sizeof(size_t) == sizeof(faiss::idx_t),
                 "sizes of size_t and faiss::idx_t are different");
 
-  thread_pool_->submit(
-    [&](int i) {
-      // Use thread pool for batch size = 1. FAISS multi-threads internally for batch size > 1.
-      index_->search(batch_size, queries, k, distances, reinterpret_cast<faiss::idx_t*>(neighbors));
-    },
-    1);
+  index_->search(batch_size, queries, k, distances, reinterpret_cast<faiss::idx_t*>(neighbors));
 }
 
 template <typename T>
@@ -327,4 +324,62 @@ class faiss_cpu_flat : public faiss_cpu<T> {
   }
 };
 
+template <typename T>
+class faiss_cpu_hnsw_flat : public faiss_cpu<T> {
+ public:
+  struct build_param : public faiss_cpu<T>::build_param {
+    int M;
+    int efConstruction;
+  };
+  struct search_param : public faiss_cpu<T>::search_param {
+    faiss::SearchParametersHNSW p;
+  };
+  faiss_cpu_hnsw_flat(Metric metric, int dim, const build_param& param)
+    : faiss_cpu<T>(metric, dim, param)
+  {
+    this->index_ = std::make_shared<faiss::IndexHNSWFlat>(dim, param.M, this->metric_type_);
+    faiss::IndexHNSWFlat* hnsw_index = static_cast<faiss::IndexHNSWFlat*>(this->index_.get());
+    hnsw_index->hnsw.efConstruction  = param.efConstruction;
+  }
+
+  void set_search_param(const typename algo<T>::search_param& param,
+                        const void* filter_bitset) override
+  {
+    if (filter_bitset != nullptr) { throw std::runtime_error("Filtering is not supported yet."); }
+    auto sp              = static_cast<const typename faiss_cpu_hnsw_flat<T>::search_param&>(param);
+    this->search_params_ = std::make_shared<faiss::SearchParametersHNSW>(sp.p);
+  };
+
+  void save(const std::string& file) const override
+  {
+    this->template save_<faiss::IndexHNSWFlat>(file);
+  }
+  void load(const std::string& file) override { this->template load_<faiss::IndexHNSWFlat>(file); }
+
+  std::unique_ptr<algo<T>> copy()
+  {
+    return std::make_unique<faiss_cpu_hnsw_flat<T>>(*this);  // use copy constructor
+  }
+
+  void search(const T* queries,
+              int batch_size,
+              int k,
+              algo_base::index_type* neighbors,
+              float* distances) const override
+  {
+    static_assert(sizeof(size_t) == sizeof(faiss::idx_t),
+                  "sizes of size_t and faiss::idx_t are different");
+
+    this->index_->search(batch_size,
+                         queries,
+                         k,
+                         distances,
+                         reinterpret_cast<faiss::idx_t*>(neighbors),
+                         search_params_.get());
+  }
+
+ private:
+  std::shared_ptr<faiss::SearchParameters> search_params_;
+};
+
 }  // namespace cuvs::bench

cpp/bench/ann/src/faiss/faiss_gpu_benchmark.cu

@@ -27,6 +27,22 @@
 #include <type_traits>
 #include <utility>
 
+namespace {
+nlohmann::json collect_conf_with_prefix(const nlohmann::json& conf,
+                                        const std::string& prefix,
+                                        bool remove_prefix = true)
+{
+  nlohmann::json out;
+  for (auto& i : conf.items()) {
+    if (i.key().compare(0, prefix.size(), prefix) == 0) {
+      auto new_key = remove_prefix ? i.key().substr(prefix.size()) : i.key();
+      out[new_key] = i.value();
+    }
+  }
+  return out;
+}
+}  // namespace
+
 namespace cuvs::bench {
 
 template <typename T>
@@ -100,6 +116,82 @@ void parse_build_param(const nlohmann::json& conf,
     param.intermediate_graph_degree = 128;
   }
   if (conf.contains("cagra_build_algo")) { param.cagra_build_algo = conf.at("cagra_build_algo"); }
+  if (conf.contains("nn_descent_niter")) {
+    param.nn_descent_niter = conf.at("nn_descent_niter");
+  } else {
+    param.nn_descent_niter = 20;
+  }
+  nlohmann::json ivf_pq_build_conf = collect_conf_with_prefix(conf, "b_");
+  if (!ivf_pq_build_conf.empty()) {
+    faiss::gpu::IVFPQBuildCagraConfig ivf_pq_build_p;
+
+    if (ivf_pq_build_conf.contains("nlist")) {
+      ivf_pq_build_p.n_lists = ivf_pq_build_conf.at("nlist");
+    }
+    if (ivf_pq_build_conf.contains("niter")) {
+      ivf_pq_build_p.kmeans_n_iters = ivf_pq_build_conf.at("niter");
+    }
+    if (ivf_pq_build_conf.contains("ratio")) {
+      ivf_pq_build_p.kmeans_trainset_fraction = 1.0 / (double)conf.at("ratio");
+    }
+    if (ivf_pq_build_conf.contains("pq_bits")) {
+      ivf_pq_build_p.pq_bits = ivf_pq_build_conf.at("pq_bits");
+    }
+    if (ivf_pq_build_conf.contains("pq_dim")) {
+      ivf_pq_build_p.pq_dim = ivf_pq_build_conf.at("pq_dim");
+    }
+    param.ivf_pq_build_params = std::make_shared<faiss::gpu::IVFPQBuildCagraConfig>(ivf_pq_build_p);
+  }
+  nlohmann::json ivf_pq_search_conf = collect_conf_with_prefix(conf, "s_");
+  if (!ivf_pq_search_conf.empty()) {
+    faiss::gpu::IVFPQSearchCagraConfig ivf_pq_search_p;
+    if (ivf_pq_search_conf.contains("nprobe")) {
+      ivf_pq_search_p.n_probes = ivf_pq_search_conf.at("nprobe");
+    }
+    if (ivf_pq_search_conf.contains("internalDistanceDtype")) {
+      std::string type = ivf_pq_search_conf.at("internalDistanceDtype");
+      if (type == "float") {
+        ivf_pq_search_p.internal_distance_dtype = CUDA_R_32F;
+      } else if (type == "half") {
+        ivf_pq_search_p.internal_distance_dtype = CUDA_R_16F;
+      } else {
+        throw std::runtime_error("internalDistanceDtype: '" + type +
+                                 "', should be either 'float' or 'half'");
+      }
+    } else {
+      // set half as default type
+      ivf_pq_search_p.internal_distance_dtype = CUDA_R_16F;
+    }
+
+    if (ivf_pq_search_conf.contains("smemLutDtype")) {
+      std::string type = ivf_pq_search_conf.at("smemLutDtype");
+      if (type == "float") {
+        ivf_pq_search_p.lut_dtype = CUDA_R_32F;
+      } else if (type == "half") {
+        ivf_pq_search_p.lut_dtype = CUDA_R_16F;
+      } else if (type == "fp8") {
+        ivf_pq_search_p.lut_dtype = CUDA_R_8U;
+      } else {
+        throw std::runtime_error("smemLutDtype: '" + type +
+                                 "', should be either 'float', 'half' or 'fp8'");
+      }
+    } else {
+      // set half as default
+      ivf_pq_search_p.lut_dtype = CUDA_R_16F;
+    }
+    param.ivf_pq_search_params =
+      std::make_shared<faiss::gpu::IVFPQSearchCagraConfig>(ivf_pq_search_p);
+  }
+}
+
+template <typename T>
+void parse_build_param(const nlohmann::json& conf,
+                       typename cuvs::bench::faiss_gpu_cagra_hnsw<T>::build_param& param)
+{
+  typename cuvs::bench::faiss_gpu_cagra<T>::build_param p;
+  parse_build_param<T>(conf, p);
+  param.p = p;
+  if (conf.contains("base_level_only")) { param.base_level_only = conf.at("base_level_only"); }
 }
 
 template <typename T>
@@ -131,7 +223,13 @@ void parse_search_param(const nlohmann::json& conf,
       THROW("Invalid value for algo: %s", tmp.c_str());
     }
   }
-  if (conf.contains("refine_ratio")) { param.refine_ratio = conf.at("refine_ratio"); }
+}
+
+template <typename T>
+void parse_search_param(const nlohmann::json& conf,
+                        typename cuvs::bench::faiss_gpu_cagra_hnsw<T>::search_param& param)
+{
+  if (conf.contains("efSearch")) { param.p.efSearch = conf.at("efSearch"); }
 }
 
 template <typename T, template <typename> class Algo>
@@ -163,6 +261,8 @@ auto create_algo(const std::string& algo_name,
       a = std::make_unique<cuvs::bench::faiss_gpu_flat<T>>(metric, dim);
     } else if (algo_name == "faiss_gpu_cagra") {
       a = make_algo<T, cuvs::bench::faiss_gpu_cagra>(metric, dim, conf);
+    } else if (algo_name == "faiss_gpu_cagra_hnsw") {
+      a = make_algo<T, cuvs::bench::faiss_gpu_cagra_hnsw>(metric, dim, conf);
     }
   }
 
@@ -187,6 +287,10 @@ auto create_search_param(const std::string& algo_name, const nlohmann::json& con
     auto param = std::make_unique<typename cuvs::bench::faiss_gpu_cagra<T>::search_param>();
     parse_search_param<T>(conf, *param);
     return param;
+  } else if (algo_name == "faiss_gpu_cagra_hnsw") {
+    auto param = std::make_unique<typename cuvs::bench::faiss_gpu_cagra_hnsw<T>::search_param>();
+    parse_search_param<T>(conf, *param);
+    return param;
   }
   // else
   throw std::runtime_error("invalid algo: '" + algo_name + "'");