uninitialized_buffer::get_resource returns a ref to an `any_resourc…

…e` that can be copied (#2431)

* `uninitialized_buffer::get_resource` returns a ref to an `any_resource` that can be copied

* Also update `uninintialized_async_buffer`

* Fix doc string

---------

Co-authored-by: Michael Schellenberger Costa <[email protected]>

cudax/include/cuda/experimental/__container/uninitialized_async_buffer.cuh

@@ -72,7 +72,8 @@ template <class _Tp, class... _Properties>
 class uninitialized_async_buffer
 {
 private:
-  ::cuda::experimental::mr::async_any_resource<_Properties...> __mr_;
+  using __async_resource = ::cuda::experimental::mr::async_any_resource<_Properties...>;
+  __async_resource __mr_;
   ::cuda::stream_ref __stream_ = {};
   size_t __count_              = 0;
   void* __buf_                 = nullptr;
@@ -127,9 +128,7 @@ public:
   //! @param __count The desired size of the buffer.
   //! @note Depending on the alignment requirements of `T` the size of the underlying allocation might be larger
   //! than `count * sizeof(T)`. Only allocates memory when \p __count > 0
-  uninitialized_async_buffer(::cuda::experimental::mr::async_any_resource<_Properties...> __mr,
-                             const ::cuda::stream_ref __stream,
-                             const size_t __count)
+  uninitialized_async_buffer(__async_resource __mr, const ::cuda::stream_ref __stream, const size_t __count)
       : __mr_(_CUDA_VSTD::move(__mr))
       , __stream_(__stream)
       , __count_(__count)
@@ -205,12 +204,12 @@ public:
   }
 
   //! @rst
-  //! Returns an :ref:`asnyc_resource_ref <libcudacxx-extended-api-memory-resources-resource-ref>` to the resource used
-  //! to allocate the buffer
+  //! Returns a \c const reference to the :ref:`any_async_resource <cudax-memory-resource-async-any-resource>`
+  //! that holds the memory resource used to allocate the buffer
   //! @endrst
-  _CCCL_NODISCARD _CUDA_VMR::async_resource_ref<_Properties...> get_resource() const noexcept
+  _CCCL_NODISCARD const __async_resource& get_resource() const noexcept
   {
-    return _CUDA_VMR::async_resource_ref<_Properties...>{const_cast<uninitialized_async_buffer*>(this)->__mr_};
+    return __mr_;
   }
 
   //! @brief Returns the stored stream

cudax/include/cuda/experimental/__container/uninitialized_buffer.cuh

@@ -63,7 +63,8 @@ template <class _Tp, class... _Properties>
 class uninitialized_buffer
 {
 private:
-  ::cuda::experimental::mr::any_resource<_Properties...> __mr_;
+  using __resource = ::cuda::experimental::mr::any_resource<_Properties...>;
+  __resource __mr_;
   size_t __count_ = 0;
   void* __buf_    = nullptr;
 
@@ -116,7 +117,7 @@ public:
   //! @note Depending on the alignment requirements of `T` the size of the underlying allocation might be larger
   //! than `count * sizeof(T)`.
   //! @note Only allocates memory when \p __count > 0
-  uninitialized_buffer(::cuda::experimental::mr::any_resource<_Properties...> __mr, const size_t __count)
+  uninitialized_buffer(__resource __mr, const size_t __count)
       : __mr_(_CUDA_VSTD::move(__mr))
       , __count_(__count)
       , __buf_(__count_ == 0 ? nullptr : __mr_.allocate(__get_allocation_size(__count_)))
@@ -188,13 +189,13 @@ public:
   }
 
   //! @rst
-  //! Returns a :ref:`resource_ref <libcudacxx-extended-api-memory-resources-resource-ref>` to the resource used to
-  //! allocate the buffer
+  //! Returns a \c const reference to the :ref:`any_resource <cudax-memory-resource-any-resource>`
+  //! that holds the memory resource used to allocate the buffer
   //! @endrst
   _CCCL_EXEC_CHECK_DISABLE
-  _CCCL_NODISCARD _CCCL_HOST_DEVICE _CUDA_VMR::resource_ref<_Properties...> get_resource() const noexcept
+  _CCCL_NODISCARD _CCCL_HOST_DEVICE const __resource& get_resource() const noexcept
   {
-    return _CUDA_VMR::resource_ref<_Properties...>{const_cast<uninitialized_buffer*>(this)->__mr_};
+    return __mr_;
   }
 
   //! @brief Swaps the contents with those of another \c uninitialized_buffer

cudax/test/containers/uninitialized_async_buffer.cu

@@ -156,3 +156,53 @@ TEMPLATE_TEST_CASE(
     }
   }
 }
+
+// A test resource that keeps track of the number of resources are
+// currently alive.
+struct test_async_memory_resource : cudax::mr::async_memory_resource
+{
+  static int count;
+
+  test_async_memory_resource()
+  {
+    ++count;
+  }
+
+  test_async_memory_resource(const test_async_memory_resource& other)
+      : cudax::mr::async_memory_resource{other}
+  {
+    ++count;
+  }
+
+  ~test_async_memory_resource()
+  {
+    --count;
+  }
+};
+
+int test_async_memory_resource::count = 0;
+
+TEST_CASE("uninitialized_async_buffer's memory resource does not dangle", "[container]")
+{
+  cuda::experimental::stream stream{};
+  cudax::uninitialized_async_buffer<int, ::cuda::mr::device_accessible> buffer{
+    cudax::mr::async_memory_resource{}, stream, 0};
+
+  {
+    CHECK(test_async_memory_resource::count == 0);
+
+    cudax::uninitialized_async_buffer<int, ::cuda::mr::device_accessible> src_buffer{
+      test_async_memory_resource{}, stream, 1024};
+
+    CHECK(test_async_memory_resource::count == 1);
+
+    cudax::uninitialized_async_buffer<int, ::cuda::mr::device_accessible> dst_buffer{
+      src_buffer.get_resource(), stream, 1024};
+
+    CHECK(test_async_memory_resource::count == 2);
+
+    buffer = ::cuda::std::move(dst_buffer);
+  }
+
+  CHECK(test_async_memory_resource::count == 1);
+}

cudax/test/containers/uninitialized_buffer.cu

@@ -200,3 +200,49 @@ TEST_CASE("uninitialized_buffer is usable with cudax::launch", "[container]")
     cudax::launch(stream, dimensions, const_kernel, buffer);
   }
 }
+
+// A test resource that keeps track of the number of resources are
+// currently alive.
+struct test_device_memory_resource : cuda::mr::device_memory_resource
+{
+  static int count;
+
+  test_device_memory_resource()
+  {
+    ++count;
+  }
+
+  test_device_memory_resource(const test_device_memory_resource& other)
+      : cuda::mr::device_memory_resource{other}
+  {
+    ++count;
+  }
+
+  ~test_device_memory_resource()
+  {
+    --count;
+  }
+};
+
+int test_device_memory_resource::count = 0;
+
+TEST_CASE("uninitialized_buffer's memory resource does not dangle", "[container]")
+{
+  cudax::uninitialized_buffer<int, ::cuda::mr::device_accessible> buffer{cuda::mr::device_memory_resource{}, 0};
+
+  {
+    CHECK(test_device_memory_resource::count == 0);
+
+    cudax::uninitialized_buffer<int, ::cuda::mr::device_accessible> src_buffer{test_device_memory_resource{}, 1024};
+
+    CHECK(test_device_memory_resource::count == 1);
+
+    cudax::uninitialized_buffer<int, ::cuda::mr::device_accessible> dst_buffer{src_buffer.get_resource(), 1024};
+
+    CHECK(test_device_memory_resource::count == 2);
+
+    buffer = ::cuda::std::move(dst_buffer);
+  }
+
+  CHECK(test_device_memory_resource::count == 1);
+}