From 32a602677f02af335da4d4ad582f6e6501e81345 Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@gmail.com>
Date: Thu, 2 Nov 2023 03:18:16 +0100
Subject: [PATCH] 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: https://github.com/rapidsai/raft/pull/1951
---
 cpp/bench/ann/src/raft/raft_benchmark.cu         | 13 ++++++++++++-
 .../run/conf/algos/raft_cagra.yaml               | 16 ++++++++--------
 2 files changed, 20 insertions(+), 9 deletions(-)

diff --git a/cpp/bench/ann/src/raft/raft_benchmark.cu b/cpp/bench/ann/src/raft/raft_benchmark.cu
index fa20c5c223..6888340b4d 100644
--- a/cpp/bench/ann/src/raft/raft_benchmark.cu
+++ b/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
diff --git a/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_cagra.yaml b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_cagra.yaml
index 927b90557a..0f80608eef 100644
--- a/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_cagra.yaml
+++ b/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]