19 #include <rmm/detail/error.hpp>
20 #include <rmm/detail/export.hpp>
27 namespace RMM_NAMESPACE {
46 template <
typename Upstream>
58 std::size_t allocation_limit,
60 : upstream_{upstream},
61 allocation_limit_{allocation_limit},
78 std::size_t allocation_limit,
81 allocation_limit_{allocation_limit},
140 auto const proposed_size =
align_up(bytes, alignment_);
141 auto const old = allocated_bytes_.fetch_add(proposed_size);
142 if (old + proposed_size <= allocation_limit_) {
144 return get_upstream_resource().allocate_async(bytes, stream);
146 allocated_bytes_ -= proposed_size;
151 allocated_bytes_ -= proposed_size;
162 void do_deallocate(
void* ptr, std::size_t bytes, cuda_stream_view stream)
override
164 std::size_t allocated_size =
align_up(bytes, alignment_);
165 get_upstream_resource().deallocate_async(ptr, bytes, stream);
166 allocated_bytes_ -= allocated_size;
176 [[nodiscard]]
bool do_is_equal(device_memory_resource
const& other)
const noexcept
override
178 if (
this == &other) {
return true; }
179 auto const* cast =
dynamic_cast<limiting_resource_adaptor<Upstream> const*
>(&other);
180 if (cast ==
nullptr) {
return false; }
181 return get_upstream_resource() == cast->get_upstream_resource();
188 std::size_t allocation_limit_;
191 std::atomic<std::size_t> allocated_bytes_;
194 std::size_t alignment_;
206 template <
typename Upstream>
208 "make_limiting_adaptor is deprecated in RMM 24.10. Use the limiting_resource_adaptor constructor "
211 std::size_t allocation_limit)
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
Resource that uses Upstream to allocate memory and limits the total allocations possible.
Definition: limiting_resource_adaptor.hpp:47
limiting_resource_adaptor(device_async_resource_ref upstream, std::size_t allocation_limit, std::size_t alignment=CUDA_ALLOCATION_ALIGNMENT)
Construct a new limiting resource adaptor using upstream to satisfy allocation requests and limiting ...
Definition: limiting_resource_adaptor.hpp:57
limiting_resource_adaptor(limiting_resource_adaptor &&) noexcept=default
Default move constructor.
std::size_t get_allocated_bytes() const
Query the number of bytes that have been allocated. Note that this can not be used to know how large ...
Definition: limiting_resource_adaptor.hpp:113
std::size_t get_allocation_limit() const
Query the maximum number of bytes that this allocator is allowed to allocate. This is the limit on th...
Definition: limiting_resource_adaptor.hpp:122
limiting_resource_adaptor(Upstream *upstream, std::size_t allocation_limit, std::size_t alignment=CUDA_ALLOCATION_ALIGNMENT)
Construct a new limiting resource adaptor using upstream to satisfy allocation requests and limiting ...
Definition: limiting_resource_adaptor.hpp:77
Exception thrown when RMM runs out of memory.
Definition: error.hpp:87
limiting_resource_adaptor< Upstream > make_limiting_adaptor(Upstream *upstream, std::size_t allocation_limit)
Convenience factory to return a limiting_resource_adaptor around the upstream resource upstream.
Definition: limiting_resource_adaptor.hpp:210
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
Management of per-device device_memory_resources.