19 #include <rmm/detail/error.hpp>
20 #include <rmm/detail/export.hpp>
21 #include <rmm/detail/format.hpp>
29 namespace RMM_NAMESPACE {
48 template <
typename Upstream>
60 std::size_t allocation_limit,
62 : upstream_{upstream},
63 allocation_limit_{allocation_limit},
80 std::size_t allocation_limit,
83 allocation_limit_{allocation_limit},
142 auto const proposed_size =
align_up(bytes, alignment_);
143 auto const old = allocated_bytes_.fetch_add(proposed_size);
144 if (old + proposed_size <= allocation_limit_) {
146 return get_upstream_resource().allocate_async(bytes, stream);
148 allocated_bytes_ -= proposed_size;
153 allocated_bytes_ -= proposed_size;
154 auto const msg = std::string(
"Exceeded memory limit (failed to allocate ") +
155 rmm::detail::format_bytes(bytes) +
")";
166 void do_deallocate(
void* ptr, std::size_t bytes, cuda_stream_view stream)
override
168 std::size_t allocated_size =
align_up(bytes, alignment_);
169 get_upstream_resource().deallocate_async(ptr, bytes, stream);
170 allocated_bytes_ -= allocated_size;
180 [[nodiscard]]
bool do_is_equal(device_memory_resource
const& other)
const noexcept
override
182 if (
this == &other) {
return true; }
183 auto const* cast =
dynamic_cast<limiting_resource_adaptor<Upstream> const*
>(&other);
184 if (cast ==
nullptr) {
return false; }
185 return get_upstream_resource() == cast->get_upstream_resource();
192 std::size_t allocation_limit_;
195 std::atomic<std::size_t> allocated_bytes_;
198 std::size_t alignment_;
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
Resource that uses Upstream to allocate memory and limits the total allocations possible.
Definition: limiting_resource_adaptor.hpp:49
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:59
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:115
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:124
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:79
Exception thrown when RMM runs out of memory.
Definition: error.hpp:87
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.