All Classes Files Functions Variables Typedefs Enumerations Enumerator Friends Modules Pages
sam_headroom_memory_resource.hpp
Go to the documentation of this file.
1 /*
2  * Copyright (c) 2024-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/cuda_device.hpp>
19 #include <rmm/cuda_stream_view.hpp>
20 #include <rmm/detail/export.hpp>
23 #include <rmm/resource_ref.hpp>
24 
25 #include <algorithm>
26 #include <cstddef>
27 
28 namespace RMM_NAMESPACE {
29 namespace mr {
49  public:
55  explicit sam_headroom_memory_resource(std::size_t headroom) : system_mr_{}, headroom_{headroom} {}
56 
58  ~sam_headroom_memory_resource() override = default;
63 
64  private:
76  void* do_allocate(std::size_t bytes, [[maybe_unused]] cuda_stream_view stream) override
77  {
78  void* pointer = system_mr_.allocate_async(bytes, rmm::CUDA_ALLOCATION_ALIGNMENT, stream);
79 
80  auto const free = rmm::available_device_memory().first;
81  auto const allocatable = free > headroom_ ? free - headroom_ : 0UL;
82  auto const gpu_portion =
83  rmm::align_down(std::min(allocatable, bytes), rmm::CUDA_ALLOCATION_ALIGNMENT);
84  auto const cpu_portion = bytes - gpu_portion;
85  if (gpu_portion != 0) {
86  RMM_CUDA_TRY(cudaMemAdvise(pointer,
87  gpu_portion,
88  cudaMemAdviseSetPreferredLocation,
89  rmm::get_current_cuda_device().value()));
90  }
91  if (cpu_portion != 0) {
92  RMM_CUDA_TRY(cudaMemAdvise(static_cast<char*>(pointer) + gpu_portion,
93  cpu_portion,
94  cudaMemAdviseSetPreferredLocation,
95  cudaCpuDeviceId));
96  }
97 
98  return pointer;
99  }
100 
110  void do_deallocate(void* ptr,
111  [[maybe_unused]] std::size_t bytes,
112  [[maybe_unused]] cuda_stream_view stream) override
113  {
114  system_mr_.deallocate_async(ptr, rmm::CUDA_ALLOCATION_ALIGNMENT, stream);
115  }
116 
124  [[nodiscard]] bool do_is_equal(device_memory_resource const& other) const noexcept override
125  {
126  if (this == &other) { return true; }
127  auto cast = dynamic_cast<sam_headroom_memory_resource const*>(&other);
128  if (cast == nullptr) { return false; }
129  return headroom_ == cast->headroom_;
130  }
131 
133  system_memory_resource system_mr_;
135  std::size_t headroom_;
136 }; // end of group
138 } // namespace mr
139 } // 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
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:216
Resource that uses system memory resource to allocate memory with a headroom.
Definition: sam_headroom_memory_resource.hpp:48
sam_headroom_memory_resource(std::size_t headroom)
Construct a headroom memory resource.
Definition: sam_headroom_memory_resource.hpp:55
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_device_id get_current_cuda_device()
Returns a cuda_device_id for the current device.
Definition: cuda_device.hpp:99
static constexpr std::size_t CUDA_ALLOCATION_ALIGNMENT
Default alignment used for CUDA memory allocation.
Definition: aligned.hpp:43
constexpr std::size_t align_down(std::size_t value, std::size_t alignment) noexcept
Align down to the nearest multiple of specified power of 2.
Definition: aligned.hpp:91