Core API#

enum class rapidsmpf::AllowOverbooking : bool#

Policy controlling whether a memory reservation is allowed to overbook.

This enum is used throughout RapidsMPF to specify the overbooking behavior of a memory reservation request. The exact semantics depend on the specific API and execution context in which it is used.

Values:

enumerator NO#

Overbooking is not allowed.

enumerator YES#

Overbooking is allowed.

enum class rapidsmpf::MemoryType : int#

Enum representing the type of memory sorted in decreasing order of preference.

Values:

enumerator DEVICE#

Device memory.

enumerator PINNED_HOST#

Pinned host memory.

enumerator HOST#

Host memory.

enum class rapidsmpf::TrimZeroFraction#

Control whether a zero fractional part is omitted when formatting values.

Values:

enumerator NO#

Always keep the fractional part.

enumerator YES#

Omit the fractional part when it consists only of zeros.

typedef std::int32_t rapidsmpf::Rank#

The rank of a node (e.g. the rank of a MPI process), or world size (total number of ranks).

Note

Ranks are always consecutive integers from zero to the total number of ranks.

typedef std::int32_t rapidsmpf::OpID#

Operation ID defined by the user. This allows users to concurrently execute multiple operations, and each operation will be identified by its OpID.

Note

Although typed as an int32, the number of distinct operations is limited to 2^20.

typedef std::int32_t rapidsmpf::StageID#

Identifier for a stage of a communication operation.

Note

Although typed as an int32, the number of distinct stages is limited to 2^3.

using rapidsmpf::any_device_resource = cuda::mr::any_resource<cuda::mr::device_accessible>#

Owning type-erased device memory resource.

using rapidsmpf::any_host_device_resource = cuda::mr::any_resource<cuda::mr::host_accessible, cuda::mr::device_accessible>#

Owning type-erased host- and device-accessible memory resource.

using rapidsmpf::any_host_resource = cuda::mr::any_resource<cuda::mr::host_accessible>#

Owning type-erased host memory resource.

using rapidsmpf::Clock = std::chrono::high_resolution_clock#

Alias for high-resolution clock from the chrono library.

using rapidsmpf::Duration = std::chrono::duration<double>#

Alias for a duration type representing time in seconds as a double.

using rapidsmpf::TimePoint = std::chrono::time_point<Clock, Duration>#

Alias for a time point with double precision in seconds.

constexpr bool rapidsmpf::COMM_HAVE_UCXX = false#

Whether RapidsMPF was built with the UCXX Communicator.

constexpr bool rapidsmpf::COMM_HAVE_MPI = false#

Whether RapidsMPF was built with the MPI Communicator.

constexpr std::array<MemoryType, 3> rapidsmpf::MEMORY_TYPES{{MemoryType::DEVICE, MemoryType::PINNED_HOST, MemoryType::HOST}}#

All memory types sorted in decreasing order of preference.

constexpr std::array<char const*, MEMORY_TYPES.size()> rapidsmpf::MEMORY_TYPE_NAMES{{"DEVICE", "PINNED_HOST", "HOST"}}#

Memory type names sorted to match MemoryType and MEMORY_TYPES.

constexpr std::array<MemoryType, 2> rapidsmpf::SPILL_TARGET_MEMORY_TYPES{{MemoryType::PINNED_HOST, MemoryType::HOST}}#

Memory types that are valid spill destinations in decreasing order of preference.

This array defines the preferred targets for spilling when device memory is insufficient. The ordering reflects the policy of spilling in RapidsMPF, where earlier entries are considered more desirable spill destinations.

constexpr std::optional<PinnedPoolProperties> rapidsmpf::PinnedMemoryDisabled = {}#

Sentinel used to disable pinned host memory.

Pass this in place of a PinnedPoolProperties (e.g. to BufferResource::create()) to disable pinned host memory allocations.

inline std::ostream &rapidsmpf::operator<<(
std::ostream &os,
Communicator const &obj
)#

Overloads the stream insertion operator for the Communicator class.

This function allows a description of a Communicator to be written to an output stream.

Parameters:
  • os – The output stream to write to.

  • obj – The object to write.

Returns:

A reference to the modified output stream.

template<detail::input_range_of<cuda::stream_ref> Range1, detail::input_range_of<cuda::stream_ref> Range2>
void rapidsmpf::cuda_stream_join(
Range1 const &downstreams,
Range2 const &upstreams,
CudaEvent *event = nullptr
)#

Make downstream CUDA streams wait on upstream CUDA streams.

This call is asynchronous with respect to the host thread; no host-side blocking occurs.

Note

If all upstream and downstream streams are identical, this function is a no-op.

Template Parameters:
  • Range1 – Iterable whose elements are cuda::stream_ref.

  • Range2 – Iterable whose elements are cuda::stream_ref.

Parameters:
  • downstreams – Streams that must not run ahead.

  • upstreams – Streams whose already-enqueued work must complete first.

  • event – Optional CUDA event used for synchronization. A unique event per call is not required; the same event may be reused. If nullptr, a temporary event is created internally. The reason to provide an event is to avoid the small overhead of constructing a temporary one.

inline void rapidsmpf::cuda_stream_join(
cuda::stream_ref downstream,
cuda::stream_ref upstream,
CudaEvent *event = nullptr
)#

Make a downstream CUDA stream wait on an upstream CUDA stream.

This call is asynchronous with respect to the host thread; no host-side blocking occurs.

Equivalent to calling the range overload with one upstream and one downstream.

Note

If downstream and upstream are identical, this function is a no-op.

Parameters:
  • downstream – Stream that must not run ahead.

  • upstream – Stream whose already-enqueued work must complete first.

  • event – Optional CUDA event used for synchronization. A unique event per call is not required; the same event may be reused. If nullptr, a temporary event is created internally to avoid the small overhead of constructing one per call site.

void rapidsmpf::buffer_copy(
std::shared_ptr<Statistics> statistics,
Buffer &dst,
Buffer const &src,
std::size_t size,
std::ptrdiff_t dst_offset = 0,
std::ptrdiff_t src_offset = 0
)#

Asynchronously copy data between buffers.

Copies size bytes from src, starting at src_offset, into dst at dst_offset.

Note

The copy is stream-ordered on dst's stream, correct cross-stream ordering between src's stream and dst's stream is provided automatically.

Parameters:
  • statisticsStatistics object used to record the copy operation. Use Statistics::disabled() to skip recording.

  • dst – Destination buffer.

  • src – Source buffer.

  • size – Number of bytes to copy.

  • dst_offset – Byte offset into the destination buffer.

  • src_offset – Byte offset into the source buffer.

Throws:

std::invalid_argument – If the requested range is out of bounds.

std::int64_t rapidsmpf::device_limit_from_options(
config::Options options
)#

Parse the spill_device_limit parameter from configuration options.

Reads the spill_device_limit option, falling back to 80% of total device memory when unset. The result is aligned down to rmm::CUDA_ALLOCATION_ALIGNMENT.

Parameters:

options – Configuration options.

Returns:

The device memory limit in bytes.

std::optional<Duration> rapidsmpf::periodic_spill_check_from_options(
config::Options options
)#

Get the periodic_spill_check parameter from configuration options.

Parameters:

options – Configuration options.

Returns:

The duration of the pause between spill checks or std::nullopt if no dedicated thread should check for spilling.

std::shared_ptr<StreamPool> rapidsmpf::stream_pool_from_options(
config::Options options
)#

Get a new CUDA stream pool from configuration options.

Parameters:

options – Configuration options.

Returns:

Pool of CUDA streams used throughout RapidsMPF for operations that do not take an explicit CUDA stream.

inline cudaError_t rapidsmpf::cuda_memcpy_batch_async(
void *const *dsts,
void const *const *srcs,
std::size_t const *sizes,
std::size_t count,
cuda::stream_ref stream
)#

Asynchronously copies a batch of buffers using the most efficient available API.

On CUDA 13.0+ with a non-default stream, uses cudaMemcpyBatchAsync with cudaMemcpySrcAccessOrderStream, which defers reading the source buffers until the stream reaches each copy. This enables true asynchronous copies from pageable host memory on modern systems with HMM/ATS support. A batch uses cudaMemcpyFlagPreferOverlapWithCompute when every copy is 128 KiB or less. If any copy is larger, the batch uses cudaMemcpyFlagDefault.

Falls back to per-copy cudaMemcpyAsync on older CUDA versions or when the default stream is used.

Parameters:
  • dsts – Host pointer to a list of destination pointers.

  • srcs – Host pointer to a list of source pointers.

  • sizes – Host pointer to a list of sizes (bytes).

  • count – Number of entries in dsts, srcs, sizes.

  • stream – CUDA stream on which copies are enqueued.

Returns:

cudaError_t CUDA error code.

inline cudaError_t rapidsmpf::cuda_memcpy_async(
void *dst,
void const *src,
std::size_t count,
cuda::stream_ref stream
)#

Asynchronously copies memory between host and/or device buffers.

The copy direction is inferred from the pointer types (cudaMemcpyDefault). The source buffer must remain valid until the stream executes the copy.

This function should be used instead of cudaMemcpyAsync, as it provides improved semantics for asynchronous copies, especially from pageable host memory.

constexpr std::span<MemoryType const> rapidsmpf::leq_memory_types(
MemoryType mem_type
) noexcept#

Get the memory types with preference lower than or equal to mem_type.

The returned span reflects the predefined ordering used in MEMORY_TYPES, which lists memory types in decreasing order of preference.

Parameters:

mem_type – The memory type used as the starting point.

Returns:

A span of memory types whose preference is lower than or equal to the given type.

constexpr char const *rapidsmpf::to_string(MemoryType mem_type)#

Get the name of a MemoryType.

Parameters:

mem_type – The memory type.

Returns:

The memory type name.

std::ostream &rapidsmpf::operator<<(
std::ostream &os,
MemoryType mem_type
)#

Overload to write type name to the output stream.

Parameters:
  • os – The output stream.

  • mem_type – The memory type to write name of to the output stream.

Returns:

The output stream.

std::istream &rapidsmpf::operator>>(
std::istream &is,
MemoryType &out
)#

Overload to read a MemoryType value from an input stream.

Parsing is case-insensitive. Supported values are: “DEVICE”, “PINNED_HOST”, “PINNED”, “PINNED-HOST”, and “HOST”.

If token extraction from the stream fails, the stream state is preserved. If extraction succeeds but the token does not represent a valid MemoryType, the stream failbit is set.

Parameters:
  • is – The input stream.

  • out – The memory type read from the input stream.

Returns:

The input stream.

inline bool rapidsmpf::is_pinned_memory_resources_supported()#

Checks if the PinnedMemoryResource is supported for the current CUDA version.

RapidsMPF requires CUDA 12.6 or newer to support pinned memory resources.

Returns:

True if the PinnedMemoryResource is supported for the current CUDA version, false otherwise.

std::optional<PinnedPoolProperties> rapidsmpf::pinned_pool_properties_from_options(
config::Options options
)#

Parse pinned memory pool properties from configuration options.

Recognized options:

  • ”pinned_memory”: enable pinned memory.

  • ”pinned_initial_pool_size” (bytes or percentage): initial pool size.

    • Byte values (e.g. “1 MiB”) are applied literally.

    • Percentages (e.g. “10%”) are relative to get_host_memory_per_gpu().

  • ”pinned_max_pool_size” (bytes, percentage, or disabled): maximum pool size.

    • Byte and percentages uses the same parsing rules as “pinned_initial_pool_size”.

    • A disabled value (e.g. “off”) leaves the pool unbounded.

Parameters:

options – Configuration options.

Returns:

The parsed PinnedPoolProperties when “pinned_memory” is enabled, otherwise std::nullopt (pinned host memory disabled).

template<typename ...Properties>
bool rapidsmpf::is_host_accessible(
cuda::mr::resource_ref<Properties...> const &mr
) noexcept#

Check whether a type-erased memory resource is host-accessible.

Queries the resource’s dynamic_accessibility_property and returns true if the reported accessibility is host-only or host-and-device.

Template Parameters:

Properties – The property pack of the resource reference.

Parameters:

mr – The memory resource reference to query.

Returns:

True if the resource is host-accessible, false otherwise.

template<typename ...Properties>
bool rapidsmpf::is_device_accessible(
cuda::mr::resource_ref<Properties...> const &mr
) noexcept#

Check whether a type-erased memory resource is device-accessible.

Queries the resource’s dynamic_accessibility_property and returns true if the reported accessibility is device-only or host-and-device.

Template Parameters:

Properties – The property pack of the resource reference.

Parameters:

mr – The memory resource reference to query.

Returns:

True if the resource is device-accessible, false otherwise.

std::vector<PackedData> rapidsmpf::spill_partitions(
std::vector<PackedData> &&partitions,
BufferResource *br
)#

Spill partitions from device memory to host memory.

Moves the buffer of each PackedData from device memory to host memory using the provided buffer resource and the buffer’s CUDA stream. Partitions that are already in host memory are passed through unchanged.

For device-resident partitions, a host memory reservation is made before moving the buffer. If the reservation fails due to insufficient host memory, an exception is thrown. Overbooking is not allowed.

Parameters:
  • partitions – The partitions to spill.

  • brBuffer resource used to reserve host memory and perform the move.

Throws:

rapidsmpf::reservation_error – If host memory reservation fails.

Returns:

A vector of PackedData, where each buffer resides in host memory.

std::vector<PackedData> rapidsmpf::unspill_partitions(
std::vector<PackedData> &&partitions,
BufferResource *br,
AllowOverbooking allow_overbooking
)#

Move spilled partitions (i.e., packed tables in host memory) back to device memory.

Each partition is inspected to determine whether its buffer resides in device memory. Buffers already in device memory are left untouched. Host-resident buffers are moved to device memory using the provided buffer resource and the buffer’s CUDA stream.

If insufficient device memory is available, the buffer resource’s spill manager is invoked to free memory. If overbooking occurs and spilling fails to reclaim enough memory, behavior depends on the allow_overbooking flag.

Parameters:
  • partitions – The partitions to unspill, potentially containing host-resident data.

  • brBuffer resource responsible for memory reservation and spills.

  • allow_overbooking – If false, ensures enough memory is freed to satisfy the reservation; otherwise, allows overbooking even if spilling was insufficient.

Throws:

rapidsmpf::reservation_error – If overbooking exceeds the amount spilled and allow_overbooking is false.

Returns:

A vector of PackedData, each with a buffer in device memory.

std::uint64_t rapidsmpf::get_total_host_memory() noexcept#

Get the total amount of system memory.

Note

On WSL and in containerized environments, the returned value reflects the memory visible to the Linux kernel instance, which may differ from the physical memory of the host.

Note

Terminates the process if sysconf(_SC_PAGE_SIZE) or sysconf(_SC_PHYS_PAGES) fails.

Returns:

Total host memory in bytes.

int rapidsmpf::get_current_numa_node() noexcept#

Get the NUMA node ID associated with the calling CPU thread.

A NUMA (Non-Uniform Memory Access) node represents a group of CPU cores and memory that have faster access to each other than to memory attached to other nodes. On NUMA systems, binding allocations and threads to the same NUMA node can significantly reduce memory access latency and improve bandwidth.

This function returns the NUMA node on which the calling thread is currently executing, as determined by the operating system’s CPU and memory topology. The value can change if the thread migrates between CPUs.

If NUMA support is not available on the system or cannot be queried, the function returns 0, which corresponds to the single implicit NUMA node on non-NUMA systems.

Returns:

The NUMA node ID of the calling thread, or 0 if NUMA is unavailable.

std::vector<int> rapidsmpf::get_current_numa_nodes() noexcept#

Get current NUMA node(s) for memory binding.

Queries the process memory policy and returns the NUMA nodes from which the process may allocate memory. This reflects bindings applied via numa_set_membind() rather than the NUMA node of the CPU currently running the caller.

If NUMA support is not available or the NUMA node cannot be determined, the function returns a vector containing a single element, 0, which corresponds to the single implicit NUMA node on non-NUMA systems.

Returns:

Vector of NUMA node IDs in the current memory policy.

std::uint64_t rapidsmpf::get_numa_node_host_memory(
int numa_id = get_current_numa_node()
) noexcept#

Get the total amount of host memory for a NUMA node.

Note

If NUMA support is not available or the node size cannot be determined, this function falls back to returning the total host memory.

Parameters:

numa_id – NUMA node for which to query the total host memory. Defaults to the current NUMA node as returned by get_current_numa_node().

Returns:

Total host memory of the NUMA node in bytes.

std::uint64_t rapidsmpf::get_host_memory_per_gpu()#

Get the amount of host memory per GPU.

This is calculated as the total host memory available for the current NUMA node divided by the number of GPUs bound to that NUMA node.

Throws:

std::runtime_error – if no GPUs are found on the current NUMA node.

Returns:

Amount of host memory per GPU in bytes.

template<typename MapType>
std::pair<typename MapType::key_type, typename MapType::mapped_type> rapidsmpf::extract_item(
MapType &map,
typename MapType::const_iterator position
)#

Extracts a key-value pair from a map, removing it from the map.

Note

Invalidates any iterators to the extracted element (notably position).

Template Parameters:

MapType – The type of the associative container.

Parameters:
  • map – The map from which to extract the key-value pair.

  • position – Const iterator pointing to a node in the map.

Throws:

std::out_of_range – If the iterator is not found in the map.

Returns:

A pair containing the extracted key and value.

template<typename MapType>
std::pair<typename MapType::key_type, typename MapType::mapped_type> rapidsmpf::extract_item(
MapType &map,
typename MapType::key_type const &key
)#

Extracts a key-value pair from a map, removing it from the map.

Template Parameters:

MapType – The type of the associative container.

Parameters:
  • map – The map from which to extract the key-value pair.

  • key – The key to extract.

Throws:

std::out_of_range – If the key is not found in the map.

Returns:

A pair containing the extracted key and value.

template<typename MapType>
MapType::mapped_type rapidsmpf::extract_value(
MapType &map,
typename MapType::key_type const &key
)#

Extracts the value associated with a specific key from a map, removing the key-value pair.

Template Parameters:

MapType – The type of the associative container.

Parameters:
  • map – The map from which to extract the value.

  • key – The key associated with the value to extract.

Throws:

std::out_of_range – If the key is not found in the map.

Returns:

The extracted value.

template<typename MapType>
MapType::mapped_type rapidsmpf::extract_value(
MapType &map,
typename MapType::const_iterator position
)#

Extracts the value associated with a specific key from a map, removing the key-value pair.

Note

Invalidates any iterators to the extracted element (notably position).

Template Parameters:

MapType – The type of the associative container.

Parameters:
  • map – The map from which to extract the value.

  • position – Const iterator pointing to a node in the map.

Throws:

std::out_of_range – If the key is not found in the map.

Returns:

The extracted value.

template<typename MapType>
MapType::key_type rapidsmpf::extract_key(
MapType &map,
typename MapType::key_type const &key
)#

Extracts a key from a map, removing the key-value pair.

Template Parameters:

MapType – The type of the associative container.

Parameters:
  • map – The map from which to extract the key.

  • key – The key to extract.

Throws:

std::out_of_range – If the key is not found in the map.

Returns:

The extracted key.

template<typename MapType>
MapType::key_type rapidsmpf::extract_key(
MapType &map,
typename MapType::const_iterator position
)#

Extracts a key from a map, removing the key-value pair.

Note

Invalidates any iterators to the extracted element (notably position).

Template Parameters:

MapType – The type of the associative container.

Parameters:
  • map – The map from which to extract the key.

  • position – Const iterator pointing to a node in the map.

Throws:

std::out_of_range – If the key is not found in the map.

Returns:

The extracted key.

template<typename MapType>
auto rapidsmpf::to_vector(MapType &&map)#

Converts a map-like associative container to a vector by moving the values and discarding the keys.

Template Parameters:

MapType – The type of the map-like associative container. Must provide a mapped_type and support range-based for-loops.

Parameters:

map – The map whose values will be moved into the resulting vector. Keys are ignored.

Returns:

A std::vector containing the moved values from the input map.

bool rapidsmpf::is_running_under_valgrind()#

Checks whether the application is running under Valgrind.

Returns:

true if the application is running under Valgrind, false otherwise.

template<typename T>
constexpr T rapidsmpf::safe_div(T x, T y)#

Performs safe division, returning 0 if the denominator is zero.

Template Parameters:

T – The numeric type of the operands.

Parameters:
  • x – The numerator.

  • y – The denominator.

Returns:

T The result of x / y, or 0 if y is zero.

template<std::integral T>
constexpr T rapidsmpf::ceil_div(T x, T y)#

Computes the ceiling of the division of two integers.

Returns the smallest integer not less than x / y. Both operands must be non-negative and the denominator must be non-zero. Computed as x / y + (x % y != 0) to avoid the overflow (and signed UB)

Template Parameters:

T – An integral type.

Parameters:
  • x – The numerator (must be non-negative).

  • y – The denominator (must be positive).

Returns:

T The ceiling of x / y.

inline auto rapidsmpf::chunk_indices(
std::size_t count,
std::size_t num_chunks
)#

Splits the index range [0, count) into exactly num_chunks contiguous chunks.

Each chunk is a half-open [begin, end) index pair. Chunks are front-loaded with size ceil(count / num_chunks); the last non-empty chunk may be smaller and, when count < num_chunks, the trailing chunks are empty (begin == end). The chunks exactly tile [0, count), so their sizes sum to count.

Unlike std::ranges::chunk_view (C++23), this always yields exactly num_chunks chunks, injecting empty trailing chunks as needed.

Parameters:
  • count – The number of elements to split.

  • num_chunks – The number of chunks to produce (must be positive).

Returns:

A lazy view of num_chunks std::pair<std::size_t, std::size_t> [begin, end) index pairs.

template<std::ranges::input_range R, typename T, typename Proj = std::identity>
constexpr bool rapidsmpf::contains(
R &&range,
T const &value,
Proj proj = {}
)#

Backport of std::ranges::contains from C++23 for C++20.

Checks whether a range contains a given value.

Template Parameters:
  • R – An input range type.

  • T – The type of the value to search for.

  • Proj – A projection function applied to each element before comparison.

Parameters:
  • range – The range to search.

  • value – The value to search for in the range.

  • proj – The projection to apply to each element before comparison.

Returns:

true if any element in the range compares equal to value after projection, false otherwise.

template<SharedOrWeakPtr A, SharedOrWeakPtr B>
bool rapidsmpf::owner_equal(
A const &a,
B const &b
) noexcept#

Backport of std::weak_ptr::owner_equal / std::owner_equal from C++26.

Returns whether a and b share ownership of the same managed object, or are both empty. Mirrors the contract of the C++26 standard utilities.

Both arguments must be std::shared_ptr<T> or std::weak_ptr<T> for the same element type T; mismatched element types are rejected at compile time.

Template Parameters:
  • Astd::shared_ptr<T> or std::weak_ptr<T>.

  • Bstd::shared_ptr<T> or std::weak_ptr<T> (same T as A).

Parameters:
  • a – First pointer.

  • b – Second pointer.

Returns:

true iff a and b own the same managed object, or are both empty.

template<typename To, typename From>
To rapidsmpf::safe_cast(
From value,
std::source_location const &loc = std::source_location::current()
)#

Safely casts a numeric value to another type with overflow checking.

For integral conversions, the value must be representable in the destination type or an exception is thrown.

For conversions involving floating point types, overflow and underflow follow standard floating point semantics. The result may become inf or -inf, or lose precision, without throwing.

Template Parameters:
  • To – The destination type.

  • From – The source type.

Parameters:
  • value – The value to cast.

  • loc – Source location (automatically captured).

Throws:

std::overflow_error – if an integral value cannot be represented in the destination type.

Returns:

To The safely cast value.

std::string rapidsmpf::trim(std::string_view text)#

Trims whitespace from both ends of the specified string.

Parameters:

text – The input string to be processed.

Returns:

The trimmed string.

std::string rapidsmpf::to_lower(std::string_view text)#

Converts the specified string to lowercase.

Parameters:

text – The input string to be processed.

Returns:

The trimmed string.

std::string rapidsmpf::to_upper(std::string_view text)#

Converts the specified string to uppercase.

Parameters:

text – The input string to be processed.

Returns:

The trimmed string.

std::string rapidsmpf::format_nbytes(
double nbytes,
int num_decimals = 2,
TrimZeroFraction trim_zero_fraction = TrimZeroFraction::YES
)#

Format a byte count as a human-readable string using IEC units.

Converts an integer byte count into a scaled string representation using binary (base-1024) units such as KiB, MiB, and GiB.

Negative values are supported and are formatted with a leading minus sign, which is useful when representing signed byte deltas or accounting values.

Decimal formatting is controlled by precision. When trim_zero_fraction is set to TrimZeroFraction::YES, the fractional part is omitted entirely if all decimal digits are zero. Otherwise, the specified number of decimal places is preserved.

Examples:

  • 1024 bytes with 2 decimals → “1.00 KiB” or “1 KiB” (trimmed)

  • 1536 bytes with 2 decimals → “1.50 KiB”

Parameters:
  • nbytes – Signed number of bytes to format, provided as a double to support any integer magnitude.

  • num_decimals – Number of decimal places to include in the formatted value.

  • trim_zero_fraction – Whether to omit the fractional part when it consists only of zeros.

Returns:

Human-readable string representation of the byte count.

std::string rapidsmpf::format_duration(
double seconds,
int precision = 2,
TrimZeroFraction trim_zero_fraction = TrimZeroFraction::YES
)#

Format a time duration as a human-readable string.

Converts a duration given in seconds into a scaled string representation using common time units such as ns, us, ms, s, min, h, and d.

The duration is accepted as a double to support both fractional seconds and very large values without overflow.

Negative values are supported and are formatted with a leading minus sign, which is useful when representing signed time deltas.

Decimal formatting is controlled by precision. When trim_zero_fraction is set to TrimZeroFraction::YES, the fractional part is omitted entirely if all decimal digits are zero. Otherwise, the specified number of decimal places is preserved.

Parameters:
  • seconds – Time duration to format, in seconds.

  • precision – Number of decimal places to include in the formatted value.

  • trim_zero_fraction – Whether to omit the fractional part when it consists only of zeros.

Returns:

Human-readable string representation of the time duration.

std::int64_t rapidsmpf::parse_nbytes(std::string_view text)#

Parse a human-readable byte count into an integer number of bytes.

Parses a numeric value followed by an optional unit suffix and converts it to a byte count. Both IEC (base-1024) and SI (base-1000) units are supported.

Supported units:

  • Bytes: B

  • IEC (base-1024): KiB, MiB, GiB, TiB, PiB, EiB, ZiB, YiB

  • SI (base-1000): KB, MB, GB, TB, PB, EB, ZB, YB

Units are case-insensitive. If no unit is provided, the value is interpreted as bytes.

The numeric portion may be specified using integer, decimal, or scientific notation (e.g. “1e6”, “2.5E-3”). The final byte count is rounded to the nearest integer, with ties rounded away from zero.

Parameters:

text – Byte count string to parse.

Throws:
  • std::invalid_argument – If the string format is invalid or the unit is not recognized.

  • std::out_of_range – If the parsed value is not finite or the resulting byte count overflows a 64-bit signed integer.

Returns:

Parsed byte count in bytes.

std::size_t rapidsmpf::parse_nbytes_unsigned(std::string_view text)#

Parse a human-readable byte count into a non-negative number of bytes.

Parses a numeric value followed by an optional unit suffix and converts it to a byte count. Both IEC (base-1024) and SI (base-1000) units are supported.

Supported units:

  • Bytes: B

  • IEC (base-1024): KiB, MiB, GiB, TiB, PiB, EiB, ZiB, YiB

  • SI (base-1000): KB, MB, GB, TB, PB, EB, ZB, YB

Units are case-insensitive. If no unit is provided, the value is interpreted as bytes.

The numeric portion may be specified using integer, decimal, or scientific notation (e.g. “1e6”, “2.5E-3”). The final byte count is rounded to the nearest integer, with ties rounded away from zero.

Negative values are not permitted.

Parameters:

text – Byte count string to parse.

Throws:
  • std::invalid_argument – If the string format is invalid, the unit is not recognized, or the parsed value is negative.

  • std::out_of_range – If the parsed value is not finite or overflows std::size_t.

Returns:

Parsed byte count in bytes.

std::size_t rapidsmpf::parse_nbytes_or_percent(
std::string_view text,
double total_bytes
)#

Parse a byte quantity or percentage into an absolute byte count.

The input may be a human-readable byte string (e.g. “1GiB”, “512MB”) or a percentage (e.g. “25%”). See parse_nbytes_unsigned for the exact parsing semantics of the numeric part.

If text ends with ‘’, the numeric part is first parsed using parse_nbytes_unsigned, then interpreted as a percentage of total_bytes.

Otherwise, text is parsed as an absolute byte value and returned as-is.

Parameters:
  • text – Input string representing a byte quantity or percentage.

  • total_bytes – Total number of bytes used when text is a percentage. Must be positive.

Throws:
  • std::invalid_argument – If the input format is invalid, the value is negative, or if total_bytes is not positive.

  • std::out_of_range – If the parsed or computed value exceeds the representable range.

Returns:

Absolute number of bytes computed from text.

Duration rapidsmpf::parse_duration(std::string_view text)#

Parse a human-readable time duration into seconds.

Parses a numeric value followed by an optional time unit suffix and converts it to a Duration, which represents a time interval in seconds as a double.

Supported units:

  • Nanoseconds: ns

  • Microseconds: µs or us

  • Milliseconds: ms

  • Seconds: s

  • Minutes: m or min

  • Hours: h

  • Days: d

Units are case-insensitive. If no unit is provided, the value is interpreted as seconds.

The numeric portion may be specified using integer, decimal, or scientific notation (e.g. “1e3”, “2.5E-2”). Negative values are supported.

Parameters:

text – Time duration string to parse.

Throws:
  • std::invalid_argument – If the string format is invalid or the unit is not recognized.

  • std::out_of_range – If the parsed value is not finite.

Returns:

Parsed duration in seconds.

template<typename T>
T rapidsmpf::parse_string(
std::string const &text
)#

Specialization of parse_string for boolean values.

Converts the input string to a boolean. This function handles common boolean representations such as true, false, on, off, yes, and no, as well as numeric representations (e.g., 0 or 1). The input is first checked for a numeric value using std::stoi; if that fails, it is lowercased and trimmed before matching against known textual representations.

Parameters:

text – String to convert to a boolean.

Throws:

std::invalid_argument – If the string cannot be interpreted as a boolean.

Returns:

The corresponding boolean value.

template<>
bool rapidsmpf::parse_string(std::string const &text)#

Specialization of parse_string for boolean values.

Converts the input string to a boolean. This function handles common boolean representations such as true, false, on, off, yes, and no, as well as numeric representations (e.g., 0 or 1). The input is first checked for a numeric value using std::stoi; if that fails, it is lowercased and trimmed before matching against known textual representations.

Parameters:

text – String to convert to a boolean.

Throws:

std::invalid_argument – If the string cannot be interpreted as a boolean.

Returns:

The corresponding boolean value.

std::optional<std::string> rapidsmpf::parse_optional(
std::string text
)#

Parse an optional string value.

Returns std::nullopt if the input string represents a disabled value. Otherwise, the input string is returned unchanged.

Disabled values are matched case-insensitively and may include surrounding whitespace. Recognized values include: false, no, off, disable, disabled, none, n/a, and na.

Parameters:

text – Input string to parse.

Returns:

std::optional<std::string> Parsed optional string.

std::vector<std::string> rapidsmpf::parse_string_list(
std::string_view text,
char delimiter = ','
)#

Parse a delimited string into a list of trimmed substrings.

Splits the input string by the specified delimiter and returns a vector of trimmed tokens. Leading and trailing whitespace is removed from each token.

If the input string is empty or contains only whitespace, an empty vector is returned.

Parameters:
  • text – Input string to parse.

  • delimiter – Character to use as the delimiter. Defaults to comma.

Returns:

Vector of trimmed strings.

class Tag#
#include <communicator.hpp>

A tag used for identifying messages in a communication operation.

The tag is a 32-bit integer, with the following layout (low bits to high bits left-to-right)

bits   | 012       | 34567 01234567 0123456 | 7 01234567
       |           |                        |
value  | stage (3) | operation (20)         | empty (9)
       |           |                        |
{}

The restriction of 23 used bits comes from empirical MPI implementation limits for MPI_TAG_UB. For example, when using UCX as a transport layer, OpenMPI is restricted to 2^23 distinct tags.

All messages in rapidsmpf over the same Communicator are disambiguated by Tags. A message sent with a given Tag can only be matched by a receive with a matching Tag.

In the same way, collective operations over a Communicator are disambiguated by an OpID: two different collectives with different OpIDs will not interfere.

Due to implementation restrictions, we are limited in the number of distinct tags we can use (the Tag constructor checks for overflow but not reuse). In particular, we support at most 2^20 distinct OpIDs, corresponding to 2^20 distinct collectives running simultaneously.

Public Types

typedef std::int32_t StorageT#

The physical data type to store the tag.

Public Functions

inline constexpr Tag(OpID const op, StageID const stage)#

Constructs a tag.

Throws:

std::overflow_error – If either the op or stage values are negative or too large.

Parameters:
  • op – The operation ID

  • stage – The stage ID

inline constexpr operator StorageT() const noexcept#

Returns the int32 view of the tag.

Returns:

int32 view of the tag

inline constexpr OpID op() const noexcept#

Extracts the operation ID from the tag.

Returns:

The operation ID

inline constexpr StageID stage() const noexcept#

Extracts the stage ID from the tag.

Returns:

The stage ID

Public Static Functions

static inline constexpr std::size_t bit_length() noexcept#

Returns the max number of bits used for the tag.

Returns:

bit length

static inline constexpr StorageT max_value() noexcept#

Returns the max value of the tag.

Returns:

max value

Public Static Attributes

static constexpr int stage_id_bits = {3}#

Number of bits for the stage ID.

static constexpr StorageT stage_id_mask = {(1 << stage_id_bits) - 1}#

Mask for the stage ID.

static constexpr int op_id_bits = {20}#

Number of bits for the operation ID.

static constexpr StorageT op_id_mask{((1 << (op_id_bits + stage_id_bits)) - 1) ^ stage_id_mask}#

Mask for the operation ID.

class Communicator#
#include <communicator.hpp>

Abstract base class for a communication mechanism between nodes.

Provides an interface for sending and receiving messages between nodes, supporting asynchronous operations, GPU data transfers, and custom logging. Implementations must define the virtual methods to enable specific communication backends.

The API of the Communicator is not stream-ordered (since the concrete libraries we use to implement communication patterns are not stream-ordered). A consequence of this is that the user must ensure that any stream-ordered work on buffers is complete before passing a buffer into send or recv. As a corollary, after completing a future, the extracted Buffer is valid on its stream (it has no stream-ordered work queued up).

Sends and receives are matched on (rank, tag) pairs. The concrete implementation must provide that there is no message overtaking.

Subclassed by rapidsmpf::MPI, rapidsmpf::Single, rapidsmpf::ucxx::UCXX

Public Functions

virtual Rank rank() const = 0#

Retrieves the rank of the current node.

Returns:

The rank of the node.

virtual Rank nranks() const = 0#

Retrieves the total number of ranks.

Returns:

The total number of ranks.

virtual std::unique_ptr<Future> send(
std::unique_ptr<std::vector<std::uint8_t>> msg,
Rank rank,
Tag tag
) = 0#

Sends a host message to a specific rank.

This is used to send data that resides in host memory and is guaranteed to be valid at the time of the call.

Use release_sync_host_data to obtain the data buffer again once the future is completed.

Parameters:
  • msg – Unique pointer to the message data (host memory).

  • rank – The destination rank.

  • tag – Message tag for identification.

Throws:

std::invalid_argument – If msg is nullptr.

Returns:

A unique pointer to a Future representing the asynchronous operation.

virtual std::unique_ptr<Future> send(
std::unique_ptr<Buffer> msg,
Rank rank,
Tag tag
) = 0#

Sends a message (device or host) to a specific rank. Use release_data to obtain the data buffer again once the future is completed.

Warning

The caller is responsible to ensure the underlying Buffer allocation and data are already valid before calling, for example, when a CUDA allocation and/or copy are done asynchronously. Specifically, the caller should ensure Buffer::is_ready() returns true before calling this function. Providing a non-ready buffer leads to an irrecoverable condition.

Parameters:
  • msg – Unique pointer to the message data (Buffer).

  • rank – The destination rank.

  • tag – Message tag for identification.

Throws:
  • std::invalid_argument – If msg is nullptr.

  • std::logic_error – If msg is not ready (see warning for more details).

Returns:

A unique pointer to a Future representing the asynchronous operation.

virtual std::unique_ptr<Future> recv(
Rank rank,
Tag tag,
std::unique_ptr<Buffer> recv_buffer
) = 0#

Receives a message from a specific rank to a buffer. Use release_data to extract the data out of the buffer once the future is completed.

Warning

The caller is responsible to ensure the underlying Buffer allocation is already valid before calling, for example, when a CUDA allocation and/or copy are done asynchronously. Specifically, the caller should ensure Buffer::is_ready() returns true before calling this function. Providing a non-ready buffer leads to an irrecoverable condition.

Parameters:
  • rank – The source rank.

  • tag – Message tag for identification.

  • recv_buffer – The receive buffer.

Throws:
  • std::invalid_argument – If recv_buffer is nullptr.

  • std::logic_error – If recv_buffer is not ready (see warning for more details).

Returns:

A unique pointer to a Future representing the asynchronous operation.

virtual std::unique_ptr<Future> recv_sync_host_data(
Rank rank,
Tag tag,
std::unique_ptr<std::vector<std::uint8_t>> synced_buffer
) = 0#

Receives a message from a specific rank to an allocated (synchronized) host buffer. Use release_sync_host_data to extract the data out of the buffer once the future is completed.

Parameters:
  • rank – The source rank.

  • tag – Message tag for identification.

  • synced_buffer – The receive buffer.

Throws:

std::invalid_argument – If synced_buffer is nullptr.

Returns:

A unique pointer to a Future representing the asynchronous operation.

virtual std::pair<std::unique_ptr<std::vector<std::uint8_t>>, Rank> recv_any(
Tag tag
) = 0#

Receives a message from any rank (blocking).

Note

If no message is available this is indicated by returning a nullptr in the first slot of the pair.

Parameters:

tag – Message tag for identification.

Returns:

A pair containing the message data (host memory) and the rank of the sender.

virtual std::unique_ptr<std::vector<std::uint8_t>> recv_from(
Rank src,
Tag tag
) = 0#

Receives a message from a specific rank (blocking).

Note

If no message is available, this function returns a nullptr.

Parameters:
  • src – The source rank from which to receive the message.

  • tag – Message tag for identification.

Returns:

A unique pointer to a vector containing the received message data (host memory).

virtual std::pair<std::vector<std::unique_ptr<Future>>, std::vector<std::size_t>> test_some(
std::vector<std::unique_ptr<Future>> &future_vector
) = 0#

Tests for completion of multiple futures.

Parameters:

future_vector[inout] Vector of Future objects. Completed futures are erased from the vector.

Returns:

Pair of completed futures and indices of input vector that were completed.

virtual std::vector<std::size_t> test_some(
std::unordered_map<std::size_t, std::unique_ptr<Communicator::Future>> const &future_map
) = 0#

Tests for completion of multiple futures in a map.

Parameters:

future_map – Map of futures identified by keys.

Returns:

Keys of completed futures.

virtual bool test(std::unique_ptr<Communicator::Future> &future) = 0#

Test for completion of a single future.

Parameters:

futureFuture to test

Returns:

True if the future is completed. After test returns true, it is safe to call release_data().

virtual std::vector<std::unique_ptr<Buffer>> wait_all(
std::vector<std::unique_ptr<Communicator::Future>> &&futures
) = 0#

Wait for completion of all futures and return their data buffers.

Parameters:

futures – Futures to wait for completion of, consumed.

Returns:

A vector of the contained data buffers.

virtual std::unique_ptr<Buffer> wait(
std::unique_ptr<Future> future
) = 0#

Wait for a future to complete and return the data buffer.

Parameters:

future – The future to wait for completion of.

Returns:

A unique pointer to the GPU data buffer (or nullptr if the future had no data).

virtual std::unique_ptr<Buffer> release_data(
std::unique_ptr<Communicator::Future> future
) = 0#

Retrieves data associated with a completed future.

Parameters:

future – The completed future.

Throws:

std::runtime_error – if the future has no data.

Returns:

A unique pointer to the data buffer.

virtual std::unique_ptr<std::vector<std::uint8_t>> release_sync_host_data(
std::unique_ptr<Communicator::Future> future
) = 0#

Retrieves synchronized host data associated with a completed future. When the future is completed, the the host data is valid, and ready, but not stream-ordered.

Parameters:

future – The completed future.

Throws:

std::runtime_error – if the future has no data.

Returns:

A unique pointer to the synchronized host data.

virtual std::shared_ptr<Logger> const &logger() = 0#

Retrieves the logger associated with this communicator.

Returns:

Shared pointer to the logger.

virtual std::shared_ptr<ProgressThread> const &progress_thread(
) const = 0#

Retrieves the progress thread associated with this communicator.

Returns:

Shared pointer to the progress thread.

virtual std::string str() const = 0#

Provides a string representation of the communicator.

Returns:

A string describing the communicator.

class Future#
#include <communicator.hpp>

Abstract base class for asynchronous operation within the communicator.

Encapsulates the concept of an asynchronous operation, allowing users to query or wait for completion.

Subclassed by rapidsmpf::MPI::Future, rapidsmpf::Single::Future, rapidsmpf::ucxx::UCXX::Future

Public Functions

Future(Future&&) = default#

Movable.

Future &operator=(Future&&) = default#

Move assignment.

Returns:

Moved this.

Future(Future const&) = delete#

Not copyable.

Future &operator=(Future const&) = delete#

Not copy-assignable.

class Logger : public std::enable_shared_from_this<Logger>#
#include <logger.hpp>

A logger base class for handling different levels of log messages.

The logger class provides various logging methods with different verbosity levels. It ensures thread-safety using a mutex and allows filtering of log messages based on the configured verbosity level.

The name used in log message prefixes can either be supplied at construction time or set later via set_name(). The latter supports communicators (such as UCXX) where the identifying name (e.g. the rank) is only known after a bootstrap handshake.

TODO: support writing to a file.

Public Types

enum class LOG_LEVEL : std::uint32_t#

Log verbosity levels.

Defines different logging levels for filtering messages.

Values:

enumerator NONE#

No logging.

enumerator PRINT#

General print messages.

enumerator WARN#

Warning messages.

enumerator INFO#

Informational messages.

enumerator DEBUG#

Debug messages.

enumerator TRACE#

Trace messages.

Public Functions

inline LOG_LEVEL verbosity_level() const#

Get the verbosity level of the logger.

Returns:

The verbosity level.

void set_name(std::string name)#

Update the name used in log message prefixes.

Thread-safe (acquires the same mutex used to serialize log output). May be called any number of times. Concurrent log calls will observe either the old or the new value, never a partially written one.

Parameters:

name – The new name.

template<typename ...Args>
inline void log(
LOG_LEVEL level,
Args const&... args
)#

Logs a message using the specified verbosity level.

Formats and outputs a message if the verbosity level is high enough.

Template Parameters:

Args – Types of the message components, must support the << operator.

Parameters:
  • level – The verbosity level of the message.

  • args – The components of the message to log.

template<typename ...Args>
inline void print(Args const&... args)#

Logs a print message.

Template Parameters:

Args – Types of the message components.

Parameters:

args – The components of the message to log.

template<typename ...Args>
inline void warn(Args const&... args)#

Logs a warning message.

Template Parameters:

Args – Types of the message components.

Parameters:

args – The components of the message to log.

template<typename ...Args>
inline void info(Args const&... args)#

Logs an informational message.

Template Parameters:

Args – Types of the message components.

Parameters:

args – The components of the message to log.

template<typename ...Args>
inline void debug(Args const&... args)#

Logs a debug message.

Template Parameters:

Args – Types of the message components.

Parameters:

args – The components of the message to log.

template<typename ...Args>
inline void trace(Args const&... args)#

Logs a trace message.

Template Parameters:

Args – Types of the message components.

Parameters:

args – The components of the message to log.

Public Static Functions

static inline constexpr char const *level_name(LOG_LEVEL level)#

Get the string name of a log level.

Parameters:

level – The log level.

Returns:

The corresponding log level name or “UNKNOWN” if out of range.

static std::shared_ptr<Logger> create(
LOG_LEVEL level = LOG_LEVEL::WARN,
std::string name = "unknown"
)#

Create a logger.

Parameters:
  • level – The verbosity level (defaults to LOG_LEVEL::WARN).

  • name – The logger name (defaults to "unknown").

Returns:

A shared pointer to the newly constructed logger.

static std::shared_ptr<Logger> from_options(
config::Options options
)#

Create a logger from configuration options.

The name defaults to "unknown" and may be updated later via set_name(). This is intended for bootstrap scenarios where the identifying name is only known after a network handshake but logging may already be required.

To control the verbosity level, set the configuration option “log” to one of following:

  • NONE: No logging.

  • PRINT: General print messages.

  • WARN: Warning messages (default)

  • INFO: Informational messages.

  • DEBUG: Debug messages.

  • TRACE: Trace messages.

Parameters:

options – Configuration options.

Returns:

A shared pointer to the newly constructed logger.

Public Static Attributes

static constexpr std::array<char const*, 6> LOG_LEVEL_NAMES{"NONE", "PRINT", "WARN", "INFO", "DEBUG", "TRACE"}#

Log level names corresponding to the LOG_LEVEL enum.

class MPI : public rapidsmpf::Communicator#
#include <mpi.hpp>

MPI communicator class that implements the Communicator interface.

This class implements communication functions using MPI, allowing for data exchange between processes in a distributed system. It supports sending and receiving data, both on the CPU and GPU, and provides asynchronous operations with support for future results.

Public Functions

MPI(
MPI_Comm comm,
std::shared_ptr<ProgressThread> progress_thread,
std::shared_ptr<Logger> logger
)#

Construct an MPI communicator.

Parameters:
  • comm – The MPI communicator to be used for communication.

  • progress_thread – Progress thread for this communicator.

  • logger – Externally provided logger. Must be non-null. The communicator will overwrite the logger’s rank to the rank reported by MPI_Comm_rank on comm.

inline virtual Rank rank() const override#

Retrieves the rank of the current node.

Returns:

The rank of the node.

inline virtual Rank nranks() const override#

Retrieves the total number of ranks.

Returns:

The total number of ranks.

virtual std::unique_ptr<Communicator::Future> send(
std::unique_ptr<std::vector<std::uint8_t>> msg,
Rank rank,
Tag tag
) override#

Sends a host message to a specific rank.

This is used to send data that resides in host memory and is guaranteed to be valid at the time of the call.

Use release_sync_host_data to obtain the data buffer again once the future is completed.

Parameters:
  • msg – Unique pointer to the message data (host memory).

  • rank – The destination rank.

  • tag – Message tag for identification.

Throws:
  • std::invalid_argument – If msg is nullptr.

  • std::runtime_error – If the message exceeds MPI size limit (2^31 bytes).

Returns:

A unique pointer to a Future representing the asynchronous operation.

virtual std::unique_ptr<Communicator::Future> send(
std::unique_ptr<Buffer> msg,
Rank rank,
Tag tag
) override#

Sends a message (device or host) to a specific rank. Use release_data to obtain the data buffer again once the future is completed.

Warning

The caller is responsible to ensure the underlying Buffer allocation and data are already valid before calling, for example, when a CUDA allocation and/or copy are done asynchronously. Specifically, the caller should ensure Buffer::is_ready() returns true before calling this function. Providing a non-ready buffer leads to an irrecoverable condition.

Parameters:
  • msg – Unique pointer to the message data (Buffer).

  • rank – The destination rank.

  • tag – Message tag for identification.

Throws:
  • std::invalid_argument – If msg is nullptr.

  • std::logic_error – If msg is not ready (see warning for more details).

  • std::runtime_error – If the message exceeds MPI size limit (2^31 bytes).

Returns:

A unique pointer to a Future representing the asynchronous operation.

virtual std::unique_ptr<Communicator::Future> recv(
Rank rank,
Tag tag,
std::unique_ptr<Buffer> recv_buffer
) override#

Receives a message from a specific rank to a buffer. Use release_data to extract the data out of the buffer once the future is completed.

Warning

The caller is responsible to ensure the underlying Buffer allocation is already valid before calling, for example, when a CUDA allocation and/or copy are done asynchronously. Specifically, the caller should ensure Buffer::is_ready() returns true before calling this function. Providing a non-ready buffer leads to an irrecoverable condition.

Parameters:
  • rank – The source rank.

  • tag – Message tag for identification.

  • recv_buffer – The receive buffer.

Throws:
  • std::invalid_argument – If recv_buffer is nullptr.

  • std::logic_error – If recv_buffer is not ready (see warning for more details).

  • std::runtime_error – If the message exceeds MPI size limit (2^31 bytes).

Returns:

A unique pointer to a Future representing the asynchronous operation.

virtual std::unique_ptr<Communicator::Future> recv_sync_host_data(
Rank rank,
Tag tag,
std::unique_ptr<std::vector<std::uint8_t>> synced_buffer
) override#

Receives a message from a specific rank to an allocated (synchronized) host buffer. Use release_sync_host_data to extract the data out of the buffer once the future is completed.

Parameters:
  • rank – The source rank.

  • tag – Message tag for identification.

  • synced_buffer – The receive buffer.

Throws:
  • std::invalid_argument – If synced_buffer is nullptr.

  • std::runtime_error – If the message exceeds MPI size limit (2^31 bytes).

Returns:

A unique pointer to a Future representing the asynchronous operation.

virtual std::pair<std::unique_ptr<std::vector<std::uint8_t>>, Rank> recv_any(
Tag tag
) override#

Receives a message from any rank (blocking).

Note

If no message is available this is indicated by returning a nullptr in the first slot of the pair.

Parameters:

tag – Message tag for identification.

Returns:

A pair containing the message data (host memory) and the rank of the sender.

virtual std::unique_ptr<std::vector<std::uint8_t>> recv_from(
Rank src,
Tag tag
) override#

Receives a message from a specific rank (blocking).

Note

If no message is available, this function returns a nullptr.

Parameters:
  • src – The source rank from which to receive the message.

  • tag – Message tag for identification.

Returns:

A unique pointer to a vector containing the received message data (host memory).

std::pair<std::vector<std::unique_ptr<Communicator::Future>>, std::vector<std::size_t>> test_some(
std::vector<std::unique_ptr<Communicator::Future>> &future_vector
) override#

Tests for completion of multiple futures.

Parameters:

future_vector[inout] Vector of Future objects. Completed futures are erased from the vector.

Returns:

Pair of completed futures and indices of input vector that were completed.

std::vector<std::size_t> test_some(
std::unordered_map<std::size_t, std::unique_ptr<Communicator::Future>> const &future_map
) override#

Tests for completion of multiple futures in a map.

Parameters:

future_map – Map of futures identified by keys.

Returns:

Keys of completed futures.

bool test(std::unique_ptr<Communicator::Future> &future) override#

Test for completion of a single future.

Parameters:

futureFuture to test

Returns:

True if the future is completed. After test returns true, it is safe to call release_data().

std::vector<std::unique_ptr<Buffer>> wait_all(
std::vector<std::unique_ptr<Communicator::Future>> &&futures
) override#

Wait for completion of all futures and return their data buffers.

Parameters:

futures – Futures to wait for completion of, consumed.

Returns:

A vector of the contained data buffers.

std::unique_ptr<Buffer> wait(
std::unique_ptr<Communicator::Future> future
) override#

Wait for a future to complete and return the data buffer.

Parameters:

future – The future to wait for completion of.

Returns:

A unique pointer to the GPU data buffer (or nullptr if the future had no data).

std::unique_ptr<Buffer> release_data(
std::unique_ptr<Communicator::Future> future
) override#

Retrieves data associated with a completed future.

Parameters:

future – The completed future.

Throws:

std::runtime_error – if the future has no data.

Returns:

A unique pointer to the data buffer.

std::unique_ptr<std::vector<std::uint8_t>> release_sync_host_data(
std::unique_ptr<Communicator::Future> future
) override#

Retrieves synchronized host data associated with a completed future. When the future is completed, the the host data is valid, and ready, but not stream-ordered.

Parameters:

future – The completed future.

Throws:

std::runtime_error – if the future has no data.

Returns:

A unique pointer to the synchronized host data.

inline virtual std::shared_ptr<Logger> const &logger() override#

Retrieves the logger associated with this communicator.

Returns:

Shared pointer to the logger.

inline virtual std::shared_ptr<ProgressThread> const &progress_thread(
) const override#

Retrieves the progress thread associated with this communicator.

Returns:

Shared pointer to the progress thread.

virtual std::string str() const override#

Provides a string representation of the communicator.

Returns:

A string describing the communicator.

class Future : public rapidsmpf::Communicator::Future#
#include <mpi.hpp>

Represents the future result of an MPI operation.

This class is used to handle the result of an MPI communication operation asynchronously.

Public Functions

inline Future(
MPI_Request req,
std::unique_ptr<Buffer> data_buffer
)#

Construct a Future from a data buffer.

Parameters:
  • req – The MPI request handle for the operation.

  • data_buffer – A unique pointer to the data buffer.

inline Future(
MPI_Request req,
std::unique_ptr<std::vector<std::uint8_t>> synced_host_data
)#

Construct a Future from synchronized host data.

This constructor is used for MPI operations where the data resides in host memory and is guaranteed to be valid at the time of the call.

Parameters:
  • req – The MPI request handle for the operation.

  • synced_host_data – A unique pointer to a vector containing host memory.

class Single : public rapidsmpf::Communicator#
#include <single.hpp>

Single process communicator class that implements the Communicator interface.

This class stubs out the Communicator interface with functions that throw. When sending to/receiving from self the internal logic should move buffers through the shuffler, rather than invoking send/recv.

Public Functions

Single(
std::shared_ptr<ProgressThread> progress_thread,
std::shared_ptr<Logger> logger
)#

Construct a single process communicator.

Parameters:
  • progress_thread – Progress thread for this communicator.

  • logger – Externally provided logger. Must be non-null. The communicator will overwrite the logger’s rank to 0.

inline virtual constexpr Rank rank() const override#

Retrieves the rank of the current node.

Returns:

The rank of the node.

inline virtual constexpr Rank nranks() const override#

Retrieves the total number of ranks.

Returns:

The total number of ranks.

virtual std::unique_ptr<Communicator::Future> send(
std::unique_ptr<std::vector<std::uint8_t>> msg,
Rank rank,
Tag tag
) override#

Sends a host message to a specific rank.

This is used to send data that resides in host memory and is guaranteed to be valid at the time of the call.

Use release_sync_host_data to obtain the data buffer again once the future is completed.

Parameters:
  • msg – Unique pointer to the message data (host memory).

  • rank – The destination rank.

  • tag – Message tag for identification.

Throws:

std::invalid_argument – If msg is nullptr.

Returns:

A unique pointer to a Future representing the asynchronous operation.

virtual std::unique_ptr<Communicator::Future> send(
std::unique_ptr<Buffer> msg,
Rank rank,
Tag tag
) override#

Sends a message (device or host) to a specific rank. Use release_data to obtain the data buffer again once the future is completed.

Warning

The caller is responsible to ensure the underlying Buffer allocation and data are already valid before calling, for example, when a CUDA allocation and/or copy are done asynchronously. Specifically, the caller should ensure Buffer::is_ready() returns true before calling this function. Providing a non-ready buffer leads to an irrecoverable condition.

Parameters:
  • msg – Unique pointer to the message data (Buffer).

  • rank – The destination rank.

  • tag – Message tag for identification.

Throws:
  • std::invalid_argument – If msg is nullptr.

  • std::logic_error – If msg is not ready (see warning for more details).

  • std::runtime_error – if called (single-process communicators should never send messages).

Returns:

A unique pointer to a Future representing the asynchronous operation.

virtual std::unique_ptr<Communicator::Future> recv(
Rank rank,
Tag tag,
std::unique_ptr<Buffer> recv_buffer
) override#

Receives a message from a specific rank to a buffer. Use release_data to extract the data out of the buffer once the future is completed.

Warning

The caller is responsible to ensure the underlying Buffer allocation is already valid before calling, for example, when a CUDA allocation and/or copy are done asynchronously. Specifically, the caller should ensure Buffer::is_ready() returns true before calling this function. Providing a non-ready buffer leads to an irrecoverable condition.

Parameters:
  • rank – The source rank.

  • tag – Message tag for identification.

  • recv_buffer – The receive buffer.

Throws:
  • std::invalid_argument – If recv_buffer is nullptr.

  • std::logic_error – If recv_buffer is not ready (see warning for more details).

  • std::runtime_error – if called (single-process communicators should never send messages).

Returns:

A unique pointer to a Future representing the asynchronous operation.

virtual std::unique_ptr<Communicator::Future> recv_sync_host_data(
Rank rank,
Tag tag,
std::unique_ptr<std::vector<std::uint8_t>> synced_buffer
) override#

Receives a message from a specific rank to an allocated (synchronized) host buffer. Use release_sync_host_data to extract the data out of the buffer once the future is completed.

Parameters:
  • rank – The source rank.

  • tag – Message tag for identification.

  • synced_buffer – The receive buffer.

Throws:
  • std::invalid_argument – If synced_buffer is nullptr.

  • std::runtime_error – if called (single-process communicators should never send messages).

Returns:

A unique pointer to a Future representing the asynchronous operation.

virtual std::pair<std::unique_ptr<std::vector<std::uint8_t>>, Rank> recv_any(
Tag tag
) override#

Receives a message from any rank (blocking).

Note

If no message is available this is indicated by returning a nullptr in the first slot of the pair.

Note

Always returns a nullptr for the received message, indicating that no message is available.

Parameters:

tag – Message tag for identification.

Returns:

A pair containing the message data (host memory) and the rank of the sender.

virtual std::unique_ptr<std::vector<std::uint8_t>> recv_from(
Rank src,
Tag tag
) override#

Receives a message from a specific rank (blocking).

Note

If no message is available, this function returns a nullptr.

Note

Always returns a nullptr for the received message, indicating that no message is available.

Parameters:
  • src – The source rank from which to receive the message.

  • tag – Message tag for identification.

Returns:

A unique pointer to a vector containing the received message data (host memory).

std::pair<std::vector<std::unique_ptr<Communicator::Future>>, std::vector<std::size_t>> test_some(
std::vector<std::unique_ptr<Communicator::Future>> &future_vector
) override#

Tests for completion of multiple futures.

Parameters:

future_vector[inout] Vector of Future objects. Completed futures are erased from the vector.

Throws:

std::runtime_error – if called (single-process communicators should never send messages).

Returns:

Pair of completed futures and indices of input vector that were completed.

std::vector<std::size_t> test_some(
std::unordered_map<std::size_t, std::unique_ptr<Communicator::Future>> const &future_map
) override#

Tests for completion of multiple futures in a map.

Parameters:

future_map – Map of futures identified by keys.

Throws:

std::runtime_error – if called (single-process communicators should never send messages).

Returns:

Keys of completed futures.

bool test(std::unique_ptr<Communicator::Future> &future) override#

Test for completion of a single future.

Parameters:

futureFuture to test

Returns:

True if the future is completed. After test returns true, it is safe to call release_data().

std::vector<std::unique_ptr<Buffer>> wait_all(
std::vector<std::unique_ptr<Communicator::Future>> &&futures
) override#

Wait for completion of all futures and return their data buffers.

Parameters:

futures – Futures to wait for completion of, consumed.

Returns:

A vector of the contained data buffers.

std::unique_ptr<Buffer> wait(
std::unique_ptr<Communicator::Future> future
) override#

Wait for a future to complete and return the data buffer.

Parameters:

future – The future to wait for completion of.

Throws:

std::runtime_error – if called (single-process communicators should never send messages)

Returns:

A unique pointer to the GPU data buffer (or nullptr if the future had no data).

std::unique_ptr<Buffer> release_data(
std::unique_ptr<Communicator::Future> future
) override#

Retrieves data associated with a completed future.

Parameters:

future – The completed future.

Throws:
  • std::runtime_error – if the future has no data.

  • std::runtime_error – if called (single-process communicators should never send messages).

Returns:

A unique pointer to the data buffer.

std::unique_ptr<std::vector<std::uint8_t>> release_sync_host_data(
std::unique_ptr<Communicator::Future> future
) override#

Retrieves synchronized host data associated with a completed future. When the future is completed, the the host data is valid, and ready, but not stream-ordered.

Parameters:

future – The completed future.

Throws:
  • std::runtime_error – if the future has no data.

  • std::runtime_error – if called (single-process communicators should never send messages).

Returns:

A unique pointer to the synchronized host data.

inline virtual std::shared_ptr<Logger> const &logger() override#

Retrieves the logger associated with this communicator.

Returns:

Shared pointer to the logger.

inline virtual std::shared_ptr<ProgressThread> const &progress_thread(
) const override#

Retrieves the progress thread associated with this communicator.

Returns:

Shared pointer to the progress thread.

virtual std::string str() const override#

Provides a string representation of the communicator.

Returns:

A string describing the communicator.

class Future : public rapidsmpf::Communicator::Future#
#include <single.hpp>

Represents the future result of an operation.

This class is used to handle the result of a communication operation asynchronously.

class CudaEvent#
#include <cuda_event.hpp>

RAII wrapper for a CUDA event with convenience methods.

Creates a CUDA event on construction and destroys it on destruction.

Note

To prevent undefined behavior due to unfinished memory operations, events should be used in the following cases if any of the operations below were performed asynchronously with respect to the host:

  1. Before addressing a device buffer’s allocation.

  2. Before accessing a device buffer’s data that has been copied from any location, or processed by a CUDA kernel.

  3. Before accessing a host buffer’s data that has been copied from device or processed by a CUDA kernel.

Note

CudaEvent objects must not have static storage duration, since CUDA resources are not guaranteed to be valid during program initialization or shutdown.

Public Functions

CudaEvent(unsigned flags = cudaEventDisableTiming)#

Construct a CUDA event.

Parameters:

flags – CUDA event creation flags.

Throws:

rapidsmpf::cuda_error – if cudaEventCreateWithFlags fails.

~CudaEvent() noexcept#

Destroy the CUDA event.

Automatically releases the underlying CUDA event resource.

CudaEvent(CudaEvent const&) = delete#

Non-copyable.

CudaEvent &operator=(CudaEvent const&) = delete#

Non-copy-assignable.

CudaEvent(CudaEvent &&other) noexcept#

Move constructor.

Parameters:

other – Source CudaEvent to move from.

CudaEvent &operator=(CudaEvent &&other)#

Move assignment operator.

The destination must be empty, for example because it has previously been moved from. Move-assigning into a CudaEvent that already owns an event is not allowed.

Parameters:

other – Source CudaEvent to move from.

Throws:

std::invalid_argument – if this object already owns an event.

Returns:

Reference to this object.

void record(cuda::stream_ref stream)#

Record the event on a CUDA stream.

Marks the event as occurring after all prior operations on the given stream.

Parameters:

stream – The CUDA stream to record the event on.

Throws:

rapidsmpf::cuda_error – if cudaEventRecord fails.

bool is_ready() const#

Check if the CUDA event has been completed.

Throws:

rapidsmpf::cuda_error – if cudaEventQuery fails.

Returns:

true if the event has been completed, false otherwise.

void host_wait() const#

Wait for the event to be completed (blocking).

Throws:

rapidsmpf::cuda_error – if cudaEventSynchronize fails.

void stream_wait(cuda::stream_ref stream) const#

Make a CUDA stream wait on this event (non-blocking).

Ensures that all operations submitted to the given stream after this call will not begin execution until this event has completed.

Parameters:

stream – CUDA stream that should wait for the event.

Throws:

rapidsmpf::cuda_error – if cudaStreamWaitEvent fails.

cudaEvent_t const &value() const noexcept#

Access the underlying CUDA event handle.

Returns:

Const reference to the underlying cudaEvent_t.

inline operator cudaEvent_t const&() const noexcept#

Implicit conversion operator to the CUDA event handle.

Returns:

Const reference to the underlying cudaEvent_t.

Public Static Functions

static std::shared_ptr<CudaEvent> make_shared_record(
cuda::stream_ref stream,
unsigned flags = cudaEventDisableTiming
)#

Create and record a CUDA event on a given stream.

Convenience factory that constructs a shared CudaEvent with the specified creation flags, immediately records it on the provided stream, and returns it as a std::shared_ptr.

Parameters:
  • stream – CUDA stream on which to record the event.

  • flags – CUDA event creation flags.

Throws:

rapidsmpf::cuda_error – if event creation or recording fails.

Returns:

A shared pointer to the newly created and recorded CudaEvent.

struct cuda_error : public std::runtime_error#
#include <error.hpp>

Exception thrown when a CUDA error is encountered.

class bad_alloc : public std::bad_alloc#
#include <error.hpp>

Exception thrown when a RapidsMPF allocation fails.

Subclassed by rapidsmpf::out_of_memory, rapidsmpf::reservation_error

Public Functions

inline explicit bad_alloc(char const *msg)#

Construct a bad_alloc with the error message.

Parameters:

msg – Message to be associated with the exception.

inline explicit bad_alloc(std::string const &msg)#

Construct a bad_alloc with the error message.

Parameters:

msg – Message to be associated with the exception.

inline char const *what() const noexcept override#

Return the explanatory string.

Returns:

The explanatory string.

class out_of_memory : public rapidsmpf::bad_alloc#
#include <error.hpp>

Exception thrown when RapidsMPF runs out of memory.

This exception should only be thrown when we know a resource is out of memory.

Public Functions

inline explicit out_of_memory(char const *msg)#

Construct an out_of_memory with the error message.

Parameters:

msg – Message to be associated with the exception.

inline explicit out_of_memory(std::string const &msg)#

Construct an out_of_memory with the error message.

Parameters:

msg – Message to be associated with the exception.

class reservation_error : public rapidsmpf::bad_alloc#
#include <error.hpp>

Exception thrown when a memory reservation fails in RapidsMPF.

This exception is thrown when attempting to reserve memory fails, or when an existing reservation is insufficient for a requested allocation. It does not necessarily indicate that the system is out of physical memory, only that the reservation contract could not be satisfied.

Public Functions

inline explicit reservation_error(char const *msg)#

Construct a reservation_error with an error message.

Parameters:

msg – Message to be associated with the exception.

inline explicit reservation_error(std::string const &msg)#

Construct a reservation_error with an error message.

Parameters:

msg – Message to be associated with the exception.

template<typename BackRef>
class BackRefMixin#
#include <back_ref_mixin.hpp>

Mixin that lets copies of the this object keep an external object reference of type BackRef alive.

Note

Copying an instance without calling set_backref() throws std::bad_weak_ptr.

Template Parameters:

BackRef – Type of the external object reference.

Public Functions

BackRefMixin() noexcept = default#

Construct a mixin with no installed back-reference.

inline BackRefMixin(BackRefMixin const &other)#

Acquire shared ownership of other's back-referenced owner.

After construction, this references the same owner as other and keeps it alive for as long as this lives.

Parameters:

other – Mixin to copy from.

Throws:

std::bad_weak_ptr – if other is uninstalled, or is installed but its owner has been destroyed.

inline BackRefMixin &operator=(BackRefMixin const &other)#

Copy assignment with the same semantics as the copy constructor.

If the assignment throws, this is left unchanged.

Parameters:

other – Mixin to copy from.

Throws:

std::bad_weak_ptr – if other is uninstalled, or is installed but its owner has been destroyed.

Returns:

Reference to this.

BackRefMixin(BackRefMixin&&) noexcept = default#

Move constructor.

BackRefMixin &operator=(BackRefMixin&&) noexcept = default#

Move assignment operator.

Returns:

Reference to this mixin.

inline bool operator==(BackRefMixin const &other) const noexcept#

Owner-based equality.

Parameters:

other – Mixin to compare against.

Returns:

true if both mixins reference the same back-referenced owner, or are both uninstalled.

inline void set_backref(std::weak_ptr<BackRef> backref)#

Install a back-reference on this instance.

After this call, every subsequent copy of this will hold shared ownership of backref's referent for the lifetime of the copy.

Installation only affects this and copies made from it afterwards; copies that already exist are not retroactively back-referenced. To guarantee the back-reference is honored by every observable copy, install it before the host object becomes reachable to any code that may copy it.

Parameters:

backref – Non-empty weak reference to the back-referenced owner.

Throws:

std::invalid_argument – if backref is empty, or if a back-reference is already installed on this instance.

class Buffer#
#include <buffer.hpp>

Buffer representing device or host memory.

A Buffer holds either device memory or host memory, determined by its memory type at construction. See device_buffer_types and host_buffer_types for the sets of memory types that result in device-backed and host-backed storage.

Buffers are stream ordered and have an associated CUDA stream (see stream()). All work that reads or writes the buffer must either be enqueued on that stream or be synchronized with it before accessing the memory. For example, when passing the buffer to a non-stream aware API (e.g., MPI or host-only code), the caller must ensure that the most recent write has completed before the hand off. This can be done by synchronizing the buffer’s stream or by checking is_latest_write_done().

To obtain an rmm::device_buffer from a Buffer, first ensure that the buffer’s memory type is one of the types listed in device_buffer_types (moving the buffer if necessary), then call release_device_buffer().

Note

The constructors are private. Buffers are created through BufferResource.

Public Types

using DeviceBufferT = std::unique_ptr<rmm::device_buffer>#

Storage type for a device buffer.

using HostBufferT = std::unique_ptr<HostBuffer>#

Storage type for a host buffer.

Public Functions

std::byte const *data() const#

Access the underlying memory buffer (host or device memory).

Throws:
  • std::logic_error – if the buffer does not manage any memory.

  • std::logic_error – If the buffer is locked.

Returns:

A const pointer to the underlying host or device memory.

template<typename F>
inline auto write_access(
F &&f
) -> std::invoke_result_t<F, std::byte*, cuda::stream_ref>#

Provides stream-ordered write access to the buffer.

Calls f with a pointer to the buffer’s memory and the buffer’s stream (i.e., this->stream()).

The callable must be invocable as:

  • R(std::byte*, cuda::stream_ref).

All work performed by f must be stream-ordered on the buffer’s stream. Enqueuing work on any other stream without synchronizing with the buffer’s stream before and after the call is undefined behavior. In other words, f must behave as a single stream-ordered operation, similar to issuing one rapidsmpf::cuda_memcpy_async on the buffer’s stream. For non-stream-aware integrations, use exclusive_data_access().

After f returns, an event is recorded on the buffer’s stream, establishing the new “latest write” for this buffer.

// Snippet: copy data from `src_ptr` into `buffer` on the buffer's stream.
buffer.write_access([&](std::byte* buffer_ptr, cuda::stream_ref stream) {
  assert(buffer.stream().get() == stream.get());
  RAPIDSMPF_CUDA_TRY(rapidsmpf::cuda_memcpy_async(
      buffer_ptr,
      src_ptr,
      num_bytes,
      stream
  ));
});

Warning

The pointer is valid only for the duration of the call. Using it outside of f is undefined behavior.

Template Parameters:

F – Callable type.

Parameters:

f – Callable that accepts (std::byte*, cuda::stream_ref).

Throws:

std::logic_error – If the buffer is locked.

Returns:

Whatever f returns (void if none).

std::byte *exclusive_data_access()#

Acquire non-stream-ordered exclusive access to the buffer’s memory.

Alternative to write_access(). Acquires an internal exclusive lock so that any other access through the Buffer API (including write_access()) will fail with std::logic_error while the lock is held. The lock remains held until unlock() is called. This lock is not a concurrency mechanism; it only prevents accidental access to the Buffer through the rest of the Buffer API while locked.

Use this when integrating with non-stream-aware consumer APIs that require a raw pointer and cannot be expressed as work on a CUDA stream (e.g., MPI, blocking host I/O).

See also

write_access(), is_locked(), unlock()

Note

Prefer write_access(...) if you can express the operation as a single callable on a stream, even if that requires manually synchronizing the stream before the callable returns.

Warning

The Buffer does not track read access to its underlying storage, and so one should be aware of write-after-read anti-dependencies when obtaining exclusive access.

Throws:
  • std::logic_error – If the buffer is already locked.

  • std::logic_error – If is_latest_write_done() != true.

Returns:

Pointer to the underlying storage.

void unlock()#

Release the exclusive lock acquired by exclusive_data_access().

inline constexpr MemoryType mem_type() const#

Get the memory type of the buffer.

Throws:

std::logic_error – if the buffer is not initialized.

Returns:

The memory type of the buffer.

inline constexpr cuda::stream_ref stream() const noexcept#

Get the associated CUDA stream.

All operations must either use this stream or synchronize with it before accessing the underlying data (both host and device memory).

Returns:

The associated CUDA stream.

inline CudaEvent const &latest_write_event() const noexcept#

Get the CUDA event that tracks the latest write into the buffer.

Returns:

The CUDA event that tracks the latest write into the buffer.

void rebind_stream(cuda::stream_ref new_stream)#

Rebind the buffer to a new CUDA stream.

Changes the buffer’s associated stream to new_stream and ensures proper synchronization: new_stream will wait for any pending work on the current stream before proceeding. The underlying storage stream (e.g., the stream of an rmm::device_buffer or HostBuffer) is also updated.

// Example: merge buffers from different streams onto a single stream.
Buffer buffer_a = ...;  // associated with stream_a
Buffer buffer_b = ...;  // associated with stream_b

buffer_a.rebind_stream(merged_stream);
buffer_b.rebind_stream(merged_stream);

// Both buffers now use merged_stream with proper synchronization
buffer_copy(buffer_a, buffer_b, size);

Parameters:

new_stream – The new CUDA stream.

Throws:

std::logic_error – If the buffer is locked.

bool is_latest_write_done() const#

Check whether the buffer’s most recent write has completed.

Returns whether the CUDA event that tracks the most recent write into this buffer has been signaled.

Use this to guard non-stream-ordered consumer-APIs that do not accept a CUDA stream (e.g., MPI sends/receives, host-side reads).

// Example: send the buffer via MPI (non-stream-ordered).
if (buffer.is_latest_write_done()) {
  MPI_Isend(buffer.data(), buffer.size(), MPI_BYTE, dst, tag, comm, &req);
} else {
  // Ensure completion before handing to MPI.
  buffer.stream().sync();
  MPI_Isend(buffer.data(), buffer.size(), MPI_BYTE, dst, tag, comm, &req);
}

Note

This is a non-blocking, point-in-time status check and is subject to TOCTOU races: another thread may enqueue additional writes after this returns true. Ensure no further writes are enqueued, or establish stronger synchronization (e.g., synchronize the buffer’s stream) before using the buffer.

Warning

This check only confirms that there are no pending writes to the Buffer. Pending stream-ordered reads from the Buffer are not tracked and therefore one should be aware of write-after-read anti-dependencies when using this check to pass from stream-ordered to non-stream-ordered code.

Throws:

std::logic_error – If the buffer is locked.

Returns:

true if the last recorded write event has completed; false otherwise.

Buffer(Buffer&&) = delete#

Delete move and copy constructors and assignment operators.

Public Members

std::size_t const size#

The size of the buffer in bytes.

Public Static Attributes

static constexpr std::array<MemoryType, 1> device_buffer_types = {MemoryType::DEVICE}#

Memory types suitable for constructing a device backed buffer.

A buffer may use DeviceBufferT only if its memory type is listed here. This ensures that the buffer is backed by memory that behaves as device accessible memory.

static constexpr std::array<MemoryType, 2> host_buffer_types = {MemoryType::HOST, MemoryType::PINNED_HOST}#

Memory types suitable for constructing a host backed buffer.

A buffer may use HostBufferT only if its memory type is listed here. This ensures that the buffer is backed by memory that behaves as host accessible memory.

class StreamPool#
#include <buffer_resource.hpp>

Pool of non-blocking CUDA streams.

The pool is backed by RMM until CUDA Core provides an owning stream-pool implementation. Its public interface returns CUDA Core stream references so callers do not need to bridge between stream abstractions.

Public Functions

inline explicit StreamPool(std::size_t num_streams)#

Construct a pool of non-blocking CUDA streams.

Parameters:

num_streams – Number of streams to create.

inline explicit StreamPool(
std::shared_ptr<rmm::cuda_stream_pool> pool
)#

Construct a stream pool backed by an existing RMM pool.

Parameters:

pool – RMM stream pool providing stream ownership.

inline cuda::stream_ref get_stream() const#

Get the next stream from the pool.

Returns:

A CUDA Core reference to the selected stream.

inline cuda::stream_ref get_stream(std::size_t stream_id) const#

Get a stream from the pool by index.

Parameters:

stream_id – Index of the stream to retrieve.

Returns:

A CUDA Core reference to the selected stream.

inline std::size_t get_pool_size() const#

Get the number of streams in the pool.

Returns:

Number of streams in the pool.

class BufferResource : public std::enable_shared_from_this<BufferResource>#
#include <buffer_resource.hpp>

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.

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.

Note

BufferResource instances must be constructed through create() or from_options(), both of which return a std::shared_ptr<BufferResource>. Direct construction is disabled.

Note

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

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

Returns:

Reference to this.

BufferResource &operator=(BufferResource&&) = delete#

BufferResource is non-movable.

Returns:

Reference to this.

rmm::device_async_resource_ref device_mr() noexcept#

Get the device memory resource.

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.

CCCL’s lifetime semantic

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.

Returns:

rmm::device_async_resource_ref to the device memory resource.

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

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.

Returns:

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

rmm::host_async_resource_ref host_mr() noexcept#

Get the RMM host memory resource.

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.

Returns:

Reference to the RMM resource used for host allocations.

rmm::host_device_async_resource_ref pinned_mr()#

Get the RMM pinned host memory resource.

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.

Throws:

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

Returns:

Reference to the RMM resource used for pinned host allocations.

std::optional<PinnedMemoryResource> 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.

std::int64_t 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.

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

std::int64_t memory_available_for_reservation(
MemoryType mem_type
) const#

Returns the memory available to a new reservation, in bytes.

A snapshot of memory_available(mem_type) minus the outstanding reservations of that memory type. May be negative.

Parameters:

mem_type – The memory type to query.

Returns:

The memory available for reservation in bytes.

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.

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.

Throws:

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

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

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

Throws:

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

Returns:

The memory reservation.

template<std::ranges::input_range Range>
inline 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.

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.

Throws:

std::runtime_error – if no memory reservation was made.

Returns:

A memory reservation.

inline MemoryReservation reserve_or_fail(
std::size_t size,
MemoryType mem_type
)#

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.

Throws:

std::runtime_error – if no memory reservation was made.

Returns:

A memory reservation.

std::size_t 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.

Throws:

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

Returns:

The remaining size of the reserved memory after consumption.

std::unique_ptr<Buffer> make_buffer(
std::size_t size,
cuda::stream_ref 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.

Throws:
  • std::invalid_argument – if the memory type does not match the reservation.

  • rapidsmpf::reservation_error – if size exceeds the size of the reservation.

Returns:

A unique pointer to the allocated Buffer.

std::unique_ptr<Buffer> make_buffer(
cuda::stream_ref 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.

std::unique_ptr<Buffer> move(
std::unique_ptr<rmm::device_buffer> data,
cuda::stream_ref 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.

std::unique_ptr<Buffer> 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:
  • bufferBuffer to move.

  • reservation – Memory reservation used if a copy is required.

Throws:

rapidsmpf::reservation_error – If the allocation size exceeds the reservation.

Returns:

Unique pointer to the resulting Buffer.

std::unique_ptr<rmm::device_buffer> 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.

Throws:
  • std::invalid_argument – If the reservation’s memory type isn’t device memory.

  • rapidsmpf::reservation_error – if the memory requirement exceeds the reservation.

Returns:

A unique pointer to the resulting device buffer.

std::unique_ptr<HostBuffer> 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:
  • bufferBuffer to move.

  • reservation – Memory reservation used if a copy is required.

Throws:
  • std::invalid_argument – If the reservation’s memory type isn’t host memory.

  • rapidsmpf::reservation_error – If the allocation size exceeds the reservation.

Returns:

Unique pointer to the resulting host buffer.

std::shared_ptr<StreamPool> const &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 CUDA stream pool.

SpillManager &spill_manager()#

Gets a reference to the spill manager used.

Returns:

Reference to the SpillManager instance.

std::shared_ptr<Statistics> statistics() const noexcept#

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

Returns:

Shared pointer the Statistics instance.

Public Static 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<StreamPool> stream_pool = std::make_shared<StreamPool>(16),
std::shared_ptr<Statistics> statistics = Statistics::disabled()
)#

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.

  • statisticsStatistics instance used for runtime metrics.

Throws:

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

Returns:

A newly constructed BufferResource owned by std::shared_ptr.

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.

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.

class ContentDescription#
#include <content_description.hpp>

Description of an object’s content.

A ContentDescription encapsulates resource-related information about an object’s content such as memory sizes, memory types, and spillability.

In RapidsMPF, an object’s content refers to the actual data associated with the object, not its metadata or auxiliary information. Typically, this content is represented by one or more Buffer instances that may reside in different memory spaces (e.g., host or device). It is also this content that is subject to spilling and that typically accounts for the majority of an object’s overall memory footprint.

The spillability state of an object is treated as an all-or-nothing property. While one could imagine an object with a mix of spillable and non-spillable content, this distinction is intentionally simplified in RapidsMPF.

Public Types

enum class Spillable : bool#

Indicates whether the content is spillable.

Values:

enumerator NO#
enumerator YES#

Public Functions

template<std::ranges::input_range Range = std::initializer_list<std::pair<MemoryType, std::size_t>>>
inline explicit ContentDescription(
Range &&sizes,
Spillable spillable
)#

Construct a content description from a range of (MemoryType, size) pairs.

Memory types omitted from the input are initialized to zero.

ContentDescription desc{
    std::array{
        std::pair{MemoryType::HOST,   1024UL},
        std::pair{MemoryType::DEVICE, 2048UL}
    },
    ContentDescription::Spillable::YES
};

Template Parameters:

Range – A range whose value type is convertible to std::pair<MemoryType, std::size_t>.

Parameters:
  • sizes – Range of (MemoryType, size) pairs representing content sizes.

  • spillable – Whether the content is spillable to slower memory tiers.

inline constexpr ContentDescription(
Spillable spillable = Spillable::NO
)#

Construct a description with all sizes zero and a given spillability.

Useful when you need a zero-sized content description or when building a content description iteratively.

Parameters:

spillable – Whether the content are spillable.

inline constexpr std::size_t &content_size(
MemoryType mem_type
) noexcept#

Access (read/write) the size for a specific memory type.

Parameters:

mem_type – The memory type entry to access.

Returns:

Reference to the size (in bytes) for the given memory type.

inline constexpr std::size_t content_size(
MemoryType mem_type
) const noexcept#

Get the size for a specific memory type.

Parameters:

mem_type – The memory type entry to access.

Returns:

Size (in bytes) for the given memory type.

inline constexpr MemoryType principal_memory_type() const noexcept#

Returns the principal memory type of the content.

The principal memory type is the first memory type with non-zero content, evaluated in the preference order defined by MEMORY_TYPES.

If all content sizes are zero, the method returns MemoryType::DEVICE. Zero-sized content usually means the Message contains only metadata, or that the content size is unknown or irrelevant. In such cases the choice of memory type has no practical effect. We default to MemoryType::DEVICE because, when spillable() == false, it is the only memory type that meaningfully applies.

Returns:

The first memory type with non-zero content, or MemoryType::DEVICE if none contain content.

inline constexpr std::size_t content_size() const noexcept#

Get the total content size across all memory types.

Computes the sum of all per-memory-type content sizes. This represents the total size (in bytes) of the object’s content across host, device, and any other memory types.

Returns:

Total size (in bytes) across all memory types.

inline constexpr bool spillable() const noexcept#
Returns:

Whether the content can be spilled.

inline constexpr bool operator==(
ContentDescription const &other
) const noexcept#

Equality comparison.

Parameters:

other – The content description to compare against.

Returns:

true if both descriptions are equal; otherwise false.

class HostBuffer#
#include <host_buffer.hpp>

Block of host memory.

Public Functions

HostBuffer(
std::size_t size,
cuda::stream_ref stream,
cuda::mr::any_resource<cuda::mr::host_accessible> mr
)#

Allocate a new host buffer.

If size is greater than zero, memory is allocated using the provided memory resource and stream. If size is zero, the buffer is created empty.

Parameters:
  • size – Number of bytes to allocate.

  • stream – CUDA stream on which allocation and deallocation occur.

  • mr – Host-accessible memory resource used for allocation. Taken by value so the buffer shares ownership of the resource (e.g. bumps the refcount when constructed from a shared-ownership resource); an implicit conversion from rmm::host_async_resource_ref is also supported.

HostBuffer(HostBuffer &&other) noexcept#

Move constructor.

Transfers ownership of the underlying memory. The moved-from object becomes empty.

Parameters:

other – The buffer to move from.

HostBuffer &operator=(HostBuffer &&other)#

Move assignment operator.

Transfers ownership of the underlying memory. The current buffer must be empty before assignment.

Parameters:

other – The buffer to move from.

Throws:

std::invalid_argument – if this buffer is already initialized.

Returns:

Reference to this object.

void deallocate_async() noexcept#

Stream-ordered deallocates the buffer, if allocated.

After deallocation the buffer becomes empty. It is safe to call this method multiple times.

cuda::stream_ref stream() const noexcept#

Get the CUDA stream associated with this buffer.

Returns:

CUDA stream view.

std::size_t size() const noexcept#

Get the size of the buffer in bytes.

Returns:

Number of bytes in the buffer.

bool empty() const noexcept#

Check whether the buffer is empty.

Returns:

True if no memory is allocated.

std::byte *data() noexcept#

Get a pointer to the buffer data.

Returns:

Pointer to the underlying memory.

std::byte const *data() const noexcept#

Get a const pointer to the buffer data.

Returns:

Const pointer to the underlying memory.

std::vector<std::uint8_t> copy_to_uint8_vector() const#

Copy the contents of the buffer into a host std::vector.

This is primarily intended for debugging or testing. It performs a stream synchronization before and after the copy.

Returns:

A vector containing the copied bytes.

void set_stream(cuda::stream_ref stream)#

Set the associated CUDA stream.

This function only updates the buffer’s associated CUDA stream. It does not perform any stream synchronization or establish ordering guarantees between the previous stream and stream. Callers must ensure that any work previously enqueued on the old stream that may access the buffer has completed or is otherwise properly synchronized before switching streams.

Parameters:

stream – The new CUDA stream.

Public Static Functions

static HostBuffer from_uint8_vector(
std::vector<std::uint8_t> const &data,
cuda::stream_ref stream,
rmm::host_async_resource_ref mr
)#

Construct a HostBuffer by copying data from a std::vector<std::uint8_t>.

A new buffer is allocated using the provided memory resource and stream. This helper is intended for tests and small debug utilities.

Parameters:
  • data – Source vector containing the bytes to copy.

  • stream – CUDA stream used for allocation and copy.

  • mr – Host memory resource used to allocate the buffer.

Returns:

A new HostBuffer containing a copy of data.

static HostBuffer from_owned_vector(
std::vector<std::uint8_t> &&data,
cuda::stream_ref stream
)#

Construct a HostBuffer by taking ownership of a std::vector<std::uint8_t>.

The buffer takes ownership of the vector’s memory. The vector is moved into internal storage and will be destroyed when the HostBuffer is destroyed.

Parameters:
  • data – Vector to take ownership of (will be moved).

  • stream – CUDA stream to associate with this buffer.

Returns:

A new HostBuffer owning the vector’s memory.

static HostBuffer from_rmm_device_buffer(
std::unique_ptr<rmm::device_buffer> pinned_host_buffer,
cuda::stream_ref stream
)#

Construct a HostBuffer by taking ownership of an rmm::device_buffer.

The buffer takes ownership of the device buffer. The caller must ensure that the device buffer contains host-accessible memory (e.g., pinned host memory allocated via a managed or pinned memory resource).

Warning

The caller is responsible for ensuring the device buffer’s memory is host-accessible. Using this with non-host-accessible device memory will result in a std::invalid_argument exception.

Warning

The caller is responsible to ensure pinned_host_buffer is of pinned host memory type. If non-pinned host buffer leads to an irrecoverable condition.

Parameters:
  • pinned_host_buffer – Device buffer to take ownership of.

  • stream – CUDA stream to associate with this buffer.

Throws:
  • std::invalid_argument – if pinned_host_buffer is null.

  • std::logic_error – if pinned_host_buffer is not of pinned memory host memory type (see warning for details).

Returns:

A new HostBuffer owning the device buffer’s memory.

class HostMemoryResource : public rapidsmpf::BackRefMixin<BufferResource>#
#include <host_memory_resource.hpp>

Host memory resource using standard CPU allocation.

This resource allocates pageable host memory using the new and delete operators. It is intended for use with cuda::mr::resource and related facilities, and advertises the cuda::mr::host_accessible property.

For sufficiently large allocations (>4 MiB), this resource also issues a best-effort request to enable Transparent Huge Pages (THP) on the allocated region. THP can improve device-host memory transfer performance for large buffers. The hint is applied via madvise(MADV_HUGEPAGE) and may be ignored by the kernel depending on system configuration or resource availability.

Public Functions

HostMemoryResource(HostMemoryResource const&) = default#

Copyable.

HostMemoryResource(HostMemoryResource&&) = default#

Movable.

HostMemoryResource &operator=(HostMemoryResource const&) = default#

Copy assignment.

Returns:

Reference to this object after assignment.

HostMemoryResource &operator=(HostMemoryResource&&) = default#

Move assignment.

Returns:

Reference to this object after assignment.

inline void *allocate_sync(std::size_t, std::size_t)#

Synchronously allocates host memory is disabled.

Always use stream-ordered allocators in RapidsMPF.

Throws:

std::invalid_argument – Always.

Returns:

N/A.

inline void deallocate_sync(void*, std::size_t, std::size_t)#

Synchronously deallocates host memory is disabled.

Throws:

std::invalid_argument – Always.

void *allocate(
cuda::stream_ref stream,
std::size_t size,
std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT
)#

Allocates host memory associated with a CUDA stream.

Parameters:
  • stream – CUDA stream associated with the allocation.

  • size – Number of bytes to at least allocate.

  • alignment – Required alignment.

Throws:
  • std::bad_alloc – If the allocation fails.

  • std::invalid_argument – If alignment is not a valid alignment.

Returns:

Pointer to the allocated memory.

void deallocate(
cuda::stream_ref stream,
void *ptr,
std::size_t size,
std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT
) noexcept#

Deallocates host memory associated with a CUDA stream.

Synchronizes stream before deallocating the memory with the delete operator.

Parameters:
  • stream – CUDA stream associated with operations that used ptr.

  • ptr – Pointer to the memory to deallocate. May be nullptr.

  • size – Number of bytes previously allocated at ptr.

  • alignment – Alignment originally used for the allocation.

inline bool operator==(
[[maybe_unused]] HostMemoryResource const &other
) const noexcept#

Compares this resource to another resource.

All instances are stateless and interchangeable, so this always returns true.

Parameters:

other – The resource to compare with.

Returns:

true

inline bool operator!=(
[[maybe_unused]] HostMemoryResource const &other
) const noexcept#

Compares this resource to another resource.

All instances are stateless and interchangeable, so this always returns true.

Parameters:

other – The resource to compare with.

Returns:

true

Friends

inline friend void get_property(
HostMemoryResource const&,
cuda::mr::host_accessible
) noexcept#

Enables the cuda::mr::host_accessible property.

This property declares that a HostMemoryResource provides host accessible memory

class MemoryReservation#
#include <memory_reservation.hpp>

Represents a reservation for future memory allocation.

A reservation is returned by BufferResource::reserve and must be used when allocating buffers through the BufferResource.

Public Functions

~MemoryReservation() noexcept#

Destructor for the memory reservation.

Cleans up resources associated with the reservation.

void clear() noexcept#

Clear the remaining size of the reservation.

MemoryReservation split(std::size_t size)#

Split off a sub-reservation of size bytes.

Reduces this reservation by size and returns a new reservation of that size on the same buffer resource and memory type. The total reserved by the buffer resource is unchanged, the bytes are only moved between the two reservations.

Useful for scoping part of a reservation to the allocation it covers, since the returned reservation releases its bytes when it goes out of scope.

Parameters:

size – The number of bytes to split off.

Throws:

rapidsmpf::reservation_error – if size exceeds the remaining size.

Returns:

The new reservation.

MemoryReservation(MemoryReservation &&o)#

Move constructor for MemoryReservation.

Parameters:

o – The memory reservation to move from.

MemoryReservation &operator=(MemoryReservation &&o) noexcept#

Move assignment operator for MemoryReservation.

Parameters:

o – The memory reservation to move from.

Returns:

A reference to the updated MemoryReservation.

MemoryReservation(MemoryReservation const&) = delete#

A memory reservation is not copyable.

inline constexpr std::size_t size() const noexcept#

Get the remaining size of the reserved memory.

Returns:

The size of the reserved memory in bytes.

inline constexpr MemoryType mem_type() const noexcept#

Get the type of memory associated with this reservation.

Returns:

The type of memory associated with this reservation.

inline constexpr BufferResource *br() const noexcept#

Get the buffer resource associated with this reservation.

Returns:

The buffer resource associated with this reservation.

struct PackedData#
#include <packed_data.hpp>

Bag of bytes with metadata suitable for sending over the wire.

Contains arbitrary GPU data and host side metadata indicating how the data should be interpreted.

Public Functions

inline PackedData(
std::unique_ptr<std::vector<std::uint8_t>> metadata,
std::unique_ptr<Buffer> data
)#

Construct from metadata and GPU data, taking ownership.

Parameters:
  • metadata – Host-side metadata describing the GPU data.

  • data – Pointer to GPU data.

PackedData(PackedData&&) = default#

PackedData is moveable.

PackedData &operator=(PackedData&&) = default#

Move assignment.

Returns:

Moved this.

inline bool empty() const#

Check if the packed data is empty.

Returns:

True if the packed data is empty, false otherwise.

inline cuda::stream_ref stream() const#

Get the stream associated with the data buffer.

Returns:

The CUDA stream.

inline PackedData copy(MemoryReservation &reservation) const#

Create a deep copy of the packed data.

Parameters:

reservation – Memory reservation to use.

Returns:

A new PackedData instance containing a deep copy of both the data buffer and metadata.

Public Members

std::unique_ptr<std::vector<std::uint8_t>> metadata#

The metadata.

std::unique_ptr<Buffer> data#

The GPU data.

struct PinnedPoolProperties#
#include <pinned_memory_resource.hpp>

Properties for configuring a pinned memory pool.

Public Members

std::size_t initial_pool_size = 0#

Initial size of the pool. Initial size is important for pinned memory performance, especially for the first allocation. (See BM_PinnedFirstAlloc_InitialPoolSize benchmark.)

std::optional<std::size_t> max_pool_size = std::nullopt#

Maximum size of the pool. std::nullopt means no limit.

int numa_id = get_current_numa_node()#

NUMA node from which pinned memory should be allocated. Defaults to the NUMA node of the calling thread.

class PinnedMemoryResource : public cuda::mr::shared_resource<detail::RmmResourceAdaptorImpl<cuda::pinned_memory_pool>>, public rapidsmpf::BackRefMixin<BufferResource>#
#include <pinned_memory_resource.hpp>

Memory resource that provides pinned (page-locked) host memory using a pool.

Inherits from cuda::mr::shared_resource<RmmResourceAdaptorImpl<cuda::pinned_memory_pool>>, which holds the pool directly inside the shared control block — no extra heap allocation for the pool itself. Copies share the same underlying pool and memory statistics.

This resource allocates and deallocates pinned host memory asynchronously through CUDA streams. It offers higher bandwidth and lower latency for device transfers compared to regular pageable host memory.

Public Functions

inline void *allocate(
cuda::stream_ref stream,
std::size_t size,
std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT
)#

Allocates pinned host memory associated with a CUDA stream.

Parameters:
  • stream – CUDA stream associated with the allocation.

  • size – Number of bytes to at least allocate.

  • alignment – Required alignment.

Throws:
  • std::bad_alloc – If the allocation fails.

  • std::invalid_argument – If alignment is not a valid alignment.

Returns:

Pointer to the allocated memory.

inline void deallocate(
cuda::stream_ref stream,
void *ptr,
std::size_t size,
std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT
) noexcept#

Deallocates pinned host memory associated with a CUDA stream.

Parameters:
  • stream – CUDA stream associated with operations that used ptr.

  • ptr – Pointer to the memory to deallocate. May be nullptr.

  • size – Number of bytes previously allocated at ptr.

  • alignment – Alignment originally used for the allocation.

inline bool operator==(
PinnedMemoryResource const &other
) const noexcept#

Equality comparison.

Parameters:

other – The other resource to compare.

Returns:

True if the two resources share the same underlying shared state.

inline std::int64_t current_allocated() const noexcept#

Returns the total number of currently allocated bytes.

Returns:

The total number of currently allocated bytes.

inline ScopedMemoryRecord get_main_memory_record() const#

Returns the main memory record for the pinned pool.

Returns:

The main memory record for the pinned pool.

inline constexpr PinnedPoolProperties const &properties(
) const noexcept#

Returns the properties used to configure the pool.

Returns:

The properties used to configure the pool.

Friends

inline friend void get_property(
PinnedMemoryResource const&,
cuda::mr::host_accessible
) noexcept#

Enables the cuda::mr::host_accessible property.

struct ScopedMemoryRecord#
#include <scoped_memory_record.hpp>

Memory statistics for a specific scope.

Note

Is trivially copyable.

Public Functions

std::int64_t num_total_allocs() const noexcept#

Returns the total number of allocations performed.

Returns:

The total number of allocations.

std::int64_t num_current_allocs() const noexcept#

Returns the number of currently active (non-deallocated) allocations.

Returns:

The number of active allocations.

std::int64_t current() const noexcept#

Returns the current memory usage in bytes.

Current usage is the total bytes currently allocated but not yet deallocated.

Returns:

The current memory usage in bytes.

std::int64_t total() const noexcept#

Returns the total number of bytes allocated.

This value accumulates over time and is not reduced by deallocations.

Returns:

The total number of bytes allocated.

std::int64_t peak() const noexcept#

Returns the peak memory usage (in bytes).

The peak represents the highest value reached by current memory usage over the lifetime of the scope, including contributions from merged subscopes.

Returns:

The peak memory usage in bytes.

std::int64_t max() const noexcept#

Returns the size of the largest single allocation (in bytes).

Returns:

The largest single allocation in bytes.

void record_allocation(std::int64_t nbytes)#

Records a memory allocation event.

Updates the allocation counters and memory usage statistics, and adjusts peak usage if the new current usage exceeds the previous peak.

Note

Is not thread-safe.

Parameters:

nbytes – The number of bytes allocated.

void record_deallocation(std::int64_t nbytes)#

Records a memory deallocation event.

Reduces the current memory usage.

Note

Is not thread-safe.

Parameters:

nbytes – The number of bytes deallocated.

ScopedMemoryRecord &add_subscope(ScopedMemoryRecord const &subscope)#

Merge the memory statistics of a subscope into this record.

Combines the memory tracking data from a nested scope (subscope) into this record, updating statistics to include the subscope’s allocations, peaks, and totals.

This method treats the given record as a child scope nested within this scope, so peak usage is updated considering the current usage plus the subscope’s peak, reflecting hierarchical (inclusive) memory usage accounting.

This design allows memory scopes to be organized hierarchically, so when querying a parent scope, its statistics are inclusive of all nested scopes — similar to hierarchical memory profiling tools. However, it assumes that the parent scope’s statistics remain constant during the execution of the subscope.

See also

add_scope()

Parameters:

subscope – The scoped memory record representing a completed nested region.

Returns:

Reference to this object after merging the subscope.

ScopedMemoryRecord &add_scope(ScopedMemoryRecord const &scope)#

Merge the memory statistics of another scope into this one.

Unlike add_subscope(), this method treats the given scope as a peer or sibling, rather than a nested child. It aggregates totals and allocation counts and updates peak usage by taking the maximum peaks independently.

This is useful for combining memory statistics across multiple independent scopes, such as from different threads or non-nested regions.

See also

add_subscope()

Parameters:

scope – The scope to combine with this one.

Returns:

Reference to this object after summing.

class SpillManager#
#include <spill_manager.hpp>

Manages memory spilling to free up device memory when needed.

The SpillManager is responsible for registering, prioritizing, and executing spill functions to ensure efficient memory management.

Public Types

using SpillFunction = std::function<std::size_t(std::size_t)>#

Spill function type.

A spill function receives a requested spill size in bytes and returns the actual number of bytes spilled.

Spill functions must not capture owning references to the BufferResource that owns this SpillManager, either directly or indirectly through objects that allocate from BufferResource::device_mr(). Doing so creates a reference cycle:

BufferResource -> SpillManager -> SpillFunction -> BufferResource

Spill functions should capture only non-owning references or raw pointers. Owners registering spill functions should additionally call remove_spill_function() during destruction before releasing any buffer-owning members.

using SpillFunctionID = std::size_t#

Represents a unique identifier for a registered spill function.

Public Functions

SpillManager(
BufferResource *br,
std::optional<Duration> periodic_spill_check = std::nullopt
)#

Constructs a SpillManager instance.

Parameters:
  • brBuffer resource used to retrieve current available memory.

  • periodic_spill_check – Enable periodic spill checks. A dedicated thread continuously checks and perform spilling based on the current available memory as reported by the buffer resource. The value of periodic_spill_check is used as the pause between checks. If std::nullopt, no periodic spill check is performed.

~SpillManager()#

Destructor for SpillManager.

Cleans up any allocated resources and stops periodic spill checks if active (this will block until all spill functions has stopped).

SpillFunctionID add_spill_function(
SpillFunction spill_function,
int priority
)#

Adds a spill function with a given priority to the spill manager.

The spill function is prioritized according to the specified priority value.

Parameters:
  • spill_function – The spill function to be added.

  • priority – The priority level of the spill function (higher values indicate higher priority).

Returns:

The id assigned to the newly added spill function.

void remove_spill_function(SpillFunctionID fid)#

Removes a spill function from the spill manager.

This method unregisters the spill function associated with the given ID and removes it from the priority list. If no more spill functions remain, the periodic spill thread is paused.

Parameters:

fid – The id of the spill function to be removed.

std::size_t spill(std::size_t amount)#

Initiates spilling to free up a specified amount of memory.

This method iterates through registered spill functions in priority order, invoking them until at least the requested amount of memory has been spilled or no more spilling is possible.

Parameters:

amount – The amount of memory (in bytes) to spill.

Returns:

The actual amount of memory spilled (in bytes), which may be more, less or equal to the requested.

std::size_t spill_to_make_headroom(std::int64_t headroom = 0)#

Attempts to free up memory by spilling data until the requested headroom is reservable.

The headroom measurement is a snapshot, so a later reserve() of headroom bytes is not guaranteed to succeed. Spilling is performed in order of the function priorities until the requested headroom is reservable or no more spilling is possible. Spilling reduces allocations, never outstanding reservations.

Parameters:

headroom – The target amount of headroom (in bytes). A negative headroom triggers spilling only once the memory available for reservation drops below headroom.

Returns:

The actual amount of memory spilled (in bytes), which may be less than requested if there is insufficient spillable data, but may also be more or equal to requested depending on the sizes of spillable data buffers.

std::optional<std::size_t> try_spill_to_make_headroom(
std::int64_t headroom = 0
)#

Non-blocking version of spill_to_make_headroom().

Returns immediately instead of waiting when the spill lock is unavailable. Intended for pollers that retry, such as the streaming layer’s memory reservation loop.

Parameters:

headroom – The target amount of headroom (in bytes). A negative headroom triggers spilling only once the memory available for reservation drops below headroom.

Returns:

The actual amount of memory spilled (in bytes), or std::nullopt if no spill was attempted. A std::nullopt result does not imply that spilling is impossible or that another spill is in progress. Callers should retry.

class OwningWrapper#
#include <owning_wrapper.hpp>

Utility class to store an arbitrary type-erased object while another object is alive.

When sending messages through Channels from Python, we typically need to keep various Python objects alive since the matching C++ objects only hold views.

For example, a C++ message payload may hold non-owning views into a Python-owned object. If we want to allow creation of such objects in Python with the ability to sink them on the C++ side we cannot rely on the Python side keeping that owner alive (the reference disappears!). Similarly when we send a message through a Channel the sender will, once pushed into the channel, drop the reference to the message payload and so, again, we need some way of keeping the payload alive.

To square this circle, such C++ objects have an OwningWrapper slot that stores a type-erased pointer with, as far as we are concerned, unique ownership semantics. When this object is destroyed, the custom deleter runs and can do whatever deallocation is necessary.

Warning

Behaviour is undefined if the unique ownership semantic is not respected. The deleter may be called from any thread at any time, the implementer of the deleter is responsible for correct synchronisation with (for example) the Python GIL. Furthermore, the deleter may not throw: if an error occurs, the only safe thing to do is std::terminate.

Warning

When using this OwningWrapper inside a C++ object, make sure it is constructed first and destructed last.

Public Types

using deleter_type = void (*)(void*)#

Callback used to delete the owned object.

Public Functions

inline explicit OwningWrapper(void *obj, deleter_type deleter)#

Take ownership and responsibility for the destruction of an object.

Parameters:
  • obj – Type-erased object to own.

  • deleter – Function called to destruct the object.

inline void *release() noexcept#

Release ownership of the underlying pointer.

Returns:

Pointer to object.

inline void *get() const noexcept#
Returns:

Get access to the underlying pointer.

class ProgressThread#
#include <progress_thread.hpp>

A progress thread that can execute arbitrary functions.

Execute each of the registered arbitrary functions in a separate thread. The functions are executed in the order they were registered, and a newly registered function will only execute for the first time in the next iteration of the progress thread.

Public Types

enum ProgressState#

The progress state of a function, can be either InProgress or Done.

Values:

enumerator InProgress#
enumerator Done#
typedef std::uint64_t FunctionIndex#

The sequential index of a function within a ProgressThread.

typedef std::uintptr_t ProgressThreadAddress#

The address of a ProgressThread instance.

typedef std::function<ProgressState()> Function#

The function type supported by ProgressThread, returning the progress state of the function.

Public Functions

ProgressThread(
std::shared_ptr<Statistics> statistics = Statistics::disabled(),
Duration sleep = std::chrono::microseconds{1}
)#

Construct a new progress thread that can handle multiple functions.

Parameters:
  • statistics – The statistics instance to use (disabled by default).

  • sleep – The duration to sleep between each progress loop iteration. If 0, the thread yields execution instead of sleeping. Anecdotally, a 1 us sleep time (the default) is sufficient to avoid starvation and get smooth progress.

void stop()#

Stop the thread, blocking until all functions are done.

FunctionID add_function(Function &&function)#

Insert a function to process as part of the event loop.

Note

This function does not need to be thread-safe if not used in multiple progress threads.

Parameters:

function – The function to register.

Returns:

The unique ID of the function that was registered.

void remove_function(FunctionID function_id)#

Remove a function and stop processing it as part of the event loop.

This function blocks until the function is done (returning ProgressState::Done).

Parameters:

function_id – The unique function ID returned by add_function.

Throws:

std::logic_error – if the function was not registered with this ProgressThread or was already removed.

void pause()#

Pause the progress thread.

Note

This blocks until the thread is actually paused.

void resume()#

Resume the progress thread.

bool is_running() const#

Check if the progress thread is currently running.

Returns:

true if the thread is running, false otherwise.

std::shared_ptr<Statistics> statistics() const noexcept#
Returns:

The statistics instance on this progress thread.

struct FunctionID#
#include <progress_thread.hpp>

The unique ID of a function registered with ProgressThread. Composed of the ProgressThread address and a sequential function index.

Public Functions

FunctionID() = default#

Construct a FunctionID with an invalid address.

For a valid object the constructor that takes thread_addr and index must be used.

Note

This is the default constructor.

inline constexpr FunctionID(
ProgressThreadAddress thread_addr,
FunctionIndex index
)#

Construct a new FunctionID.

Parameters:
  • thread_addr – The address of the ProgressThread instance

  • index – The sequential index of the function

inline constexpr bool is_valid() const#

Check if the FunctionID is valid.

Returns:

True if the FunctionID is valid, false otherwise.

Public Members

ProgressThreadAddress thread_address = {ProgressThreadAddress(0)}#

The address of the ProgressThread instance.

FunctionIndex function_index = {0}#

The sequential index of the function.

class FunctionState#
#include <progress_thread.hpp>

Store state of a function.

Public Functions

explicit FunctionState(Function &&function)#

Construct state of a function.

Parameters:

function – The function to execute.

void operator()()#

Execute the function.

Note

Calling this from multiple threads is not allowed.

Public Members

Function function#

The function to execute.

bool is_done = {false}#

Whether the function has completed.

class RmmResourceAdaptor : public cuda::mr::shared_resource<detail::RmmResourceAdaptorImpl<cuda::mr::any_resource<cuda::mr::device_accessible>>>, public rapidsmpf::BackRefMixin<BufferResource>#
#include <rmm_resource_adaptor.hpp>

A RMM memory resource adaptor tailored to RapidsMPF.

This adaptor wraps a primary device memory resource and adds memory usage tracking (lifetime stats plus per-thread scoped records).

This class is copyable and shares ownership of its internal state via cuda::mr::shared_resource.

It is not possible to construct an RmmResourceAdaptor directly. Instead, obtain one by creating a BufferResource with the memory resource you wish to use for allocations and then obtain the adaptor with device_mr_adaptor().

Public Functions

inline bool operator==(
RmmResourceAdaptor const &other
) const noexcept#

Equality comparison.

Two adaptors are equal iff they share the same underlying shared state. Because adaptors are privately constructed by BufferResource and a back-reference cannot be overwritten, sharing the same shared state implies referencing the same owning BufferResource.

Parameters:

other – The other adaptor to compare.

Returns:

True if both adaptors refer to the same shared resource instance.

rmm::device_async_resource_ref get_upstream_resource() const noexcept#

Get a reference to the primary upstream resource.

Returns:

Reference to the RMM memory resource.

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

std::int64_t current_allocated() const noexcept#

Get the total current allocated memory through this resource.

Returns:

Total number of currently allocated bytes.

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

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

Throws:

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

Returns:

The scope that was just ended.

Friends

inline friend void get_property(
RmmResourceAdaptor const&,
cuda::mr::device_accessible
) noexcept#

Tag this resource as device-accessible for the CCCL concept.

class Statistics : public std::enable_shared_from_this<Statistics>#
#include <statistics.hpp>

Tracks statistics across rapidsmpf operations.

Two naming concepts are used throughout this class:

  • Stat name: identifies an individual Stat accumulator, as passed to add_stat(), get_stat(), add_bytes_stat(), and add_duration_stat(). Stats are pure numeric accumulators with no associated rendering information. Examples: "spill-time", "spill-bytes".

  • Report entry name: the label of a formatted line in report(), passed to add_report_entry(). An entry names one or more stats and a Formatter that selects how those stats are rendered. When the entry covers a single stat, the report entry name and stat name are typically identical. Example: "spill" (aggregating "spill-bytes" and "spill-time").

Formatters are a fixed, predefined set (see Statistics::Formatter).

Statistics stats;

// Associate two stats with a predefined multi-stat formatter.
stats.add_report_entry(
    "copy-device-to-host",                // report entry name
    {"copy-device-to-host-bytes",
     "copy-device-to-host-time",
     "copy-device-to-host-stream-delay"},
    Statistics::Formatter::MemoryThroughput
);

stats.add_bytes_stat("spill-bytes", 1024);    // helper: registers Bytes entry
stats.add_duration_stat("spill-time", 0.5s);  // helper: registers Duration entry

auto s = stats.get_stat("spill-bytes");  // retrieve without formatter
std::cout << stats.report();

Public Types

enum class Formatter : std::uint8_t#

Identifies a predefined formatter used by report().

Each formatter consumes a fixed number of Stat entries and renders them into a human-readable string.

Available formatters (examples):

  • Default (1 stat): “123”

  • Bytes (1 stat): “1.2 GiB | avg 300 MiB”

  • Duration (1 stat): “2.5 ms | avg 600 us”

  • HitRate (1 stat): “42/100 (hits/lookups)”

  • MemoryThroughput (3 stats: bytes, time, stream-delay), where stream-delay is the wall-clock gap between CPU submission and GPU execution of the operation: “1.2 GiB | 2.5 ms | 480 GiB/s | avg-stream-delay 10 us”

_Count is an internal sentinel — always keep it last.

Values:

enumerator Default#
enumerator Bytes#
enumerator Duration#
enumerator HitRate#
enumerator MemoryThroughput#
enumerator _Count#

Sentinel; must remain last.

enum class Mode : std::uint8_t#

Selects whether a newly constructed Statistics instance tracks data or is a no-op.

Values:

enumerator Enabled#

Statistics tracking is active.

enumerator Disabled#

All operations are no-ops; can be toggled on via enable().

Public Functions

inline bool enabled() const noexcept#

Checks if statistics tracking is enabled.

Returns:

True if statistics tracking is active, otherwise False.

inline void enable() noexcept#

Enable statistics tracking for this instance.

inline void disable() noexcept#

Disable statistics tracking for this instance.

std::string report(ReportArgs report_args) const#

Generates a formatted report of all collected statistics.

Every registered report entry always produces a line. If all the stats it references have been recorded, the entry’s Formatter renders the values; otherwise the line reads “No data collected”. Statistics not covered by any report entry are shown with Formatter::Default (raw numeric value, optionally annotated with the count). All entries are sorted alphabetically.

Note

If any statistics are collected via stream-ordered timing (e.g. through record_copy()), all relevant CUDA streams must be synchronized before calling this method. Otherwise, some timing statistics may not yet have been recorded, causing entries to read “No data collected” or imprecise statistics.

Parameters:

report_args – Report options. See ReportArgs.

Returns:

Formatted statistics report.

inline std::string report() const#

Overload with all-default options. Equivalent to report(ReportArgs{}).

Returns:

Formatted statistics report.

void write_json(std::ostream &os) const#

Writes a JSON representation of all collected statistics to a stream.

Values are written as raw numbers (count, sum, max). Formatter metadata is not emitted — use report() for the human-readable rendering.

Parameters:

os – Output stream to write to.

Throws:

std::invalid_argument – If any stat name or memory record name contains characters that require JSON escaping (double quotes, backslashes, or ASCII control characters 0x00–0x1F).

void write_json(std::filesystem::path const &filepath) const#

Writes a JSON report of all collected statistics to a file.

Parameters:

filepath – Path to the output file. Created or overwritten.

Throws:

std::ios_base::failure – If the file cannot be opened or writing fails.

std::shared_ptr<Statistics> copy() const#

Creates a deep copy of this Statistics object.

Note

Memory records are not copied.

Returns:

A shared pointer to the new copy.

std::vector<std::uint8_t> serialize() const#

Serializes the stats and report entries to a binary byte vector.

Note

Memory records are not serialized.

Returns:

A vector of bytes containing the serialized statistics.

Stat get_stat(std::string const &name) const#

Retrieves a statistic by name.

Parameters:

name – Name of the statistic.

Returns:

The requested statistic.

void add_stat(std::string const &name, double value)#

Adds a numeric value to the named statistic.

Creates the statistic if it doesn’t exist. Does not associate any formatter with the stat — use add_report_entry() (or a helper like add_bytes_stat()) for that.

Parameters:
  • name – Name of the statistic.

  • value – Value to add.

void add_report_entry(
std::string const &report_entry_name,
std::initializer_list<std::string_view> stat_names,
Formatter formatter
)#

Associate a formatter with one or more named statistics for report rendering.

First-wins: if a report entry is already registered under report_entry_name, this call has no effect. The entry appears in report() as a single line; if any stat it references is missing, the line reads “No data collected”.

Parameters:
  • report_entry_name – Report entry name.

  • stat_names – Names of the stats this entry aggregates. Caller is responsible for passing the number of stats the chosen formatter expects; a mismatch surfaces as std::out_of_range when report() renders the entry.

  • formatter – Predefined formatter to render the entry with.

void add_report_entry(
std::string const &report_entry_name,
std::vector<std::string> stat_names,
Formatter formatter
)#

Associate a formatter with one or more named statistics for report rendering.

First-wins: if a report entry is already registered under report_entry_name, this call has no effect. The entry appears in report() as a single line; if any stat it references is missing, the line reads “No data collected”.

Overload for callers whose stat names come from a runtime container (e.g. the Python bindings).

Parameters:
  • report_entry_name – Report entry name.

  • stat_names – Names of the stats this entry aggregates. Caller is responsible for passing the number of stats the chosen formatter expects; a mismatch surfaces as std::out_of_range when report() renders the entry.

  • formatter – Predefined formatter to render the entry with.

void add_bytes_stat(std::string const &name, std::size_t nbytes)#

Adds a byte count to the named statistic.

Registers a Formatter::Bytes report entry named name if no report entry already exists under that name, then adds nbytes to the named statistic.

Parameters:
  • name – Name of the statistic.

  • nbytes – Number of bytes to add.

void add_duration_stat(std::string const &name, Duration seconds)#

Adds a duration to the named statistic.

Registers a Formatter::Duration report entry named name if no report entry already exists under that name, then adds seconds to the named statistic.

Parameters:
  • name – Name of the statistic.

  • seconds – Duration in seconds to add.

void record_copy(
MemoryType src,
MemoryType dst,
std::size_t nbytes,
StreamOrderedTiming &&timing
)#

Record byte count and wall-clock duration for a memory copy operation.

Records three statistics entries for "copy-{src}-to-{dst}":

  • "-bytes" — the number of bytes copied.

  • "-time" — the copy duration, recorded in stream order.

  • "-stream-delay" — time between CPU submission and GPU execution of the copy, recorded in stream order.

All three entries are aggregated into a single combined report line under the name "copy-{src}-to-{dst}", showing total bytes, total time, bandwidth, and average stream delay.

Parameters:
  • src – Source memory type.

  • dst – Destination memory type.

  • nbytes – Number of bytes copied.

  • timing – A StreamOrderedTiming that should be started just before the copy was enqueued on the stream. Its stop_and_record() is called here to enqueue the stop callback.

void record_alloc(
MemoryType mem_type,
std::size_t nbytes,
StreamOrderedTiming &&timing
)#

Record size and wall-clock duration for a buffer allocation.

Records three statistics entries for "alloc-{memtype}":

  • "-bytes" — the number of bytes allocated.

  • "-time" — the allocation duration, recorded in stream order.

  • "-stream-delay" — time between CPU submission and GPU execution, recorded in stream order.

All three entries are aggregated into a single combined report line showing total bytes, total time, throughput, and average stream delay.

Parameters:
  • mem_type – Memory type of the allocation.

  • nbytes – Number of bytes allocated.

  • timing – A StreamOrderedTiming constructed just before the allocation was issued. Its stop_and_record() is called here.

std::vector<std::string> list_stat_names() const#

Get the names of all statistics.

Returns:

A vector of all statistic names.

void clear()#

Clears all statistics.

Note

Memory profiling records and report entries are not cleared.

MemoryRecorder create_memory_recorder(
any_device_resource mr,
std::string name
)#

Creates a scoped memory recorder for the given name.

Parameters:
  • mr – Type-erased device memory resource. Recording is only active when the underlying resource is an RmmResourceAdaptor.

  • name – Name of the scope.

Returns:

A MemoryRecorder instance. If !enabled() or mr is not backed by an RmmResourceAdaptor, returns a no-op recorder.

std::unordered_map<std::string, MemoryRecord> const &get_memory_records(
) const#

Retrieves all memory profiling records stored by this instance.

Returns:

A reference to a map from record name to memory usage data.

Public Static Functions

static std::shared_ptr<Statistics> create(Mode mode = Mode::Enabled)#

Creates a Statistics instance.

Parameters:

mode – Selects whether tracking starts enabled or disabled. See Mode. Disabled instances can be toggled on later via enable().

Returns:

A shared pointer to a newly constructed Statistics instance.

static inline std::shared_ptr<Statistics> disabled()#

Returns a disabled Statistics instance which can be enabled later.

Returns:

A Statistics instance with tracking disabled.

static std::shared_ptr<Statistics> from_options(
config::Options options
)#

Construct from configuration options.

Parameters:

options – Configuration options.

Returns:

A shared pointer to the constructed Statistics instance.

static std::shared_ptr<Statistics> deserialize(
std::span<std::uint8_t const> data
)#

Deserializes a Statistics object from a binary byte vector.

Note

The resulting object has no memory records.

Parameters:

data – The serialized statistics data.

Throws:

std::invalid_argument – If the data is malformed or truncated.

Returns:

A shared pointer to the reconstructed Statistics object.

static std::shared_ptr<Statistics> merge(
std::span<std::shared_ptr<Statistics> const> stats
)#

Merge a set of Statistics into a new instance.

For each stat name present across the inputs, the result contains the summed count, summed value, and the maximum of the recorded maxima. The result’s enabled() is true if any input is enabled. Memory records are not merged.

Report entries are unified by name. If multiple inputs contain the same report-entry name, their Formatter and stat_names must match; otherwise, this function throws std::invalid_argument to prevent silent rendering inconsistencies (especially across serialize/deserialize boundaries).

Parameters:

stats – Non-empty span of non-null Statistics instances to merge.

Throws:

std::invalid_argument – If stats is empty, contains a null pointer, or if inputs disagree on the formatter or stat-name set for a shared report entry.

Returns:

A new Statistics instance containing the merged data.

struct MemoryRecord#
#include <statistics.hpp>

Holds memory profiling information for a named scope.

Public Members

ScopedMemoryRecord scoped#

Scoped memory stats.

std::int64_t global_peak = {0}#

Peak global memory usage during the scope.

std::uint64_t num_calls = {0}#

Number of times the scope was invoked.

class MemoryRecorder#
#include <statistics.hpp>

RAII-style object for scoped memory usage tracking.

Automatically tracks memory usage between construction and destruction.

Public Functions

MemoryRecorder() = default#

Constructs a no-op MemoryRecorder.

MemoryRecorder(
std::shared_ptr<Statistics> stats,
RmmResourceAdaptor mr,
std::string name
)#

Constructs an active MemoryRecorder. Pushes a scoped record at construction; the destructor pops it and (if stats is still enabled) publishes it under name.

Parameters:
  • stats – Owning Statistics. Must not be null.

  • mr – RMM resource adaptor providing scoped memory statistics.

  • name – Name of the scope.

struct ReportArgs#
#include <statistics.hpp>

Named-argument struct for report().

All fields carry defaults so any subset may be supplied using designated initialisers:

stats.report({.mr = my_mr, .header = "Run 1:"});

Public Members

std::optional<any_device_resource> mr = std::nullopt#

Optional RMM resource adaptor used for memory profiling. When provided, a memory profiling section is included in the report. When std::nullopt, the memory profiling section shows “Disabled”.

std::optional<any_host_device_resource> pinned_mr = std::nullopt#

Optional pinned memory resource. When provided, a pinned memory section is included in the report.

std::string_view header = "Statistics:"#

Header line prepended to the report.

class Stat#
#include <statistics.hpp>

Represents a single tracked statistic.

Note

Stat is not thread-safe. Thread safety is provided by the enclosing Statistics object’s mutex.

Public Functions

Stat() = default#

Default-constructs a Stat.

Stat(std::size_t count, double value, double max)#

Constructs a Stat with explicit field values.

Parameters:
  • count – Number of updates.

  • value – Total accumulated value.

  • max – Maximum value seen.

auto operator<=>(Stat const&) const noexcept = default#

Three-way comparison operator.

Performs memberwise comparison of all data members.

Returns:

The ordering result of the memberwise comparison.

void add(double value)#

Adds a value to this statistic.

Parameters:

value – The value to add.

std::size_t count() const noexcept#

Returns the number of updates applied to this statistic.

Returns:

The number of times add() was called.

double value() const noexcept#

Returns the total accumulated value.

Returns:

The sum of all values added.

double max() const noexcept#

Returns the maximum value seen across all add() calls.

Returns:

The maximum value added, or negative infinity if add() was never called.

std::uint8_t *serialize(std::uint8_t *out) const#

Serializes this Stat to a byte buffer.

Parameters:

out – Pointer to the output buffer. Must have at least serialized_size() bytes available.

Returns:

Pointer past the last byte written.

Stat merge(Stat const &other) const#

Merges another Stat into this one, returning the combined result.

Counts and values are summed; the maximum is taken.

Parameters:

other – The Stat to merge with.

Returns:

A new Stat containing the merged result.

Public Static Functions

static inline constexpr std::size_t serialized_size() noexcept#

Returns the serialized size of this Stat in bytes.

We size each field individually rather than using sizeof(Stat) to avoid platform-dependent struct padding.

Returns:

The number of bytes needed to serialize this Stat.

static std::pair<Stat, std::span<std::uint8_t const>> deserialize(
std::span<std::uint8_t const> data
)#

Deserializes a Stat from a byte buffer.

Parameters:

data – The input buffer. Must contain at least serialized_size() bytes.

Throws:

std::invalid_argument – If the data is truncated.

Returns:

A pair of the deserialized Stat and the remaining unconsumed bytes.

class StreamOrderedTiming#
#include <stream_ordered_timing.hpp>

Stream-ordered wall-clock timer that records its result into Statistics.

Marks a start position in the CUDA stream on construction and a stop position when stop_and_record is called. The elapsed wall-clock time between those two stream positions is recorded into the supplied Statistics object once the stream reaches the stop marker — guaranteeing that the measurement covers exactly the work enqueued between the two calls, in stream order.

If statistics is disabled (e.g. constructed with Statistics::Mode::Disabled), the entire class is a no-op.

StreamOrderedTiming timing{stream, stats};
// ... enqueue GPU work on `stream` ...
timing.stop_and_record("my-operation-time");

Public Functions

StreamOrderedTiming(
cuda::stream_ref stream,
std::shared_ptr<Statistics> statistics
)#

Constructs a StreamOrderedTiming and marks the start position in the stream.

If statistics is disabled (e.g. constructed with Statistics::Mode::Disabled), this is a no-op and subsequent calls to stop_and_record will also be no-ops.

Parameters:
  • stream – The CUDA stream to time.

  • statistics – The Statistics object that will receive the duration entry.

void stop_and_record(
std::string const &name,
std::optional<std::string> stream_delay_name = std::nullopt
)#

Marks the stop position in the stream and schedules recording of the duration.

The stream-ordered duration (time between the start and stop stream positions) is recorded under name. If stream_delay_name is set, the stream delay — the wall-clock time between object construction and when the stream actually executed the start callback — is also recorded under that name. The stream delay reveals how far ahead the CPU is running relative to the GPU stream.

Both values are written to the Statistics object in stream order — i.e. only after all work enqueued between construction and this call has been reached by the stream. If the Statistics object is destroyed before that point, the recording is silently skipped.

Behaviour is undefined if this method is called more than once per StreamOrderedTiming instance.

Parameters:
  • name – Name of the stream-ordered duration statistic.

  • stream_delay_name – Name of the stream-delay statistic. If std::nullopt (the default), no stream-delay entry is written.

Public Static Functions

static void cancel_inflight_timings(Statistics const *statistics)#

Cancel all in-flight timings associated with a Statistics object.

Should be called when a Statistics object is about to be destroyed, to prevent dangling references from any in-flight stream callbacks. It is safe to call when no in-flight timings are present.

Note

If a stop callback has already executed before this function is called, the associated statistic may still be recorded. The guarantee is only that any still-pending callbacks are cancelled.

Parameters:

statistics – The Statistics object whose in-flight timings should be cancelled.

template<class ...Ts>
struct overloaded : public rapidsmpf::Ts#
#include <misc.hpp>

Helper for overloaded lambdas using std::visit.

MPI Utilities#

void rapidsmpf::mpi::init(int *argc, char ***argv)#

Helper to initialize MPI with threading support.

Parameters:
  • argc – Pointer to the number of arguments passed to the program.

  • argv – Pointer to the argument vector passed to the program.

bool rapidsmpf::mpi::is_initialized()#

Check if MPI is initialized.

Returns:

true If MPI is initialized.