rmm - RAPIDS内存管理器

导航

验证详情

这些详情已由PyPI验证
维护者
Avatar for nvidia from gravatar.com nvidia

未验证详情

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

项目描述

 RMM：RAPIDS内存管理器

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

资源

概述

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

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

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

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

安装

Conda

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

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依赖于您的路径上或定义在CUDACXX环境变量中的nvcc可执行文件。
$ 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依赖于您的路径上或定义在CUDACXX环境变量中的nvcc可执行文件。
$ ./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 定义了两个抽象接口类

这些类基于 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 支持的 bin 和一个用于大于最大 bin 大小分配的单个 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 需要仔细设置当前设备，在创建每个资源之前，并保持资源的生命周期，只要它们作为按设备资源设置。以下是一个示例循环，它为每个设备创建 unique_ptr 到 pool_memory_resource 对象，并将它们设置为该设备的按设备资源。

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，表示 (de) 分配发生的流。

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，它有两个关键的（de）分配函数

  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。第二个 stream 参数是用来执行 Thrust 算法的。这两个参数必须相同。

日志记录

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

内存事件日志和 logging_resource_adaptor

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

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

logging_resource_adaptor 的 CSV 输出文件可以用作 REPLAY_BENCHMARK 的输入，这是从源代码构建 RMM 时在构建目录中的 gbenchmarks 文件夹中可用的。这个日志重放器可以用于分析和调试分配器问题。

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

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 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:必须在任何设备上分配任何设备内存之前设置默认资源。在设备分配之后设置或更改资源可能会导致意外的行为或崩溃。请参阅多个设备

作为另一个示例，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”大小的块

MemoryResource具有高度的可配置性，可以以不同的方式组合在一起。请参阅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++ 对象的所有权。

当从 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。

项目详情

验证详情

这些详情已由PyPI验证
维护者
Avatar for nvidia from gravatar.com nvidia

未验证详情

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

发布历史 发布通知 | RSS 源

此版本

24.8.2

24.6.0

下载文件

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

源分布

librmm_cu12-24.8.2.tar.gz (13.8 kB 查看哈希值)

上传时间 源

由以下支持

AWS AWS 云计算和安全赞助商 Datadog Datadog 监控 Fastly Fastly CDN Google Google 下载分析 Microsoft Microsoft PSF 赞助商 Pingdom Pingdom 监控 Sentry Sentry 错误记录 StatusPage StatusPage 状态页