FAISS with cuVS enabled in cuvs-bench

cpp/bench/ann/CMakeLists.txt

@@ -20,6 +20,7 @@ list(APPEND CMAKE_MODULE_PATH "${CUVS_SOURCE_DIR}")
 option(CUVS_ANN_BENCH_USE_FAISS_GPU_FLAT "Include faiss' brute-force knn algorithm in benchmark" ON)
 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_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
@@ -124,6 +125,7 @@ function(ConfigureAnnBench)
   target_link_libraries(
     ${BENCH_NAME}
     PRIVATE ${ConfigureAnnBench_LINKS}
+            raft::raft
             nlohmann_json::nlohmann_json
             Threads::Threads
             $<TARGET_NAME_IF_EXISTS:raft::raft_logger>
@@ -290,6 +292,12 @@ if(CUVS_ANN_BENCH_USE_FAISS_GPU_FLAT AND CUVS_FAISS_ENABLE_GPU)
   )
 endif()
 
+if(CUVS_ANN_BENCH_USE_FAISS_GPU_CAGRA AND CUVS_FAISS_ENABLE_GPU)
+  ConfigureAnnBench(
+    NAME FAISS_GPU_CAGRA PATH src/faiss/faiss_gpu_benchmark.cu LINKS ${CUVS_FAISS_TARGETS}
+  )
+endif()
+
 if(CUVS_ANN_BENCH_USE_GGNN)
   include(cmake/thirdparty/get_glog)
   ConfigureAnnBench(

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

@@ -42,6 +42,11 @@ void parse_build_param(const nlohmann::json& conf,
                        typename cuvs::bench::faiss_gpu_ivf_flat<T>::build_param& param)
 {
   parse_base_build_param<T>(conf, param);
+  if (conf.contains("use_cuvs")) {
+    param.use_cuvs = conf.at("use_cuvs");
+  } else {
+    param.use_cuvs = false;
+  }
 }
 
 template <typename T>
@@ -60,6 +65,16 @@ void parse_build_param(const nlohmann::json& conf,
   } else {
     param.use_float16 = false;
   }
+  if (conf.contains("use_cuvs")) {
+    param.use_cuvs = conf.at("use_cuvs");
+  } else {
+    param.use_cuvs = false;
+  }
+  if (conf.contains("bitsPerCode")) {
+    param.bitsPerCode = conf.at("bitsPerCode");
+  } else {
+    param.bitsPerCode = 8;
+  }
 }
 
 template <typename T>
@@ -70,11 +85,52 @@ 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_gpu_cagra<T>::build_param& param)
+{
+  if (conf.contains("graph_degree")) {
+    param.graph_degree = conf.at("graph_degree");
+  } else {
+    param.graph_degree = 64;
+  }
+  if (conf.contains("intermediate_graph_degree")) {
+    param.intermediate_graph_degree = conf.at("intermediate_graph_degree");
+  } else {
+    param.intermediate_graph_degree = 128;
+  }
+  if (conf.contains("cagra_build_algo")) { param.cagra_build_algo = conf.at("cagra_build_algo"); }
+}
+
 template <typename T>
 void parse_search_param(const nlohmann::json& conf,
                         typename cuvs::bench::faiss_gpu<T>::search_param& param)
 {
-  param.nprobe = conf.at("nprobe");
+  if (conf.contains("nprobe")) { param.nprobe = conf.at("nprobe"); }
+  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<T>::search_param& param)
+{
+  if (conf.contains("itopk")) { param.p.itopk_size = conf.at("itopk"); }
+  if (conf.contains("search_width")) { param.p.search_width = conf.at("search_width"); }
+  if (conf.contains("max_iterations")) { param.p.max_iterations = conf.at("max_iterations"); }
+  if (conf.contains("algo")) {
+    if (conf.at("algo") == "single_cta") {
+      param.p.algo = faiss::gpu::search_algo::SINGLE_CTA;
+    } else if (conf.at("algo") == "multi_cta") {
+      param.p.algo = faiss::gpu::search_algo::MULTI_CTA;
+    } else if (conf.at("algo") == "multi_kernel") {
+      param.p.algo = faiss::gpu::search_algo::MULTI_KERNEL;
+    } else if (conf.at("algo") == "auto") {
+      param.p.algo = faiss::gpu::search_algo::AUTO;
+    } else {
+      std::string tmp = conf.at("algo");
+      THROW("Invalid value for algo: %s", tmp.c_str());
+    }
+  }
   if (conf.contains("refine_ratio")) { param.refine_ratio = conf.at("refine_ratio"); }
 }
 
@@ -105,6 +161,8 @@ auto create_algo(const std::string& algo_name,
       a = make_algo<T, cuvs::bench::faiss_gpu_ivfsq>(metric, dim, conf);
     } else if (algo_name == "faiss_gpu_flat") {
       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);
     }
   }
 
@@ -125,6 +183,10 @@ auto create_search_param(const std::string& algo_name, const nlohmann::json& con
   } else if (algo_name == "faiss_gpu_flat") {
     auto param = std::make_unique<typename cuvs::bench::faiss_gpu<T>::search_param>();
     return param;
+  } else if (algo_name == "faiss_gpu_cagra") {
+    auto param = std::make_unique<typename cuvs::bench::faiss_gpu_cagra<T>::search_param>();
+    parse_search_param<T>(conf, *param);
+    return param;
   }
   // else
   throw std::runtime_error("invalid algo: '" + algo_name + "'");