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 <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...
RMM_EXPORT auto &	rmm::mr::detail::get_map ()
	Reference to the map from device id -> resource. 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...

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() will automatically use the current CUDA device id from cudaGetDevice().

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());
}

Function Documentation

get_map()

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

inline

Reference to the map from device id -> resource.

Returns

Reference to the map from device id -> resource

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