Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[FEA] Add RAFT wrappers for RMM current_device_resource functions. #2423

Open
harrism opened this issue Aug 28, 2024 · 0 comments
Open

[FEA] Add RAFT wrappers for RMM current_device_resource functions. #2423

harrism opened this issue Aug 28, 2024 · 0 comments
Assignees
Labels
feature request New feature or request

Comments

@harrism
Copy link
Member

harrism commented Aug 28, 2024

Is your feature request related to a problem? Please describe.
I'm about to undertake yet another change to every function that takes an MR parameter. To prevent doing this too many times in the future, I would like to provide a wrapper in RAFT to centralize calls to rmm::get_current_device_resource_ref and aliases for RMM resource_ref types. hopefully this can reduce the number of files that need to change in the future when this functionality changes (e.g. if it moves out of RMM to CCCL).

Describe the solution you'd like
Either add the following to raft/core/resource/device_memory_resource.hpp, or add a new header, say raft/core/resource/memory_resource.hpp. Please provide suggestions on the best place for this. Contents something like this:

namespace raft {

using device_async_resource_ref = rmm::device_async_resource_ref;

inline device_async_resource_ref get_current_device_resource_ref() {
  return rmm::get_current_device_resource_ref();
}

inline device_async_resource_ref set_current_device_resource_ref(device_async_resource_ref mr) {
  return rmm::set_current_device_resource_ref(mr);
}

inline device_async_resource_ref reset_current_device_resource_ref() {
  return rmm::reset_current_device_resource_ref();
}

Describe alternatives you've considered

Additional context
rapidsai/rmm#1598

related cudf PR: rapidsai/cudf#16679

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
feature request New feature or request
Projects
Status: Todo
Development

Successfully merging a pull request may close this issue.

1 participant