Adopt rmm::device_async_resource_ref alias

rapidsai · Feb 1, 2024 · 58d276e · 58d276e
1 parent 274f0b5
commit 58d276e
Showing 1 changed file with 4 additions and 7 deletions.
diff --git a/include/rmm/device_scalar.hpp b/include/rmm/device_scalar.hpp
@@ -19,8 +19,7 @@
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/device_uvector.hpp>
 #include <rmm/mr/device/per_device_resource.hpp>
-
-#include <cuda/memory_resource>
+#include <rmm/resource_ref.hpp>
 
 #include <type_traits>
 
@@ -40,8 +39,6 @@ namespace rmm {
  */
 template <typename T>
 class device_scalar {
-  using async_resource_ref = cuda::mr::async_resource_ref<cuda::mr::device_accessible>;
-
  public:
   static_assert(std::is_trivially_copyable<T>::value, "Scalar type must be trivially copyable");
 
@@ -96,7 +93,7 @@ class device_scalar {
    * @param mr Optional, resource with which to allocate.
    */
   explicit device_scalar(cuda_stream_view stream,
-                         async_resource_ref mr = rmm::mr::get_current_device_resource())
+                         rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource())
     : _storage{1, stream, mr}
   {
   }
@@ -119,7 +116,7 @@ class device_scalar {
    */
   explicit device_scalar(value_type const& initial_value,
                          cuda_stream_view stream,
-                         async_resource_ref mr = rmm::mr::get_current_device_resource())
+                         rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource())
     : _storage{1, stream, mr}
   {
     set_value_async(initial_value, stream);
@@ -139,7 +136,7 @@ class device_scalar {
    */
   device_scalar(device_scalar const& other,
                 cuda_stream_view stream,
-                async_resource_ref mr = rmm::mr::get_current_device_resource())
+                rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource())
     : _storage{other._storage, stream, mr}
   {
   }