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-2025, 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 <cassert>
27 #include <map>
28 #include <memory>
29 #include <optional>
30 #include <vector>
31 
32 namespace RMM_NAMESPACE {
33 namespace mr {
46 template <typename Upstream>
48  public:
58  : upstream_mr_{upstream_resource}
59  {
60  }
61 
72  explicit binning_memory_resource(Upstream* upstream_resource)
73  : upstream_mr_{to_device_async_resource_ref_checked(upstream_resource)}
74  {
75  }
76 
90  int8_t min_size_exponent, // NOLINT(bugprone-easily-swappable-parameters)
91  int8_t max_size_exponent)
92  : upstream_mr_{upstream_resource}
93  {
94  for (auto i = min_size_exponent; i <= max_size_exponent; i++) {
95  add_bin(1 << i);
96  }
97  }
98 
113  binning_memory_resource(Upstream* upstream_resource,
114  int8_t min_size_exponent, // NOLINT(bugprone-easily-swappable-parameters)
115  int8_t max_size_exponent)
116  : upstream_mr_{to_device_async_resource_ref_checked(upstream_resource)}
117  {
118  for (auto i = min_size_exponent; i <= max_size_exponent; i++) {
119  add_bin(1 << i);
120  }
121  }
122 
127  ~binning_memory_resource() override = default;
128 
129  binning_memory_resource() = delete;
132  binning_memory_resource& operator=(binning_memory_resource const&) = delete;
133  binning_memory_resource& operator=(binning_memory_resource&&) = delete;
134 
138  [[nodiscard]] device_async_resource_ref get_upstream_resource() const noexcept
139  {
140  return upstream_mr_;
141  }
142 
158  void add_bin(std::size_t allocation_size,
159  std::optional<device_async_resource_ref> bin_resource = std::nullopt)
160  {
161  allocation_size = align_up(allocation_size, CUDA_ALLOCATION_ALIGNMENT);
162 
163  if (bin_resource.has_value()) {
164  resource_bins_.insert({allocation_size, bin_resource.value()});
165  } else if (resource_bins_.count(allocation_size) == 0) { // do nothing if bin already exists
166  owned_bin_resources_.push_back(
167  std::make_unique<fixed_size_memory_resource<Upstream>>(upstream_mr_, allocation_size));
168  resource_bins_.insert({allocation_size, owned_bin_resources_.back().get()});
169  }
170  }
171 
172  private:
181  device_async_resource_ref get_resource_ref(std::size_t bytes)
182  {
183  auto iter = resource_bins_.lower_bound(bytes);
184  return (iter != resource_bins_.cend()) ? iter->second : get_upstream_resource();
185  }
186 
196  void* do_allocate(std::size_t bytes, cuda_stream_view stream) override
197  {
198  if (bytes <= 0) { return nullptr; }
199  return get_resource_ref(bytes).allocate_async(bytes, stream);
200  }
201 
210  void do_deallocate(void* ptr, std::size_t bytes, cuda_stream_view stream) override
211  {
212  get_resource_ref(bytes).deallocate_async(ptr, bytes, stream);
213  }
214 
216  upstream_mr_; // The upstream memory_resource from which to allocate blocks.
217 
218  std::vector<std::unique_ptr<fixed_size_memory_resource<Upstream>>> owned_bin_resources_;
219 
220  std::map<std::size_t, device_async_resource_ref> resource_bins_;
221 };
222  // end of group
224 } // namespace mr
225 } // namespace RMM_NAMESPACE
Strongly-typed non-owning wrapper for CUDA streams with default constructor.
Definition: cuda_stream_view.hpp:39
Allocates memory from upstream resources associated with bin sizes.
Definition: binning_memory_resource.hpp:47
binning_memory_resource(Upstream *upstream_resource)
Construct a new binning memory resource object.
Definition: binning_memory_resource.hpp:72
device_async_resource_ref get_upstream_resource() const noexcept
device_async_resource_ref to the upstream resource
Definition: binning_memory_resource.hpp:138
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:89
~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:113
binning_memory_resource(device_async_resource_ref upstream_resource)
Construct a new binning memory resource object.
Definition: binning_memory_resource.hpp:57
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:158
Base class for all librmm device memory allocation.
Definition: device_memory_resource.hpp:93
A device_memory_resource which allocates memory blocks of a single fixed size.
Definition: fixed_size_memory_resource.hpp:53
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