Diskann Benchmarking Wrapper

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

@@ -25,6 +25,8 @@ dependencies:
 - glog>=0.6.0
 - h5py>=3.8.0
 - hnswlib=0.6.2
+- libaio
+- libboost-devel
 - libcublas-dev=11.11.3.6
 - libcublas=11.11.3.6
 - libcurand-dev=10.3.0.86
@@ -35,6 +37,7 @@ dependencies:
 - libcusparse=11.7.5.86
 - librmm==24.12.*,>=0.0.0a0
 - matplotlib
+- mkl-devel
 - nccl>=2.19
 - ninja
 - nlohmann_json>=3.11.2

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

@@ -25,6 +25,8 @@ dependencies:
 - glog>=0.6.0
 - h5py>=3.8.0
 - hnswlib=0.6.2
+- libaio
+- libboost-devel
 - libcublas-dev=11.11.3.6
 - libcublas=11.11.3.6
 - libcurand-dev=10.3.0.86
@@ -35,6 +37,7 @@ dependencies:
 - libcusparse=11.7.5.86
 - librmm==24.12.*,>=0.0.0a0
 - matplotlib
+- mkl-devel
 - nccl>=2.19
 - ninja
 - nlohmann_json>=3.11.2

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

@@ -26,12 +26,15 @@ dependencies:
 - glog>=0.6.0
 - h5py>=3.8.0
 - hnswlib=0.6.2
+- libaio
+- libboost-devel
 - libcublas-dev
 - libcurand-dev
 - libcusolver-dev
 - libcusparse-dev
 - librmm==24.12.*,>=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

@@ -26,12 +26,15 @@ dependencies:
 - glog>=0.6.0
 - h5py>=3.8.0
 - hnswlib=0.6.2
+- libaio
+- libboost-devel
 - libcublas-dev
 - libcurand-dev
 - libcusolver-dev
 - libcusparse-dev
 - librmm==24.12.*,>=0.0.0a0
 - matplotlib
+- mkl-devel
 - nccl>=2.19
 - ninja
 - nlohmann_json>=3.11.2

conda/recipes/cuvs_bench/meta.yaml

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

conda/recipes/cuvs_bench_cpu/meta.yaml

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

cpp/bench/ann/CMakeLists.txt

@@ -32,6 +32,8 @@ 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)
 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
@@ -69,10 +71,15 @@ 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()
 
+if(CUVS_ANN_BENCH_USE_CUVS_VAMANA)
+  set(CUVS_ANN_BENCH_USE_DISKANN ON)
+endif()
+
 # ##################################################################################################
 # * Fetch requirements -------------------------------------------------------------
 
@@ -90,6 +97,10 @@ if(CUVS_ANN_BENCH_USE_FAISS)
   include(cmake/thirdparty/get_faiss)
 endif()
 
+if(CUVS_ANN_BENCH_USE_DISKANN)
+  include(cmake/thirdparty/get_diskann)
+endif()
+
 # ##################################################################################################
 # * Target function -------------------------------------------------------------
 
@@ -290,6 +301,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 cuvs diskann::diskann
+  )
+  ConfigureAnnBench(
+    NAME DISKANN_SSD PATH src/diskann/diskann_benchmark.cpp LINKS cuvs diskann::diskann
+  )
+endif()
+
+if(CUVS_ANN_BENCH_USE_CUVS_VAMANA)
+  ConfigureAnnBench(NAME CUVS_VAMANA PATH src/cuvs/cuvs_vamana.cpp LINKS cuvs diskann::diskann)
+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_cagra_diskann_wrapper.h

@@ -0,0 +1,108 @@
+/*
+ * Copyright (c) 2023-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.
+ */
+#pragma once
+
+#include "../hnswlib/hnswlib_wrapper.h"
+#include "cuvs_cagra_wrapper.h"
+
+#include <memory>
+
+namespace cuvs::bench {
+
+template <typename T, typename IdxT>
+class cuvs_cagra_diskann : public algo<T>, public algo_gpu {
+ public:
+  using search_param_base = typename algo<T>::search_param;
+  using build_param       = typename cuvs_cagra<T, IdxT>::build_param;
+  using search_param      = typename diskann_mem<T>::search_param;
+
+  cuvs_cagra_diskann(Metric metric, int dim, const build_param& param)
+    : algo<T>(metric, dim),
+      cagra_build_{metric, dim, param},
+      // hnsw_lib param values don't matter since we don't build with hnsw_lib
+      diskann_mem_search_{metric, dim, typename diskann_mem<T>::build_param{50, 100}}
+  {
+  }
+
+  void build(const T* dataset, size_t nrow) final;
+
+  void set_search_param(const search_param_base& param) override;
+
+  void search(const T* queries,
+              int batch_size,
+              int k,
+              algo_base::index_type* neighbors,
+              float* distances) const override;
+
+  [[nodiscard]] auto get_sync_stream() const noexcept -> cudaStream_t override
+  {
+    return cagra_build_.get_sync_stream();
+  }
+
+  // to enable dataset access from GPU memory
+  [[nodiscard]] auto get_preference() const -> algo_property override
+  {
+    algo_property property;
+    property.dataset_memory_type = MemoryType::kHostMmap;
+    property.query_memory_type   = MemoryType::kHost;
+    return property;
+  }
+
+  void save(const std::string& file) const override;
+  void load(const std::string&) override;
+  std::unique_ptr<algo<T>> copy() override
+  {
+    return std::make_unique<cuvs_cagra_hnswlib<T, IdxT>>(*this);
+  }
+
+ private:
+  cuvs_cagra<T, IdxT> cagra_build_;
+  hnsw_lib<T> hnswlib_search_;
+};
+
+template <typename T, typename IdxT>
+void cuvs_cagra_hnswlib<T, IdxT>::build(const T* dataset, size_t nrow)
+{
+  cagra_build_.build(dataset, nrow);
+}
+
+template <typename T, typename IdxT>
+void cuvs_cagra_hnswlib<T, IdxT>::set_search_param(const search_param_base& param_)
+{
+  hnswlib_search_.set_search_param(param_);
+}
+
+template <typename T, typename IdxT>
+void cuvs_cagra_hnswlib<T, IdxT>::save(const std::string& file) const
+{
+  cagra_build_.save_to_hnswlib(file);
+}
+
+template <typename T, typename IdxT>
+void cuvs_cagra_hnswlib<T, IdxT>::load(const std::string& file)
+{
+  hnswlib_search_.load(file);
+  hnswlib_search_.set_base_layer_only();
+}
+
+template <typename T, typename IdxT>
+void cuvs_cagra_hnswlib<T, IdxT>::search(
+  const T* queries, int batch_size, int k, algo_base::index_type* neighbors, float* distances) const
+{
+  hnswlib_search_.search(queries, batch_size, k, neighbors, distances);
+}
+
+}  // namespace cuvs::bench