per_device_resource.hpp File Reference

Management of per-device device_memory_resources. 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>

Include dependency graph for per_device_resource.hpp:

This graph shows which files directly or indirectly include this file:

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...

Detailed Description

Management of per-device device_memory_resources.

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.

Note

device_memory_resources 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_ptrs to pool_memory_resource objects for each device and sets them as the per-device resource for that device.

std::vector<unique_ptr<pool_memory_resource>> per_device_pools;
for(int i = 0; i < N; ++i) {
  cudaSetDevice(i);
  per_device_pools.push_back(std::make_unique<pool_memory_resource>());
  set_per_device_resource(cuda_device_id{i}, &per_device_pools.back());
}

using pool_mr = rmm::mr::pool_memory_resource<rmm::mr::cuda_memory_resource>;
std::vector<unique_ptr<pool_mr>> per_device_pools;
for(int i = 0; i < N; ++i) {
  cudaSetDevice(i);
  // Note: for brevity, omitting creation of upstream and computing initial_size
  per_device_pools.push_back(std::make_unique<pool_mr>(upstream, initial_size));
  set_per_device_resource(cuda_device_id{i}, &per_device_pools.back());
}

Function Documentation

get_map()

auto& rmm::mr::detail::get_map ( )

inline

Reference to the map from device id -> resource.

Returns

Reference to the map from device id -> resource

get_ref_map()

auto& rmm::mr::detail::get_ref_map ( )

inline

Reference to the map from device id -> resource_ref.

Returns

Reference to the map from device id -> resource_ref

initial_resource()

device_memory_resource* rmm::mr::detail::initial_resource ( )

inline

Returns a pointer to the initial resource.

Returns a global instance of a cuda_memory_resource as a function local static.

Returns

Pointer to the static cuda_memory_resource used as the initial, default resource

map_lock()

std::mutex& rmm::mr::detail::map_lock ( )

inline

Reference to the lock.

Returns

Reference to the lock

ref_map_lock()

std::mutex& rmm::mr::detail::ref_map_lock ( )

inline

Reference to the lock.

Returns

Reference to the lock