raft icon indicating copy to clipboard operation
raft copied to clipboard

[FEA] Add RAFT wrappers for RMM current_device_resource functions.

Open harrism opened this issue 5 months ago • 0 comments

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 https://github.com/rapidsai/rmm/pull/1598

related cudf PR: https://github.com/rapidsai/cudf/pull/16679

harrism avatar Aug 28 '24 03:08 harrism