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=1.87
 - libclang==16.0.6
 - libcublas-dev=11.11.3.6
 - libcublas=11.11.3.6
@@ -38,6 +40,7 @@ dependencies:
 - libcuvs==25.4.*,>=0.0.0a0
 - librmm==25.4.*,>=0.0.0a0
 - matplotlib-base
+- mkl-devel=2023
 - nccl>=2.19
 - ninja
 - nlohmann_json>=3.11.2

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

@@ -27,6 +27,8 @@ dependencies:
 - gcc_linux-64=13.*
 - glog>=0.6.0
 - h5py>=3.8.0
+- libaio
+- libboost-devel=1.87
 - libclang==16.0.6
 - libcublas-dev
 - libcurand-dev
@@ -35,6 +37,7 @@ dependencies:
 - libcuvs==25.4.*,>=0.0.0a0
 - librmm==25.4.*,>=0.0.0a0
 - matplotlib-base
+- mkl-devel=2023
 - nccl>=2.19
 - ninja
 - nlohmann_json>=3.11.2

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

@@ -46,6 +46,9 @@ requirements:
   host:
     - benchmark
     - glog {{ glog_version }}
+    - libaio  # [linux64]
+    - libboost-devel=1.87  # [linux64]
+    - mkl-devel=2023  # [linux64]
     - nlohmann_json {{ nlohmann_json_version }}
     - openblas
     - python
@@ -62,6 +65,9 @@ requirements:
     - pyyaml
     - python
     - requests
+    - libaio  # [linux64]
+    - libboost-devel=1.87  # [linux64]
+    - mkl-devel=2023  # [linux64]
     - sklearn>=1.5
 about:
   home: https://rapids.ai/

conda/recipes/cuvs-bench/meta.yaml

@@ -70,7 +70,10 @@ requirements:
     - libcublas-dev
     {% endif %}
     - glog {{ glog_version }}
+    - libaio  # [linux64]
+    - libboost-devel=1.87  # [linux64]
     - libcuvs {{ version }}
+    - mkl-devel=2023  # [linux64]
     - nlohmann_json {{ nlohmann_json_version }}
     - openblas
     # rmm is needed to determine if package is gpu-enabled
@@ -101,6 +104,9 @@ requirements:
     - python
     - requests
     - rmm ={{ minor_version }}
+    - libaio  # [linux64]
+    - libboost-devel=1.87  # [linux64]
+    - mkl-devel=2023  # [linux64]
     - sklearn>=1.5
 about:
   home: https://rapids.ai/

cpp/bench/ann/CMakeLists.txt

@@ -1,5 +1,5 @@
 # =============================================================================
-# Copyright (c) 2024, NVIDIA CORPORATION.
+# Copyright (c) 2024-2025, 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
@@ -33,6 +33,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
@@ -58,6 +64,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()
@@ -70,6 +77,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()
@@ -87,6 +95,10 @@ if(CUVS_ANN_BENCH_USE_GGNN)
   include(cmake/thirdparty/get_ggnn)
 endif()
 
+if(CUVS_ANN_BENCH_USE_DISKANN OR CUVS_ANN_BENCH_USE_CUVS_VAMANA)
+  include(cmake/thirdparty/get_diskann)
+endif()
+
 if(CUVS_ANN_BENCH_USE_FAISS)
   include(cmake/thirdparty/get_faiss)
 endif()
@@ -299,6 +311,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

@@ -118,7 +118,7 @@ inline auto parse_algo_property(algo_property prop, const nlohmann::json& conf)
 template <typename T>
 void bench_build(::benchmark::State& state,
                  std::shared_ptr<const dataset<T>> dataset,
-                 configuration::index index,
+                 const configuration::index& index,
                  bool force_overwrite,
                  bool no_lap_sync)
 {
@@ -185,7 +185,7 @@ void bench_build(::benchmark::State& state,
 
 template <typename T>
 void bench_search(::benchmark::State& state,
-                  configuration::index index,
+                  const configuration::index& index,
                   std::size_t search_param_ix,
                   std::shared_ptr<const dataset<T>> dataset,
                   bool no_lap_sync)
@@ -469,11 +469,11 @@ inline void printf_usage()
 
 template <typename T>
 void register_build(std::shared_ptr<const dataset<T>> dataset,
-                    std::vector<configuration::index> indices,
+                    std::vector<configuration::index>& indices,
                     bool force_overwrite,
                     bool no_lap_sync)
 {
-  for (auto index : indices) {
+  for (auto& index : indices) {
     auto suf      = static_cast<std::string>(index.build_param["override_suffix"]);
     auto file_suf = suf;
     index.build_param.erase("override_suffix");
@@ -489,12 +489,12 @@ void register_build(std::shared_ptr<const dataset<T>> dataset,
 
 template <typename T>
 void register_search(std::shared_ptr<const dataset<T>> dataset,
-                     std::vector<configuration::index> indices,
+                     std::vector<configuration::index>& indices,
                      Mode metric_objective,
                      const std::vector<int>& threads,
                      bool no_lap_sync)
 {
-  for (auto index : indices) {
+  for (auto& index : indices) {
     for (std::size_t i = 0; i < index.search_params.size(); i++) {
       auto suf = static_cast<std::string>(index.search_params[i]["override_suffix"]);
       index.search_params[i].erase("override_suffix");

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

@@ -0,0 +1,93 @@
+/*
+ * Copyright (c) 2025, 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