From de3a1df4767da1435afc521bd0f1ed7ccad1ce74 Mon Sep 17 00:00:00 2001
From: Mark Harris <mharris@nvidia.com>
Date: Wed, 4 Sep 2024 02:04:51 +0000
Subject: [PATCH] Add missing namespace and include

---
 cpp/bench/prims/neighbors/refine.cuh | 3 ++-
 cpp/bench/prims/random/subsample.cu  | 3 ++-
 cpp/template/src/ivf_flat_example.cu | 3 ++-
 cpp/template/src/ivf_pq_example.cu   | 3 ++-
 cpp/test/mr/device/buffer.cpp        | 5 +++--
 5 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/cpp/bench/prims/neighbors/refine.cuh b/cpp/bench/prims/neighbors/refine.cuh
index 14a01321ac..c29349927a 100644
--- a/cpp/bench/prims/neighbors/refine.cuh
+++ b/cpp/bench/prims/neighbors/refine.cuh
@@ -20,6 +20,7 @@
 
 #include <raft/core/device_mdspan.hpp>
 #include <raft/core/device_resources.hpp>
+#include <raft/core/resource/device_memory_resource.hpp>
 #include <raft/distance/distance_types.hpp>
 #include <raft/neighbors/detail/refine.cuh>
 #include <raft/neighbors/refine.cuh>
@@ -84,7 +85,7 @@ class RefineAnn : public fixture {
                                                               data.p.metric);
       });
     }
-    raft::set_current_device_resource(old_mr);
+    raft::resource::set_current_device_resource(old_mr);
   }
 
  private:
diff --git a/cpp/bench/prims/random/subsample.cu b/cpp/bench/prims/random/subsample.cu
index 7eef2b747c..ca0aa79101 100644
--- a/cpp/bench/prims/random/subsample.cu
+++ b/cpp/bench/prims/random/subsample.cu
@@ -20,6 +20,7 @@
 #include <raft/core/device_resources.hpp>
 #include <raft/core/host_mdarray.hpp>
 #include <raft/core/operators.hpp>
+#include <raft/core/resource/device_memory_resource.hpp>
 #include <raft/random/permute.cuh>
 #include <raft/random/rng.cuh>
 #include <raft/random/sample_without_replacement.cuh>
@@ -73,7 +74,7 @@ struct sample : public fixture {
     raft::random::RngState r(123456ULL);
   }
 
-  ~sample() { raft::set_current_device_resource(old_mr); }
+  ~sample() { raft::resource::set_current_device_resource(old_mr); }
   void run_benchmark(::benchmark::State& state) override
   {
     std::ostringstream label_stream;
diff --git a/cpp/template/src/ivf_flat_example.cu b/cpp/template/src/ivf_flat_example.cu
index 9fb0bfc708..1f6fe7dff8 100644
--- a/cpp/template/src/ivf_flat_example.cu
+++ b/cpp/template/src/ivf_flat_example.cu
@@ -18,6 +18,7 @@
 
 #include <raft/core/device_mdarray.hpp>
 #include <raft/core/device_resources.hpp>
+#include <raft/core/resource/device_memory_resource.hpp>
 #include <raft/core/resource/thrust_policy.hpp>
 #include <raft/neighbors/ivf_flat.cuh>
 #include <raft/util/cudart_utils.hpp>
@@ -133,7 +134,7 @@ int main()
   // Set pool memory resource with 1 GiB initial pool size. All allocations use the same pool.
   rmm::mr::pool_memory_resource<rmm::mr::device_memory_resource> pool_mr(
     raft::resource::get_current_device_resource(), 1024 * 1024 * 1024ull);
-  raft::set_current_device_resource(&pool_mr);
+  raft::resource::set_current_device_resource(&pool_mr);
 
   // Alternatively, one could define a pool allocator for temporary arrays (used within RAFT
   // algorithms). In that case only the internal arrays would use the pool, any other allocation
diff --git a/cpp/template/src/ivf_pq_example.cu b/cpp/template/src/ivf_pq_example.cu
index 8ded79ba2e..98fdc76c12 100644
--- a/cpp/template/src/ivf_pq_example.cu
+++ b/cpp/template/src/ivf_pq_example.cu
@@ -18,6 +18,7 @@
 
 #include <raft/core/device_mdarray.hpp>
 #include <raft/core/device_resources.hpp>
+#include <raft/core/resource/device_memory_resource.hpp>
 #include <raft/neighbors/ivf_pq.cuh>
 #include <raft/neighbors/refine.cuh>
 
@@ -93,7 +94,7 @@ int main()
   // Set pool memory resource with 1 GiB initial pool size. All allocations use the same pool.
   rmm::mr::pool_memory_resource<rmm::mr::device_memory_resource> pool_mr(
     raft::resource::get_current_device_resource(), 1024 * 1024 * 1024ull);
-  raft::set_current_device_resource(&pool_mr);
+  raft::resource::set_current_device_resource(&pool_mr);
 
   // Alternatively, one could define a pool allocator for temporary arrays (used within RAFT
   // algorithms). In that case only the internal arrays would use the pool, any other allocation
diff --git a/cpp/test/mr/device/buffer.cpp b/cpp/test/mr/device/buffer.cpp
index 2ba6ca112a..6c9a968c86 100644
--- a/cpp/test/mr/device/buffer.cpp
+++ b/cpp/test/mr/device/buffer.cpp
@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 
+#include <raft/core/resource/device_memory_resource.hpp>
 #include <raft/util/cudart_utils.hpp>
 
 #include <rmm/device_uvector.hpp>
@@ -64,7 +65,7 @@ TEST(Raft, DeviceBufferZeroResize)
     std::make_shared<rmm::mr::limiting_resource_adaptor<rmm::mr::cuda_memory_resource>>(curr_mr,
                                                                                         1000);
 
-  raft::set_current_device_resource(limit_mr.get());
+  raft::resource::set_current_device_resource(limit_mr.get());
 
   cudaStream_t stream;
   RAFT_CUDA_TRY(cudaStreamCreate(&stream));
@@ -84,7 +85,7 @@ TEST(Raft, DeviceBufferZeroResize)
   // Now check that there is no memory left. (Used to not be true)
   ASSERT_EQ(0, limit_mr->get_allocated_bytes());
 
-  raft::set_current_device_resource(curr_mr);
+  raft::resource::set_current_device_resource(curr_mr);
 
   RAFT_CUDA_TRY(cudaStreamSynchronize(stream));
   RAFT_CUDA_TRY(cudaStreamDestroy(stream));