From d4d4c0af7e29a6c6e3ea522ba6fb5f3a50ac968d Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sun, 17 May 2026 17:23:22 +0000 Subject: [PATCH 1/2] Forward failure callback alignment --- ...failure_callback_resource_adaptor_impl.hpp | 8 +++---- cpp/tests/mr/failure_callback_mr_tests.cpp | 24 +++++++++++++++++++ 2 files changed, 28 insertions(+), 4 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 9b3e61f14..4cd07e29b 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 @@ -62,12 +62,12 @@ class failure_callback_resource_adaptor_impl { void* allocate(cuda::stream_ref stream, std::size_t bytes, - std::size_t /*alignment*/ = alignof(std::max_align_t)) + std::size_t alignment = alignof(std::max_align_t)) { void* ret{}; while (true) { try { - ret = upstream_mr_.allocate(stream, bytes, rmm::CUDA_ALLOCATION_ALIGNMENT); + ret = upstream_mr_.allocate(stream, bytes, alignment); break; } catch (ExceptionType const&) { if (!callback_(bytes, callback_arg_)) { throw; } @@ -79,9 +79,9 @@ class failure_callback_resource_adaptor_impl { void deallocate(cuda::stream_ref stream, void* ptr, std::size_t bytes, - std::size_t /*alignment*/ = alignof(std::max_align_t)) noexcept + std::size_t alignment = alignof(std::max_align_t)) noexcept { - upstream_mr_.deallocate(stream, ptr, bytes, rmm::CUDA_ALLOCATION_ALIGNMENT); + upstream_mr_.deallocate(stream, ptr, bytes, alignment); } void* allocate_sync(std::size_t bytes, std::size_t alignment = alignof(std::max_align_t)) diff --git a/cpp/tests/mr/failure_callback_mr_tests.cpp b/cpp/tests/mr/failure_callback_mr_tests.cpp index bd835680c..0e561d393 100644 --- a/cpp/tests/mr/failure_callback_mr_tests.cpp +++ b/cpp/tests/mr/failure_callback_mr_tests.cpp @@ -4,6 +4,7 @@ */ #include "../byte_literals.hpp" +#include "../mock_resource.hpp" #include #include @@ -24,6 +25,9 @@ namespace { template using failure_callback_adaptor = rmm::mr::failure_callback_resource_adaptor; +using ::testing::_; +using ::testing::Return; + bool failure_handler(std::size_t /*bytes*/, void* arg) { // NOLINTNEXTLINE(cppcoreguidelines-pro-type-reinterpret-cast) @@ -78,6 +82,26 @@ TEST(FailureCallbackTest, RetryAllocationOnce) EXPECT_EQ(retried, true); } +TEST(FailureCallbackTest, ForwardsAlignment) +{ + mock_resource mock; + mock_resource_wrapper wrapper{&mock}; + bool retried{false}; + failure_callback_adaptor<> mr{device_async_resource_ref{wrapper}, failure_handler, &retried}; + + cuda_stream_view stream; + std::byte pointer_value{}; + void* const pointer = &pointer_value; + auto const size{1024}; + auto const alignment{512}; + + EXPECT_CALL(mock, allocate(_, size, alignment)).WillOnce(Return(pointer)); + EXPECT_CALL(mock, deallocate(_, pointer, size, alignment)).Times(1); + + EXPECT_EQ(mr.allocate(stream, size, alignment), pointer); + mr.deallocate(stream, pointer, size, alignment); +} + TEST(FailureCallbackTest, DifferentExceptionTypes) { always_throw_memory_resource bad_alloc_mr; From 65960340d85994c3fb4126788750842e3d034737 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Thu, 21 May 2026 16:47:46 +0000 Subject: [PATCH 2/2] Default failure callback alignment to CUDA alignment --- .../mr/detail/failure_callback_resource_adaptor_impl.hpp | 8 ++++---- 1 file changed, 4 insertions(+), 4 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 4cd07e29b..2cd9e2457 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 @@ -62,7 +62,7 @@ class failure_callback_resource_adaptor_impl { void* allocate(cuda::stream_ref stream, std::size_t bytes, - std::size_t alignment = alignof(std::max_align_t)) + std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) { void* ret{}; while (true) { @@ -79,19 +79,19 @@ class failure_callback_resource_adaptor_impl { void deallocate(cuda::stream_ref stream, void* ptr, std::size_t bytes, - std::size_t alignment = alignof(std::max_align_t)) noexcept + std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) noexcept { upstream_mr_.deallocate(stream, ptr, bytes, alignment); } - void* allocate_sync(std::size_t bytes, std::size_t alignment = alignof(std::max_align_t)) + void* allocate_sync(std::size_t bytes, std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) { return allocate(cuda_stream_view{}, bytes, alignment); } void deallocate_sync(void* ptr, std::size_t bytes, - std::size_t alignment = alignof(std::max_align_t)) noexcept + std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) noexcept { deallocate(cuda_stream_view{}, ptr, bytes, alignment); }