RMM: RAPIDS Memory Manager

Achieving optimal performance in GPU-centric workflows frequently requires customizing how host and device memory are allocated. For example, using "pinned" host memory for asynchronous host <-> device memory transfers, or using a device memory pool sub-allocator to reduce the cost of dynamic device memory allocation.

The goal of the RAPIDS Memory Manager (RMM) is to provide: - A common interface that allows customizing device and host memory allocation - A collection of implementations of the interface - A collection of data structures that use the interface for memory allocation

For information on the interface RMM provides and how to use RMM in your C++ code, see below.

NOTE: For the latest stable README.md ensure you are on the

main

Installation

Conda

RMM can be installed with Conda (miniconda, or the full Anaconda distribution) from the

rapidsai

# for CUDA 11.2
conda install -c nvidia -c rapidsai -c conda-forge \
    rmm cudatoolkit=11.2
# for CUDA 11.1
conda install -c nvidia -c rapidsai -c conda-forge \
    rmm cudatoolkit=11.1
# for CUDA 11.0
conda install -c nvidia -c rapidsai -c conda-forge \
    rmm cudatoolkit=11.0

We also provide nightly Conda packages built from the HEAD of our latest development branch.

Note: RMM is supported only on Linux, and with Python versions 3.7 and later.

Note: The RMM package from Conda requires building with GCC 9 or later. Otherwise, your application may fail to build.

See the Get RAPIDS version picker for more OS and version info.

Building from Source

Get RMM Dependencies

Compiler requirements:

• gcc

  version 9.3+

• nvcc

  version 11.0+

• cmake

  version 3.20.1+

CUDA/GPU requirements:

• CUDA 11.0+
• NVIDIA driver 450.51+
• Pascal architecture or better

You can obtain CUDA from https://developer.nvidia.com/cuda-downloads

Script to build RMM from source

To install RMM from source, ensure the dependencies are met and follow the steps below:

• Clone the repository and submodules

  bash
  $ git clone --recurse-submodules https://github.com/rapidsai/rmm.git
  $ cd rmm
  

• Create the conda development environment

  rmm_dev

  ```bash

  create the conda environment (assuming in base

  rmm

  directory)

  $ conda env create --name rmmdev --file conda/environments/rmmdev_cuda11.0.yml

  activate the environment

  $ conda activate rmm_dev ```

• Build and install

  librmm

  using cmake & make. CMake depends on the

  nvcc

  executable being on your path or defined in

  CUDACXX

  environment variable.

$ mkdir build                                       # make a build directory
$ cd build                                          # enter the build directory
$ cmake .. -DCMAKE_INSTALL_PREFIX=/install/path     # configure cmake ... use $CONDA_PREFIX if you're using Anaconda
$ make -j                                           # compile the library librmm.so ... '-j' will start a parallel job using the number of physical cores available on your system
$ make install                                      # install the library librmm.so to '/install/path'

• Building and installing

  librmm

  and

  rmm

  using build.sh. Build.sh creates build dir at root of git repository. build.sh depends on the

  nvcc

  executable being on your path or defined in

  CUDACXX

  environment variable.

$ ./build.sh -h                                     # Display help and exit
$ ./build.sh -n librmm                              # Build librmm without installing
$ ./build.sh -n rmm                                 # Build rmm without installing
$ ./build.sh -n librmm rmm                          # Build librmm and rmm without installing
$ ./build.sh librmm rmm                             # Build and install librmm and rmm

• To run tests (Optional):

  bash
  $ cd build (if you are not already in build directory)
  $ make test
  

• Build, install, and test the

  rmm

  python package, in the

  python

  folder:

  bash
  $ python setup.py build_ext --inplace
  $ python setup.py install
  $ pytest -v
  

Done! You are ready to develop for the RMM OSS project.

Caching third-party dependencies

RMM uses CPM.cmake to handle third-party dependencies like spdlog, Thrust, GoogleTest, GoogleBenchmark. In general you won't have to worry about it. If CMake finds an appropriate version on your system, it uses it (you can help it along by setting

CMAKE_PREFIX_PATH

If you frequently start new builds from scratch, consider setting the environment variable

CPM_SOURCE_CACHE

Using RMM in a downstream CMake project

The installed RMM library provides a set of config files that makes it easy to integrate RMM into your own CMake project. In your

CMakeLists.txt

find_package(rmm [VERSION])
# ...
target_link_libraries( (PRIVATE|PUBLIC) rmm::rmm)

Since RMM is a header-only library, this does not actually link RMM, but it makes the headers available and pulls in transitive dependencies. If RMM is not installed in a default location, use

CMAKE_PREFIX_PATH

rmm_ROOT

One of RMM's dependencies is the Thrust library, so the above automatically pulls in

Thrust

rmm::Thrust

THRUST_HOST_SYSTEM

THRUST_DEVICE_SYSTEM

Using RMM in C++

The first goal of RMM is to provide a common interface for device and host memory allocation. This allows both users and implementers of custom allocation logic to program to a single interface.

To this end, RMM defines two abstract interface classes: -

rmm::mr::device_memory_resource

rmm::mr::host_memory_resource

These classes are based on the

std::pmr::memory_resource

device_memory_resource

rmm::mr::device_memory_resource

It has two key functions:

1. void* device_memory_resource::allocate(std::size_t bytes, cuda_stream_view s)

   • Returns a pointer to an allocation of at least

     bytes

     bytes.

2. void device_memory_resource::deallocate(void* p, std::size_t bytes, cuda_stream_view s)

   • Reclaims a previous allocation of size

     bytes

     pointed to by

     p

     .

   • p

     must have been returned by a previous call to

     allocate(bytes)

     , otherwise behavior is undefined

It is up to a derived class to provide implementations of these functions. See available resources for example

device_memory_resource

Unlike

std::pmr::memory_resource

rmm::mr::device_memory_resource

device_memory_resource

cuda_stream_view

cuda_stream_view

and

cuda_stream

rmm::cuda_stream_view

cudaStream_t

cudaStream_t

0

rmm::cuda_stream_view

rmm::cuda_stream

cudaStream_t

rmm::cuda_stream

rmm::cuda_stream

cuda_stream_pool

rmm::cuda_stream_pool

cuda_stream

cuda_stream_pool

cuda_stream_pool::get_stream()

rmm::cuda_stream_view

Thread Safety

All current device memory resources are thread safe unless documented otherwise. More specifically, calls to memory resource

allocate()

deallocate()

Note that a class

thread_safe_resource_adapter

Stream-ordered Memory Allocation

rmm::mr::device_memory_resource

A call to

device_memory_resource::allocate(bytes, stream_a)

stream_a

stream_b

cudaStreamSynchronize(stream_a)

stream_a

cudaStreamWaitEvent(stream_b, event)

The stream specified to

device_memory_resource::deallocate

deallocate

device_memory_resource

cudaStreamSynchronize()

For this reason, it is Undefined Behavior to destroy a CUDA stream that is passed to

device_memory_resource::deallocate

deallocate

deallocate

Note that device memory data structures such as

rmm::device_buffer

rmm::device_uvector

Available Resources

RMM provides several

device_memory_resource

cuda_memory_resource

Allocates and frees device memory using

cudaMalloc

cudaFree

managed_memory_resource

Allocates and frees device memory using

cudaMallocManaged

cudaFree

Note that

managed_memory_resource

pool_memory_resource

A coalescing, best-fit pool sub-allocator.

fixed_size_memory_resource

A memory resource that can only allocate a single fixed size. Average allocation and deallocation cost is constant.

binning_memory_resource

Configurable to use multiple upstream memory resources for allocations that fall within different bin sizes. Often configured with multiple bins backed by

fixed_size_memory_resource

pool_memory_resource

Default Resources and Per-device Resources

RMM users commonly need to configure a

device_memory_resource

pool_memory_resource

To enable this use case, RMM provides the concept of a "default"

device_memory_resource

Accessing and modifying the default resource is done through two functions: -

device_memory_resource* get_current_device_resource()

cuda_memory_resource

set_current_device_resource()

get_per_device_resource()

• device_memory_resource* set_current_device_resource(device_memory_resource* new_mr)

  • Updates the default memory resource pointer for the current CUDA device to

    new_mr

  • Returns the previous default resource pointer
  • If

    new_mr

    is

    nullptr

    , then resets the default resource to

    cuda_memory_resource

  • This function is thread safe with respect to concurrent calls to it and

    get_current_device_resource()

  • For more explicit control, you can use

    set_per_device_resource()

    , which takes a device ID.

Example

rmm::mr::cuda_memory_resource cuda_mr;
// Construct a resource that uses a coalescing best-fit pool allocator
rmm::mr::pool_memory_resource<:mr::cuda_memory_resource> pool_mr{&cuda_mr};
rmm::mr::set_current_device_resource(&pool_mr); // Updates the current device resource pointer to `pool_mr`
rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource(); // Points to `pool_mr`

Multiple Devices

A

device_memory_resource

device_memory_resource

If a

device_memory_resource

Creating a

device_memory_resource

unique_ptr

pool_memory_resource

std::vector> per_device_pools;
for(int i = 0; i < N; ++i) {
    cudaSetDevice(i); // set device i before creating MR
    // Use a vector of unique_ptr to maintain the lifetime of the MRs
    per_device_pools.push_back(std::make_unique());
    // Set the per-device resource for device i
    set_per_device_resource(cuda_device_id{i}, &per_device_pools.back());
}

Allocators

C++ interfaces commonly allow customizable memory allocation through an

Allocator

Allocator

Allocator

polymorphic_allocator

A stream-ordered allocator similar to

std::pmr::polymorphic_allocator

Allocator

allocate

deallocate

cuda_stream_view

stream_allocator_adaptor

stream_allocator_adaptor

Allocator

Example: ```c++ rmm::cudastream stream; rmm::mr::polymorphicallocator stream_alloc;

// Constructs an adaptor that forwards all (de)allocations to

stream_alloc

stream

// Allocates 100 bytes using

stream_alloc

stream

stream_alloc

stream

thrust_allocator

thrust_allocator

thrust::device_ptr

thrust::device_vector

See below for more information on using RMM with Thrust.

Device Data Structures

device_buffer

An untyped, uninitialized RAII class for stream ordered device memory allocation.

Example

cuda_stream_view s{...};
// Allocates at least 100 bytes on stream `s` using the *default* resource
rmm::device_buffer b{100,s}; 
void* p = b.data();                   // Raw, untyped pointer to underlying device memory

kernel<<<... s.value>>>(b.data()); // b is only safe to use on s
rmm::mr::device_memory_resource * mr = new my_custom_resource{...};
// Allocates at least 100 bytes on stream s using the resource mr
rmm::device_buffer b2{100, s, mr};
</...>

device_uvector

A typed, uninitialized RAII class for allocation of a contiguous set of elements in device memory. Similar to a

thrust::device_vector

T

Example

cuda_stream_view s{...};
// Allocates uninitialized storage for 100 `int32_t` elements on stream `s` using the
// default resource
rmm::device_uvector v(100, s);
// Initializes the elements to 0
thrust::uninitialized_fill(thrust::cuda::par.on(s.value()), v.begin(), v.end(), int32_t{0}); 

rmm::mr::device_memory_resource * mr = new my_custom_resource{...};
// Allocates uninitialized storage for 100 int32_t elements on stream s using the resource mr
rmm::device_uvector v2{100, s, mr}; 

device_scalar

A typed, RAII class for allocation of a single element in device memory. This is similar to a

device_uvector

Example

cuda_stream_view s{...};
// Allocates uninitialized storage for a single `int32_t` in device memory
rmm::device_scalar a{s}; 
a.set_value(42, s); // Updates the value in device memory to `42` on stream `s`

kernel<<<...>>>(a.data()); // Pass raw pointer to underlying element in device memory
int32_t v = a.value(s); // Retrieves the value from device to host on stream s
</...>

host_memory_resource

rmm::mr::host_memory_resource

Similar to

device_memory_resource

1. void* host_memory_resource::allocate(std::size_t bytes, std::size_t alignment)

   • Returns a pointer to an allocation of at least

     bytes

     bytes aligned to the specified

     alignment

2. void host_memory_resource::deallocate(void* p, std::size_t bytes, std::size_t alignment)

   • Reclaims a previous allocation of size

     bytes

     pointed to by

     p

     .

Unlike

device_memory_resource

host_memory_resource

std::pmr::memory_resource

Available Resources

new_delete_resource

Uses the global

operator new

operator delete

pinned_memory_resource

Allocates "pinned" host memory using

cuda(Malloc/Free)Host

Host Data Structures

RMM does not currently provide any data structures that interface with

host_memory_resource

device_buffer

Using RMM with Thrust

RAPIDS and other CUDA libraries make heavy use of Thrust. Thrust uses CUDA device memory in two situations:

1. As the backing store for

   thrust::device_vector

   , and
2. As temporary storage inside some algorithms, such as

   thrust::sort

   .

RMM provides

rmm::mr::thrust_allocator

device_memory_resource

Thrust Algorithms

To instruct a Thrust algorithm to use

rmm::mr::thrust_allocator

rmm::exec_policy(stream)

thrust::sort(rmm::exec_policy(stream, ...);

The first

stream

stream

rmm::mr::thrust_allocator

stream

Logging

RMM includes two forms of logging. Memory event logging and debug logging.

Memory Event Logging and

logging_resource_adaptor

Memory event logging writes details of every allocation or deallocation to a CSV (comma-separated value) file. In C++, Memory Event Logging is enabled by using the

logging_resource_adaptor

device_memory_resource

Each row in the log represents either an allocation or a deallocation. The columns of the file are "Thread, Time, Action, Pointer, Size, Stream".

The CSV output files of the

logging_resource_adaptor

REPLAY_BENCHMARK

gbenchmarks

The following C++ example creates a logging version of a

cuda_memory_resource

std::string filename{"logs/test1.csv"};
rmm::mr::cuda_memory_resource upstream;
rmm::mr::logging_resource_adaptor<:mr::cuda_memory_resource> log_mr{&upstream, filename};

If a file name is not specified, the environment variable

RMM_LOG_FILE

RMM_LOG_FILE

logging_resource_adaptor

In Python, memory event logging is enabled when the

logging

rmm.reinitialize()

True

log_file_name

help(rmm.reinitialize)

Debug Logging

RMM includes a debug logger which can be enabled to log trace and debug information to a file. This information can show when errors occur, when additional memory is allocated from upstream resources, etc. The default log file is

rmm_log.txt

RMM_DEBUG_LOG_FILE

There is a CMake configuration variable

RMM_LOGGING_LEVEL

INFO

TRACE

DEBUG

INFO

WARN

ERROR

CRITICAL

OFF

The log relies on the spdlog library.

Note that to see logging below the

INFO

rmm::logger().set_level()

TRACE

rmm::logger().set_level(spdlog::level::trace)

-DRMM_LOGGING_LEVEL=TRACE

Note that debug logging is different from the CSV memory allocation logging provided by

rmm::mr::logging_resource_adapter

RMM and CUDA Memory Bounds Checking

Memory allocations taken from a memory resource that allocates a pool of memory (such as

pool_memory_resource

arena_memory_resource

Exceptions to this are

cuda_memory_resource

cudaMalloc

cuda_async_memory_resource

cudaMallocAsync

It may be possible in the future to add support for memory bounds checking with other memory resources using NVTX APIs.

Using RMM in Python Code

There are two ways to use RMM in Python code:

1. Using the

   rmm.DeviceBuffer

   API to explicitly create and manage device memory allocations
2. Transparently via external libraries such as CuPy and Numba

RMM provides a

MemoryResource

DeviceBuffers

A DeviceBuffer represents an untyped, uninitialized device memory allocation. DeviceBuffers can be created by providing the size of the allocation in bytes:

>>> import rmm
>>> buf = rmm.DeviceBuffer(size=100)

The size of the allocation and the memory address associated with it can be accessed via the

.size

.ptr

>>> buf.size
100
>>> buf.ptr
140202544726016

DeviceBuffers can also be created by copying data from host memory:

>>> import rmm
>>> import numpy as np
>>> a = np.array([1, 2, 3], dtype='float64')
>>> buf = rmm.to_device(a.tobytes())
>>> buf.size
24

Conversely, the data underlying a DeviceBuffer can be copied to the host:

>>> np.frombuffer(buf.tobytes())
array([1., 2., 3.])

MemoryResource objects

MemoryResource

By default if a

MemoryResource

CudaMemoryResource

cudaMalloc

rmm.reinitialize()

help(rmm.reinitialize)

For lower-level control, the

rmm.mr.set_current_device_resource()

ManagedMemoryResource

cudaMallocManaged

cudaMalloc

>>> import rmm
>>> rmm.mr.set_current_device_resource(rmm.mr.ManagedMemoryResource())

:warning: The default resource must be set for any device before allocating any device memory on that device. Setting or changing the resource after device allocations have been made can lead to unexpected behaviour or crashes. See Multiple Devices

As another example,

PoolMemoryResource

CudaMemoryResource

>>> import rmm
>>> pool = rmm.mr.PoolMemoryResource(
...     rmm.mr.CudaMemoryResource(),
...     initial_pool_size=2**30,
...     maximum_pool_size=2**32
... )
>>> rmm.mr.set_current_device_resource(pool)

Other MemoryResources include:

• FixedSizeMemoryResource

  for allocating fixed blocks of memory

• BinningMemoryResource

  for allocating blocks within specified "bin" sizes from different memory resources

MemoryResources are highly configurable and can be composed together in different ways. See

help(rmm.mr)

Using RMM with CuPy

You can configure CuPy to use RMM for memory allocations by setting the CuPy CUDA allocator to

rmm_cupy_allocator

>>> import rmm
>>> import cupy
>>> cupy.cuda.set_allocator(rmm.rmm_cupy_allocator)

Using RMM with Numba

You can configure Numba to use RMM for memory allocations using the Numba EMM Plugin.

This can be done in two ways:

1. Setting the environment variable

   NUMBA_CUDA_MEMORY_MANAGER

   :

  $ NUMBA_CUDA_MEMORY_MANAGER=rmm python (args)

1. Using the

   set_memory_manager()

   function provided by Numba:

  >>> from numba import cuda
  >>> import rmm
  >>> cuda.set_memory_manager(rmm.RMMNumbaManager)