Benchmark CAGRA + DiskANN

README.md

@@ -271,8 +271,8 @@ The easiest way to install RAFT is through conda and several packages are provid
 - `libraft` (optional) C++ shared library containing pre-compiled template instantiations and runtime API.
 - `pylibraft` (optional) Python library
 - `raft-dask` (optional) Python library for deployment of multi-node multi-GPU algorithms that use the RAFT `raft::comms` abstraction layer in Dask clusters.
-- `raft-ann-bench` (optional) Benchmarking tool for easily producing benchmarks that compare RAFT's vector search algorithms against other state-of-the-art implementations.
-- `raft-ann-bench-cpu` (optional) Reproducible benchmarking tool similar to above, but doesn't require CUDA to be installed on the machine. Can be used to test in environments with competitive CPUs.
+- `raft_ann_bench` (optional) Benchmarking tool for easily producing benchmarks that compare RAFT's vector search algorithms against other state-of-the-art implementations.
+- `raft_ann_bench-cpu` (optional) Reproducible benchmarking tool similar to above, but doesn't require CUDA to be installed on the machine. Can be used to test in environments with competitive CPUs.
 
 Use the following command, depending on your CUDA version, to install all of the RAFT packages with conda (replace `rapidsai` with `rapidsai-nightly` to install more up-to-date but less stable nightly packages). `mamba` is preferred over the `conda` command.
 ```bash

build.sh

@@ -505,9 +505,9 @@ if (( ${NUMARGS} == 0 )) || hasArg raft-dask; then
         python -m pip install --no-build-isolation --no-deps --config-settings rapidsai.disable-cuda=true ${REPODIR}/python/raft-dask
 fi
 
-# Build and (optionally) install the raft-ann-bench Python package
+# Build and (optionally) install the raft_ann_bench Python package
 if (( ${NUMARGS} == 0 )) || hasArg bench-ann; then
-    python -m pip install --no-build-isolation --no-deps --config-settings rapidsai.disable-cuda=true ${REPODIR}/python/raft-ann-bench -vvv
+    python -m pip install --no-build-isolation --no-deps --config-settings rapidsai.disable-cuda=true ${REPODIR}/python/raft_ann_bench -vvv
 fi
 
 if hasArg docs; then

ci/build_python.sh

@@ -40,7 +40,7 @@ rapids-conda-retry mambabuild \
   --no-test \
   --channel "${CPP_CHANNEL}" \
   --channel "${RAPIDS_CONDA_BLD_OUTPUT_DIR}" \
-  conda/recipes/raft-ann-bench
+  conda/recipes/raft_ann_bench
 
 # Build ann-bench-cpu only in CUDA 11 jobs since it only depends on python
 # version
@@ -50,7 +50,7 @@ if [[ ${RAPIDS_CUDA_MAJOR} == "11" ]]; then
   --no-test \
   --channel "${CPP_CHANNEL}" \
   --channel "${RAPIDS_CONDA_BLD_OUTPUT_DIR}" \
-  conda/recipes/raft-ann-bench-cpu
+  conda/recipes/raft_ann_bench-cpu
 fi
 
 rapids-upload-conda-to-s3 python

conda/recipes/raft-ann-bench-cpu/meta.yaml

@@ -9,7 +9,7 @@
 {% set date_string = environ['RAPIDS_DATE_STRING'] %}
 
 package:
-  name: raft-ann-bench-cpu
+  name: raft_ann_bench-cpu
   version: {{ version }}
   script: build.sh
 

conda/recipes/raft-ann-bench/meta.yaml

@@ -10,7 +10,7 @@
 {% set date_string = environ['RAPIDS_DATE_STRING'] %}
 
 package:
-  name: raft-ann-bench
+  name: raft_ann_bench
   version: {{ version }}
   script: build.sh
 

cpp/bench/ann/CMakeLists.txt

@@ -31,8 +31,9 @@ option(RAFT_ANN_BENCH_USE_RAFT_IVF_PQ "Include raft's ivf pq algorithm in benchm
 option(RAFT_ANN_BENCH_USE_RAFT_CAGRA "Include raft's CAGRA in benchmark" ON)
 option(RAFT_ANN_BENCH_USE_RAFT_BRUTE_FORCE "Include raft's brute force knn in benchmark" ON)
 option(RAFT_ANN_BENCH_USE_RAFT_CAGRA_HNSWLIB "Include raft's CAGRA in benchmark" ON)
-option(RAFT_ANN_BENCH_USE_HNSWLIB "Include hnsw algorithm in benchmark" ON)
-option(RAFT_ANN_BENCH_USE_GGNN "Include ggnn algorithm in benchmark" ON)
+option(RAFT_ANN_BENCH_USE_HNSWLIB "Include hnsw algorithm in benchmark" OFF)
+option(RAFT_ANN_BENCH_USE_GGNN "Include ggnn algorithm in benchmark" OFF)
+option(RAFT_ANN_BENCH_USE_DISKANN "Include diskann algorithm in benchmark" ON)
 option(RAFT_ANN_BENCH_SINGLE_EXE
        "Make a single executable with benchmark as shared library modules" OFF
 )
@@ -51,13 +52,15 @@ if(BUILD_CPU_ONLY)
   # Include necessary logging dependencies
   include(cmake/thirdparty/get_fmt)
   include(cmake/thirdparty/get_spdlog)
+
   set(RAFT_FAISS_ENABLE_GPU OFF)
   set(RAFT_ANN_BENCH_USE_RAFT_IVF_FLAT OFF)
   set(RAFT_ANN_BENCH_USE_RAFT_IVF_PQ OFF)
   set(RAFT_ANN_BENCH_USE_RAFT_CAGRA OFF)
   set(RAFT_ANN_BENCH_USE_RAFT_BRUTE_FORCE OFF)
   set(RAFT_ANN_BENCH_USE_RAFT_CAGRA_HNSWLIB OFF)
   set(RAFT_ANN_BENCH_USE_GGNN OFF)
+  set(RAFT_ANN_BENCH_USE_DISKANN OFF)
 elseif(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0.0)
   # Disable faiss benchmarks on CUDA 12 since faiss is not yet CUDA 12-enabled.
   # https://github.com/rapidsai/raft/issues/1627
@@ -91,6 +94,10 @@ if(RAFT_ANN_BENCH_USE_FAISS)
   include(cmake/thirdparty/get_faiss)
 endif()
 
+if(RAFT_ANN_BENCH_USE_DISKANN)
+  include(cmake/thirdparty/get_diskann)
+endif()
+
 # ##################################################################################################
 # * Enable NVTX if available
 
@@ -149,6 +156,8 @@ function(ConfigureAnnBench)
             $<$<BOOL:${GPU_BUILD}>:${RAFT_CTK_MATH_DEPENDENCIES}>
             $<TARGET_NAME_IF_EXISTS:OpenMP::OpenMP_CXX>
             $<TARGET_NAME_IF_EXISTS:conda_env>
+            -static-libgcc
+            -static-libstdc++
             $<$<BOOL:${BUILD_CPU_ONLY}>:fmt::fmt-header-only>
             $<$<BOOL:${BUILD_CPU_ONLY}>:spdlog::spdlog_header_only>
   )
@@ -198,40 +207,25 @@ endfunction()
 # * Configure tests-------------------------------------------------------------
 
 if(RAFT_ANN_BENCH_USE_HNSWLIB)
-  ConfigureAnnBench(
-    NAME HNSWLIB PATH src/hnswlib/hnswlib_benchmark.cpp LINKS hnswlib::hnswlib
-  )
+  ConfigureAnnBench(NAME HNSWLIB PATH src/hnswlib/hnswlib_benchmark.cpp LINKS hnswlib::hnswlib)
 
 endif()
 
 if(RAFT_ANN_BENCH_USE_RAFT_IVF_PQ)
   ConfigureAnnBench(
-    NAME
-    RAFT_IVF_PQ
-    PATH
-    src/raft/raft_benchmark.cu
-    src/raft/raft_ivf_pq.cu
-    LINKS
-    raft::compiled
+    NAME RAFT_IVF_PQ PATH src/raft/raft_benchmark.cu src/raft/raft_ivf_pq.cu LINKS raft::compiled
   )
 endif()
 
 if(RAFT_ANN_BENCH_USE_RAFT_IVF_FLAT)
   ConfigureAnnBench(
-    NAME
-    RAFT_IVF_FLAT
-    PATH
-    src/raft/raft_benchmark.cu
-    src/raft/raft_ivf_flat.cu
-    LINKS
+    NAME RAFT_IVF_FLAT PATH src/raft/raft_benchmark.cu src/raft/raft_ivf_flat.cu LINKS
     raft::compiled
   )
 endif()
 
 if(RAFT_ANN_BENCH_USE_RAFT_BRUTE_FORCE)
-  ConfigureAnnBench(
-    NAME RAFT_BRUTE_FORCE PATH src/raft/raft_benchmark.cu LINKS raft::compiled
-  )
+  ConfigureAnnBench(NAME RAFT_BRUTE_FORCE PATH src/raft/raft_benchmark.cu LINKS raft::compiled)
 endif()
 
 if(RAFT_ANN_BENCH_USE_RAFT_CAGRA)
@@ -260,36 +254,31 @@ message("RAFT_FAISS_TARGETS: ${RAFT_FAISS_TARGETS}")
 message("CUDAToolkit_LIBRARY_DIR: ${CUDAToolkit_LIBRARY_DIR}")
 if(RAFT_ANN_BENCH_USE_FAISS_CPU_FLAT)
   ConfigureAnnBench(
-    NAME FAISS_CPU_FLAT PATH src/faiss/faiss_cpu_benchmark.cpp LINKS
-    ${RAFT_FAISS_TARGETS}
+    NAME FAISS_CPU_FLAT PATH src/faiss/faiss_cpu_benchmark.cpp LINKS ${RAFT_FAISS_TARGETS}
   )
 endif()
 
 if(RAFT_ANN_BENCH_USE_FAISS_CPU_IVF_FLAT)
   ConfigureAnnBench(
-    NAME FAISS_CPU_IVF_FLAT PATH src/faiss/faiss_cpu_benchmark.cpp LINKS
-    ${RAFT_FAISS_TARGETS}
+    NAME FAISS_CPU_IVF_FLAT PATH src/faiss/faiss_cpu_benchmark.cpp LINKS ${RAFT_FAISS_TARGETS}
   )
 endif()
 
 if(RAFT_ANN_BENCH_USE_FAISS_CPU_IVF_PQ)
   ConfigureAnnBench(
-    NAME FAISS_CPU_IVF_PQ PATH src/faiss/faiss_cpu_benchmark.cpp LINKS
-    ${RAFT_FAISS_TARGETS}
+    NAME FAISS_CPU_IVF_PQ PATH src/faiss/faiss_cpu_benchmark.cpp LINKS ${RAFT_FAISS_TARGETS}
   )
 endif()
 
 if(RAFT_ANN_BENCH_USE_FAISS_GPU_IVF_FLAT AND RAFT_FAISS_ENABLE_GPU)
   ConfigureAnnBench(
-    NAME FAISS_GPU_IVF_FLAT PATH src/faiss/faiss_gpu_benchmark.cu LINKS
-    ${RAFT_FAISS_TARGETS}
+    NAME FAISS_GPU_IVF_FLAT PATH src/faiss/faiss_gpu_benchmark.cu LINKS ${RAFT_FAISS_TARGETS}
   )
 endif()
 
 if(RAFT_ANN_BENCH_USE_FAISS_GPU_IVF_PQ AND RAFT_FAISS_ENABLE_GPU)
   ConfigureAnnBench(
-    NAME FAISS_GPU_IVF_PQ PATH src/faiss/faiss_gpu_benchmark.cu LINKS
-    ${RAFT_FAISS_TARGETS}
+    NAME FAISS_GPU_IVF_PQ PATH src/faiss/faiss_gpu_benchmark.cu LINKS ${RAFT_FAISS_TARGETS}
   )
 endif()
 
@@ -304,6 +293,10 @@ if(RAFT_ANN_BENCH_USE_GGNN)
   ConfigureAnnBench(NAME GGNN PATH src/ggnn/ggnn_benchmark.cu LINKS glog::glog ggnn::ggnn)
 endif()
 
+if(RAFT_ANN_BENCH_USE_DISKANN)
+  ConfigureAnnBench(NAME DISKANN PATH src/diskann/diskann_benchmark.cu LINKS diskann::diskann)
+endif()
+
 # ##################################################################################################
 # * Dynamically-loading ANN_BENCH executable -------------------------------------------------------
 if(RAFT_ANN_BENCH_SINGLE_EXE)

cpp/bench/ann/src/diskann/diskann_benchmark.cu

@@ -0,0 +1,123 @@
+/*
+ * 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.
+ */
+
+#include "../common/ann_types.hpp"
+#include "diskann_wrapper.cuh"
+
+#define JSON_DIAGNOSTICS 1
+#include <nlohmann/json.hpp>
+
+#include <algorithm>
+#include <cmath>
+#include <memory>
+#include <stdexcept>
+#include <string>
+#include <type_traits>
+#include <utility>
+
+namespace raft::bench::ann {
+
+template <typename T>
+void parse_build_param(const nlohmann::json& conf,
+                       typename raft::bench::ann::DiskANNMemory<T>::BuildParam& param)
+{
+  param.R       = conf.at("R");
+  param.L_build = conf.at("Lb");
+  param.alpha   = conf.at("alpha");
+  if (conf.contains("numThreads")) { param.num_threads = conf.at("numThreads"); }
+  param.use_cagra_graph = conf.at("use_cagra_graph");
+  if (param.use_cagra_graph) {
+    param.cagra_graph_degree              = conf.at("cagra_graph_degree");
+    param.cagra_intermediate_graph_degree = conf.at("cagra_intermediate_graph_degree");
+  }
+}
+
+template <typename T>
+void parse_search_param(const nlohmann::json& conf,
+                        typename raft::bench::ann::DiskANNMemory<T>::SearchParam& param)
+{
+  param.L_search = conf.at("L_search");
+}
+
+template <typename T, template <typename> class Algo>
+std::unique_ptr<raft::bench::ann::ANN<T>> make_algo(raft::bench::ann::Metric metric,
+                                                    int dim,
+                                                    const nlohmann::json& conf)
+{
+  typename Algo<T>::BuildParam param;
+  parse_build_param<T>(conf, param);
+  return std::make_unique<Algo<T>>(metric, dim, param);
+}
+
+template <typename T, template <typename> class Algo>
+std::unique_ptr<raft::bench::ann::ANN<T>> make_algo(raft::bench::ann::Metric metric,
+                                                    int dim,
+                                                    const nlohmann::json& conf,
+                                                    const std::vector<int>& dev_list)
+{
+  typename Algo<T>::BuildParam param;
+  parse_build_param<T>(conf, param);
+
+  (void)dev_list;
+  return std::make_unique<Algo<T>>(metric, dim, param);
+}
+
+template <typename T>
+std::unique_ptr<raft::bench::ann::ANN<T>> create_algo(const std::string& algo,
+                                                      const std::string& distance,
+                                                      int dim,
+                                                      const nlohmann::json& conf,
+                                                      const std::vector<int>& dev_list)
+{
+  // stop compiler warning; not all algorithms support multi-GPU so it may not be used
+  (void)dev_list;
+
+  raft::bench::ann::Metric metric = parse_metric(distance);
+  std::unique_ptr<raft::bench::ann::ANN<T>> ann;
+
+  if constexpr (std::is_same_v<T, float> || std::is_same_v<T, uint8_t> ||
+                std::is_same_v<T, int8_t>) {
+    if (algo == "diskann") {
+      ann = make_algo<T, raft::bench::ann::DiskANNMemory>(metric, dim, conf);
+    }
+  }
+  if (!ann) { throw std::runtime_error("invalid algo: '" + algo + "'"); }
+
+  return ann;
+}
+
+template <typename T>
+std::unique_ptr<typename raft::bench::ann::ANN<T>::AnnSearchParam> create_search_param(
+  const std::string& algo, const nlohmann::json& conf)
+{
+  if (algo == "diskann") {
+    auto param = std::make_unique<typename raft::bench::ann::DiskANNMemory<T>::SearchParam>();
+    parse_search_param<T>(conf, *param);
+    return param;
+  }
+  throw std::runtime_error("invalid algo: '" + algo + "'");
+}
+
+};  // namespace raft::bench::ann
+
+REGISTER_ALGO_INSTANCE(float);
+REGISTER_ALGO_INSTANCE(std::int8_t);
+REGISTER_ALGO_INSTANCE(std::uint8_t);
+
+#ifdef ANN_BENCH_BUILD_MAIN
+#include "../common/benchmark.hpp"
+int main(int argc, char** argv) { return raft::bench::ann::run_main(argc, argv); }
+#endif