A RMM memory resource adaptor tailored to RapidsMPF. More...
#include <rmm_resource_adaptor.hpp>
Public Member Functions | |
| RmmResourceAdaptor (cuda::mr::any_resource< cuda::mr::device_accessible > primary_mr, std::optional< cuda::mr::any_resource< cuda::mr::device_accessible >> fallback_mr=std::nullopt) | |
| Construct with specified primary and optional fallback memory resource. More... | |
| rmm::device_async_resource_ref | get_upstream_resource () const noexcept |
| Get a reference to the primary upstream resource. More... | |
| std::optional< rmm::device_async_resource_ref > | get_fallback_resource () const noexcept |
| Get a reference to the fallback upstream resource. More... | |
| ScopedMemoryRecord | get_main_record () const |
| Returns a copy of the main memory record. More... | |
| std::int64_t | current_allocated () const noexcept |
| Get the total current allocated memory from both primary and fallback. More... | |
| void | begin_scoped_memory_record () |
| Begin recording a new scoped memory usage record for the current thread. More... | |
| ScopedMemoryRecord | end_scoped_memory_record () |
| End the current scoped memory record and return it. More... | |
Friends | |
| void | get_property (RmmResourceAdaptor const &, cuda::mr::device_accessible) noexcept |
| Tag this resource as device-accessible for the CCCL concept. | |
A RMM memory resource adaptor tailored to RapidsMPF.
This adaptor implements:
This class is copyable and shares ownership of its internal state via cuda::mr::shared_resource.
Definition at line 31 of file rmm_resource_adaptor.hpp.
| rapidsmpf::RmmResourceAdaptor::RmmResourceAdaptor | ( | cuda::mr::any_resource< cuda::mr::device_accessible > | primary_mr, |
| std::optional< cuda::mr::any_resource< cuda::mr::device_accessible >> | fallback_mr = std::nullopt |
||
| ) |
Construct with specified primary and optional fallback memory resource.
| primary_mr | The primary memory resource. |
| fallback_mr | Optional fallback memory resource. |
| void rapidsmpf::RmmResourceAdaptor::begin_scoped_memory_record | ( | ) |
Begin recording a new scoped memory usage record for the current thread.
This method pushes a new empty ScopedMemoryRecord onto the thread-local record stack, allowing for nested memory tracking scopes.
Must be paired with a matching call to end_scoped_memory_record().
|
noexcept |
Get the total current allocated memory from both primary and fallback.
| ScopedMemoryRecord rapidsmpf::RmmResourceAdaptor::end_scoped_memory_record | ( | ) |
End the current scoped memory record and return it.
Pops the top ScopedMemoryRecord from the thread-local stack and returns it. If this scope was nested within another (i.e. if begin_scoped_memory_record() was called multiple times in a row), the returned scope is automatically added as a subscope to the next scope remaining on the stack.
This allows nesting of scoped memory tracking, where each scope can contain one or more subscopes. When analyzing or reporting memory statistics, the memory usage of each scope can be calculated inclusive of its subscopes. This behavior mimics standard hierarchical memory profilers, where the total memory attributed to a scope includes all allocations made within it, plus those made in its nested regions.
| std::out_of_range | if called without a matching begin_scoped_memory_record(). |
|
noexcept |
Get a reference to the fallback upstream resource.
This resource is used if the primary resource throws rmm::out_of_memory.
| ScopedMemoryRecord rapidsmpf::RmmResourceAdaptor::get_main_record | ( | ) | const |
Returns a copy of the main memory record.
The main record tracks memory statistics for the lifetime of the resource.
|
noexcept |
Get a reference to the primary upstream resource.