All Classes Files Functions Variables Typedefs Enumerations Enumerator Friends Modules Pages
binning_memory_resource.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/aligned.hpp>
19 #include <rmm/detail/export.hpp>
22 #include <rmm/resource_ref.hpp>
23 
24 #include <cuda_runtime_api.h>
25 
26 #include <algorithm>
27 #include <cassert>
28 #include <map>
29 #include <memory>
30 #include <optional>
31 #include <vector>
32 
33 namespace RMM_NAMESPACE {
34 namespace mr {
47 template <typename Upstream>
49  public:
59  : upstream_mr_{upstream_resource}
60  {
61  }
62 
73  explicit binning_memory_resource(Upstream* upstream_resource)
74  : upstream_mr_{to_device_async_resource_ref_checked(upstream_resource)}
75  {
76  }
77 
91  int8_t min_size_exponent, // NOLINT(bugprone-easily-swappable-parameters)
92  int8_t max_size_exponent)
93  : upstream_mr_{upstream_resource}
94  {
95  for (auto i = min_size_exponent; i <= max_size_exponent; i++) {
96  add_bin(1 << i);
97  }
98  }
99 
114  binning_memory_resource(Upstream* upstream_resource,
115  int8_t min_size_exponent, // NOLINT(bugprone-easily-swappable-parameters)
116  int8_t max_size_exponent)
117  : upstream_mr_{to_device_async_resource_ref_checked(upstream_resource)}
118  {
119  for (auto i = min_size_exponent; i <= max_size_exponent; i++) {
120  add_bin(1 << i);
121  }
122  }
123 
128  ~binning_memory_resource() override = default;
129 
130  binning_memory_resource() = delete;
133  binning_memory_resource& operator=(binning_memory_resource const&) = delete;
134  binning_memory_resource& operator=(binning_memory_resource&&) = delete;
135 
139  [[nodiscard]] device_async_resource_ref get_upstream_resource() const noexcept
140  {
141  return upstream_mr_;
142  }
143 
159  void add_bin(std::size_t allocation_size,
160  std::optional<device_async_resource_ref> bin_resource = std::nullopt)
161  {
162  allocation_size = align_up(allocation_size, CUDA_ALLOCATION_ALIGNMENT);
163 
164  if (bin_resource.has_value()) {
165  resource_bins_.insert({allocation_size, bin_resource.value()});
166  } else if (resource_bins_.count(allocation_size) == 0) { // do nothing if bin already exists
167  owned_bin_resources_.push_back(
168  std::make_unique<fixed_size_memory_resource<Upstream>>(upstream_mr_, allocation_size));
169  resource_bins_.insert({allocation_size, owned_bin_resources_.back().get()});
170  }
171  }
172 
173  private:
182  device_async_resource_ref get_resource_ref(std::size_t bytes)
183  {
184  auto iter = resource_bins_.lower_bound(bytes);
185  return (iter != resource_bins_.cend()) ? iter->second : get_upstream_resource();
186  }
187 
197  void* do_allocate(std::size_t bytes, cuda_stream_view stream) override
198  {
199  if (bytes <= 0) { return nullptr; }
200  return get_resource_ref(bytes).allocate_async(bytes, stream);
201  }
202 
211  void do_deallocate(void* ptr, std::size_t bytes, cuda_stream_view stream) override
212  {
213  get_resource_ref(bytes).deallocate_async(ptr, bytes, stream);
214  }
215 
217  upstream_mr_; // The upstream memory_resource from which to allocate blocks.
218 
219  std::vector<std::unique_ptr<fixed_size_memory_resource<Upstream>>> owned_bin_resources_;
220 
221  std::map<std::size_t, device_async_resource_ref> resource_bins_;
222 };
223  // end of group
225 } // namespace mr
226 } // namespace RMM_NAMESPACE
Strongly-typed non-owning wrapper for CUDA streams with default constructor.
Definition: cuda_stream_view.hpp:41
Allocates memory from upstream resources associated with bin sizes.
Definition: binning_memory_resource.hpp:48
binning_memory_resource(Upstream *upstream_resource)
Construct a new binning memory resource object.
Definition: binning_memory_resource.hpp:73
device_async_resource_ref get_upstream_resource() const noexcept
device_async_resource_ref to the upstream resource
Definition: binning_memory_resource.hpp:139
binning_memory_resource(device_async_resource_ref upstream_resource, int8_t min_size_exponent, int8_t max_size_exponent)
Construct a new binning memory resource object with a range of initial bins.
Definition: binning_memory_resource.hpp:90
~binning_memory_resource() override=default
Destroy the binning_memory_resource and free all memory allocated from the upstream resource.
binning_memory_resource(Upstream *upstream_resource, int8_t min_size_exponent, int8_t max_size_exponent)
Construct a new binning memory resource object with a range of initial bins.
Definition: binning_memory_resource.hpp:114
binning_memory_resource(device_async_resource_ref upstream_resource)
Construct a new binning memory resource object.
Definition: binning_memory_resource.hpp:58
void add_bin(std::size_t allocation_size, std::optional< device_async_resource_ref > bin_resource=std::nullopt)
Add a bin allocator to this resource.
Definition: binning_memory_resource.hpp:159
Base class for all libcudf device memory allocation.
Definition: device_memory_resource.hpp:94
A device_memory_resource which allocates memory blocks of a single fixed size.
Definition: fixed_size_memory_resource.hpp:55
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
static constexpr std::size_t CUDA_ALLOCATION_ALIGNMENT
Default alignment used for CUDA memory allocation.
Definition: aligned.hpp:43
constexpr std::size_t align_up(std::size_t value, std::size_t alignment) noexcept
Align up to nearest multiple of specified power of 2.
Definition: aligned.hpp:77