A RMM memory resource adaptor tailored to RapidsMPF. More...

#include <rmm_resource_adaptor.hpp>

Inheritance diagram for rapidsmpf::RmmResourceAdaptor:

Public Member Functions
	RmmResourceAdaptor (rmm::device_async_resource_ref primary_mr, std::optional< rmm::device_async_resource_ref > 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...

Public Member Functions inherited from rmm::mr::device_memory_resource
	device_memory_resource (device_memory_resource const &)=default

	device_memory_resource (device_memory_resource &&) noexcept=default

device_memory_resource &	operator= (device_memory_resource const &)=default

device_memory_resource &	operator= (device_memory_resource &&) noexcept=default

void *	allocate_sync (std::size_t bytes, std::size_t alignment=rmm::CUDA_ALLOCATION_ALIGNMENT)

void	deallocate_sync (void *ptr, std::size_t bytes,[[maybe_unused]] std::size_t alignment=rmm::CUDA_ALLOCATION_ALIGNMENT) noexcept

void *	allocate (cuda_stream_view stream, std::size_t bytes, std::size_t alignment=rmm::CUDA_ALLOCATION_ALIGNMENT)

void	deallocate (cuda_stream_view stream, void *ptr, std::size_t bytes,[[maybe_unused]] std::size_t alignment=rmm::CUDA_ALLOCATION_ALIGNMENT) noexcept

bool	is_equal (device_memory_resource const &other) const noexcept

bool	operator== (device_memory_resource const &other) const noexcept

bool	operator!= (device_memory_resource const &other) const noexcept

Detailed Description

A RMM memory resource adaptor tailored to RapidsMPF.

This adaptor implements:

Memory usage tracking.
Fallback memory resource support upon out-of-memory in the primary resource.

Definition at line 193 of file rmm_resource_adaptor.hpp.

Constructor & Destructor Documentation

◆ RmmResourceAdaptor()

rapidsmpf::RmmResourceAdaptor::RmmResourceAdaptor	(	rmm::device_async_resource_ref	primary_mr,
		std::optional< rmm::device_async_resource_ref >	fallback_mr = `std::nullopt`
	)

inline

Construct with specified primary and optional fallback memory resource.

Parameters

primary_mr	The primary memory resource.
fallback_mr	Optional fallback memory resource.

Definition at line 201 of file rmm_resource_adaptor.hpp.

Member Function Documentation

◆ begin_scoped_memory_record()

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().

See also: end_scoped_memory_record()

◆ current_allocated()

std::int64_t rapidsmpf::RmmResourceAdaptor::current_allocated ( ) const

noexcept

Get the total current allocated memory from both primary and fallback.

Returns: Total number of currently allocated bytes.

◆ end_scoped_memory_record()

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.

Returns: The scope that was just ended.

Exceptions

std::out_of_range if called without a matching begin_scoped_memory_record().

See also: begin_scoped_memory_record()

◆ get_fallback_resource()

std::optional<rmm::device_async_resource_ref> rapidsmpf::RmmResourceAdaptor::get_fallback_resource ( ) const

inlinenoexcept

Get a reference to the fallback upstream resource.

This resource is used if the primary resource throws rmm::out_of_memory.

Returns: Optional reference to the fallback RMM memory resource.

Definition at line 227 of file rmm_resource_adaptor.hpp.

◆ get_main_record()

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.

Returns: A copy of the current main memory record.

◆ get_upstream_resource()

rmm::device_async_resource_ref rapidsmpf::RmmResourceAdaptor::get_upstream_resource ( ) const