system_memory_resource.hpp

转到此文件的文档。

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

#include <rmm/cuda_device.hpp>
#include <rmm/cuda_stream_view.hpp>
#include <rmm/detail/error.hpp>
#include <rmm/detail/export.hpp>
#include <rmm/detail/format.hpp>
#include <rmm/mr/device/device_memory_resource.hpp>

#include <cstddef>
#include <string>

namespace RMM_NAMESPACE {
namespace mr {

namespace detail {
// 检查指定设备上是否支持系统分配内存 (SAM)。static bool is_system_memory_supported(cuda_device_id device_id)
{
 int pageableMemoryAccess;
 RMM_CUDA_TRY(cudaDeviceGetAttribute(
 &pageableMemoryAccess, cudaDevAttrPageableMemoryAccess, device_id.value()));
 return pageableMemoryAccess == 1;
}
} // namespace detail

// system_memory_resource 是 device_memory_resource 的派生类，使用 malloc/free 进行分配/释放。class system_memory_resource final : public device_memory_resource {
 public
 system_memory_resource()
 {
 RMM_EXPECTS(rmm::mr::detail::is_system_memory_supported(rmm::get_current_cuda_device()),
 "此硬件/软件版本不支持系统内存分配器。");
 }
 ~system_memory_resource() override = default;
 // 默认拷贝构造函数。system_memory_resource(system_memory_resource const&) = default;
 system_memory_resource(system_memory_resource&&) = default;
 // 默认拷贝赋值运算符。system_memory_resource& operator=(system_memory_resource const&) =
 default;
 // 默认移动赋值运算符。system_memory_resource& operator=(system_memory_resource&&) =
 default;

 private
 // 使用 `malloc` 分配内存并返回指向已分配内存的指针。void* do_allocate(std::size_t bytes, [[maybe_unused]] cuda_stream_view stream) override
 {
 try {
 return rmm::detail::aligned_host_allocate(
 bytes, CUDA_ALLOCATION_ALIGNMENT, [](std::size_t size) { return ::operator new(size); });
 } catch (std::bad_alloc const& e) {
 auto const msg = std::string("分配失败 ") + rmm::detail::format_bytes(bytes) +
 std::string("内存：") + e.what();
 RMM_FAIL(msg.c_str(), rmm::out_of_memory);
 }
 }

 // 使用 `free` 释放由 `do_allocate` 分配的内存。void do_deallocate(void* ptr,
 [[maybe_unused]] std::size_t bytes,
 cuda_stream_view stream) override
 {
 // 使用 cudaFree 时，CUDA 运行时会跟踪相关的操作并进行隐式同步。然而，对于 SAM，由于 free 是即时的，我们需要等待进行中的
 // CUDA 操作完成后再释放内存，以避免潜在的 use-after-free 错误或竞态条件。
 //
 // In `cudaFree`, the CUDA runtime keeps track of dependent operations and does implicit synchronization. However, with SAM, since `free` is immediate, we need to wait for in-flight CUDA operations to finish before freeing the memory, to avoid potential use-after-free errors or race conditions.stream.synchronize();


 rmm::detail::aligned_host_deallocate(
 ptr, bytes, CUDA_ALLOCATION_ALIGNMENT, [](void* ptr) { ::operator delete(ptr); });
 }

 // 比较是否相等。system_memory_resource 对象仅当它们是同一类型时才相等。[[nodiscard]] bool do_is_equal(device_memory_resource const& other) const noexcept override
 {
 return dynamic_cast<system_memory_resource const*>(&other) != nullptr;
 }
 // 启用 cuda::mr::device_accessible 属性。friend void get_property(system_memory_resource const&, cuda::mr::device_accessible) noexcept {}

 // 启用 cuda::mr::host_accessible 属性。friend void get_property(system_memory_resource const&, cuda::mr::host_accessible) noexcept {}
};

// 静态属性检查
static_assert(cuda::mr::async_resource_with<system_memory_resource, cuda::mr::device_accessible>);
static_assert(cuda::mr::async_resource_with<system_memory_resource, cuda::mr::host_accessible>); // end of group
} // namespace mr
} // namespace RMM_NAMESPACE