From d001223250c1e6d26e6e4f1e917e0f8069345283 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Wed, 17 Jun 2026 16:49:56 +0100 Subject: [PATCH 1/3] Correctly synchronise all _sync allocate/deallocations - Closes #2448 --- ...failure_callback_resource_adaptor_impl.hpp | 12 ++++++++-- .../detail/stream_ordered_memory_resource.hpp | 8 ++++--- .../detail/aligned_resource_adaptor_impl.cpp | 12 ++++++++-- .../mr/detail/arena_memory_resource_impl.cpp | 10 +++++++-- .../detail/binning_memory_resource_impl.cpp | 4 ++-- .../detail/callback_memory_resource_impl.cpp | 13 +++++++++-- .../detail/limiting_resource_adaptor_impl.cpp | 12 ++++++++-- .../detail/logging_resource_adaptor_impl.cpp | 22 ++++++++++++------- .../detail/prefetch_resource_adaptor_impl.cpp | 13 +++++++++-- .../sam_headroom_memory_resource_impl.cpp | 5 ++++- .../statistics_resource_adaptor_impl.cpp | 13 +++++++++-- .../thread_safe_resource_adaptor_impl.cpp | 13 +++++++++-- .../detail/tracking_resource_adaptor_impl.cpp | 13 +++++++++-- 13 files changed, 118 insertions(+), 32 deletions(-) diff --git a/cpp/include/rmm/mr/detail/failure_callback_resource_adaptor_impl.hpp b/cpp/include/rmm/mr/detail/failure_callback_resource_adaptor_impl.hpp index 2cd9e2457..c324bde52 100644 --- a/cpp/include/rmm/mr/detail/failure_callback_resource_adaptor_impl.hpp +++ b/cpp/include/rmm/mr/detail/failure_callback_resource_adaptor_impl.hpp @@ -6,11 +6,14 @@ #include #include +#include #include #include #include #include +#include +#include #include #include @@ -86,14 +89,19 @@ class failure_callback_resource_adaptor_impl { void* allocate_sync(std::size_t bytes, std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) { - return allocate(cuda_stream_view{}, bytes, alignment); + auto const stream = cuda::stream_ref{cudaStream_t{nullptr}}; + auto* ptr = allocate(stream, bytes, alignment); + RMM_CUDA_TRY(cudaStreamSynchronize(stream.get())); + return ptr; } void deallocate_sync(void* ptr, std::size_t bytes, std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) noexcept { - deallocate(cuda_stream_view{}, ptr, bytes, alignment); + auto const stream = cuda::stream_ref{cudaStream_t{nullptr}}; + deallocate(stream, ptr, bytes, alignment); + RMM_ASSERT_CUDA_SUCCESS_SAFE_SHUTDOWN(cudaStreamSynchronize(stream.get())); } RMM_CONSTEXPR_FRIEND void get_property(failure_callback_resource_adaptor_impl const&, diff --git a/cpp/include/rmm/mr/detail/stream_ordered_memory_resource.hpp b/cpp/include/rmm/mr/detail/stream_ordered_memory_resource.hpp index 17c51ab6b..60dd88f44 100644 --- a/cpp/include/rmm/mr/detail/stream_ordered_memory_resource.hpp +++ b/cpp/include/rmm/mr/detail/stream_ordered_memory_resource.hpp @@ -164,9 +164,9 @@ class stream_ordered_memory_resource : public crtp { [[nodiscard]] void* allocate_sync(std::size_t bytes, std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) { - auto const stream = cuda_stream_view{}; + auto const stream = cuda::stream_ref{cudaStream_t{nullptr}}; void* ptr = allocate(stream, bytes, alignment); - stream.synchronize(); + RMM_CUDA_TRY(cudaStreamSynchronize(stream.get())); return ptr; } @@ -182,7 +182,9 @@ class stream_ordered_memory_resource : public crtp { std::size_t bytes, [[maybe_unused]] std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) noexcept { - deallocate(cuda_stream_view{}, ptr, bytes, alignment); + auto const stream = cuda::stream_ref{cudaStream_t{nullptr}}; + deallocate(stream, ptr, bytes, alignment); + RMM_ASSERT_CUDA_SUCCESS_SAFE_SHUTDOWN(cudaStreamSynchronize(stream.get())); } protected: diff --git a/cpp/src/mr/detail/aligned_resource_adaptor_impl.cpp b/cpp/src/mr/detail/aligned_resource_adaptor_impl.cpp index b99dc5750..52369e7e3 100644 --- a/cpp/src/mr/detail/aligned_resource_adaptor_impl.cpp +++ b/cpp/src/mr/detail/aligned_resource_adaptor_impl.cpp @@ -8,6 +8,9 @@ #include #include +#include +#include + #include #include @@ -89,14 +92,19 @@ void aligned_resource_adaptor_impl::deallocate(cuda::stream_ref stream, void* aligned_resource_adaptor_impl::allocate_sync(std::size_t bytes, std::size_t alignment) { - return allocate(cuda_stream_view{}, bytes, alignment); + auto const stream = cuda::stream_ref{cudaStream_t{nullptr}}; + auto* ptr = allocate(stream, bytes, alignment); + RMM_CUDA_TRY(cudaStreamSynchronize(stream.get())); + return ptr; } void aligned_resource_adaptor_impl::deallocate_sync(void* ptr, std::size_t bytes, std::size_t alignment) noexcept { - deallocate(cuda_stream_view{}, ptr, bytes, alignment); + auto const stream = cuda::stream_ref{cudaStream_t{nullptr}}; + deallocate(stream, ptr, bytes, alignment); + RMM_ASSERT_CUDA_SUCCESS_SAFE_SHUTDOWN(cudaStreamSynchronize(stream.get())); } } // namespace detail diff --git a/cpp/src/mr/detail/arena_memory_resource_impl.cpp b/cpp/src/mr/detail/arena_memory_resource_impl.cpp index d3f53eb8b..8d80d33e3 100644 --- a/cpp/src/mr/detail/arena_memory_resource_impl.cpp +++ b/cpp/src/mr/detail/arena_memory_resource_impl.cpp @@ -9,6 +9,7 @@ #include #include +#include #include namespace RMM_NAMESPACE { @@ -89,14 +90,19 @@ void arena_memory_resource_impl::deallocate(cuda::stream_ref stream, void* arena_memory_resource_impl::allocate_sync(std::size_t bytes, std::size_t alignment) { - return allocate(cuda_stream_view{}, bytes, alignment); + auto const stream = cuda::stream_ref{cudaStream_t{nullptr}}; + auto* ptr = allocate(stream, bytes, alignment); + RMM_CUDA_TRY(cudaStreamSynchronize(stream.get())); + return ptr; } void arena_memory_resource_impl::deallocate_sync(void* ptr, std::size_t bytes, std::size_t alignment) noexcept { - deallocate(cuda_stream_view{}, ptr, bytes, alignment); + auto const stream = cuda::stream_ref{cudaStream_t{nullptr}}; + deallocate(stream, ptr, bytes, alignment); + RMM_ASSERT_CUDA_SUCCESS_SAFE_SHUTDOWN(cudaStreamSynchronize(stream.get())); } void arena_memory_resource_impl::defragment() diff --git a/cpp/src/mr/detail/binning_memory_resource_impl.cpp b/cpp/src/mr/detail/binning_memory_resource_impl.cpp index 6cb9e7bc0..57f45d634 100644 --- a/cpp/src/mr/detail/binning_memory_resource_impl.cpp +++ b/cpp/src/mr/detail/binning_memory_resource_impl.cpp @@ -78,14 +78,14 @@ void binning_memory_resource_impl::deallocate(cuda::stream_ref stream, void* binning_memory_resource_impl::allocate_sync(std::size_t bytes, std::size_t alignment) { - return get_resource_ref(bytes).allocate(cuda_stream_view{}, bytes, alignment); + return get_resource_ref(bytes).allocate_sync(bytes, alignment); } void binning_memory_resource_impl::deallocate_sync(void* ptr, std::size_t bytes, std::size_t alignment) noexcept { - get_resource_ref(bytes).deallocate(cuda_stream_view{}, ptr, bytes, alignment); + return get_resource_ref(bytes).deallocate_sync(ptr, bytes, alignment); } } // namespace detail diff --git a/cpp/src/mr/detail/callback_memory_resource_impl.cpp b/cpp/src/mr/detail/callback_memory_resource_impl.cpp index 2f82ff7bf..d3948fa65 100644 --- a/cpp/src/mr/detail/callback_memory_resource_impl.cpp +++ b/cpp/src/mr/detail/callback_memory_resource_impl.cpp @@ -3,8 +3,12 @@ * SPDX-License-Identifier: Apache-2.0 */ +#include #include +#include +#include + #include namespace RMM_NAMESPACE { @@ -40,14 +44,19 @@ void callback_memory_resource_impl::deallocate(cuda::stream_ref stream, void* callback_memory_resource_impl::allocate_sync(std::size_t bytes, std::size_t alignment) { - return allocate(cuda_stream_view{}, bytes, alignment); + auto const stream = cuda::stream_ref{cudaStream_t{nullptr}}; + auto* ptr = allocate(stream, bytes, alignment); + RMM_CUDA_TRY(cudaStreamSynchronize(stream.get())); + return ptr; } void callback_memory_resource_impl::deallocate_sync(void* ptr, std::size_t bytes, std::size_t alignment) noexcept { - deallocate(cuda_stream_view{}, ptr, bytes, alignment); + auto const stream = cuda::stream_ref{cudaStream_t{nullptr}}; + deallocate(stream, ptr, bytes, alignment); + RMM_ASSERT_CUDA_SUCCESS_SAFE_SHUTDOWN(cudaStreamSynchronize(stream.get())); } } // namespace detail diff --git a/cpp/src/mr/detail/limiting_resource_adaptor_impl.cpp b/cpp/src/mr/detail/limiting_resource_adaptor_impl.cpp index ccc89d6d4..28539d608 100644 --- a/cpp/src/mr/detail/limiting_resource_adaptor_impl.cpp +++ b/cpp/src/mr/detail/limiting_resource_adaptor_impl.cpp @@ -8,6 +8,9 @@ #include #include +#include +#include + namespace RMM_NAMESPACE { namespace mr { namespace detail { @@ -69,14 +72,19 @@ void limiting_resource_adaptor_impl::deallocate(cuda::stream_ref stream, void* limiting_resource_adaptor_impl::allocate_sync(std::size_t bytes, std::size_t alignment) { - return allocate(cuda_stream_view{}, bytes, alignment); + auto const stream = cuda::stream_ref{cudaStream_t{nullptr}}; + auto* ptr = allocate(stream, bytes, alignment); + RMM_CUDA_TRY(cudaStreamSynchronize(stream.get())); + return ptr; } void limiting_resource_adaptor_impl::deallocate_sync(void* ptr, std::size_t bytes, std::size_t alignment) noexcept { - deallocate(cuda_stream_view{}, ptr, bytes, alignment); + auto const stream = cuda::stream_ref{cudaStream_t{nullptr}}; + deallocate(stream, ptr, bytes, alignment); + RMM_ASSERT_CUDA_SUCCESS_SAFE_SHUTDOWN(cudaStreamSynchronize(stream.get())); } } // namespace detail diff --git a/cpp/src/mr/detail/logging_resource_adaptor_impl.cpp b/cpp/src/mr/detail/logging_resource_adaptor_impl.cpp index 36b8f47e3..9c225e7f1 100644 --- a/cpp/src/mr/detail/logging_resource_adaptor_impl.cpp +++ b/cpp/src/mr/detail/logging_resource_adaptor_impl.cpp @@ -5,9 +5,13 @@ #include #include +#include #include #include +#include +#include + namespace RMM_NAMESPACE { namespace mr { namespace detail { @@ -26,24 +30,26 @@ logging_resource_adaptor_impl::logging_resource_adaptor_impl( void* logging_resource_adaptor_impl::allocate_sync(std::size_t bytes, std::size_t alignment) { - auto const stream = cuda_stream_view{}; + auto const stream = cuda::stream_ref{cudaStream_t{nullptr}}; + void* ptr{nullptr}; try { - auto const ptr = upstream_mr_.allocate(stream, bytes, alignment); - logger_->info("allocate,%p,%zu,%s", ptr, bytes, rmm::detail::format_stream(stream)); - return ptr; + ptr = allocate(stream, bytes, alignment); + RMM_CUDA_TRY(cudaStreamSynchronize(stream.get())); } catch (...) { - logger_->info("allocate failure,%p,%zu,%s", nullptr, bytes, rmm::detail::format_stream(stream)); + // TODO: Do we need this one? + RMM_CUDA_TRY(cudaStreamSynchronize(stream.get())); throw; } + return ptr; } void logging_resource_adaptor_impl::deallocate_sync(void* ptr, std::size_t bytes, std::size_t alignment) noexcept { - auto const stream = cuda_stream_view{}; - logger_->info("free,%p,%zu,%s", ptr, bytes, rmm::detail::format_stream(stream)); - upstream_mr_.deallocate(stream, ptr, bytes, alignment); + auto const stream = cuda::stream_ref{cudaStream_t{nullptr}}; + deallocate(stream, ptr, bytes, alignment); + RMM_ASSERT_CUDA_SUCCESS_SAFE_SHUTDOWN(cudaStreamSynchronize(stream.get())); } void* logging_resource_adaptor_impl::allocate(cuda::stream_ref stream, diff --git a/cpp/src/mr/detail/prefetch_resource_adaptor_impl.cpp b/cpp/src/mr/detail/prefetch_resource_adaptor_impl.cpp index 31be9763f..9fbc001f8 100644 --- a/cpp/src/mr/detail/prefetch_resource_adaptor_impl.cpp +++ b/cpp/src/mr/detail/prefetch_resource_adaptor_impl.cpp @@ -4,9 +4,13 @@ */ #include +#include #include #include +#include +#include + namespace RMM_NAMESPACE { namespace mr { namespace detail { @@ -42,14 +46,19 @@ void prefetch_resource_adaptor_impl::deallocate(cuda::stream_ref stream, void* prefetch_resource_adaptor_impl::allocate_sync(std::size_t bytes, std::size_t alignment) { - return allocate(cuda_stream_view{}, bytes, alignment); + auto const stream = cuda::stream_ref{cudaStream_t{nullptr}}; + auto* ptr = allocate(stream, bytes, alignment); + RMM_CUDA_TRY(cudaStreamSynchronize(stream.get())); + return ptr; } void prefetch_resource_adaptor_impl::deallocate_sync(void* ptr, std::size_t bytes, std::size_t alignment) noexcept { - deallocate(cuda_stream_view{}, ptr, bytes, alignment); + auto const stream = cuda::stream_ref{cudaStream_t{nullptr}}; + deallocate(stream, ptr, bytes, alignment); + RMM_ASSERT_CUDA_SUCCESS_SAFE_SHUTDOWN(cudaStreamSynchronize(stream.get())); } } // namespace detail diff --git a/cpp/src/mr/detail/sam_headroom_memory_resource_impl.cpp b/cpp/src/mr/detail/sam_headroom_memory_resource_impl.cpp index 0f1d180e6..a87a79881 100644 --- a/cpp/src/mr/detail/sam_headroom_memory_resource_impl.cpp +++ b/cpp/src/mr/detail/sam_headroom_memory_resource_impl.cpp @@ -9,6 +9,7 @@ #include #include +#include #include #include @@ -83,7 +84,9 @@ void sam_headroom_memory_resource_impl::deallocate_sync(void* ptr, std::size_t bytes, std::size_t alignment) noexcept { - deallocate(cuda::stream_ref{cudaStream_t{nullptr}}, ptr, bytes, alignment); + auto const stream = cuda::stream_ref{cudaStream_t{nullptr}}; + deallocate(stream, ptr, bytes, alignment); + RMM_ASSERT_CUDA_SUCCESS_SAFE_SHUTDOWN(cudaStreamSynchronize(stream.get())); } } // namespace detail diff --git a/cpp/src/mr/detail/statistics_resource_adaptor_impl.cpp b/cpp/src/mr/detail/statistics_resource_adaptor_impl.cpp index 094194c6e..233f42253 100644 --- a/cpp/src/mr/detail/statistics_resource_adaptor_impl.cpp +++ b/cpp/src/mr/detail/statistics_resource_adaptor_impl.cpp @@ -4,8 +4,12 @@ */ #include +#include #include +#include +#include + #include namespace RMM_NAMESPACE { @@ -87,14 +91,19 @@ void statistics_resource_adaptor_impl::deallocate(cuda::stream_ref stream, void* statistics_resource_adaptor_impl::allocate_sync(std::size_t bytes, std::size_t alignment) { - return allocate(cuda_stream_view{}, bytes, alignment); + auto const stream = cuda::stream_ref{cudaStream_t{nullptr}}; + auto* ptr = allocate(stream, bytes, alignment); + RMM_CUDA_TRY(cudaStreamSynchronize(stream.get())); + return ptr; } void statistics_resource_adaptor_impl::deallocate_sync(void* ptr, std::size_t bytes, std::size_t alignment) noexcept { - deallocate(cuda_stream_view{}, ptr, bytes, alignment); + auto const stream = cuda::stream_ref{cudaStream_t{nullptr}}; + deallocate(stream, ptr, bytes, alignment); + RMM_ASSERT_CUDA_SUCCESS_SAFE_SHUTDOWN(cudaStreamSynchronize(stream.get())); } } // namespace detail diff --git a/cpp/src/mr/detail/thread_safe_resource_adaptor_impl.cpp b/cpp/src/mr/detail/thread_safe_resource_adaptor_impl.cpp index 91c47e2d4..8b4bbbaae 100644 --- a/cpp/src/mr/detail/thread_safe_resource_adaptor_impl.cpp +++ b/cpp/src/mr/detail/thread_safe_resource_adaptor_impl.cpp @@ -3,8 +3,12 @@ * SPDX-License-Identifier: Apache-2.0 */ +#include #include +#include +#include + namespace RMM_NAMESPACE { namespace mr { namespace detail { @@ -40,14 +44,19 @@ void thread_safe_resource_adaptor_impl::deallocate(cuda::stream_ref stream, void* thread_safe_resource_adaptor_impl::allocate_sync(std::size_t bytes, std::size_t alignment) { - return allocate(cuda_stream_view{}, bytes, alignment); + auto const stream = cuda::stream_ref{cudaStream_t{nullptr}}; + auto* ptr = allocate(stream, bytes, alignment); + RMM_CUDA_TRY(cudaStreamSynchronize(stream.get())); + return ptr; } void thread_safe_resource_adaptor_impl::deallocate_sync(void* ptr, std::size_t bytes, std::size_t alignment) noexcept { - deallocate(cuda_stream_view{}, ptr, bytes, alignment); + auto const stream = cuda::stream_ref{cudaStream_t{nullptr}}; + deallocate(stream, ptr, bytes, alignment); + RMM_ASSERT_CUDA_SUCCESS_SAFE_SHUTDOWN(cudaStreamSynchronize(stream.get())); } } // namespace detail diff --git a/cpp/src/mr/detail/tracking_resource_adaptor_impl.cpp b/cpp/src/mr/detail/tracking_resource_adaptor_impl.cpp index b4b4c58be..6cd00395c 100644 --- a/cpp/src/mr/detail/tracking_resource_adaptor_impl.cpp +++ b/cpp/src/mr/detail/tracking_resource_adaptor_impl.cpp @@ -4,9 +4,13 @@ */ #include +#include #include #include +#include +#include + #include #include @@ -103,14 +107,19 @@ void tracking_resource_adaptor_impl::deallocate(cuda::stream_ref stream, void* tracking_resource_adaptor_impl::allocate_sync(std::size_t bytes, std::size_t alignment) { - return allocate(cuda_stream_view{}, bytes, alignment); + auto const stream = cuda::stream_ref{cudaStream_t{nullptr}}; + auto* ptr = allocate(stream, bytes, alignment); + RMM_CUDA_TRY(cudaStreamSynchronize(stream.get())); + return ptr; } void tracking_resource_adaptor_impl::deallocate_sync(void* ptr, std::size_t bytes, std::size_t alignment) noexcept { - deallocate(cuda_stream_view{}, ptr, bytes, alignment); + auto const stream = cuda::stream_ref{cudaStream_t{nullptr}}; + deallocate(stream, ptr, bytes, alignment); + RMM_ASSERT_CUDA_SUCCESS_SAFE_SHUTDOWN(cudaStreamSynchronize(stream.get())); } } // namespace detail From db44906e4e591861e92d7b8169ddf0f8a4563867 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Mon, 22 Jun 2026 09:32:44 +0100 Subject: [PATCH 2/3] Better limiting resource allocation breach error --- cpp/src/mr/detail/limiting_resource_adaptor_impl.cpp | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/cpp/src/mr/detail/limiting_resource_adaptor_impl.cpp b/cpp/src/mr/detail/limiting_resource_adaptor_impl.cpp index 28539d608..3dd4da839 100644 --- a/cpp/src/mr/detail/limiting_resource_adaptor_impl.cpp +++ b/cpp/src/mr/detail/limiting_resource_adaptor_impl.cpp @@ -55,9 +55,10 @@ void* limiting_resource_adaptor_impl::allocate(cuda::stream_ref stream, } allocated_bytes_ -= proposed_size; - auto const msg = std::string("Exceeded memory limit (failed to allocate ") + - rmm::detail::format_bytes(bytes) + ")"; - RMM_FAIL(msg.c_str(), rmm::out_of_memory); + std::stringstream msg; + msg << "Exceeded memory limit " << allocation_limit_ << "; Allocated bytes " << allocated_bytes_ + << "; Requested bytes " << bytes << "\n"; + RMM_FAIL(msg.str(), rmm::out_of_memory); } void limiting_resource_adaptor_impl::deallocate(cuda::stream_ref stream, From f117eb636c97dc17f333a96da18317277f8458c4 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Fri, 10 Jul 2026 11:14:24 +0100 Subject: [PATCH 3/3] Review changes --- cpp/src/mr/detail/binning_memory_resource_impl.cpp | 2 +- cpp/src/mr/detail/logging_resource_adaptor_impl.cpp | 11 ++--------- 2 files changed, 3 insertions(+), 10 deletions(-) diff --git a/cpp/src/mr/detail/binning_memory_resource_impl.cpp b/cpp/src/mr/detail/binning_memory_resource_impl.cpp index 57f45d634..c6fd2e09e 100644 --- a/cpp/src/mr/detail/binning_memory_resource_impl.cpp +++ b/cpp/src/mr/detail/binning_memory_resource_impl.cpp @@ -85,7 +85,7 @@ void binning_memory_resource_impl::deallocate_sync(void* ptr, std::size_t bytes, std::size_t alignment) noexcept { - return get_resource_ref(bytes).deallocate_sync(ptr, bytes, alignment); + get_resource_ref(bytes).deallocate_sync(ptr, bytes, alignment); } } // namespace detail diff --git a/cpp/src/mr/detail/logging_resource_adaptor_impl.cpp b/cpp/src/mr/detail/logging_resource_adaptor_impl.cpp index 9c225e7f1..edca8e310 100644 --- a/cpp/src/mr/detail/logging_resource_adaptor_impl.cpp +++ b/cpp/src/mr/detail/logging_resource_adaptor_impl.cpp @@ -31,15 +31,8 @@ logging_resource_adaptor_impl::logging_resource_adaptor_impl( void* logging_resource_adaptor_impl::allocate_sync(std::size_t bytes, std::size_t alignment) { auto const stream = cuda::stream_ref{cudaStream_t{nullptr}}; - void* ptr{nullptr}; - try { - ptr = allocate(stream, bytes, alignment); - RMM_CUDA_TRY(cudaStreamSynchronize(stream.get())); - } catch (...) { - // TODO: Do we need this one? - RMM_CUDA_TRY(cudaStreamSynchronize(stream.get())); - throw; - } + void* ptr = allocate(stream, bytes, alignment); + RMM_CUDA_TRY(cudaStreamSynchronize(stream.get())); return ptr; }