Skip to content

Remove memory resources from Statistics#1003

Merged
rapids-bot[bot] merged 15 commits intorapidsai:mainfrom
nirandaperera:remove-mrs-from-statistics
May 6, 2026
Merged

Remove memory resources from Statistics#1003
rapids-bot[bot] merged 15 commits intorapidsai:mainfrom
nirandaperera:remove-mrs-from-statistics

Conversation

@nirandaperera
Copy link
Copy Markdown
Contributor

@nirandaperera nirandaperera commented Apr 29, 2026

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

@nirandaperera nirandaperera requested review from a team as code owners April 29, 2026 20:46
@nirandaperera nirandaperera added breaking Introduces a breaking change improvement Improves an existing functionality labels Apr 29, 2026
@nirandaperera nirandaperera marked this pull request as draft April 29, 2026 20:47
@copy-pr-bot
Copy link
Copy Markdown

copy-pr-bot Bot commented Apr 29, 2026

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>
Signed-off-by: niranda perera <niranda.perera@gmail.com>
Signed-off-by: niranda perera <niranda.perera@gmail.com>
@nirandaperera nirandaperera force-pushed the remove-mrs-from-statistics branch from a7bce3f to ac296bc Compare April 30, 2026 18:32
@nirandaperera nirandaperera marked this pull request as ready for review April 30, 2026 18:32
Signed-off-by: niranda perera <niranda.perera@gmail.com>
Signed-off-by: niranda perera <niranda.perera@gmail.com>
@nirandaperera nirandaperera mentioned this pull request Apr 30, 2026
3 tasks
Comment thread cmake/cpm_override_packages.json Outdated
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 think this file should be removed.

@@ -135,6 +153,16 @@ class BufferResource {
*/
[[nodiscard]] rmm::host_async_resource_ref pinned_mr();
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.

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

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;
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
return resource ? std::make_optional(*resource) : std::nullopt;
return resource ? *resource : std::nullopt;

Comment on lines +161 to +164
[[nodiscard]] constexpr const std::optional<PinnedMemoryResource>&
concrete_pinned_mr() const noexcept {
return pinned_mr_;
}
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 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>
Signed-off-by: niranda perera <niranda.perera@gmail.com>
@nirandaperera nirandaperera requested a review from wence- May 5, 2026 04:54
@nirandaperera
Copy link
Copy Markdown
Contributor Author

@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?

Comment thread cmake/cpm_override_packages.json Outdated
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.

Comment on lines +460 to +468
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):"
));
}
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.

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

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.

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

Comment thread cpp/include/rapidsmpf/statistics.hpp Outdated
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,
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
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);

?

Comment thread cpp/src/memory/buffer_resource.cpp Outdated
}

std::optional<any_host_device_resource> BufferResource::try_pinned_mr() const noexcept {
return pinned_mr_.has_value() ? std::make_optional(*pinned_mr_) : std::nullopt;
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.

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?

Comment thread cpp/src/statistics.cpp
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));
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.

question: What happens if the resource_cast fails?

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.

it returns a nullptr

Comment thread cpp/CMakeLists.txt Outdated
Comment on lines +126 to +128
include("${rapids-cmake-dir}/cpm/package_override.cmake")
rapids_cpm_package_override("${CMAKE_CURRENT_SOURCE_DIR}/../cmake/cpm_override_packages.json")

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 think you are not using this.

def report(
self,
mr: RmmResourceAdaptor | None = None,
pinned_mr: PinnedMemoryResource | None = None,
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.

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

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.

I havent seen any uses of it so far.

Comment thread cufile.log Outdated
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.

Please remove.

Comment thread python/rapidsmpf/CMakeLists.txt Outdated
Comment on lines +12 to +15

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

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.

Not needed.

Signed-off-by: niranda perera <niranda.perera@gmail.com>
@nirandaperera nirandaperera requested a review from wence- May 5, 2026 14:55
Comment on lines +207 to +215
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{});
}
Copy link
Copy Markdown
Contributor

@wence- wence- May 5, 2026

Choose a reason for hiding this comment

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

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;

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.

clang didnt like this. 🙁

Copy link
Copy Markdown
Contributor

@wence- wence- left a comment

Choose a reason for hiding this comment

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

Tiny suggested cleanup

Comment thread python/rapidsmpf/rapidsmpf/statistics.pyx
self,
RmmResourceAdaptor mr = None,
PinnedMemoryResource pinned_mr = None,
):
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.

Please allow caller to provide header as well.

Comment thread python/rapidsmpf/rapidsmpf/statistics.pyi
nirandaperera and others added 3 commits May 5, 2026 09:48
Co-authored-by: Lawrence Mitchell <wence@gmx.li>
Signed-off-by: niranda perera <niranda.perera@gmail.com>
@nirandaperera
Copy link
Copy Markdown
Contributor Author

/merge

@rapids-bot rapids-bot Bot merged commit 94e6bb8 into rapidsai:main May 6, 2026
66 checks passed
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.

Remove memory resources from Statistics object, and pass them along when creating memory records.

2 participants