Remove memory resources from Statistics by nirandaperera · Pull Request #1003 · rapidsai/rapidsmpf

nirandaperera · 2026-04-29T20:46:29Z

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

copy-pr-bot · 2026-04-29T20:47:21Z

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>

wence- · 2026-05-01T08:16:00Z

I think this file should be removed.

wence- · 2026-05-01T08:19:05Z

@@ -135,6 +153,16 @@ class BufferResource {
     */
    [[nodiscard]] rmm::host_async_resource_ref pinned_mr();


perhaps not in this PR, but this should return a host_device_async_resource_ref, I think.

wence- · 2026-05-01T08:19:59Z

+        requires cuda::mr::resource_with<T, cuda::mr::device_accessible>
+    [[nodiscard]] std::optional<T> device_mr_as() const noexcept {
+        auto resource = cuda::mr::resource_cast<T>(&device_mr_);
+        return resource ? std::make_optional(*resource) : std::nullopt;


Suggested change

return resource ? std::make_optional(*resource) : std::nullopt;

return resource ? *resource : std::nullopt;

wence- · 2026-05-01T08:26:52Z

+    [[nodiscard]] constexpr const std::optional<PinnedMemoryResource>&
+    concrete_pinned_mr() const noexcept {
+        return pinned_mr_;
+    }


I am bikeshedding about the names. concrete is perhaps not the best prefix?

We have device_mr_as<T> why not for symmetry pinned_mr_as<T>?

Also then, this return a const & std::optional whereas device_mr_as returns a std::optional.

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

…/rapidsmpf into remove-mrs-from-statistics

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

…from-statistics

nirandaperera · 2026-05-05T04:57:12Z

@wence- I removed those adhoc templates from the public API now. Now, methods return refs and accept resources as any_resource<...> by value.
Could you please take another look?

This reverts commit 2817ab6.

wence- · 2026-05-05T09:56:31Z

Should still be removed.

wence- · 2026-05-05T09:57:55Z

+    if (args.enable_memory_profiler) {
+        log.print(ctx->statistics()->report(
+            stat_enabled_mr, pinned_mr, "Statistics (of the last run):"
+        ));
+    } else {
+        log.print(ctx->statistics()->report(
+            std::nullopt, std::nullopt, "Statistics (of the last run):"
+        ));
+    }


Let's make the MRs arguments after the header line, I think?

Most common use was .report() or report(device_mr). I added a struct ReportArgs that gives the feeling of named-args.

wence- · 2026-05-05T10:00:02Z

-    std::string report(std::string const& header = "Statistics:") const;
+    std::string report(
+        std::optional<any_device_resource> mr = std::nullopt,
+        std::optional<any_host_device_resource> pinned_mr = std::nullopt,


Suggested change

std::optional<any_host_device_resource> pinned_mr = std::nullopt,

std::optional<any_host_device_resource> pinned_mr = PinnedMemoryResource::Disabled,

?

And, as suggested above, how about:

std::string report(std::string const& header = "Statistics:", std::optional<any_device_resource> mr = std::nullopt, std::optional<any_host_device_resource> pinned_mr = PinnedMemoryResource::Disabled);

?

wence- · 2026-05-05T10:03:37Z

 }

+std::optional<any_host_device_resource> BufferResource::try_pinned_mr() const noexcept {
+    return pinned_mr_.has_value() ? std::make_optional(*pinned_mr_) : std::nullopt;


nit: do we need the make_optional, or is:

Suggested change

return pinned_mr_.has_value() ? std::make_optional(*pinned_mr_) : std::nullopt;

return pinned_mr_.has_value() ? *pinned_mr_ : std::nullopt;

sufficient?

wence- · 2026-05-05T10:04:34Z

+template <typename T, typename... Properties>
+T* get_optional_resource_as(std::optional<cuda::mr::any_resource<Properties...>>& mr) {
+    if (mr.has_value()) {
+        return cuda::mr::resource_cast<T>(&(*mr));


question: What happens if the resource_cast fails?

it returns a nullptr

wence- · 2026-05-05T10:05:07Z

+include("${rapids-cmake-dir}/cpm/package_override.cmake")
+rapids_cpm_package_override("${CMAKE_CURRENT_SOURCE_DIR}/../cmake/cpm_override_packages.json")
+


I think you are not using this.

wence- · 2026-05-05T10:05:40Z

+    def report(
+        self,
+        mr: RmmResourceAdaptor | None = None,
+        pinned_mr: PinnedMemoryResource | None = None,


question: While here should we offer access to the "header" argument from python?

I havent seen any uses of it so far.

wence- · 2026-05-05T10:05:53Z

Please remove.

wence- · 2026-05-05T10:05:59Z

+
+include("${rapids-cmake-dir}/cpm/package_override.cmake")
+rapids_cpm_package_override("${CMAKE_CURRENT_SOURCE_DIR}/../../cmake/cpm_override_packages.json")
+


Not needed.

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

wence- · 2026-05-05T15:58:38Z

+    std::string report(ReportArgs report_args) const;
+
+    /**
+     * @brief Overload with all-default options. Equivalent to `report(ReportArgs{})`.
+     * @return Formatted statistics report.
+     */
+    std::string report() const {
+        return report(ReportArgs{});
+    }


Suggested change

std::string report(ReportArgs report_args) const;

/**

* @brief Overload with all-default options. Equivalent to `report(ReportArgs{})`.

* @return Formatted statistics report.

*/

std::string report() const {

return report(ReportArgs{});

}

std::string report(ReportArgs report_args = {}) const;

clang didnt like this. 🙁

wence-

Tiny suggested cleanup

wence- · 2026-05-05T16:00:16Z

+        self,
+        RmmResourceAdaptor mr = None,
+        PinnedMemoryResource pinned_mr = None,
+    ):


Please allow caller to provide header as well.

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

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

nirandaperera · 2026-05-06T17:10:22Z

/merge

`Statistics` was constructible directly (`make_shared<Statistics>(false)` was not equal to `Statistics::disabled()`), `MemoryRecorder` held a raw `Statistics*` (lifetime hazard noted by an existing TODO), provider `statistics()` accessors returned `shared_ptr` by value (atomic refcount bumps on hot paths), and secondary objects like `AllGather` redundantly carried their own `statistics_` member already reachable through the buffer resource / context. In this PR, following changes were made - Make `Statistics` constructible only through factory methods (`create()`, `from_options()`, `disabled()`); the public ctor is removed and `disabled()` is now the canonical singleton for the disabled case. - Inherit `Statistics` from `std::enable_shared_from_this` so `MemoryRecorder` owns a `shared_ptr<Statistics>` instead of a raw pointer; recorder no-ops when stats are disabled. - Introduce a `StatisticsProvider` concept and have providers (`BufferResource`, `ProgressThread`, `streaming::Context`) return `std::shared_ptr<Statistics> const&` from `statistics()`; each provider `static_assert`s the concept. - Drop the redundant `statistics` parameter and `statistics_` member from `coll::AllGather` (and its Python binding); secondary classes now derive their `Statistics` from the provider they already receive. Related to #1008 Depends on #1003 Authors: - Niranda Perera (https://github.com/nirandaperera) - Peter Andreas Entschev (https://github.com/pentschev) Approvers: - Peter Andreas Entschev (https://github.com/pentschev) - Lawrence Mitchell (https://github.com/wence-) URL: #1009

nirandaperera requested review from a team as code owners April 29, 2026 20:46

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

nirandaperera marked this pull request as draft April 29, 2026 20:47

nirandaperera added 3 commits April 30, 2026 11:26

remove memory resources from stats

1ba1c28

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

todo

0a08600

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

report change

ac296bc

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

nirandaperera force-pushed the remove-mrs-from-statistics branch from a7bce3f to ac296bc Compare April 30, 2026 18:32

nirandaperera marked this pull request as ready for review April 30, 2026 18:32

nirandaperera added 2 commits April 30, 2026 12:41

remove cccl custom version

6c83c42

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

precmmit

0c34b67

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

nirandaperera mentioned this pull request Apr 30, 2026

Stats refactor #1008

Open

3 tasks

Merge branch 'main' into remove-mrs-from-statistics

27b8958

nirandaperera mentioned this pull request May 1, 2026

Initial stats cleanup #1009

Merged

wence- reviewed May 1, 2026

View reviewed changes

nirandaperera added 4 commits May 4, 2026 21:40

remove adhoc templates

98bc32a

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

Merge branch 'remove-mrs-from-statistics' of github.com:nirandaperera…

7110a2b

…/rapidsmpf into remove-mrs-from-statistics

enable shared from this

2817ab6

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

Merge branch 'main' of github.com:rapidsai/rapidsmpf into remove-mrs-…

20f51da

…from-statistics

nirandaperera requested a review from wence- May 5, 2026 04:54

Revert "enable shared from this"

e2b7187

This reverts commit 2817ab6.

wence- reviewed May 5, 2026

View reviewed changes

Comment thread cmake/cpm_override_packages.json Outdated

wence- May 5, 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.

Should still be removed.

wence- requested changes May 5, 2026

View reviewed changes

addressing comments

14b6450

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

nirandaperera requested a review from wence- May 5, 2026 14:55

wence- reviewed May 5, 2026

View reviewed changes

wence- approved these changes May 5, 2026

View reviewed changes

nirandaperera and others added 3 commits May 5, 2026 09:48

Apply suggestions from code review

c2e1454

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

addressing comments

e11ee85

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

Merge branch 'main' into remove-mrs-from-statistics

b295870

rapids-bot Bot merged commit 94e6bb8 into rapidsai:main May 6, 2026
66 checks passed

		@@ -135,6 +153,16 @@ class BufferResource {
		*/
		[[nodiscard]] rmm::host_async_resource_ref pinned_mr();

	return resource ? std::make_optional(*resource) : std::nullopt;
	return resource ? *resource : std::nullopt;

	std::optional<any_host_device_resource> pinned_mr = std::nullopt,
	std::optional<any_host_device_resource> pinned_mr = PinnedMemoryResource::Disabled,

	return pinned_mr_.has_value() ? std::make_optional(*pinned_mr_) : std::nullopt;
	return pinned_mr_.has_value() ? *pinned_mr_ : std::nullopt;

		include("${rapids-cmake-dir}/cpm/package_override.cmake")
		rapids_cpm_package_override("${CMAKE_CURRENT_SOURCE_DIR}/../cmake/cpm_override_packages.json")


		include("${rapids-cmake-dir}/cpm/package_override.cmake")
		rapids_cpm_package_override("${CMAKE_CURRENT_SOURCE_DIR}/../../cmake/cpm_override_packages.json")

Uh oh!

Conversation

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

Uh oh!

Remove memory resources from Statistics construction

Solution

Uh oh!

copy-pr-bot Bot commented Apr 29, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

nirandaperera commented May 5, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

wence- May 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

wence- left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

nirandaperera commented May 6, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

nirandaperera commented Apr 29, 2026 •

edited

Loading

Remove memory resources from `Statistics` construction

wence- May 5, 2026 •

edited

Loading