18 #include <rmm/detail/export.hpp>
25 #include <shared_mutex>
28 namespace RMM_NAMESPACE {
57 template <
typename Upstream>
61 std::shared_lock<std::shared_mutex>;
63 std::unique_lock<std::shared_mutex>;
82 peak = std::max(value, peak);
111 peak = std::max(value + val.
peak, peak);
166 return counter_stack_.top().first;
180 return counter_stack_.top().second;
193 auto ret = counter_stack_.top();
208 if (counter_stack_.size() < 2) {
throw std::out_of_range(
"cannot pop the last counter pair"); }
209 auto ret = counter_stack_.top();
210 counter_stack_.pop();
212 counter_stack_.top().first.add_counters_from_tracked_sub_block(ret.first);
213 counter_stack_.top().second.add_counters_from_tracked_sub_block(ret.second);
233 void* ptr = get_upstream_resource().allocate_async(bytes, stream);
237 write_lock_t lock(mtx_);
240 counter_stack_.top().first += bytes;
241 counter_stack_.top().second += 1;
254 void do_deallocate(
void* ptr, std::size_t bytes, cuda_stream_view stream)
override
256 get_upstream_resource().deallocate_async(ptr, bytes, stream);
259 write_lock_t lock(mtx_);
262 counter_stack_.top().first -= bytes;
263 counter_stack_.top().second -= 1;
274 bool do_is_equal(device_memory_resource
const& other)
const noexcept
override
276 if (
this == &other) {
return true; }
277 auto cast =
dynamic_cast<statistics_resource_adaptor<Upstream> const*
>(&other);
278 if (cast ==
nullptr) {
return false; }
279 return get_upstream_resource() == cast->get_upstream_resource();
284 std::stack<std::pair<counter, counter>> counter_stack_{{std::make_pair(counter{}, counter{})}};
285 std::shared_mutex
mutable mtx_;
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 tracks statistics on memory allocations.
Definition: statistics_resource_adaptor.hpp:58
statistics_resource_adaptor(device_async_resource_ref upstream)
Construct a new statistics resource adaptor using upstream to satisfy allocation requests.
Definition: statistics_resource_adaptor.hpp:123
statistics_resource_adaptor(statistics_resource_adaptor &&) noexcept=default
Default move constructor.
std::pair< counter, counter > push_counters()
Push a pair of zero counters on the stack, which becomes the new counters returned by get_bytes_count...
Definition: statistics_resource_adaptor.hpp:190
std::unique_lock< std::shared_mutex > write_lock_t
Type of lock used to synchronize write access.
Definition: statistics_resource_adaptor.hpp:63
statistics_resource_adaptor(Upstream *upstream)
Construct a new statistics resource adaptor using upstream to satisfy allocation requests.
Definition: statistics_resource_adaptor.hpp:133
std::shared_lock< std::shared_mutex > read_lock_t
Type of lock used to synchronize read access.
Definition: statistics_resource_adaptor.hpp:61
std::pair< counter, counter > pop_counters()
Pop a pair of counters from the stack.
Definition: statistics_resource_adaptor.hpp:205
counter get_allocations_counter() const noexcept
Returns a counter struct for this adaptor containing the current, peak, and total number of allocatio...
Definition: statistics_resource_adaptor.hpp:176
counter get_bytes_counter() const noexcept
Returns a counter struct for this adaptor containing the current, peak, and total number of allocated...
Definition: statistics_resource_adaptor.hpp:162
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
Management of per-device device_memory_resources.
Utility struct for counting the current, peak, and total value of a number.
Definition: statistics_resource_adaptor.hpp:67
counter & operator-=(int64_t val)
Subtract val from the current value and update the peak value if necessary.
Definition: statistics_resource_adaptor.hpp:92
int64_t value
Current value.
Definition: statistics_resource_adaptor.hpp:68
int64_t peak
Max value of value
Definition: statistics_resource_adaptor.hpp:69
void add_counters_from_tracked_sub_block(const counter &val)
Add val to the current value and update the peak value if necessary.
Definition: statistics_resource_adaptor.hpp:109
counter & operator+=(int64_t val)
Add val to the current value and update the peak value if necessary.
Definition: statistics_resource_adaptor.hpp:78
int64_t total
Sum of all added values.
Definition: statistics_resource_adaptor.hpp:70