|
| 1 | +# RMM - the RAPIDS Memory Manager |
| 2 | + |
| 3 | +Achieving optimal performance in GPU-centric workflows frequently requires |
| 4 | +customizing how GPU ("device") memory is allocated. |
| 5 | + |
| 6 | +RMM is a package that enables you to allocate device memory |
| 7 | +in a highly configurable way. For example, it enables you to |
| 8 | +allocate and use pools of GPU memory, or to use |
| 9 | +[managed memory](https://developer.nvidia.com/blog/unified-memory-cuda-beginners/) |
| 10 | +for allocations. |
| 11 | + |
| 12 | +You can also easily configure other libraries like Numba and CuPy |
| 13 | +to use RMM for allocating device memory. |
| 14 | + |
| 15 | +## Installation |
| 16 | + |
| 17 | +See the project [README](https://github.com/rapidsai/rmm) for how to install RMM. |
| 18 | + |
| 19 | +## Using RMM |
| 20 | + |
| 21 | +There are two ways to use RMM in Python code: |
| 22 | + |
| 23 | +1. Using the `rmm.DeviceBuffer` API to explicitly create and manage |
| 24 | + device memory allocations |
| 25 | +2. Transparently via external libraries such as CuPy and Numba |
| 26 | + |
| 27 | +RMM provides a `MemoryResource` abstraction to control _how_ device |
| 28 | +memory is allocated in both the above uses. |
| 29 | + |
| 30 | +### DeviceBuffers |
| 31 | + |
| 32 | +A DeviceBuffer represents an **untyped, uninitialized device memory |
| 33 | +allocation**. DeviceBuffers can be created by providing the |
| 34 | +size of the allocation in bytes: |
| 35 | + |
| 36 | +```python |
| 37 | +>>> import rmm |
| 38 | +>>> buf = rmm.DeviceBuffer(size=100) |
| 39 | +``` |
| 40 | + |
| 41 | +The size of the allocation and the memory address associated with it |
| 42 | +can be accessed via the `.size` and `.ptr` attributes respectively: |
| 43 | + |
| 44 | +```python |
| 45 | +>>> buf.size |
| 46 | +100 |
| 47 | +>>> buf.ptr |
| 48 | +140202544726016 |
| 49 | +``` |
| 50 | + |
| 51 | +DeviceBuffers can also be created by copying data from host memory: |
| 52 | + |
| 53 | +```python |
| 54 | +>>> import rmm |
| 55 | +>>> import numpy as np |
| 56 | +>>> a = np.array([1, 2, 3], dtype='float64') |
| 57 | +>>> buf = rmm.to_device(a.tobytes()) |
| 58 | +>>> buf.size |
| 59 | +24 |
| 60 | +``` |
| 61 | + |
| 62 | +Conversely, the data underlying a DeviceBuffer can be copied to the |
| 63 | +host: |
| 64 | + |
| 65 | +```python |
| 66 | +>>> np.frombuffer(buf.tobytes()) |
| 67 | +array([1., 2., 3.]) |
| 68 | +``` |
| 69 | + |
| 70 | +### MemoryResource objects |
| 71 | + |
| 72 | +`MemoryResource` objects are used to configure how device memory allocations are made by |
| 73 | +RMM. |
| 74 | + |
| 75 | +By default if a `MemoryResource` is not set explicitly, RMM uses the `CudaMemoryResource`, which |
| 76 | +uses `cudaMalloc` for allocating device memory. |
| 77 | + |
| 78 | +`rmm.reinitialize()` provides an easy way to initialize RMM with specific memory resource options |
| 79 | +across multiple devices. See `help(rmm.reinitialize)` for full details. |
| 80 | + |
| 81 | +For lower-level control, the `rmm.mr.set_current_device_resource()` function can be |
| 82 | +used to set a different MemoryResource for the current CUDA device. For |
| 83 | +example, enabling the `ManagedMemoryResource` tells RMM to use |
| 84 | +`cudaMallocManaged` instead of `cudaMalloc` for allocating memory: |
| 85 | + |
| 86 | +```python |
| 87 | +>>> import rmm |
| 88 | +>>> rmm.mr.set_current_device_resource(rmm.mr.ManagedMemoryResource()) |
| 89 | +``` |
| 90 | + |
| 91 | +> :warning: The default resource must be set for any device **before** |
| 92 | +> allocating any device memory on that device. Setting or changing the |
| 93 | +> resource after device allocations have been made can lead to unexpected |
| 94 | +> behaviour or crashes. See [Multiple Devices](#multiple-devices) |
| 95 | +
|
| 96 | +As another example, `PoolMemoryResource` allows you to allocate a |
| 97 | +large "pool" of device memory up-front. Subsequent allocations will |
| 98 | +draw from this pool of already allocated memory. The example |
| 99 | +below shows how to construct a PoolMemoryResource with an initial size |
| 100 | +of 1 GiB and a maximum size of 4 GiB. The pool uses |
| 101 | +`CudaMemoryResource` as its underlying ("upstream") memory resource: |
| 102 | + |
| 103 | +```python |
| 104 | +>>> import rmm |
| 105 | +>>> pool = rmm.mr.PoolMemoryResource( |
| 106 | +... upstream=rmm.mr.CudaMemoryResource(), |
| 107 | +... initial_pool_size=2**30, |
| 108 | +... maximum_pool_size=2**32 |
| 109 | +... ) |
| 110 | +>>> rmm.mr.set_current_device_resource(pool) |
| 111 | +``` |
| 112 | + |
| 113 | +Similarly, to use a pool of managed memory: |
| 114 | + |
| 115 | +```python |
| 116 | +>>> import rmm |
| 117 | +>>> pool = rmm.mr.PoolMemoryResource( |
| 118 | +... upstream=rmm.mr.ManagedMemoryResource(), |
| 119 | +... initial_pool_size=2**30, |
| 120 | +... maximum_pool_size=2**32 |
| 121 | +... ) |
| 122 | +>>> rmm.mr.set_current_device_resource(pool) |
| 123 | +``` |
| 124 | + |
| 125 | +Other MemoryResources include: |
| 126 | + |
| 127 | +* `FixedSizeMemoryResource` for allocating fixed blocks of memory |
| 128 | +* `BinningMemoryResource` for allocating blocks within specified "bin" sizes from different memory |
| 129 | +resources |
| 130 | + |
| 131 | +MemoryResources are highly configurable and can be composed together in different ways. |
| 132 | +See `help(rmm.mr)` for more information. |
| 133 | + |
| 134 | +### Using RMM with CuPy |
| 135 | + |
| 136 | +You can configure [CuPy](https://cupy.dev/) to use RMM for memory |
| 137 | +allocations by setting the CuPy CUDA allocator to |
| 138 | +`rmm_cupy_allocator`: |
| 139 | + |
| 140 | +```python |
| 141 | +>>> import rmm |
| 142 | +>>> import cupy |
| 143 | +>>> cupy.cuda.set_allocator(rmm.rmm_cupy_allocator) |
| 144 | +``` |
| 145 | + |
| 146 | +### Using RMM with Numba |
| 147 | + |
| 148 | +You can configure Numba to use RMM for memory allocations using the |
| 149 | +Numba [EMM Plugin](http://numba.pydata.org/numba-doc/latest/cuda/external-memory.html#setting-the-emm-plugin). |
| 150 | + |
| 151 | +This can be done in two ways: |
| 152 | + |
| 153 | +1. Setting the environment variable `NUMBA_CUDA_MEMORY_MANAGER`: |
| 154 | + |
| 155 | + ```python |
| 156 | + $ NUMBA_CUDA_MEMORY_MANAGER=rmm python (args) |
| 157 | + ``` |
| 158 | + |
| 159 | +2. Using the `set_memory_manager()` function provided by Numba: |
| 160 | + |
| 161 | + ```python |
| 162 | + >>> from numba import cuda |
| 163 | + >>> import rmm |
| 164 | + >>> cuda.set_memory_manager(rmm.RMMNumbaManager) |
| 165 | + ``` |
0 commit comments