rmm - RAPIDS内存管理器

导航

验证详情

这些详情已由PyPI 验证
维护者
来自gravatar.com的nvidia头像 nvidia

未验证详情

这些详情尚未由PyPI 验证
项目链接
元数据
  • 许可证: Apache软件许可证 (Apache 2.0)
  • 作者: NVIDIA公司
分类

项目描述

 RMM: RAPIDS内存管理器

注意: 要获取最新的稳定版README.md，请确保您处于main分支。

资源

概述

在以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要求

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，并且您可以为您项目的RMM依赖项使用CPM。

使用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定义了两个抽象接口类

这些类基于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 之前已被销毁，或者已知它将被销毁，那么在销毁之前同步流（例如使用 cudaStreamSynchronize()）然后传递不同的流给 deallocate（例如默认流）可能是更好的选择。

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

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

可用设备资源

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

cuda_memory_resource

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

managed_memory_resource

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

请注意，由于NVIDIA CUDA统一内存不支持NVIDIA vGPU，因此managed_memory_resource不能与NVIDIA虚拟GPU软件（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
}

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

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

示例

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 与 Thrust

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 参数是用于执行 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的重放Benchmark一起重放。

RMM和CUDA内存边界检查

从分配内存池的内存资源（如pool_memory_resource和arena_memory_resource）分配的内存是同一底层CUDA内存分配的一部分。因此，对这些分配的越界或未对齐访问不太可能被CUDA工具（如CUDA Compute Sanitizer memcheck）检测到。

例外的是cuda_memory_resource，它包装了cudaMalloc，以及使用CUDA内置内存池功能（CUDA 11.2或更高版本）的cuda_async_memory_resource。使用这些资源分配的内存的非法内存访问可以使用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)

其他MemoryResources包括

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

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)
  1. 使用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++对象的控制权。

在与使用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返回的内容产生分歧。因此，在使用简化的接口时，唯一安全的事情是确保set_current_device_resource仅在Python端调用。

项目详情

验证详情

这些详情已由PyPI 验证
维护者
来自gravatar.com的nvidia头像 nvidia

未验证详情

这些详情尚未由PyPI 验证
项目链接
元数据
  • 许可证: Apache软件许可证 (Apache 2.0)
  • 作者: NVIDIA公司
分类

发布历史 发布通知 | RSS源

此版本

24.8.2

24.6.0

下载文件

下载适合您平台的文件。如果您不确定选择哪个，请了解更多关于安装包的信息。

源分布

librmm_cu11-24.8.2.tar.gz (13.8 kB 查看散列)

上传时间 源

支持者