Class managing buffer resources. More...

#include <buffer_resource.hpp>

Inheritance diagram for rapidsmpf::BufferResource:

Public Member Functions
	BufferResource (BufferResource const &)=delete
	`BufferResource` is non-copyable, it is owned by `std::shared_ptr`.

	BufferResource (BufferResource &&)=delete
	`BufferResource` is non-movable, it is owned by `std::shared_ptr`.

BufferResource &	operator= (BufferResource const &)=delete
	`BufferResource` is non-copyable. More...

BufferResource &	operator= (BufferResource &&)=delete
	`BufferResource` is non-movable. More...

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

RmmResourceAdaptor &	device_mr_adaptor () noexcept
	Access the concrete device memory resource adaptor. More...

rmm::host_async_resource_ref	host_mr () noexcept
	Get the RMM host memory resource. More...

rmm::host_device_async_resource_ref	pinned_mr ()
	Get the RMM pinned host memory resource. More...

std::optional< PinnedMemoryResource >	try_pinned_mr () const
	Get the pinned host memory resource if available. More...

std::int64_t	memory_available (MemoryType mem_type) const noexcept
	Returns the currently available memory for a given memory type, in bytes. More...

void	set_memory_limit (MemoryType mem_type, std::int64_t limit) noexcept
	Updates the memory limit for a given memory type at runtime. More...

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

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

MemoryReservation	reserve_device_memory_and_spill (std::size_t size, AllowOverbooking allow_overbooking)
	Reserve device 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 (std::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 (std::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 >	make_buffer (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 >	make_buffer (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 or pinned host 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< HostBuffer >	move_to_host_buffer (std::unique_ptr< Buffer > buffer, MemoryReservation &reservation)
	Move a Buffer into a host buffer. More...

std::shared_ptr< 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 () const noexcept
	Gets a shared pointer to the statistics associated with this buffer resource. More...

Static Public Member Functions
static std::shared_ptr< BufferResource >	create (cuda::mr::any_resource< cuda::mr::device_accessible > device_mr, std::optional< PinnedPoolProperties > pinned_pool_properties=PinnedMemoryDisabled, std::unordered_map< MemoryType, std::int64_t > memory_limits={}, 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())
	Construct a `BufferResource` managed by `std::shared_ptr`. More...

static std::shared_ptr< BufferResource >	from_options (cuda::mr::any_resource< cuda::mr::device_accessible > mr, config::Options options, std::shared_ptr< Statistics > statistics=Statistics::disabled())
	Construct a BufferResource from configuration options. 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: BufferResource instances must be constructed through create() or from_options(), both of which return a std::shared_ptr<BufferResource>. Direct construction is disabled.; Allocation tracking only applies to allocations routed through this BufferResource. The constructor wraps the supplied device memory resource in an internal adaptor that records all allocations and deallocations; that adaptor is exposed via device_mr().

Allocations made through the original, unwrapped memory resource bypass this tracking and are therefore invisible to memory-limit accounting and statistics.

To ensure all CUDA allocations count against the BufferResource budget, use br->device_mr() everywhere instead of the underlying memory resource passed to the constructor.

Tracking allocations made outside BufferResource, for example allocations performed before construction or through code paths that use a raw memory resource directly, is a separate design concern and is not handled by this class.

Definition at line 76 of file buffer_resource.hpp.

Member Function Documentation

◆ create()

static std::shared_ptr<BufferResource> rapidsmpf::BufferResource::create	(	cuda::mr::any_resource< cuda::mr::device_accessible >	device_mr,
		std::optional< PinnedPoolProperties >	pinned_pool_properties = `PinnedMemoryDisabled`,
		std::unordered_map< MemoryType, std::int64_t >	memory_limits = `{}`,
		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()`
	)

static

Construct a BufferResource managed by std::shared_ptr.

Available memory is computed per MemoryType as limit - allocated.

Device and pinned-host allocations routed through this BufferResource are tracked automatically. Host memory allocations are not tracked and therefore always report the configured limit as available memory.

If pinned-host memory is disabled, available pinned-host memory is always reported as zero regardless of the configured limit.

Parameters

device_mr	Device memory resource used for device allocations. To ensure allocations are tracked for memory-limit accounting and statistics, use `BufferResource::device_mr()` instead of the original memory resource after construction.
pinned_pool_properties	Configuration for the pinned host memory pool used for `MemoryType::PINNED_HOST` allocations, or `PinnedMemoryDisabled` to disable pinned allocations. The pinned resource is constructed internally and owned by the `BufferResource`. When a value is provided, pinned host memory must be supported on the system (see `is_pinned_memory_resources_supported()`); otherwise a `std::runtime_error` is thrown.
memory_limits	Maximum allocation limits in bytes per `MemoryType`. Missing entries are treated as unlimited.
periodic_spill_check	Interval between periodic spill checks. `std::nullopt` disables the dedicated spill-check thread.
stream_pool	CUDA stream pool used for operations that do not take an explicit CUDA stream.
statistics	Statistics instance used for runtime metrics.

Returns: A newly constructed BufferResource owned by std::shared_ptr.

Exceptions

std::runtime_error if pinned_pool_properties has a value but pinned host memory is not supported on this system.

◆ device_mr()

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

noexcept

Get the device memory resource.

Returns: rmm::device_async_resource_ref to the device memory resource.

CCCL's lifetime semantic

The returned rmm::device_async_resource_ref is a non-owning cuda::mr::resource_ref, so callers must take care to avoid use-after-free issues.

When working directly with the returned reference, the caller must ensure that this BufferResource remains alive for the full duration of that use:

auto br = BufferResource::create(...);
auto mr = br->device_mr();
mr.allocate_async(...);  // direct use through a non-owning ref
br.reset();              // do not destroy `br` while `mr` is in use

To store the resource beyond the immediate call, promote the ref to an owning cuda::mr::any_resource:

auto br = BufferResource::create(...);
cuda::mr::any_resource<cuda::mr::device_accessible> mr = br->device_mr();
br.reset();       // safe: `mr` keeps the BufferResource alive
mr.allocate(...); // safe

In the common case, no explicit promotion is needed because RMM containers that store a memory resource do this internally:

auto br = BufferResource::create(...);
rmm::device_buffer buf{1024, stream, br->device_mr()};
br.reset();  // safe: `buf` keeps the BufferResource alive internally

Note: Device memory resource provided to the constructor is wrapped in an RmmResourceAdaptor for allocation tracking, and concretely the returned resource_ref points to that adaptor. See device_mr_adaptor() for a more convenient way to access the adaptor.

◆ device_mr_adaptor()

RmmResourceAdaptor& rapidsmpf::BufferResource::device_mr_adaptor ( )

noexcept

Access the concrete device memory resource adaptor.

BufferResource wraps the device memory resource in an internal RmmResourceAdaptor for allocation tracking. This exposes that adaptor directly, e.g. to query allocation statistics via get_main_record() or current_allocated().

Returns: Reference to the internal device RmmResourceAdaptor. The reference is valid for as long as this BufferResource is alive.

Note: To ensure that the allocations are properly tracked, use device_mr() or device_mr_adaptor() instead of the original memory resource passed to the constructor.

◆ from_options()

static std::shared_ptr<BufferResource> rapidsmpf::BufferResource::from_options	(	cuda::mr::any_resource< cuda::mr::device_accessible >	mr,
		config::Options	options,
		std::shared_ptr< Statistics >	statistics = `Statistics::disabled()`
	)

static

Construct a BufferResource from configuration options.

This factory method creates a BufferResource using configuration options to initialize all components. The supplied device memory resource is wrapped in an internal RmmResourceAdaptor for allocation tracking.

Parameters

mr	A device-accessible RMM memory resource.
options	Configuration options.
statistics	The statistics instance to use (disabled by default).

Returns: A shared pointer to a BufferResource instance configured according to the options.

◆ host_mr()

rmm::host_async_resource_ref rapidsmpf::BufferResource::host_mr ( )

noexcept

Get the RMM host memory resource.

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

Note: Lifetime semantics are identical to device_mr(). See its @par CCCL lifetime semantics section for details. In brief, the returned resource_ref is non-owning. Promote it to a any_host_resource to extend the BufferResource lifetime.

◆ make_buffer() [1/2]

std::unique_ptr<Buffer> rapidsmpf::BufferResource::make_buffer	(	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.

◆ make_buffer() [2/2]

std::unique_ptr<Buffer> rapidsmpf::BufferResource::make_buffer	(	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.
rapidsmpf::reservation_error	if `size` exceeds the size of the reservation.

◆ memory_available()

std::int64_t rapidsmpf::BufferResource::memory_available ( MemoryType mem_type ) const

noexcept

Returns the currently available memory for a given memory type, in bytes.

Computed as limit - allocated, where allocated is reported by the memory type's allocation counter (see the constructor documentation for how each memory type is tracked). The value may be negative when allocations exceed the configured limit.

Parameters

mem_type The memory type to query.

Returns: The available memory in bytes.

◆ memory_reserved()

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 290 of file buffer_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

rapidsmpf::reservation_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 or pinned host buffer data into a Buffer.

This operation is cheap; no copy is performed.

The resulting Buffer's memory type is inferred from data's memory resource: if the resource is host-accessible (e.g. pinned host memory), the Buffer is created with MemoryType::PINNED_HOST; otherwise it is created with MemoryType::DEVICE.

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 or pinned host 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.
rapidsmpf::reservation_error	if the memory requirement exceeds the reservation.

◆ move_to_host_buffer()

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

Move a Buffer into a host buffer.

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 buffer.

Exceptions

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

◆ operator=() [1/2]

BufferResource& rapidsmpf::BufferResource::operator= ( BufferResource && )

delete

BufferResource is non-movable.

Returns: Reference to this.

◆ operator=() [2/2]

BufferResource& rapidsmpf::BufferResource::operator= ( BufferResource const & )

delete

BufferResource is non-copyable.

Returns: Reference to this.

◆ pinned_mr()

rmm::host_device_async_resource_ref rapidsmpf::BufferResource::pinned_mr ( )

Get the RMM pinned host memory resource.

Exceptions

std::invalid_argument if no pinned memory resource is available.

Returns: Reference to the RMM resource used for pinned host allocations.

Note: Lifetime semantics are identical to device_mr(). See its @par CCCL lifetime semantics section for details. In brief, the returned resource_ref is non-owning. Promote it to a any_host_device_resource to extend the BufferResource lifetime.

◆ 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

rapidsmpf::reservation_error if the released size exceeds the size of the reservation.

◆ reserve()

std::pair<MemoryReservation, std::size_t> rapidsmpf::BufferResource::reserve	(	MemoryType	mem_type,
		std::size_t	size,
		AllowOverbooking	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).

Exceptions

std::invalid_argument if the memory type is MemoryType::PINNED_HOST and the pinned memory resource is not available.

◆ reserve_device_memory_and_spill()

MemoryReservation rapidsmpf::BufferResource::reserve_device_memory_and_spill	(	std::size_t	size,
		AllowOverbooking	allow_overbooking
	)

Reserve device memory and spill if necessary.

Attempts to reserve the requested amount of device memory. If insufficient memory is available, spilling is triggered to free up space. When overbooking is allowed, the reservation may succeed even if spilling was not sufficient to fully satisfy the request.

Parameters

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

rapidsmpf::reservation_error if allow_overbooking is false and the buffer resource cannot reserve and spill enough device memory.

◆ reserve_or_fail() [1/2]

MemoryReservation rapidsmpf::BufferResource::reserve_or_fail	(	std::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 381 of file buffer_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	(	std::size_t	size,
		Range	mem_types
	)

inline

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

The function attempts to reserve memory by iterating over mem_types in the given order of preference. For each memory type, it requests a reservation without overbooking. If no memory type can satisfy the request, the function throws.

Parameters

size	The size of the buffer to allocate.
mem_types	Range of memory types to try to reserve memory from.

Returns: A memory reservation.

Exceptions

std::runtime_error if no memory reservation was made.

Definition at line 356 of file buffer_resource.hpp.

◆ set_memory_limit()

void rapidsmpf::BufferResource::set_memory_limit	(	MemoryType	mem_type,
		std::int64_t	limit
	)

noexcept

Updates the memory limit for a given memory type at runtime.

The store is atomic, but readers (e.g. memory_available() and reserve()) observe the limit and the allocation count independently. A concurrent set_memory_limit() call can change the limit between a caller's read of memory_available() and a subsequent allocation decision; callers that need a coherent view must serialize updates with higher-level synchronization.

Parameters

mem_type	The memory type whose limit is being updated.
limit	The new byte limit. Negative values are permitted; they make `memory_available(mem_type)` always negative and so trigger continuous spilling.

◆ 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 ( ) const

noexcept

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

Returns: Shared pointer the Statistics instance.

◆ stream_pool()

std::shared_ptr<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: Shared pointer to the underlying CUDA stream pool.

◆ try_pinned_mr()

std::optional<PinnedMemoryResource> rapidsmpf::BufferResource::try_pinned_mr ( ) const

Get the pinned host memory resource if available.

Returns: The PinnedMemoryResource is available, or std::nullopt if pinned host memory is not available. The returned handle keeps this BufferResource alive as long as the handle (or any copy) exists.

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

/__w/rapidsmpf/rapidsmpf/cpp/include/rapidsmpf/memory/buffer_resource.hpp

Public Member Functions

Static Public Member Functions

Detailed Description

Member Function Documentation

◆ create()

◆ device_mr()

◆ device_mr_adaptor()

◆ from_options()

◆ host_mr()

◆ make_buffer() [1/2]

◆ make_buffer() [2/2]

◆ memory_available()

◆ memory_reserved()

◆ move() [1/2]

◆ move() [2/2]

◆ move_to_device_buffer()

◆ move_to_host_buffer()

◆ operator=() [1/2]

◆ operator=() [2/2]

◆ pinned_mr()

◆ release()

◆ reserve()

◆ reserve_device_memory_and_spill()

◆ reserve_or_fail() [1/2]

◆ reserve_or_fail() [2/2]

◆ set_memory_limit()

◆ spill_manager()

◆ statistics()

◆ stream_pool()

◆ try_pinned_mr()