20 #include <rmm/detail/export.hpp>
28 namespace RMM_NAMESPACE {
76 void* do_allocate(std::size_t bytes, [[maybe_unused]]
cuda_stream_view stream)
override
81 auto const allocatable = free > headroom_ ? free - headroom_ : 0UL;
82 auto const gpu_portion =
84 auto const cpu_portion = bytes - gpu_portion;
85 if (gpu_portion != 0) {
86 RMM_CUDA_TRY(cudaMemAdvise(pointer,
88 cudaMemAdviseSetPreferredLocation,
91 if (cpu_portion != 0) {
92 RMM_CUDA_TRY(cudaMemAdvise(
static_cast<char*
>(pointer) + gpu_portion,
94 cudaMemAdviseSetPreferredLocation,
110 void do_deallocate(
void* ptr,
111 [[maybe_unused]] std::size_t bytes,
112 [[maybe_unused]] cuda_stream_view stream)
override
124 [[nodiscard]]
bool do_is_equal(device_memory_resource
const& other)
const noexcept
override
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_;
133 system_memory_resource system_mr_;
135 std::size_t headroom_;
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