Adding pool back in and fixing cagra benchmark params (#1951)

Authors:
  - Corey J. Nolet (https://github.com/cjnolet)

Approvers:
  - Divye Gala (https://github.com/divyegala)

URL: #1951

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

@@ -19,6 +19,7 @@
 #include <algorithm>
 #include <cmath>
 #include <memory>
+#include <rmm/mr/device/pool_memory_resource.hpp>
 #include <stdexcept>
 #include <string>
 #include <type_traits>
@@ -272,5 +273,15 @@ 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); }
+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`
+  rmm::mr::device_memory_resource* mr =
+    rmm::mr::get_current_device_resource();  // Points to `pool_mr`
+  return raft::bench::ann::run_main(argc, argv);
+}
 #endif

python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_cagra.yaml

@@ -4,14 +4,14 @@ constraints:
 groups:
   base:
     build:
-      graph_degree: [32, 64]
-      intermediate_graph_degree: [64, 96]
+      graph_degree: [32, 64, 128, 256]
+      intermediate_graph_degree: [32, 64, 96, 128]
       graph_build_algo: ["NN_DESCENT"]
     search:
-      itopk: [32, 64, 128]
+      itopk: [32, 64, 128, 256, 512]
+      search_width: [1, 2, 4, 8, 16, 32, 64]
+
+
+
+
 
-  large:
-    build:
-      graph_degree: [32, 64]
-    search:
-      itopk: [32, 64, 128]