binning_memory_resource.hpp

前往此文件的文档。

/*
 * 版权所有 (c) 2020-2025, 英伟达公司。
 *
 * 根据 Apache 许可证 2.0 版本（“许可证”）获得许可；
 * 除非遵守许可证，否则您不得使用此文件。
 * 您可以在以下网址获得许可证副本：
 *
 * https://apache.ac.cn/licenses/LICENSE-2.0
 *
 * 除非适用法律要求或书面同意，否则根据“按原样”基础分发的软件
 * 不提供任何明示或默示的保证或条件。
 * 请参阅许可证，了解特定语言的权限和限制。
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#pragma once

#include <rmm/aligned.hpp>
#include <rmm/detail/export.hpp>
#include <rmm/mr/device/device_memory_resource.hpp>
#include <rmm/mr/device/fixed_size_memory_resource.hpp>
#include <rmm/resource_ref.hpp>

#include <cuda_runtime_api.h>

#include <cassert>
#include <map>
#include <memory>
#include <optional>
#include <vector>

namespace RMM_NAMESPACE {
namespace mr {
/// @brief 从与 bin 大小关联的上游资源分配内存。
template <typename Upstream>
class binning_memory_resource final : public device_memory_resource {
 public
 /// @brief 构造一个新的 binning 内存资源对象。
 /// @param upstream_resource 指向上游内存资源的 device_async_resource_ref。
 explicit binning_memory_resource(device_async_resource_ref upstream_resource)
 : upstream_mr_{upstream_resource}
 {
 }
 /// @brief 构造一个新的 binning 内存资源对象。
 /// @param upstream_resource 指向上游内存资源的指针。
 explicit binning_memory_resource(Upstream* upstream_resource)
 : upstream_mr_{to_device_async_resource_ref_checked(upstream_resource)}
 {
 }/// @brief 构造一个新的具有初始 bin 范围的 binning 内存资源对象。
 /// @param upstream_resource 指向上游内存资源的 device_async_resource_ref。
 /// @param min_size_exponent 最小 bin 大小，以 2 的幂为底。
 /// @param max_size_exponent 最大 bin 大小，以 2 的幂为底。
 binning_memory_resource(device_async_resource_ref upstream_resource,
 int8_t min_size_exponent, // NOLINT(bugprone-easily-swappable-parameters)
 int8_t max_size_exponent)
 : upstream_mr_{upstream_resource}
 {
 for (auto i = min_size_exponent; i <= max_size_exponent; i++) {
 add_bin(1 << i);
 }
 }/// @brief 构造一个新的具有初始 bin 范围的 binning 内存资源对象。
 /// @param upstream_resource 指向上游内存资源的指针。
 /// @param min_size_exponent 最小 bin 大小，以 2 的幂为底。
 /// @param max_size_exponent 最大 bin 大小，以 2 的幂为底。
 binning_memory_resource(Upstream* upstream_resource,
 int8_t min_size_exponent, // NOLINT(bugprone-easily-swappable-parameters)
 int8_t max_size_exponent)
 : upstream_mr_{to_device_async_resource_ref_checked(upstream_resource)}
 {
 for (auto i = min_size_exponent; i <= max_size_exponent; i++) {
 add_bin(1 << i);
 }
 }/// @brief 销毁 binning_memory_resource 并释放从上游资源分配的所有内存。
 ~binning_memory_resource() override = default;

 binning_memory_resource() = delete;
 binning_memory_resource(binning_memory_resource const&) = delete;
 binning_memory_resource(binning_memory_resource&&) = delete;
 binning_memory_resource& operator=(binning_memory_resource const&) = delete;
 binning_memory_resource& operator=(binning_memory_resource&&) = delete;

 /// @brief 指向上游资源的 device_async_resource_ref
 /// @return device_async_resource_ref 指向上游资源的 device_async_resource_ref
 [[nodiscard]] device_async_resource_ref get_upstream_resource() const noexcept
 {
 return upstream_mr_;
 }

 /// @brief 向此资源添加一个 bin 分配器。
 /// @param allocation_size 此 bin 分配器分配的固定大小。
 /// @param bin_resource 可选地提供要在此 bin 中使用的现有内存资源。如果提供，则不会在此 bin 中创建新的 fixed_size_memory_resource。
 void add_bin(std::size_t allocation_size,
 std::optional<device_async_resource_ref> bin_resource = std::nullopt)
 {
 allocation_size = align_up(allocation_size, CUDA_ALLOCATION_ALIGNMENT);

 if (bin_resource.has_value()) {
 resource_bins_.insert({allocation_size, bin_resource.value()});
 } else if (resource_bins_.count(allocation_size) == 0) { // 如果 bin 已存在则不执行任何操作
 owned_bin_resources_.push_back(
 std::make_unique<fixed_size_memory_resource<Upstream>>(upstream_mr_, allocation_size));
 resource_bins_.insert({allocation_size, owned_bin_resources_.back().get()});
 }
 }

 private
 /// @brief 查找能够容纳所请求分配大小的最小尺寸的 bin 资源。
 /// @param bytes 要分配的大小。
 /// @return device_async_resource_ref 能够容纳所请求分配大小的最小尺寸的 bin 资源。
 device_async_resource_ref get_resource_ref(std::size_t bytes)
 {
 auto iter = resource_bins_.lower_bound(bytes);
 return (iter != resource_bins_.cend()) ? iter->second : get_upstream_resource();
 }/// @brief 使用适当的 bin 资源或上游资源分配大小为 \p bytes 的内存。
 /// @param bytes 要分配的大小。
 /// @param stream 要执行分配的 CUDA 流。
 /// @return void* 指向分配内存的指针。
 void* do_allocate(std::size_t bytes, cuda_stream_view stream) override
 {
 if (bytes <= 0) { return nullptr; }
 return get_resource_ref(bytes).allocate_async(bytes, stream);