prefetch_resource_adaptor.hpp
转到此文件的文档。
1 /*
2  * 版权所有 (c) 2024-2025, NVIDIA CORPORATION.
3  *
4  * 根据 Apache 许可版本 2.0 (以下简称“许可”) 获得许可；
5  * 除非遵守许可，否则您不得使用此文件。
6  * 您可以在以下网址获得许可副本：
7  *
8  * https://apache.ac.cn/licenses/LICENSE-2.0
9  *
10  * 除非适用法律要求或书面同意，否则根据许可分发的软件
11  * 按“原样”分发，不附带任何明示或暗示的保证或条件。
12  * 有关管理权限和限制的特定语言，请参阅许可。
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 #pragma once
17 
18 #include <rmm/detail/export.hpp>
19 #include <rmm/mr/device/device_memory_resource.hpp>
20 #include <rmm/prefetch.hpp>
21 #include <rmm/resource_ref.hpp>
22 
23 #include <cstddef>
24 
25 namespace RMM_NAMESPACE {
26 namespace mr {
38 /// @addtogroup memory_resources
39 template <typename Upstream>
40 /// Resource that adds prefetching all memory allocations to the CUDA device
41 /// that allocated the memory.
42 /// @tparam Upstream The type of the upstream resource.
43  class prefetch_resource_adaptor final : public device_memory_resource {
44  public
45  /// @brief Constructor
46  /// @param upstream The upstream resource used to satisfy allocation requests.
47  prefetch_resource_adaptor(device_async_resource_ref upstream) : upstream_{upstream} {}
48 
49  /// @brief Constructor
50  /// @param upstream The upstream resource used to satisfy allocation requests.
51  /// @throw rmm::bad_alloc if constructing a device_async_resource_ref from `upstream`
52  /// fails (e.g., because `upstream` is `nullptr`).
53  /// @todo Consider removing this overload because it's not generally safe to store
54  /// raw pointers to resources that might have a shorter lifetime than the adaptor.
55  /// However, it is convenient when used with `get_default_resource()`.
56   template <typename U = Upstream, typename = std::enable_if_t<!std::is_base_of_v<device_async_resource_ref, U>>>
57  prefetch_resource_adaptor(Upstream* upstream)
58  : upstream_{to_device_async_resource_ref_checked(upstream)}
59  {
60  }
61 
62  prefetch_resource_adaptor() = delete;
63  ~prefetch_resource_adaptor() override = default;
64  prefetch_resource_adaptor(prefetch_resource_adaptor const&) = delete;
65  prefetch_resource_adaptor& operator=(prefetch_resource_adaptor const&) = delete;
66  prefetch_resource_adaptor(prefetch_resource_adaptor&&) noexcept =
67  default;
68  prefetch_resource_adaptor& operator=(prefetch_resource_adaptor&&) noexcept =
69  default;
70 
71  /// @brief Get the upstream resource used by this adaptor.
72  /// @return The upstream resource used by this adaptor.
73   // @todo Remove if no longer used outside testing.
74  [[nodiscard]] rmm::device_async_resource_ref get_upstream_resource() const noexcept
75  {
76  return upstream_;
77  }
78 
79  private
80  /// @brief Implements allocation.
81  ///
82  /// Satisfies the allocation request using the upstream resource and then
83  /// prefetches the allocated memory to the device that allocated it.
84  ///
85  /// @param bytes The number of bytes to allocate.
86  /// @param stream The stream to perform allocation and prefetch on.
87  /// @return Pointer to the allocated memory.
88  /// @throw rmm::bad_alloc if the requested allocation size cannot be met.
89   // Documented in base class
90  // NOLINTNEXTLINE(google-default-arguments)
91  // NOLINTNEXTLINE(rmm/require-stream-aware-allocator)
92   void* do_allocate(std::size_t bytes, cuda_stream_view stream) override
93  {
94  void* ptr = get_upstream_resource().allocate_async(bytes, stream);
95  rmm::prefetch(ptr, bytes, rmm::get_current_cuda_device(), stream);
96  return ptr;
97  }
98 
99  /// @brief Implements deallocation.
100  ///
101  /// Satisfies the deallocation request using the upstream resource.
102  ///
103  /// @param ptr The pointer to the memory to deallocate.
104  /// @param bytes The size of the memory to deallocate.
102  ///
105  /// @param stream The stream to perform deallocation on.
106   // Documented in base class
107  void do_deallocate(void* ptr, std::size_t bytes, cuda_stream_view stream) override
102  ///
108  {
106   // Documented in base class
109  get_upstream_resource().deallocate_async(ptr, bytes, stream);
110  }
111 
112  /// @brief Compares two prefetch resources for equality.
113  ///
114  /// Prefetch resources are equal if and only if they are the same object
115  /// or if they are both prefetch resources using the same upstream resource.
116  /// @param other The other resource to compare to.
117  /// @return True if the two resources are equal, false otherwise.
118   // Documented in base class
119  bool do_is_equal(device_memory_resource const& other) const noexcept override
120  {
121  if (this == &other) { return true; }
122  auto cast = dynamic_cast<prefetch_resource_adaptor<Upstream> const*>(&other);
123  if (cast == nullptr) { return false; }
124  return get_upstream_resource() == cast->get_upstream_resource();
125  }
126 
127  // the upstream resource used for satisfying allocation requests
128  device_async_resource_ref upstream_;
129 };