tracking_resource_adaptor.hpp
Go to the documentation of this file.
1 /*
2  * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION.
3  * SPDX-License-Identifier: Apache-2.0
4  */
5 #pragma once
6 
7 #include <rmm/detail/export.hpp>
8 #include <rmm/mr/detail/tracking_resource_adaptor_impl.hpp>
9 #include <rmm/resource_ref.hpp>
10 
11 #include <cuda/memory_resource>
12 
13 #include <cstddef>
14 #include <map>
15 #include <memory>
16 #include <string>
17 
18 namespace RMM_NAMESPACE {
19 namespace mr {
34 class RMM_EXPORT tracking_resource_adaptor
35  : public cuda::mr::shared_resource<detail::tracking_resource_adaptor_impl> {
36  using shared_base = cuda::mr::shared_resource<detail::tracking_resource_adaptor_impl>;
37 
38  public:
40  using allocation_info = detail::tracking_resource_adaptor_impl::allocation_info;
42  using read_lock_t = detail::tracking_resource_adaptor_impl::read_lock_t;
44  using write_lock_t = detail::tracking_resource_adaptor_impl::write_lock_t;
45 
49  RMM_CONSTEXPR_FRIEND void get_property(tracking_resource_adaptor const&,
50  cuda::mr::device_accessible) noexcept
51  {
52  }
53 
60  tracking_resource_adaptor(cuda::mr::any_resource<cuda::mr::device_accessible> upstream,
61  bool capture_stacks = false);
62 
63  ~tracking_resource_adaptor() = default;
64 
68  [[nodiscard]] device_async_resource_ref get_upstream_resource() const noexcept;
69 
75  [[nodiscard]] std::map<void*, allocation_info> const& get_outstanding_allocations()
76  const noexcept;
77 
83  [[nodiscard]] std::size_t get_allocated_bytes() const noexcept;
84 
90  [[nodiscard]] std::string get_outstanding_allocations_str() const;
91 
95  void log_outstanding_allocations() const;
96 };
97 
98 static_assert(cuda::mr::resource_with<tracking_resource_adaptor, cuda::mr::device_accessible>,
99  "tracking_resource_adaptor does not satisfy the cuda::mr::resource concept");
100  // end of group
102 } // namespace mr
103 } // namespace RMM_NAMESPACE
Resource that uses an upstream resource to allocate memory and tracks allocations.
Definition: tracking_resource_adaptor.hpp:35
detail::tracking_resource_adaptor_impl::allocation_info allocation_info
Allocation info type (pointer, size, optional stack trace).
Definition: tracking_resource_adaptor.hpp:40
device_async_resource_ref get_upstream_resource() const noexcept
rmm::device_async_resource_ref to the upstream resource
detail::tracking_resource_adaptor_impl::read_lock_t read_lock_t
Shared-reader lock type used to protect the allocations map.
Definition: tracking_resource_adaptor.hpp:42
friend void get_property(tracking_resource_adaptor const &, cuda::mr::device_accessible) noexcept
Enables the cuda::mr::device_accessible property.
Definition: tracking_resource_adaptor.hpp:49
detail::tracking_resource_adaptor_impl::write_lock_t write_lock_t
Exclusive-writer lock type used to protect the allocations map.
Definition: tracking_resource_adaptor.hpp:44
tracking_resource_adaptor(cuda::mr::any_resource< cuda::mr::device_accessible > upstream, bool capture_stacks=false)
Construct a tracking resource adaptor using upstream to satisfy allocation requests.
cuda::mr::resource_ref< cuda::mr::device_accessible > device_async_resource_ref
Alias for a cuda::mr::resource_ref with the property cuda::mr::device_accessible.
Definition: resource_ref.hpp:30