Skip to content

Correctly synchronise all _sync allocate/deallocations #2449

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
rapids-bot merged 4 commits into from
Jul 10, 2026
Merged

Correctly synchronise all _sync allocate/deallocations #2449

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
d001223
Correctly synchronise all _sync allocate/deallocations
wence- Jun 17, 2026
db44906
Better limiting resource allocation breach error
wence- Jun 22, 2026
542bf4b
Merge branch 'main' into wence/fix/2448
lamarrr Jul 9, 2026
f117eb6
Review changes
wence- Jul 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 10 additions & 2 deletions cpp/include/rmm/mr/detail/failure_callback_resource_adaptor_impl.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,11 +6,14 @@

#include <rmm/aligned.hpp>
#include <rmm/cuda_stream_view.hpp>
#include <rmm/detail/error.hpp>
#include <rmm/detail/export.hpp>
#include <rmm/mr/failure_callback_t.hpp>
#include <rmm/resource_ref.hpp>

#include <cuda/memory_resource>
#include <cuda/stream_ref>
#include <cuda_runtime_api.h>

#include <cstddef>
#include <utility>
Expand Down Expand Up @@ -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&,
Expand Down
8 changes: 5 additions & 3 deletions cpp/include/rmm/mr/detail/stream_ordered_memory_resource.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -164,9 +164,9 @@ class stream_ordered_memory_resource : public crtp<PoolResource> {
[[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;
}

Expand All @@ -182,7 +182,9 @@ class stream_ordered_memory_resource : public crtp<PoolResource> {
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:
Expand Down
12 changes: 10 additions & 2 deletions cpp/src/mr/detail/aligned_resource_adaptor_impl.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,9 @@
#include <rmm/detail/error.hpp>
#include <rmm/mr/detail/aligned_resource_adaptor_impl.hpp>

#include <cuda/stream_ref>
#include <cuda_runtime_api.h>

#include <algorithm>
#include <cstddef>

Expand Down Expand Up @@ -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
Expand Down
10 changes: 8 additions & 2 deletions cpp/src/mr/detail/arena_memory_resource_impl.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
#include <rmm/logger.hpp>
#include <rmm/mr/detail/arena_memory_resource_impl.hpp>

#include <cuda/stream_ref>
#include <cuda_runtime_api.h>

namespace RMM_NAMESPACE {
Expand Down Expand Up @@ -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()
Expand Down
4 changes: 2 additions & 2 deletions cpp/src/mr/detail/binning_memory_resource_impl.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
get_resource_ref(bytes).deallocate_sync(ptr, bytes, alignment);
}

} // namespace detail
Expand Down
13 changes: 11 additions & 2 deletions cpp/src/mr/detail/callback_memory_resource_impl.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,12 @@
* SPDX-License-Identifier: Apache-2.0
*/

#include <rmm/detail/error.hpp>
#include <rmm/mr/detail/callback_memory_resource_impl.hpp>

#include <cuda/stream_ref>
#include <cuda_runtime_api.h>

#include <utility>

namespace RMM_NAMESPACE {
Expand Down Expand Up @@ -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
Expand Down
19 changes: 14 additions & 5 deletions cpp/src/mr/detail/limiting_resource_adaptor_impl.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,9 @@
#include <rmm/detail/format.hpp>
#include <rmm/mr/detail/limiting_resource_adaptor_impl.hpp>

#include <cuda/stream_ref>
#include <cuda_runtime_api.h>

namespace RMM_NAMESPACE {
namespace mr {
namespace detail {
Expand Down Expand Up @@ -52,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);
Comment thread
wence- marked this conversation as resolved.
}

void limiting_resource_adaptor_impl::deallocate(cuda::stream_ref stream,
Expand All @@ -69,14 +73,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
Expand Down
23 changes: 11 additions & 12 deletions cpp/src/mr/detail/logging_resource_adaptor_impl.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,9 +5,13 @@

#include <rmm/aligned.hpp>
#include <rmm/cuda_stream_view.hpp>
#include <rmm/detail/error.hpp>
#include <rmm/detail/format.hpp>
#include <rmm/mr/detail/logging_resource_adaptor_impl.hpp>

#include <cuda/stream_ref>
#include <cuda_runtime_api.h>

namespace RMM_NAMESPACE {
namespace mr {
namespace detail {
Expand All @@ -26,24 +30,19 @@ 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{};
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;
} catch (...) {
logger_->info("allocate failure,%p,%zu,%s", nullptr, bytes, rmm::detail::format_stream(stream));
throw;
}
auto const stream = cuda::stream_ref{cudaStream_t{nullptr}};
void* ptr = allocate(stream, bytes, alignment);
RMM_CUDA_TRY(cudaStreamSynchronize(stream.get()));
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));
Comment thread
wence- marked this conversation as resolved.
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,
Expand Down
13 changes: 11 additions & 2 deletions cpp/src/mr/detail/prefetch_resource_adaptor_impl.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,9 +4,13 @@
*/

#include <rmm/cuda_device.hpp>
#include <rmm/detail/error.hpp>
#include <rmm/mr/detail/prefetch_resource_adaptor_impl.hpp>
#include <rmm/prefetch.hpp>

#include <cuda/stream_ref>
#include <cuda_runtime_api.h>

namespace RMM_NAMESPACE {
namespace mr {
namespace detail {
Expand Down Expand Up @@ -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
Expand Down
5 changes: 4 additions & 1 deletion cpp/src/mr/detail/sam_headroom_memory_resource_impl.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
#include <rmm/detail/error.hpp>
#include <rmm/mr/detail/sam_headroom_memory_resource_impl.hpp>

#include <cuda/stream_ref>
#include <cuda_runtime_api.h>

#include <algorithm>
Expand Down Expand Up @@ -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
Expand Down
13 changes: 11 additions & 2 deletions cpp/src/mr/detail/statistics_resource_adaptor_impl.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,12 @@
*/

#include <rmm/cuda_stream_view.hpp>
#include <rmm/detail/error.hpp>
#include <rmm/mr/detail/statistics_resource_adaptor_impl.hpp>

#include <cuda/stream_ref>
#include <cuda_runtime_api.h>

#include <stdexcept>

namespace RMM_NAMESPACE {
Expand Down Expand Up @@ -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
Expand Down
13 changes: 11 additions & 2 deletions cpp/src/mr/detail/thread_safe_resource_adaptor_impl.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,12 @@
* SPDX-License-Identifier: Apache-2.0
*/

#include <rmm/detail/error.hpp>
#include <rmm/mr/detail/thread_safe_resource_adaptor_impl.hpp>

#include <cuda/stream_ref>
#include <cuda_runtime_api.h>

namespace RMM_NAMESPACE {
namespace mr {
namespace detail {
Expand Down Expand Up @@ -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
Expand Down
Loading
Loading