[FEA] Support specifying a host mr for `rmm::device_scalar`

[vuule]

The host resource would be used for the bounce buffer in value(). The potential optimization would be to pass a pinned memory resource to avoid a pageable copy in value(). cudf::detail::device_scalar already uses a similar implementation.

[JigaoLuo]

@vuule Thank you for opening this issue. As an external contributor, I discovered this issue in cuDF, documented it in rapidsai/cudf#18967, and proposed a performance-optimized fix in rapidsai/cudf#18968.

Once my current PRs in cuDF are merged, I’d be happy to volunteer to implement this FEA in RMM. I understand that external contributions require additional code review passes, but I’m committed to ensuring the changes meet RMM standards. I’ve already contributed to cuDF and am happy to continue collaborating with you.

[bdice]

@JigaoLuo Thanks for your interest in this! Yes, happy to work with you.

[JigaoLuo]

@bdice We’ve decided to drop the PR in cuDF rapidsai/cudf#19119 and instead add pinned memory as an intermediate copy buffer in rmm::device_scalar, which would benefit multiple projects beyond cuDF.

I’d appreciate guidance in RMM, as it has fewer utility functions than cuDF. Specifically, I’m considering a built-in bounce_buffer as in https://github.com/rapidsai/cudf/blob/branch-25.08/cpp/include/cudf/detail/device_scalar.hpp, but I’m unsure if forcing a single pinned buffer in RMM is the right approach.

Let me know if you’d prefer to discuss this here or on Slack - I’m available either way! Thanks

[bdice]

We can discuss here. That proposal seems fine to me, if it can be made non-breaking for existing rmm::device_scalar users.

[vuule]

my two cents: cudf's version forces a pinned bounce buffer because there is a pinned pool by default. I don't think rmm should default to a pinned allocation for the bounce buffer, because (AFAIK) the cost of a pinned allocation is higher than the overhead of the pageable copy.

[JigaoLuo]

@vuule Thanks for raising this valuable point. Actually, I have the same concern.

• In cuDF: I also did a PR to"force" use of a pinned pool (pinned vector) in cudf::detail::device_scalar.
• But in RMM: inside rmm::device_scalar, a pinned_host_memory_resource might be needed. Currently, there is no host memory resource inside rmm::device_scalar.

[vuule]

There might be two pieces required here:

1. A default host resource (global function that returns a rmm::host_resource_ref, with the default probably being a rmm::mr::new_delete_resource.
2. Changes to rmm::device_scalar, including added rmm::host_resource_ref paramater that defaults to the global function above.

[bdice]

I was thinking we could add an optional host_resource_ref to the device_scalar constructor without needing to add a global host resource. Adding more global state opens all kinds of new design issues in RMM and I want to avoid that if possible.

Applications (cuDF) can manage their own pinned pool allocator, but that host resource should be passed explicitly when constructing the device_scalar. It will be up to the caller of the device_scalar constructor to ensure that the passed host resource remains alive for the duration of the scalar's lifetime.

Note: lifetime management of resources will get easier in the future, as there is semi-shared ownership of memory resources in the CCCL experimental memory resources. We will adopt that in the future as that code stabilizes.

[vuule]

I was thinking we could add an optional host_resource_ref to the device_scalar constructor without needing to add a global host resource.

Doesn't it need a default value to be optional? Whatever the default references needs to be stored somewhere.
Unless you're suggestion a ctor overload that takes a host_resource_ref?

[JigaoLuo]

After a quick code review, I agree with the first step suggested by @vuule . Specifically, a function like mr::get_current_host_resource_ref() seems necessary to me. With this function, we can safely overload the constructor of device_scalar without breaking changes.

I understand the priority to avoid breaking changes and global state. But, currently, the existing constructors need a default value (or a call) to be optional.

A copy of the existing constructors

  explicit device_scalar(cuda_stream_view stream,
                         device_async_resource_ref mr = mr::get_current_device_resource_ref())
    : _storage{1, stream, mr} {}

  explicit device_scalar(value_type const& initial_value,
                         cuda_stream_view stream,
                         device_async_resource_ref mr = mr::get_current_device_resource_ref())
    : _storage{1, stream, mr}
  {
    set_value_async(initial_value, stream);
  }

  device_scalar(device_scalar const& other,
                cuda_stream_view stream,
                device_async_resource_ref mr = mr::get_current_device_resource_ref())
    : _storage{other._storage, stream, mr} {}