Management of per-device device_memory_resource
s.
More...
#include <rmm/cuda_device.hpp>
#include <rmm/detail/export.hpp>
#include <rmm/mr/device/cuda_memory_resource.hpp>
#include <rmm/mr/device/device_memory_resource.hpp>
#include <rmm/resource_ref.hpp>
#include <map>
#include <mutex>
Go to the source code of this file.
Functions | |
device_memory_resource * | rmm::mr::detail::initial_resource () |
Returns a pointer to the initial resource. More... | |
std::mutex & | rmm::mr::detail::map_lock () |
Reference to the lock. More... | |
auto & | rmm::mr::detail::get_map () |
Reference to the map from device id -> resource. More... | |
std::mutex & | rmm::mr::detail::ref_map_lock () |
Reference to the lock. More... | |
auto & | rmm::mr::detail::get_ref_map () |
Reference to the map from device id -> resource_ref. More... | |
device_memory_resource * | rmm::mr::get_per_device_resource (cuda_device_id device_id) |
Get the resource for the specified device. More... | |
device_memory_resource * | rmm::mr::set_per_device_resource (cuda_device_id device_id, device_memory_resource *new_mr) |
Set the device_memory_resource for the specified device. More... | |
device_memory_resource * | rmm::mr::get_current_device_resource () |
Get the memory resource for the current device. More... | |
device_memory_resource * | rmm::mr::set_current_device_resource (device_memory_resource *new_mr) |
Set the memory resource for the current device. More... | |
device_async_resource_ref | rmm::mr::get_per_device_resource_ref (cuda_device_id device_id) |
Get the device_async_resource_ref for the specified device. More... | |
device_async_resource_ref | rmm::mr::set_per_device_resource_ref (cuda_device_id device_id, device_async_resource_ref new_resource_ref) |
Set the device_async_resource_ref for the specified device to new_resource_ref More... | |
device_async_resource_ref | rmm::mr::get_current_device_resource_ref () |
Get the device_async_resource_ref for the current device. More... | |
device_async_resource_ref | rmm::mr::set_current_device_resource_ref (device_async_resource_ref new_resource_ref) |
Set the device_async_resource_ref for the current device. More... | |
device_async_resource_ref | rmm::mr::reset_per_device_resource_ref (cuda_device_id device_id) |
Reset the device_async_resource_ref for the specified device to the initial resource. More... | |
device_async_resource_ref | rmm::mr::reset_current_device_resource_ref () |
Reset the device_async_resource_ref for the current device to the initial resource. More... | |
Management of per-device device_memory_resource
s.
One might wish to construct a device_memory_resource
and use it for (de)allocation without explicit dependency injection, i.e., passing a reference to that object to all places it is to be used. Instead, one might want to set their resource as a "default" and have it be used in all places where another resource has not been explicitly specified. In applications with multiple GPUs in the same process, it may also be necessary to maintain independent default resources for each device. To this end, the set_per_device_resource
and get_per_device_resource
functions enable mapping a CUDA device id to a device_memory_resource
pointer.
For example, given a pointer, mr
, to a device_memory_resource
object, calling set_per_device_resource(cuda_device_id{0}, mr)
will establish a mapping between CUDA device 0 and mr
such that all future calls to get_per_device_resource(cuda_device_id{0})
will return the same pointer, mr
. In this way, all places that use the resource returned from get_per_device_resource
for (de)allocation will use the user provided resource, mr
.
device_memory_resource
s make CUDA API calls without setting the current CUDA device. Therefore a memory resource should only be used with the current CUDA device set to the device that was active when the memory resource was created. Calling set_per_device_resource(id, mr)
is only valid if id
refers to the CUDA device that was active when mr
was created.If no resource was explicitly set for a given device specified by id
, then get_per_device_resource(id)
will return a pointer to a cuda_memory_resource
.
To fetch and modify the resource for the current CUDA device, get_current_device_resource()
and set_current_device_resource()
automatically use the current CUDA device id from cudaGetDevice()
.
RMM is in transition to use cuda::mr::async_resource_ref
in place of raw pointers to device_memory_resource
. The set_per_device_resource_ref
, get_per_device_resource_ref
, get_current_device_resource_ref
, set_current_device_resource_ref
, and reset_current_device_resource_ref
functions provide the same functionality as their device_memory_resource
counterparts, but with device_async_resource_ref
objects. The raw pointer versions and the resource_ref
versions maintain distinc state and are not interchangeable. The raw pointer versions are expected to be deprecated and removed in a future release.
Creating a device_memory_resource for each device requires care to set the current device before creating each resource, and to maintain the lifetime of the resources as long as they are set as per-device resources. Here is an example loop that creates unique_ptr
s to pool_memory_resource objects for each device and sets them as the per-device resource for that device.
|
inline |
Reference to the map from device id -> resource.
|
inline |
Reference to the map from device id -> resource_ref.
|
inline |
Returns a pointer to the initial resource.
Returns a global instance of a cuda_memory_resource
as a function local static.
|
inline |
Reference to the lock.
|
inline |
Reference to the lock.