用户指南

在以 GPU 为中心的工作流程中实现最佳性能通常需要定制 GPU（“设备”）内存的分配方式。

RMM 是一个软件包，使您能够以高度可配置的方式分配设备内存。例如，它使您能够分配和使用 GPU 内存池，或使用托管内存进行分配。

您还可以轻松配置 Numba 和 CuPy 等其他库，使其使用 RMM 来分配设备内存。

安装

有关如何安装 RMM，请参阅项目README。

使用 RMM

在 Python 代码中有两种使用 RMM 的方式

使用rmm.DeviceBuffer API 显式创建和管理设备内存分配
通过 CuPy 和 Numba 等外部库透明使用

RMM 提供了一个 MemoryResource 抽象，用于控制在上述两种使用方式中如何分配设备内存。

`DeviceBuffer` 对象

一个 DeviceBuffer 表示一个无类型、未初始化的设备内存分配。DeviceBuffer 可以通过提供以字节为单位的分配大小来创建

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

分配的大小以及与其关联的内存地址可以通过 .size 和 .ptr 属性分别访问

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

DeviceBuffer 也可以通过从主机内存复制数据来创建

>>> import rmm
>>> import numpy as np
>>> a = np.array([1, 2, 3], dtype='float64')
>>> buf = rmm.DeviceBuffer.to_device(a.view("int8"))  # to_device expects an 8-bit type or `bytes`
>>> buf.size
24

相反，DeviceBuffer 底层的数据可以复制到主机

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

预取 `DeviceBuffer`

CUDA Unified Memory，也称为托管内存，可以使用 rmm.mr.ManagedMemoryResource 显式分配，或者通过调用 rmm.reinitialize 并设置 managed_memory=True 来分配。

由托管内存或其他可迁移内存（例如 HMM/ATS 内存）支持的 DeviceBuffer 可以预取到指定的设备，例如以减少或消除页面错误。

>>> import rmm
>>> rmm.reinitialize(managed_memory=True)
>>> buf = rmm.DeviceBuffer(size=100)
>>> buf.prefetch()

上面的示例将 DeviceBuffer 内存预取到 DeviceBuffer 最后使用的流上的当前 CUDA 设备（例如在构建时）。目标设备 ID 和流是可选参数。

>>> import rmm
>>> rmm.reinitialize(managed_memory=True)
>>> from rmm.pylibrmm.stream import Stream
>>> stream = Stream()
>>> buf = rmm.DeviceBuffer(size=100, stream=stream)
>>> buf.prefetch(device=3, stream=stream) # prefetch to device on stream.

如果 DeviceBuffer 不是由可迁移内存支持的，则 DeviceBuffer.prefetch() 是一个无操作。

`MemoryResource` 对象

MemoryResource 对象用于配置 RMM 如何进行设备内存分配。

默认情况下，如果没有显式设置 MemoryResource，RMM 会使用 CudaMemoryResource，它使用 cudaMalloc 来分配设备内存。

rmm.reinitialize() 提供了一种简单的方法，可以在多个设备上使用特定的内存资源选项初始化 RMM。有关详细信息，请参阅 help(rmm.reinitialize)。

对于更低级别的控制，可以使用 rmm.mr.set_current_device_resource() 函数为当前 CUDA 设备设置不同的 MemoryResource。例如，启用 ManagedMemoryResource 会告诉 RMM 使用 cudaMallocManaged 而不是 cudaMalloc 来分配内存

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

:warning: 必须在任何设备上分配任何设备内存之前为该设备设置默认资源。在进行设备分配后设置或更改资源可能导致意外行为或崩溃。

另一个例子是，PoolMemoryResource 允许您预先分配一个大的设备内存“池”。随后的分配将从这个已分配的内存池中提取。下面的示例展示了如何构造一个初始大小为 1 GiB、最大大小为 4 GiB 的 PoolMemoryResource。该池使用 CudaMemoryResource 作为其底层（“上游”）内存资源

>>> import rmm
>>> pool = rmm.mr.PoolMemoryResource(
...     rmm.mr.CudaMemoryResource(),
...     initial_pool_size="1GiB", # equivalent to initial_pool_size=2**30
...     maximum_pool_size="4GiB"
... )
>>> rmm.mr.set_current_device_resource(pool)

类似地，要使用托管内存池

>>> import rmm
>>> pool = rmm.mr.PoolMemoryResource(
...     rmm.mr.ManagedMemoryResource(),
...     initial_pool_size="1GiB",
...     maximum_pool_size="4GiB"
... )
>>> rmm.mr.set_current_device_resource(pool)

其他 MemoryResource 包括

FixedSizeMemoryResource 用于分配固定大小的内存块
BinningMemoryResource 用于从不同的内存资源中分配指定“bin”大小内的块

MemoryResource 具有高度可配置性，并且可以通过不同的方式组合使用。有关更多信息，请参阅 help(rmm.mr)。

将 RMM 与第三方库一起使用

许多库提供控制其设备分配的钩子。RMM 在 rmm.allocators 子模块中为 CuPy、numba 和 PyTorch 提供了这些实现的版本。所有这些方法都配置库使用当前的 RMM 内存资源进行设备分配。

将 RMM 与 CuPy 一起使用

您可以通过将 CuPy CUDA 分配器设置为 rmm.allocators.cupy.rmm_cupy_allocator 来配置 CuPy 使用 RMM 进行内存分配

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

将 RMM 与 Numba 一起使用

您可以使用 Numba 的 EMM 插件配置 Numba 使用 RMM 进行内存分配。

这可以通过两种方式完成

设置环境变量 NUMBA_CUDA_MEMORY_MANAGER

$ NUMBA_CUDA_MEMORY_MANAGER=rmm.allocators.numba python (args)

使用 Numba 提供的 set_memory_manager() 函数

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

将 RMM 与 PyTorch 一起使用

您可以通过配置当前分配器来配置 PyTorch 使用 RMM 进行内存分配。

>>> from rmm.allocators.torch import rmm_torch_allocator
>>> import torch

>>> torch.cuda.memory.change_current_allocator(rmm_torch_allocator)

内存统计和性能分析

RMM 可以通过使用以下任一方式对内存使用进行性能分析并跟踪内存统计信息

使用上下文管理器 rmm.statistics.statistics() 为特定的代码块启用统计信息跟踪。
调用 rmm.statistics.enable_statistics() 全局启用统计信息跟踪。

这两种用法都有一个共同点，即它们会修改当前活动的 RMM 内存资源。当前设备资源被封装在一个 StatisticsResourceAdaptor 中，该封装必须在整个统计信息跟踪过程中保持在最顶层。

>>> import rmm
>>> import rmm.statistics

>>> # We start with the default cuda memory resource
>>> rmm.mr.get_current_device_resource()
<rmm.pylibrmm.memory_resource.CudaMemoryResource object at 0x7fa0da48a8e0>

>>> # When using statistics, we get a StatisticsResourceAdaptor with the context
>>> with rmm.statistics.statistics():
...     rmm.mr.get_current_device_resource()
<rmm.pylibrmm.memory_resource.StatisticsResourceAdaptor object at 0x7fa0dd6e4a40>

>>> # We can also enable statistics globally
>>> rmm.statistics.enable_statistics()
>>> print(rmm.mr.get_current_device_resource())
<rmm.pylibrmm.memory_resource.StatisticsResourceAdaptor object at 0x7f9a11340a40>

启用统计信息后，您可以查询当前 RMM 内存资源执行的当前和峰值字节数以及分配次数的统计信息

>>> buf = rmm.DeviceBuffer(size=10)
>>> rmm.statistics.get_statistics()
Statistics(current_bytes=16, current_count=1, peak_bytes=16, peak_count=1, total_bytes=16, total_count=1)

内存性能分析器

要对特定的代码块进行性能分析，首先通过调用 rmm.statistics.enable_statistics() 启用内存统计信息。要对函数进行性能分析，请将 profiler 用作函数装饰器

>>> @rmm.statistics.profiler()
... def f(size):
...   rmm.DeviceBuffer(size=size)
>>> f(1000)

>>> # By default, the profiler write to rmm.statistics.default_profiler_records
>>> print(rmm.statistics.default_profiler_records.report())
Memory Profiling
================

Legends:
  ncalls       - number of times the function or code block was called
  memory_peak  - peak memory allocated in function or code block (in bytes)
  memory_total - total memory allocated in function or code block (in bytes)

Ordered by: memory_peak

ncalls     memory_peak    memory_total  filename:lineno(function)
     1           1,008           1,008  <ipython-input-11-5fc63161ac29>:1(f)

要对代码块进行性能分析，请将 profiler 用作上下文管理器

>>> with rmm.statistics.profiler(name="my code block"):
...     rmm.DeviceBuffer(size=20)
>>> print(rmm.statistics.default_profiler_records.report())
Memory Profiling
================

Legends:
  ncalls       - number of times the function or code block was called
  memory_peak  - peak memory allocated in function or code block (in bytes)
  memory_total - total memory allocated in function or code block (in bytes)

Ordered by: memory_peak

ncalls     memory_peak    memory_total  filename:lineno(function)
     1           1,008           1,008  <ipython-input-11-5fc63161ac29>:1(f)
     1              32              32  my code block

profiler 支持嵌套

>>> with rmm.statistics.profiler(name="outer"):
...     buf1 = rmm.DeviceBuffer(size=10)
...     with rmm.statistics.profiler(name="inner"):
...         buf2 = rmm.DeviceBuffer(size=10)
>>> print(rmm.statistics.default_profiler_records.report())
Memory Profiling
================

Legends:
  ncalls       - number of times the function or code block was called
  memory_peak  - peak memory allocated in function or code block (in bytes)
  memory_total - total memory allocated in function or code block (in bytes)

Ordered by: memory_peak

ncalls     memory_peak    memory_total  filename:lineno(function)
     1           1,008           1,008  <ipython-input-4-865fbe04e29f>:1(f)
     1              32              32  my code block
     1              32              32  outer
     1              16              16  inner