RMM: RAPIDS内存管理器

注意: 要查看最新的README.md，请确保您位于main分支。

资源

• RMM参考文档：Python API参考、教程和主题指南。
• librmm参考文档：C/C++ CUDA库API参考。
• 入门：安装RMM的说明。
• RAPIDS社区：获取帮助、贡献和协作。
• GitHub仓库：下载RMM源代码。
• 问题跟踪器：报告问题或请求功能。

概述

在以GPU为中心的工作流程中实现最佳性能通常需要自定义主机和设备内存的分配方式。例如，使用“固定”的主机内存进行异步主机 <-> 设备内存传输，或使用设备内存池子分配器来降低动态设备内存分配的成本。

RAPIDS内存管理器（RMM）的目标是提供：

• 一个通用接口，允许自定义设备和主机内存分配
• 接口的实现集合
• 使用该接口进行内存分配的数据结构集合

有关RMM提供的接口和使用RMM在C++代码中使用的详细信息，请参阅下文。

有关RAPIDS内存管理器设计的教程，请阅读NVIDIA开发者博客上的使用RAPIDS内存管理器为NVIDIA CUDA实现快速灵活的分配。

安装

Conda

您可以使用Conda（miniconda或完整的Anaconda发行版）从rapidsai频道安装RMM。

conda install -c rapidsai -c conda-forge -c nvidia rmm cuda-version=12.0

我们还提供了夜间Conda包，这些包是从我们最新开发分支的HEAD构建的。

注意：RMM仅支持Linux，并且仅在Python版本3.9、3.10和3.11上进行了测试。

注意：从Conda安装的RMM软件包需要使用GCC 9或更高版本进行构建。否则，您的应用程序可能无法构建。

有关更多操作系统和版本信息，请参阅RAPIDS版本选择器。

从源码构建

获取RMM依赖项

编译器要求

• gcc版本9.3+
• nvcc版本11.4+
• cmake版本3.26.4+

CUDA/GPU要求

• CUDA 11.4+. 您可以从https://developer.nvidia.com/cuda-downloads获取CUDA。

GPU支持

• RMM仅在Volta架构及更高版本（计算能力7.0+）上进行了测试和支持。它可能在较早的架构上也能工作。

Python要求

• rapids-build-backend（从PyPI或rapidsai conda频道获取）
• scikit-build-core
• cuda-python
• cython

有关更多详细信息，请参阅pyproject.toml

从源码构建RMM的脚本

要从源码安装RMM，请确保满足依赖项，并按照以下步骤操作：

• 克隆存储库和子模块

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

• 创建conda开发环境rmm_dev

# create the conda environment (assuming in base `rmm` directory)
$ conda env create --name rmm_dev --file conda/environments/all_cuda-118_arch-x86_64.yaml
# activate the environment
$ conda activate rmm_dev

• 使用cmake & make构建和安装librmm。CMake依赖于您的路径上的nvcc可执行文件或定义在CUDACXX环境变量中。

$ 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'

• 使用build.sh构建和安装librmm和rmm。build.sh在git存储库的根目录中创建构建目录。build.sh依赖于您的路径上的nvcc可执行文件或定义在CUDACXX环境变量中。

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

• 运行测试（可选）

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

• 在python文件夹中构建、安装和测试rmm Python包

# In the root rmm directory
$ python -m pip install -e ./python/rmm
$ pytest -v

完成！您现在可以为RMM OSS项目开发了。

缓存第三方依赖项

RMM使用CPM.cmake来处理类似spdlog、Thrust、GoogleTest、GoogleBenchmark的第三方依赖项。通常您不需要担心它。如果CMake在您的系统上找到合适的版本，它就会使用它（您可以通过设置CMAKE_PREFIX_PATH来指向安装位置来帮助它）。否则，这些依赖项将在构建过程中下载。

如果您经常从头开始新建，请考虑将环境变量CPM_SOURCE_CACHE设置为外部下载目录，以避免重复下载第三方依赖项。

在下游CMake项目中使用RMM

安装的RMM库提供了一组配置文件，使您能够轻松地将RMM集成到自己的CMake项目中。在您的CMakeLists.txt中，只需添加：

find_package(rmm [VERSION])
# ...
target_link_libraries(<your-target> (PRIVATE|PUBLIC|INTERFACE) rmm::rmm)

由于RMM是仅头文件的库，这实际上并没有链接RMM，但它使头文件可用并引入了传递依赖项。如果RMM未安装在任何默认位置，请使用CMAKE_PREFIX_PATH或rmm_ROOT来指向其位置。

RMM的依赖之一是Thrust库，因此上述操作通过依赖于rmm::Thrust目标自动引入了Thrust。默认情况下，它使用Thrust的标准配置。如果您想自定义它，可以设置变量THRUST_HOST_SYSTEM和THRUST_DEVICE_SYSTEM；请参阅Thrust的CMake文档。

使用CPM管理RMM

RMM使用CPM.cmake来管理其依赖项，包括CCCL，您也可以使用CPM来管理项目对RMM的依赖。

使用CPM的单参数紧凑语法为RMM/CCCL创建问题时，它会递归地将目标标记为SYSTEM依赖项。这导致通过CPM引入的CCCL头文件在预处理器中的优先级低于由CUDA SDK提供的（可能过时的）CCCL头文件。为了避免此问题，请使用CPM的多参数语法。

CPMAddPackage(NAME rmm [VERSION]
              GITHUB_REPOSITORY rapidsai/rmm
              SYSTEM Off)
# ...
target_link_libraries(<your-target> (PRIVATE|PUBLIC|INTERFACE) rmm::rmm)

在C++中使用RMM

RMM的第一个目标是提供一个设备内存和主机内存分配的通用接口。这允许用户和实现者自定义分配逻辑以使用单个接口进行编程。

为此，RMM定义了两个抽象接口类

• rmm::mr::device_memory_resource用于设备内存分配
• rmm::mr::host_memory_resource用于主机内存分配

这些类基于C++17中引入的用于多态内存分配的接口类std::pmr::memory_resource。

device_memory_resource

rmm::mr::device_memory_resource是定义分配和释放设备内存接口的基类。

它有两个关键函数

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

   • 返回至少bytes字节的分配的指针。

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

   • 回收由p指向的先前大小为bytes的分配。
   • p 必须是先前对allocate(bytes)的调用返回的，否则行为是未定义的

由派生类提供这些函数的实现。请参阅可用资源示例device_memory_resource派生类。

与std::pmr::memory_resource不同，rmm::mr::device_memory_resource不允许指定对齐参数。所有分配都需要至少对齐到256B。此外，device_memory_resource添加了一个额外的cuda_stream_view参数，允许指定执行（解）分配的流。

流顺序内存分配

rmm::mr::device_memory_resource是一个提供流顺序内存分配的基类。这允许像在不进行同步开销的情况下重用同一流上释放的内存这样的优化。

对device_memory_resource::allocate(bytes, stream_a)的调用返回一个在stream_a上有效的指针。在除stream_a以外的流上使用内存是未定义行为，除非首先同步这两个流，例如使用cudaStreamSynchronize(stream_a)或通过在stream_a上记录CUDA事件然后调用cudaStreamWaitEvent(stream_b, event)。

指定给device_memory_resource::deallocate的流应该是可以立即用于另一个分配的流。通常这是在调用deallocate之前最后使用分配的流。传递的流可能由device_memory_resource内部使用来以最小的同步管理可用内存，并且它可能在稍后同步，例如使用cudaStreamSynchronize()调用。

因此，销毁传递给 device_memory_resource::deallocate 的 CUDA 流是未定义行为。如果分配最后一次使用所在的流在调用 deallocate 之前已被销毁，或者已知它将被销毁，那么在销毁之前同步流（例如，使用默认流）然后传递不同的流给 deallocate 可能会更好。

请注意，如 rmm::device_buffer 和 rmm::device_uvector 这样的设备内存数据结构遵循这些流顺序内存分配的语义和规则。

有关流顺序内存分配语义的更多信息，请阅读 NVIDIA 开发者博客上的 使用 NVIDIA CUDA 流顺序内存分配器。

可用设备资源

RMM 提供了几个从 device_memory_resource 派生的类，以满足各种用户需求。有关这些资源的更详细信息，请参阅它们的相应文档。

cuda_memory_resource

使用 cudaMalloc 和 cudaFree 分配和释放设备内存。

managed_memory_resource

使用 cudaMallocManaged 和 cudaFree 分配和释放设备内存。

请注意，managed_memory_resource 不能与 NVIDIA 虚拟 GPU 软件（vGPU，用于虚拟机或虚拟机管理程序）一起使用，因为 NVIDIA CUDA 统一内存不支持 NVIDIA vGPU。

pool_memory_resource

一个合并、最佳匹配池子子分配器。

fixed_size_memory_resource

只能分配单个固定大小的内存资源。平均分配和释放成本是恒定的。

binning_memory_resource

可根据不同的分箱大小配置，用于分配。通常配置多个分箱，由 fixed_size_memory_resource 支持，并为大于最大分箱大小的分配配置一个 pool_memory_resource。

默认资源和设备资源

RMM 用户通常需要配置一个用于所有分配的 device_memory_resource 对象，其中未明确提供其他资源。一个常见的例子是将 pool_memory_resource 配置为用于所有分配以获得快速动态分配。

为了启用此用例，RMM 提供了“默认” device_memory_resource 的概念。当没有明确提供其他资源时，将使用此资源。

访问和修改默认资源是通过两个函数完成的

• device_memory_resource* get_current_device_resource()

  • 返回当前 CUDA 设备的默认资源的指针。
  • 初始默认内存资源是 cuda_memory_resource 的一个实例。
  • 此函数在并发调用它和 set_current_device_resource() 方面是线程安全的。
  • 为了更明确的控制，您可以使用 get_per_device_resource()，它接受一个设备 ID。

• device_memory_resource* set_current_device_resource(device_memory_resource* new_mr)

  • 更新当前 CUDA 设备的默认内存资源指针为 new_mr
  • 返回上一个默认资源指针
  • 如果 new_mr 是 nullptr，则将默认资源重置为 cuda_memory_resource
  • 此函数在并发调用它和 get_current_device_resource() 方面是线程安全的。
  • 为了更明确的控制，您可以使用 set_per_device_resource()，它接受一个设备 ID。

示例

rmm::mr::cuda_memory_resource cuda_mr;
// Construct a resource that uses a coalescing best-fit pool allocator
// With the pool initially half of available device memory
auto initial_size = rmm::percent_of_free_device_memory(50);
rmm::mr::pool_memory_resource<rmm::mr::cuda_memory_resource> pool_mr{&cuda_mr, initial_size};
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`

多个设备

当活动的 CUDA 设备与创建 device_memory_resource 时活动的设备相同，才应使用 device_memory_resource。否则，行为是未定义的。

如果使用与创建内存资源的设备不同的 CUDA 设备的流使用 device_memory_resource，行为是未定义的。

为每个设备创建一个 device_memory_resource 需要小心设置当前设备，在创建每个资源之前，并保持资源的生命周期，只要它们被设置为按设备资源。以下是一个示例循环，它为每个设备创建指向 pool_memory_resource 对象的 unique_ptr，并将它们设置为该设备的按设备资源。

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); // set device i before creating MR
  // Use a vector of unique_ptr to maintain the lifetime of the MRs
  // 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 the per-device resource for device i
  set_per_device_resource(cuda_device_id{i}, &per_device_pools.back());
}

请注意，在创建 device_memory_resource 时当前活动的 CUDA 设备，在 device_memory_resource 用来释放内存的任何时候（包括在析构函数中）也必须处于当前活动状态。RAII 类 rmm::device_buffer 和使用它作为后端存储的类（如 rmm::device_scalar 和 rmm::device_uvector）通过在构造函数调用时存储活动设备，并确保在执行分配或释放操作（包括在析构函数中）时存储的设备处于活动状态来处理这一点。因此，用户必须确保在 rmm::device_buffer 的 创建 期间活动的设备与正在使用的内存资源的活动设备相匹配。

以下是一个不正确的示例，它在设备零上创建内存资源，然后使用它来在设备一上分配一个 device_buffer

{
  RMM_CUDA_TRY(cudaSetDevice(0));
  auto mr = rmm::mr::cuda_memory_resource{};
  {
    RMM_CUDA_TRY(cudaSetDevice(1));
    // Invalid, current device is 1, but MR is only valid for device 0
    rmm::device_buffer buf(16, rmm::cuda_stream_default, &mr);
  }
}

正确示例是当设备零处于活动状态时创建设备缓冲区。之后，可以安全地切换设备，并允许缓冲区在活动设备不同的状态下超出范围并析构。例如，以下代码是正确的

{
  RMM_CUDA_TRY(cudaSetDevice(0));
  auto mr = rmm::mr::cuda_memory_resource{};
  rmm::device_buffer buf(16, rmm::cuda_stream_default, &mr);
  RMM_CUDA_TRY(cudaSetDevice(1));
  ...
  // No need to switch back to device 0 before ~buf runs
}

使用具有多个设备的 rmm::device_vector

rmm:device_vector 使用 rmm::mr::thrust_allocator 来使 thrust::device_vector 能够使用 RMM 来分配和释放内存。因此，适用于后端内存资源使用的常规规则适用：活动设备必须与资源构造时的活动设备相匹配。为了便于在 RAII 设置中使用，rmm::mr::thrust_allocator 在构造时记录活动设备，并确保在分配或释放内存时该设备处于活动状态。因此，使用具有多个设备的 rmm::device_vector 与 rmm::device_buffer 相同。必须使用正确的设备 创建 device_vector，但可以安全地使用不同的活动设备来销毁它们。

例如，使用 rmm::device_vector 重述前面的示例

{
  RMM_CUDA_TRY(cudaSetDevice(0));
  auto mr = rmm::mr::cuda_memory_resource{};
  rmm::device_vector<int> vec(16, rmm::mr::thrust_allocator<int>(rmm::cuda_stream_default, &mr));
  RMM_CUDA_TRY(cudaSetDevice(1));
  ...
  // No need to switch back to device 0 before ~vec runs
}

[!NOTE] 虽然在 thrust_allocator 中的分配和释放操作以正确的活动设备运行，但修改 rmm::device_vector 可能需要启动内核，并且这必须在正确的设备处于活动状态时运行。例如，.resize() 可能会分配 并 启动内核以初始化新元素：用户必须安排在内存资源的正确设备处于活动状态时发生此内核启动。

cuda_stream_view 和 cuda_stream

rmm::cuda_stream_view 是一个简单的非拥有包装器，包装了 CUDA 的 cudaStream_t。此包装器的目的是提供流类型的强类型安全性。（cudaStream_t 是一个指针的别名，当它被赋值为 0 时可能会在 API 中引起歧义。）所有 RMM 流顺序 API 都接受一个 rmm::cuda_stream_view 参数。

rmm::cuda_stream 是一个简单的拥有包装器，包装了 CUDA 的 cudaStream_t。此类提供 RAII 语义（构造函数创建 CUDA 流，析构函数销毁它）。rmm::cuda_stream 从不能表示 CUDA 默认流或线程默认流；它只代表一个非默认流。不能复制 rmm::cuda_stream，但可以移动它。

cuda_stream_pool

rmm::cuda_stream_pool 提供了对 CUDA 流池的快速访问。此类可用于创建一组生命周期与 cuda_stream_pool 相等的 cuda_stream 对象。使用流池可能比动态创建流更快。池的大小可配置。根据此大小，多次调用 cuda_stream_pool::get_stream() 可能会返回代表相同 CUDA 流的 rmm::cuda_stream_view 实例。

线程安全

除非另有说明，所有当前的设备内存资源都是线程安全的。具体来说，对内存资源 allocate() 和 deallocate() 方法的调用是安全的，无论是从其他线程调用这些函数中的任何一个。它们在内存资源对象的构建和销毁方面不是线程安全的。

请注意，提供了一个名为 thread_safe_resource_adapter 的类，可以用于将非线程安全的内存资源转换为线程安全（如上所述）。在使用任何当前 RMM 设备内存资源时，不需要此适配器。

分配器

C++ 接口通常允许通过一个 Allocator 对象来定制内存分配。RMM 提供了多个 Allocator 和 Allocator-like 类。

polymorphic_allocator

类似于 std::pmr::polymorphic_allocator 的流顺序分配器。与标准 C++ Allocator 接口不同，allocate 和 deallocate 函数接受一个 cuda_stream_view，指示分配（或释放）发生的流。

stream_allocator_adaptor

stream_allocator_adaptor 可用于将流顺序分配器适配为向可能未设计为与流顺序接口一起工作的消费者提供标准 Allocator 接口。

示例

rmm::cuda_stream stream;
rmm::mr::polymorphic_allocator<int> stream_alloc;

// Constructs an adaptor that forwards all (de)allocations to `stream_alloc` on `stream`.
auto adapted = rmm::mr::make_stream_allocator_adaptor(stream_alloc, stream);

// Allocates 100 bytes using `stream_alloc` on `stream`
auto p = adapted.allocate(100);
...
// Deallocates using `stream_alloc` on `stream`
adapted.deallocate(p,100);

thrust_allocator

thrust_allocator 是一个使用强类型 thrust::device_ptr 的设备内存分配器，使其可用于像 thrust::device_vector 这样的容器。

有关使用 RMM 与 Thrust 一起使用的更多信息，请参阅以下内容。

设备数据结构

device_buffer

一个无类型、未初始化的 RAII 类，用于流顺序设备内存分配。

示例

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

一个有类型、未初始化的 RAII 类，用于在设备内存中分配连续元素集。类似于 thrust::device_vector，但作为一个优化，不默认初始化包含的元素。此优化将类型 T 限制为可 trivial 复制的类型。

示例

cuda_stream_view s{...};
// Allocates uninitialized storage for 100 `int32_t` elements on stream `s` using the
// default resource
rmm::device_uvector<int32_t> 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<int32_t> v2{100, s, mr};

device_scalar

一个有类型、RAII 类，用于在设备内存中分配单个元素。这类似于具有单个元素的 device_uvector，但提供了从主机修改设备内存中的值或从设备检索值的便利函数。

示例

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

kernel<<<...,s.value()>>>(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 是定义分配和释放主机内存接口的基类。

与 device_memory_resource 类似，它有两个关键的（释放）分配函数

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

   • 返回一个指向至少 bytes 字节的分配，该分配按指定的 alignment 对齐

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

   • 回收由p指向的先前大小为bytes的分配。

与 device_memory_resource 不同，host_memory_resource 接口和行为与 std::pmr::memory_resource 相同。

可用主机资源

new_delete_resource

使用全局 operator new 和 operator delete 分配主机内存。

pinned_memory_resource

使用 cuda(Malloc/Free)Host 分配“固定”的主机内存。

主机数据结构

RMM 目前不提供任何与 host_memory_resource 接口的数据结构。未来，RMM 将提供类似于 device_buffer 的类似主机端结构以及可以与 STL 容器一起使用的分配器。

使用带推力的RMM

RAPIDS和其他CUDA库大量使用推力（Thrust）。推力在两种情况下使用CUDA设备内存：

1. 作为thrust::device_vector的后备存储，
2. 以及作为某些算法（如thrust::sort）中的临时存储。

RMM提供了rmm::mr::thrust_allocator作为符合规范的推力分配器，它使用device_memory_resource。

推力算法

要指示推力算法使用rmm::mr::thrust_allocator来分配临时存储，您可以使用自定义的推力CUDA设备执行策略：rmm::exec_policy(stream)。

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

第一个stream参数是要用于rmm::mr::thrust_allocator的流。第二个stream参数是用来执行推力算法的。这两个参数必须相同。

日志记录

RMM包含两种日志记录形式：内存事件日志和调试日志。

内存事件日志和logging_resource_adaptor

内存事件日志将每次分配或释放的详细信息写入CSV（逗号分隔值）文件。在C++中，通过将logging_resource_adaptor用作任何其他device_memory_resource对象的包装器来启用内存事件日志。

日志中的每一行代表一个分配或释放操作。文件中的列是“线程、时间、操作、指针、大小、流”。

logging_resource_adaptor的CSV输出文件可以用作REPLAY_BENCHMARK的输入，该工具在从源代码构建RMM时可用，位于构建目录中的gbenchmarks文件夹内。此日志重放器对于分析和调试分配器问题非常有用。

以下C++示例创建了一个将日志输出到“logs/test1.csv”文件的日志版本的cuda_memory_resource。

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

如果没有指定文件名，将查询环境变量RMM_LOG_FILE以获取文件名。如果未设置RMM_LOG_FILE，则logging_resource_adaptor构造函数将抛出异常。

在Python中，当将rmm.reinitialize()的logging参数设置为True时启用内存事件日志。可以使用log_file_name参数设置日志文件名。有关详细信息，请参阅help(rmm.reinitialize)。

调试日志

RMM包含一个调试记录器，可以启用以将跟踪和调试信息记录到文件。这些信息可以显示错误发生的时间、从上游资源分配额外内存的时间等。默认日志文件是当前工作目录中的rmm_log.txt，但可以将环境变量RMM_DEBUG_LOG_FILE设置为指定路径和文件名。

有一个CMake配置变量RMM_LOGGING_LEVEL，可以设置为启用更详细的日志记录的编译。默认值为INFO。可用的级别是TRACE、DEBUG、INFO、WARN、ERROR、CRITICAL和OFF。

日志依赖于spdlog库。

请注意，要查看低于INFO级别的日志，应用程序还必须在运行时设置日志级别。C++应用程序必须调用rmm::logger().set_level()，例如，要启用所有级别的日志记录直到TRACE，请调用rmm::logger().set_level(spdlog::level::trace)（并使用-DRMM_LOGGING_LEVEL=TRACE编译librmm）。Python应用程序必须调用rmm.set_logging_level()，例如，要启用所有级别的日志记录直到TRACE，请调用rmm.set_logging_level("trace")（并使用-DRMM_LOGGING_LEVEL=TRACE编译RMM Python模块）。

请注意，调试日志与由rmm::mr::logging_resource_adapter提供的CSV内存分配日志不同。后者用于记录分配/释放操作的历史，这对于使用RMM的重放基准进行重放非常有用。

资源管理器（RMM）和CUDA内存边界检查

从分配内存池的资源（如 pool_memory_resource 和 arena_memory_resource）中获取的内存分配属于同一低级CUDA内存分配。因此，对这些分配超出范围或不正确对齐的访问不太可能被CUDA工具（如 CUDA 计算安全器 memcheck）检测到。

此例外的是 cuda_memory_resource，它包装了 cudaMalloc，以及 cuda_async_memory_resource，它使用 cudaMallocAsync 并启用CUDA内置的内存池功能（CUDA 11.2或更高版本）。这些资源分配的非法内存访问可以通过Compute Sanitizer Memcheck检测到。

将来可能会添加对其他内存资源使用NVTX API进行内存边界检查的支持。

在Python中使用RMM

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

1. 通过使用 rmm.DeviceBuffer API显式创建和管理设备内存分配
2. 通过外部库如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.tobytes())
>>> buf.size
24

相反，DeviceBuffer下载数据可以复制到主机

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

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=2**30,
...     maximum_pool_size=2**32
... )
>>> rmm.mr.set_current_device_resource(pool)

其他MemoryResource包括

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

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

与第三方库一起使用RMM

与CuPy一起使用RMM

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

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

注意： 这仅配置CuPy使用当前RMM资源进行分配。它不会初始化或更改当前资源，例如启用内存池。有关更改当前内存资源的更多信息，请参阅此处。

在Numba中使用RMM

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

这可以通过两种方式实现

1. 设置环境变量NUMBA_CUDA_MEMORY_MANAGER

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

2. 使用Numba提供的set_memory_manager()函数

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

注意： 这仅配置Numba使用当前RMM资源进行分配。它不会初始化或更改当前资源，例如启用内存池。有关更改当前内存资源的更多信息，请参阅此处。

在PyTorch中使用RMM

PyTorch可以使用RMM进行内存分配。例如，要将PyTorch配置为使用RMM管理的池

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

rmm.reinitialize(pool_allocator=True)
torch.cuda.memory.change_current_allocator(rmm_torch_allocator)

PyTorch和RMM现在将共享同一个内存池。

当然，您也可以使用PyTorch自定义内存资源

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

# note that you can configure PyTorch to use RMM either before or
# after changing RMM's memory resource.  PyTorch will use whatever
# memory resource is configured to be the "current" memory resource at
# the time of allocation.
torch.cuda.change_current_allocator(rmm_torch_allocator)

# configure RMM to use a managed memory resource, wrapped with a
# statistics resource adaptor that can report information about the
# amount of memory allocated:
mr = rmm.mr.StatisticsResourceAdaptor(rmm.mr.ManagedMemoryResource())
rmm.mr.set_current_device_resource(mr)

x = torch.tensor([1, 2]).cuda()

# the memory resource reports information about PyTorch allocations:
mr.allocation_counts
Out[6]:
{'current_bytes': 16,
 'current_count': 1,
 'peak_bytes': 16,
 'peak_count': 1,
 'total_bytes': 16,
 'total_count': 1}

从Python获取C++对象的所有权。

当从Python与使用RMM的C++库交互时，在Python端获取rmm::device_buffer对象的所有权时必须小心。`rmm::device_buffer`不包含用于其分配的内存资源的所有者引用（仅包含`device_async_resource_ref`），并且期望分配用户至少保持此内存资源存活，直到缓冲区的生命周期。当在Python中获取此类缓冲区的所有权时，我们（在一般情况下）没有方法确保内存资源将比我们现在持有的缓冲区存活时间更长。

为了避免任何问题，我们需要两件事

1. 我们与之交互的C++库应接受用于返回给用户的分配的内存资源。
2. 当我们从Python调用库时，我们应该提供一个我们控制的内存资源。然后，在获取任何分配的rmm::device_buffer的所有权时，应提供此内存资源。

例如，假设我们有一个分配device_buffer的C++函数，该函数有一个默认为当前设备资源的实用重载

std::unique_ptr<rmm::device_buffer> allocate(
  std::size_t size,
  rmm::mr::device_async_resource_ref mr = get_current_device_resource())
{
    return std::make_unique<rmm::device_buffer>(size, rmm::cuda_stream_default, mr);
}

Python DeviceBuffer类有一个方便的Cython函数，c_from_unique_ptr，可以从`unique_ptr<rmm::device_buffer>`构造DeviceBuffer，并获取其所有权。为了安全地执行此操作，我们必须确保C++端执行的分配使用我们控制的内存资源。因此

# Bad, doesn't control lifetime
buffer_bad = DeviceBuffer.c_from_unique_ptr(allocate(10))

# Good, allocation happens with a memory resource we control
# mr is a DeviceMemoryResource
buffer_good = DeviceBuffer.c_from_unique_ptr(
    allocate(10, mr.get_mr()),
    mr=mr,
)

注意好坏两种情况的两个区别

1. 在好情况下，我们向分配函数传递内存资源。
2. 在好情况下，我们传递相同的内存资源到DeviceBuffer构造函数，以便其生命周期与缓冲区绑定。

依赖get_current_device_resource的潜在风险

在C++和Python API中执行分配的函数通常将内存资源参数默认为get_current_device_resource的值。这是为了简化调用者的接口。当从Python使用C++库时，这种默认值是安全的，只要只有Python进程调用set_current_device_resource。

这是因为在C++侧当前设备资源有预期由用户管理的生命周期。由rmm::mr::set_current_device_resource设置的资源存储在一个静态的std::map中，其键是设备ID，值是内存资源的原始指针。因此，rmm::mr::get_current_device_resource返回一个没有生命周期来源的对象。这就是为什么它不能从Python中使用。为了在Python侧处理这个问题，Python级别的set_current_device_resource会将C++资源和将Python对象存储在静态的全局字典中。然后，Python的get_current_device_resource将不使用rmm::mr::get_current_device_resource，而是在这个全局字典中查找当前设备资源。

因此，如果我们交互的C++库调用rmm::mr::set_current_device_resource，C++和Python侧的程序可能会对get_current_device_resource返回的内容产生分歧。因此，如果使用简化接口，唯一安全的事情是确保只在Python侧调用set_current_device_resource。