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, 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 namespace RMM_NAMESPACE {
26 namespace mr {
46  public:
52  explicit sam_headroom_memory_resource(std::size_t headroom) : system_mr_{}, headroom_{headroom} {}
53 
55  ~sam_headroom_memory_resource() override = default;
60 
61  private:
73  void* do_allocate(std::size_t bytes, [[maybe_unused]] cuda_stream_view stream) override
74  {
75  void* pointer = system_mr_.allocate_async(bytes, rmm::CUDA_ALLOCATION_ALIGNMENT, stream);
76 
77  auto const free = rmm::available_device_memory().first;
78  auto const allocatable = free > headroom_ ? free - headroom_ : 0UL;
79  auto const gpu_portion =
80  rmm::align_down(std::min(allocatable, bytes), rmm::CUDA_ALLOCATION_ALIGNMENT);
81  auto const cpu_portion = bytes - gpu_portion;
82  if (gpu_portion != 0) {
83  RMM_CUDA_TRY(cudaMemAdvise(pointer,
84  gpu_portion,
85  cudaMemAdviseSetPreferredLocation,
86  rmm::get_current_cuda_device().value()));
87  }
88  if (cpu_portion != 0) {
89  RMM_CUDA_TRY(cudaMemAdvise(static_cast<char*>(pointer) + gpu_portion,
90  cpu_portion,
91  cudaMemAdviseSetPreferredLocation,
92  cudaCpuDeviceId));
93  }
94 
95  return pointer;
96  }
97 
107  void do_deallocate(void* ptr,
108  [[maybe_unused]] std::size_t bytes,
109  [[maybe_unused]] cuda_stream_view stream) override
110  {
111  system_mr_.deallocate_async(ptr, rmm::CUDA_ALLOCATION_ALIGNMENT, stream);
112  }
113 
121  [[nodiscard]] bool do_is_equal(device_memory_resource const& other) const noexcept override
122  {
123  if (this == &other) { return true; }
124  auto cast = dynamic_cast<sam_headroom_memory_resource const*>(&other);
125  if (cast == nullptr) { return false; }
126  return headroom_ == cast->headroom_;
127  }
128 
130  system_memory_resource system_mr_;
132  std::size_t headroom_;
133 }; // end of group
135 } // namespace mr
136 } // 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
Resource that uses system memory resource to allocate memory with a headroom.
Definition: sam_headroom_memory_resource.hpp:45
sam_headroom_memory_resource(std::size_t headroom)
Construct a headroom memory resource.
Definition: sam_headroom_memory_resource.hpp:52
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:120
cuda_device_id get_current_cuda_device()
Returns a cuda_device_id for the current device.
Definition: cuda_device.hpp:96
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