Skip to content

Add a strict Buffer Resource back-reference to Host and Pinned memory resources - #1106

Merged
rapids-bot[bot] merged 20 commits into
rapidsai:release/26.08from
nirandaperera:backref-host-pinned-mr
Jul 19, 2026
Merged

Add a strict Buffer Resource back-reference to Host and Pinned memory resources#1106
rapids-bot[bot] merged 20 commits into
rapidsai:release/26.08from
nirandaperera:backref-host-pinned-mr

Conversation

@nirandaperera

Copy link
Copy Markdown
Contributor

HostMemoryResource and PinnedMemoryResource are now BufferResource-only resources that carry a strict BackRefMixin<BufferResource> back-reference. When a copy of one of these resources is made (e.g. when CCCL promotes host_mr() / pinned_mr() into an owning cuda::mr::any_resource inside a HostBuffer), the copy promotes the stored weak_ptr to a shared_ptr, keeping the owning BufferResource alive for as long as any derived buffer lives. This closes a lifetime gap where a Buffer could outlive the BufferResource that produced its memory resource.

The mixin is strict: copying an instance without an installed back-reference throws std::bad_weak_ptr. To make that safe, these resources can no longer be constructed as standalone objects — they are created only by BufferResource, which installs the back-reference immediately after construction. This mirrors how RmmResourceAdaptor already works on this branch.

C++ changes

  • HostMemoryResource: inherits BackRefMixin<BufferResource>; constructor is private with friend class BufferResource. Equality remains stateless (always true) since instances are interchangeable and the back-reference is installed exactly once.
  • PinnedMemoryResource: inherits BackRefMixin<BufferResource> as a second base; constructor is private (friend class BufferResource). The make_if_available / from_options factories were removed so no back-reference-less instance can ever escape. Equality compares the shared pool state only.
  • PinnedPoolProperties: gained a numa_id field (defaults to the calling thread's NUMA node) so all pinned configuration flows as one struct. Added a free helper pinned_pool_properties_from_options(...).
  • PinnedMemoryDisabled: new constexpr std::nullopt_t sentinel used to disable pinned host memory (replaces the old PinnedMemoryResource::Disabled).
  • BufferResource::create() / from_options(): now take std::optional<PinnedPoolProperties> pinned_pool_properties (default PinnedMemoryDisabled) instead of a pre-built resource. The pinned resource is constructed internally and the back-reference is installed on owning_mr_, host_mr_, and pinned_mr_ before create() returns.
  • try_pinned_mr(): now returns std::optional<PinnedMemoryResource> (a back-referenced handle) so Python/Statistics get a concrete, lifetime-safe handle.

Python changes

  • PinnedMemoryResource: now an opaque, non-constructible handle (__init__ raises TypeError); obtained via BufferResource.pinned_mr.
  • PinnedPoolProperties: new @dataclass (initial_pool_size, max_pool_size, numa_id) used to configure pinned memory.
  • BufferResource(...): replaces the old pinned_mr= argument with pinned_pool_properties: PinnedPoolProperties | None = None (None disables pinned host memory). from_options derives this from the config options.
  • Statistics.report(pinned_mr=...): unchanged signature; the handle is now sourced from BufferResource.pinned_mr.

Tests

  • Added HostMrKeepsBufferResourceAlive and PinnedMrKeepsBufferResourceAlive regression tests mirroring DeviceMrKeepsBufferResourceAlive.
  • Routed all C++ tests/benchmarks through a BufferResource (no standalone Host/Pinned MR construction); reworked test_host_buffer, test_config, test_memory_resources, and bench_memory_resources.
  • Updated Python test_config.py to configure pinned memory via BufferResource and assert on BufferResource.pinned_mr.

Breaking changes

  • C++: HostMemoryResource / PinnedMemoryResource can no longer be constructed directly; PinnedMemoryResource::make_if_available / from_options and the Disabled sentinel are removed. BufferResource::create() takes PinnedPoolProperties instead of a PinnedMemoryResource.
  • Python: PinnedMemoryResource(...) is no longer constructible; the BufferResource(pinned_mr=...) argument is replaced by pinned_pool_properties=.

Depends on #1078

Closes #1070

@nirandaperera nirandaperera self-assigned this Jun 22, 2026
@nirandaperera nirandaperera added breaking Introduces a breaking change improvement Improves an existing functionality labels Jun 22, 2026
@copy-pr-bot

copy-pr-bot Bot commented Jun 22, 2026

Copy link
Copy Markdown

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

Signed-off-by: niranda perera <niranda.perera@gmail.com>
@nirandaperera
nirandaperera force-pushed the backref-host-pinned-mr branch from ebc6a70 to ec86fe5 Compare June 22, 2026 23:34
@nirandaperera
nirandaperera marked this pull request as ready for review June 22, 2026 23:34
@nirandaperera
nirandaperera requested review from a team as code owners June 22, 2026 23:34
@nirandaperera
nirandaperera force-pushed the backref-host-pinned-mr branch from ec86fe5 to 2db591e Compare June 24, 2026 19:38
…t-pinned-mr

Signed-off-by: niranda perera <niranda.perera@gmail.com>
@nirandaperera
nirandaperera force-pushed the backref-host-pinned-mr branch from 2db591e to dfd225d Compare July 1, 2026 22:18
Signed-off-by: niranda perera <niranda.perera@gmail.com>

@madsbk madsbk left a comment

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.

Overall looks good

Comment thread python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pyx Outdated
Comment thread cpp/include/rapidsmpf/memory/buffer_resource.hpp Outdated
Comment thread cpp/include/rapidsmpf/memory/buffer_resource.hpp Outdated
Comment thread cpp/include/rapidsmpf/memory/host_memory_resource.hpp Outdated
Comment thread cpp/include/rapidsmpf/memory/pinned_memory_resource.hpp Outdated
wence-
wence- previously requested changes Jul 13, 2026

@wence- wence- left a comment

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.

Code changes broadly look good, but please clean up the docstrings and comments throughout.

In general, please try and elide as much implementation detail in comments and, if it is necessary, have it only in one place rather than everywhere. Docstrings and comments should generally talk about observable behaviour and semantics, not implementation specifics.

Comment thread cpp/benchmarks/bench_memory_resources.cpp Outdated
Comment thread cpp/include/rapidsmpf/memory/buffer_resource.hpp Outdated
Comment thread cpp/include/rapidsmpf/memory/host_memory_resource.hpp Outdated
Comment thread cpp/include/rapidsmpf/memory/pinned_memory_resource.hpp Outdated
Comment thread cpp/include/rapidsmpf/memory/pinned_memory_resource.hpp Outdated
Comment thread cpp/include/rapidsmpf/memory/pinned_memory_resource.hpp Outdated
Comment thread cpp/src/memory/buffer_resource.cpp Outdated
Comment thread cpp/src/memory/pinned_memory_resource.cpp Outdated
Comment thread cpp/tests/test_memory_resources.cpp Outdated
Comment thread python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx Outdated
nirandaperera and others added 3 commits July 13, 2026 08:58
Co-authored-by: Mads R. B. Kristensen <madsbk@gmail.com>
Co-authored-by: Lawrence Mitchell <wence@gmx.li>
Signed-off-by: niranda perera <niranda.perera@gmail.com>
@nirandaperera
nirandaperera requested review from madsbk and wence- July 13, 2026 18:42
Signed-off-by: niranda perera <niranda.perera@gmail.com>
Signed-off-by: niranda perera <niranda.perera@gmail.com>
Signed-off-by: niranda perera <niranda.perera@gmail.com>
@nirandaperera
nirandaperera requested a review from a team as a code owner July 13, 2026 21:21

@madsbk madsbk left a comment

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.

Looks good. The only thing we should consider is the enablement semantics. Currently, setting pinned_pool_properties to a non-None value does not guarantee that pinned memory is enabled, since that still depends on system support.

It would be helpful to have an option that means: enable pinned memory, or raise an error if it is unavailable.

Alternatively, this option could be implemented at the cudf-polars level. After creating the memory resource, cudf-polars could verify that pinned memory was actually enabled and raise an error otherwise.

Signed-off-by: niranda perera <niranda.perera@gmail.com>
Signed-off-by: niranda perera <niranda.perera@gmail.com>
@nirandaperera

Copy link
Copy Markdown
Contributor Author

Looks good. The only thing we should consider is the enablement semantics. Currently, setting pinned_pool_properties to a non-None value does not guarantee that pinned memory is enabled, since that still depends on system support.

It would be helpful to have an option that means: enable pinned memory, or raise an error if it is unavailable.

Alternatively, this option could be implemented at the cudf-polars level. After creating the memory resource, cudf-polars could verify that pinned memory was actually enabled and raise an error otherwise.

I made BufferResource::create throw, if PinnedProperties provided and pinned memory is not supported

Signed-off-by: niranda perera <niranda.perera@gmail.com>
@nirandaperera
nirandaperera changed the base branch from main to branch-25.06 July 18, 2026 16:12
@nirandaperera
nirandaperera requested review from a team as code owners July 18, 2026 16:12
@nirandaperera
nirandaperera changed the base branch from branch-25.06 to release/26.08 July 18, 2026 16:17
…ckref-host-pinned-mr

Signed-off-by: niranda perera <niranda.perera@gmail.com>
Signed-off-by: niranda perera <niranda.perera@gmail.com>
@nirandaperera

Copy link
Copy Markdown
Contributor Author

/merge

@nirandaperera
nirandaperera dismissed wence-’s stale review July 19, 2026 07:46

dismissing request as all comments have been addressed

@rapids-bot
rapids-bot Bot merged commit 02c5654 into rapidsai:release/26.08 Jul 19, 2026
66 checks passed
vyasr pushed a commit to NVIDIA/cudf that referenced this pull request Jul 20, 2026
## Description
Adds PR changes to Host and Pinned MRs. 

Depends on rapidsai/rapidsmpf#1106

## Checklist
- [x] I am familiar with the [Contributing
Guidelines](https://github.com/rapidsai/cudf/blob/HEAD/CONTRIBUTING.md).
- [x] New or existing tests cover these changes.
- [x] The documentation is up to date with these changes.

---------

Signed-off-by: niranda perera <niranda.perera@gmail.com>
davidwendt pushed a commit to wjxiz1992/cudf that referenced this pull request Jul 21, 2026
## Description
Adds PR changes to Host and Pinned MRs. 

Depends on rapidsai/rapidsmpf#1106

## Checklist
- [x] I am familiar with the [Contributing
Guidelines](https://github.com/rapidsai/cudf/blob/HEAD/CONTRIBUTING.md).
- [x] New or existing tests cover these changes.
- [x] The documentation is up to date with these changes.

---------

Signed-off-by: niranda perera <niranda.perera@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

breaking Introduces a breaking change improvement Improves an existing functionality

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Make BufferResource::host_mr() and pinned_mr() return owning MRs

3 participants