From ec97f1a7b1309b07398fa281997493cf3ded2e8b Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sun, 17 May 2026 17:04:37 +0000 Subject: [PATCH] Synchronize cuda async view deallocate_sync --- cpp/include/rmm/mr/cuda_async_view_memory_resource.hpp | 4 +++- cpp/src/mr/detail/cuda_async_managed_memory_resource_impl.cpp | 4 +++- cpp/src/mr/detail/cuda_async_memory_resource_impl.cpp | 4 +++- 3 files changed, 9 insertions(+), 3 deletions(-) diff --git a/cpp/include/rmm/mr/cuda_async_view_memory_resource.hpp b/cpp/include/rmm/mr/cuda_async_view_memory_resource.hpp index e2aa54b45..718603ba8 100644 --- a/cpp/include/rmm/mr/cuda_async_view_memory_resource.hpp +++ b/cpp/include/rmm/mr/cuda_async_view_memory_resource.hpp @@ -131,7 +131,9 @@ class cuda_async_view_memory_resource final { std::size_t bytes, std::size_t alignment = rmm::CUDA_ALLOCATION_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())); } /** diff --git a/cpp/src/mr/detail/cuda_async_managed_memory_resource_impl.cpp b/cpp/src/mr/detail/cuda_async_managed_memory_resource_impl.cpp index 86d39a335..fb945fb9b 100644 --- a/cpp/src/mr/detail/cuda_async_managed_memory_resource_impl.cpp +++ b/cpp/src/mr/detail/cuda_async_managed_memory_resource_impl.cpp @@ -66,7 +66,9 @@ void cuda_async_managed_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/cuda_async_memory_resource_impl.cpp b/cpp/src/mr/detail/cuda_async_memory_resource_impl.cpp index a2a66b5d7..0b939a345 100644 --- a/cpp/src/mr/detail/cuda_async_memory_resource_impl.cpp +++ b/cpp/src/mr/detail/cuda_async_memory_resource_impl.cpp @@ -101,7 +101,9 @@ void cuda_async_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