limiting_resource_adaptor.hpp

前往此文件的文档。

/*
 * Copyright (c) 2021-2025, NVIDIA CORPORATION.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * https://apache.ac.cn/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#pragma once

#include <rmm/aligned.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 <rmm/mr/device/per_device_resource.hpp>
#include <rmm/resource_ref.hpp>

#include <atomic>
#include <cstddef>

namespace RMM_NAMESPACE {
namespace mr {
/*!
 * @brief Resource that uses Upstream to allocate memory and limits the total allocations possible.
 */
 template <typename Upstream>
 class limiting_resource_adaptor final : public device_memory_resource {
 public
 /*!
 * @brief Construct a new limiting resource adaptor using upstream to satisfy allocation requests and
 * limiting total allocations by allocation_limit.
 *
 * @throws rmm::bad_arg if allocation_limit == 0
 * @param upstream The upstream resource used for satisfying allocation requests.
 * @param allocation_limit The maximum number of bytes this resource is allowed to allocate.
 * @param alignment The alignment to use for allocations. Defaults to CUDA_ALLOCATION_ALIGNMENT.
 */
 limiting_resource_adaptor(device_async_resource_ref upstream,
 std::size_t allocation_limit,
 std::size_t alignment = CUDA_ALLOCATION_ALIGNMENT)
 : upstream_{upstream},
 allocation_limit_{allocation_limit},
 allocated_bytes_(0),
 alignment_(alignment)
 {
 }

 /*!
 * @brief Construct a new limiting resource adaptor using upstream to satisfy allocation requests and
 * limiting total allocations by allocation_limit.
 *
 * @throws rmm::bad_arg if allocation_limit == 0 or upstream is nullptr
 * @param upstream The upstream resource used for satisfying allocation requests.
 * @param allocation_limit The maximum number of bytes this resource is allowed to allocate.
 * @param alignment The alignment to use for allocations. Defaults to CUDA_ALLOCATION_ALIGNMENT.
 */
 limiting_resource_adaptor(Upstream* upstream,
 std::size_t allocation_limit,
 std::size_t alignment = CUDA_ALLOCATION_ALIGNMENT)
 : upstream_{to_device_async_resource_ref_checked(upstream)},
 allocation_limit_{allocation_limit},
 allocated_bytes_(0),
 alignment_(alignment)
 {
 }

 limiting_resource_adaptor() = delete;
 ~limiting_resource_adaptor() override = default;
 limiting_resource_adaptor(limiting_resource_adaptor const&) = delete;
 limiting_resource_adaptor(limiting_resource_adaptor&&) noexcept =
 default;
 limiting_resource_adaptor& operator=(limiting_resource_adaptor const&) = delete;
 limiting_resource_adaptor& operator=(limiting_resource_adaptor&&) noexcept =
 default;

 /*!
 * @brief Get the upstream resource used for satisfying allocation requests.
 */
 [[nodiscard]] device_async_resource_ref get_upstream_resource() const noexcept
 {
 return upstream_;
 }

 /*!
 * @brief Query the number of bytes that have been allocated.
 *
 * Note that this can not be used to know how large a particular allocation was. The return value is the sum
 * of all sizes passed to `do_allocate` rounded up to the nearest multiple of the alignment.
 */
 [[nodiscard]] std::size_t get_allocated_bytes() const { return allocated_bytes_; }

 /*!
 * @brief Query the maximum number of bytes that this allocator is allowed to allocate.
 *
 * This is the limit on the sum of all sizes passed to `do_allocate` rounded up to the nearest multiple of
 * the alignment.
 */
 [[nodiscard]] std::size_t get_allocation_limit() const { return allocation_limit_; }

 private
 /*!
 * @brief Performs memory allocation.
 *
 * @param bytes The number of bytes to allocate.
 * @param stream The stream to use for the allocation.
 * @return void* Pointer to the newly allocated memory.
 * @throws rmm::out_of_memory if allocating size bytes would exceed the allocation limit.
 */
 void* do_allocate(std::size_t bytes, cuda_stream_view stream) override
 {
 auto const proposed_size = align_up(bytes, alignment_);
 auto const old = allocated_bytes_.fetch_add(proposed_size);
 if (old + proposed_size <= allocation_limit_) {