Change rmm::exec_policy to take async_resource_ref (#1449)

Authors:
  - Michael Schellenberger Costa (https://github.com/miscco)

Approvers:
  - Mark Harris (https://github.com/harrism)

URL: #1449

include/rmm/exec_policy.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2021, NVIDIA CORPORATION.
+ * Copyright (c) 2020-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -23,6 +23,7 @@
 
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/mr/device/thrust_allocator_adaptor.hpp>
+#include <rmm/resource_ref.hpp>
 
 #include <rmm/detail/thrust_namespace.h>
 #include <thrust/system/cuda/execution_policy.h>
@@ -39,7 +40,7 @@ namespace rmm {
  * @brief Synchronous execution policy for allocations using thrust
  */
 using thrust_exec_policy_t =
-  thrust::detail::execute_with_allocator<rmm::mr::thrust_allocator<char>,
+  thrust::detail::execute_with_allocator<mr::thrust_allocator<char>,
                                          thrust::cuda_cub::execute_on_stream_base>;
 
 /**
@@ -54,10 +55,10 @@ class exec_policy : public thrust_exec_policy_t {
    * @param stream The stream on which to allocate temporary memory
    * @param mr The resource to use for allocating temporary memory
    */
-  explicit exec_policy(cuda_stream_view stream             = cuda_stream_default,
-                       rmm::mr::device_memory_resource* mr = mr::get_current_device_resource())
+  explicit exec_policy(cuda_stream_view stream      = cuda_stream_default,
+                       device_async_resource_ref mr = mr::get_current_device_resource())
     : thrust_exec_policy_t(
-        thrust::cuda::par(rmm::mr::thrust_allocator<char>(stream, mr)).on(stream.value()))
+        thrust::cuda::par(mr::thrust_allocator<char>(stream, mr)).on(stream.value()))
   {
   }
 };
@@ -68,7 +69,7 @@ class exec_policy : public thrust_exec_policy_t {
  * @brief Asynchronous execution policy for allocations using thrust
  */
 using thrust_exec_policy_nosync_t =
-  thrust::detail::execute_with_allocator<rmm::mr::thrust_allocator<char>,
+  thrust::detail::execute_with_allocator<mr::thrust_allocator<char>,
                                          thrust::cuda_cub::execute_on_stream_nosync_base>;
 /**
  * @brief Helper class usable as a Thrust CUDA execution policy
@@ -78,11 +79,10 @@ using thrust_exec_policy_nosync_t =
  */
 class exec_policy_nosync : public thrust_exec_policy_nosync_t {
  public:
-  explicit exec_policy_nosync(
-    cuda_stream_view stream             = cuda_stream_default,
-    rmm::mr::device_memory_resource* mr = mr::get_current_device_resource())
+  explicit exec_policy_nosync(cuda_stream_view stream      = cuda_stream_default,
+                              device_async_resource_ref mr = mr::get_current_device_resource())
     : thrust_exec_policy_nosync_t(
-        thrust::cuda::par_nosync(rmm::mr::thrust_allocator<char>(stream, mr)).on(stream.value()))
+        thrust::cuda::par_nosync(mr::thrust_allocator<char>(stream, mr)).on(stream.value()))
   {
   }
 };