Add YAML config files to run parameter sweeps for ANN benchmarks

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

@@ -12,6 +12,7 @@ dependencies:
 - clang-tools=16.0.6
 - clang==16.0.6
 - cmake>=3.26.4
+- cuda-nvtx=11.8
 - cuda-profiler-api=11.8.86
 - cuda-python>=11.7.1,<12.0a0
 - cuda-version=11.8

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

@@ -12,6 +12,7 @@ dependencies:
 - clang-tools=16.0.6
 - clang==16.0.6
 - cmake>=3.26.4
+- cuda-nvtx=11.8
 - cuda-profiler-api=11.8.86
 - cuda-python>=11.7.1,<12.0a0
 - cuda-version=11.8

conda/environments/all_cuda-120_arch-aarch64.yaml

@@ -14,6 +14,7 @@ dependencies:
 - cmake>=3.26.4
 - cuda-cudart-dev
 - cuda-nvcc
+- cuda-nvtx-dev
 - cuda-profiler-api
 - cuda-python>=12.0,<13.0a0
 - cuda-version=12.0

conda/environments/all_cuda-120_arch-x86_64.yaml

@@ -14,6 +14,7 @@ dependencies:
 - cmake>=3.26.4
 - cuda-cudart-dev
 - cuda-nvcc
+- cuda-nvtx-dev
 - cuda-profiler-api
 - cuda-python>=12.0,<13.0a0
 - cuda-version=12.0

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

@@ -12,6 +12,7 @@ dependencies:
 - clang-tools=16.0.6
 - clang==16.0.6
 - cmake>=3.26.4
+- cuda-nvtx=11.8
 - cuda-profiler-api=11.8.86
 - cuda-version=11.8
 - cudatoolkit

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

@@ -12,6 +12,7 @@ dependencies:
 - clang-tools=16.0.6
 - clang==16.0.6
 - cmake>=3.26.4
+- cuda-nvtx=11.8
 - cuda-profiler-api=11.8.86
 - cuda-version=11.8
 - cudatoolkit

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

@@ -0,0 +1,40 @@
+# This file is generated by `rapids-dependency-file-generator`.
+# To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.
+channels:
+- rapidsai
+- rapidsai-nightly
+- dask/label/dev
+- conda-forge
+- nvidia
+dependencies:
+- benchmark>=1.8.2
+- c-compiler
+- clang-tools=16.0.6
+- clang==16.0.6
+- cmake>=3.26.4
+- cuda-cudart-dev
+- cuda-nvcc
+- cuda-nvtx-dev
+- cuda-profiler-api
+- cuda-version=12.0
+- cxx-compiler
+- cython>=3.0.0
+- gcc_linux-aarch64=11.*
+- glog>=0.6.0
+- h5py>=3.8.0
+- hnswlib=0.7.0
+- libcublas-dev
+- libcurand-dev
+- libcusolver-dev
+- libcusparse-dev
+- matplotlib
+- nccl>=2.9.9
+- ninja
+- nlohmann_json>=3.11.2
+- openblas
+- pandas
+- pyyaml
+- rmm==23.12.*
+- scikit-build>=0.13.1
+- sysroot_linux-aarch64==2.17
+name: bench_ann_cuda-120_arch-aarch64

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

@@ -0,0 +1,40 @@
+# This file is generated by `rapids-dependency-file-generator`.
+# To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.
+channels:
+- rapidsai
+- rapidsai-nightly
+- dask/label/dev
+- conda-forge
+- nvidia
+dependencies:
+- benchmark>=1.8.2
+- c-compiler
+- clang-tools=16.0.6
+- clang==16.0.6
+- cmake>=3.26.4
+- cuda-cudart-dev
+- cuda-nvcc
+- cuda-nvtx-dev
+- cuda-profiler-api
+- cuda-version=12.0
+- cxx-compiler
+- cython>=3.0.0
+- gcc_linux-64=11.*
+- glog>=0.6.0
+- h5py>=3.8.0
+- hnswlib=0.7.0
+- libcublas-dev
+- libcurand-dev
+- libcusolver-dev
+- libcusparse-dev
+- matplotlib
+- nccl>=2.9.9
+- ninja
+- nlohmann_json>=3.11.2
+- openblas
+- pandas
+- pyyaml
+- rmm==23.12.*
+- scikit-build>=0.13.1
+- sysroot_linux-64==2.17
+name: bench_ann_cuda-120_arch-x86_64

cpp/bench/ann/src/faiss/faiss_cpu_benchmark.cpp

@@ -49,7 +49,7 @@ void parse_build_param(const nlohmann::json& conf,
                        typename raft::bench::ann::FaissCpuIVFPQ<T>::BuildParam& param)
 {
   parse_base_build_param<T>(conf, param);
-  param.M = conf.at("M");
+  param.M_ratio = conf.at("M_ratio");
   if (conf.contains("usePrecomputed")) {
     param.usePrecomputed = conf.at("usePrecomputed");
   } else {

cpp/bench/ann/src/faiss/faiss_cpu_wrapper.h

@@ -229,16 +229,20 @@ template <typename T>
 class FaissCpuIVFPQ : public FaissCpu<T> {
  public:
   struct BuildParam : public FaissCpu<T>::BuildParam {
-    int M;
+    int M_ratio;
     int bitsPerCode;
     bool usePrecomputed;
   };
 
   FaissCpuIVFPQ(Metric metric, int dim, const BuildParam& param) : FaissCpu<T>(metric, dim, param)
   {
     this->init_quantizer(dim);
-    this->index_ = std::make_unique<faiss::IndexIVFPQ>(
-      this->quantizer_.get(), dim, param.nlist, param.M, param.bitsPerCode, this->metric_type_);
+    this->index_ = std::make_unique<faiss::IndexIVFPQ>(this->quantizer_.get(),
+                                                       dim,
+                                                       param.nlist,
+                                                       dim / param.M_ratio,
+                                                       param.bitsPerCode,
+                                                       this->metric_type_);
   }
 
   void save(const std::string& file) const override

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

@@ -50,7 +50,7 @@ void parse_build_param(const nlohmann::json& conf,
                        typename raft::bench::ann::FaissGpuIVFPQ<T>::BuildParam& param)
 {
   parse_base_build_param<T>(conf, param);
-  param.M = conf.at("M");
+  param.M_ratio = conf.at("M_ratio");
   if (conf.contains("usePrecomputed")) {
     param.usePrecomputed = conf.at("usePrecomputed");
   } else {

cpp/bench/ann/src/faiss/faiss_gpu_wrapper.h

@@ -263,7 +263,7 @@ template <typename T>
 class FaissGpuIVFPQ : public FaissGpu<T> {
  public:
   struct BuildParam : public FaissGpu<T>::BuildParam {
-    int M;
+    int M_ratio;
     bool useFloat16;
     bool usePrecomputed;
   };
@@ -274,11 +274,12 @@ class FaissGpuIVFPQ : public FaissGpu<T> {
     config.useFloat16LookupTables = param.useFloat16;
     config.usePrecomputedTables   = param.usePrecomputed;
     config.device                 = this->device_;
+
     this->index_ =
       std::make_unique<faiss::gpu::GpuIndexIVFPQ>(&(this->gpu_resource_),
                                                   dim,
                                                   param.nlist,
-                                                  param.M,
+                                                  dim / param.M_ratio,
                                                   8,  // FAISS only supports bitsPerCode=8
                                                   this->metric_type_,
                                                   config);

cpp/bench/ann/src/raft/raft_benchmark.cu

@@ -272,13 +272,5 @@ REGISTER_ALGO_INSTANCE(std::uint8_t);
 
 #ifdef ANN_BENCH_BUILD_MAIN
 #include "../common/benchmark.hpp"
-int main(int argc, char** argv)
-{
-  rmm::mr::cuda_memory_resource cuda_mr;
-  // Construct a resource that uses a coalescing best-fit pool allocator
-  rmm::mr::pool_memory_resource<rmm::mr::cuda_memory_resource> pool_mr{&cuda_mr};
-  rmm::mr::set_current_device_resource(
-    &pool_mr);  // Updates the current device resource pointer to `pool_mr`
-  return raft::bench::ann::run_main(argc, argv);
-}
+int main(int argc, char** argv) { return raft::bench::ann::run_main(argc, argv); }
 #endif

dependencies.yaml

@@ -22,7 +22,7 @@ files:
   bench_ann:
     output: conda
     matrix:
-      cuda: ["11.8"]
+      cuda: ["11.8", "12.0"]
       arch: [x86_64, aarch64]
     includes:
       - build
@@ -246,6 +246,7 @@ dependencies:
               cuda: "12.0"
             packages:
               - cuda-version=12.0
+              - cuda-nvtx-dev
               - cuda-cudart-dev
               - cuda-profiler-api
               - libcublas-dev
@@ -257,6 +258,7 @@ dependencies:
             packages:
               - cuda-version=11.8
               - cudatoolkit
+              - cuda-nvtx=11.8
               - cuda-profiler-api=11.8.86
               - libcublas-dev=11.11.3.6
               - libcublas=11.11.3.6
@@ -271,6 +273,7 @@ dependencies:
             packages:
               - cuda-version=11.5
               - cudatoolkit
+              - cuda-nvtx=11.5
               - cuda-profiler-api>=11.4.240,<=11.8.86 # use any `11.x` version since pkg is missing several CUDA/arch packages
               - libcublas-dev>=11.7.3.1,<=11.7.4.6
               - libcublas>=11.7.3.1,<=11.7.4.6
@@ -285,6 +288,7 @@ dependencies:
             packages:
               - cuda-version=11.4
               - cudatoolkit
+              - &cudanvtx114 cuda-nvtx=11.4
               - cuda-profiler-api>=11.4.240,<=11.8.86 # use any `11.x` version since pkg is missing several CUDA/arch packages
               - &libcublas_dev114 libcublas-dev>=11.5.2.43,<=11.6.5.2
               - &libcublas114 libcublas>=11.5.2.43,<=11.6.5.2
@@ -299,6 +303,7 @@ dependencies:
             packages:
               - cuda-version=11.2
               - cudatoolkit
+              - *cudanvtx114
               - cuda-profiler-api>=11.4.240,<=11.8.86 # use any `11.x` version since pkg is missing several CUDA/arch packages
               # The NVIDIA channel doesn't publish pkgs older than 11.4 for these libs,
               # so 11.2 uses 11.4 packages (the oldest available).