All Classes Files Functions Variables Typedefs Enumerations Enumerator Friends Modules Pages
failure_callback_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/detail/error.hpp>
19 #include <rmm/detail/export.hpp>
22 #include <rmm/resource_ref.hpp>
23 
24 #include <cstddef>
25 #include <functional>
26 #include <utility>
27 
28 namespace RMM_NAMESPACE {
29 namespace mr {
51 using failure_callback_t = std::function<bool(std::size_t, void*)>;
52 
94 template <typename Upstream, typename ExceptionType = rmm::out_of_memory>
96  public:
97  using exception_type = ExceptionType;
98 
108  failure_callback_t callback,
109  void* callback_arg)
110  : upstream_{upstream}, callback_{std::move(callback)}, callback_arg_{callback_arg}
111  {
112  }
113 
125  failure_callback_t callback,
126  void* callback_arg)
127  : upstream_{to_device_async_resource_ref_checked(upstream)},
128  callback_{std::move(callback)},
129  callback_arg_{callback_arg}
130  {
131  }
132 
134  ~failure_callback_resource_adaptor() override = default;
138  default;
140  default;
141 
145  [[nodiscard]] rmm::device_async_resource_ref get_upstream_resource() const noexcept
146  {
147  return upstream_;
148  }
149 
150  private:
162  void* do_allocate(std::size_t bytes, cuda_stream_view stream) override
163  {
164  void* ret{};
165 
166  while (true) {
167  try {
168  ret = get_upstream_resource().allocate_async(bytes, stream);
169  break;
170  } catch (exception_type const& e) {
171  if (!callback_(bytes, callback_arg_)) { throw; }
172  }
173  }
174  return ret;
175  }
176 
184  void do_deallocate(void* ptr, std::size_t bytes, cuda_stream_view stream) override
185  {
186  get_upstream_resource().deallocate_async(ptr, bytes, stream);
187  }
188 
196  [[nodiscard]] bool do_is_equal(device_memory_resource const& other) const noexcept override
197  {
198  if (this == &other) { return true; }
199  auto cast = dynamic_cast<failure_callback_resource_adaptor<Upstream> const*>(&other);
200  if (cast == nullptr) { return false; }
201  return get_upstream_resource() == cast->get_upstream_resource();
202  }
203 
204  // the upstream resource used for satisfying allocation requests
205  device_async_resource_ref upstream_;
206  failure_callback_t callback_;
207  void* callback_arg_;
208 };
209  // end of group
211 } // namespace mr
212 } // 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
A device memory resource that calls a callback function when allocations throw a specified exception ...
Definition: failure_callback_resource_adaptor.hpp:95
failure_callback_resource_adaptor(device_async_resource_ref upstream, failure_callback_t callback, void *callback_arg)
Construct a new failure_callback_resource_adaptor using upstream to satisfy allocation requests.
Definition: failure_callback_resource_adaptor.hpp:107
ExceptionType exception_type
The type of exception this object catches/throws.
Definition: failure_callback_resource_adaptor.hpp:97
failure_callback_resource_adaptor(Upstream *upstream, failure_callback_t callback, void *callback_arg)
Construct a new failure_callback_resource_adaptor using upstream to satisfy allocation requests.
Definition: failure_callback_resource_adaptor.hpp:124
failure_callback_resource_adaptor(failure_callback_resource_adaptor &&) noexcept=default
Default move constructor.
std::function< bool(std::size_t, void *)> failure_callback_t
Callback function type used by failure_callback_resource_adaptor.
Definition: failure_callback_resource_adaptor.hpp:51
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
Management of per-device device_memory_resources.