Skip to content

Simplify HostBuffer memory ownership#1013

Open
nirandaperera wants to merge 4 commits intorapidsai:mainfrom
nirandaperera:minor-host-buffer-refactor
Open

Simplify HostBuffer memory ownership#1013
nirandaperera wants to merge 4 commits intorapidsai:mainfrom
nirandaperera:minor-host-buffer-refactor

Conversation

@nirandaperera
Copy link
Copy Markdown
Contributor

@nirandaperera nirandaperera commented May 5, 2026

HostBuffer tracked two redundant deallocation paths (mr_ + owned_storage_), where mr_ was stored unconditionally but never used when owned_storage_ was present — and as a non-owning ref it could silently outlive a PinnedMemoryResource.

Solution

  • Replace both members with a single std::function<void(rmm::cuda_stream_view)> deallocate_fn_, invoked with the live stream at deallocation time to respect set_stream().
  • In the RMM-allocated path, wrap mr in cuda::mr::any_resource before capturing it — deep-copying the resource and incrementing its refcount if it is a PinnedMemoryResource.
  • Remove the vestigial mr parameter from from_owned_vector and from_rmm_device_buffer (it was always unused for deallocation).
  • Remove the leaked OwnedStorageDeleter public type alias.

Signed-off-by: niranda perera <niranda.perera@gmail.com>
@nirandaperera nirandaperera requested a review from a team as a code owner May 5, 2026 04:27
@nirandaperera nirandaperera added improvement Improves an existing functionality non-breaking Introduces a non-breaking change labels May 5, 2026
: stream_{stream}, mr_{std::move(mr)} {
: stream_{stream} {
if (size > 0) {
auto owned_mr = cuda::mr::any_resource<cuda::mr::host_accessible>{mr};
Copy link
Copy Markdown
Member

@madsbk madsbk May 5, 2026

Choose a reason for hiding this comment

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

I don’t think owned_mr actually owns the memory resource. AFAIK, cuda::mr::any_resource<P> only owns the rmm::host_async_resource_ref mr, which is itself a non-owning reference. So owned_mr ends up owning a copy of a non-owning reference, not the underlying memory resource?

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.

This is a reasonable thought but incorrect. any_resource(resource_ref) magically converts the ref into an owning type.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Thanks @wence-, good to know! @nirandaperera could you add a comment noting this?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Yeah, any_resource takes in a copy of the underlying resource in the ref. Idea is to increment the refcount if that resource is shared.

Copy link
Copy Markdown
Member

@madsbk madsbk left a comment

Choose a reason for hiding this comment

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

@nirandaperera, it would be good to add a test that constructs a HostBuffer from a transient memory resource, for example a PinnedMemoryResource that goes out of scope before the buffer is deallocated.

@nirandaperera nirandaperera changed the title [MINOR] Simplify HostBuffer memory ownership Simplify HostBuffer memory ownership May 5, 2026
Signed-off-by: niranda perera <niranda.perera@gmail.com>
@nirandaperera nirandaperera requested a review from madsbk May 5, 2026 15:11
@nirandaperera
Copy link
Copy Markdown
Contributor Author

@madsbk done

@nirandaperera
Copy link
Copy Markdown
Contributor Author

/merge

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

improvement Improves an existing functionality non-breaking Introduces a non-breaking change

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants