Refactor memory resource ownership and type safety by nirandaperera · Pull Request #985 · rapidsai/rapidsmpf

nirandaperera · 2026-04-22T04:56:54Z

`RmmResourceAdaptorImpl` — templated on primary and fallback resources

Add PrimaryMR and FallbackMR template parameters (both constrained to
cuda::mr::resource_with<device_accessible>; FallbackMR defaults to
any_resource<device_accessible>)
Move all method implementations into the header to support full template instantiation
Add explicit RmmResourceAdaptorImpl(std::in_place_t, Args&&...) constructor to
support in-place construction of non-movable types (e.g. cuda::pinned_memory_pool)

`PinnedMemoryResource` — eliminate redundant heap allocation

Inherit from cuda::mr::shared_resource<RmmResourceAdaptorImpl<cuda::pinned_memory_pool>>
directly, storing the pool inside the shared control block rather than through a
separate RmmResourceAdaptor member
Remove virtual methods from HostMemoryResource, making it a standalone value-semantic type
Add default-alignment allocate/deallocate overloads so call sites omitting the
alignment argument continue to work

Value-semantic factory methods

PinnedMemoryResource::make_if_available and from_options now return
std::optional<PinnedMemoryResource> instead of std::shared_ptr<PinnedMemoryResource>
PinnedMemoryResource::Disabled changed from nullptr to std::nullopt_t — existing
== Disabled comparisons remain valid via optional::operator==(nullopt_t)
BufferResource and Statistics members/parameters updated from shared_ptr to
std::optional
BufferResource::host_mr_ reverted to the concrete HostMemoryResource type

Cython bindings

PinnedMemoryResource._handle changed from shared_ptr[cpp_PinnedMemoryResource] to
optional[cpp_PinnedMemoryResource], matching the C++ value semantics throughout
buffer_resource.pyx and statistics.pyx updated accordingly

Closes #978

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

madsbk

Thanks @nirandaperera, looks good. I only have one change request

wence- · 2026-04-22T14:02:20Z

+        auto* ptr = allocate(cuda::stream_ref{cudaStream_t{nullptr}}, bytes, alignment);
+        RAPIDSMPF_CUDA_TRY(cudaStreamSynchronize(cudaStream_t{nullptr}));


I think we should use an internal stream (stored in the adaptor) rather than the null stream.

wence- · 2026-04-22T14:04:02Z

+ * holds the pool directly inside the shared control block — no extra heap allocation for
+ * the pool itself. Copies share the same underlying pool and memory statistics.


This aspect seems like it is just not worth it.

Co-authored-by: Mads R. B. Kristensen <madsbk@gmail.com>

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

…_cccl_mr

…mpf into host_mrs_to_cccl_mr

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

…_cccl_mr

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

…_cccl_mr

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

…mpf into host_mrs_to_cccl_mr

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

Co-authored-by: Lawrence Mitchell <wence@gmx.li>

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

nirandaperera · 2026-04-28T16:19:57Z

@wence- Could you please go through this again? I am inclined to keep the templates in the detail::AdapterImpl, because IINM, we dont have a good way to create a resource in-place other than, wrapping it with a shared_resource.

wence-

I have some minor questions/suggestions, but I don't think they are blocking. I am just about convinced that the added complexity of templating here is worth it.

wence- · 2026-04-30T08:40:33Z

+class PinnedMemoryResource final
+    : public cuda::mr::shared_resource<
+          detail::RmmResourceAdaptorImpl<cuda::pinned_memory_pool>> {
+    using shared_base = cuda::mr::shared_resource<
+        detail::RmmResourceAdaptorImpl<cuda::pinned_memory_pool>>;


OK, so the only question I now have here is that effectively this bakes that allocation with the pinned pool always uses the RMMResourceAdaptor. I suppose we're going to be relying on that, and it's not different to the device MR which also has to have this wrapping?

Yes, I think so. We can provide a mem-tracking-enabled device memory pool inside rapidsmpf that bakes in cuda mempool into the adapter. Do you think we should provide that?

wence- · 2026-04-30T08:42:57Z

+        auto tracker = *this;  // shared ownership — same underlying pool
+        return [tracker, limit]() { return limit - tracker.get().current_allocated(); };
    }


What does dereferencing this buy us here over return [this, limit]() {...};?

this would simply pass by ptr, isnt it? We want to increment the refcount before passing it in IINM.

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

nirandaperera · 2026-04-30T17:59:02Z

/merge

## Remove memory resources from `Statistics` construction `Statistics` held `RmmResourceAdaptor` and `PinnedMemoryResource` as instance fields, tying resource lifetime to the stats object and preventing a single `Statistics` instance from being used with different resources at report time. ### Solution - Remove `mr_` / `pinned_mr_` fields and the `Statistics(RmmResourceAdaptor, ...)` constructor; `report()` and `create_memory_recorder()` now accept explicit `std::optional<any_device_resource>` / `std::optional<any_host_device_resource>` parameters — callers supply resources at use time. - Add `try_pinned_mr()` returning `std::optional<any_host_device_resource>` non-throwingly; correct `pinned_mr()` return type to `host_device_async_resource_ref`. - Add `memory/resource_types.hpp` with `any_device_resource` / `any_host_device_resource` aliases; update `RAPIDSMPF_MEMORY_PROFILE` macro to require an explicit `mr` argument. Closes #979 Depends on rapidsai/rapids-cmake#1008 and #985 Authors: - Niranda Perera (https://github.com/nirandaperera) Approvers: - Lawrence Mitchell (https://github.com/wence-) URL: #1003

nirandaperera added 6 commits April 21, 2026 16:53

remove inheritance

a1c6be0

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

remove inheritence

5753ac3

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

remove virtual methods

301fa5c

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

inherit from rmm resource adapter impl

08eadab

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

revert host mr

26be0d1

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

move pinned mr to optional

8baea65

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

nirandaperera requested review from a team as code owners April 22, 2026 04:56

nirandaperera added breaking Introduces a breaking change improvement Improves an existing functionality labels Apr 22, 2026

private ctr

7ca6c84

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

madsbk approved these changes Apr 22, 2026

View reviewed changes

Comment thread cpp/include/rapidsmpf/memory/pinned_memory_resource.hpp Outdated

wence- requested changes Apr 22, 2026

View reviewed changes

nirandaperera and others added 4 commits April 22, 2026 11:25

Update cpp/include/rapidsmpf/memory/pinned_memory_resource.hpp

640c0f9

Co-authored-by: Mads R. B. Kristensen <madsbk@gmail.com>

addressing comments

377153a

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

Merge branch 'main' of github.com:rapidsai/rapidsmpf into host_mrs_to…

e3ffe03

…_cccl_mr

Merge branch 'host_mrs_to_cccl_mr' of github.com:nirandaperera/rapids…

4b77c81

…mpf into host_mrs_to_cccl_mr

nirandaperera requested a review from wence- April 22, 2026 18:40

nirandaperera added 4 commits April 22, 2026 11:57

precommit

44fada4

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

remove sys mem resource

3a583c7

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

Merge branch 'main' of github.com:rapidsai/rapidsmpf into host_mrs_to…

fe7d027

…_cccl_mr

Merge branch 'main' into host_mrs_to_cccl_mr

949a7ae

wence- reviewed Apr 23, 2026

View reviewed changes

Comment thread cpp/include/rapidsmpf/detail/rmm_resource_adaptor_impl.hpp

wence- reviewed Apr 23, 2026

View reviewed changes

Comment thread cpp/include/rapidsmpf/detail/rmm_resource_adaptor_impl.hpp Outdated

nirandaperera and others added 6 commits April 23, 2026 11:39

fix equality

b95b6d4

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

Merge branch 'main' of github.com:rapidsai/rapidsmpf into host_mrs_to…

feab18f

…_cccl_mr

equality fix

915782e

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

Merge branch 'host_mrs_to_cccl_mr' of github.com:nirandaperera/rapids…

774be35

…mpf into host_mrs_to_cccl_mr

precommit

edec665

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

Update cpp/include/rapidsmpf/detail/rmm_resource_adaptor_impl.hpp

a5c9b05

Co-authored-by: Lawrence Mitchell <wence@gmx.li>

pinned mem equality

a446264

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

nirandaperera requested a review from wence- April 24, 2026 16:40

nirandaperera added 2 commits April 24, 2026 09:40

Merge branch 'main' into host_mrs_to_cccl_mr

48a4786

Merge branch 'main' into host_mrs_to_cccl_mr

f75a9ba

nirandaperera mentioned this pull request Apr 29, 2026

Remove memory resources from Statistics #1003

Merged

wence- approved these changes Apr 30, 2026

View reviewed changes

nirandaperera added 2 commits April 30, 2026 07:55

Merge branch 'main' into host_mrs_to_cccl_mr

b8be3a3

minor comments

2bd3e46

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

rapids-bot Bot merged commit 82ba989 into rapidsai:main Apr 30, 2026
66 checks passed

nirandaperera mentioned this pull request Apr 30, 2026

Stats refactor #1008

Open

3 tasks

		auto* ptr = allocate(cuda::stream_ref{cudaStream_t{nullptr}}, bytes, alignment);
		RAPIDSMPF_CUDA_TRY(cudaStreamSynchronize(cudaStream_t{nullptr}));

		* holds the pool directly inside the shared control block — no extra heap allocation for
		* the pool itself. Copies share the same underlying pool and memory statistics.

Uh oh!

Conversation

nirandaperera commented Apr 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

RmmResourceAdaptorImpl — templated on primary and fallback resources

PinnedMemoryResource — eliminate redundant heap allocation

Value-semantic factory methods

Cython bindings

Uh oh!

madsbk left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

wence- Apr 22, 2026

Choose a reason for hiding this comment

Uh oh!

wence- Apr 22, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

nirandaperera commented Apr 28, 2026

Uh oh!

wence- left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

wence- Apr 30, 2026

Choose a reason for hiding this comment

Uh oh!

nirandaperera Apr 30, 2026

Choose a reason for hiding this comment

Uh oh!

wence- Apr 30, 2026

Choose a reason for hiding this comment

Uh oh!

nirandaperera Apr 30, 2026

Choose a reason for hiding this comment

Uh oh!

nirandaperera commented Apr 30, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

nirandaperera commented Apr 22, 2026 •

edited

Loading

`RmmResourceAdaptorImpl` — templated on primary and fallback resources

`PinnedMemoryResource` — eliminate redundant heap allocation