thrust_allocator_adaptor.hpp

转到此文件的文档。

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

#pragma once

#include <rmm/cuda_device.hpp>
#include <rmm/detail/cuda_memory_resource.hpp>
#include <rmm/detail/export.hpp>
#include <rmm/detail/thrust_namespace.h>
#include <rmm/mr/device/per_device_resource.hpp>
#include <rmm/resource_ref.hpp>

#include <thrust/device_malloc_allocator.h>
#include <thrust/device_ptr.h>
#include <thrust/memory.h>

namespace RMM_NAMESPACE {
namespace mr {
/**
 * @brief 与 Thrust 容器和算法兼容的分配器。
 *
 * 此分配器使用 `rmm::device_async_resource_ref` 包装上游设备内存资源，用于在设备上分配内存。
 *
 * 分配器通过 `T` 参数化，`T` 是将为其分配内存的元素的类型。
 *
 * 此分配器满足 `thrust::Allocator` 和 `std::allocator` 的要求（带有流/资源的扩展）。
 *
 * 与标准的 `thrust::device_malloc_allocator` 或 `thrust::cuda::malloc_allocator` 不同，此
 * 分配器使用指定的流和资源执行分配和释放。
 */
template <typename T>
class thrust_allocator : public thrust::device_malloc_allocator<T> {
 public
 using Base = thrust::device_malloc_allocator<T>;
 using pointer = typename Base::pointer;
 using size_type = typename Base::size_type;

 /** @brief 提供使用另一种类型实例化的 thrust_allocator 的类型。*/
 template <typename U>
 struct rebind {
 using other = thrust_allocator<U>;
 };

 /** @brief 默认构造函数使用默认内存资源和默认流创建分配器。 */
 thrust_allocator() = default;

 /** @brief 使用默认设备内存资源和指定流构造 thrust_allocator。 */
 explicit thrust_allocator(cuda_stream_view stream) : _stream{stream} {}

 /** @brief 使用设备内存资源和流构造 thrust_allocator。 */
 thrust_allocator(cuda_stream_view stream, rmm::device_async_resource_ref mr)
 : _stream{stream}, _mr(mr)
 {
 }

 /** @brief 复制构造函数。复制资源指针和流。 */
 template <typename U>
 thrust_allocator(thrust_allocator<U> const& other)
 : _mr(other.resource()), _stream{other.stream()}, _device{other._device}
 {
 }

 /** @brief 分配类型 T 的对象
 */
 pointer allocate(size_type num)
 {
 cuda_set_device_raii dev{_device};
 return thrust::device_pointer_cast(
 static_cast<T*>(_mr.allocate_async(num * sizeof(T), _stream)));
 }

 /** @brief 释放类型 T 的对象
 */
 void deallocate(pointer ptr, size_type num)