From 445fd71ad66435ff7db7df166b6cd7e36a7c7f9c Mon Sep 17 00:00:00 2001 From: Eric Niebler Date: Thu, 19 Sep 2024 07:15:47 -0700 Subject: [PATCH] `uninitialized_buffer::get_resource` returns a ref to an `any_resource` 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 --- .../uninitialized_async_buffer.cuh | 15 +++--- .../__container/uninitialized_buffer.cuh | 13 ++--- .../containers/uninitialized_async_buffer.cu | 50 +++++++++++++++++++ cudax/test/containers/uninitialized_buffer.cu | 46 +++++++++++++++++ 4 files changed, 110 insertions(+), 14 deletions(-) diff --git a/cudax/include/cuda/experimental/__container/uninitialized_async_buffer.cuh b/cudax/include/cuda/experimental/__container/uninitialized_async_buffer.cuh index eea30c1b69..54115e4ccd 100644 --- a/cudax/include/cuda/experimental/__container/uninitialized_async_buffer.cuh +++ b/cudax/include/cuda/experimental/__container/uninitialized_async_buffer.cuh @@ -72,7 +72,8 @@ template 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 ` to the resource used - //! to allocate the buffer + //! Returns a \c const reference to the :ref:`any_async_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(this)->__mr_}; + return __mr_; } //! @brief Returns the stored stream diff --git a/cudax/include/cuda/experimental/__container/uninitialized_buffer.cuh b/cudax/include/cuda/experimental/__container/uninitialized_buffer.cuh index 9c88df1d95..9cce7e706e 100644 --- a/cudax/include/cuda/experimental/__container/uninitialized_buffer.cuh +++ b/cudax/include/cuda/experimental/__container/uninitialized_buffer.cuh @@ -63,7 +63,8 @@ template 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 ` to the resource used to - //! allocate the buffer + //! Returns a \c const reference to the :ref:`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(this)->__mr_}; + return __mr_; } //! @brief Swaps the contents with those of another \c uninitialized_buffer diff --git a/cudax/test/containers/uninitialized_async_buffer.cu b/cudax/test/containers/uninitialized_async_buffer.cu index b32d322827..8e1f6e304b 100644 --- a/cudax/test/containers/uninitialized_async_buffer.cu +++ b/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 buffer{ + cudax::mr::async_memory_resource{}, stream, 0}; + + { + CHECK(test_async_memory_resource::count == 0); + + cudax::uninitialized_async_buffer src_buffer{ + test_async_memory_resource{}, stream, 1024}; + + CHECK(test_async_memory_resource::count == 1); + + cudax::uninitialized_async_buffer 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); +} diff --git a/cudax/test/containers/uninitialized_buffer.cu b/cudax/test/containers/uninitialized_buffer.cu index c924750a8a..e8aecad470 100644 --- a/cudax/test/containers/uninitialized_buffer.cu +++ b/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 buffer{cuda::mr::device_memory_resource{}, 0}; + + { + CHECK(test_device_memory_resource::count == 0); + + cudax::uninitialized_buffer src_buffer{test_device_memory_resource{}, 1024}; + + CHECK(test_device_memory_resource::count == 1); + + cudax::uninitialized_buffer 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); +}