RMM: RAPIDS内存管理器

注意：要确保您在main分支上，请阅读最新的README.md。

资源

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

概述

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

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

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

有关RMM提供的接口以及如何在C++代码中使用RMM的信息，请参阅以下内容。

有关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

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

注意：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_b）上使用该内存是未定义行为，除非这两个流首先同步，例如通过使用 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`

多个设备

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

如果使用与创建内存资源时设备不同的设备关联的流来使用 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
}

[!注意]尽管 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 的类。

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 限制为可以 trivially copyable 的类型。

示例

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 容器一起使用的分配器。

与 Thrust 一起使用 RMM

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

1. 作为 thrust::device_vector 的后端存储，
2. 作为某些算法（如 thrust::sort）内部的临时存储。

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

Thrust 算法

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

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

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

日志记录

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

内存事件日志和 logging_resource_adaptor

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

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

logging_resource_adaptor 的 CSV 输出文件可以用作 REPLAY_BENCHMARK 的输入，该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)（并在编译 librmm 时使用 -DRMM_LOGGING_LEVEL=TRACE）。Python 应用程序必须调用 rmm.set_logging_level()，例如，要启用所有级别的日志（包括 TRACE），请调用 rmm.set_logging_level("trace")（并在编译 RMM Python 模块时使用 -DRMM_LOGGING_LEVEL=TRACE）。

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

RMM 和 CUDA 内存边界检查

从分配内存池的内存资源（如 pool_memory_resource 和 arena_memory_resource）中取出的内存分配是同一低级 CUDA 内存分配的一部分。因此，对这些分配的越界或未对齐访问不太可能被 CUDA 工具（如 CUDA Compute Sanitizer 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: 在任何设备上分配任何设备内存之前，必须为该设备设置默认资源。在设备分配之后设置或更改资源可能导致意外的行为或崩溃。有关详细信息，请参阅 Multiple Devices

作为另一个示例，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)

其他 MemoryResources 包括

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

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

使用 RMM 与第三方库

使用 RMM 与 CuPy

您可以通过将 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 资源进行分配。它不会初始化或更改当前资源，例如启用内存池。有关更改当前内存资源的信息，请参阅此处。

使用 RMM 与 Numba

您可以使用 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 资源进行分配。它不会初始化或更改当前资源，例如启用内存池。有关更改当前内存资源的信息，请参阅此处。

使用 RMM 与 PyTorch

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。