All Classes Files Functions Variables Typedefs Enumerations Enumerator Friends Modules Pages
prefetch_resource_adaptor.hpp
Go to the documentation of this file.
1 /*
2  * Copyright (c) 2024, NVIDIA CORPORATION.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  * http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
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>
20 #include <rmm/prefetch.hpp>
21 #include <rmm/resource_ref.hpp>
22 
23 #include <cstddef>
24 #include <mutex>
25 #include <shared_mutex>
26 #include <stack>
27 
28 namespace RMM_NAMESPACE {
29 namespace mr {
41 template <typename Upstream>
43  public:
50  prefetch_resource_adaptor(device_async_resource_ref upstream) : upstream_{upstream} {}
51 
60  prefetch_resource_adaptor(Upstream* upstream)
61  : upstream_{to_device_async_resource_ref_checked(upstream)}
62  {
63  }
64 
65  prefetch_resource_adaptor() = delete;
66  ~prefetch_resource_adaptor() override = default;
68  prefetch_resource_adaptor& operator=(prefetch_resource_adaptor const&) = delete;
70  default;
72  default;
73 
77  [[nodiscard]] rmm::device_async_resource_ref get_upstream_resource() const noexcept
78  {
79  return upstream_;
80  }
81 
82  private:
96  void* do_allocate(std::size_t bytes, cuda_stream_view stream) override
97  {
98  void* ptr = get_upstream_resource().allocate_async(bytes, stream);
99  rmm::prefetch(ptr, bytes, rmm::get_current_cuda_device(), stream);
100  return ptr;
101  }
102 
110  void do_deallocate(void* ptr, std::size_t bytes, cuda_stream_view stream) override
111  {
112  get_upstream_resource().deallocate_async(ptr, bytes, stream);
113  }
114 
122  bool do_is_equal(device_memory_resource const& other) const noexcept override
123  {
124  if (this == &other) { return true; }
125  auto cast = dynamic_cast<prefetch_resource_adaptor<Upstream> const*>(&other);
126  if (cast == nullptr) { return false; }
127  return get_upstream_resource() == cast->get_upstream_resource();
128  }
129 
130  // the upstream resource used for satisfying allocation requests
131  device_async_resource_ref upstream_;
132 };
133  // end of group
135 } // namespace mr
136 } // namespace RMM_NAMESPACE
Strongly-typed non-owning wrapper for CUDA streams with default constructor.
Definition: cuda_stream_view.hpp:41
Base class for all libcudf device memory allocation.
Definition: device_memory_resource.hpp:94
void * allocate_async(std::size_t bytes, std::size_t alignment, cuda_stream_view stream)
Allocates memory of size at least bytes.
Definition: device_memory_resource.hpp:217
Resource that prefetches all memory allocations.
Definition: prefetch_resource_adaptor.hpp:42
prefetch_resource_adaptor(prefetch_resource_adaptor &&) noexcept=default
Default move constructor.
prefetch_resource_adaptor(Upstream *upstream)
Construct a new prefetch resource adaptor using upstream to satisfy allocation requests.
Definition: prefetch_resource_adaptor.hpp:60
prefetch_resource_adaptor(device_async_resource_ref upstream)
Construct a new prefetch resource adaptor using upstream to satisfy allocation requests.
Definition: prefetch_resource_adaptor.hpp:50
cuda_device_id get_current_cuda_device()
Returns a cuda_device_id for the current device.
Definition: cuda_device.hpp:96
cuda::mr::async_resource_ref< cuda::mr::device_accessible > device_async_resource_ref
Alias for a cuda::mr::async_resource_ref with the property cuda::mr::device_accessible.
Definition: resource_ref.hpp:41
device_async_resource_ref to_device_async_resource_ref_checked(Resource *res)
Convert pointer to memory resource into device_async_resource_ref, checking for nullptr
Definition: resource_ref.hpp:79
void prefetch(void const *ptr, std::size_t size, rmm::cuda_device_id device, rmm::cuda_stream_view stream)
Prefetch memory to the specified device on the specified stream.
Definition: prefetch.hpp:46