All Classes Files Functions Variables Typedefs Enumerations Enumerator Friends Modules Pages
thread_safe_resource_adaptor.hpp
Go to the documentation of this file.
1 /*
2  * Copyright (c) 2020-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/cuda_stream_view.hpp>
19 #include <rmm/detail/error.hpp>
20 #include <rmm/detail/export.hpp>
22 #include <rmm/resource_ref.hpp>
23 
24 #include <cstddef>
25 #include <mutex>
26 
27 namespace RMM_NAMESPACE {
28 namespace mr {
43 template <typename Upstream>
45  public:
46  using lock_t = std::lock_guard<std::mutex>;
47 
56  thread_safe_resource_adaptor(device_async_resource_ref upstream) : upstream_{upstream} {}
57 
68  thread_safe_resource_adaptor(Upstream* upstream)
69  : upstream_{to_device_async_resource_ref_checked(upstream)}
70  {
71  }
72 
74  ~thread_safe_resource_adaptor() override = default;
79 
84  {
85  return upstream_;
86  }
87 
88  private:
100  void* do_allocate(std::size_t bytes, cuda_stream_view stream) override
101  {
102  lock_t lock(mtx);
103  return get_upstream_resource().allocate_async(bytes, stream);
104  }
105 
113  void do_deallocate(void* ptr, std::size_t bytes, cuda_stream_view stream) override
114  {
115  lock_t lock(mtx);
116  get_upstream_resource().deallocate_async(ptr, bytes, stream);
117  }
118 
126  bool do_is_equal(device_memory_resource const& other) const noexcept override
127  {
128  if (this == &other) { return true; }
129  auto cast = dynamic_cast<thread_safe_resource_adaptor<Upstream> const*>(&other);
130  if (cast == nullptr) { return false; }
131  return get_upstream_resource() == cast->get_upstream_resource();
132  }
133 
134  std::mutex mutable mtx; // mutex for thread safe access to upstream
136  upstream_;
137 };
138  // end of group
140 } // namespace mr
141 } // namespace RMM_NAMESPACE
Strongly-typed non-owning wrapper for CUDA streams with default constructor.
Definition: cuda_stream_view.hpp:39
Base class for all librmm device memory allocation.
Definition: device_memory_resource.hpp:93
Resource that adapts Upstream memory resource adaptor to be thread safe.
Definition: thread_safe_resource_adaptor.hpp:44
std::lock_guard< std::mutex > lock_t
Type of lock used to synchronize access.
Definition: thread_safe_resource_adaptor.hpp:46
thread_safe_resource_adaptor(device_async_resource_ref upstream)
Construct a new thread safe resource adaptor using upstream to satisfy allocation requests.
Definition: thread_safe_resource_adaptor.hpp:56
rmm::device_async_resource_ref get_upstream_resource() const noexcept
rmm::device_async_resource_ref to the upstream resource
Definition: thread_safe_resource_adaptor.hpp:83
thread_safe_resource_adaptor(Upstream *upstream)
Construct a new thread safe resource adaptor using upstream to satisfy allocation requests.
Definition: thread_safe_resource_adaptor.hpp:68
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