Diskann Benchmarking Wrapper

conda/environments/bench_ann_cuda-118_arch-x86_64.yaml

@@ -26,6 +26,8 @@ dependencies:
 - gcc_linux-64=11.*
 - glog>=0.6.0
 - h5py>=3.8.0
+- libaio
+- libboost-devel
 - libcublas-dev=11.11.3.6
 - libcublas=11.11.3.6
 - libcurand-dev=10.3.0.86
@@ -37,6 +39,7 @@ dependencies:
 - libcuvs==25.2.*,>=0.0.0a0
 - librmm==25.2.*,>=0.0.0a0
 - matplotlib
+- mkl-devel
 - nccl>=2.19
 - ninja
 - nlohmann_json>=3.11.2

conda/environments/bench_ann_cuda-125_arch-x86_64.yaml

@@ -27,13 +27,16 @@ dependencies:
 - gcc_linux-64=11.*
 - glog>=0.6.0
 - h5py>=3.8.0
+- libaio
+- libboost-devel
 - libcublas-dev
 - libcurand-dev
 - libcusolver-dev
 - libcusparse-dev
 - libcuvs==25.2.*,>=0.0.0a0
 - librmm==25.2.*,>=0.0.0a0
 - matplotlib
+- mkl-devel
 - nccl>=2.19
 - ninja
 - nlohmann_json>=3.11.2

conda/recipes/cuvs-bench-cpu/meta.yaml

@@ -47,6 +47,9 @@ requirements:
     - benchmark
     - fmt {{ fmt_version }}
     - glog {{ glog_version }}
+    - libaio
+    - libboost-devel
+    - mkl-devel  # [linux64]
     - nlohmann_json {{ nlohmann_json_version }}
     - openblas
     - python

conda/recipes/cuvs-bench/meta.yaml

@@ -72,7 +72,10 @@ requirements:
     - libcublas-dev
     {% endif %}
     - glog {{ glog_version }}
+    - libaio
+    - libboost-devel
     - libcuvs {{ version }}
+    - mkl-devel  # [linux64]
     - nlohmann_json {{ nlohmann_json_version }}
     - openblas
     # rmm is needed to determine if package is gpu-enabled

cpp/bench/ann/CMakeLists.txt

@@ -32,6 +32,12 @@ option(CUVS_ANN_BENCH_USE_CUVS_BRUTE_FORCE "Include cuVS brute force knn in benc
 option(CUVS_ANN_BENCH_USE_CUVS_CAGRA_HNSWLIB "Include cuVS CAGRA with HNSW search in benchmark" ON)
 option(CUVS_ANN_BENCH_USE_HNSWLIB "Include hnsw algorithm in benchmark" ON)
 option(CUVS_ANN_BENCH_USE_GGNN "Include ggnn algorithm in benchmark" OFF)
+option(CUVS_ANN_BENCH_USE_DISKANN "Include DISKANN search in benchmark" ON)
+option(CUVS_ANN_BENCH_USE_CUVS_VAMANA "Include cuVS Vamana with DiskANN search in benchmark" ON)
+if(CMAKE_SYSTEM_PROCESSOR MATCHES "(ARM|arm|aarch64)")
+  set(CUVS_ANN_BENCH_USE_DISKANN OFF)
+  set(CUVS_ANN_BENCH_USE_CUVS_VAMANA OFF)
+endif()
 option(CUVS_ANN_BENCH_USE_CUVS_MG "Include cuVS ann mg algorithm in benchmark" ${BUILD_MG_ALGOS})
 option(CUVS_ANN_BENCH_SINGLE_EXE
        "Make a single executable with benchmark as shared library modules" OFF
@@ -57,6 +63,7 @@ if(BUILD_CPU_ONLY)
   set(CUVS_ANN_BENCH_USE_GGNN OFF)
   set(CUVS_KNN_BENCH_USE_CUVS_BRUTE_FORCE OFF)
   set(CUVS_ANN_BENCH_USE_CUVS_MG OFF)
+  set(CUVS_ANN_BENCH_USE_CUVS_VAMANA OFF)
 else()
   set(CUVS_FAISS_ENABLE_GPU ON)
 endif()
@@ -69,6 +76,7 @@ if(CUVS_ANN_BENCH_USE_CUVS_IVF_PQ
    OR CUVS_ANN_BENCH_USE_CUVS_CAGRA_HNSWLIB
    OR CUVS_KNN_BENCH_USE_CUVS_BRUTE_FORCE
    OR CUVS_ANN_BENCH_USE_CUVS_MG
+   OR CUVS_ANN_BENCH_USE_CUVS_VAMANA
 )
   set(CUVS_ANN_BENCH_USE_CUVS ON)
 endif()
@@ -90,6 +98,10 @@ if(CUVS_ANN_BENCH_USE_FAISS)
   include(cmake/thirdparty/get_faiss)
 endif()
 
+if(CUVS_ANN_BENCH_USE_DISKANN OR CUVS_ANN_BENCH_USE_CUVS_VAMANA)
+  include(cmake/thirdparty/get_diskann)
+endif()
+
 # ##################################################################################################
 # * Target function -------------------------------------------------------------
 
@@ -294,6 +306,19 @@ if(CUVS_ANN_BENCH_USE_GGNN)
   )
 endif()
 
+if(CUVS_ANN_BENCH_USE_DISKANN)
+  ConfigureAnnBench(
+    NAME DISKANN_MEMORY PATH src/diskann/diskann_benchmark.cpp LINKS diskann::diskann aio
+  )
+  ConfigureAnnBench(
+    NAME DISKANN_SSD PATH src/diskann/diskann_benchmark.cpp LINKS diskann::diskann aio
+  )
+endif()
+
+if(CUVS_ANN_BENCH_USE_CUVS_VAMANA)
+  ConfigureAnnBench(NAME CUVS_VAMANA PATH src/cuvs/cuvs_vamana.cu LINKS cuvs diskann::diskann aio)
+endif()
+
 # ##################################################################################################
 # * Dynamically-loading ANN_BENCH executable -------------------------------------------------------
 if(CUVS_ANN_BENCH_SINGLE_EXE)

cpp/bench/ann/src/common/benchmark.hpp

@@ -135,6 +135,12 @@ void bench_build(::benchmark::State& state,
     }
   }
 
+  if (index.algo == "diskann_ssd") {
+    make_sure_parent_dir_exists(index.file);
+    index.build_param["dataset_file"]  = dataset->base_filename();
+    index.build_param["path_to_index"] = index.file;
+  }
+
   std::unique_ptr<algo<T>> algo;
   try {
     algo = create_algo<T>(index.algo, dataset->distance(), dataset->dim(), index.build_param);
@@ -144,7 +150,8 @@ void bench_build(::benchmark::State& state,
 
   const auto algo_property = parse_algo_property(algo->get_preference(), index.build_param);
 
-  const T* base_set      = dataset->base_set(algo_property.dataset_memory_type);
+  const T* base_set = nullptr;
+  if (index.algo != "diskann_ssd") base_set = dataset->base_set(algo_property.dataset_memory_type);
   std::size_t index_size = dataset->base_set_size();
 
   cuda_timer gpu_timer{algo};
@@ -223,7 +230,12 @@ void bench_search(::benchmark::State& state,
 
   const T* query_set = nullptr;
 
-  if (!file_exists(index.file)) {
+  std::string filename;
+  if (index.algo != "diskann_ssd")
+    filename = index.file;
+  else
+    filename = index.file + "_disk.index";
+  if (!file_exists(filename)) {
     state.SkipWithError("Index file is missing. Run the benchmark in the build mode first.");
     return;
   }

cpp/bench/ann/src/common/dataset.hpp

@@ -114,6 +114,8 @@ class bin_file {
     }
   }
 
+  std::string file() const { return file_; }
+
  private:
   void check_suffix();
   void open_file() const;
@@ -253,10 +255,11 @@ class dataset {
 
   auto name() const -> std::string { return name_; }
   auto distance() const -> std::string { return distance_; }
-  virtual auto dim() const -> int               = 0;
-  virtual auto max_k() const -> uint32_t        = 0;
-  virtual auto base_set_size() const -> size_t  = 0;
-  virtual auto query_set_size() const -> size_t = 0;
+  virtual auto dim() const -> int                   = 0;
+  virtual auto max_k() const -> uint32_t            = 0;
+  virtual auto base_set_size() const -> size_t      = 0;
+  virtual auto query_set_size() const -> size_t     = 0;
+  virtual auto base_filename() const -> std::string = 0;
 
   // load data lazily, so don't pay the overhead of reading unneeded set
   // e.g. don't load base set when searching
@@ -424,6 +427,7 @@ class bin_dataset : public dataset<T> {
   auto max_k() const -> uint32_t override;
   auto base_set_size() const -> size_t override;
   auto query_set_size() const -> size_t override;
+  std::string base_filename() const override;
 
  private:
   void load_base_set() const;
@@ -541,4 +545,10 @@ void bin_dataset<T>::map_base_set() const
   this->mapped_base_set_ = base_file_.map();
 }
 
+template <typename T>
+std::string bin_dataset<T>::base_filename() const
+{
+  return base_file_.file();
+}
+
 }  // namespace  cuvs::bench

cpp/bench/ann/src/cuvs/cuvs_vamana.cu

@@ -0,0 +1,93 @@
+/*
+ * Copyright (c) 2024, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "../common/ann_types.hpp"
+#include "cuvs_vamana_wrapper.h"
+
+#include <rmm/cuda_device.hpp>
+#include <rmm/mr/device/pool_memory_resource.hpp>
+#include <rmm/resource_ref.hpp>
+
+namespace cuvs::bench {
+
+template <typename T, typename IdxT>
+void parse_build_param(const nlohmann::json& conf,
+                       typename cuvs::bench::cuvs_vamana<T, IdxT>::build_param& param)
+{
+  if (conf.contains("graph_degree")) { param.graph_degree = conf.at("graph_degree"); }
+  if (conf.contains("visited_size")) { param.visited_size = conf.at("visited_size"); }
+  if (conf.contains("alpha")) { param.alpha = conf.at("alpha"); }
+}
+
+template <typename T, typename IdxT>
+void parse_search_param(const nlohmann::json& conf,
+                        typename cuvs::bench::cuvs_vamana<T, IdxT>::search_param& param)
+{
+  if (conf.contains("L_search")) { param.L_search = conf.at("L_search"); }
+  if (conf.contains("num_threads")) { param.num_threads = conf.at("num_threads"); }
+}
+
+template <typename T>
+auto create_algo(const std::string& algo_name,
+                 const std::string& distance,
+                 int dim,
+                 const nlohmann::json& conf) -> std::unique_ptr<cuvs::bench::algo<T>>
+{
+  [[maybe_unused]] cuvs::bench::Metric metric = parse_metric(distance);
+  std::unique_ptr<cuvs::bench::algo<T>> a;
+
+  if constexpr (std::is_same_v<T, float> or std::is_same_v<T, std::uint8_t>) {
+    if (algo_name == "cuvs_vamana") {
+      typename cuvs::bench::cuvs_vamana<T, uint32_t>::build_param param;
+      parse_build_param<T, uint32_t>(conf, param);
+      a = std::make_unique<cuvs::bench::cuvs_vamana<T, uint32_t>>(metric, dim, param);
+    }
+  }
+
+  if (!a) { throw std::runtime_error("invalid algo: '" + algo_name + "'"); }
+
+  return a;
+}
+
+template <typename T>
+auto create_search_param(const std::string& algo_name, const nlohmann::json& conf)
+  -> std::unique_ptr<typename cuvs::bench::algo<T>::search_param>
+{
+  if (algo_name == "cuvs_vamana") {
+    auto param = std::make_unique<typename cuvs::bench::cuvs_vamana<T, uint32_t>::search_param>();
+    parse_search_param<T, uint32_t>(conf, *param);
+    return param;
+  }
+
+  throw std::runtime_error("invalid algo: '" + algo_name + "'");
+}
+
+}  // namespace cuvs::bench
+
+REGISTER_ALGO_INSTANCE(float);
+
+#ifdef ANN_BENCH_BUILD_MAIN
+#include "../common/benchmark.hpp"
+/*
+[NOTE] Dear developer,
+
+Please don't modify the content of the `main` function; this will make the behavior of the benchmark
+executable differ depending on the cmake flags and will complicate the debugging. In particular,
+don't try to setup an RMM memory resource here; it will anyway be modified by the memory resource
+set on per-algorithm basis. For example, see `cuvs/cuvs_ann_bench_utils.h`.
+*/
+int main(int argc, char** argv) { return cuvs::bench::run_main(argc, argv); }
+#endif