device_memory_resource.hpp

转到此文件的文档。

/*
 * 版权所有 (c) 2019-2025, NVIDIA CORPORATION。
 *
 * 根据 Apache 许可，版本 2.0（“许可”）获得许可；
 * 除非符合许可的规定，否则不得使用此文件。
 * 您可以在以下位置获得许可的副本：
 *
 * https://apache.ac.cn/licenses/LICENSE-2.0
 *
 * 除非适用法律要求或书面同意，否则软件
 * 在许可下分发，按“原样”提供，
 * 不附带任何明示或暗示的保证或条件。
 * 请参阅许可了解特定语言的管理权限和
 * 许可下的限制。
 */
#pragma once

#include <rmm/cuda_stream_view.hpp>
#include <rmm/detail/aligned.hpp>
#include <rmm/detail/cuda_memory_resource.hpp>
#include <rmm/detail/export.hpp>
#include <rmm/detail/nvtx/ranges.hpp>

#include <cstddef>

namespace RMM_NAMESPACE {
namespace mr {
/**
 * @brief 所有 librmm 设备内存分配的基类。
 *
 * device_memory_resource 概念支持异步内存分配（即
 * 不需要与 host 同步）。
 *
 * 与 C++ 标准库中的 `std::pmr::memory_resource` 类似，
 * `device_memory_resource` 是多态的，它定义了分配和释放
 * GPU 设备内存所需的接口。
 *
 * `device_memory_resource` 的派生类必须实现纯虚函数
 * `do_allocate` 和 `do_deallocate`。
 *
 * RMM 提供了几种 `device_memory_resource` 的实现，例如
 * `cuda_memory_resource` 和 `pool_memory_resource`。
 *
 * 由于所有的 `device_memory_resource` 都允许异步操作，因此它们
 * 都支持 `cuda::mr::device_accessible` 属性。这个属性表示资源
 * 可以从 CUDA 设备代码访问。
 *
 * 用户定义的内存资源只需要实现 `do_allocate` 和 `do_deallocate`
 * 即可与 RMM 一起使用。
 *
 * @ingroup memory_resources
 */
class device_memory_resource {
 public
 device_memory_resource() = default;
 virtual ~device_memory_resource() = default;
 /// @brief 默认拷贝构造函数。
 device_memory_resource(device_memory_resource const&) = default;
 /// @brief 默认移动构造函数。
 device_memory_resource(device_memory_resource&&) noexcept =
 default;
 /// @brief 默认拷贝赋值运算符。
 device_memory_resource& operator=(device_memory_resource const&) =
 default;
 /// @brief 默认移动赋值运算符。
 device_memory_resource& operator=(device_memory_resource&&) noexcept =
 default;

 /**
 * @brief 分配至少 `bytes` 大小的内存。
 *
 * 请求 `bytes` 大小的内存分配。返回的指针满足 `cuda::mr::device_accessible`
 * 属性。
 *
 * @throws `rmm::bad_alloc` 如果无法满足分配请求。
 *
 * @param bytes 请求分配的字节数。
 * @param stream 用于分配操作的 CUDA 流。
 * @return 指向新分配内存的指针。
 */
 void* allocate(std::size_t bytes, cuda_stream_view stream = cuda_stream_view{})
 {
 RMM_FUNC_RANGE();
 return do_allocate(bytes, stream);
 }

 /**
 * @brief 释放指针 `ptr` 指向的内存。
 *
 * 要求 `ptr` 是之前由同一个 `device_memory_resource` 的 `allocate` 方法返回的。
 * `bytes` 参数必须等于用于获取 `ptr` 的分配请求中的 `bytes` 参数。
 *
 * @param ptr 要释放的内存指针。
 * @param bytes 要释放的字节数。
 * @param stream 用于释放操作的 CUDA 流。
 */
 void deallocate(void* ptr, std::size_t bytes, cuda_stream_view stream = cuda_stream_view{})
 {
 RMM_FUNC_RANGE();
 do_deallocate(ptr, bytes, stream);
 }

 /**
 * @brief 将此资源与另一个进行比较。
 *
 * @param other 要比较的 device_memory_resource。
 * @return 如果两个资源等价，则为 true，否则为 false。
 */
 [[nodiscard]] bool is_equal(device_memory_resource const& other) const noexcept
 {
 return do_is_equal(other);
 }

 /**
 * @brief 分配至少 `bytes` 大小的内存，对齐到 `alignment`。
 *
 * @throws `rmm::bad_alloc` 如果无法满足分配请求。
 *
 * @param bytes 请求分配的字节数。
 * @param alignment 内存的对齐要求。