All Classes Files Functions Variables Typedefs Enumerations Enumerator Friends Modules Pages
fixed_size_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/cuda_stream_view.hpp>
20 #include <rmm/detail/error.hpp>
21 #include <rmm/detail/export.hpp>
22 #include <rmm/detail/logging_assert.hpp>
23 #include <rmm/detail/thrust_namespace.h>
24 #include <rmm/mr/device/detail/fixed_size_free_list.hpp>
25 #include <rmm/mr/device/detail/stream_ordered_memory_resource.hpp>
26 #include <rmm/resource_ref.hpp>
27 
28 #include <cuda_runtime_api.h>
29 #include <thrust/iterator/counting_iterator.h>
30 #include <thrust/iterator/transform_iterator.h>
31 
32 #include <algorithm>
33 #include <cstddef>
34 #include <utility>
35 #include <vector>
36 
37 namespace RMM_NAMESPACE {
38 namespace mr {
50 template <typename Upstream>
52  : public detail::stream_ordered_memory_resource<fixed_size_memory_resource<Upstream>,
53  detail::fixed_size_free_list> {
54  public:
55  friend class detail::stream_ordered_memory_resource<fixed_size_memory_resource<Upstream>,
56  detail::fixed_size_free_list>;
57 
58  static constexpr std::size_t default_block_size = 1 << 20;
59 
62  static constexpr std::size_t default_blocks_to_preallocate = 128;
63 
75  explicit fixed_size_memory_resource(
76  device_async_resource_ref upstream_mr,
77  // NOLINTNEXTLINE bugprone-easily-swappable-parameters
78  std::size_t block_size = default_block_size,
79  std::size_t blocks_to_preallocate = default_blocks_to_preallocate)
80  : upstream_mr_{upstream_mr},
81  block_size_{align_up(block_size, CUDA_ALLOCATION_ALIGNMENT)},
82  upstream_chunk_size_{block_size_ * blocks_to_preallocate}
83  {
84  // allocate initial blocks and insert into free list
85  this->insert_blocks(std::move(blocks_from_upstream(cuda_stream_legacy)), cuda_stream_legacy);
86  }
87 
99  explicit fixed_size_memory_resource(
100  Upstream* upstream_mr,
101  // NOLINTNEXTLINE bugprone-easily-swappable-parameters
102  std::size_t block_size = default_block_size,
103  std::size_t blocks_to_preallocate = default_blocks_to_preallocate)
104  : upstream_mr_{to_device_async_resource_ref_checked(upstream_mr)},
105  block_size_{align_up(block_size, CUDA_ALLOCATION_ALIGNMENT)},
106  upstream_chunk_size_{block_size_ * blocks_to_preallocate}
107  {
108  // allocate initial blocks and insert into free list
109  this->insert_blocks(std::move(blocks_from_upstream(cuda_stream_legacy)), cuda_stream_legacy);
110  }
111 
116  ~fixed_size_memory_resource() override { release(); }
117 
118  fixed_size_memory_resource() = delete;
121  fixed_size_memory_resource& operator=(fixed_size_memory_resource const&) = delete;
123 
127  [[nodiscard]] device_async_resource_ref get_upstream_resource() const noexcept
128  {
129  return upstream_mr_;
130  }
131 
137  [[nodiscard]] std::size_t get_block_size() const noexcept { return block_size_; }
138 
139  protected:
140  using free_list = detail::fixed_size_free_list;
141  using block_type = free_list::block_type;
142  using typename detail::stream_ordered_memory_resource<fixed_size_memory_resource<Upstream>,
143  detail::fixed_size_free_list>::split_block;
144  using lock_guard = std::lock_guard<std::mutex>;
145 
152  [[nodiscard]] std::size_t get_maximum_allocation_size() const { return get_block_size(); }
153 
165  block_type expand_pool(std::size_t size, free_list& blocks, cuda_stream_view stream)
166  {
167  blocks.insert(std::move(blocks_from_upstream(stream)));
168  return blocks.get_block(size);
169  }
170 
178  {
179  void* ptr = get_upstream_resource().allocate_async(upstream_chunk_size_, stream);
180  block_type block{ptr};
181  upstream_blocks_.push_back(block);
182 
183  auto num_blocks = upstream_chunk_size_ / block_size_;
184 
185  auto block_gen = [ptr, this](int index) {
186  // NOLINTNEXTLINE(cppcoreguidelines-pro-bounds-pointer-arithmetic)
187  return block_type{static_cast<char*>(ptr) + index * block_size_};
188  };
189  auto first =
190  thrust::make_transform_iterator(thrust::make_counting_iterator(std::size_t{0}), block_gen);
191  return free_list(first, first + num_blocks);
192  }
193 
204  split_block allocate_from_block(block_type const& block, std::size_t size)
205  {
206  return {block, block_type{nullptr}};
207  }
208 
217  block_type free_block(void* ptr, std::size_t size) noexcept
218  {
219  // Deallocating a fixed-size block just inserts it in the free list, which is
220  // handled by the parent class
221  RMM_LOGGING_ASSERT(align_up(size, CUDA_ALLOCATION_ALIGNMENT) <= block_size_);
222  return block_type{ptr};
223  }
224 
229  void release()
230  {
231  lock_guard lock(this->get_mutex());
232 
233  for (auto block : upstream_blocks_) {
234  get_upstream_resource().deallocate(block.pointer(), upstream_chunk_size_);
235  }
236  upstream_blocks_.clear();
237  }
238 
239 #ifdef RMM_DEBUG_PRINT
240  void print()
241  {
242  lock_guard lock(this->get_mutex());
243 
244  auto const [free, total] = rmm::available_device_memory();
245  std::cout << "GPU free memory: " << free << " total: " << total << "\n";
246 
247  std::cout << "upstream_blocks: " << upstream_blocks_.size() << "\n";
248  std::size_t upstream_total{0};
249 
250  for (auto blocks : upstream_blocks_) {
251  blocks.print();
252  upstream_total += upstream_chunk_size_;
253  }
254  std::cout << "total upstream: " << upstream_total << " B\n";
255 
256  this->print_free_blocks();
257  }
258 #endif
259 
268  std::pair<std::size_t, std::size_t> free_list_summary(free_list const& blocks)
269  {
270  return blocks.is_empty() ? std::make_pair(std::size_t{0}, std::size_t{0})
271  : std::make_pair(block_size_, blocks.size() * block_size_);
272  }
273 
274  private:
275  device_async_resource_ref upstream_mr_; // The resource from which to allocate new blocks
276 
277  std::size_t block_size_; // size of blocks this MR allocates
278  std::size_t upstream_chunk_size_; // size of chunks allocated from heap MR
279 
280  // blocks allocated from heap: so they can be easily freed
281  std::vector<block_type> upstream_blocks_;
282 };
283  // end of group
285 } // namespace mr
286 } // namespace RMM_NAMESPACE
Strongly-typed non-owning wrapper for CUDA streams with default constructor.
Definition: cuda_stream_view.hpp:39
A device_memory_resource which allocates memory blocks of a single fixed size.
Definition: fixed_size_memory_resource.hpp:53
std::pair< std::size_t, std::size_t > free_list_summary(free_list const &blocks)
Get the largest available block size and total free size in the specified free list.
Definition: fixed_size_memory_resource.hpp:268
detail::fixed_size_free_list free_list
The free list type.
Definition: fixed_size_memory_resource.hpp:140
block_type free_block(void *ptr, std::size_t size) noexcept
Finds, frees and returns the block associated with pointer.
Definition: fixed_size_memory_resource.hpp:217
std::size_t get_block_size() const noexcept
Get the size of blocks allocated by this memory resource.
Definition: fixed_size_memory_resource.hpp:137
std::size_t get_maximum_allocation_size() const
Get the (fixed) size of allocations supported by this memory resource.
Definition: fixed_size_memory_resource.hpp:152
free_list::block_type block_type
The type of block managed by the free list.
Definition: fixed_size_memory_resource.hpp:141
device_async_resource_ref get_upstream_resource() const noexcept
device_async_resource_ref to the upstream resource
Definition: fixed_size_memory_resource.hpp:127
block_type expand_pool(std::size_t size, free_list &blocks, cuda_stream_view stream)
Allocate a block from upstream to supply the suballocation pool.
Definition: fixed_size_memory_resource.hpp:165
free_list blocks_from_upstream(cuda_stream_view stream)
Allocate blocks from upstream to expand the suballocation pool.
Definition: fixed_size_memory_resource.hpp:177
std::lock_guard< std::mutex > lock_guard
Type of lock used to synchronize access.
Definition: fixed_size_memory_resource.hpp:144
void release()
free all memory allocated using the upstream resource.
Definition: fixed_size_memory_resource.hpp:229
split_block allocate_from_block(block_type const &block, std::size_t size)
Splits block if necessary to return a pointer to memory of size bytes.
Definition: fixed_size_memory_resource.hpp:204
std::pair< std::size_t, std::size_t > available_device_memory()
Returns the available and total device memory in bytes for the current device.
Definition: cuda_device.hpp:123
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
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