Class managing buffer resources. More...

#include <resource.hpp>

Public Types
using	MemoryAvailable = std::function< std::int64_t()>
	Callback function to determine available memory. More...

Public Member Functions
	BufferResource (rmm::device_async_resource_ref device_mr, std::unordered_map< MemoryType, MemoryAvailable > memory_available={}, std::optional< Duration > periodic_spill_check=std::chrono::milliseconds{1}, std::shared_ptr< rmm::cuda_stream_pool > stream_pool=std::make_shared< rmm::cuda_stream_pool >(16, rmm::cuda_stream::flags::non_blocking), std::shared_ptr< Statistics > statistics=Statistics::disabled())
	Constructs a buffer resource. More...

rmm::device_async_resource_ref	device_mr () const noexcept
	Get the RMM device memory resource. More...

MemoryAvailable const &	memory_available (MemoryType mem_type) const
	Retrieves the memory availability function for a given memory type. More...

std::size_t	memory_reserved (MemoryType mem_type) const
	Get the current reserved memory of the specified memory type. More...

std::size_t &	memory_reserved (MemoryType mem_type)
	Get a reference to the current reserved memory of the specified memory type. More...

std::pair< MemoryReservation, std::size_t >	reserve (MemoryType mem_type, size_t size, bool allow_overbooking)
	Reserve an amount of the specified memory type. More...

MemoryReservation	reserve_and_spill (MemoryType mem_type, size_t size, bool allow_overbooking)
	Reserve memory and spill if necessary. More...

template<std::ranges::input_range Range>
requires std::convertible_to< std::ranges::range_value_t< Range >, MemoryType > MemoryReservation	reserve_or_fail (size_t size, Range mem_types)
	Make a memory reservation or fail based on the given order of memory types. More...

MemoryReservation	reserve_or_fail (size_t size, MemoryType mem_type)
	Make a memory reservation or fail. More...

std::size_t	release (MemoryReservation &reservation, std::size_t size)
	Consume a portion of the reserved memory. More...

std::unique_ptr< Buffer >	allocate (std::size_t size, rmm::cuda_stream_view stream, MemoryReservation &reservation)
	Allocate a buffer of the specified memory type by the reservation. More...

std::unique_ptr< Buffer >	allocate (rmm::cuda_stream_view stream, MemoryReservation &&reservation)
	Allocate a buffer consuming the entire reservation. More...

std::unique_ptr< Buffer >	move (std::unique_ptr< rmm::device_buffer > data, rmm::cuda_stream_view stream)
	Move device buffer data into a Buffer. More...

std::unique_ptr< Buffer >	move (std::unique_ptr< Buffer > buffer, MemoryReservation &reservation)
	Move a Buffer to the memory type specified by the reservation. More...

std::unique_ptr< rmm::device_buffer >	move_to_device_buffer (std::unique_ptr< Buffer > buffer, MemoryReservation &reservation)
	Move a Buffer to a device buffer. More...

std::unique_ptr< std::vector< uint8_t > >	move_to_host_vector (std::unique_ptr< Buffer > buffer, MemoryReservation &reservation)
	Move a Buffer into a host vector. More...

rmm::cuda_stream_pool const &	stream_pool () const
	Returns the CUDA stream pool used by this buffer resource. More...

SpillManager &	spill_manager ()
	Gets a reference to the spill manager used. More...

std::shared_ptr< Statistics >	statistics ()
	Gets a shared pointer to the statistics associated with this buffer resource. More...

Detailed Description

Class managing buffer resources.

This class handles memory allocation and transfers between different memory types (e.g., host and device). All memory operations in rapidsmpf, such as those performed by the Shuffler, rely on a buffer resource for memory management.

Note: Similar to RMM's memory resource, the BufferResource instance must outlive all allocated buffers and memory reservations.

Definition at line 133 of file resource.hpp.

Member Typedef Documentation

◆ MemoryAvailable

using rapidsmpf::BufferResource::MemoryAvailable = std::function<std::int64_t()>

Callback function to determine available memory.

The function should return the current available memory of a specific type and must be thread-safe iff used by multiple BufferResource instances concurrently.

Warning: Calling any BufferResource instance methods in the function might result in a deadlock. This is because the buffer resource is locked when the function is called.

Definition at line 145 of file resource.hpp.

Constructor & Destructor Documentation

◆ BufferResource()

rapidsmpf::BufferResource::BufferResource	(	rmm::device_async_resource_ref	device_mr,
		std::unordered_map< MemoryType, MemoryAvailable >	memory_available = `{}`,
		std::optional< Duration >	periodic_spill_check = `std::chrono::milliseconds{1}`,
		std::shared_ptr< rmm::cuda_stream_pool >	stream_pool = `std::make_shared< rmm::cuda_stream_pool >(16, rmm::cuda_stream::flags::non_blocking)`,
		std::shared_ptr< Statistics >	statistics = `Statistics::disabled()`
	)

Constructs a buffer resource.

Parameters

device_mr	Reference to the RMM device memory resource used for device allocations.
memory_available	Optional memory availability functions mapping memory types to available memory checkers. Memory types without availability functions are assumed to have unlimited memory.
periodic_spill_check	Enable periodic spill checks. A dedicated thread continuously checks and perform spilling based on the memory availability functions. The value of `periodic_spill_check` is used as the pause between checks. If `std::nullopt`, no periodic spill check is performed.
stream_pool	Pool of CUDA streams. Used throughout RapidsMPF for operations that do not take an explicit CUDA stream.
statistics	The statistics instance to use (disabled by default).

Member Function Documentation

◆ allocate() [1/2]

std::unique_ptr<Buffer> rapidsmpf::BufferResource::allocate	(	rmm::cuda_stream_view	stream,
		MemoryReservation &&	reservation
	)

Allocate a buffer consuming the entire reservation.

This overload allocates a buffer that matches the full size and memory type of the provided reservation. The reservation is consumed by the call.

Parameters

stream	CUDA stream to use for device allocations.
reservation	The memory reservation to consume for the allocation.

Returns: A unique pointer to the allocated Buffer.

◆ allocate() [2/2]

std::unique_ptr<Buffer> rapidsmpf::BufferResource::allocate	(	std::size_t	size,
		rmm::cuda_stream_view	stream,
		MemoryReservation &	reservation
	)

Allocate a buffer of the specified memory type by the reservation.

Parameters

size	The size of the buffer in bytes.
stream	CUDA stream to use for device allocations.
reservation	The reservation to use for memory allocations.

Returns: A unique pointer to the allocated Buffer.

Exceptions

std::invalid_argument	if the memory type does not match the reservation.
std::overflow_error	if `size` exceeds the size of the reservation.

◆ device_mr()

rmm::device_async_resource_ref rapidsmpf::BufferResource::device_mr ( ) const

inlinenoexcept

Get the RMM device memory resource.

Returns: Reference to the RMM resource used for device allocations.

Definition at line 179 of file resource.hpp.

◆ memory_available()

MemoryAvailable const& rapidsmpf::BufferResource::memory_available ( MemoryType mem_type ) const

inline

Retrieves the memory availability function for a given memory type.

This function returns the callback function used to determine the available memory for the specified memory type.

Parameters

mem_type The type of memory whose availability function is requested.

Returns: Reference to the memory availability function associated with mem_type.

Definition at line 192 of file resource.hpp.

◆ memory_reserved() [1/2]

std::size_t& rapidsmpf::BufferResource::memory_reserved ( MemoryType mem_type )

inline

Get a reference to the current reserved memory of the specified memory type.

Parameters

mem_type The target memory type.

Returns: A reference to the memory reserved.

Definition at line 212 of file resource.hpp.

◆ memory_reserved() [2/2]

std::size_t rapidsmpf::BufferResource::memory_reserved ( MemoryType mem_type ) const

inline

Get the current reserved memory of the specified memory type.

Parameters

mem_type The target memory type.

Returns: The memory reserved.

Definition at line 202 of file resource.hpp.

◆ move() [1/2]

std::unique_ptr<Buffer> rapidsmpf::BufferResource::move	(	std::unique_ptr< Buffer >	buffer,
		MemoryReservation &	reservation
	)

Move a Buffer to the memory type specified by the reservation.

If the Buffer already resides in the target memory type, a cheap move is performed. Otherwise, the Buffer is copied to the target memory using its own CUDA stream.

Parameters

buffer	Buffer to move.
reservation	Memory reservation used if a copy is required.

Returns: Unique pointer to the resulting Buffer.

Exceptions

std::overflow_error If the allocation size exceeds the reservation.

◆ move() [2/2]

std::unique_ptr<Buffer> rapidsmpf::BufferResource::move	(	std::unique_ptr< rmm::device_buffer >	data,
		rmm::cuda_stream_view	stream
	)

Move device buffer data into a Buffer.

This operation is cheap; no copy is performed. The resulting Buffer resides in device memory.

If stream differs from the device buffer's current stream:

stream is synchronized with the device buffer's current stream, and
the device buffer's current stream is updated to stream.

Parameters

data	Unique pointer to the device buffer.
stream	CUDA stream associated with the new Buffer. Use or synchronize with this stream when operating on the Buffer.

Returns: Unique pointer to the resulting Buffer.

◆ move_to_device_buffer()

std::unique_ptr<rmm::device_buffer> rapidsmpf::BufferResource::move_to_device_buffer	(	std::unique_ptr< Buffer >	buffer,
		MemoryReservation &	reservation
	)

Move a Buffer to a device buffer.

If the Buffer already resides in device memory, a cheap move is performed. Otherwise, the Buffer is copied to device memory using its own CUDA stream.

Parameters

buffer	The buffer to move.
reservation	Memory reservation used if a copy is required.

Returns: A unique pointer to the resulting device buffer.

Exceptions

std::invalid_argument	If the reservation's memory type isn't device memory.
std::overflow_error	if the memory requirement exceeds the reservation.

◆ move_to_host_vector()

std::unique_ptr<std::vector<uint8_t> > rapidsmpf::BufferResource::move_to_host_vector	(	std::unique_ptr< Buffer >	buffer,
		MemoryReservation &	reservation
	)

Move a Buffer into a host vector.

If the Buffer already resides in host memory, a cheap move is performed. Otherwise, the Buffer is copied to host memory using its own CUDA stream.

Parameters

buffer	Buffer to move.
reservation	Memory reservation used if a copy is required.

Returns: Unique pointer to the resulting host vector.

Exceptions

std::invalid_argument	If the reservation's memory type isn't host memory.
std::overflow_error	If the allocation size exceeds the reservation.

◆ release()

std::size_t rapidsmpf::BufferResource::release	(	MemoryReservation &	reservation,
		std::size_t	size
	)

Consume a portion of the reserved memory.

Reduces the remaining size of the reserved memory by the specified amount.

Parameters

reservation	The reservation to release.
size	The size to consume in bytes.

Returns: The remaining size of the reserved memory after consumption.

Exceptions

std::overflow_error if the released size exceeds the size of the reservation.

◆ reserve()

std::pair<MemoryReservation, std::size_t> rapidsmpf::BufferResource::reserve	(	MemoryType	mem_type,
		size_t	size,
		bool	allow_overbooking
	)

Reserve an amount of the specified memory type.

Creates a new reservation of the specified size and type to inform about upcoming buffer allocations.

If overbooking is allowed, a reservation of size is returned even when the amount of memory isn't available. In this case, the caller must promise to free buffers corresponding to (at least) the amount of overbooking before using the reservation.

If overbooking isn't allowed, a reservation of size zero is returned on failure.

Parameters

mem_type	The target memory type.
size	The number of bytes to reserve.
allow_overbooking	Whether overbooking is allowed.

Returns: A pair containing the reservation and the amount of overbooking. On success the size of the reservation always equals size and on failure the size always equals zero (a zero-sized reservation never fails).

◆ reserve_and_spill()

MemoryReservation rapidsmpf::BufferResource::reserve_and_spill	(	MemoryType	mem_type,
		size_t	size,
		bool	allow_overbooking
	)

Reserve memory and spill if necessary.

Parameters

mem_type	The memory type to reserve.
size	The size of the memory to reserve.
allow_overbooking	Whether to allow overbooking. If false, ensures enough memory is freed to satisfy the reservation; otherwise, allows overbooking even if spilling was insufficient.

Returns: The memory reservation.

Exceptions

std::overflow_error if allow_overbooking is false and the buffer resource cannot reserve and spill enough memory.

◆ reserve_or_fail() [1/2]

MemoryReservation rapidsmpf::BufferResource::reserve_or_fail	(	size_t	size,
		MemoryType	mem_type
	)

inline

Make a memory reservation or fail.

Parameters

size	The size of the buffer to allocate.
mem_type	The memory type to try to reserve memory from.

Returns: A memory reservation.

Exceptions

std::runtime_error if no memory reservation was made.

Definition at line 286 of file resource.hpp.

◆ reserve_or_fail() [2/2]

template<std::ranges::input_range Range>

requires std::convertible_to<std::ranges::range_value_t<Range>, MemoryType> MemoryReservation rapidsmpf::BufferResource::reserve_or_fail	(	size_t	size,
		Range	mem_types
	)

inline

Make a memory reservation or fail based on the given order of memory types.

Parameters

size	The size of the buffer to allocate.
mem_types	Range of memory types to try to reserve memory from. If not provided, all memory types will be tried in the order they appear in `MEMORY_TYPES`.

Returns: A memory reservation.

Exceptions

std::runtime_error if no memory reservation was made.

Definition at line 267 of file resource.hpp.

◆ spill_manager()

SpillManager& rapidsmpf::BufferResource::spill_manager ( )

Gets a reference to the spill manager used.

Returns: Reference to the SpillManager instance.

◆ statistics()

std::shared_ptr<Statistics> rapidsmpf::BufferResource::statistics ( )

Gets a shared pointer to the statistics associated with this buffer resource.

Returns: Shared pointer the Statistics instance.

◆ stream_pool()

rmm::cuda_stream_pool const& rapidsmpf::BufferResource::stream_pool ( ) const

Returns the CUDA stream pool used by this buffer resource.

Use this pool for operations that do not take an explicit CUDA stream.

Returns: Reference to the underlying CUDA stream pool.

The documentation for this class was generated from the following file:

/__w/rapidsmpf/rapidsmpf/cpp/include/rapidsmpf/buffer/resource.hpp

Public Types

Public Member Functions

Detailed Description

Member Typedef Documentation

◆ MemoryAvailable

Constructor & Destructor Documentation

◆ BufferResource()

Member Function Documentation

◆ allocate() [1/2]

◆ allocate() [2/2]

◆ device_mr()

◆ memory_available()

◆ memory_reserved() [1/2]

◆ memory_reserved() [2/2]

◆ move() [1/2]

◆ move() [2/2]

◆ move_to_device_buffer()

◆ move_to_host_vector()

◆ release()

◆ reserve()

◆ reserve_and_spill()

◆ reserve_or_fail() [1/2]

◆ reserve_or_fail() [2/2]

◆ spill_manager()

◆ statistics()

◆ stream_pool()