Skip to content

Honor requested aligned adaptor alignment #2396

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

Open
bdice wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Honor requested aligned adaptor alignment #2396

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f996e36
Honor requested aligned adaptor alignment
bdice May 17, 2026
aaf9cc9
Merge branch 'main' into fix-aligned-adaptor-request-alignment
bdice May 21, 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
2 changes: 0 additions & 2 deletions cpp/include/rmm/mr/detail/aligned_resource_adaptor_impl.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -77,8 +77,6 @@ class aligned_resource_adaptor_impl {
}

private:
[[nodiscard]] std::size_t upstream_allocation_size(std::size_t bytes) const;

cuda::mr::any_resource<cuda::mr::device_accessible> upstream_mr_;
std::unordered_map<void*, void*> pointers_;
std::size_t alignment_;
Expand Down
45 changes: 32 additions & 13 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 @@ -14,6 +14,26 @@
namespace RMM_NAMESPACE {
namespace mr {
namespace detail {
namespace {

[[nodiscard]] std::size_t upstream_allocation_size(std::size_t bytes, std::size_t alignment)

@wence- wence- May 26, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
[[nodiscard]] std::size_t upstream_allocation_size(std::size_t bytes, std::size_t alignment)
[[nodiscard]] constexpr std::size_t upstream_allocation_size(std::size_t bytes, std::size_t alignment) noexcept

(Although align_up is not constexpr because in NDEBUG mode it asserts)

{
auto const aligned_size = rmm::align_up(bytes, alignment);
return aligned_size + (alignment - rmm::CUDA_ALLOCATION_ALIGNMENT);
}

[[nodiscard]] std::size_t effective_alignment(std::size_t bytes,

@wence- wence- May 26, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
[[nodiscard]] std::size_t effective_alignment(std::size_t bytes,
[[nodiscard]] constexpr std::size_t effective_alignment(std::size_t bytes,

std::size_t requested_alignment,
std::size_t configured_alignment,
std::size_t alignment_threshold) noexcept
{
return std::max(
{requested_alignment,
bytes >= alignment_threshold ? configured_alignment : rmm::CUDA_ALLOCATION_ALIGNMENT,
rmm::CUDA_ALLOCATION_ALIGNMENT});
}

} // namespace

aligned_resource_adaptor_impl::aligned_resource_adaptor_impl(
cuda::mr::any_resource<cuda::mr::device_accessible> upstream,
Expand All @@ -39,24 +59,21 @@ std::size_t aligned_resource_adaptor_impl::get_alignment_threshold() const noexc
return alignment_threshold_;
}

std::size_t aligned_resource_adaptor_impl::upstream_allocation_size(std::size_t bytes) const
{
auto const aligned_size = rmm::align_up(bytes, alignment_);
return aligned_size + (alignment_ - rmm::CUDA_ALLOCATION_ALIGNMENT);
}

void* aligned_resource_adaptor_impl::allocate(cuda::stream_ref stream,
std::size_t bytes,
std::size_t /*alignment*/)
std::size_t alignment)
{
if (bytes == 0 || alignment_ == rmm::CUDA_ALLOCATION_ALIGNMENT || bytes < alignment_threshold_) {
RMM_EXPECTS(rmm::is_supported_alignment(alignment), "Allocation alignment is not a power of 2.");
auto const effective_align =
effective_alignment(bytes, alignment, alignment_, alignment_threshold_);
if (bytes == 0 || effective_align == rmm::CUDA_ALLOCATION_ALIGNMENT) {
return upstream_mr_.allocate(stream, bytes, 1);
Comment on lines +69 to 70

@wence- wence- May 26, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

note: this implicitly encodes that every RMM memory resource must provide allocations that are CUDA_ALLOCATION_ALIGNMENT aligned. Should we (at least in debug mode) assert that?

Imagine I am a writing my own resource, and I forget about this restriction.

I think we should also push (again) for the "related" request in this cccl issue NVIDIA/cccl#8157

We can at least implement said properties on all RMM resources today, I think.

}
auto const size = upstream_allocation_size(bytes);
auto const size = upstream_allocation_size(bytes, effective_align);
void* pointer = upstream_mr_.allocate(stream, size, 1);
// NOLINTNEXTLINE(cppcoreguidelines-pro-type-reinterpret-cast)
auto const address = reinterpret_cast<std::size_t>(pointer);
auto const aligned_address = rmm::align_up(address, alignment_);
auto const aligned_address = rmm::align_up(address, effective_align);
// NOLINTNEXTLINE(cppcoreguidelines-pro-type-reinterpret-cast,performance-no-int-to-ptr)
void* aligned_pointer = reinterpret_cast<void*>(aligned_address);
if (pointer != aligned_pointer) {
Expand All @@ -70,9 +87,11 @@ void* aligned_resource_adaptor_impl::allocate(cuda::stream_ref stream,
void aligned_resource_adaptor_impl::deallocate(cuda::stream_ref stream,
void* ptr,
std::size_t bytes,
std::size_t /*alignment*/) noexcept
std::size_t alignment) noexcept
{
if (bytes == 0 || alignment_ == rmm::CUDA_ALLOCATION_ALIGNMENT || bytes < alignment_threshold_) {
auto const effective_align =
effective_alignment(bytes, alignment, alignment_, alignment_threshold_);
if (bytes == 0 || effective_align == rmm::CUDA_ALLOCATION_ALIGNMENT) {
upstream_mr_.deallocate(stream, ptr, bytes, 1);
} else {
{
Expand All @@ -83,7 +102,7 @@ void aligned_resource_adaptor_impl::deallocate(cuda::stream_ref stream,
pointers_.erase(iter);
}
}
upstream_mr_.deallocate(stream, ptr, upstream_allocation_size(bytes), 1);
upstream_mr_.deallocate(stream, ptr, upstream_allocation_size(bytes, effective_align), 1);
}
}

Expand Down
57 changes: 57 additions & 0 deletions cpp/tests/mr/aligned_mr_tests.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,15 @@ TEST(AlignedTest, ThrowOnInvalidAllocationAlignment)
EXPECT_THROW(construct_alignment(wrapper, 768), rmm::logic_error);
}

TEST(AlignedTest, ThrowOnInvalidRequestedAlignment)
{
mock_resource mock;
mock_resource_wrapper wrapper{&mock};
aligned_adaptor mr{device_async_resource_ref{wrapper}};

EXPECT_THROW((void)mr.allocate(cuda_stream_view{}, 1024, 768), rmm::logic_error);
}

TEST(AlignedTest, SupportsGetMemInfo)
{
mock_resource mock;
Expand Down Expand Up @@ -197,6 +206,54 @@ TEST(AlignedTest, AlignUpstreamAddress)
}
}

TEST(AlignedTest, AlignsCallerRequestedAlignment)
{
mock_resource mock;
mock_resource_wrapper wrapper{&mock};
aligned_adaptor mr{device_async_resource_ref{wrapper}};

cuda_stream_view stream;
auto const alignment{4096};
{
void* const pointer = int_to_address(256);
auto const size{7936};
EXPECT_CALL(mock, allocate(_, size, _)).WillOnce(Return(pointer));
EXPECT_CALL(mock, deallocate(_, pointer, size, _)).Times(1);
}

{
void* const expected_pointer = int_to_address(4096);
auto const size{1024};
EXPECT_EQ(mr.allocate(stream, size, alignment), expected_pointer);
mr.deallocate(stream, expected_pointer, size, alignment);
Comment on lines +225 to +228

@wence- wence- May 26, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't understand how gtest's mock methods work. But I would have thought that the thing to test here is:

void *ptr = mr.allocate(...);
EXPECT_EQ(static_cast<std::uintptr_t>(ptr) % alignment, 0);
...

}
}

TEST(AlignedTest, CallerRequestedAlignmentOverridesThreshold)
{
mock_resource mock;
mock_resource_wrapper wrapper{&mock};
auto const alignment{4096};
auto const threshold{65536};
aligned_adaptor mr{device_async_resource_ref{wrapper}, alignment, threshold};

cuda_stream_view stream;
auto const requested_alignment{8192};
{
void* const pointer = int_to_address(256);
auto const size{16128};
EXPECT_CALL(mock, allocate(_, size, _)).WillOnce(Return(pointer));
EXPECT_CALL(mock, deallocate(_, pointer, size, _)).Times(1);
}

{
void* const expected_pointer = int_to_address(8192);
auto const size{1024};
EXPECT_EQ(mr.allocate(stream, size, requested_alignment), expected_pointer);
mr.deallocate(stream, expected_pointer, size, requested_alignment);
Comment on lines +250 to +253

@wence- wence- May 26, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

again here.

}
}

TEST(AlignedTest, AlignMultiple)
{
mock_resource mock;
Expand Down
Loading