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 9ea266077..59a9f0fd3 100644 --- a/cpp/include/rmm/mr/cuda_async_view_memory_resource.hpp +++ b/cpp/include/rmm/mr/cuda_async_view_memory_resource.hpp @@ -133,7 +133,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 e965922a2..fd9d20f09 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 @@ -71,7 +71,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 3509f9c4d..79dd65fc7 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