Set pool memory resource for raft IVF ANN benchmarks (#1625)

This PR fixes perf issues by #1558 Authors: - Tamas Bela Feher (https://github.com/tfeher) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: #1625
rapidsai · Jun 30, 2023 · 58631df · 58631df
1 parent 503fc43
commit 58631df
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/cpp/bench/ann/src/raft/raft_ivf_flat_wrapper.h b/cpp/bench/ann/src/raft/raft_ivf_flat_wrapper.h
@@ -97,6 +97,7 @@ RaftIvfFlatGpu<T, IdxT>::RaftIvfFlatGpu(Metric metric, int dim, const BuildParam
     mr_(rmm::mr::get_current_device_resource(), 1024 * 1024 * 1024ull)
 {
   index_params_.metric = parse_metric_type(metric);
+  rmm::mr::set_current_device_resource(&mr_);
   RAFT_CUDA_TRY(cudaGetDevice(&device_));
 }
 

diff --git a/cpp/bench/ann/src/raft/raft_ivf_pq_wrapper.h b/cpp/bench/ann/src/raft/raft_ivf_pq_wrapper.h
@@ -98,6 +98,7 @@ RaftIvfPQ<T, IdxT>::RaftIvfPQ(Metric metric, int dim, const BuildParam& param, f
     refine_ratio_(refine_ratio),
     mr_(rmm::mr::get_current_device_resource(), 1024 * 1024 * 1024ull)
 {
+  rmm::mr::set_current_device_resource(&mr_);
   index_params_.metric = parse_metric_type(metric);
   RAFT_CUDA_TRY(cudaGetDevice(&device_));
 }