Use the RMM ref aliases

rapidsai · Sep 4, 2024 · 5683243 · 5683243
1 parent 7a8141e
commit 5683243
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 10 deletions.
diff --git a/cpp/src/neighbors/detail/ann_utils.cuh b/cpp/src/neighbors/detail/ann_utils.cuh
@@ -557,14 +557,14 @@ struct batch_load_iterator {
    * @param mr a custom memory resource for the intermediate buffer, if applicable.
    * @param prefetch enable prefetch feature in order to achieve kernel/copy overlapping.
    */
-  batch_load_iterator(const T* source,
-                      size_type n_rows,
-                      size_type row_width,
-                      size_type batch_size,
-                      rmm::cuda_stream_view stream,
-                      raft::resource::device_async_resource_ref mr =
-                        raft::resource::get_current_device_resource_ref(),
-                      bool prefetch = false)
+  batch_load_iterator(
+    const T* source,
+    size_type n_rows,
+    size_type row_width,
+    size_type batch_size,
+    rmm::cuda_stream_view stream,
+    rmm::device_async_resource_ref mr = raft::resource::get_current_device_resource_ref(),
+    bool prefetch                     = false)
     : cur_batch_(new batch(source, n_rows, row_width, batch_size, stream, mr, prefetch)),
       cur_pos_(0),
       cur_prefetch_pos_(0)

diff --git a/cpp/src/neighbors/detail/cagra/utils.hpp b/cpp/src/neighbors/detail/cagra/utils.hpp
@@ -271,7 +271,7 @@ void copy_with_padding(
   raft::resources const& res,
   raft::device_matrix<T, int64_t, raft::row_major>& dst,
   raft::mdspan<const T, raft::matrix_extent<int64_t>, raft::row_major, data_accessor> src,
-  raft::resource::device_async_resource_ref mr = raft::resource::get_current_device_resource_ref())
+  rmm::device_async_resource_ref mr = raft::resource::get_current_device_resource_ref())
 {
   size_t padded_dim = raft::round_up_safe<size_t>(src.extent(1) * sizeof(T), 16) / sizeof(T);
 

diff --git a/cpp/test/neighbors/naive_knn.cuh b/cpp/test/neighbors/naive_knn.cuh
@@ -100,7 +100,7 @@ void naive_knn(raft::resources const& handle,
                uint32_t k,
                cuvs::distance::DistanceType type)
 {
-  raft::resource::device_async_resource_ref mr = rmm::mr::get_current_device_resource_ref();
+  rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource_ref();
 
   auto stream = raft::resource::get_cuda_stream(handle);
   dim3 block_dim(16, 32, 1);