From 5ee3a88f64b1a70feefbd8326d542c5a4a381edc Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Fri, 29 May 2026 08:55:19 +0200 Subject: [PATCH 01/15] Add and use BufferResource::device_mr_ref() --- cpp/benchmarks/bench_comm.cpp | 2 +- cpp/benchmarks/bench_shuffle.cpp | 4 +- cpp/benchmarks/streaming/data_generator.hpp | 2 +- cpp/benchmarks/streaming/ndsh/bench_read.cpp | 2 +- cpp/benchmarks/streaming/ndsh/concatenate.cpp | 2 +- cpp/benchmarks/streaming/ndsh/groupby.cpp | 2 +- cpp/benchmarks/streaming/ndsh/join.cpp | 14 +++---- cpp/benchmarks/streaming/ndsh/q01.cpp | 16 ++++---- cpp/benchmarks/streaming/ndsh/q03.cpp | 20 +++++----- cpp/benchmarks/streaming/ndsh/q04.cpp | 6 +-- cpp/benchmarks/streaming/ndsh/q09.cpp | 14 +++---- cpp/benchmarks/streaming/ndsh/q21.cpp | 28 +++++++------- cpp/benchmarks/streaming/ndsh/sort.cpp | 4 +- .../rapidsmpf/memory/buffer_resource.hpp | 38 +++++++++++++++++-- cpp/src/integrations/cudf/partition.cpp | 20 +++++----- cpp/src/memory/buffer_resource.cpp | 7 +++- cpp/src/streaming/cudf/bloom_filter.cpp | 6 +-- cpp/src/streaming/cudf/channel_metadata.cpp | 4 +- cpp/src/streaming/cudf/parquet.cpp | 2 +- cpp/src/streaming/cudf/table_chunk.cpp | 6 +-- cpp/tests/streaming/test_leaf_actor.cpp | 2 +- cpp/tests/streaming/test_partition.cpp | 2 +- cpp/tests/streaming/test_read_parquet.cpp | 2 +- cpp/tests/streaming/test_shuffler.cpp | 2 +- cpp/tests/test_allgather.cpp | 2 +- cpp/tests/test_buffer.cpp | 2 +- cpp/tests/utils.hpp | 2 +- 27 files changed, 125 insertions(+), 88 deletions(-) diff --git a/cpp/benchmarks/bench_comm.cpp b/cpp/benchmarks/bench_comm.cpp index 37a5039a1..076d88a59 100644 --- a/cpp/benchmarks/bench_comm.cpp +++ b/cpp/benchmarks/bench_comm.cpp @@ -220,7 +220,7 @@ Duration run( auto [res, _] = br->reserve(MemoryType::DEVICE, args.msg_size * 2, AllowOverbooking::YES); auto buf = br->make_buffer(args.msg_size, stream, res); - random_fill(*buf, br->device_mr()); + random_fill(*buf, br->device_mr_ref()); send_bufs.push_back(std::move(buf)); recv_bufs.push_back(br->make_buffer(args.msg_size, stream, res)); } diff --git a/cpp/benchmarks/bench_shuffle.cpp b/cpp/benchmarks/bench_shuffle.cpp index cd30993e3..9921bb271 100644 --- a/cpp/benchmarks/bench_shuffle.cpp +++ b/cpp/benchmarks/bench_shuffle.cpp @@ -302,7 +302,7 @@ rapidsmpf::Duration do_run( { RAPIDSMPF_NVTX_SCOPED_RANGE("Shuffling", total_num_partitions); if (args.enable_memory_profiler) { - RAPIDSMPF_MEMORY_PROFILE(statistics, br->device_mr(), "shuffling"); + RAPIDSMPF_MEMORY_PROFILE(statistics, br->device_mr_ref(), "shuffling"); } rapidsmpf::shuffler::Shuffler shuffler( comm, @@ -390,7 +390,7 @@ std::vector generate_input_partitions( size_lb, args.input_data_allow_overbooking ); cudf::table table = random_table( - num_columns, num_local_rows, min_val, max_val, stream, br->device_mr() + num_columns, num_local_rows, min_val, max_val, stream, br->device_mr_ref() ); input_partitions.emplace_back(transform_fn(std::move(table))); } diff --git a/cpp/benchmarks/streaming/data_generator.hpp b/cpp/benchmarks/streaming/data_generator.hpp index 9bb53795c..531c153d4 100644 --- a/cpp/benchmarks/streaming/data_generator.hpp +++ b/cpp/benchmarks/streaming/data_generator.hpp @@ -62,7 +62,7 @@ inline Actor random_table_generator( seq, std::make_unique( std::make_unique(random_table( - ncolumns, nrows, min_val, max_val, stream, ctx->br()->device_mr() + ncolumns, nrows, min_val, max_val, stream, ctx->br()->device_mr_ref() )), stream ) diff --git a/cpp/benchmarks/streaming/ndsh/bench_read.cpp b/cpp/benchmarks/streaming/ndsh/bench_read.cpp index 7d75f9687..53d32c9c0 100644 --- a/cpp/benchmarks/streaming/ndsh/bench_read.cpp +++ b/cpp/benchmarks/streaming/ndsh/bench_read.cpp @@ -413,7 +413,7 @@ int main(int argc, char** argv) { timings.push_back(compute.count()); auto statistics = ctx->statistics(); comm->logger()->print(statistics->report( - {.mr = ctx->br()->device_mr(), .pinned_mr = ctx->br()->try_pinned_mr()} + {.mr = ctx->br()->device_mr_ref(), .pinned_mr = ctx->br()->try_pinned_mr()} )); statistics->clear(); } diff --git a/cpp/benchmarks/streaming/ndsh/concatenate.cpp b/cpp/benchmarks/streaming/ndsh/concatenate.cpp index 39ce71720..1637f6ede 100644 --- a/cpp/benchmarks/streaming/ndsh/concatenate.cpp +++ b/cpp/benchmarks/streaming/ndsh/concatenate.cpp @@ -70,7 +70,7 @@ streaming::Actor concatenate( chunks.push_back(std::move(chunk)); } auto result = std::make_unique( - cudf::concatenate(views, concat_stream, ctx->br()->device_mr()), concat_stream + cudf::concatenate(views, concat_stream, ctx->br()->device_mr_ref()), concat_stream ); cuda_stream_join( chunks | std::views::transform([](auto&& chunk) { return chunk.stream(); }), diff --git a/cpp/benchmarks/streaming/ndsh/groupby.cpp b/cpp/benchmarks/streaming/ndsh/groupby.cpp index 5a13054c2..fd4234613 100644 --- a/cpp/benchmarks/streaming/ndsh/groupby.cpp +++ b/cpp/benchmarks/streaming/ndsh/groupby.cpp @@ -55,7 +55,7 @@ streaming::Actor chunkwise_group_by( cudf::groupby::groupby(table.select(keys), null_policy, cudf::sorted::NO); auto [keys, aggregated] = - grouper.aggregate(agg_requests, stream, ctx->br()->device_mr()); + grouper.aggregate(agg_requests, stream, ctx->br()->device_mr_ref()); std::ignore = std::move(chunk); auto result = keys->release(); for (auto&& a : aggregated) { diff --git a/cpp/benchmarks/streaming/ndsh/join.cpp b/cpp/benchmarks/streaming/ndsh/join.cpp index 65ff8be52..8aa6c35fa 100644 --- a/cpp/benchmarks/streaming/ndsh/join.cpp +++ b/cpp/benchmarks/streaming/ndsh/join.cpp @@ -69,7 +69,7 @@ coro::task broadcast( ); } else { RAPIDSMPF_EXPECTS(chunks.size() > 0, "No chunks in broadcast"); - auto result = cudf::concatenate(views, gather_stream, ctx->br()->device_mr()); + auto result = cudf::concatenate(views, gather_stream, ctx->br()->device_mr_ref()); // So that deallocation of the consitutent tables is stream-ordered wrt the // concatenation. cuda_stream_join( @@ -94,7 +94,7 @@ coro::task broadcast( auto chunk = co_await msg.release().make_available(ctx); auto pack = - cudf::pack(chunk.table_view(), chunk.stream(), ctx->br()->device_mr()); + cudf::pack(chunk.table_view(), chunk.stream(), ctx->br()->device_mr_ref()); auto packed_data = PackedData( std::move(pack.metadata), ctx->br()->move(std::move(pack.gpu_data), chunk.stream()) @@ -183,7 +183,7 @@ streaming::Message semi_join_chunk( ); auto match = joiner.semi_join( - left_chunk.table_view().select(left_on), chunk_stream, ctx->br()->device_mr() + left_chunk.table_view().select(left_on), chunk_stream, ctx->br()->device_mr_ref() ); ctx->logger()->debug( @@ -197,7 +197,7 @@ streaming::Message semi_join_chunk( indices, cudf::out_of_bounds_policy::DONT_CHECK, chunk_stream, - ctx->br()->device_mr() + ctx->br()->device_mr_ref() ) ->release(); @@ -244,7 +244,7 @@ streaming::Message inner_join_chunk( auto probe_table = right_chunk.table_view(); auto probe_keys = probe_table.select(right_on); auto [probe_match, build_match] = - joiner.inner_join(probe_keys, std::nullopt, chunk_stream, ctx->br()->device_mr()); + joiner.inner_join(probe_keys, std::nullopt, chunk_stream, ctx->br()->device_mr_ref()); cudf::column_view build_indices = cudf::device_span(*build_match); @@ -258,7 +258,7 @@ streaming::Message inner_join_chunk( build_indices, cudf::out_of_bounds_policy::DONT_CHECK, chunk_stream, - ctx->br()->device_mr() + ctx->br()->device_mr_ref() ) ->release(); // drop key columns from probe table. @@ -274,7 +274,7 @@ streaming::Message inner_join_chunk( probe_indices, cudf::out_of_bounds_policy::DONT_CHECK, chunk_stream, - ctx->br()->device_mr() + ctx->br()->device_mr_ref() ) ->release(), std::back_inserter(result_columns) diff --git a/cpp/benchmarks/streaming/ndsh/q01.cpp b/cpp/benchmarks/streaming/ndsh/q01.cpp index d4dd6975f..68ef2e298 100644 --- a/cpp/benchmarks/streaming/ndsh/q01.cpp +++ b/cpp/benchmarks/streaming/ndsh/q01.cpp @@ -136,7 +136,7 @@ rapidsmpf::streaming::Actor postprocess_group_by( co_await msg.release().make_available(ctx); auto stream = chunk.stream(); auto columns = - cudf::table{chunk.table_view(), stream, ctx->br()->device_mr()}.release(); + cudf::table{chunk.table_view(), stream, ctx->br()->device_mr_ref()}.release(); std::ignore = std::move(chunk); auto count = std::move(columns.back()); columns.pop_back(); @@ -150,7 +150,7 @@ rapidsmpf::streaming::Actor postprocess_group_by( cudf::binary_operator::TRUE_DIV, cudf::data_type(cudf::type_id::FLOAT64), stream, - ctx->br()->device_mr() + ctx->br()->device_mr_ref() ) ); } @@ -161,7 +161,7 @@ rapidsmpf::streaming::Actor postprocess_group_by( cudf::binary_operator::TRUE_DIV, cudf::data_type(cudf::type_id::FLOAT64), stream, - ctx->br()->device_mr() + ctx->br()->device_mr_ref() ) ); columns.push_back(std::move(count)); @@ -202,7 +202,7 @@ rapidsmpf::streaming::Actor select_columns_for_groupby( auto table = chunk.table_view(); // l_returnflag, l_linestatus, l_quantity, l_extendedprice auto result = - cudf::table(table.select({0, 1, 2, 3}), chunk_stream, ctx->br()->device_mr()) + cudf::table(table.select({0, 1, 2, 3}), chunk_stream, ctx->br()->device_mr_ref()) .release(); result.reserve(7); auto extendedprice = table.column(3); @@ -233,7 +233,7 @@ static __device__ void calculate_charge(double *charge, double discprice, double std::nullopt, cudf::output_nullability::PRESERVE, chunk_stream, - ctx->br()->device_mr() + ctx->br()->device_mr_ref() ) ); // charge @@ -248,12 +248,12 @@ static __device__ void calculate_charge(double *charge, double discprice, double std::nullopt, cudf::output_nullability::PRESERVE, chunk_stream, - ctx->br()->device_mr() + ctx->br()->device_mr_ref() ) ); // l_discount result.push_back( - std::make_unique(discount, chunk_stream, ctx->br()->device_mr()) + std::make_unique(discount, chunk_stream, ctx->br()->device_mr_ref()) ); co_await ch_out->send( rapidsmpf::streaming::to_message( @@ -435,7 +435,7 @@ int main(int argc, char** argv) { timings.push_back(compute.count()); auto statistics = ctx->statistics(); comm->logger()->print(statistics->report( - {.mr = ctx->br()->device_mr(), .pinned_mr = ctx->br()->try_pinned_mr()} + {.mr = ctx->br()->device_mr_ref(), .pinned_mr = ctx->br()->try_pinned_mr()} )); statistics->clear(); } diff --git a/cpp/benchmarks/streaming/ndsh/q03.cpp b/cpp/benchmarks/streaming/ndsh/q03.cpp index 89a922a80..81748dfaf 100644 --- a/cpp/benchmarks/streaming/ndsh/q03.cpp +++ b/cpp/benchmarks/streaming/ndsh/q03.cpp @@ -234,19 +234,19 @@ rapidsmpf::streaming::Actor select_columns_for_groupby( // o_orderkey result.push_back( std::make_unique( - table.column(0), chunk_stream, ctx->br()->device_mr() + table.column(0), chunk_stream, ctx->br()->device_mr_ref() ) ); // o_orderdate result.push_back( std::make_unique( - table.column(1), chunk_stream, ctx->br()->device_mr() + table.column(1), chunk_stream, ctx->br()->device_mr_ref() ) ); // o_shippriority result.push_back( std::make_unique( - table.column(2), chunk_stream, ctx->br()->device_mr() + table.column(2), chunk_stream, ctx->br()->device_mr_ref() ) ); auto extendedprice = table.column(3); @@ -270,7 +270,7 @@ static __device__ void calculate_revenue(double *revenue, double extprice, doubl std::nullopt, cudf::output_nullability::PRESERVE, chunk_stream, - ctx->br()->device_mr() + ctx->br()->device_mr_ref() ) ); co_await ch_out->send( @@ -311,7 +311,7 @@ rapidsmpf::streaming::Actor top_k_by( order, {}, chunk.stream(), - ctx->br()->device_mr() + ctx->br()->device_mr_ref() ); partials.push_back( cudf::gather( @@ -319,7 +319,7 @@ rapidsmpf::streaming::Actor top_k_by( cudf::split(indices->view(), {k}, chunk.stream()).front(), cudf::out_of_bounds_policy::DONT_CHECK, chunk.stream(), - ctx->br()->device_mr() + ctx->br()->device_mr_ref() ) ); chunk_streams.push_back(chunk.stream()); @@ -336,11 +336,11 @@ rapidsmpf::streaming::Actor top_k_by( std::ranges::transform(partials, std::back_inserter(views), [](auto& t) { return t->view(); }); - auto merged = cudf::merge(views, keys, order, {}, out_stream, ctx->br()->device_mr()); + auto merged = cudf::merge(views, keys, order, {}, out_stream, ctx->br()->device_mr_ref()); auto result = std::make_unique( cudf::slice(merged->view(), {0, 10}, out_stream), out_stream, - ctx->br()->device_mr() + ctx->br()->device_mr_ref() ); co_await ch_out->send( rapidsmpf::streaming::to_message( @@ -381,7 +381,7 @@ rapidsmpf::streaming::Actor fanout_bounded( std::make_unique( chunk.table_view().select(ch1_cols), chunk.stream(), - ctx->br()->device_mr() + ctx->br()->device_mr_ref() ), chunk.stream() ) @@ -697,7 +697,7 @@ int main(int argc, char** argv) { timings.push_back(compute.count()); auto statistics = ctx->statistics(); comm->logger()->print(statistics->report( - {.mr = ctx->br()->device_mr(), .pinned_mr = ctx->br()->try_pinned_mr()} + {.mr = ctx->br()->device_mr_ref(), .pinned_mr = ctx->br()->try_pinned_mr()} )); statistics->clear(); } diff --git a/cpp/benchmarks/streaming/ndsh/q04.cpp b/cpp/benchmarks/streaming/ndsh/q04.cpp index a76c3fd80..52ef7a178 100644 --- a/cpp/benchmarks/streaming/ndsh/q04.cpp +++ b/cpp/benchmarks/streaming/ndsh/q04.cpp @@ -151,7 +151,7 @@ rapidsmpf::streaming::Actor filter_lineitem( ) { rapidsmpf::streaming::ShutdownAtExit c{ch_in, ch_out}; co_await ctx->executor()->schedule(); - auto mr = ctx->br()->device_mr(); + auto mr = ctx->br()->device_mr_ref(); while (!ch_out->is_shutdown()) { auto msg = co_await ch_in->receive(); @@ -218,7 +218,7 @@ rapidsmpf::streaming::Actor fanout_bounded( std::make_unique( chunk.table_view().select(ch1_cols), chunk.stream(), - ctx->br()->device_mr() + ctx->br()->device_mr_ref() ), chunk.stream() ) @@ -520,7 +520,7 @@ int main(int argc, char** argv) { timings.push_back(compute.count()); auto statistics = ctx->statistics(); comm->logger()->print(statistics->report( - {.mr = ctx->br()->device_mr(), .pinned_mr = ctx->br()->try_pinned_mr()} + {.mr = ctx->br()->device_mr_ref(), .pinned_mr = ctx->br()->try_pinned_mr()} )); statistics->clear(); } diff --git a/cpp/benchmarks/streaming/ndsh/q09.cpp b/cpp/benchmarks/streaming/ndsh/q09.cpp index a52ae40f2..cf450fae2 100644 --- a/cpp/benchmarks/streaming/ndsh/q09.cpp +++ b/cpp/benchmarks/streaming/ndsh/q09.cpp @@ -175,7 +175,7 @@ rapidsmpf::streaming::Actor filter_part( std::shared_ptr ch_out ) { rapidsmpf::streaming::ShutdownAtExit c{ch_in, ch_out}; - auto mr = ctx->br()->device_mr(); + auto mr = ctx->br()->device_mr_ref(); while (true) { auto msg = co_await ch_in->receive(); if (msg.empty()) { @@ -233,7 +233,7 @@ rapidsmpf::streaming::Actor select_columns( // n_name result.push_back( std::make_unique( - table.column(0), chunk_stream, ctx->br()->device_mr() + table.column(0), chunk_stream, ctx->br()->device_mr_ref() ) ); result.push_back( @@ -241,7 +241,7 @@ rapidsmpf::streaming::Actor select_columns( table.column(5), cudf::datetime::datetime_component::YEAR, chunk_stream, - ctx->br()->device_mr() + ctx->br()->device_mr_ref() ) ); auto discount = table.column(2); @@ -267,7 +267,7 @@ static __device__ void calculate_amount(double *amount, double discount, double std::nullopt, cudf::output_nullability::PRESERVE, chunk_stream, - ctx->br()->device_mr() + ctx->br()->device_mr_ref() ) ); co_await ch_out->send( @@ -312,7 +312,7 @@ rapidsmpf::streaming::Actor round_sum_profit( 2, cudf::rounding_method::HALF_EVEN, chunk.stream(), - ctx->br()->device_mr() + ctx->br()->device_mr_ref() ); #pragma GCC diagnostic pop auto result = rapidsmpf::streaming::to_message( @@ -321,7 +321,7 @@ rapidsmpf::streaming::Actor round_sum_profit( std::make_unique( cudf::table_view({table.column(0), table.column(1), rounded->view()}), chunk.stream(), - ctx->br()->device_mr() + ctx->br()->device_mr_ref() ), chunk.stream() ) @@ -655,7 +655,7 @@ int main(int argc, char** argv) { timings.push_back(compute.count()); auto statistics = ctx->statistics(); comm->logger()->print(statistics->report( - {.mr = ctx->br()->device_mr(), .pinned_mr = ctx->br()->try_pinned_mr()} + {.mr = ctx->br()->device_mr_ref(), .pinned_mr = ctx->br()->try_pinned_mr()} )); statistics->clear(); } diff --git a/cpp/benchmarks/streaming/ndsh/q21.cpp b/cpp/benchmarks/streaming/ndsh/q21.cpp index 3b7e6a052..369ba46a1 100644 --- a/cpp/benchmarks/streaming/ndsh/q21.cpp +++ b/cpp/benchmarks/streaming/ndsh/q21.cpp @@ -98,7 +98,7 @@ rapidsmpf::streaming::Actor read_nation( constexpr auto name = "SAUDI ARABIA"; owner->push_back( std::make_shared( - name, /* is_valid = */ true, stream, ctx->br()->device_mr() + name, /* is_valid = */ true, stream, ctx->br()->device_mr_ref() ) ); owner->push_back( @@ -156,7 +156,7 @@ rapidsmpf::streaming::Actor read_orders( constexpr auto status = "F"; owner->push_back( std::make_shared( - status, /* is_valid = */ true, stream, ctx->br()->device_mr() + status, /* is_valid = */ true, stream, ctx->br()->device_mr_ref() ) ); owner->push_back( @@ -279,7 +279,7 @@ rapidsmpf::streaming::Actor filter_lineitem( cudf::binary_operator::GREATER, cudf::data_type(cudf::type_id::BOOL8), chunk.stream(), - ctx->br()->device_mr() + ctx->br()->device_mr_ref() ); co_await ch_out->send( rapidsmpf::streaming::to_message( @@ -289,7 +289,7 @@ rapidsmpf::streaming::Actor filter_lineitem( chunk.table_view().select({0, 1}), mask->view(), chunk.stream(), - ctx->br()->device_mr() + ctx->br()->device_mr_ref() ), chunk.stream() ) @@ -318,12 +318,12 @@ rapidsmpf::streaming::Actor filter_grouped_greater( auto mask = cudf::binary_operation( chunk.table_view().column(1), cudf::numeric_scalar( - 1, /* is_valid = */ true, chunk.stream(), ctx->br()->device_mr() + 1, /* is_valid = */ true, chunk.stream(), ctx->br()->device_mr_ref() ), cudf::binary_operator::GREATER, cudf::data_type(cudf::type_id::BOOL8), chunk.stream(), - ctx->br()->device_mr() + ctx->br()->device_mr_ref() ); latch->count_down(); co_await ch_out->send( @@ -334,7 +334,7 @@ rapidsmpf::streaming::Actor filter_grouped_greater( chunk.table_view().select({0}), mask->view(), chunk.stream(), - ctx->br()->device_mr() + ctx->br()->device_mr_ref() ), chunk.stream() ) @@ -362,12 +362,12 @@ rapidsmpf::streaming::Actor filter_grouped_equal( auto mask = cudf::binary_operation( chunk.table_view().column(1), cudf::numeric_scalar( - 1, /* is_valid = */ true, chunk.stream(), ctx->br()->device_mr() + 1, /* is_valid = */ true, chunk.stream(), ctx->br()->device_mr_ref() ), cudf::binary_operator::EQUAL, cudf::data_type(cudf::type_id::BOOL8), chunk.stream(), - ctx->br()->device_mr() + ctx->br()->device_mr_ref() ); co_await ch_out->send( rapidsmpf::streaming::to_message( @@ -377,7 +377,7 @@ rapidsmpf::streaming::Actor filter_grouped_equal( chunk.table_view().select({0}), mask->view(), chunk.stream(), - ctx->br()->device_mr() + ctx->br()->device_mr_ref() ), chunk.stream() ) @@ -415,7 +415,7 @@ rapidsmpf::streaming::Actor fanout_bounded( std::make_unique( chunk.table_view().select(ch1_cols), chunk.stream(), - ctx->br()->device_mr() + ctx->br()->device_mr_ref() ), chunk.stream() ) @@ -505,7 +505,7 @@ rapidsmpf::streaming::Actor slice( auto sliced_table = std::make_unique( cudf::slice(chunk.table_view(), {local_start, local_end})[0], chunk.stream(), - ctx->br()->device_mr() + ctx->br()->device_mr_ref() ); co_await ch_out->send( rapidsmpf::streaming::to_message( @@ -574,7 +574,7 @@ rapidsmpf::streaming::Actor populate_bloom_filter( ); auto stream = chunk.stream(); auto out = std::make_unique( - chunk.table_view().select(keys), stream, ctx->br()->device_mr() + chunk.table_view().select(keys), stream, ctx->br()->device_mr_ref() ); std::ignore = std::move(chunk); co_await passthrough->send( @@ -997,7 +997,7 @@ int main(int argc, char** argv) { timings.push_back(compute.count()); auto statistics = ctx->statistics(); comm->logger()->print(statistics->report( - {.mr = ctx->br()->device_mr(), .pinned_mr = ctx->br()->try_pinned_mr()} + {.mr = ctx->br()->device_mr_ref(), .pinned_mr = ctx->br()->try_pinned_mr()} )); statistics->clear(); } diff --git a/cpp/benchmarks/streaming/ndsh/sort.cpp b/cpp/benchmarks/streaming/ndsh/sort.cpp index 756d2bc65..b300ce4ed 100644 --- a/cpp/benchmarks/streaming/ndsh/sort.cpp +++ b/cpp/benchmarks/streaming/ndsh/sort.cpp @@ -36,7 +36,7 @@ rapidsmpf::streaming::Actor chunkwise_sort_by( order, null_order, chunk.stream(), - ctx->br()->device_mr() + ctx->br()->device_mr_ref() ); } else { return cudf::sort_by_key( @@ -45,7 +45,7 @@ rapidsmpf::streaming::Actor chunkwise_sort_by( order, null_order, chunk.stream(), - ctx->br()->device_mr() + ctx->br()->device_mr_ref() ); } }; diff --git a/cpp/include/rapidsmpf/memory/buffer_resource.hpp b/cpp/include/rapidsmpf/memory/buffer_resource.hpp index 5b01ef1fc..f0efd4efc 100644 --- a/cpp/include/rapidsmpf/memory/buffer_resource.hpp +++ b/cpp/include/rapidsmpf/memory/buffer_resource.hpp @@ -115,11 +115,43 @@ class BufferResource { ~BufferResource() noexcept = default; /** - * @brief Get the RMM device memory resource. + * @brief Get the device memory resource as an owning, type-erased handle. + * + * Returns a copy of the internally stored + * `cuda::mr::any_resource`. The copy is cheap — + * `any_resource` of a `shared_resource`-based adaptor shares state via an + * internal `shared_ptr`. Because the returned object owns its own + * refcounted handle on the underlying adaptor, it remains valid even if + * this `BufferResource` is destroyed. + * + * Use this when you need to **store** the device MR for the lifetime of + * another object (e.g. as the `mr` argument to `rmm::device_buffer`, which + * keeps its MR by value). For transient cuDF/RMM calls that take an + * `rmm::device_async_resource_ref`, prefer `device_mr_ref()`. + * + * @return An owning `cuda::mr::any_resource` handle on + * the device memory resource. + */ + [[nodiscard]] cuda::mr::any_resource + device_mr() const noexcept; + + /** + * @brief Get a non-owning reference to the device memory resource. + * + * Returns an `rmm::device_async_resource_ref` pointing into this + * `BufferResource`'s internal storage. The ref is cheap to obtain and + * binds directly to cuDF/RMM APIs that take + * `rmm::device_async_resource_ref` parameters. The caller must keep this + * `BufferResource` alive for as long as the returned ref is in use. + * + * This is the right accessor for the common case of passing the device + * MR to a cuDF/RMM function call. Use `device_mr()` (owning) when the MR + * needs to outlive this `BufferResource`. * - * @return Reference to the RMM resource used for device allocations. + * @return Non-owning `rmm::device_async_resource_ref` into BR-owned + * storage. */ - [[nodiscard]] rmm::device_async_resource_ref device_mr() const noexcept; + [[nodiscard]] rmm::device_async_resource_ref device_mr_ref() const noexcept; /** * @brief Get the RMM host memory resource. diff --git a/cpp/src/integrations/cudf/partition.cpp b/cpp/src/integrations/cudf/partition.cpp index b423c0803..a3f2c6241 100644 --- a/cpp/src/integrations/cudf/partition.cpp +++ b/cpp/src/integrations/cudf/partition.cpp @@ -36,10 +36,10 @@ partition_and_split( BufferResource* br, AllowOverbooking allow_overbooking ) { - RAPIDSMPF_MEMORY_PROFILE(br->statistics(), br->device_mr()); + RAPIDSMPF_MEMORY_PROFILE(br->statistics(), br->device_mr_ref()); if (table.num_rows() == 0) { // Return views of a copy of the empty `table`. - auto owner = std::make_unique(table, stream, br->device_mr()); + auto owner = std::make_unique(table, stream, br->device_mr_ref()); return { std::vector( safe_cast(num_partitions), owner->view() @@ -60,7 +60,7 @@ partition_and_split( hash_function, seed, stream, - br->device_mr() + br->device_mr_ref() ); reservation.clear(); @@ -88,7 +88,7 @@ std::unordered_map partition_and_pack( AllowOverbooking allow_overbooking ) { RAPIDSMPF_NVTX_FUNC_RANGE(); - RAPIDSMPF_MEMORY_PROFILE(br->statistics(), br->device_mr()); + RAPIDSMPF_MEMORY_PROFILE(br->statistics(), br->device_mr_ref()); RAPIDSMPF_EXPECTS(num_partitions > 0, "Need to split to at least one partition"); if (table.num_rows() == 0) { auto splits = @@ -108,7 +108,7 @@ std::unordered_map partition_and_pack( hash_function, seed, stream, - br->device_mr() + br->device_mr_ref() ); reservation.clear(); std::vector splits(split_points.begin() + 1, split_points.end() - 1); @@ -123,15 +123,15 @@ std::unordered_map split_and_pack( AllowOverbooking allow_overbooking ) { RAPIDSMPF_NVTX_FUNC_RANGE(); - RAPIDSMPF_MEMORY_PROFILE(br->statistics(), br->device_mr()); + RAPIDSMPF_MEMORY_PROFILE(br->statistics(), br->device_mr_ref()); std::unordered_map ret; // contiguous split does a deep-copy. Therefore, we need to reserve memory for // at least the size of the table. auto reservation = br->reserve_device_memory_and_spill( - cudf::packed_size(table, stream, br->device_mr()), allow_overbooking + cudf::packed_size(table, stream, br->device_mr_ref()), allow_overbooking ); - auto packed = cudf::contiguous_split(table, splits, stream, br->device_mr()); + auto packed = cudf::contiguous_split(table, splits, stream, br->device_mr_ref()); reservation.clear(); ret.reserve(packed.size()); for (shuffler::PartID i = 0; safe_cast(i) < packed.size(); i++) { @@ -153,7 +153,7 @@ std::unique_ptr unpack_and_concat( AllowOverbooking allow_overbooking ) { RAPIDSMPF_NVTX_FUNC_RANGE(); - RAPIDSMPF_MEMORY_PROFILE(br->statistics(), br->device_mr()); + RAPIDSMPF_MEMORY_PROFILE(br->statistics(), br->device_mr_ref()); // Let's find the total size of the partitions and how much of the packed data we // need to move to device memory (unspill). @@ -205,7 +205,7 @@ std::unique_ptr unpack_and_concat( } reservation = br->reserve_device_memory_and_spill(total_size, allow_overbooking); - return cudf::concatenate(unpacked, stream, br->device_mr()); + return cudf::concatenate(unpacked, stream, br->device_mr_ref()); } std::vector spill_partitions( diff --git a/cpp/src/memory/buffer_resource.cpp b/cpp/src/memory/buffer_resource.cpp index ba3025c8a..de02c9450 100644 --- a/cpp/src/memory/buffer_resource.cpp +++ b/cpp/src/memory/buffer_resource.cpp @@ -94,7 +94,12 @@ void BufferResource::set_memory_limit(MemoryType mem_type, std::int64_t limit) n ); } -rmm::device_async_resource_ref BufferResource::device_mr() const noexcept { +cuda::mr::any_resource +BufferResource::device_mr() const noexcept { + return device_mr_; +} + +rmm::device_async_resource_ref BufferResource::device_mr_ref() const noexcept { return rmm::device_async_resource_ref{ const_cast&>(device_mr_) }; diff --git a/cpp/src/streaming/cudf/bloom_filter.cpp b/cpp/src/streaming/cudf/bloom_filter.cpp index ac41e10ee..37b054f9d 100644 --- a/cpp/src/streaming/cudf/bloom_filter.cpp +++ b/cpp/src/streaming/cudf/bloom_filter.cpp @@ -26,7 +26,7 @@ Actor BloomFilter::build( co_await ch_in->shutdown_metadata(); co_await ch_out->shutdown_metadata(); auto const& br = ctx_->br(); - auto mr = br->device_mr(); + auto mr = br->device_mr_ref(); auto filter_stream = br->stream_pool().get_stream(); CudaEvent event; auto storage = rapidsmpf::BloomFilter::storage(num_filter_blocks_, filter_stream, mr); @@ -127,7 +127,7 @@ Actor BloomFilter::apply( 0 ); auto mask = filter.contains( - chunk.table_view().select(keys), chunk_stream, ctx_->br()->device_mr() + chunk.table_view().select(keys), chunk_stream, ctx_->br()->device_mr_ref() ); cuda_stream_join(stream, chunk_stream, &event); RAPIDSMPF_EXPECTS( @@ -142,7 +142,7 @@ Actor BloomFilter::apply( 0 }; auto result = cudf::apply_boolean_mask( - chunk.table_view(), mask_view, chunk_stream, ctx_->br()->device_mr() + chunk.table_view(), mask_view, chunk_stream, ctx_->br()->device_mr_ref() ); std::ignore = std::move(chunk); std::ignore = std::move(res); diff --git a/cpp/src/streaming/cudf/channel_metadata.cpp b/cpp/src/streaming/cudf/channel_metadata.cpp index 8d623cb1d..9bf0c1aa9 100644 --- a/cpp/src/streaming/cudf/channel_metadata.cpp +++ b/cpp/src/streaming/cudf/channel_metadata.cpp @@ -89,14 +89,14 @@ bool OrderScheme::boundaries_aligned_with( cudf::binary_operator::NULL_EQUALS, cudf::data_type{cudf::type_id::BOOL8}, stream, - br.device_mr() + br.device_mr_ref() ); auto result = cudf::reduce( eq->view(), *cudf::make_all_aggregation(), cudf::data_type{cudf::type_id::BOOL8}, stream, - br.device_mr() + br.device_mr_ref() ); auto& scalar = static_cast&>(*result); if (!scalar.value(stream)) { diff --git a/cpp/src/streaming/cudf/parquet.cpp b/cpp/src/streaming/cudf/parquet.cpp index 2edd16625..b7609ec5f 100644 --- a/cpp/src/streaming/cudf/parquet.cpp +++ b/cpp/src/streaming/cudf/parquet.cpp @@ -200,7 +200,7 @@ Message read_parquet_chunk( return to_message( sequence_number, std::make_unique( - cudf::io::read_parquet(options, stream, ctx->br()->device_mr()).tbl, + cudf::io::read_parquet(options, stream, ctx->br()->device_mr_ref()).tbl, stream ) ); diff --git a/cpp/src/streaming/cudf/table_chunk.cpp b/cpp/src/streaming/cudf/table_chunk.cpp index afef6385f..a08af79b8 100644 --- a/cpp/src/streaming/cudf/table_chunk.cpp +++ b/cpp/src/streaming/cudf/table_chunk.cpp @@ -181,7 +181,7 @@ TableChunk TableChunk::copy(MemoryReservation& reservation) const { auto statistics = br->statistics(); StreamOrderedTiming timing{stream(), statistics}; auto table = std::make_unique( - table_view(), stream(), br->device_mr() + table_view(), stream(), br->device_mr_ref() ); statistics->record_copy( MemoryType::DEVICE, MemoryType::DEVICE, nbytes, std::move(timing) @@ -217,7 +217,7 @@ TableChunk TableChunk::copy(MemoryReservation& reservation) const { // new host buffer. // TODO: use `cudf::chunked_pack()` with a bounce buffer. Currently, // `cudf::pack()` allocates device memory we haven't reserved. - auto packed_columns = cudf::pack(table_view(), stream(), br->device_mr()); + auto packed_columns = cudf::pack(table_view(), stream(), br->device_mr_ref()); auto packed_data = std::make_unique( std::move(packed_columns.metadata), br->move(std::move(packed_columns.gpu_data), stream()) @@ -269,7 +269,7 @@ std::unique_ptr TableChunk::into_packed_data(BufferResource* br) && ); // TODO: use `cudf::chunked_pack()` with a bounce buffer. Currently, // `cudf::pack()` allocates device memory we haven't reserved. - auto packed_columns = cudf::pack(table_view_.value(), stream_, br->device_mr()); + auto packed_columns = cudf::pack(table_view_.value(), stream_, br->device_mr_ref()); table_view_ = std::nullopt; return std::make_unique( std::move(packed_columns.metadata), diff --git a/cpp/tests/streaming/test_leaf_actor.cpp b/cpp/tests/streaming/test_leaf_actor.cpp index 37ec6b088..18e8c5492 100644 --- a/cpp/tests/streaming/test_leaf_actor.cpp +++ b/cpp/tests/streaming/test_leaf_actor.cpp @@ -52,7 +52,7 @@ TEST_F(StreamingLeafTasks, PushAndPullChunks) { i, std::make_unique( std::make_unique( - expects[i], stream, ctx->br()->device_mr() + expects[i], stream, ctx->br()->device_mr_ref() ), stream ) diff --git a/cpp/tests/streaming/test_partition.cpp b/cpp/tests/streaming/test_partition.cpp index b43dc62d2..3ff0e8fe6 100644 --- a/cpp/tests/streaming/test_partition.cpp +++ b/cpp/tests/streaming/test_partition.cpp @@ -43,7 +43,7 @@ TEST_F(StreamingPartition, PackUnpackRoundTrip) { inputs.emplace_back(to_message( i, std::make_unique( - std::make_unique(expects[i], stream, ctx->br()->device_mr()), + std::make_unique(expects[i], stream, ctx->br()->device_mr_ref()), stream ) )); diff --git a/cpp/tests/streaming/test_read_parquet.cpp b/cpp/tests/streaming/test_read_parquet.cpp index 1586bb7d1..dfb394736 100644 --- a/cpp/tests/streaming/test_read_parquet.cpp +++ b/cpp/tests/streaming/test_read_parquet.cpp @@ -261,7 +261,7 @@ TEST_P(StreamingReadParquetParams, ReadParquet) { ); chunk = chunk.make_available(reservation); auto packed_columns = - cudf::pack(chunk.table_view(), chunk.stream(), br->device_mr()); + cudf::pack(chunk.table_view(), chunk.stream(), br->device_mr_ref()); auto packed_data = PackedData{ std::move(packed_columns.metadata), br->move(std::move(packed_columns.gpu_data), chunk.stream()) diff --git a/cpp/tests/streaming/test_shuffler.cpp b/cpp/tests/streaming/test_shuffler.cpp index 66d899eb9..f00b37df2 100644 --- a/cpp/tests/streaming/test_shuffler.cpp +++ b/cpp/tests/streaming/test_shuffler.cpp @@ -88,7 +88,7 @@ class StreamingShuffler : public BaseStreamingShuffle, ) .at(0), stream, - ctx->br()->device_mr() + ctx->br()->device_mr_ref() ), stream ) diff --git a/cpp/tests/test_allgather.cpp b/cpp/tests/test_allgather.cpp index 667c8b4aa..f1758abe5 100644 --- a/cpp/tests/test_allgather.cpp +++ b/cpp/tests/test_allgather.cpp @@ -295,7 +295,7 @@ TEST_F(BaseAllGatherTest, opid_reuse) { if (this_rank == 0) { // Recreate the buffer resource and allgather with the delayed MR. delay_br = std::make_unique( - DelayedMemoryResource{br->device_mr(), std::chrono::milliseconds(500)} + DelayedMemoryResource{br->device_mr_ref(), std::chrono::milliseconds(500)} ); allgather = std::make_unique(GlobalEnvironment->comm_, op_id, delay_br.get()); diff --git a/cpp/tests/test_buffer.cpp b/cpp/tests/test_buffer.cpp index 770629aeb..80191620b 100644 --- a/cpp/tests/test_buffer.cpp +++ b/cpp/tests/test_buffer.cpp @@ -91,7 +91,7 @@ TEST_P(BufferRebindStreamTest, RebindStreamAndCopy) { ASSERT_NE(stream1.value(), stream2.value()); auto rmm_buffer = std::make_unique( - random_data.data(), buffer_size, stream1, br->device_mr() + random_data.data(), buffer_size, stream1, br->device_mr_ref() ); auto [reserve1, overbooking1] = diff --git a/cpp/tests/utils.hpp b/cpp/tests/utils.hpp index 0aa23ff85..a55507806 100644 --- a/cpp/tests/utils.hpp +++ b/cpp/tests/utils.hpp @@ -204,7 +204,7 @@ template std::memcpy(metadata->data(), values.data(), n_elements * sizeof(int)); auto data = std::make_unique( - values.data(), n_elements * sizeof(int), stream, br.device_mr() + values.data(), n_elements * sizeof(int), stream, br.device_mr_ref() ); return {std::move(metadata), br.move(std::move(data), stream)}; From 107b4cf54c4dc5b90dbd35748151128ffbd4c575 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Fri, 29 May 2026 09:34:35 +0200 Subject: [PATCH 02/15] Make BR a shared pointer everywhere --- cpp/benchmarks/bench_comm.cpp | 6 +- cpp/benchmarks/bench_partition.cpp | 8 +- cpp/benchmarks/bench_shuffle.cpp | 11 +- .../streaming/bench_streaming_shuffle.cpp | 2 +- cpp/benchmarks/streaming/ndsh/concatenate.cpp | 3 +- cpp/benchmarks/streaming/ndsh/join.cpp | 13 +- cpp/benchmarks/streaming/ndsh/q01.cpp | 8 +- cpp/benchmarks/streaming/ndsh/q03.cpp | 3 +- cpp/benchmarks/streaming/ndsh/utils.cpp | 2 +- cpp/examples/example_shuffle.cpp | 10 +- .../rapidsmpf/memory/buffer_resource.hpp | 86 ++++-- .../rapidsmpf/memory/spill_manager.hpp | 12 + .../rapidsmpf/rmm_resource_adaptor.hpp | 33 +++ cpp/src/memory/buffer_resource.cpp | 30 ++- cpp/src/rmm_resource_adaptor.cpp | 8 + cpp/src/streaming/cudf/table_chunk.cpp | 3 +- .../streaming/base_streaming_fixture.hpp | 2 +- cpp/tests/streaming/test_channel_metadata.cpp | 11 +- cpp/tests/streaming/test_fanout.cpp | 2 +- cpp/tests/streaming/test_message.cpp | 4 +- cpp/tests/streaming/test_partition.cpp | 4 +- cpp/tests/streaming/test_table_chunk.cpp | 2 +- cpp/tests/test_allgather.cpp | 10 +- cpp/tests/test_allreduce.cu | 4 +- cpp/tests/test_buffer.cpp | 4 +- cpp/tests/test_buffer_resource.cpp | 244 ++++++++++-------- cpp/tests/test_chunk.cpp | 4 +- cpp/tests/test_communicator.cpp | 4 +- cpp/tests/test_metadata_payload_exchange.cpp | 4 +- cpp/tests/test_partition.cpp | 18 +- cpp/tests/test_shuffler.cpp | 73 +++--- cpp/tests/test_shuffler_many_streams.cpp | 2 +- cpp/tests/test_sparse_alltoall.cpp | 4 +- cpp/tests/test_spill_manager.cpp | 38 +-- .../rapidsmpf/memory/buffer_resource.pxd | 15 +- .../rapidsmpf/memory/buffer_resource.pyx | 4 +- 36 files changed, 443 insertions(+), 248 deletions(-) diff --git a/cpp/benchmarks/bench_comm.cpp b/cpp/benchmarks/bench_comm.cpp index 076d88a59..62875b2db 100644 --- a/cpp/benchmarks/bench_comm.cpp +++ b/cpp/benchmarks/bench_comm.cpp @@ -317,7 +317,7 @@ int main(int argc, char** argv) { set_current_rmm_resource(args.rmm_mr); rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref(); - BufferResource br{ + auto br = BufferResource::create( mr, PinnedMemoryResource::Disabled, {}, @@ -326,7 +326,7 @@ int main(int argc, char** argv) { 16, rmm::cuda_stream::flags::non_blocking ), stats - }; + ); // Print benchmark/hardware info. { @@ -367,7 +367,7 @@ int main(int argc, char** argv) { if (i == args.num_warmups + args.num_runs - 1) { stats->enable(); } - auto const elapsed = run(comm, args, stream, &br, stats).count(); + auto const elapsed = run(comm, args, stream, br.get(), stats).count(); std::stringstream ss; ss << "elapsed: " << format_duration(elapsed) << " | local comm: " << format_nbytes(local_messages_send / elapsed) diff --git a/cpp/benchmarks/bench_partition.cpp b/cpp/benchmarks/bench_partition.cpp index ca286885b..909744f24 100644 --- a/cpp/benchmarks/bench_partition.cpp +++ b/cpp/benchmarks/bench_partition.cpp @@ -60,7 +60,7 @@ static void BM_PartitionAndPack(benchmark::State& state) { // Create a pool memory resource with 50% of GPU memory rmm::mr::pool_memory_resource pool_mr{rmm::mr::cuda_memory_resource{}, pool_size}; - rapidsmpf::BufferResource br{pool_mr}; + auto br = rapidsmpf::BufferResource::create(pool_mr); // Create input table auto table = create_int_table(num_rows, stream); @@ -76,7 +76,7 @@ static void BM_PartitionAndPack(benchmark::State& state) { cudf::hash_id::HASH_MURMUR3, cudf::DEFAULT_HASH_SEED, stream, - &br + br.get() ); benchmark::DoNotOptimize(pack_partitions); cudaStreamSynchronize(stream); @@ -111,7 +111,7 @@ static void BM_PartitionAndPackCurrentImpl(benchmark::State& state) { // Create a pool memory resource with 50% of GPU memory rmm::mr::pool_memory_resource pool_mr{rmm::mr::cuda_memory_resource{}, pool_size}; - rapidsmpf::BufferResource br{pool_mr}; + auto br = rapidsmpf::BufferResource::create(pool_mr); // Create input table auto table = create_int_table(num_rows, stream); @@ -128,7 +128,7 @@ static void BM_PartitionAndPackCurrentImpl(benchmark::State& state) { cudf::hash_id::HASH_MURMUR3, cudf::DEFAULT_HASH_SEED, stream, - &br + br.get() ); benchmark::DoNotOptimize(pack_partitions); } diff --git a/cpp/benchmarks/bench_shuffle.cpp b/cpp/benchmarks/bench_shuffle.cpp index 9921bb271..070eedf9c 100644 --- a/cpp/benchmarks/bench_shuffle.cpp +++ b/cpp/benchmarks/bench_shuffle.cpp @@ -547,7 +547,7 @@ int main(int argc, char** argv) { // We're only going to measure the last run, so disable initially. stats->disable(); - rapidsmpf::BufferResource br{ + auto br = rapidsmpf::BufferResource::create( stat_enabled_mr, args.pinned_mem_disable ? rapidsmpf::PinnedMemoryResource::Disabled : rapidsmpf::PinnedMemoryResource::make_if_available(), @@ -557,7 +557,7 @@ int main(int argc, char** argv) { 16, rmm::cuda_stream::flags::non_blocking ), stats - }; + ); std::shared_ptr comm; auto progress_thread = std::make_shared(stats); @@ -631,10 +631,11 @@ int main(int argc, char** argv) { } double elapsed; if (args.hash_partition_with_datagen) { - elapsed = - run_hash_partition_with_datagen(comm, args, stream, &br, stats).count(); + elapsed = run_hash_partition_with_datagen(comm, args, stream, br.get(), stats) + .count(); } else { - elapsed = run_hash_partition_inline(comm, args, stream, &br, stats).count(); + elapsed = + run_hash_partition_inline(comm, args, stream, br.get(), stats).count(); } std::stringstream ss; ss << "elapsed: " << rapidsmpf::format_duration(elapsed) diff --git a/cpp/benchmarks/streaming/bench_streaming_shuffle.cpp b/cpp/benchmarks/streaming/bench_streaming_shuffle.cpp index 5836b11ff..8b3b20cd8 100644 --- a/cpp/benchmarks/streaming/bench_streaming_shuffle.cpp +++ b/cpp/benchmarks/streaming/bench_streaming_shuffle.cpp @@ -369,7 +369,7 @@ int main(int argc, char** argv) { auto pinned_mr = args.pinned_mem_disable ? rapidsmpf::PinnedMemoryResource::Disabled : rapidsmpf::PinnedMemoryResource::make_if_available(); - auto br = std::make_shared( + auto br = rapidsmpf::BufferResource::create( stat_enabled_mr, pinned_mr, std::move(memory_limits), diff --git a/cpp/benchmarks/streaming/ndsh/concatenate.cpp b/cpp/benchmarks/streaming/ndsh/concatenate.cpp index 1637f6ede..e0e9661d8 100644 --- a/cpp/benchmarks/streaming/ndsh/concatenate.cpp +++ b/cpp/benchmarks/streaming/ndsh/concatenate.cpp @@ -70,7 +70,8 @@ streaming::Actor concatenate( chunks.push_back(std::move(chunk)); } auto result = std::make_unique( - cudf::concatenate(views, concat_stream, ctx->br()->device_mr_ref()), concat_stream + cudf::concatenate(views, concat_stream, ctx->br()->device_mr_ref()), + concat_stream ); cuda_stream_join( chunks | std::views::transform([](auto&& chunk) { return chunk.stream(); }), diff --git a/cpp/benchmarks/streaming/ndsh/join.cpp b/cpp/benchmarks/streaming/ndsh/join.cpp index 8aa6c35fa..3123e1901 100644 --- a/cpp/benchmarks/streaming/ndsh/join.cpp +++ b/cpp/benchmarks/streaming/ndsh/join.cpp @@ -69,7 +69,8 @@ coro::task broadcast( ); } else { RAPIDSMPF_EXPECTS(chunks.size() > 0, "No chunks in broadcast"); - auto result = cudf::concatenate(views, gather_stream, ctx->br()->device_mr_ref()); + auto result = + cudf::concatenate(views, gather_stream, ctx->br()->device_mr_ref()); // So that deallocation of the consitutent tables is stream-ordered wrt the // concatenation. cuda_stream_join( @@ -93,8 +94,9 @@ coro::task broadcast( // TODO: If this chunk is already in pack form, this is unnecessary. auto chunk = co_await msg.release().make_available(ctx); - auto pack = - cudf::pack(chunk.table_view(), chunk.stream(), ctx->br()->device_mr_ref()); + auto pack = cudf::pack( + chunk.table_view(), chunk.stream(), ctx->br()->device_mr_ref() + ); auto packed_data = PackedData( std::move(pack.metadata), ctx->br()->move(std::move(pack.gpu_data), chunk.stream()) @@ -243,8 +245,9 @@ streaming::Message inner_join_chunk( build_event->stream_wait(chunk_stream); auto probe_table = right_chunk.table_view(); auto probe_keys = probe_table.select(right_on); - auto [probe_match, build_match] = - joiner.inner_join(probe_keys, std::nullopt, chunk_stream, ctx->br()->device_mr_ref()); + auto [probe_match, build_match] = joiner.inner_join( + probe_keys, std::nullopt, chunk_stream, ctx->br()->device_mr_ref() + ); cudf::column_view build_indices = cudf::device_span(*build_match); diff --git a/cpp/benchmarks/streaming/ndsh/q01.cpp b/cpp/benchmarks/streaming/ndsh/q01.cpp index 68ef2e298..1c6d4b138 100644 --- a/cpp/benchmarks/streaming/ndsh/q01.cpp +++ b/cpp/benchmarks/streaming/ndsh/q01.cpp @@ -202,7 +202,9 @@ rapidsmpf::streaming::Actor select_columns_for_groupby( auto table = chunk.table_view(); // l_returnflag, l_linestatus, l_quantity, l_extendedprice auto result = - cudf::table(table.select({0, 1, 2, 3}), chunk_stream, ctx->br()->device_mr_ref()) + cudf::table( + table.select({0, 1, 2, 3}), chunk_stream, ctx->br()->device_mr_ref() + ) .release(); result.reserve(7); auto extendedprice = table.column(3); @@ -253,7 +255,9 @@ static __device__ void calculate_charge(double *charge, double discprice, double ); // l_discount result.push_back( - std::make_unique(discount, chunk_stream, ctx->br()->device_mr_ref()) + std::make_unique( + discount, chunk_stream, ctx->br()->device_mr_ref() + ) ); co_await ch_out->send( rapidsmpf::streaming::to_message( diff --git a/cpp/benchmarks/streaming/ndsh/q03.cpp b/cpp/benchmarks/streaming/ndsh/q03.cpp index 81748dfaf..190995682 100644 --- a/cpp/benchmarks/streaming/ndsh/q03.cpp +++ b/cpp/benchmarks/streaming/ndsh/q03.cpp @@ -336,7 +336,8 @@ rapidsmpf::streaming::Actor top_k_by( std::ranges::transform(partials, std::back_inserter(views), [](auto& t) { return t->view(); }); - auto merged = cudf::merge(views, keys, order, {}, out_stream, ctx->br()->device_mr_ref()); + auto merged = + cudf::merge(views, keys, order, {}, out_stream, ctx->br()->device_mr_ref()); auto result = std::make_unique( cudf::slice(merged->view(), {0, 10}, out_stream), out_stream, diff --git a/cpp/benchmarks/streaming/ndsh/utils.cpp b/cpp/benchmarks/streaming/ndsh/utils.cpp index fb20dd8b1..8716da3be 100644 --- a/cpp/benchmarks/streaming/ndsh/utils.cpp +++ b/cpp/benchmarks/streaming/ndsh/utils.cpp @@ -156,7 +156,7 @@ create_context( std::invalid_argument ); - auto br = std::make_shared( + auto br = BufferResource::create( std::move(mr), arguments.no_pinned_host_memory ? PinnedMemoryResource::Disabled : PinnedMemoryResource::make_if_available(), diff --git a/cpp/examples/example_shuffle.cpp b/cpp/examples/example_shuffle.cpp index 7ba73f0cf..144c3bd59 100644 --- a/cpp/examples/example_shuffle.cpp +++ b/cpp/examples/example_shuffle.cpp @@ -46,7 +46,7 @@ int main(int argc, char** argv) { // We will use the same stream, memory, and buffer resource throughout the example. rmm::cuda_stream_view stream = cudf::get_default_stream(); rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref(); - rapidsmpf::BufferResource br{mr}; + auto br = rapidsmpf::BufferResource::create(mr); // As input data, we use a helper function from the benchmark suite. It creates a // random cudf table with 2 columns and 100 rows. In this example, each MPI rank @@ -66,7 +66,7 @@ int main(int argc, char** argv) { comm, 0, // op_id total_num_partitions, - &br, + br.get(), rapidsmpf::shuffler::Shuffler::round_robin // partition owner ); @@ -83,7 +83,7 @@ int main(int argc, char** argv) { cudf::hash_id::HASH_MURMUR3, cudf::DEFAULT_HASH_SEED, stream, - &br + br.get() ); // Now, we can insert the packed partitions into the shuffler. This operation is @@ -112,10 +112,10 @@ int main(int argc, char** argv) { local_outputs.push_back( rapidsmpf::unpack_and_concat( rapidsmpf::unspill_partitions( - std::move(packed_chunks), &br, rapidsmpf::AllowOverbooking::YES + std::move(packed_chunks), br.get(), rapidsmpf::AllowOverbooking::YES ), stream, - &br + br.get() ) ); } diff --git a/cpp/include/rapidsmpf/memory/buffer_resource.hpp b/cpp/include/rapidsmpf/memory/buffer_resource.hpp index f0efd4efc..6c17d6bb8 100644 --- a/cpp/include/rapidsmpf/memory/buffer_resource.hpp +++ b/cpp/include/rapidsmpf/memory/buffer_resource.hpp @@ -54,8 +54,19 @@ enum class AllowOverbooking : bool { * @note Similar to RMM's memory resource, the `BufferResource` instance must outlive all * allocated buffers and memory reservations. */ -class BufferResource { +class BufferResource : public std::enable_shared_from_this { public: + /** + * @brief Tag struct used to make `BufferResource`'s constructor effectively + * private. The tag is public so that `BufferResource::create` can + * forward it, but user code can't write `BufferResource{...}` directly. + * All construction must go through `create()` or `from_options()`, which + * guarantees the result is owned by a `std::shared_ptr`. + */ + struct PrivateTag { + explicit PrivateTag() = default; + }; + /** * @brief Constructs a buffer resource. * @@ -83,6 +94,35 @@ class BufferResource { * @param statistics The statistics instance to use (disabled by default). */ BufferResource( + PrivateTag, + cuda::mr::any_resource device_mr, + std::optional pinned_mr = PinnedMemoryResource::Disabled, + std::unordered_map memory_limits = {}, + std::optional periodic_spill_check = std::chrono::milliseconds{1}, + std::shared_ptr stream_pool = std::make_shared< + rmm::cuda_stream_pool>(16, rmm::cuda_stream::flags::non_blocking), + std::shared_ptr statistics = Statistics::disabled() + ); + + /** + * @brief Factory method to construct a `BufferResource` as a `shared_ptr`. + * + * This is the standard way to construct a `BufferResource`. The returned + * handle is the canonical owner; pass `std::shared_ptr` + * (or a raw pointer to it) wherever a `BufferResource` is needed. + * + * @param device_mr The RMM device memory resource used for device allocations. + * @param pinned_mr Pinned host memory resource (or `PinnedMemoryResource::Disabled`). + * @param memory_limits Per-`MemoryType` byte limits. Missing entries default to + * unlimited. + * @param periodic_spill_check Pause between periodic spill checks + * (`std::nullopt` disables the dedicated thread). + * @param stream_pool CUDA stream pool used for operations without an explicit stream. + * @param statistics Statistics instance (defaults to disabled). + * + * @return A `std::shared_ptr` owning a newly constructed `BufferResource`. + */ + [[nodiscard]] static std::shared_ptr create( cuda::mr::any_resource device_mr, std::optional pinned_mr = PinnedMemoryResource::Disabled, std::unordered_map memory_limits = {}, @@ -114,26 +154,38 @@ class BufferResource { ~BufferResource() noexcept = default; + /// @brief `BufferResource` is non-copyable, it is owned by `std::shared_ptr`. + BufferResource(BufferResource const&) = delete; + /// @brief `BufferResource` is non-movable, it is owned by `std::shared_ptr`. + BufferResource(BufferResource&&) = delete; + /// @brief `BufferResource` is non-copyable. + /// @return Reference to this. + BufferResource& operator=(BufferResource const&) = delete; + /// @brief `BufferResource` is non-movable. + /// @return Reference to this. + BufferResource& operator=(BufferResource&&) = delete; + /** * @brief Get the device memory resource as an owning, type-erased handle. * - * Returns a copy of the internally stored - * `cuda::mr::any_resource`. The copy is cheap — - * `any_resource` of a `shared_resource`-based adaptor shares state via an - * internal `shared_ptr`. Because the returned object owns its own - * refcounted handle on the underlying adaptor, it remains valid even if - * this `BufferResource` is destroyed. - * - * Use this when you need to **store** the device MR for the lifetime of - * another object (e.g. as the `mr` argument to `rmm::device_buffer`, which - * keeps its MR by value). For transient cuDF/RMM calls that take an - * `rmm::device_async_resource_ref`, prefer `device_mr_ref()`. - * - * @return An owning `cuda::mr::any_resource` handle on - * the device memory resource. + * Returns a `cuda::mr::any_resource` that owns a copy + * of the internal `RmmResourceAdaptor` carrying an owning back-pointer + * (`shared_ptr`) to this `BufferResource`. Anything that + * stores the returned object — e.g. `rmm::device_buffer`'s internal + * `any_resource` member, or a Python `DeviceMemoryResource` wrapper — + * keeps this `BufferResource` (including its stream pool, spill + * manager, and statistics) alive for as long as the returned MR exists. + * + * Use this when you need the MR to **outlive** the caller's + * `shared_ptr` (e.g. allocating a device buffer that + * is returned out of a scope where the BR is dropped). For transient + * cuDF/RMM calls that take an `rmm::device_async_resource_ref` and + * complete before the caller's BR ref drops, prefer `device_mr_ref()`. + * + * @return An owning `cuda::mr::any_resource` that + * keeps this `BufferResource` alive. */ - [[nodiscard]] cuda::mr::any_resource - device_mr() const noexcept; + [[nodiscard]] cuda::mr::any_resource device_mr(); /** * @brief Get a non-owning reference to the device memory resource. diff --git a/cpp/include/rapidsmpf/memory/spill_manager.hpp b/cpp/include/rapidsmpf/memory/spill_manager.hpp index 43892fb7a..2969a3bfd 100644 --- a/cpp/include/rapidsmpf/memory/spill_manager.hpp +++ b/cpp/include/rapidsmpf/memory/spill_manager.hpp @@ -29,6 +29,18 @@ class SpillManager { * * A spill function takes a requested spill amount as input and returns the actual * amount of memory (in bytes) that was spilled. + * + * @warning Spill functions must not capture an owning reference (a + * `shared_ptr`, a `BufferResource` by value, etc.) to + * the `BufferResource` that owns this `SpillManager`. They must also not + * capture any object — `Buffer`, `rmm::device_buffer`, or other holder + * of memory allocated through `BufferResource::device_mr()` — that + * transitively owns the same `BufferResource`. Either case closes a + * refcount cycle (BR → SpillManager → SpillFunction → BR), which is + * invisible to Python's garbage collector and leaks the resource + * unconditionally. Capture raw pointers/references only, and unregister + * via `remove_spill_function()` from the owner's destructor before + * dropping any buffer-holding members. */ using SpillFunction = std::function; diff --git a/cpp/include/rapidsmpf/rmm_resource_adaptor.hpp b/cpp/include/rapidsmpf/rmm_resource_adaptor.hpp index 408820492..4d256918c 100644 --- a/cpp/include/rapidsmpf/rmm_resource_adaptor.hpp +++ b/cpp/include/rapidsmpf/rmm_resource_adaptor.hpp @@ -7,6 +7,7 @@ #include #include +#include #include @@ -17,6 +18,8 @@ namespace rapidsmpf { +class BufferResource; // defined in memory/buffer_resource.hpp + /** * @brief A RMM memory resource adaptor tailored to RapidsMPF. * @@ -120,6 +123,36 @@ class RmmResourceAdaptor * @see begin_scoped_memory_record() */ ScopedMemoryRecord end_scoped_memory_record(); + + /** + * @brief Return a copy of this adaptor that owns a back-pointer to @p br. + * + * The returned adaptor shares allocation-tracking state with `*this` via + * the underlying `shared_resource` (so allocations through either copy + * are accumulated into the same record); it additionally holds an + * owning `shared_ptr` so that anything storing the + * returned object (e.g. an `rmm::device_buffer`'s internal + * `any_resource`) keeps @p br — including its stream pool, spill + * manager, and statistics — alive. + * + * Used by `BufferResource::device_mr()` to produce a lifetime-safe + * owning device memory resource. The original adaptor stored inside + * `BufferResource` itself does **not** carry the back-pointer, so no + * self-cycle is created. + * + * @param br The `BufferResource` to keep alive via the returned adaptor. + * @return A copy of this adaptor with @p br stored as a back-pointer. + */ + [[nodiscard]] RmmResourceAdaptor with_buffer_resource_ref( + std::shared_ptr br + ) const; + + private: + /// @brief Optional owning back-pointer to the `BufferResource` containing + /// this adaptor. Null in the original adaptor stored inside `BufferResource` + /// (to avoid a self-cycle). Populated in copies returned via + /// `with_buffer_resource_ref()` / `BufferResource::device_mr()`. + std::shared_ptr br_; }; static_assert(cuda::mr::resource_with); diff --git a/cpp/src/memory/buffer_resource.cpp b/cpp/src/memory/buffer_resource.cpp index de02c9450..57ce87d62 100644 --- a/cpp/src/memory/buffer_resource.cpp +++ b/cpp/src/memory/buffer_resource.cpp @@ -22,6 +22,7 @@ namespace rapidsmpf { BufferResource::BufferResource( + PrivateTag, cuda::mr::any_resource device_mr, std::optional pinned_mr, std::unordered_map memory_limits, @@ -49,6 +50,25 @@ BufferResource::BufferResource( RAPIDSMPF_EXPECTS(statistics_ != nullptr, "the statistics pointer cannot be NULL"); } +std::shared_ptr BufferResource::create( + cuda::mr::any_resource device_mr, + std::optional pinned_mr, + std::unordered_map memory_limits, + std::optional periodic_spill_check, + std::shared_ptr stream_pool, + std::shared_ptr statistics +) { + return std::make_shared( + PrivateTag{}, + std::move(device_mr), + std::move(pinned_mr), + std::move(memory_limits), + periodic_spill_check, + std::move(stream_pool), + std::move(statistics) + ); +} + std::shared_ptr BufferResource::from_options( cuda::mr::any_resource mr, config::Options options, @@ -58,8 +78,7 @@ std::shared_ptr BufferResource::from_options( std::unordered_map memory_limits{ {MemoryType::DEVICE, device_limit_from_options(options)} }; - - return std::make_shared( + return create( std::move(mr), std::move(pinned_mr), std::move(memory_limits), @@ -94,9 +113,10 @@ void BufferResource::set_memory_limit(MemoryType mem_type, std::int64_t limit) n ); } -cuda::mr::any_resource -BufferResource::device_mr() const noexcept { - return device_mr_; +cuda::mr::any_resource BufferResource::device_mr() { + return cuda::mr::any_resource{ + device_adaptor_.with_buffer_resource_ref(shared_from_this()) + }; } rmm::device_async_resource_ref BufferResource::device_mr_ref() const noexcept { diff --git a/cpp/src/rmm_resource_adaptor.cpp b/cpp/src/rmm_resource_adaptor.cpp index f7a223e8c..1d9716728 100644 --- a/cpp/src/rmm_resource_adaptor.cpp +++ b/cpp/src/rmm_resource_adaptor.cpp @@ -40,4 +40,12 @@ ScopedMemoryRecord RmmResourceAdaptor::end_scoped_memory_record() { return get().end_scoped_memory_record(); } +RmmResourceAdaptor RmmResourceAdaptor::with_buffer_resource_ref( + std::shared_ptr br +) const { + RmmResourceAdaptor copy{*this}; + copy.br_ = std::move(br); + return copy; +} + } // namespace rapidsmpf diff --git a/cpp/src/streaming/cudf/table_chunk.cpp b/cpp/src/streaming/cudf/table_chunk.cpp index a08af79b8..f296f99a4 100644 --- a/cpp/src/streaming/cudf/table_chunk.cpp +++ b/cpp/src/streaming/cudf/table_chunk.cpp @@ -217,7 +217,8 @@ TableChunk TableChunk::copy(MemoryReservation& reservation) const { // new host buffer. // TODO: use `cudf::chunked_pack()` with a bounce buffer. Currently, // `cudf::pack()` allocates device memory we haven't reserved. - auto packed_columns = cudf::pack(table_view(), stream(), br->device_mr_ref()); + auto packed_columns = + cudf::pack(table_view(), stream(), br->device_mr_ref()); auto packed_data = std::make_unique( std::move(packed_columns.metadata), br->move(std::move(packed_columns.gpu_data), stream()) diff --git a/cpp/tests/streaming/base_streaming_fixture.hpp b/cpp/tests/streaming/base_streaming_fixture.hpp index d2dc7ea84..23cbf2b6e 100644 --- a/cpp/tests/streaming/base_streaming_fixture.hpp +++ b/cpp/tests/streaming/base_streaming_fixture.hpp @@ -43,7 +43,7 @@ class BaseStreamingFixture : public ::testing::Test { rapidsmpf::config::Options options(std::move(env_vars)); stream = cudf::get_default_stream(); - br = std::make_shared( + br = rapidsmpf::BufferResource::create( mr_cuda, rapidsmpf::PinnedMemoryResource::Disabled, std::move(memory_limits) ); ctx = std::make_shared( diff --git a/cpp/tests/streaming/test_channel_metadata.cpp b/cpp/tests/streaming/test_channel_metadata.cpp index 603ae2006..6c97ce3c5 100644 --- a/cpp/tests/streaming/test_channel_metadata.cpp +++ b/cpp/tests/streaming/test_channel_metadata.cpp @@ -166,7 +166,8 @@ TEST_F(StreamingChannelMetadata, MessageRoundTrip) { class StreamingChannelMetadataGPU : public ::testing::Test { protected: rmm::cuda_stream_view stream{cudf::get_default_stream()}; - rapidsmpf::BufferResource br{cudf::get_current_device_resource_ref()}; + std::shared_ptr br = + rapidsmpf::BufferResource::create(cudf::get_current_device_resource_ref()); std::shared_ptr make_chunk(std::vector vals) { rmm::device_buffer buf(vals.data(), vals.size() * sizeof(int32_t), stream); @@ -209,16 +210,16 @@ TEST_F(StreamingChannelMetadataGPU, OrderSchemeBoundariesAlignedWith) { OrderScheme o1({k0}, make_chunk({100, 200})); OrderScheme o2({k0}, make_chunk({100, 200})); - EXPECT_TRUE(o1.boundaries_aligned_with(o2, br)); + EXPECT_TRUE(o1.boundaries_aligned_with(o2, *br)); OrderScheme o_shifted({k3}, make_chunk({100, 200})); - EXPECT_TRUE(o1.boundaries_aligned_with(o_shifted, br)); + EXPECT_TRUE(o1.boundaries_aligned_with(o_shifted, *br)); OrderScheme o_strict({k0}, make_chunk({100, 200}), /*strict=*/true); - EXPECT_FALSE(o1.boundaries_aligned_with(o_strict, br)); + EXPECT_FALSE(o1.boundaries_aligned_with(o_strict, *br)); OrderScheme o_diff({k0}, make_chunk({100, 300})); - EXPECT_FALSE(o1.boundaries_aligned_with(o_diff, br)); + EXPECT_FALSE(o1.boundaries_aligned_with(o_diff, *br)); } TEST_F(StreamingChannelMetadataGPU, PartitioningSpecOrder) { diff --git a/cpp/tests/streaming/test_fanout.cpp b/cpp/tests/streaming/test_fanout.cpp index 4d77dc966..3653f643e 100644 --- a/cpp/tests/streaming/test_fanout.cpp +++ b/cpp/tests/streaming/test_fanout.cpp @@ -547,7 +547,7 @@ class SpillingStreamingFanout : public BaseStreamingFixture { std::unordered_map memory_limits = { {MemoryType::DEVICE, 0}, }; - br = std::make_shared( + br = rapidsmpf::BufferResource::create( mr_cuda, rapidsmpf::PinnedMemoryResource::Disabled, memory_limits ); auto options = ctx->options(); diff --git a/cpp/tests/streaming/test_message.cpp b/cpp/tests/streaming/test_message.cpp index 165730e1d..ba2a9454f 100644 --- a/cpp/tests/streaming/test_message.cpp +++ b/cpp/tests/streaming/test_message.cpp @@ -18,11 +18,11 @@ using namespace rapidsmpf::streaming; class StreamingMessage : public ::testing::Test { protected: void SetUp() override { - br = std::make_unique(cudf::get_current_device_resource_ref()); + br = BufferResource::create(cudf::get_current_device_resource_ref()); stream = cudf::get_default_stream(); } - std::unique_ptr br; + std::shared_ptr br; rmm::cuda_stream_view stream; }; diff --git a/cpp/tests/streaming/test_partition.cpp b/cpp/tests/streaming/test_partition.cpp index 3ff0e8fe6..a3519b1df 100644 --- a/cpp/tests/streaming/test_partition.cpp +++ b/cpp/tests/streaming/test_partition.cpp @@ -43,7 +43,9 @@ TEST_F(StreamingPartition, PackUnpackRoundTrip) { inputs.emplace_back(to_message( i, std::make_unique( - std::make_unique(expects[i], stream, ctx->br()->device_mr_ref()), + std::make_unique( + expects[i], stream, ctx->br()->device_mr_ref() + ), stream ) )); diff --git a/cpp/tests/streaming/test_table_chunk.cpp b/cpp/tests/streaming/test_table_chunk.cpp index 3b66e5e9e..5d42285bd 100644 --- a/cpp/tests/streaming/test_table_chunk.cpp +++ b/cpp/tests/streaming/test_table_chunk.cpp @@ -41,7 +41,7 @@ class StreamingTableChunk : public BaseStreamingFixture, 16, rmm::cuda_stream::flags::non_blocking ); stream = cudf::get_default_stream(); - br = std::make_shared( + br = rapidsmpf::BufferResource::create( mr_cuda, // device_mr rapidsmpf::PinnedMemoryResource::make_if_available(), // pinned_mr memory_limits, // memory_limits diff --git a/cpp/tests/test_allgather.cpp b/cpp/tests/test_allgather.cpp index f1758abe5..bf72659ee 100644 --- a/cpp/tests/test_allgather.cpp +++ b/cpp/tests/test_allgather.cpp @@ -34,7 +34,7 @@ class BaseAllGatherTest : public ::testing::Test { protected: void SetUp() override { stream = cudf::get_default_stream(); - br = std::make_unique(rmm::mr::cuda_memory_resource{}); + br = rapidsmpf::BufferResource::create(rmm::mr::cuda_memory_resource{}); } void TearDown() override { @@ -42,7 +42,7 @@ class BaseAllGatherTest : public ::testing::Test { } rmm::cuda_stream_view stream; - std::unique_ptr br; + std::shared_ptr br; }; TEST_F(BaseAllGatherTest, timeout) { @@ -289,12 +289,12 @@ TEST_F(BaseAllGatherTest, opid_reuse) { auto this_rank = comm->rank(); // On rank 0, wrap the device MR with a delayed version. - std::unique_ptr delay_br; + std::shared_ptr delay_br; std::unique_ptr allgather; constexpr rapidsmpf::OpID op_id = 0; if (this_rank == 0) { // Recreate the buffer resource and allgather with the delayed MR. - delay_br = std::make_unique( + delay_br = rapidsmpf::BufferResource::create( DelayedMemoryResource{br->device_mr_ref(), std::chrono::milliseconds(500)} ); allgather = @@ -363,7 +363,7 @@ TEST_F(BaseAllGatherTest, opid_reuse) { TEST(PostBox, spill_uses_remaining_amount) { auto stream = cudf::get_default_stream(); auto mr = std::make_unique(); - auto br = std::make_unique(*mr); + auto br = rapidsmpf::BufferResource::create(*mr); rapidsmpf::coll::detail::PostBox postbox; diff --git a/cpp/tests/test_allreduce.cu b/cpp/tests/test_allreduce.cu index 6c20b26c6..5d4b9bd93 100644 --- a/cpp/tests/test_allreduce.cu +++ b/cpp/tests/test_allreduce.cu @@ -207,7 +207,7 @@ class BaseAllReduceTest : public ::testing::Test { protected: void SetUp() override { mr = std::make_unique(); - br = std::make_unique(*mr); + br = rapidsmpf::BufferResource::create(*mr); comm = GlobalEnvironment->comm_.get(); } @@ -217,7 +217,7 @@ class BaseAllReduceTest : public ::testing::Test { } rapidsmpf::Communicator* comm; - std::unique_ptr br; + std::shared_ptr br; std::unique_ptr mr; }; diff --git a/cpp/tests/test_buffer.cpp b/cpp/tests/test_buffer.cpp index 80191620b..f5fa5a0c1 100644 --- a/cpp/tests/test_buffer.cpp +++ b/cpp/tests/test_buffer.cpp @@ -53,7 +53,7 @@ class BufferRebindStreamTest : public ::testing::TestWithParam { GTEST_SKIP() << "Pinned memory resources are not supported on this system"; } - br = std::make_unique( + br = BufferResource::create( cudf::get_current_device_resource_ref(), PinnedMemoryResource::make_if_available(), std::unordered_map{}, @@ -73,7 +73,7 @@ class BufferRebindStreamTest : public ::testing::TestWithParam { static constexpr std::size_t chunk_size = 1_MiB; std::shared_ptr stream_pool; - std::unique_ptr br; + std::shared_ptr br; std::vector random_data; }; diff --git a/cpp/tests/test_buffer_resource.cpp b/cpp/tests/test_buffer_resource.cpp index cea208c88..da23077df 100644 --- a/cpp/tests/test_buffer_resource.cpp +++ b/cpp/tests/test_buffer_resource.cpp @@ -19,6 +19,7 @@ #include #include #include +#include #include #include @@ -52,119 +53,119 @@ std::unique_ptr zeros( TEST(BufferResource, ReservationOverbooking) { // Create a buffer resource that always reports 10 KiB of available device memory. - BufferResource br{ + auto br = BufferResource::create( cudf::get_current_device_resource_ref(), PinnedMemoryResource::Disabled, {{MemoryType::DEVICE, 10_KiB}} - }; - EXPECT_EQ(br.memory_reserved(MemoryType::DEVICE), 0); - EXPECT_EQ(br.memory_reserved(MemoryType::HOST), 0); + ); + EXPECT_EQ(br->memory_reserved(MemoryType::DEVICE), 0); + EXPECT_EQ(br->memory_reserved(MemoryType::HOST), 0); // Book all available memory. auto [reserve1, overbooking1] = - br.reserve(MemoryType::DEVICE, 10_KiB, AllowOverbooking::NO); + br->reserve(MemoryType::DEVICE, 10_KiB, AllowOverbooking::NO); EXPECT_EQ(reserve1.size(), 10_KiB); EXPECT_EQ(overbooking1, 0); - EXPECT_EQ(br.memory_reserved(MemoryType::DEVICE), 10_KiB); - EXPECT_EQ(br.memory_reserved(MemoryType::HOST), 0); + EXPECT_EQ(br->memory_reserved(MemoryType::DEVICE), 10_KiB); + EXPECT_EQ(br->memory_reserved(MemoryType::HOST), 0); // Try to overbook. auto [reserve2, overbooking2] = - br.reserve(MemoryType::DEVICE, 10_KiB, AllowOverbooking::NO); + br->reserve(MemoryType::DEVICE, 10_KiB, AllowOverbooking::NO); EXPECT_EQ(reserve2.size(), 0); // Reservation failed. EXPECT_EQ(overbooking2, 10_KiB); - EXPECT_EQ(br.memory_reserved(MemoryType::DEVICE), 10_KiB); - EXPECT_EQ(br.memory_reserved(MemoryType::HOST), 0); + EXPECT_EQ(br->memory_reserved(MemoryType::DEVICE), 10_KiB); + EXPECT_EQ(br->memory_reserved(MemoryType::HOST), 0); // Allow overbooking. auto [reserve3, overbooking3] = - br.reserve(MemoryType::DEVICE, 10_KiB, AllowOverbooking::YES); + br->reserve(MemoryType::DEVICE, 10_KiB, AllowOverbooking::YES); EXPECT_EQ(reserve3.size(), 10_KiB); EXPECT_EQ(overbooking3, 10_KiB); - EXPECT_EQ(br.memory_reserved(MemoryType::DEVICE), 20_KiB); - EXPECT_EQ(br.memory_reserved(MemoryType::HOST), 0); + EXPECT_EQ(br->memory_reserved(MemoryType::DEVICE), 20_KiB); + EXPECT_EQ(br->memory_reserved(MemoryType::HOST), 0); // No host limit. auto [reserve4, overbooking4] = - br.reserve(MemoryType::HOST, 10_KiB, AllowOverbooking::NO); + br->reserve(MemoryType::HOST, 10_KiB, AllowOverbooking::NO); EXPECT_EQ(reserve4.size(), 10_KiB); EXPECT_EQ(overbooking4, 0); - EXPECT_EQ(br.memory_reserved(MemoryType::DEVICE), 20_KiB); - EXPECT_EQ(br.memory_reserved(MemoryType::HOST), 10_KiB); + EXPECT_EQ(br->memory_reserved(MemoryType::DEVICE), 20_KiB); + EXPECT_EQ(br->memory_reserved(MemoryType::HOST), 10_KiB); // Cannot release the wrong memory type. - EXPECT_EQ(br.memory_reserved(MemoryType::DEVICE), 20_KiB); - EXPECT_EQ(br.memory_reserved(MemoryType::HOST), 10_KiB); + EXPECT_EQ(br->memory_reserved(MemoryType::DEVICE), 20_KiB); + EXPECT_EQ(br->memory_reserved(MemoryType::HOST), 10_KiB); // Cannot release more than the size of the reservation. - EXPECT_THROW(br.release(reserve1, 20_KiB), rapidsmpf::reservation_error); - EXPECT_EQ(br.memory_reserved(MemoryType::DEVICE), 20_KiB); - EXPECT_EQ(br.memory_reserved(MemoryType::HOST), 10_KiB); + EXPECT_THROW(br->release(reserve1, 20_KiB), rapidsmpf::reservation_error); + EXPECT_EQ(br->memory_reserved(MemoryType::DEVICE), 20_KiB); + EXPECT_EQ(br->memory_reserved(MemoryType::HOST), 10_KiB); // Partial releasing a reservation. - EXPECT_EQ(br.release(reserve1, 5_KiB), 5_KiB); + EXPECT_EQ(br->release(reserve1, 5_KiB), 5_KiB); EXPECT_EQ(reserve1.size(), 5_KiB); - EXPECT_EQ(br.memory_reserved(MemoryType::DEVICE), 15_KiB); - EXPECT_EQ(br.memory_reserved(MemoryType::HOST), 10_KiB); + EXPECT_EQ(br->memory_reserved(MemoryType::DEVICE), 15_KiB); + EXPECT_EQ(br->memory_reserved(MemoryType::HOST), 10_KiB); // We are still overbooking. auto [reserve5, overbooking5] = - br.reserve(MemoryType::DEVICE, 5_KiB, AllowOverbooking::YES); + br->reserve(MemoryType::DEVICE, 5_KiB, AllowOverbooking::YES); EXPECT_EQ(reserve5.size(), 5_KiB); EXPECT_EQ(overbooking5, 10_KiB); - EXPECT_EQ(br.memory_reserved(MemoryType::DEVICE), 20_KiB); - EXPECT_EQ(br.memory_reserved(MemoryType::HOST), 10_KiB); + EXPECT_EQ(br->memory_reserved(MemoryType::DEVICE), 20_KiB); + EXPECT_EQ(br->memory_reserved(MemoryType::HOST), 10_KiB); } TEST(BufferResource, ReservationReleasing) { // Create a buffer resource that always reports 10 KiB of available host and device // memory. - BufferResource br{ + auto br = BufferResource::create( cudf::get_current_device_resource_ref(), PinnedMemoryResource::Disabled, {{MemoryType::DEVICE, 10_KiB}, {MemoryType::HOST, 10_KiB}} - }; - EXPECT_EQ(br.memory_reserved(MemoryType::DEVICE), 0); - EXPECT_EQ(br.memory_reserved(MemoryType::HOST), 0); + ); + EXPECT_EQ(br->memory_reserved(MemoryType::DEVICE), 0); + EXPECT_EQ(br->memory_reserved(MemoryType::HOST), 0); // Reserve all available host and device memory. auto [reserve1, overbooking1] = - br.reserve(MemoryType::DEVICE, 10_KiB, AllowOverbooking::NO); + br->reserve(MemoryType::DEVICE, 10_KiB, AllowOverbooking::NO); auto [reserve2, overbooking2] = - br.reserve(MemoryType::HOST, 10_KiB, AllowOverbooking::NO); + br->reserve(MemoryType::HOST, 10_KiB, AllowOverbooking::NO); EXPECT_EQ(reserve1.size(), 10_KiB); EXPECT_EQ(overbooking1, 0); EXPECT_EQ(reserve2.size(), 10_KiB); EXPECT_EQ(overbooking2, 0); - EXPECT_EQ(br.memory_reserved(MemoryType::DEVICE), 10_KiB); - EXPECT_EQ(br.memory_reserved(MemoryType::HOST), 10_KiB); + EXPECT_EQ(br->memory_reserved(MemoryType::DEVICE), 10_KiB); + EXPECT_EQ(br->memory_reserved(MemoryType::HOST), 10_KiB); // Cannot release the wrong memory type. - EXPECT_EQ(br.memory_reserved(MemoryType::DEVICE), 10_KiB); - EXPECT_EQ(br.memory_reserved(MemoryType::HOST), 10_KiB); + EXPECT_EQ(br->memory_reserved(MemoryType::DEVICE), 10_KiB); + EXPECT_EQ(br->memory_reserved(MemoryType::HOST), 10_KiB); // Cannot release more than the size of the reservation. - EXPECT_THROW(br.release(reserve1, 20_KiB), rapidsmpf::reservation_error); - EXPECT_EQ(br.memory_reserved(MemoryType::DEVICE), 10_KiB); - EXPECT_EQ(br.memory_reserved(MemoryType::HOST), 10_KiB); + EXPECT_THROW(br->release(reserve1, 20_KiB), rapidsmpf::reservation_error); + EXPECT_EQ(br->memory_reserved(MemoryType::DEVICE), 10_KiB); + EXPECT_EQ(br->memory_reserved(MemoryType::HOST), 10_KiB); // Partial releasing a reservation. - EXPECT_EQ(br.release(reserve1, 5_KiB), 5_KiB); + EXPECT_EQ(br->release(reserve1, 5_KiB), 5_KiB); EXPECT_EQ(reserve1.size(), 5_KiB); - EXPECT_EQ(br.memory_reserved(MemoryType::DEVICE), 5_KiB); - EXPECT_EQ(br.memory_reserved(MemoryType::HOST), 10_KiB); + EXPECT_EQ(br->memory_reserved(MemoryType::DEVICE), 5_KiB); + EXPECT_EQ(br->memory_reserved(MemoryType::HOST), 10_KiB); // A reservation is released when it goes out of scope. { auto [reserve, overbooking] = - br.reserve(MemoryType::HOST, 10_KiB, AllowOverbooking::YES); + br->reserve(MemoryType::HOST, 10_KiB, AllowOverbooking::YES); EXPECT_EQ(reserve.size(), 10_KiB); EXPECT_EQ(overbooking, 10_KiB); - EXPECT_EQ(br.memory_reserved(MemoryType::DEVICE), 5_KiB); - EXPECT_EQ(br.memory_reserved(MemoryType::HOST), 20_KiB); + EXPECT_EQ(br->memory_reserved(MemoryType::DEVICE), 5_KiB); + EXPECT_EQ(br->memory_reserved(MemoryType::HOST), 20_KiB); } - EXPECT_EQ(br.memory_reserved(MemoryType::DEVICE), 5_KiB); - EXPECT_EQ(br.memory_reserved(MemoryType::HOST), 10_KiB); + EXPECT_EQ(br->memory_reserved(MemoryType::DEVICE), 5_KiB); + EXPECT_EQ(br->memory_reserved(MemoryType::HOST), 10_KiB); } TEST(BufferResource, MemoryLimit) { @@ -172,69 +173,69 @@ TEST(BufferResource, MemoryLimit) { auto stream = cudf::get_default_stream(); // Create a buffer resource that limits available device memory to 10 KiB. - BufferResource br{ + auto br = BufferResource::create( mr_cuda, PinnedMemoryResource::Disabled, {{MemoryType::DEVICE, 10_KiB}} - }; - EXPECT_EQ(br.memory_available(MemoryType::DEVICE), 10_KiB); - EXPECT_EQ(br.memory_reserved(MemoryType::DEVICE), 0); - EXPECT_EQ(br.memory_reserved(MemoryType::HOST), 0); + ); + EXPECT_EQ(br->memory_available(MemoryType::DEVICE), 10_KiB); + EXPECT_EQ(br->memory_reserved(MemoryType::DEVICE), 0); + EXPECT_EQ(br->memory_reserved(MemoryType::HOST), 0); // Book all available device memory. auto [reserve1, overbooking1] = - br.reserve(MemoryType::DEVICE, 10_KiB, AllowOverbooking::NO); + br->reserve(MemoryType::DEVICE, 10_KiB, AllowOverbooking::NO); EXPECT_EQ(reserve1.size(), 10_KiB); EXPECT_EQ(overbooking1, 0); - EXPECT_EQ(br.memory_reserved(MemoryType::DEVICE), 10_KiB); - EXPECT_EQ(br.memory_reserved(MemoryType::HOST), 0); + EXPECT_EQ(br->memory_reserved(MemoryType::DEVICE), 10_KiB); + EXPECT_EQ(br->memory_reserved(MemoryType::HOST), 0); // Allocating a Buffer also requires a reservation, which are then released. - auto dev_buf1 = zeros(br, 10_KiB, stream, reserve1); + auto dev_buf1 = zeros(*br, 10_KiB, stream, reserve1); EXPECT_EQ(dev_buf1->mem_type(), MemoryType::DEVICE); EXPECT_EQ(dev_buf1->size, 10_KiB); EXPECT_EQ(reserve1.size(), 0); - EXPECT_EQ(br.memory_reserved(MemoryType::DEVICE), 0_KiB); - EXPECT_EQ(br.memory_reserved(MemoryType::HOST), 0_KiB); - EXPECT_EQ(br.memory_available(MemoryType::DEVICE), 0); + EXPECT_EQ(br->memory_reserved(MemoryType::DEVICE), 0_KiB); + EXPECT_EQ(br->memory_reserved(MemoryType::HOST), 0_KiB); + EXPECT_EQ(br->memory_available(MemoryType::DEVICE), 0); // Insufficent reservation for the allocation. - EXPECT_THROW(zeros(br, 10_KiB, stream, reserve1), rapidsmpf::reservation_error); + EXPECT_THROW(zeros(*br, 10_KiB, stream, reserve1), rapidsmpf::reservation_error); // Freeing a buffer increases the available but the reserved memory is unchanged. dev_buf1.reset(); - EXPECT_EQ(br.memory_available(MemoryType::DEVICE), 10_KiB); - EXPECT_EQ(br.memory_reserved(MemoryType::DEVICE), 0_KiB); - EXPECT_EQ(br.memory_reserved(MemoryType::HOST), 0_KiB); + EXPECT_EQ(br->memory_available(MemoryType::DEVICE), 10_KiB); + EXPECT_EQ(br->memory_reserved(MemoryType::DEVICE), 0_KiB); + EXPECT_EQ(br->memory_reserved(MemoryType::HOST), 0_KiB); // Moving buffers between memory types requires a reservation. auto [reserve2, overbooking2] = - br.reserve(MemoryType::DEVICE, 10_KiB, AllowOverbooking::YES); - auto dev_buf2 = zeros(br, 10_KiB, stream, reserve2); + br->reserve(MemoryType::DEVICE, 10_KiB, AllowOverbooking::YES); + auto dev_buf2 = zeros(*br, 10_KiB, stream, reserve2); EXPECT_EQ(dev_buf2->mem_type(), MemoryType::DEVICE); auto [reserve3, overbooking3] = - br.reserve(MemoryType::HOST, 10_KiB, AllowOverbooking::YES); - EXPECT_EQ(br.memory_reserved(MemoryType::DEVICE), 0_KiB); - EXPECT_EQ(br.memory_reserved(MemoryType::HOST), 10_KiB); - EXPECT_EQ(br.memory_available(MemoryType::DEVICE), 0); + br->reserve(MemoryType::HOST, 10_KiB, AllowOverbooking::YES); + EXPECT_EQ(br->memory_reserved(MemoryType::DEVICE), 0_KiB); + EXPECT_EQ(br->memory_reserved(MemoryType::HOST), 10_KiB); + EXPECT_EQ(br->memory_available(MemoryType::DEVICE), 0); - auto host_buf2 = br.move(std::move(dev_buf2), reserve3); + auto host_buf2 = br->move(std::move(dev_buf2), reserve3); EXPECT_EQ(host_buf2->mem_type(), MemoryType::HOST); - EXPECT_EQ(br.memory_reserved(MemoryType::DEVICE), 0_KiB); - EXPECT_EQ(br.memory_reserved(MemoryType::HOST), 0_KiB); - EXPECT_EQ(br.memory_available(MemoryType::DEVICE), 10_KiB); + EXPECT_EQ(br->memory_reserved(MemoryType::DEVICE), 0_KiB); + EXPECT_EQ(br->memory_reserved(MemoryType::HOST), 0_KiB); + EXPECT_EQ(br->memory_available(MemoryType::DEVICE), 10_KiB); // Moving buffers to the same memory type accepts an empty reservation. - auto host_buf3 = br.move(std::move(host_buf2), reserve3); + auto host_buf3 = br->move(std::move(host_buf2), reserve3); EXPECT_EQ(host_buf3->mem_type(), MemoryType::HOST); - EXPECT_EQ(br.memory_reserved(MemoryType::DEVICE), 0_KiB); - EXPECT_EQ(br.memory_reserved(MemoryType::HOST), 0_KiB); - EXPECT_EQ(br.memory_available(MemoryType::DEVICE), 10_KiB); + EXPECT_EQ(br->memory_reserved(MemoryType::DEVICE), 0_KiB); + EXPECT_EQ(br->memory_reserved(MemoryType::HOST), 0_KiB); + EXPECT_EQ(br->memory_available(MemoryType::DEVICE), 10_KiB); // The reservation must be of the correct memory type. auto [reserve4, overbooking4] = - br.reserve(MemoryType::HOST, 10_KiB, AllowOverbooking::YES); + br->reserve(MemoryType::HOST, 10_KiB, AllowOverbooking::YES); EXPECT_EQ(reserve4.size(), 10_KiB); - EXPECT_EQ(br.memory_reserved(MemoryType::DEVICE), 0_KiB); - EXPECT_EQ(br.memory_reserved(MemoryType::HOST), 10_KiB); + EXPECT_EQ(br->memory_reserved(MemoryType::DEVICE), 0_KiB); + EXPECT_EQ(br->memory_reserved(MemoryType::HOST), 10_KiB); } class PinnedMaxPoolSizeReservationLimitTest @@ -262,15 +263,15 @@ TEST_P(PinnedMaxPoolSizeReservationLimitTest, TwoReservations) { if (max_pool_size.has_value() && *max_pool_size > 0) { memory_limits[MemoryType::PINNED_HOST] = safe_cast(*max_pool_size); } - BufferResource br{cuda_mr, pinned_mr, std::move(memory_limits)}; + auto br = BufferResource::create(cuda_mr, pinned_mr, std::move(memory_limits)); // First 1 KiB reservation always succeeds. - auto [r1, ob1] = br.reserve(MemoryType::PINNED_HOST, 1_KiB, AllowOverbooking::NO); + auto [r1, ob1] = br->reserve(MemoryType::PINNED_HOST, 1_KiB, AllowOverbooking::NO); EXPECT_EQ(r1.size(), 1_KiB); EXPECT_EQ(ob1, 0); // Second 1 KiB reservation succeeds only when the pool is unbounded. - auto [r2, ob2] = br.reserve(MemoryType::PINNED_HOST, 1_KiB, AllowOverbooking::NO); + auto [r2, ob2] = br->reserve(MemoryType::PINNED_HOST, 1_KiB, AllowOverbooking::NO); EXPECT_EQ(r2.size(), expect_second_succeeds() ? 1_KiB : 0); EXPECT_EQ(ob2, expect_second_succeeds() ? 0 : 1_KiB); } @@ -289,14 +290,14 @@ TEST(BufferResource, AllocStatistics) { rmm::mr::cuda_memory_resource mr_cuda; auto stats = Statistics::create(); auto pinned_mr = PinnedMemoryResource::make_if_available(); - BufferResource br{ + auto br = BufferResource::create( mr_cuda, pinned_mr, {}, std::nullopt, std::make_shared(1, rmm::cuda_stream::flags::non_blocking), stats - }; + ); auto stream = cudf::get_default_stream(); constexpr std::size_t device_size = 4_KiB; @@ -305,23 +306,23 @@ TEST(BufferResource, AllocStatistics) { // Allocate device memory twice. { - auto [r, _] = br.reserve(MemoryType::DEVICE, device_size, AllowOverbooking::YES); - br.make_buffer(device_size, stream, r); + auto [r, _] = br->reserve(MemoryType::DEVICE, device_size, AllowOverbooking::YES); + br->make_buffer(device_size, stream, r); } { - auto [r, _] = br.reserve(MemoryType::DEVICE, device_size, AllowOverbooking::YES); - br.make_buffer(device_size, stream, r); + auto [r, _] = br->reserve(MemoryType::DEVICE, device_size, AllowOverbooking::YES); + br->make_buffer(device_size, stream, r); } // Allocate pinned_host memory once (if available). if (pinned_mr != PinnedMemoryResource::Disabled) { auto [r, _] = - br.reserve(MemoryType::PINNED_HOST, pinned_size, AllowOverbooking::YES); - br.make_buffer(pinned_size, stream, r); + br->reserve(MemoryType::PINNED_HOST, pinned_size, AllowOverbooking::YES); + br->make_buffer(pinned_size, stream, r); } // Allocate host memory once. { - auto [r, _] = br.reserve(MemoryType::HOST, host_size, AllowOverbooking::YES); - br.make_buffer(host_size, stream, r); + auto [r, _] = br->reserve(MemoryType::HOST, host_size, AllowOverbooking::YES); + br->make_buffer(host_size, stream, r); } stream.synchronize(); @@ -354,7 +355,7 @@ class BufferResourceReserveOrFailTest : public ::testing::Test { void SetUp() override { // Create a buffer resource with limited device memory (10 KiB) and unlimited // host memory. BufferResource auto-wraps mr_cuda for allocation tracking. - br = std::make_unique( + br = BufferResource::create( mr_cuda, PinnedMemoryResource::Disabled, std::unordered_map{{MemoryType::DEVICE, 10_KiB}} @@ -362,7 +363,7 @@ class BufferResourceReserveOrFailTest : public ::testing::Test { } rmm::mr::cuda_memory_resource mr_cuda; - std::unique_ptr br; + std::shared_ptr br; }; // Static assertions to verify that various container types can be used with @@ -417,7 +418,7 @@ TEST_F(BufferResourceReserveOrFailTest, MultipleTypes) { class BaseBufferResourceCopyTest : public ::testing::Test { protected: void SetUp() override { - br = std::make_unique(cudf::get_current_device_resource_ref()); + br = BufferResource::create(cudf::get_current_device_resource_ref()); stream = cudf::get_default_stream(); // initialize the host pattern @@ -445,7 +446,7 @@ class BaseBufferResourceCopyTest : public ::testing::Test { static constexpr std::size_t buffer_size = 1024; // 1 KiB - std::unique_ptr br; + std::shared_ptr br; rmm::cuda_stream_view stream; std::vector host_pattern; // a predefined pattern for testing @@ -639,11 +640,11 @@ class BufferResourceDifferentResourcesTest : public ::testing::Test { // Setup br1 with statistics for its device memory mr1 = std::make_unique(rmm::mr::cuda_memory_resource{}); - br1 = std::make_unique(*mr1); + br1 = BufferResource::create(*mr1); // Setup br2 with statistics for its device memory mr2 = std::make_unique(rmm::mr::cuda_memory_resource{}); - br2 = std::make_unique(*mr2); + br2 = BufferResource::create(*mr2); } std::unique_ptr create_source_buffer() { @@ -677,8 +678,8 @@ class BufferResourceDifferentResourcesTest : public ::testing::Test { std::unique_ptr mr1; std::unique_ptr mr2; - std::unique_ptr br1; - std::unique_ptr br2; + std::shared_ptr br1; + std::shared_ptr br2; }; TEST_F(BufferResourceDifferentResourcesTest, CopySlice) { @@ -786,3 +787,42 @@ TEST_F(BufferCopyEdgeCases, SameBufferIsDisallowed) { buffer_copy(br->statistics(), *buf, *buf, 16, 0, 0), std::invalid_argument ); } + +// The owning `any_resource` returned by `device_mr()` carries an owning +// back-ref to the BR via the adaptor copy. Dropping the original BR while +// the returned MR (and any buffer allocated through it) is still alive +// must not free the BR — its stream pool must remain valid for the +// buffer's stream-ordered destruction. Conversely, once everything that +// observed the back-ref drops it, BR must destruct (no refcount cycle). +TEST(BufferResource, DeviceMrOutlivesBufferResource) { + constexpr std::size_t N = 1024; + + auto br = BufferResource::create(cudf::get_current_device_resource_ref()); + std::weak_ptr weak_br = br; + auto stream = cudf::get_default_stream(); + + auto owning_mr = br->device_mr(); + auto* adaptor_in_mr = cuda::mr::resource_cast(&owning_mr); + ASSERT_NE(adaptor_in_mr, nullptr); + + auto buf = std::make_unique(N, stream, std::move(owning_mr)); + + // The adaptor in the buffer's `any_resource` shares allocation tracking + // with the original adaptor stored in BR, so the allocation is recorded. + EXPECT_GE(adaptor_in_mr->current_allocated(), static_cast(N)); + + // Drop the original BR. The adaptor copy inside `buf`'s `any_resource` + // owns the only remaining `shared_ptr`, so BR (and its + // stream pool) must stay alive. + br.reset(); + EXPECT_FALSE(weak_br.expired()) << "BR freed while buffer still holds it"; + + // Destroy the buffer — its stream-ordered deallocator runs against the + // still-alive stream pool. Must not crash. + EXPECT_NO_THROW(buf.reset()); + + // After the buffer is gone, no shared_ptr
exists anywhere. If there + // were a refcount cycle (e.g. adaptor-copy back-ref leaking back into + // the original BR), `weak_br` would still be observable here. + EXPECT_TRUE(weak_br.expired()) << "BR not destructed — refcount cycle?"; +} diff --git a/cpp/tests/test_chunk.cpp b/cpp/tests/test_chunk.cpp index 77af94a82..79ea8ab9d 100644 --- a/cpp/tests/test_chunk.cpp +++ b/cpp/tests/test_chunk.cpp @@ -24,11 +24,11 @@ using namespace rapidsmpf::shuffler::detail; class ChunkTest : public ::testing::Test { protected: void SetUp() override { - br = std::make_unique(cudf::get_current_device_resource_ref()); + br = BufferResource::create(cudf::get_current_device_resource_ref()); stream = cudf::get_default_stream(); } - std::unique_ptr br; + std::shared_ptr br; rmm::cuda_stream_view stream; }; diff --git a/cpp/tests/test_communicator.cpp b/cpp/tests/test_communicator.cpp index dc7d88aa0..d1c2eea54 100644 --- a/cpp/tests/test_communicator.cpp +++ b/cpp/tests/test_communicator.cpp @@ -24,7 +24,7 @@ class BaseCommunicatorTest : public ::testing::Test { void SetUp() override { comm = GlobalEnvironment->comm_.get(); mr = std::make_unique(); - br = std::make_unique(*mr); + br = rapidsmpf::BufferResource::create(*mr); stream = rmm::cuda_stream_default; } @@ -33,7 +33,7 @@ class BaseCommunicatorTest : public ::testing::Test { rapidsmpf::Communicator* comm; std::unique_ptr mr; rmm::cuda_stream_view stream; - std::unique_ptr br; + std::shared_ptr br; }; TEST_F(BaseCommunicatorTest, TagConstruction) { diff --git a/cpp/tests/test_metadata_payload_exchange.cpp b/cpp/tests/test_metadata_payload_exchange.cpp index bab869315..fe6d647e4 100644 --- a/cpp/tests/test_metadata_payload_exchange.cpp +++ b/cpp/tests/test_metadata_payload_exchange.cpp @@ -27,7 +27,7 @@ class MetadataPayloadExchangeTest : public ::testing::Test { void SetUp() override { comm = GlobalEnvironment->comm_.get(); mr = std::make_unique(); - br = std::make_unique(*mr); + br = BufferResource::create(*mr); stream = rmm::cuda_stream_default; statistics = Statistics::create(); @@ -101,7 +101,7 @@ class MetadataPayloadExchangeTest : public ::testing::Test { Communicator* comm; std::unique_ptr mr; rmm::cuda_stream_view stream; - std::unique_ptr br; + std::shared_ptr br; std::shared_ptr statistics; std::unique_ptr comm_interface; }; diff --git a/cpp/tests/test_partition.cpp b/cpp/tests/test_partition.cpp index 9bae901b4..b4e68a247 100644 --- a/cpp/tests/test_partition.cpp +++ b/cpp/tests/test_partition.cpp @@ -44,13 +44,13 @@ TEST_P(NumOfPartitions, partition_and_pack) { std::int64_t const seed = 42; cudf::hash_id const hash_fn = cudf::hash_id::HASH_MURMUR3; auto stream = cudf::get_default_stream(); - rapidsmpf::BufferResource br{mr()}; + auto br = rapidsmpf::BufferResource::create(mr()); cudf::table expect = random_table_with_index(seed, static_cast(num_rows), 0, 10); auto chunks = rapidsmpf::partition_and_pack( - expect, {1}, num_partitions, hash_fn, seed, stream, &br + expect, {1}, num_partitions, hash_fn, seed, stream, br.get() ); // Convert to a vector @@ -60,7 +60,8 @@ TEST_P(NumOfPartitions, partition_and_pack) { } EXPECT_EQ(chunks_vector.size(), num_partitions); - auto result = rapidsmpf::unpack_and_concat(std::move(chunks_vector), stream, &br); + auto result = + rapidsmpf::unpack_and_concat(std::move(chunks_vector), stream, br.get()); // Compare the input table with the result. We ignore the row order by // sorting by their index (first column). @@ -72,7 +73,7 @@ TEST_P(NumOfPartitions, split_and_pack) { int const num_rows = std::get<1>(GetParam()); std::int64_t const seed = 42; auto stream = cudf::get_default_stream(); - rapidsmpf::BufferResource br{cudf::get_current_device_resource_ref()}; + auto br = rapidsmpf::BufferResource::create(cudf::get_current_device_resource_ref()); cudf::table expect = random_table_with_index(seed, num_rows, 0, 10); @@ -81,7 +82,7 @@ TEST_P(NumOfPartitions, split_and_pack) { splits.emplace_back(i * num_rows / num_partitions); } - auto chunks = rapidsmpf::split_and_pack(expect, splits, stream, &br); + auto chunks = rapidsmpf::split_and_pack(expect, splits, stream, br.get()); // Convert to a vector (restoring the original order). std::vector chunks_vector; @@ -90,7 +91,8 @@ TEST_P(NumOfPartitions, split_and_pack) { } EXPECT_EQ(chunks_vector.size(), num_partitions); - auto result = rapidsmpf::unpack_and_concat(std::move(chunks_vector), stream, &br); + auto result = + rapidsmpf::unpack_and_concat(std::move(chunks_vector), stream, br.get()); // Compare the input table with the result. CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expect, *result); @@ -99,11 +101,11 @@ TEST_P(NumOfPartitions, split_and_pack) { class SpillingTest : public ::testing::Test { protected: void SetUp() override { - br = std::make_unique(cudf::get_current_device_resource_ref()); + br = BufferResource::create(cudf::get_current_device_resource_ref()); stream = cudf::get_default_stream(); } - std::unique_ptr br; + std::shared_ptr br; rmm::cuda_stream_view stream; }; diff --git a/cpp/tests/test_shuffler.cpp b/cpp/tests/test_shuffler.cpp index df89b50d5..024820a90 100644 --- a/cpp/tests/test_shuffler.cpp +++ b/cpp/tests/test_shuffler.cpp @@ -32,7 +32,7 @@ extern Environment* GlobalEnvironment; TEST(ReceivedChunks, spill_skips_control_messages) { auto mr = cudf::get_current_device_resource_ref(); - auto br = std::make_unique(mr); + auto br = rapidsmpf::BufferResource::create(mr); rapidsmpf::shuffler::detail::ReceivedChunks received; @@ -49,7 +49,7 @@ TEST(ReceivedChunks, spill_skips_control_messages) { TEST(ReceivedChunks, spill_respects_amount) { auto mr = cudf::get_current_device_resource_ref(); - auto br = std::make_unique(mr); + auto br = rapidsmpf::BufferResource::create(mr); auto stream = cudf::get_default_stream(); rapidsmpf::shuffler::detail::ReceivedChunks received; @@ -75,7 +75,7 @@ TEST(ReceivedChunks, spill_respects_amount) { TEST(MetadataMessage, round_trip) { auto stream = cudf::get_default_stream(); auto mr = cudf::get_current_device_resource_ref(); - auto br = std::make_unique(mr); + auto br = rapidsmpf::BufferResource::create(mr); auto metadata = iota_vector(100); @@ -237,7 +237,7 @@ class MemoryLimits_NumPartition memory_limits = std::get<0>(GetParam()); total_num_partitions = std::get<1>(GetParam()); total_num_rows = std::get<2>(GetParam()); - br = std::make_unique( + br = rapidsmpf::BufferResource::create( mr(), rapidsmpf::PinnedMemoryResource::Disabled, memory_limits ); @@ -260,7 +260,7 @@ class MemoryLimits_NumPartition std::int64_t seed = 42; cudf::hash_id hash_fn = cudf::hash_id::HASH_MURMUR3; rmm::cuda_stream_view stream; - std::unique_ptr br; + std::shared_ptr br; std::unique_ptr shuffler; }; @@ -309,7 +309,7 @@ class ConcurrentShuffleTest static_cast(std::get<1>(GetParam())); // these resources will be used by multiple threads to instantiate shufflers - br = std::make_unique(mr()); + br = rapidsmpf::BufferResource::create(mr()); stream = cudf::get_default_stream(); } @@ -371,7 +371,7 @@ class ConcurrentShuffleTest rapidsmpf::shuffler::PartID total_num_partitions; rmm::cuda_stream_view stream; - std::unique_ptr br; + std::shared_ptr br; }; TEST_P(ConcurrentShuffleTest, round_trip) { @@ -414,12 +414,12 @@ TEST(Shuffler, SpillOnInsertAndExtraction) { // (force spill), regardless of how many bytes are currently allocated from `mr`. constexpr std::int64_t k_no_spill_limit = (1LL << 40); constexpr std::int64_t k_force_spill_limit = -(1LL << 40); - rapidsmpf::BufferResource br{ + auto br = rapidsmpf::BufferResource::create( mr, rapidsmpf::PinnedMemoryResource::Disabled, {{rapidsmpf::MemoryType::DEVICE, k_no_spill_limit}}, std::nullopt // disable periodic spill check - }; + ); // Create a communicator of size 1, such that each shuffler will run locally. auto comm = GlobalEnvironment->split_comm(); @@ -430,7 +430,7 @@ TEST(Shuffler, SpillOnInsertAndExtraction) { comm, 0, // op_id total_num_partitions, - &br + br.get() ); cudf::table input_table = random_table_with_index(seed, 1000, 0, 10); auto input_chunks = rapidsmpf::partition_and_pack( @@ -440,7 +440,7 @@ TEST(Shuffler, SpillOnInsertAndExtraction) { hash_fn, seed, stream, - &br, + br.get(), rapidsmpf::AllowOverbooking::YES ); // with overbooking @@ -452,12 +452,12 @@ TEST(Shuffler, SpillOnInsertAndExtraction) { EXPECT_EQ(mr.get_main_record().num_current_allocs(), 2); // Let's force spilling. - br.set_memory_limit(rapidsmpf::MemoryType::DEVICE, k_force_spill_limit); + br->set_memory_limit(rapidsmpf::MemoryType::DEVICE, k_force_spill_limit); { // Now extract triggers spilling of the partition not being extracted. std::vector output_chunks = rapidsmpf::unspill_partitions( - shuffler.extract(0), &br, rapidsmpf::AllowOverbooking::YES + shuffler.extract(0), br.get(), rapidsmpf::AllowOverbooking::YES ); EXPECT_EQ(mr.get_main_record().num_current_allocs(), 1); @@ -470,16 +470,16 @@ TEST(Shuffler, SpillOnInsertAndExtraction) { // Extract and unspill both partitions. std::vector out0 = rapidsmpf::unspill_partitions( - shuffler.extract(0), &br, rapidsmpf::AllowOverbooking::YES + shuffler.extract(0), br.get(), rapidsmpf::AllowOverbooking::YES ); EXPECT_EQ(mr.get_main_record().num_current_allocs(), 1); std::vector out1 = rapidsmpf::unspill_partitions( - shuffler.extract(1), &br, rapidsmpf::AllowOverbooking::YES + shuffler.extract(1), br.get(), rapidsmpf::AllowOverbooking::YES ); EXPECT_EQ(mr.get_main_record().num_current_allocs(), 2); // Disable spilling and insert the first partition. - br.set_memory_limit(rapidsmpf::MemoryType::DEVICE, k_no_spill_limit); + br->set_memory_limit(rapidsmpf::MemoryType::DEVICE, k_no_spill_limit); { std::unordered_map chunk; chunk.emplace(0, std::move(out0.at(0))); @@ -490,7 +490,7 @@ TEST(Shuffler, SpillOnInsertAndExtraction) { // Enable spilling and insert the second partition, which should trigger spilling // of both the first partition already in the shuffler and the second partition // that are being inserted. - br.set_memory_limit(rapidsmpf::MemoryType::DEVICE, k_force_spill_limit); + br->set_memory_limit(rapidsmpf::MemoryType::DEVICE, k_force_spill_limit); { std::unordered_map chunk; chunk.emplace(1, std::move(out1.at(0))); @@ -676,7 +676,7 @@ TEST(Shuffler, ShutdownWhilePaused) { auto progress_thread = GlobalEnvironment->comm_->progress_thread(); auto mr = cudf::get_current_device_resource_ref(); - auto br = std::make_unique(mr); + auto br = rapidsmpf::BufferResource::create(mr); auto shuffler = rapidsmpf::shuffler::Shuffler(GlobalEnvironment->comm_, 0, 1, br.get()); @@ -711,7 +711,7 @@ class ExtractEmptyPartitionsTest : public cudf::test::BaseFixture { void SetUp() override { stream = cudf::get_default_stream(); - br = std::make_unique(mr()); + br = rapidsmpf::BufferResource::create(mr()); shuffler = std::make_unique( GlobalEnvironment->comm_, 0, nparts, br.get() @@ -761,7 +761,7 @@ class ExtractEmptyPartitionsTest : public cudf::test::BaseFixture { } rmm::cuda_stream_view stream; - std::unique_ptr br; + std::shared_ptr br; std::unique_ptr shuffler; }; @@ -810,9 +810,10 @@ TEST_F(ExtractEmptyPartitionsTest, SomeEmptyAndNonEmptyInsertions) { TEST(ShufflerTest, multiple_shutdowns) { auto& comm = GlobalEnvironment->comm_; - rapidsmpf::BufferResource br(cudf::get_current_device_resource_ref()); - auto shuffler = - std::make_unique(comm, 0, comm->nranks(), &br); + auto br = rapidsmpf::BufferResource::create(cudf::get_current_device_resource_ref()); + auto shuffler = std::make_unique( + comm, 0, comm->nranks(), br.get() + ); shuffler->insert_finished(); EXPECT_NO_THROW(shuffler->wait(std::chrono::seconds(30))); @@ -834,7 +835,7 @@ TEST(ShufflerTest, multiple_shutdowns) { TEST(Shuffler, concurrent_wait) { auto const& comm = GlobalEnvironment->comm_; auto stream = cudf::get_default_stream(); - rapidsmpf::BufferResource br(cudf::get_current_device_resource_ref()); + auto br = rapidsmpf::BufferResource::create(cudf::get_current_device_resource_ref()); // Use more partitions than ranks so each rank owns multiple partitions, ensuring // multiple threads call wait() concurrently on the same shuffler. @@ -845,7 +846,7 @@ TEST(Shuffler, concurrent_wait) { constexpr std::int64_t seed = 42; constexpr auto wait_timeout = std::chrono::seconds{30}; - rapidsmpf::shuffler::Shuffler shuffler(comm, 0, total_num_partitions, &br); + rapidsmpf::shuffler::Shuffler shuffler(comm, 0, total_num_partitions, br.get()); cudf::table full_input = random_table_with_index(seed, total_num_rows, 0, 10); auto [expected, owner] = rapidsmpf::partition_and_split( @@ -855,7 +856,7 @@ TEST(Shuffler, concurrent_wait) { hash_fn, seed, stream, - &br, + br.get(), rapidsmpf::AllowOverbooking::YES ); @@ -881,8 +882,8 @@ TEST(Shuffler, concurrent_wait) { static_cast(total_num_partitions), hash_fn, seed, - br.stream_pool().get_stream(), - &br, + br->stream_pool().get_stream(), + br.get(), rapidsmpf::AllowOverbooking::YES ) ); @@ -902,10 +903,10 @@ TEST(Shuffler, concurrent_wait) { auto chunks = shuffler.extract(pid); auto result = rapidsmpf::unpack_and_concat( rapidsmpf::unspill_partitions( - std::move(chunks), &br, rapidsmpf::AllowOverbooking::YES + std::move(chunks), br.get(), rapidsmpf::AllowOverbooking::YES ), stream, - &br, + br.get(), rapidsmpf::AllowOverbooking::YES ); CUDF_TEST_EXPECT_TABLES_EQUIVALENT( @@ -938,16 +939,16 @@ TEST(Shuffler, opid_reuse) { constexpr auto wait_timeout = std::chrono::seconds{30}; rmm::mr::cuda_memory_resource mr; - auto br = std::make_unique(mr); + auto br = rapidsmpf::BufferResource::create(mr); // On rank 0, wrap the device MR with a delayed version for the shuffler. std::unique_ptr delayed_mr; - std::unique_ptr delayed_br; + std::shared_ptr delayed_br; rapidsmpf::BufferResource* shuffler_br = br.get(); if (comm->rank() == 0) { delayed_mr = std::make_unique(mr, std::chrono::milliseconds(500)); - delayed_br = std::make_unique(*delayed_mr); + delayed_br = rapidsmpf::BufferResource::create(*delayed_mr); shuffler_br = delayed_br.get(); } @@ -1045,16 +1046,16 @@ TEST(Shuffler, opid_reuse_with_empty_partitions) { constexpr auto wait_timeout = std::chrono::seconds{30}; rmm::mr::cuda_memory_resource mr; - auto br = std::make_unique(mr); + auto br = rapidsmpf::BufferResource::create(mr); // On rank 0, wrap the device MR with a delayed version for the shuffler. std::unique_ptr delayed_mr; - std::unique_ptr delayed_br; + std::shared_ptr delayed_br; rapidsmpf::BufferResource* shuffler_br = br.get(); if (comm->rank() == 0) { delayed_mr = std::make_unique(mr, std::chrono::milliseconds(500)); - delayed_br = std::make_unique(*delayed_mr); + delayed_br = rapidsmpf::BufferResource::create(*delayed_mr); shuffler_br = delayed_br.get(); } diff --git a/cpp/tests/test_shuffler_many_streams.cpp b/cpp/tests/test_shuffler_many_streams.cpp index a42fea215..43c3f6402 100644 --- a/cpp/tests/test_shuffler_many_streams.cpp +++ b/cpp/tests/test_shuffler_many_streams.cpp @@ -50,7 +50,7 @@ TEST(ShufflerManyStreams, Test) { std::mt19937 random_generator{42}; constexpr std::size_t chunksize = 1 << 20; constexpr int num_partitions = 100; - auto br = std::make_unique(cudf::get_current_device_resource_ref()); + auto br = BufferResource::create(cudf::get_current_device_resource_ref()); // Create a CUDA stream for each partition. // To stress-test stream handling, assign random priorities so streams are more diff --git a/cpp/tests/test_sparse_alltoall.cpp b/cpp/tests/test_sparse_alltoall.cpp index 69b4d9e80..325060937 100644 --- a/cpp/tests/test_sparse_alltoall.cpp +++ b/cpp/tests/test_sparse_alltoall.cpp @@ -105,10 +105,10 @@ int decode_payload(rapidsmpf::PackedData const& packed_data) { class SparseAlltoallTest : public ::testing::Test { protected: void SetUp() override { - br = std::make_unique(rmm::mr::cuda_memory_resource{}); + br = rapidsmpf::BufferResource::create(rmm::mr::cuda_memory_resource{}); } - std::unique_ptr br; + std::shared_ptr br; }; TEST_F(SparseAlltoallTest, validate_constructor) { diff --git a/cpp/tests/test_spill_manager.cpp b/cpp/tests/test_spill_manager.cpp index f4fb6bf8c..401b26f2f 100644 --- a/cpp/tests/test_spill_manager.cpp +++ b/cpp/tests/test_spill_manager.cpp @@ -28,23 +28,23 @@ TEST(SpillManager, SpillFunction) { // No real allocations occur in this test, so memory_available equals the // currently configured limit. std::int64_t mem_available = 10_KiB; - BufferResource br{ + auto br = BufferResource::create( cudf::get_current_device_resource_ref(), rapidsmpf::PinnedMemoryResource::Disabled, {{MemoryType::DEVICE, mem_available}} - }; - EXPECT_EQ(br.memory_available(MemoryType::DEVICE), 10_KiB); + ); + EXPECT_EQ(br->memory_available(MemoryType::DEVICE), 10_KiB); // Spill function that increases the available memory perfectly. SpillManager::SpillFunction func1 = [&br, &mem_available](std::size_t amount) -> std::size_t { mem_available += safe_cast(amount); - br.set_memory_limit(MemoryType::DEVICE, mem_available); + br->set_memory_limit(MemoryType::DEVICE, mem_available); return amount; }; - br.spill_manager().add_spill_function(func1, /* priority = */ 1); - EXPECT_EQ(br.spill_manager().spill(10_KiB), 10_KiB); - EXPECT_EQ(br.memory_available(MemoryType::DEVICE), 20_KiB); + br->spill_manager().add_spill_function(func1, /* priority = */ 1); + EXPECT_EQ(br->spill_manager().spill(10_KiB), 10_KiB); + EXPECT_EQ(br->memory_available(MemoryType::DEVICE), 20_KiB); // Spill function that never spill any memory but has a higher priority. bool func2_called = false; @@ -52,27 +52,27 @@ TEST(SpillManager, SpillFunction) { func2_called = true; return 0; }; - auto fid2 = br.spill_manager().add_spill_function(func2, /* priority = */ 2); - EXPECT_EQ(br.spill_manager().spill(10_KiB), 10_KiB); + auto fid2 = br->spill_manager().add_spill_function(func2, /* priority = */ 2); + EXPECT_EQ(br->spill_manager().spill(10_KiB), 10_KiB); EXPECT_TRUE(func2_called); func2_called = false; - EXPECT_EQ(br.memory_available(MemoryType::DEVICE), 30_KiB); + EXPECT_EQ(br->memory_available(MemoryType::DEVICE), 30_KiB); // Removing `func2` means it shouldn't run. - br.spill_manager().remove_spill_function(fid2); - EXPECT_EQ(br.spill_manager().spill(10_KiB), 10_KiB); + br->spill_manager().remove_spill_function(fid2); + EXPECT_EQ(br->spill_manager().spill(10_KiB), 10_KiB); EXPECT_FALSE(func2_called); - EXPECT_EQ(br.memory_available(MemoryType::DEVICE), 40_KiB); + EXPECT_EQ(br->memory_available(MemoryType::DEVICE), 40_KiB); // If the headroom is already there, no spilling should be happening. - EXPECT_EQ(br.spill_manager().spill_to_make_headroom(10_KiB), 0); - EXPECT_EQ(br.memory_available(MemoryType::DEVICE), 40_KiB); + EXPECT_EQ(br->spill_manager().spill_to_make_headroom(10_KiB), 0); + EXPECT_EQ(br->memory_available(MemoryType::DEVICE), 40_KiB); // If the headroom isn't there, we should spill to get the headroom. - EXPECT_EQ(br.spill_manager().spill_to_make_headroom(100_KiB), 60_KiB); - EXPECT_EQ(br.memory_available(MemoryType::DEVICE), 100_KiB); + EXPECT_EQ(br->spill_manager().spill_to_make_headroom(100_KiB), 60_KiB); + EXPECT_EQ(br->memory_available(MemoryType::DEVICE), 100_KiB); // A negative headroom is allowed. - EXPECT_EQ(br.spill_manager().spill_to_make_headroom(-100_KiB), 0); - EXPECT_EQ(br.memory_available(MemoryType::DEVICE), 100_KiB); + EXPECT_EQ(br->spill_manager().spill_to_make_headroom(-100_KiB), 0); + EXPECT_EQ(br->memory_available(MemoryType::DEVICE), 100_KiB); } diff --git a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pxd b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pxd index 29fefa2a6..7ca4a06ed 100644 --- a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pxd +++ b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pxd @@ -5,7 +5,10 @@ from libc.stddef cimport size_t from libc.stdint cimport int64_t from libcpp cimport bool as bool_t from libcpp.memory cimport shared_ptr +from libcpp.optional cimport optional +from libcpp.unordered_map cimport unordered_map from rmm.librmm.cuda_stream_pool cimport cuda_stream_pool +from rmm.librmm.memory_resource cimport any_resource, device_accessible from rmm.pylibrmm.cuda_stream_pool cimport CudaStreamPool from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource @@ -13,7 +16,8 @@ from rapidsmpf._detail.exception_handling cimport ex_handler from rapidsmpf.config cimport Options, cpp_Options from rapidsmpf.memory.buffer cimport MemoryType from rapidsmpf.memory.memory_reservation cimport cpp_MemoryReservation -from rapidsmpf.memory.pinned_memory_resource cimport PinnedMemoryResource +from rapidsmpf.memory.pinned_memory_resource cimport (PinnedMemoryResource, + cpp_PinnedMemoryResource) from rapidsmpf.memory.spill_manager cimport SpillManager, cpp_SpillManager from rapidsmpf.statistics cimport Statistics, cpp_Statistics from rapidsmpf.utils.time cimport cpp_Duration @@ -26,6 +30,15 @@ cdef extern from "" nogil: cdef extern from "" nogil: cdef cppclass cpp_BufferResource "rapidsmpf::BufferResource": + @staticmethod + shared_ptr[cpp_BufferResource] create( + any_resource[device_accessible], + optional[cpp_PinnedMemoryResource], + unordered_map[MemoryType, int64_t], + optional[cpp_Duration], + shared_ptr[cuda_stream_pool], + shared_ptr[cpp_Statistics], + ) except +ex_handler size_t memory_reserved(MemoryType mem_type) except +ex_handler int64_t memory_available(MemoryType mem_type) except +ex_handler void set_memory_limit(MemoryType mem_type, int64_t limit) except +ex_handler diff --git a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx index 3dfef2171..1fcac4c71 100644 --- a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx +++ b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx @@ -5,7 +5,7 @@ from cython cimport no_gc_clear from cython.operator cimport dereference as deref from libc.stdint cimport int64_t from libcpp cimport bool as bool_t -from libcpp.memory cimport make_shared, shared_ptr, unique_ptr +from libcpp.memory cimport shared_ptr, unique_ptr from libcpp.optional cimport optional from libcpp.pair cimport pair from libcpp.unordered_map cimport unordered_map @@ -197,7 +197,7 @@ cdef class BufferResource: if self._pinned_mr is not None: cpp_pinned_mr = self._pinned_mr._handle with nogil: - self._handle = make_shared[cpp_BufferResource]( + self._handle = cpp_BufferResource.create( make_any_device_resource(device_mr.get_mr()), cpp_pinned_mr, move(_mem_limits), From 4665282e57b80e67bf642d6970a03266d3ef9649 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Fri, 29 May 2026 11:33:23 +0200 Subject: [PATCH 03/15] OwningResourceAdaptor --- .../rapidsmpf/memory/buffer_resource.hpp | 156 +++++++----------- .../memory/owning_resource_adaptor.hpp | 153 +++++++++++++++++ .../rapidsmpf/rmm_resource_adaptor.hpp | 33 ---- cpp/src/memory/buffer_resource.cpp | 11 +- cpp/src/rmm_resource_adaptor.cpp | 8 - cpp/tests/test_buffer_resource.cpp | 26 ++- 6 files changed, 232 insertions(+), 155 deletions(-) create mode 100644 cpp/include/rapidsmpf/memory/owning_resource_adaptor.hpp diff --git a/cpp/include/rapidsmpf/memory/buffer_resource.hpp b/cpp/include/rapidsmpf/memory/buffer_resource.hpp index 6c17d6bb8..1df83c45c 100644 --- a/cpp/include/rapidsmpf/memory/buffer_resource.hpp +++ b/cpp/include/rapidsmpf/memory/buffer_resource.hpp @@ -51,76 +51,36 @@ enum class AllowOverbooking : bool { * (e.g., host and device). All memory operations in rapidsmpf, such as those performed * by the Shuffler, rely on a buffer resource for memory management. * - * @note Similar to RMM's memory resource, the `BufferResource` instance must outlive all - * allocated buffers and memory reservations. + * @note `BufferResource` instances must be constructed through `create()` or + * `from_options()`, both of which return a `std::shared_ptr`. Direct + * construction is disabled. */ class BufferResource : public std::enable_shared_from_this { public: /** - * @brief Tag struct used to make `BufferResource`'s constructor effectively - * private. The tag is public so that `BufferResource::create` can - * forward it, but user code can't write `BufferResource{...}` directly. - * All construction must go through `create()` or `from_options()`, which - * guarantees the result is owned by a `std::shared_ptr`. - */ - struct PrivateTag { - explicit PrivateTag() = default; - }; - - /** - * @brief Constructs a buffer resource. + * @brief Construct a `BufferResource` managed by `std::shared_ptr`. * - * Memory availability is computed per `MemoryType` as `limit - allocated`. + * 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, so the available memory - * always equals the configured limit. + * 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. * - * @param device_mr The RMM device memory resource used for device allocations. - * @param pinned_mr The pinned host memory resource used for `MemoryType::PINNED_HOST` - * allocations. If disabled, pinned host allocations are unavailable regardless of - * @p memory_limits. - * @param memory_limits Maximum bytes per memory type before overbooking or spilling. - * Missing entries default to unlimited. - * @param periodic_spill_check Enable periodic spill checks. A dedicated thread - * continuously checks and performs spilling based on memory availability. The value - * of `periodic_spill_check` controls the pause between checks. If `std::nullopt`, no - * periodic spill checking is performed. - * @param stream_pool Pool of CUDA streams used throughout RapidsMPF for operations - * that do not take an explicit CUDA stream. - * @param statistics The statistics instance to use (disabled by default). - */ - BufferResource( - PrivateTag, - cuda::mr::any_resource device_mr, - std::optional pinned_mr = PinnedMemoryResource::Disabled, - std::unordered_map memory_limits = {}, - std::optional periodic_spill_check = std::chrono::milliseconds{1}, - std::shared_ptr stream_pool = std::make_shared< - rmm::cuda_stream_pool>(16, rmm::cuda_stream::flags::non_blocking), - std::shared_ptr statistics = Statistics::disabled() - ); - - /** - * @brief Factory method to construct a `BufferResource` as a `shared_ptr`. - * - * This is the standard way to construct a `BufferResource`. The returned - * handle is the canonical owner; pass `std::shared_ptr` - * (or a raw pointer to it) wherever a `BufferResource` is needed. - * - * @param device_mr The RMM device memory resource used for device allocations. - * @param pinned_mr Pinned host memory resource (or `PinnedMemoryResource::Disabled`). - * @param memory_limits Per-`MemoryType` byte limits. Missing entries default to - * unlimited. - * @param periodic_spill_check Pause between periodic spill checks - * (`std::nullopt` disables the dedicated thread). - * @param stream_pool CUDA stream pool used for operations without an explicit stream. - * @param statistics Statistics instance (defaults to disabled). - * - * @return A `std::shared_ptr` owning a newly constructed `BufferResource`. + * @param device_mr Device memory resource used for device allocations. + * @param pinned_mr Pinned host memory resource used for + * `MemoryType::PINNED_HOST` allocations, or `PinnedMemoryResource::Disabled` to + * disable pinned allocations. + * @param memory_limits Maximum allocation limits in bytes per `MemoryType`. Missing + * entries are treated as unlimited. + * @param periodic_spill_check Interval between periodic spill checks. `std::nullopt` + * disables the dedicated spill-check thread. + * @param stream_pool CUDA stream pool used for operations that do not take an + * explicit CUDA stream. + * @param statistics Statistics instance used for runtime metrics. + * @return A newly constructed `BufferResource` owned by `std::shared_ptr`. */ [[nodiscard]] static std::shared_ptr create( cuda::mr::any_resource device_mr, @@ -142,7 +102,6 @@ class BufferResource : public std::enable_shared_from_this { * @param mr A device-accessible RMM memory resource. * @param options Configuration options. * @param statistics The statistics instance to use (disabled by default). - * * @return A shared pointer to a BufferResource instance configured according to the * options. */ @@ -154,54 +113,55 @@ class BufferResource : public std::enable_shared_from_this { ~BufferResource() noexcept = default; - /// @brief `BufferResource` is non-copyable, it is owned by `std::shared_ptr`. + /** + * @brief `BufferResource` is non-copyable, it is owned by `std::shared_ptr`. + */ BufferResource(BufferResource const&) = delete; - /// @brief `BufferResource` is non-movable, it is owned by `std::shared_ptr`. + /** + * @brief `BufferResource` is non-movable, it is owned by `std::shared_ptr`. + */ BufferResource(BufferResource&&) = delete; - /// @brief `BufferResource` is non-copyable. - /// @return Reference to this. + /** + * @brief `BufferResource` is non-copyable. + * @return Reference to this. + */ BufferResource& operator=(BufferResource const&) = delete; - /// @brief `BufferResource` is non-movable. - /// @return Reference to this. + /** + * @brief `BufferResource` is non-movable. + * @return Reference to this. + */ BufferResource& operator=(BufferResource&&) = delete; /** * @brief Get the device memory resource as an owning, type-erased handle. * - * Returns a `cuda::mr::any_resource` that owns a copy - * of the internal `RmmResourceAdaptor` carrying an owning back-pointer - * (`shared_ptr`) to this `BufferResource`. Anything that - * stores the returned object — e.g. `rmm::device_buffer`'s internal - * `any_resource` member, or a Python `DeviceMemoryResource` wrapper — - * keeps this `BufferResource` (including its stream pool, spill - * manager, and statistics) alive for as long as the returned MR exists. - * - * Use this when you need the MR to **outlive** the caller's - * `shared_ptr` (e.g. allocating a device buffer that - * is returned out of a scope where the BR is dropped). For transient - * cuDF/RMM calls that take an `rmm::device_async_resource_ref` and - * complete before the caller's BR ref drops, prefer `device_mr_ref()`. - * - * @return An owning `cuda::mr::any_resource` that - * keeps this `BufferResource` alive. + * The returned memory resource keeps this `BufferResource` alive for as long + * as the returned handle exists. + * + * Use this accessor when the memory resource may outlive the caller's + * `std::shared_ptr`, for example when storing it in a buffer, + * wrapper, or other object with an independent lifetime. + * + * For transient cuDF/RMM calls that only need a non-owning + * `rmm::device_async_resource_ref`, prefer `device_mr_ref()`. + * + * @return Owning device memory resource handle. */ [[nodiscard]] cuda::mr::any_resource device_mr(); /** * @brief Get a non-owning reference to the device memory resource. * - * Returns an `rmm::device_async_resource_ref` pointing into this - * `BufferResource`'s internal storage. The ref is cheap to obtain and - * binds directly to cuDF/RMM APIs that take - * `rmm::device_async_resource_ref` parameters. The caller must keep this - * `BufferResource` alive for as long as the returned ref is in use. + * Returns an `rmm::device_async_resource_ref` referring to the device memory + * resource managed by this `BufferResource`. The returned reference is inexpensive + * to copy and integrates directly with cuDF and RMM APIs that accept + * `rmm::device_async_resource_ref`. * - * This is the right accessor for the common case of passing the device - * MR to a cuDF/RMM function call. Use `device_mr()` (owning) when the MR - * needs to outlive this `BufferResource`. + * The returned reference does not extend the lifetime of this `BufferResource`. The + * caller must ensure the `BufferResource` remains alive for as long as the reference + * is used. * - * @return Non-owning `rmm::device_async_resource_ref` into BR-owned - * storage. + * @return Non-owning reference to the device memory resource. */ [[nodiscard]] rmm::device_async_resource_ref device_mr_ref() const noexcept; @@ -506,6 +466,16 @@ class BufferResource : public std::enable_shared_from_this { std::shared_ptr statistics() const noexcept; private: + /** @brief Private constructor, use `create()` or `from_options()`. */ + BufferResource( + cuda::mr::any_resource device_mr, + std::optional pinned_mr, + std::unordered_map memory_limits, + std::optional periodic_spill_check, + std::shared_ptr stream_pool, + std::shared_ptr statistics + ); + std::mutex mutex_; // The internal RmmResourceAdaptor wraps the user's device MR so that // allocations are tracked for the DEVICE memory_available calculation. diff --git a/cpp/include/rapidsmpf/memory/owning_resource_adaptor.hpp b/cpp/include/rapidsmpf/memory/owning_resource_adaptor.hpp new file mode 100644 index 000000000..1ccdfc85d --- /dev/null +++ b/cpp/include/rapidsmpf/memory/owning_resource_adaptor.hpp @@ -0,0 +1,153 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include + +#include +#include + +#include + +namespace rapidsmpf { + +/** + * @brief CCCL-compatible memory resource adaptor that wraps another + * resource and carries additional owned state. + * + * Allocation, deallocation, and property queries are forwarded to the + * wrapped resource. The adaptor additionally stores an arbitrary + * value-semantic object whose lifetime is tied to the adaptor itself. + * + * This is useful when a resource must keep related state alive for as long + * as the resource object exists, for example shared ownership of runtime + * state, pools, contexts, or other resources required by the wrapped + * allocator implementation. + * + * Properties exposed by the wrapped resource (via `get_property` ADL) are + * forwarded transparently through `cuda::forward_property`. + * + * @tparam Resource The wrapped CCCL-compatible resource type. Must satisfy + * the `cuda::mr::resource` concept and be copyable. + * @tparam Owner The owned state type carried alongside the resource. + */ +template +class OwningResourceAdaptor + : public cuda::forward_property, Resource> { + public: + /** + * @brief Construct an owning wrapper. + * + * @param resource The wrapped resource (typically a refcount-shared handle, e.g. + * `RmmResourceAdaptor`, copied cheaply). + * @param back_ref Owning back-reference kept alive for the lifetime of the wrapper. + */ + OwningResourceAdaptor(Resource resource, BackRef back_ref) noexcept + : resource_{std::move(resource)}, back_ref_{std::move(back_ref)} {} + + /** + * @brief Accessor used by `cuda::forward_property` to query properties on the wrapped + * resource. + * + * @return Reference to the wrapped resource. + */ + [[nodiscard]] Resource& get() noexcept { + return resource_; + } + + /** @copydoc get() */ + [[nodiscard]] Resource const& get() const noexcept { + return resource_; + } + + /** + * @brief Allocate memory on a CUDA stream. + * + * Forwards the allocation request to the wrapped resource. + * + * @param stream CUDA stream associated with the allocation. + * @param bytes Number of bytes to allocate. + * @param alignment Required allocation alignment. + * @return Pointer to the allocated memory. + */ + void* allocate( + cuda::stream_ref stream, + std::size_t bytes, + std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT + ) { + return resource_.allocate(stream, bytes, alignment); + } + + /** + * @brief Deallocate memory on a CUDA stream. + * + * Forwards the deallocation request to the wrapped resource. + * + * @param stream CUDA stream associated with the deallocation. + * @param ptr Pointer to the memory to deallocate. + * @param bytes Number of bytes originally allocated. + * @param alignment Alignment of the original allocation. + */ + void deallocate( + cuda::stream_ref stream, + void* ptr, + std::size_t bytes, + std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT + ) noexcept { + resource_.deallocate(stream, ptr, bytes, alignment); + } + + /** + * @brief Allocate memory synchronously. + * + * Forwards the allocation request to the wrapped resource. + * + * @param bytes Number of bytes to allocate. + * @param alignment Required allocation alignment. + * @return Pointer to the allocated memory. + */ + void* allocate_sync( + std::size_t bytes, std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT + ) { + return resource_.allocate_sync(bytes, alignment); + } + + /** + * @brief Deallocate memory synchronously. + * + * Forwards the deallocation request to the wrapped resource. + * + * @param ptr Pointer to the memory to deallocate. + * @param bytes Number of bytes originally allocated. + * @param alignment Alignment of the original allocation. + */ + void deallocate_sync( + void* ptr, + std::size_t bytes, + std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT + ) noexcept { + resource_.deallocate_sync(ptr, bytes, alignment); + } + + /** + * @brief Compare wrapped resource identity. + * + * Two adaptors compare equal if their wrapped resources compare equal. + * + * @param other Adaptor to compare against. + * @return `true` if both adaptors refer to the same underlying resource. + */ + [[nodiscard]] bool operator==(OwningResourceAdaptor const& other) const noexcept { + return resource_ == other.resource_; + } + + private: + Resource resource_; + BackRef back_ref_; +}; + +} // namespace rapidsmpf diff --git a/cpp/include/rapidsmpf/rmm_resource_adaptor.hpp b/cpp/include/rapidsmpf/rmm_resource_adaptor.hpp index 4d256918c..408820492 100644 --- a/cpp/include/rapidsmpf/rmm_resource_adaptor.hpp +++ b/cpp/include/rapidsmpf/rmm_resource_adaptor.hpp @@ -7,7 +7,6 @@ #include #include -#include #include @@ -18,8 +17,6 @@ namespace rapidsmpf { -class BufferResource; // defined in memory/buffer_resource.hpp - /** * @brief A RMM memory resource adaptor tailored to RapidsMPF. * @@ -123,36 +120,6 @@ class RmmResourceAdaptor * @see begin_scoped_memory_record() */ ScopedMemoryRecord end_scoped_memory_record(); - - /** - * @brief Return a copy of this adaptor that owns a back-pointer to @p br. - * - * The returned adaptor shares allocation-tracking state with `*this` via - * the underlying `shared_resource` (so allocations through either copy - * are accumulated into the same record); it additionally holds an - * owning `shared_ptr` so that anything storing the - * returned object (e.g. an `rmm::device_buffer`'s internal - * `any_resource`) keeps @p br — including its stream pool, spill - * manager, and statistics — alive. - * - * Used by `BufferResource::device_mr()` to produce a lifetime-safe - * owning device memory resource. The original adaptor stored inside - * `BufferResource` itself does **not** carry the back-pointer, so no - * self-cycle is created. - * - * @param br The `BufferResource` to keep alive via the returned adaptor. - * @return A copy of this adaptor with @p br stored as a back-pointer. - */ - [[nodiscard]] RmmResourceAdaptor with_buffer_resource_ref( - std::shared_ptr br - ) const; - - private: - /// @brief Optional owning back-pointer to the `BufferResource` containing - /// this adaptor. Null in the original adaptor stored inside `BufferResource` - /// (to avoid a self-cycle). Populated in copies returned via - /// `with_buffer_resource_ref()` / `BufferResource::device_mr()`. - std::shared_ptr br_; }; static_assert(cuda::mr::resource_with); diff --git a/cpp/src/memory/buffer_resource.cpp b/cpp/src/memory/buffer_resource.cpp index 57ce87d62..1f7a8cfc0 100644 --- a/cpp/src/memory/buffer_resource.cpp +++ b/cpp/src/memory/buffer_resource.cpp @@ -15,6 +15,7 @@ #include #include #include +#include #include #include #include @@ -22,7 +23,6 @@ namespace rapidsmpf { BufferResource::BufferResource( - PrivateTag, cuda::mr::any_resource device_mr, std::optional pinned_mr, std::unordered_map memory_limits, @@ -58,15 +58,16 @@ std::shared_ptr BufferResource::create( std::shared_ptr stream_pool, std::shared_ptr statistics ) { - return std::make_shared( - PrivateTag{}, + // Use `new` rather than `std::make_shared` so the private ctor can be + // called via this member function (`make_shared` would need access). + return std::shared_ptr{new BufferResource{ std::move(device_mr), std::move(pinned_mr), std::move(memory_limits), periodic_spill_check, std::move(stream_pool), std::move(statistics) - ); + }}; } std::shared_ptr BufferResource::from_options( @@ -115,7 +116,7 @@ void BufferResource::set_memory_limit(MemoryType mem_type, std::int64_t limit) n cuda::mr::any_resource BufferResource::device_mr() { return cuda::mr::any_resource{ - device_adaptor_.with_buffer_resource_ref(shared_from_this()) + OwningResourceAdaptor{device_adaptor_, shared_from_this()} }; } diff --git a/cpp/src/rmm_resource_adaptor.cpp b/cpp/src/rmm_resource_adaptor.cpp index 1d9716728..f7a223e8c 100644 --- a/cpp/src/rmm_resource_adaptor.cpp +++ b/cpp/src/rmm_resource_adaptor.cpp @@ -40,12 +40,4 @@ ScopedMemoryRecord RmmResourceAdaptor::end_scoped_memory_record() { return get().end_scoped_memory_record(); } -RmmResourceAdaptor RmmResourceAdaptor::with_buffer_resource_ref( - std::shared_ptr br -) const { - RmmResourceAdaptor copy{*this}; - copy.br_ = std::move(br); - return copy; -} - } // namespace rapidsmpf diff --git a/cpp/tests/test_buffer_resource.cpp b/cpp/tests/test_buffer_resource.cpp index da23077df..515b098c7 100644 --- a/cpp/tests/test_buffer_resource.cpp +++ b/cpp/tests/test_buffer_resource.cpp @@ -19,6 +19,7 @@ #include #include #include +#include #include #include #include @@ -788,12 +789,6 @@ TEST_F(BufferCopyEdgeCases, SameBufferIsDisallowed) { ); } -// The owning `any_resource` returned by `device_mr()` carries an owning -// back-ref to the BR via the adaptor copy. Dropping the original BR while -// the returned MR (and any buffer allocated through it) is still alive -// must not free the BR — its stream pool must remain valid for the -// buffer's stream-ordered destruction. Conversely, once everything that -// observed the back-ref drops it, BR must destruct (no refcount cycle). TEST(BufferResource, DeviceMrOutlivesBufferResource) { constexpr std::size_t N = 1024; @@ -802,17 +797,16 @@ TEST(BufferResource, DeviceMrOutlivesBufferResource) { auto stream = cudf::get_default_stream(); auto owning_mr = br->device_mr(); - auto* adaptor_in_mr = cuda::mr::resource_cast(&owning_mr); - ASSERT_NE(adaptor_in_mr, nullptr); + auto* wrapper = cuda::mr::resource_cast< + OwningResourceAdaptor>>( + &owning_mr + ); + ASSERT_NE(wrapper, nullptr); auto buf = std::make_unique(N, stream, std::move(owning_mr)); - // The adaptor in the buffer's `any_resource` shares allocation tracking - // with the original adaptor stored in BR, so the allocation is recorded. - EXPECT_GE(adaptor_in_mr->current_allocated(), static_cast(N)); - - // Drop the original BR. The adaptor copy inside `buf`'s `any_resource` - // owns the only remaining `shared_ptr`, so BR (and its + // Drop the original BR. The wrapper inside `buf`'s `any_resource` owns + // the only remaining `shared_ptr`, so BR (and its // stream pool) must stay alive. br.reset(); EXPECT_FALSE(weak_br.expired()) << "BR freed while buffer still holds it"; @@ -822,7 +816,7 @@ TEST(BufferResource, DeviceMrOutlivesBufferResource) { EXPECT_NO_THROW(buf.reset()); // After the buffer is gone, no shared_ptr
exists anywhere. If there - // were a refcount cycle (e.g. adaptor-copy back-ref leaking back into + // were a refcount cycle (e.g. the wrapper's back-ref leaking back into // the original BR), `weak_br` would still be observable here. - EXPECT_TRUE(weak_br.expired()) << "BR not destructed — refcount cycle?"; + EXPECT_TRUE(weak_br.expired()) << "BR not destructed, refcount cycle?"; } From 58ee92aecd0fa90a352fc8936cb1a97f18568a56 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Fri, 29 May 2026 12:38:59 +0200 Subject: [PATCH 04/15] cleanup --- .../memory/owning_resource_adaptor.hpp | 2 +- .../rapidsmpf/memory/spill_manager.hpp | 24 +++++++++---------- cpp/src/memory/buffer_resource.cpp | 2 -- cpp/tests/test_buffer_resource.cpp | 16 ++++++------- 4 files changed, 20 insertions(+), 24 deletions(-) diff --git a/cpp/include/rapidsmpf/memory/owning_resource_adaptor.hpp b/cpp/include/rapidsmpf/memory/owning_resource_adaptor.hpp index 1ccdfc85d..5ea5886b4 100644 --- a/cpp/include/rapidsmpf/memory/owning_resource_adaptor.hpp +++ b/cpp/include/rapidsmpf/memory/owning_resource_adaptor.hpp @@ -1,5 +1,5 @@ /** - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/include/rapidsmpf/memory/spill_manager.hpp b/cpp/include/rapidsmpf/memory/spill_manager.hpp index 2969a3bfd..79dfe5372 100644 --- a/cpp/include/rapidsmpf/memory/spill_manager.hpp +++ b/cpp/include/rapidsmpf/memory/spill_manager.hpp @@ -27,20 +27,18 @@ class SpillManager { /** * @brief Spill function type. * - * A spill function takes a requested spill amount as input and returns the actual - * amount of memory (in bytes) that was spilled. + * A spill function receives a requested spill size in bytes and returns the actual + * number of bytes spilled. * - * @warning Spill functions must not capture an owning reference (a - * `shared_ptr`, a `BufferResource` by value, etc.) to - * the `BufferResource` that owns this `SpillManager`. They must also not - * capture any object — `Buffer`, `rmm::device_buffer`, or other holder - * of memory allocated through `BufferResource::device_mr()` — that - * transitively owns the same `BufferResource`. Either case closes a - * refcount cycle (BR → SpillManager → SpillFunction → BR), which is - * invisible to Python's garbage collector and leaks the resource - * unconditionally. Capture raw pointers/references only, and unregister - * via `remove_spill_function()` from the owner's destructor before - * dropping any buffer-holding members. + * 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 SpillFunction = std::function; diff --git a/cpp/src/memory/buffer_resource.cpp b/cpp/src/memory/buffer_resource.cpp index 1f7a8cfc0..679217f7c 100644 --- a/cpp/src/memory/buffer_resource.cpp +++ b/cpp/src/memory/buffer_resource.cpp @@ -58,8 +58,6 @@ std::shared_ptr BufferResource::create( std::shared_ptr stream_pool, std::shared_ptr statistics ) { - // Use `new` rather than `std::make_shared` so the private ctor can be - // called via this member function (`make_shared` would need access). return std::shared_ptr{new BufferResource{ std::move(device_mr), std::move(pinned_mr), diff --git a/cpp/tests/test_buffer_resource.cpp b/cpp/tests/test_buffer_resource.cpp index 515b098c7..29f2522a7 100644 --- a/cpp/tests/test_buffer_resource.cpp +++ b/cpp/tests/test_buffer_resource.cpp @@ -805,18 +805,18 @@ TEST(BufferResource, DeviceMrOutlivesBufferResource) { auto buf = std::make_unique(N, stream, std::move(owning_mr)); - // Drop the original BR. The wrapper inside `buf`'s `any_resource` owns - // the only remaining `shared_ptr`, so BR (and its - // stream pool) must stay alive. + // Drop the original BR. The wrapper inside `buf`'s `any_resource` owns the only + // remaining `shared_ptr`, so BR (and its stream pool) must stay + // alive. br.reset(); EXPECT_FALSE(weak_br.expired()) << "BR freed while buffer still holds it"; - // Destroy the buffer — its stream-ordered deallocator runs against the - // still-alive stream pool. Must not crash. + // Destroy the buffer: its stream-ordered deallocator runs against the still-alive + // stream pool. Must not crash. EXPECT_NO_THROW(buf.reset()); - // After the buffer is gone, no shared_ptr
exists anywhere. If there - // were a refcount cycle (e.g. the wrapper's back-ref leaking back into - // the original BR), `weak_br` would still be observable here. + // After the buffer is gone, no shared_ptr
exists anywhere. If there were a + // refcount cycle (e.g. the wrapper's back-ref leaking back into the original BR), + // `weak_br` would still be observable here. EXPECT_TRUE(weak_br.expired()) << "BR not destructed, refcount cycle?"; } From 5919f12acdf5b3a4b30c3125ed8f2f65eee6cc01 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Fri, 29 May 2026 13:07:35 +0200 Subject: [PATCH 05/15] docs --- cpp/include/rapidsmpf/memory/buffer_resource.hpp | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/cpp/include/rapidsmpf/memory/buffer_resource.hpp b/cpp/include/rapidsmpf/memory/buffer_resource.hpp index 1df83c45c..28240b77d 100644 --- a/cpp/include/rapidsmpf/memory/buffer_resource.hpp +++ b/cpp/include/rapidsmpf/memory/buffer_resource.hpp @@ -142,9 +142,6 @@ class BufferResource : public std::enable_shared_from_this { * `std::shared_ptr`, for example when storing it in a buffer, * wrapper, or other object with an independent lifetime. * - * For transient cuDF/RMM calls that only need a non-owning - * `rmm::device_async_resource_ref`, prefer `device_mr_ref()`. - * * @return Owning device memory resource handle. */ [[nodiscard]] cuda::mr::any_resource device_mr(); @@ -153,13 +150,13 @@ class BufferResource : public std::enable_shared_from_this { * @brief Get a non-owning reference to the device memory resource. * * Returns an `rmm::device_async_resource_ref` referring to the device memory - * resource managed by this `BufferResource`. The returned reference is inexpensive - * to copy and integrates directly with cuDF and RMM APIs that accept - * `rmm::device_async_resource_ref`. + * resource managed by this `BufferResource`. The returned reference does not + * extend the lifetime of the `BufferResource`. * - * The returned reference does not extend the lifetime of this `BufferResource`. The - * caller must ensure the `BufferResource` remains alive for as long as the reference - * is used. + * This accessor is intended for transient use, for example when passing a + * memory resource into a cuDF or RMM API call. It must not be used in storage + * paths where the memory resource may outlive the owning `BufferResource`. + * For such cases, use `device_mr()` instead. * * @return Non-owning reference to the device memory resource. */ From 07cb0c20e4f47bf092bc7c69c7dfa2cf5bc30701 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Sun, 31 May 2026 21:12:21 +0200 Subject: [PATCH 06/15] only device_mr() --- cpp/benchmarks/bench_comm.cpp | 2 +- cpp/benchmarks/bench_shuffle.cpp | 4 +- cpp/benchmarks/streaming/data_generator.hpp | 2 +- cpp/benchmarks/streaming/ndsh/bench_read.cpp | 2 +- cpp/benchmarks/streaming/ndsh/concatenate.cpp | 3 +- cpp/benchmarks/streaming/ndsh/groupby.cpp | 2 +- cpp/benchmarks/streaming/ndsh/join.cpp | 21 +-- cpp/benchmarks/streaming/ndsh/q01.cpp | 20 +-- cpp/benchmarks/streaming/ndsh/q03.cpp | 21 ++- cpp/benchmarks/streaming/ndsh/q04.cpp | 6 +- cpp/benchmarks/streaming/ndsh/q09.cpp | 14 +- cpp/benchmarks/streaming/ndsh/q21.cpp | 28 +-- cpp/benchmarks/streaming/ndsh/sort.cpp | 4 +- .../rapidsmpf/memory/buffer_resource.hpp | 85 ++++++--- .../memory/owning_resource_adaptor.hpp | 169 ++++++++++++++---- .../streaming/cudf/channel_metadata.hpp | 2 +- cpp/src/integrations/cudf/partition.cpp | 20 +-- cpp/src/memory/buffer_resource.cpp | 24 +-- cpp/src/streaming/cudf/bloom_filter.cpp | 6 +- cpp/src/streaming/cudf/channel_metadata.cpp | 6 +- cpp/src/streaming/cudf/parquet.cpp | 2 +- cpp/src/streaming/cudf/table_chunk.cpp | 7 +- cpp/tests/streaming/test_leaf_actor.cpp | 2 +- cpp/tests/streaming/test_partition.cpp | 4 +- cpp/tests/streaming/test_read_parquet.cpp | 2 +- cpp/tests/streaming/test_shuffler.cpp | 2 +- cpp/tests/test_allgather.cpp | 2 +- cpp/tests/test_buffer.cpp | 2 +- cpp/tests/test_buffer_resource.cpp | 83 +++++++-- cpp/tests/utils.hpp | 2 +- 30 files changed, 356 insertions(+), 193 deletions(-) diff --git a/cpp/benchmarks/bench_comm.cpp b/cpp/benchmarks/bench_comm.cpp index 62875b2db..afc70faf7 100644 --- a/cpp/benchmarks/bench_comm.cpp +++ b/cpp/benchmarks/bench_comm.cpp @@ -220,7 +220,7 @@ Duration run( auto [res, _] = br->reserve(MemoryType::DEVICE, args.msg_size * 2, AllowOverbooking::YES); auto buf = br->make_buffer(args.msg_size, stream, res); - random_fill(*buf, br->device_mr_ref()); + random_fill(*buf, br->device_mr()); send_bufs.push_back(std::move(buf)); recv_bufs.push_back(br->make_buffer(args.msg_size, stream, res)); } diff --git a/cpp/benchmarks/bench_shuffle.cpp b/cpp/benchmarks/bench_shuffle.cpp index 070eedf9c..36710c4fb 100644 --- a/cpp/benchmarks/bench_shuffle.cpp +++ b/cpp/benchmarks/bench_shuffle.cpp @@ -302,7 +302,7 @@ rapidsmpf::Duration do_run( { RAPIDSMPF_NVTX_SCOPED_RANGE("Shuffling", total_num_partitions); if (args.enable_memory_profiler) { - RAPIDSMPF_MEMORY_PROFILE(statistics, br->device_mr_ref(), "shuffling"); + RAPIDSMPF_MEMORY_PROFILE(statistics, br->device_mr(), "shuffling"); } rapidsmpf::shuffler::Shuffler shuffler( comm, @@ -390,7 +390,7 @@ std::vector generate_input_partitions( size_lb, args.input_data_allow_overbooking ); cudf::table table = random_table( - num_columns, num_local_rows, min_val, max_val, stream, br->device_mr_ref() + num_columns, num_local_rows, min_val, max_val, stream, br->device_mr() ); input_partitions.emplace_back(transform_fn(std::move(table))); } diff --git a/cpp/benchmarks/streaming/data_generator.hpp b/cpp/benchmarks/streaming/data_generator.hpp index 531c153d4..9bb53795c 100644 --- a/cpp/benchmarks/streaming/data_generator.hpp +++ b/cpp/benchmarks/streaming/data_generator.hpp @@ -62,7 +62,7 @@ inline Actor random_table_generator( seq, std::make_unique( std::make_unique(random_table( - ncolumns, nrows, min_val, max_val, stream, ctx->br()->device_mr_ref() + ncolumns, nrows, min_val, max_val, stream, ctx->br()->device_mr() )), stream ) diff --git a/cpp/benchmarks/streaming/ndsh/bench_read.cpp b/cpp/benchmarks/streaming/ndsh/bench_read.cpp index 53d32c9c0..7d75f9687 100644 --- a/cpp/benchmarks/streaming/ndsh/bench_read.cpp +++ b/cpp/benchmarks/streaming/ndsh/bench_read.cpp @@ -413,7 +413,7 @@ int main(int argc, char** argv) { timings.push_back(compute.count()); auto statistics = ctx->statistics(); comm->logger()->print(statistics->report( - {.mr = ctx->br()->device_mr_ref(), .pinned_mr = ctx->br()->try_pinned_mr()} + {.mr = ctx->br()->device_mr(), .pinned_mr = ctx->br()->try_pinned_mr()} )); statistics->clear(); } diff --git a/cpp/benchmarks/streaming/ndsh/concatenate.cpp b/cpp/benchmarks/streaming/ndsh/concatenate.cpp index e0e9661d8..39ce71720 100644 --- a/cpp/benchmarks/streaming/ndsh/concatenate.cpp +++ b/cpp/benchmarks/streaming/ndsh/concatenate.cpp @@ -70,8 +70,7 @@ streaming::Actor concatenate( chunks.push_back(std::move(chunk)); } auto result = std::make_unique( - cudf::concatenate(views, concat_stream, ctx->br()->device_mr_ref()), - concat_stream + cudf::concatenate(views, concat_stream, ctx->br()->device_mr()), concat_stream ); cuda_stream_join( chunks | std::views::transform([](auto&& chunk) { return chunk.stream(); }), diff --git a/cpp/benchmarks/streaming/ndsh/groupby.cpp b/cpp/benchmarks/streaming/ndsh/groupby.cpp index fd4234613..5a13054c2 100644 --- a/cpp/benchmarks/streaming/ndsh/groupby.cpp +++ b/cpp/benchmarks/streaming/ndsh/groupby.cpp @@ -55,7 +55,7 @@ streaming::Actor chunkwise_group_by( cudf::groupby::groupby(table.select(keys), null_policy, cudf::sorted::NO); auto [keys, aggregated] = - grouper.aggregate(agg_requests, stream, ctx->br()->device_mr_ref()); + grouper.aggregate(agg_requests, stream, ctx->br()->device_mr()); std::ignore = std::move(chunk); auto result = keys->release(); for (auto&& a : aggregated) { diff --git a/cpp/benchmarks/streaming/ndsh/join.cpp b/cpp/benchmarks/streaming/ndsh/join.cpp index 3123e1901..65ff8be52 100644 --- a/cpp/benchmarks/streaming/ndsh/join.cpp +++ b/cpp/benchmarks/streaming/ndsh/join.cpp @@ -69,8 +69,7 @@ coro::task broadcast( ); } else { RAPIDSMPF_EXPECTS(chunks.size() > 0, "No chunks in broadcast"); - auto result = - cudf::concatenate(views, gather_stream, ctx->br()->device_mr_ref()); + auto result = cudf::concatenate(views, gather_stream, ctx->br()->device_mr()); // So that deallocation of the consitutent tables is stream-ordered wrt the // concatenation. cuda_stream_join( @@ -94,9 +93,8 @@ coro::task broadcast( // TODO: If this chunk is already in pack form, this is unnecessary. auto chunk = co_await msg.release().make_available(ctx); - auto pack = cudf::pack( - chunk.table_view(), chunk.stream(), ctx->br()->device_mr_ref() - ); + auto pack = + cudf::pack(chunk.table_view(), chunk.stream(), ctx->br()->device_mr()); auto packed_data = PackedData( std::move(pack.metadata), ctx->br()->move(std::move(pack.gpu_data), chunk.stream()) @@ -185,7 +183,7 @@ streaming::Message semi_join_chunk( ); auto match = joiner.semi_join( - left_chunk.table_view().select(left_on), chunk_stream, ctx->br()->device_mr_ref() + left_chunk.table_view().select(left_on), chunk_stream, ctx->br()->device_mr() ); ctx->logger()->debug( @@ -199,7 +197,7 @@ streaming::Message semi_join_chunk( indices, cudf::out_of_bounds_policy::DONT_CHECK, chunk_stream, - ctx->br()->device_mr_ref() + ctx->br()->device_mr() ) ->release(); @@ -245,9 +243,8 @@ streaming::Message inner_join_chunk( build_event->stream_wait(chunk_stream); auto probe_table = right_chunk.table_view(); auto probe_keys = probe_table.select(right_on); - auto [probe_match, build_match] = joiner.inner_join( - probe_keys, std::nullopt, chunk_stream, ctx->br()->device_mr_ref() - ); + auto [probe_match, build_match] = + joiner.inner_join(probe_keys, std::nullopt, chunk_stream, ctx->br()->device_mr()); cudf::column_view build_indices = cudf::device_span(*build_match); @@ -261,7 +258,7 @@ streaming::Message inner_join_chunk( build_indices, cudf::out_of_bounds_policy::DONT_CHECK, chunk_stream, - ctx->br()->device_mr_ref() + ctx->br()->device_mr() ) ->release(); // drop key columns from probe table. @@ -277,7 +274,7 @@ streaming::Message inner_join_chunk( probe_indices, cudf::out_of_bounds_policy::DONT_CHECK, chunk_stream, - ctx->br()->device_mr_ref() + ctx->br()->device_mr() ) ->release(), std::back_inserter(result_columns) diff --git a/cpp/benchmarks/streaming/ndsh/q01.cpp b/cpp/benchmarks/streaming/ndsh/q01.cpp index 1c6d4b138..d4dd6975f 100644 --- a/cpp/benchmarks/streaming/ndsh/q01.cpp +++ b/cpp/benchmarks/streaming/ndsh/q01.cpp @@ -136,7 +136,7 @@ rapidsmpf::streaming::Actor postprocess_group_by( co_await msg.release().make_available(ctx); auto stream = chunk.stream(); auto columns = - cudf::table{chunk.table_view(), stream, ctx->br()->device_mr_ref()}.release(); + cudf::table{chunk.table_view(), stream, ctx->br()->device_mr()}.release(); std::ignore = std::move(chunk); auto count = std::move(columns.back()); columns.pop_back(); @@ -150,7 +150,7 @@ rapidsmpf::streaming::Actor postprocess_group_by( cudf::binary_operator::TRUE_DIV, cudf::data_type(cudf::type_id::FLOAT64), stream, - ctx->br()->device_mr_ref() + ctx->br()->device_mr() ) ); } @@ -161,7 +161,7 @@ rapidsmpf::streaming::Actor postprocess_group_by( cudf::binary_operator::TRUE_DIV, cudf::data_type(cudf::type_id::FLOAT64), stream, - ctx->br()->device_mr_ref() + ctx->br()->device_mr() ) ); columns.push_back(std::move(count)); @@ -202,9 +202,7 @@ rapidsmpf::streaming::Actor select_columns_for_groupby( auto table = chunk.table_view(); // l_returnflag, l_linestatus, l_quantity, l_extendedprice auto result = - cudf::table( - table.select({0, 1, 2, 3}), chunk_stream, ctx->br()->device_mr_ref() - ) + cudf::table(table.select({0, 1, 2, 3}), chunk_stream, ctx->br()->device_mr()) .release(); result.reserve(7); auto extendedprice = table.column(3); @@ -235,7 +233,7 @@ static __device__ void calculate_charge(double *charge, double discprice, double std::nullopt, cudf::output_nullability::PRESERVE, chunk_stream, - ctx->br()->device_mr_ref() + ctx->br()->device_mr() ) ); // charge @@ -250,14 +248,12 @@ static __device__ void calculate_charge(double *charge, double discprice, double std::nullopt, cudf::output_nullability::PRESERVE, chunk_stream, - ctx->br()->device_mr_ref() + ctx->br()->device_mr() ) ); // l_discount result.push_back( - std::make_unique( - discount, chunk_stream, ctx->br()->device_mr_ref() - ) + std::make_unique(discount, chunk_stream, ctx->br()->device_mr()) ); co_await ch_out->send( rapidsmpf::streaming::to_message( @@ -439,7 +435,7 @@ int main(int argc, char** argv) { timings.push_back(compute.count()); auto statistics = ctx->statistics(); comm->logger()->print(statistics->report( - {.mr = ctx->br()->device_mr_ref(), .pinned_mr = ctx->br()->try_pinned_mr()} + {.mr = ctx->br()->device_mr(), .pinned_mr = ctx->br()->try_pinned_mr()} )); statistics->clear(); } diff --git a/cpp/benchmarks/streaming/ndsh/q03.cpp b/cpp/benchmarks/streaming/ndsh/q03.cpp index 190995682..89a922a80 100644 --- a/cpp/benchmarks/streaming/ndsh/q03.cpp +++ b/cpp/benchmarks/streaming/ndsh/q03.cpp @@ -234,19 +234,19 @@ rapidsmpf::streaming::Actor select_columns_for_groupby( // o_orderkey result.push_back( std::make_unique( - table.column(0), chunk_stream, ctx->br()->device_mr_ref() + table.column(0), chunk_stream, ctx->br()->device_mr() ) ); // o_orderdate result.push_back( std::make_unique( - table.column(1), chunk_stream, ctx->br()->device_mr_ref() + table.column(1), chunk_stream, ctx->br()->device_mr() ) ); // o_shippriority result.push_back( std::make_unique( - table.column(2), chunk_stream, ctx->br()->device_mr_ref() + table.column(2), chunk_stream, ctx->br()->device_mr() ) ); auto extendedprice = table.column(3); @@ -270,7 +270,7 @@ static __device__ void calculate_revenue(double *revenue, double extprice, doubl std::nullopt, cudf::output_nullability::PRESERVE, chunk_stream, - ctx->br()->device_mr_ref() + ctx->br()->device_mr() ) ); co_await ch_out->send( @@ -311,7 +311,7 @@ rapidsmpf::streaming::Actor top_k_by( order, {}, chunk.stream(), - ctx->br()->device_mr_ref() + ctx->br()->device_mr() ); partials.push_back( cudf::gather( @@ -319,7 +319,7 @@ rapidsmpf::streaming::Actor top_k_by( cudf::split(indices->view(), {k}, chunk.stream()).front(), cudf::out_of_bounds_policy::DONT_CHECK, chunk.stream(), - ctx->br()->device_mr_ref() + ctx->br()->device_mr() ) ); chunk_streams.push_back(chunk.stream()); @@ -336,12 +336,11 @@ rapidsmpf::streaming::Actor top_k_by( std::ranges::transform(partials, std::back_inserter(views), [](auto& t) { return t->view(); }); - auto merged = - cudf::merge(views, keys, order, {}, out_stream, ctx->br()->device_mr_ref()); + auto merged = cudf::merge(views, keys, order, {}, out_stream, ctx->br()->device_mr()); auto result = std::make_unique( cudf::slice(merged->view(), {0, 10}, out_stream), out_stream, - ctx->br()->device_mr_ref() + ctx->br()->device_mr() ); co_await ch_out->send( rapidsmpf::streaming::to_message( @@ -382,7 +381,7 @@ rapidsmpf::streaming::Actor fanout_bounded( std::make_unique( chunk.table_view().select(ch1_cols), chunk.stream(), - ctx->br()->device_mr_ref() + ctx->br()->device_mr() ), chunk.stream() ) @@ -698,7 +697,7 @@ int main(int argc, char** argv) { timings.push_back(compute.count()); auto statistics = ctx->statistics(); comm->logger()->print(statistics->report( - {.mr = ctx->br()->device_mr_ref(), .pinned_mr = ctx->br()->try_pinned_mr()} + {.mr = ctx->br()->device_mr(), .pinned_mr = ctx->br()->try_pinned_mr()} )); statistics->clear(); } diff --git a/cpp/benchmarks/streaming/ndsh/q04.cpp b/cpp/benchmarks/streaming/ndsh/q04.cpp index 52ef7a178..a76c3fd80 100644 --- a/cpp/benchmarks/streaming/ndsh/q04.cpp +++ b/cpp/benchmarks/streaming/ndsh/q04.cpp @@ -151,7 +151,7 @@ rapidsmpf::streaming::Actor filter_lineitem( ) { rapidsmpf::streaming::ShutdownAtExit c{ch_in, ch_out}; co_await ctx->executor()->schedule(); - auto mr = ctx->br()->device_mr_ref(); + auto mr = ctx->br()->device_mr(); while (!ch_out->is_shutdown()) { auto msg = co_await ch_in->receive(); @@ -218,7 +218,7 @@ rapidsmpf::streaming::Actor fanout_bounded( std::make_unique( chunk.table_view().select(ch1_cols), chunk.stream(), - ctx->br()->device_mr_ref() + ctx->br()->device_mr() ), chunk.stream() ) @@ -520,7 +520,7 @@ int main(int argc, char** argv) { timings.push_back(compute.count()); auto statistics = ctx->statistics(); comm->logger()->print(statistics->report( - {.mr = ctx->br()->device_mr_ref(), .pinned_mr = ctx->br()->try_pinned_mr()} + {.mr = ctx->br()->device_mr(), .pinned_mr = ctx->br()->try_pinned_mr()} )); statistics->clear(); } diff --git a/cpp/benchmarks/streaming/ndsh/q09.cpp b/cpp/benchmarks/streaming/ndsh/q09.cpp index cf450fae2..a52ae40f2 100644 --- a/cpp/benchmarks/streaming/ndsh/q09.cpp +++ b/cpp/benchmarks/streaming/ndsh/q09.cpp @@ -175,7 +175,7 @@ rapidsmpf::streaming::Actor filter_part( std::shared_ptr ch_out ) { rapidsmpf::streaming::ShutdownAtExit c{ch_in, ch_out}; - auto mr = ctx->br()->device_mr_ref(); + auto mr = ctx->br()->device_mr(); while (true) { auto msg = co_await ch_in->receive(); if (msg.empty()) { @@ -233,7 +233,7 @@ rapidsmpf::streaming::Actor select_columns( // n_name result.push_back( std::make_unique( - table.column(0), chunk_stream, ctx->br()->device_mr_ref() + table.column(0), chunk_stream, ctx->br()->device_mr() ) ); result.push_back( @@ -241,7 +241,7 @@ rapidsmpf::streaming::Actor select_columns( table.column(5), cudf::datetime::datetime_component::YEAR, chunk_stream, - ctx->br()->device_mr_ref() + ctx->br()->device_mr() ) ); auto discount = table.column(2); @@ -267,7 +267,7 @@ static __device__ void calculate_amount(double *amount, double discount, double std::nullopt, cudf::output_nullability::PRESERVE, chunk_stream, - ctx->br()->device_mr_ref() + ctx->br()->device_mr() ) ); co_await ch_out->send( @@ -312,7 +312,7 @@ rapidsmpf::streaming::Actor round_sum_profit( 2, cudf::rounding_method::HALF_EVEN, chunk.stream(), - ctx->br()->device_mr_ref() + ctx->br()->device_mr() ); #pragma GCC diagnostic pop auto result = rapidsmpf::streaming::to_message( @@ -321,7 +321,7 @@ rapidsmpf::streaming::Actor round_sum_profit( std::make_unique( cudf::table_view({table.column(0), table.column(1), rounded->view()}), chunk.stream(), - ctx->br()->device_mr_ref() + ctx->br()->device_mr() ), chunk.stream() ) @@ -655,7 +655,7 @@ int main(int argc, char** argv) { timings.push_back(compute.count()); auto statistics = ctx->statistics(); comm->logger()->print(statistics->report( - {.mr = ctx->br()->device_mr_ref(), .pinned_mr = ctx->br()->try_pinned_mr()} + {.mr = ctx->br()->device_mr(), .pinned_mr = ctx->br()->try_pinned_mr()} )); statistics->clear(); } diff --git a/cpp/benchmarks/streaming/ndsh/q21.cpp b/cpp/benchmarks/streaming/ndsh/q21.cpp index 369ba46a1..3b7e6a052 100644 --- a/cpp/benchmarks/streaming/ndsh/q21.cpp +++ b/cpp/benchmarks/streaming/ndsh/q21.cpp @@ -98,7 +98,7 @@ rapidsmpf::streaming::Actor read_nation( constexpr auto name = "SAUDI ARABIA"; owner->push_back( std::make_shared( - name, /* is_valid = */ true, stream, ctx->br()->device_mr_ref() + name, /* is_valid = */ true, stream, ctx->br()->device_mr() ) ); owner->push_back( @@ -156,7 +156,7 @@ rapidsmpf::streaming::Actor read_orders( constexpr auto status = "F"; owner->push_back( std::make_shared( - status, /* is_valid = */ true, stream, ctx->br()->device_mr_ref() + status, /* is_valid = */ true, stream, ctx->br()->device_mr() ) ); owner->push_back( @@ -279,7 +279,7 @@ rapidsmpf::streaming::Actor filter_lineitem( cudf::binary_operator::GREATER, cudf::data_type(cudf::type_id::BOOL8), chunk.stream(), - ctx->br()->device_mr_ref() + ctx->br()->device_mr() ); co_await ch_out->send( rapidsmpf::streaming::to_message( @@ -289,7 +289,7 @@ rapidsmpf::streaming::Actor filter_lineitem( chunk.table_view().select({0, 1}), mask->view(), chunk.stream(), - ctx->br()->device_mr_ref() + ctx->br()->device_mr() ), chunk.stream() ) @@ -318,12 +318,12 @@ rapidsmpf::streaming::Actor filter_grouped_greater( auto mask = cudf::binary_operation( chunk.table_view().column(1), cudf::numeric_scalar( - 1, /* is_valid = */ true, chunk.stream(), ctx->br()->device_mr_ref() + 1, /* is_valid = */ true, chunk.stream(), ctx->br()->device_mr() ), cudf::binary_operator::GREATER, cudf::data_type(cudf::type_id::BOOL8), chunk.stream(), - ctx->br()->device_mr_ref() + ctx->br()->device_mr() ); latch->count_down(); co_await ch_out->send( @@ -334,7 +334,7 @@ rapidsmpf::streaming::Actor filter_grouped_greater( chunk.table_view().select({0}), mask->view(), chunk.stream(), - ctx->br()->device_mr_ref() + ctx->br()->device_mr() ), chunk.stream() ) @@ -362,12 +362,12 @@ rapidsmpf::streaming::Actor filter_grouped_equal( auto mask = cudf::binary_operation( chunk.table_view().column(1), cudf::numeric_scalar( - 1, /* is_valid = */ true, chunk.stream(), ctx->br()->device_mr_ref() + 1, /* is_valid = */ true, chunk.stream(), ctx->br()->device_mr() ), cudf::binary_operator::EQUAL, cudf::data_type(cudf::type_id::BOOL8), chunk.stream(), - ctx->br()->device_mr_ref() + ctx->br()->device_mr() ); co_await ch_out->send( rapidsmpf::streaming::to_message( @@ -377,7 +377,7 @@ rapidsmpf::streaming::Actor filter_grouped_equal( chunk.table_view().select({0}), mask->view(), chunk.stream(), - ctx->br()->device_mr_ref() + ctx->br()->device_mr() ), chunk.stream() ) @@ -415,7 +415,7 @@ rapidsmpf::streaming::Actor fanout_bounded( std::make_unique( chunk.table_view().select(ch1_cols), chunk.stream(), - ctx->br()->device_mr_ref() + ctx->br()->device_mr() ), chunk.stream() ) @@ -505,7 +505,7 @@ rapidsmpf::streaming::Actor slice( auto sliced_table = std::make_unique( cudf::slice(chunk.table_view(), {local_start, local_end})[0], chunk.stream(), - ctx->br()->device_mr_ref() + ctx->br()->device_mr() ); co_await ch_out->send( rapidsmpf::streaming::to_message( @@ -574,7 +574,7 @@ rapidsmpf::streaming::Actor populate_bloom_filter( ); auto stream = chunk.stream(); auto out = std::make_unique( - chunk.table_view().select(keys), stream, ctx->br()->device_mr_ref() + chunk.table_view().select(keys), stream, ctx->br()->device_mr() ); std::ignore = std::move(chunk); co_await passthrough->send( @@ -997,7 +997,7 @@ int main(int argc, char** argv) { timings.push_back(compute.count()); auto statistics = ctx->statistics(); comm->logger()->print(statistics->report( - {.mr = ctx->br()->device_mr_ref(), .pinned_mr = ctx->br()->try_pinned_mr()} + {.mr = ctx->br()->device_mr(), .pinned_mr = ctx->br()->try_pinned_mr()} )); statistics->clear(); } diff --git a/cpp/benchmarks/streaming/ndsh/sort.cpp b/cpp/benchmarks/streaming/ndsh/sort.cpp index b300ce4ed..756d2bc65 100644 --- a/cpp/benchmarks/streaming/ndsh/sort.cpp +++ b/cpp/benchmarks/streaming/ndsh/sort.cpp @@ -36,7 +36,7 @@ rapidsmpf::streaming::Actor chunkwise_sort_by( order, null_order, chunk.stream(), - ctx->br()->device_mr_ref() + ctx->br()->device_mr() ); } else { return cudf::sort_by_key( @@ -45,7 +45,7 @@ rapidsmpf::streaming::Actor chunkwise_sort_by( order, null_order, chunk.stream(), - ctx->br()->device_mr_ref() + ctx->br()->device_mr() ); } }; diff --git a/cpp/include/rapidsmpf/memory/buffer_resource.hpp b/cpp/include/rapidsmpf/memory/buffer_resource.hpp index 28240b77d..7924d0131 100644 --- a/cpp/include/rapidsmpf/memory/buffer_resource.hpp +++ b/cpp/include/rapidsmpf/memory/buffer_resource.hpp @@ -23,6 +23,7 @@ #include #include #include +#include #include #include #include @@ -133,34 +134,55 @@ class BufferResource : public std::enable_shared_from_this { BufferResource& operator=(BufferResource&&) = delete; /** - * @brief Get the device memory resource as an owning, type-erased handle. - * - * The returned memory resource keeps this `BufferResource` alive for as long - * as the returned handle exists. - * - * Use this accessor when the memory resource may outlive the caller's - * `std::shared_ptr`, for example when storing it in a buffer, - * wrapper, or other object with an independent lifetime. - * - * @return Owning device memory resource handle. - */ - [[nodiscard]] cuda::mr::any_resource device_mr(); - - /** - * @brief Get a non-owning reference to the device memory resource. - * - * Returns an `rmm::device_async_resource_ref` referring to the device memory - * resource managed by this `BufferResource`. The returned reference does not - * extend the lifetime of the `BufferResource`. - * - * This accessor is intended for transient use, for example when passing a - * memory resource into a cuDF or RMM API call. It must not be used in storage - * paths where the memory resource may outlive the owning `BufferResource`. - * For such cases, use `device_mr()` instead. - * - * @return Non-owning reference to the device memory resource. + * @brief Get the device memory resource. + * + * @return `rmm::device_async_resource_ref` to the device memory resource. + * + * @par CCCL's lifetime semantic + * + * The returned `rmm::device_async_resource_ref` is a non-owning + * `cuda::mr::resource_ref`, so callers must take care to avoid use-after-free issues. + * + * When working directly with the returned reference, the caller must ensure that this + * `BufferResource` remains alive for the full duration of that use: + * @code + * auto br = std::make_shared(...); + * 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 + * @endcode + * + * To store the resource beyond the immediate call, promote the ref to an + * owning `cuda::mr::any_resource`: + * @code + * auto br = std::make_shared(...); + * cuda::mr::any_resource mr = br->device_mr(); + * br.reset(); // safe: `mr` keeps the BufferResource alive + * mr.allocate(...); // safe + * @endcode + * + * In the common case, no explicit promotion is needed because RMM and cuDF containers + * that store a memory resource do this internally: + * @code + * auto br = std::make_shared(...); + * rmm::device_buffer buf{1024, stream, br->device_mr()}; + * br.reset(); // safe: `buf` keeps the BufferResource alive internally + * @endcode + * + * Returned objects from cuDF APIs typically behave the same way: + * @code + * auto br = std::make_shared(...); + * auto col = cudf::make_numeric_column( + * cudf::data_type{cudf::type_id::INT32}, + * 1000, + * cudf::mask_state::UNALLOCATED, + * stream, + * br->device_mr() + * ); + * br.reset(); // safe: `col` keeps the BufferResource alive internally + * @endcode */ - [[nodiscard]] rmm::device_async_resource_ref device_mr_ref() const noexcept; + [[nodiscard]] rmm::device_async_resource_ref device_mr(); /** * @brief Get the RMM host memory resource. @@ -476,9 +498,14 @@ class BufferResource : public std::enable_shared_from_this { std::mutex mutex_; // The internal RmmResourceAdaptor wraps the user's device MR so that // allocations are tracked for the DEVICE memory_available calculation. - // Declared before device_mr_ because device_mr_ is initialized from it. RmmResourceAdaptor device_adaptor_; - cuda::mr::any_resource device_mr_; + // Stable storage for the device MR returned by `device_mr()`. Holds a + // `weak_ptr` to *this; copies promote that to a + // `shared_ptr` on the deep-copy path. Installed by + // `create()` after construction so that `weak_from_this()` is valid; + // wrapped in `optional` so it can be default-constructed empty + // (`RmmResourceAdaptor` itself is not default-constructible). + std::optional> owning_mr_; std::optional pinned_mr_; HostMemoryResource host_mr_; std::array, MEMORY_TYPES.size()> memory_limits_; diff --git a/cpp/include/rapidsmpf/memory/owning_resource_adaptor.hpp b/cpp/include/rapidsmpf/memory/owning_resource_adaptor.hpp index 5ea5886b4..d51df15e8 100644 --- a/cpp/include/rapidsmpf/memory/owning_resource_adaptor.hpp +++ b/cpp/include/rapidsmpf/memory/owning_resource_adaptor.hpp @@ -6,6 +6,7 @@ #pragma once #include +#include #include #include @@ -16,42 +17,147 @@ namespace rapidsmpf { /** - * @brief CCCL-compatible memory resource adaptor that wraps another - * resource and carries additional owned state. + * @brief CCCL-compatible memory resource adaptor that makes copies keep an + * external owner alive. * - * Allocation, deallocation, and property queries are forwarded to the - * wrapped resource. The adaptor additionally stores an arbitrary - * value-semantic object whose lifetime is tied to the adaptor itself. + * Wraps another CCCL-compatible resource and forwards allocation, + * deallocation, and property queries to it. The adaptor is useful when an + * object owns runtime state needed by the wrapped resource, but exposes that + * resource through a non-owning `cuda::mr::resource_ref`. * - * This is useful when a resource must keep related state alive for as long - * as the resource object exists, for example shared ownership of runtime - * state, pools, contexts, or other resources required by the wrapped - * allocator implementation. + * This makes it safe for downstream code to promote that ref to an owning + * `cuda::mr::any_resource`: the copied adaptor stored in the + * `any_resource` keeps the back-referenced owner alive. * - * Properties exposed by the wrapped resource (via `get_property` ADL) are - * forwarded transparently through `cuda::forward_property`. + * Properties exposed by the wrapped resource through `get_property` ADL are + * forwarded transparently via `cuda::forward_property`. * - * @tparam Resource The wrapped CCCL-compatible resource type. Must satisfy - * the `cuda::mr::resource` concept and be copyable. - * @tparam Owner The owned state type carried alongside the resource. + * @tparam Resource Wrapped CCCL-compatible resource type. Must satisfy the + * `cuda::mr::resource` concept and be copyable. + * @tparam BackRef Type of the back-referenced owner object. + * + * @par CCCL's lifetime semantic + * + * `cuda::mr::resource_ref<...>` is a non-owning type-erased view. It stores a + * pointer to a concrete resource object together with a vtable. Copying a + * `resource_ref` only copies that pointer and vtable. It does not copy the + * resource and does not extend its lifetime. The caller must ensure that the + * referenced resource object remains alive for every direct use of the ref. + * + * `cuda::mr::any_resource<...>` is an owning type-erased resource. Constructing + * an `any_resource` from a `resource_ref` performs a deep copy of the + * referenced concrete resource into storage owned by the `any_resource`. That + * deep copy invokes the concrete resource's copy constructor. + * + * A deep copy of the resource is not necessarily enough to make the result + * self-contained. If the copied resource depends on state owned outside the + * resource object, for example through a back-pointer, reference, or shared + * runtime state, then the copy preserves that dependency. Without additional + * lifetime management, the original owner can be destroyed while the copied + * resource inside the `any_resource` still depends on it. + * + * This adaptor makes that copy step acquire ownership of the back-referenced + * owner. The original adaptor normally lives inside `*back_ref` and stores + * only a `std::weak_ptr`, avoiding a reference cycle. When CCCL copies + * the adaptor, such as during `any_resource(resource_ref)` construction, the + * adaptor's copy constructor promotes the weak reference to a + * `std::shared_ptr`. The copied adaptor therefore keeps the owner + * alive for as long as the copied resource exists. + * + * If the owner has already been destroyed when the adaptor is copied, the copy + * constructor throws `std::bad_weak_ptr` instead of producing a dangling + * adaptor. + * + * This does not make direct use of a bare `resource_ref` safe after the owner + * has been destroyed. A call such as `ref.allocate(...)` still dispatches + * through the `resource_ref` into the original resource object, so the owner + * must outlive all direct uses of the non-owning ref. */ template class OwningResourceAdaptor : public cuda::forward_property, Resource> { public: /** - * @brief Construct an owning wrapper. + * @brief Construct an adaptor that weakly references @p back_ref. + * + * The original adaptor instance typically lives inside `*back_ref` and stores only a + * `std::weak_ptr`, avoiding a reference cycle. Copies of the adaptor promote + * that weak reference to a `std::shared_ptr`, allowing copied resources to + * keep the owner alive. + * + * @param resource Wrapped resource. + * @param back_ref Weak reference to the back-referenced owner object. + */ + OwningResourceAdaptor(Resource resource, std::weak_ptr back_ref) noexcept + : resource_{std::move(resource)}, weak_{std::move(back_ref)} {} + + /** + * @brief Copy constructor that promotes the back-reference to shared ownership. + * + * This constructor is the key lifetime hook used by `any_resource(resource_ref)`. + * When CCCL deep-copies the wrapped resource into an `any_resource`, it invokes this + * copy constructor. The copied adaptor acquires a `std::shared_ptr`, + * ensuring the back-referenced owner stays alive for as long as the copied resource + * exists. + * + * If @p other already holds a strong back-reference, it is reused. Otherwise the weak + * reference is promoted via `std::shared_ptr` construction. + * + * @param other Adaptor to copy from. + * @throws std::bad_weak_ptr if the weak back-reference has expired. + */ + OwningResourceAdaptor(OwningResourceAdaptor const& other) + : resource_{other.resource_}, + weak_{other.weak_}, + strong_{ + other.strong_ ? other.strong_ + : std::shared_ptr{other.weak_} // throws if expired + } {} + + /** + * @brief Copy assignment operator that promotes the back-reference to shared + * ownership. + * + * Uses the same lifetime semantics as the copy constructor. After the assignment, + * this adaptor holds a strong reference to the back-referenced owner object. * - * @param resource The wrapped resource (typically a refcount-shared handle, e.g. - * `RmmResourceAdaptor`, copied cheaply). - * @param back_ref Owning back-reference kept alive for the lifetime of the wrapper. + * @param other Adaptor to copy from. + * @return Reference to this adaptor. + * @throws std::bad_weak_ptr if the weak back-reference has expired. */ - OwningResourceAdaptor(Resource resource, BackRef back_ref) noexcept - : resource_{std::move(resource)}, back_ref_{std::move(back_ref)} {} + OwningResourceAdaptor& operator=(OwningResourceAdaptor const& other) { + if (this != &other) { + resource_ = other.resource_; + weak_ = other.weak_; + strong_ = other.strong_ + ? other.strong_ + : std::shared_ptr{other.weak_}; // throws if expired + } + return *this; + } /** - * @brief Accessor used by `cuda::forward_property` to query properties on the wrapped - * resource. + * @brief Move constructor. + * + * @param other Adaptor to move from. + */ + OwningResourceAdaptor(OwningResourceAdaptor&& other) noexcept = default; + + /** + * @brief Move assignment operator. + * + * @param other Adaptor to move from. + * @return Reference to this adaptor. + */ + OwningResourceAdaptor& operator=(OwningResourceAdaptor&& other) noexcept = default; + + ~OwningResourceAdaptor() = default; + + /** + * @brief Accessor used by `cuda::forward_property`. + * + * Property queries received through `cuda::forward_property` are forwarded to the + * wrapped resource returned by this accessor. * * @return Reference to the wrapped resource. */ @@ -65,9 +171,7 @@ class OwningResourceAdaptor } /** - * @brief Allocate memory on a CUDA stream. - * - * Forwards the allocation request to the wrapped resource. + * @brief Allocate memory on a CUDA stream. Forwards to the wrapped resource. * * @param stream CUDA stream associated with the allocation. * @param bytes Number of bytes to allocate. @@ -83,9 +187,7 @@ class OwningResourceAdaptor } /** - * @brief Deallocate memory on a CUDA stream. - * - * Forwards the deallocation request to the wrapped resource. + * @brief Deallocate memory on a CUDA stream. Forwards to the wrapped resource. * * @param stream CUDA stream associated with the deallocation. * @param ptr Pointer to the memory to deallocate. @@ -102,9 +204,7 @@ class OwningResourceAdaptor } /** - * @brief Allocate memory synchronously. - * - * Forwards the allocation request to the wrapped resource. + * @brief Allocate memory synchronously. Forwards to the wrapped resource. * * @param bytes Number of bytes to allocate. * @param alignment Required allocation alignment. @@ -117,9 +217,7 @@ class OwningResourceAdaptor } /** - * @brief Deallocate memory synchronously. - * - * Forwards the deallocation request to the wrapped resource. + * @brief Deallocate memory synchronously. Forwards to the wrapped resource. * * @param ptr Pointer to the memory to deallocate. * @param bytes Number of bytes originally allocated. @@ -147,7 +245,8 @@ class OwningResourceAdaptor private: Resource resource_; - BackRef back_ref_; + std::weak_ptr weak_; + std::shared_ptr strong_; }; } // namespace rapidsmpf diff --git a/cpp/include/rapidsmpf/streaming/cudf/channel_metadata.hpp b/cpp/include/rapidsmpf/streaming/cudf/channel_metadata.hpp index a94aee21f..0e27c3f63 100644 --- a/cpp/include/rapidsmpf/streaming/cudf/channel_metadata.hpp +++ b/cpp/include/rapidsmpf/streaming/cudf/channel_metadata.hpp @@ -125,7 +125,7 @@ struct OrderScheme { * attributes, and the schemes are otherwise compatible (same order and null_order). */ [[nodiscard]] bool boundaries_aligned_with( - OrderScheme const& other, rapidsmpf::BufferResource const& br + OrderScheme const& other, rapidsmpf::BufferResource& br ) const; }; diff --git a/cpp/src/integrations/cudf/partition.cpp b/cpp/src/integrations/cudf/partition.cpp index a3f2c6241..b423c0803 100644 --- a/cpp/src/integrations/cudf/partition.cpp +++ b/cpp/src/integrations/cudf/partition.cpp @@ -36,10 +36,10 @@ partition_and_split( BufferResource* br, AllowOverbooking allow_overbooking ) { - RAPIDSMPF_MEMORY_PROFILE(br->statistics(), br->device_mr_ref()); + RAPIDSMPF_MEMORY_PROFILE(br->statistics(), br->device_mr()); if (table.num_rows() == 0) { // Return views of a copy of the empty `table`. - auto owner = std::make_unique(table, stream, br->device_mr_ref()); + auto owner = std::make_unique(table, stream, br->device_mr()); return { std::vector( safe_cast(num_partitions), owner->view() @@ -60,7 +60,7 @@ partition_and_split( hash_function, seed, stream, - br->device_mr_ref() + br->device_mr() ); reservation.clear(); @@ -88,7 +88,7 @@ std::unordered_map partition_and_pack( AllowOverbooking allow_overbooking ) { RAPIDSMPF_NVTX_FUNC_RANGE(); - RAPIDSMPF_MEMORY_PROFILE(br->statistics(), br->device_mr_ref()); + RAPIDSMPF_MEMORY_PROFILE(br->statistics(), br->device_mr()); RAPIDSMPF_EXPECTS(num_partitions > 0, "Need to split to at least one partition"); if (table.num_rows() == 0) { auto splits = @@ -108,7 +108,7 @@ std::unordered_map partition_and_pack( hash_function, seed, stream, - br->device_mr_ref() + br->device_mr() ); reservation.clear(); std::vector splits(split_points.begin() + 1, split_points.end() - 1); @@ -123,15 +123,15 @@ std::unordered_map split_and_pack( AllowOverbooking allow_overbooking ) { RAPIDSMPF_NVTX_FUNC_RANGE(); - RAPIDSMPF_MEMORY_PROFILE(br->statistics(), br->device_mr_ref()); + RAPIDSMPF_MEMORY_PROFILE(br->statistics(), br->device_mr()); std::unordered_map ret; // contiguous split does a deep-copy. Therefore, we need to reserve memory for // at least the size of the table. auto reservation = br->reserve_device_memory_and_spill( - cudf::packed_size(table, stream, br->device_mr_ref()), allow_overbooking + cudf::packed_size(table, stream, br->device_mr()), allow_overbooking ); - auto packed = cudf::contiguous_split(table, splits, stream, br->device_mr_ref()); + auto packed = cudf::contiguous_split(table, splits, stream, br->device_mr()); reservation.clear(); ret.reserve(packed.size()); for (shuffler::PartID i = 0; safe_cast(i) < packed.size(); i++) { @@ -153,7 +153,7 @@ std::unique_ptr unpack_and_concat( AllowOverbooking allow_overbooking ) { RAPIDSMPF_NVTX_FUNC_RANGE(); - RAPIDSMPF_MEMORY_PROFILE(br->statistics(), br->device_mr_ref()); + RAPIDSMPF_MEMORY_PROFILE(br->statistics(), br->device_mr()); // Let's find the total size of the partitions and how much of the packed data we // need to move to device memory (unspill). @@ -205,7 +205,7 @@ std::unique_ptr unpack_and_concat( } reservation = br->reserve_device_memory_and_spill(total_size, allow_overbooking); - return cudf::concatenate(unpacked, stream, br->device_mr_ref()); + return cudf::concatenate(unpacked, stream, br->device_mr()); } std::vector spill_partitions( diff --git a/cpp/src/memory/buffer_resource.cpp b/cpp/src/memory/buffer_resource.cpp index 679217f7c..bd6db16ce 100644 --- a/cpp/src/memory/buffer_resource.cpp +++ b/cpp/src/memory/buffer_resource.cpp @@ -31,7 +31,6 @@ BufferResource::BufferResource( std::shared_ptr statistics ) : device_adaptor_{std::move(device_mr)}, - device_mr_{device_adaptor_}, // any_resource shares state via shared_resource pinned_mr_{std::move(pinned_mr)}, host_mr_{}, stream_pool_{std::move(stream_pool)}, @@ -58,7 +57,8 @@ std::shared_ptr BufferResource::create( std::shared_ptr stream_pool, std::shared_ptr statistics ) { - return std::shared_ptr{new BufferResource{ + // Construct via `new` (private ctor visible here as a member function). + std::shared_ptr br{new BufferResource{ std::move(device_mr), std::move(pinned_mr), std::move(memory_limits), @@ -66,6 +66,14 @@ std::shared_ptr BufferResource::create( std::move(stream_pool), std::move(statistics) }}; + // Install the owning adaptor *after* construction so that + // `weak_from_this()` is valid. The adaptor holds only a `weak_ptr` to + // `*br` so no refcount cycle. Copies of the adaptor (made by CCCL's + // `any_resource(resource_ref)` deep-copy in downstream consumers) promote + // that weak ref to a `shared_ptr`, transferring ownership + // to the copy. + br->owning_mr_.emplace(br->device_adaptor_, br->weak_from_this()); + return br; } std::shared_ptr BufferResource::from_options( @@ -112,16 +120,8 @@ void BufferResource::set_memory_limit(MemoryType mem_type, std::int64_t limit) n ); } -cuda::mr::any_resource BufferResource::device_mr() { - return cuda::mr::any_resource{ - OwningResourceAdaptor{device_adaptor_, shared_from_this()} - }; -} - -rmm::device_async_resource_ref BufferResource::device_mr_ref() const noexcept { - return rmm::device_async_resource_ref{ - const_cast&>(device_mr_) - }; +rmm::device_async_resource_ref BufferResource::device_mr() { + return rmm::device_async_resource_ref{*owning_mr_}; } rmm::host_async_resource_ref BufferResource::host_mr() noexcept { diff --git a/cpp/src/streaming/cudf/bloom_filter.cpp b/cpp/src/streaming/cudf/bloom_filter.cpp index 37b054f9d..ac41e10ee 100644 --- a/cpp/src/streaming/cudf/bloom_filter.cpp +++ b/cpp/src/streaming/cudf/bloom_filter.cpp @@ -26,7 +26,7 @@ Actor BloomFilter::build( co_await ch_in->shutdown_metadata(); co_await ch_out->shutdown_metadata(); auto const& br = ctx_->br(); - auto mr = br->device_mr_ref(); + auto mr = br->device_mr(); auto filter_stream = br->stream_pool().get_stream(); CudaEvent event; auto storage = rapidsmpf::BloomFilter::storage(num_filter_blocks_, filter_stream, mr); @@ -127,7 +127,7 @@ Actor BloomFilter::apply( 0 ); auto mask = filter.contains( - chunk.table_view().select(keys), chunk_stream, ctx_->br()->device_mr_ref() + chunk.table_view().select(keys), chunk_stream, ctx_->br()->device_mr() ); cuda_stream_join(stream, chunk_stream, &event); RAPIDSMPF_EXPECTS( @@ -142,7 +142,7 @@ Actor BloomFilter::apply( 0 }; auto result = cudf::apply_boolean_mask( - chunk.table_view(), mask_view, chunk_stream, ctx_->br()->device_mr_ref() + chunk.table_view(), mask_view, chunk_stream, ctx_->br()->device_mr() ); std::ignore = std::move(chunk); std::ignore = std::move(res); diff --git a/cpp/src/streaming/cudf/channel_metadata.cpp b/cpp/src/streaming/cudf/channel_metadata.cpp index 9bf0c1aa9..fc6bec396 100644 --- a/cpp/src/streaming/cudf/channel_metadata.cpp +++ b/cpp/src/streaming/cudf/channel_metadata.cpp @@ -57,7 +57,7 @@ OrderScheme OrderScheme::with_keys(std::vector new_keys) const { } bool OrderScheme::boundaries_aligned_with( - OrderScheme const& other, rapidsmpf::BufferResource const& br + OrderScheme const& other, rapidsmpf::BufferResource& br ) const { if (strict_boundaries != other.strict_boundaries || boundaries->shape() != other.boundaries->shape()) @@ -89,14 +89,14 @@ bool OrderScheme::boundaries_aligned_with( cudf::binary_operator::NULL_EQUALS, cudf::data_type{cudf::type_id::BOOL8}, stream, - br.device_mr_ref() + br.device_mr() ); auto result = cudf::reduce( eq->view(), *cudf::make_all_aggregation(), cudf::data_type{cudf::type_id::BOOL8}, stream, - br.device_mr_ref() + br.device_mr() ); auto& scalar = static_cast&>(*result); if (!scalar.value(stream)) { diff --git a/cpp/src/streaming/cudf/parquet.cpp b/cpp/src/streaming/cudf/parquet.cpp index b7609ec5f..2edd16625 100644 --- a/cpp/src/streaming/cudf/parquet.cpp +++ b/cpp/src/streaming/cudf/parquet.cpp @@ -200,7 +200,7 @@ Message read_parquet_chunk( return to_message( sequence_number, std::make_unique( - cudf::io::read_parquet(options, stream, ctx->br()->device_mr_ref()).tbl, + cudf::io::read_parquet(options, stream, ctx->br()->device_mr()).tbl, stream ) ); diff --git a/cpp/src/streaming/cudf/table_chunk.cpp b/cpp/src/streaming/cudf/table_chunk.cpp index f296f99a4..afef6385f 100644 --- a/cpp/src/streaming/cudf/table_chunk.cpp +++ b/cpp/src/streaming/cudf/table_chunk.cpp @@ -181,7 +181,7 @@ TableChunk TableChunk::copy(MemoryReservation& reservation) const { auto statistics = br->statistics(); StreamOrderedTiming timing{stream(), statistics}; auto table = std::make_unique( - table_view(), stream(), br->device_mr_ref() + table_view(), stream(), br->device_mr() ); statistics->record_copy( MemoryType::DEVICE, MemoryType::DEVICE, nbytes, std::move(timing) @@ -217,8 +217,7 @@ TableChunk TableChunk::copy(MemoryReservation& reservation) const { // new host buffer. // TODO: use `cudf::chunked_pack()` with a bounce buffer. Currently, // `cudf::pack()` allocates device memory we haven't reserved. - auto packed_columns = - cudf::pack(table_view(), stream(), br->device_mr_ref()); + auto packed_columns = cudf::pack(table_view(), stream(), br->device_mr()); auto packed_data = std::make_unique( std::move(packed_columns.metadata), br->move(std::move(packed_columns.gpu_data), stream()) @@ -270,7 +269,7 @@ std::unique_ptr TableChunk::into_packed_data(BufferResource* br) && ); // TODO: use `cudf::chunked_pack()` with a bounce buffer. Currently, // `cudf::pack()` allocates device memory we haven't reserved. - auto packed_columns = cudf::pack(table_view_.value(), stream_, br->device_mr_ref()); + auto packed_columns = cudf::pack(table_view_.value(), stream_, br->device_mr()); table_view_ = std::nullopt; return std::make_unique( std::move(packed_columns.metadata), diff --git a/cpp/tests/streaming/test_leaf_actor.cpp b/cpp/tests/streaming/test_leaf_actor.cpp index 18e8c5492..37ec6b088 100644 --- a/cpp/tests/streaming/test_leaf_actor.cpp +++ b/cpp/tests/streaming/test_leaf_actor.cpp @@ -52,7 +52,7 @@ TEST_F(StreamingLeafTasks, PushAndPullChunks) { i, std::make_unique( std::make_unique( - expects[i], stream, ctx->br()->device_mr_ref() + expects[i], stream, ctx->br()->device_mr() ), stream ) diff --git a/cpp/tests/streaming/test_partition.cpp b/cpp/tests/streaming/test_partition.cpp index a3519b1df..b43dc62d2 100644 --- a/cpp/tests/streaming/test_partition.cpp +++ b/cpp/tests/streaming/test_partition.cpp @@ -43,9 +43,7 @@ TEST_F(StreamingPartition, PackUnpackRoundTrip) { inputs.emplace_back(to_message( i, std::make_unique( - std::make_unique( - expects[i], stream, ctx->br()->device_mr_ref() - ), + std::make_unique(expects[i], stream, ctx->br()->device_mr()), stream ) )); diff --git a/cpp/tests/streaming/test_read_parquet.cpp b/cpp/tests/streaming/test_read_parquet.cpp index dfb394736..1586bb7d1 100644 --- a/cpp/tests/streaming/test_read_parquet.cpp +++ b/cpp/tests/streaming/test_read_parquet.cpp @@ -261,7 +261,7 @@ TEST_P(StreamingReadParquetParams, ReadParquet) { ); chunk = chunk.make_available(reservation); auto packed_columns = - cudf::pack(chunk.table_view(), chunk.stream(), br->device_mr_ref()); + cudf::pack(chunk.table_view(), chunk.stream(), br->device_mr()); auto packed_data = PackedData{ std::move(packed_columns.metadata), br->move(std::move(packed_columns.gpu_data), chunk.stream()) diff --git a/cpp/tests/streaming/test_shuffler.cpp b/cpp/tests/streaming/test_shuffler.cpp index f00b37df2..66d899eb9 100644 --- a/cpp/tests/streaming/test_shuffler.cpp +++ b/cpp/tests/streaming/test_shuffler.cpp @@ -88,7 +88,7 @@ class StreamingShuffler : public BaseStreamingShuffle, ) .at(0), stream, - ctx->br()->device_mr_ref() + ctx->br()->device_mr() ), stream ) diff --git a/cpp/tests/test_allgather.cpp b/cpp/tests/test_allgather.cpp index bf72659ee..99472b499 100644 --- a/cpp/tests/test_allgather.cpp +++ b/cpp/tests/test_allgather.cpp @@ -295,7 +295,7 @@ TEST_F(BaseAllGatherTest, opid_reuse) { if (this_rank == 0) { // Recreate the buffer resource and allgather with the delayed MR. delay_br = rapidsmpf::BufferResource::create( - DelayedMemoryResource{br->device_mr_ref(), std::chrono::milliseconds(500)} + DelayedMemoryResource{br->device_mr(), std::chrono::milliseconds(500)} ); allgather = std::make_unique(GlobalEnvironment->comm_, op_id, delay_br.get()); diff --git a/cpp/tests/test_buffer.cpp b/cpp/tests/test_buffer.cpp index f5fa5a0c1..742339c6f 100644 --- a/cpp/tests/test_buffer.cpp +++ b/cpp/tests/test_buffer.cpp @@ -91,7 +91,7 @@ TEST_P(BufferRebindStreamTest, RebindStreamAndCopy) { ASSERT_NE(stream1.value(), stream2.value()); auto rmm_buffer = std::make_unique( - random_data.data(), buffer_size, stream1, br->device_mr_ref() + random_data.data(), buffer_size, stream1, br->device_mr() ); auto [reserve1, overbooking1] = diff --git a/cpp/tests/test_buffer_resource.cpp b/cpp/tests/test_buffer_resource.cpp index 29f2522a7..2533097a9 100644 --- a/cpp/tests/test_buffer_resource.cpp +++ b/cpp/tests/test_buffer_resource.cpp @@ -789,34 +789,83 @@ TEST_F(BufferCopyEdgeCases, SameBufferIsDisallowed) { ); } -TEST(BufferResource, DeviceMrOutlivesBufferResource) { +// `device_mr()` returns a `device_async_resource_ref` whose underlying +// any_resource holds an `OwningResourceAdaptor` carrying a +// `shared_ptr`. When the ref is consumed by an API that +// stores the MR (e.g. `rmm::device_buffer`), CCCL's +// `any_resource(resource_ref)` deep-copy captures the OwningResourceAdaptor +// — including its `shared_ptr` — so the storing object keeps +// BR (and its stream pool) alive for its entire lifetime. Conversely, once +// the storing object drops, BR must destruct (no refcount cycle). +TEST(BufferResource, DeviceMrKeepsBufferResourceAlive) { constexpr std::size_t N = 1024; auto br = BufferResource::create(cudf::get_current_device_resource_ref()); std::weak_ptr weak_br = br; auto stream = cudf::get_default_stream(); - auto owning_mr = br->device_mr(); - auto* wrapper = cuda::mr::resource_cast< - OwningResourceAdaptor>>( - &owning_mr - ); - ASSERT_NE(wrapper, nullptr); - - auto buf = std::make_unique(N, stream, std::move(owning_mr)); + // Construct a device_buffer through the BR. The buffer's internal + // any_resource deep-copies the OwningResourceAdaptor (with its + // shared_ptr
) out of the ref. + auto buf = std::make_unique(N, stream, br->device_mr()); - // Drop the original BR. The wrapper inside `buf`'s `any_resource` owns the only - // remaining `shared_ptr`, so BR (and its stream pool) must stay - // alive. + // Drop the original BR shared_ptr. The buffer's any_resource owns the + // only remaining `shared_ptr`, so BR (and its stream + // pool) must stay alive. br.reset(); EXPECT_FALSE(weak_br.expired()) << "BR freed while buffer still holds it"; - // Destroy the buffer: its stream-ordered deallocator runs against the still-alive - // stream pool. Must not crash. + // Destroy the buffer: its stream-ordered deallocator runs against the + // still-alive stream pool. Must not crash. EXPECT_NO_THROW(buf.reset()); - // After the buffer is gone, no shared_ptr
exists anywhere. If there were a - // refcount cycle (e.g. the wrapper's back-ref leaking back into the original BR), - // `weak_br` would still be observable here. + // After the buffer is gone, no shared_ptr
exists anywhere. If there + // were a refcount cycle (e.g. the wrapper's back-ref leaking back into + // the original BR), `weak_br` would still be observable here. EXPECT_TRUE(weak_br.expired()) << "BR not destructed, refcount cycle?"; } + +// `OwningResourceAdaptor`'s copy ctor promotes its `weak_ptr` to a +// `shared_ptr`. If the weak ptr has expired (the back-ref object is +// already destroyed), the promotion throws `std::bad_weak_ptr`. This guards +// against silent UAF in the deep-copy path. +TEST(OwningResourceAdaptor, CopyThrowsWhenBackRefExpired) { + rmm::mr::cuda_memory_resource cuda_mr; + RmmResourceAdaptor adaptor{ + cuda::mr::any_resource{cuda_mr} + }; + + std::weak_ptr weak_br; + { + auto br = BufferResource::create(cudf::get_current_device_resource_ref()); + weak_br = br; + } + ASSERT_TRUE(weak_br.expired()); + + // Manually construct an adaptor with an already-expired weak ref. The + // adaptor itself is fine to hold; copying it triggers the promotion + // attempt, which must throw. + OwningResourceAdaptor original{adaptor, weak_br}; + + using AdaptorT = OwningResourceAdaptor; + EXPECT_THROW({ AdaptorT copy{original}; }, std::bad_weak_ptr); + + AdaptorT target{adaptor, weak_br}; + EXPECT_THROW(target = original, std::bad_weak_ptr); +} + +// The original `OwningResourceAdaptor` stored as BR's `owning_mr_` member +// must hold only a `weak_ptr` (no strong ref). Otherwise the +// member would close a refcount cycle BR -> owning_mr_ -> shared_ptr
-> BR +// and BR would leak. Verified by checking that destroying the only +// `shared_ptr` to a freshly created BR actually destroys BR. +TEST(BufferResource, OwningAdaptorMemberHasNoStrongCycle) { + std::weak_ptr weak_br; + { + auto br = BufferResource::create(cudf::get_current_device_resource_ref()); + weak_br = br; + EXPECT_FALSE(weak_br.expired()); + } + EXPECT_TRUE(weak_br.expired()) << "BR leaked - the owning_mr_ member holds a " + "strong ref to BR, closing a cycle."; +} diff --git a/cpp/tests/utils.hpp b/cpp/tests/utils.hpp index a55507806..0aa23ff85 100644 --- a/cpp/tests/utils.hpp +++ b/cpp/tests/utils.hpp @@ -204,7 +204,7 @@ template std::memcpy(metadata->data(), values.data(), n_elements * sizeof(int)); auto data = std::make_unique( - values.data(), n_elements * sizeof(int), stream, br.device_mr_ref() + values.data(), n_elements * sizeof(int), stream, br.device_mr() ); return {std::move(metadata), br.move(std::move(data), stream)}; From 03c5c96dfc4b6ddb86ac6a3bce7a6b1cbfec068c Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Sun, 31 May 2026 22:10:16 +0200 Subject: [PATCH 07/15] docs --- cpp/src/memory/buffer_resource.cpp | 16 +++++---- cpp/tests/test_buffer_resource.cpp | 54 +++++++++++++----------------- 2 files changed, 33 insertions(+), 37 deletions(-) diff --git a/cpp/src/memory/buffer_resource.cpp b/cpp/src/memory/buffer_resource.cpp index bd6db16ce..8313ac2bc 100644 --- a/cpp/src/memory/buffer_resource.cpp +++ b/cpp/src/memory/buffer_resource.cpp @@ -66,12 +66,16 @@ std::shared_ptr BufferResource::create( std::move(stream_pool), std::move(statistics) }}; - // Install the owning adaptor *after* construction so that - // `weak_from_this()` is valid. The adaptor holds only a `weak_ptr` to - // `*br` so no refcount cycle. Copies of the adaptor (made by CCCL's - // `any_resource(resource_ref)` deep-copy in downstream consumers) promote - // that weak ref to a `shared_ptr`, transferring ownership - // to the copy. + + // Install the owning adaptor *after* construction so that `weak_from_this()` is + // valid. + // + // The adaptor stored inside `BufferResource` itself holds only a `weak_ptr`, avoiding + // a reference cycle. When downstream code promotes `device_mr()` from a non-owning + // `resource_ref` to an owning `cuda::mr::any_resource` (for example internally in + // `rmm::device_buffer`), CCCL deep-copies the adaptor. That copy promotes the weak + // ref to a `shared_ptr`, ensuring the `BufferResource` stays alive + // for as long as the owning resource exists. br->owning_mr_.emplace(br->device_adaptor_, br->weak_from_this()); return br; } diff --git a/cpp/tests/test_buffer_resource.cpp b/cpp/tests/test_buffer_resource.cpp index 2533097a9..ffa1cf1a6 100644 --- a/cpp/tests/test_buffer_resource.cpp +++ b/cpp/tests/test_buffer_resource.cpp @@ -789,14 +789,6 @@ TEST_F(BufferCopyEdgeCases, SameBufferIsDisallowed) { ); } -// `device_mr()` returns a `device_async_resource_ref` whose underlying -// any_resource holds an `OwningResourceAdaptor` carrying a -// `shared_ptr`. When the ref is consumed by an API that -// stores the MR (e.g. `rmm::device_buffer`), CCCL's -// `any_resource(resource_ref)` deep-copy captures the OwningResourceAdaptor -// — including its `shared_ptr` — so the storing object keeps -// BR (and its stream pool) alive for its entire lifetime. Conversely, once -// the storing object drops, BR must destruct (no refcount cycle). TEST(BufferResource, DeviceMrKeepsBufferResourceAlive) { constexpr std::size_t N = 1024; @@ -804,31 +796,29 @@ TEST(BufferResource, DeviceMrKeepsBufferResourceAlive) { std::weak_ptr weak_br = br; auto stream = cudf::get_default_stream(); - // Construct a device_buffer through the BR. The buffer's internal - // any_resource deep-copies the OwningResourceAdaptor (with its - // shared_ptr
) out of the ref. + // Construct a device_buffer using the BR memory resource. Internally, + // `rmm::device_buffer` stores the resource as an owning `cuda::mr::any_resource`, + // which deep-copies the OwningResourceAdaptor. The copied adaptor acquires a + // `shared_ptr`. auto buf = std::make_unique(N, stream, br->device_mr()); - // Drop the original BR shared_ptr. The buffer's any_resource owns the - // only remaining `shared_ptr`, so BR (and its stream - // pool) must stay alive. + // Drop the original shared_ptr. The buffer's internally stored owning memory + // resource now holds the only remaining `shared_ptr`, so the + // BR must remain alive. br.reset(); EXPECT_FALSE(weak_br.expired()) << "BR freed while buffer still holds it"; - // Destroy the buffer: its stream-ordered deallocator runs against the - // still-alive stream pool. Must not crash. + // Destroy the buffer. This releases the owning memory resource and the final + // `shared_ptr`. The stream-ordered deallocation must complete + // successfully while the BR is still alive. EXPECT_NO_THROW(buf.reset()); - // After the buffer is gone, no shared_ptr
exists anywhere. If there - // were a refcount cycle (e.g. the wrapper's back-ref leaking back into - // the original BR), `weak_br` would still be observable here. + // After the buffer is destroyed, no shared ownership should remain. If the + // original adaptor held a strong back-reference, or another ownership cycle + // existed, the BR would still be alive here. EXPECT_TRUE(weak_br.expired()) << "BR not destructed, refcount cycle?"; } -// `OwningResourceAdaptor`'s copy ctor promotes its `weak_ptr` to a -// `shared_ptr`. If the weak ptr has expired (the back-ref object is -// already destroyed), the promotion throws `std::bad_weak_ptr`. This guards -// against silent UAF in the deep-copy path. TEST(OwningResourceAdaptor, CopyThrowsWhenBackRefExpired) { rmm::mr::cuda_memory_resource cuda_mr; RmmResourceAdaptor adaptor{ @@ -842,23 +832,25 @@ TEST(OwningResourceAdaptor, CopyThrowsWhenBackRefExpired) { } ASSERT_TRUE(weak_br.expired()); - // Manually construct an adaptor with an already-expired weak ref. The - // adaptor itself is fine to hold; copying it triggers the promotion - // attempt, which must throw. + // Construct an adaptor with an expired weak back-reference. The adaptor + // itself is valid to create and hold because it stores only the weak + // ref. The failure occurs only when a copy attempts to promote that weak + // ref to shared ownership. OwningResourceAdaptor original{adaptor, weak_br}; using AdaptorT = OwningResourceAdaptor; + + // Copy construction promotes the weak back-reference to a + // `shared_ptr`. Since the BR has already been destroyed, + // the promotion must fail with `std::bad_weak_ptr`. EXPECT_THROW({ AdaptorT copy{original}; }, std::bad_weak_ptr); + // Copy assignment uses the same promotion semantics and must likewise + // fail when the back-reference has expired. AdaptorT target{adaptor, weak_br}; EXPECT_THROW(target = original, std::bad_weak_ptr); } -// The original `OwningResourceAdaptor` stored as BR's `owning_mr_` member -// must hold only a `weak_ptr` (no strong ref). Otherwise the -// member would close a refcount cycle BR -> owning_mr_ -> shared_ptr
-> BR -// and BR would leak. Verified by checking that destroying the only -// `shared_ptr` to a freshly created BR actually destroys BR. TEST(BufferResource, OwningAdaptorMemberHasNoStrongCycle) { std::weak_ptr weak_br; { From f2663dd478d860a75df7344321cb1198292a9c70 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Mon, 1 Jun 2026 09:04:44 +0200 Subject: [PATCH 08/15] docs --- .../rapidsmpf/memory/buffer_resource.hpp | 23 ++++++++++++++- .../rapidsmpf/memory/buffer_resource.pyx | 28 ++++++++++++++++--- 2 files changed, 46 insertions(+), 5 deletions(-) diff --git a/cpp/include/rapidsmpf/memory/buffer_resource.hpp b/cpp/include/rapidsmpf/memory/buffer_resource.hpp index 7924d0131..7103bb968 100644 --- a/cpp/include/rapidsmpf/memory/buffer_resource.hpp +++ b/cpp/include/rapidsmpf/memory/buffer_resource.hpp @@ -55,6 +55,24 @@ enum class AllowOverbooking : bool { * @note `BufferResource` instances must be constructed through `create()` or * `from_options()`, both of which return a `std::shared_ptr`. 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()`. + * + * 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. */ class BufferResource : public std::enable_shared_from_this { public: @@ -70,7 +88,10 @@ class BufferResource : public std::enable_shared_from_this { * If pinned-host memory is disabled, available pinned-host memory is always reported * as zero regardless of the configured limit. * - * @param device_mr Device memory resource used for device allocations. + * @param 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. * @param pinned_mr Pinned host memory resource used for * `MemoryType::PINNED_HOST` allocations, or `PinnedMemoryResource::Disabled` to * disable pinned allocations. diff --git a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx index 1fcac4c71..f0b26d992 100644 --- a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx +++ b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx @@ -114,10 +114,10 @@ cdef class BufferResource: Parameters ---------- device_mr - The RMM device memory resource used for device allocations. The - BufferResource transparently wraps this resource in an internal RMM - adaptor for allocation tracking — callers don't need to wrap it - themselves. + The RMM 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 ``device_mr`` after + construction. pinned_mr The pinned host memory resource used for :attr:`~.MemoryType.PINNED_HOST` allocations. If None, pinned host allocations are disabled. In that case, @@ -138,6 +138,26 @@ cdef class BufferResource: statistics The statistics instance to use. If None, a disabled statistics instance will be created. + + Notes + ----- + Allocation tracking only applies to allocations routed through this + ``BufferResource``. The constructor wraps the supplied ``device_mr`` in + an internal RMM adaptor that records all allocations and deallocations; + that adaptor is exposed via ``BufferResource.device_mr``. + + Allocations made through the original, unwrapped ``device_mr`` 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 original + ``device_mr`` passed to the constructor. + + Tracking allocations made outside ``BufferResource``, for example + allocations performed before construction or through code paths that use + a raw RMM memory resource directly, is a separate design concern and is + not handled by this class. """ def __cinit__( self, From 3d067ebde13d2e2ee8db87e36deca838abe55dab Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Mon, 1 Jun 2026 09:46:17 +0200 Subject: [PATCH 09/15] based on reviews --- .../rapidsmpf/memory/buffer_resource.hpp | 11 +++----- .../memory/owning_resource_adaptor.hpp | 24 +++++++++++------ cpp/src/memory/buffer_resource.cpp | 9 +++---- cpp/tests/test_buffer_resource.cpp | 27 +------------------ 4 files changed, 25 insertions(+), 46 deletions(-) diff --git a/cpp/include/rapidsmpf/memory/buffer_resource.hpp b/cpp/include/rapidsmpf/memory/buffer_resource.hpp index 7103bb968..9b1c3136f 100644 --- a/cpp/include/rapidsmpf/memory/buffer_resource.hpp +++ b/cpp/include/rapidsmpf/memory/buffer_resource.hpp @@ -203,7 +203,7 @@ class BufferResource : public std::enable_shared_from_this { * br.reset(); // safe: `col` keeps the BufferResource alive internally * @endcode */ - [[nodiscard]] rmm::device_async_resource_ref device_mr(); + [[nodiscard]] rmm::device_async_resource_ref device_mr() noexcept; /** * @brief Get the RMM host memory resource. @@ -520,12 +520,9 @@ class BufferResource : public std::enable_shared_from_this { // The internal RmmResourceAdaptor wraps the user's device MR so that // allocations are tracked for the DEVICE memory_available calculation. RmmResourceAdaptor device_adaptor_; - // Stable storage for the device MR returned by `device_mr()`. Holds a - // `weak_ptr` to *this; copies promote that to a - // `shared_ptr` on the deep-copy path. Installed by - // `create()` after construction so that `weak_from_this()` is valid; - // wrapped in `optional` so it can be default-constructed empty - // (`RmmResourceAdaptor` itself is not default-constructible). + // Stable storage for the device MR returned by `device_mr()`. See + // `OwningResourceAdaptor` for the lifetime semantics. Installed by + // `create()` after construction. std::optional> owning_mr_; std::optional pinned_mr_; HostMemoryResource host_mr_; diff --git a/cpp/include/rapidsmpf/memory/owning_resource_adaptor.hpp b/cpp/include/rapidsmpf/memory/owning_resource_adaptor.hpp index d51df15e8..b71fc192d 100644 --- a/cpp/include/rapidsmpf/memory/owning_resource_adaptor.hpp +++ b/cpp/include/rapidsmpf/memory/owning_resource_adaptor.hpp @@ -64,10 +64,15 @@ namespace rapidsmpf { * `std::shared_ptr`. The copied adaptor therefore keeps the owner * alive for as long as the copied resource exists. * - * If the owner has already been destroyed when the adaptor is copied, the copy - * constructor throws `std::bad_weak_ptr` instead of producing a dangling + * If the owner has already been destroyed when the adaptor is copied, the + * copy constructor throws `std::bad_weak_ptr` instead of creating a dangling * adaptor. * + * This only works while the source adaptor itself is still alive, which is + * the common case for transient ref/`any_resource` consumers where CCCL + * performs the deep copy during construction. If both the owner and source + * adaptor have already been destroyed, the behavior is undefined. + * * This does not make direct use of a bare `resource_ref` safe after the owner * has been destroyed. A call such as `ref.allocate(...)` still dispatches * through the `resource_ref` into the original resource object, so the owner @@ -232,21 +237,24 @@ class OwningResourceAdaptor } /** - * @brief Compare wrapped resource identity. + * @brief Compare adaptor identity. * - * Two adaptors compare equal if their wrapped resources compare equal. + * Two adaptors compare equal if both their wrapped resources and their + * back-references compare equal. * * @param other Adaptor to compare against. - * @return `true` if both adaptors refer to the same underlying resource. + * @return `true` if both adaptors refer to the same wrapped resource and + * the same back-referenced object. */ [[nodiscard]] bool operator==(OwningResourceAdaptor const& other) const noexcept { - return resource_ == other.resource_; + return resource_ == other.resource_ && !weak_.owner_before(other.weak_) + && !other.weak_.owner_before(weak_); } private: Resource resource_; - std::weak_ptr weak_; - std::shared_ptr strong_; + std::weak_ptr weak_{nullptr}; + std::shared_ptr strong_{nullptr}; }; } // namespace rapidsmpf diff --git a/cpp/src/memory/buffer_resource.cpp b/cpp/src/memory/buffer_resource.cpp index 8313ac2bc..4ce779505 100644 --- a/cpp/src/memory/buffer_resource.cpp +++ b/cpp/src/memory/buffer_resource.cpp @@ -57,7 +57,6 @@ std::shared_ptr BufferResource::create( std::shared_ptr stream_pool, std::shared_ptr statistics ) { - // Construct via `new` (private ctor visible here as a member function). std::shared_ptr br{new BufferResource{ std::move(device_mr), std::move(pinned_mr), @@ -73,9 +72,9 @@ std::shared_ptr BufferResource::create( // The adaptor stored inside `BufferResource` itself holds only a `weak_ptr`, avoiding // a reference cycle. When downstream code promotes `device_mr()` from a non-owning // `resource_ref` to an owning `cuda::mr::any_resource` (for example internally in - // `rmm::device_buffer`), CCCL deep-copies the adaptor. That copy promotes the weak - // ref to a `shared_ptr`, ensuring the `BufferResource` stays alive - // for as long as the owning resource exists. + // `rmm::device_buffer`), CCCL deep-copies the adaptor. The copy promotes the weak + // ref to a `shared_ptr`, acquiring shared ownership so the + // `BufferResource` stays alive for as long as the owning resource exists. br->owning_mr_.emplace(br->device_adaptor_, br->weak_from_this()); return br; } @@ -124,7 +123,7 @@ void BufferResource::set_memory_limit(MemoryType mem_type, std::int64_t limit) n ); } -rmm::device_async_resource_ref BufferResource::device_mr() { +rmm::device_async_resource_ref BufferResource::device_mr() noexcept { return rmm::device_async_resource_ref{*owning_mr_}; } diff --git a/cpp/tests/test_buffer_resource.cpp b/cpp/tests/test_buffer_resource.cpp index ffa1cf1a6..95cb74edc 100644 --- a/cpp/tests/test_buffer_resource.cpp +++ b/cpp/tests/test_buffer_resource.cpp @@ -803,14 +803,10 @@ TEST(BufferResource, DeviceMrKeepsBufferResourceAlive) { auto buf = std::make_unique(N, stream, br->device_mr()); // Drop the original shared_ptr. The buffer's internally stored owning memory - // resource now holds the only remaining `shared_ptr`, so the - // BR must remain alive. + // resource should be keeping the BR alive. br.reset(); EXPECT_FALSE(weak_br.expired()) << "BR freed while buffer still holds it"; - // Destroy the buffer. This releases the owning memory resource and the final - // `shared_ptr`. The stream-ordered deallocation must complete - // successfully while the BR is still alive. EXPECT_NO_THROW(buf.reset()); // After the buffer is destroyed, no shared ownership should remain. If the @@ -832,32 +828,11 @@ TEST(OwningResourceAdaptor, CopyThrowsWhenBackRefExpired) { } ASSERT_TRUE(weak_br.expired()); - // Construct an adaptor with an expired weak back-reference. The adaptor - // itself is valid to create and hold because it stores only the weak - // ref. The failure occurs only when a copy attempts to promote that weak - // ref to shared ownership. OwningResourceAdaptor original{adaptor, weak_br}; using AdaptorT = OwningResourceAdaptor; - - // Copy construction promotes the weak back-reference to a - // `shared_ptr`. Since the BR has already been destroyed, - // the promotion must fail with `std::bad_weak_ptr`. EXPECT_THROW({ AdaptorT copy{original}; }, std::bad_weak_ptr); - // Copy assignment uses the same promotion semantics and must likewise - // fail when the back-reference has expired. AdaptorT target{adaptor, weak_br}; EXPECT_THROW(target = original, std::bad_weak_ptr); } - -TEST(BufferResource, OwningAdaptorMemberHasNoStrongCycle) { - std::weak_ptr weak_br; - { - auto br = BufferResource::create(cudf::get_current_device_resource_ref()); - weak_br = br; - EXPECT_FALSE(weak_br.expired()); - } - EXPECT_TRUE(weak_br.expired()) << "BR leaked - the owning_mr_ member holds a " - "strong ref to BR, closing a cycle."; -} From 2cc12fd54ef7a1adff2e46bba29bde1b4d48f45b Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Mon, 1 Jun 2026 14:42:46 +0200 Subject: [PATCH 10/15] update docstring examples --- cpp/include/rapidsmpf/memory/buffer_resource.hpp | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/cpp/include/rapidsmpf/memory/buffer_resource.hpp b/cpp/include/rapidsmpf/memory/buffer_resource.hpp index 9b1c3136f..84f71b5f7 100644 --- a/cpp/include/rapidsmpf/memory/buffer_resource.hpp +++ b/cpp/include/rapidsmpf/memory/buffer_resource.hpp @@ -167,7 +167,7 @@ class BufferResource : public std::enable_shared_from_this { * When working directly with the returned reference, the caller must ensure that this * `BufferResource` remains alive for the full duration of that use: * @code - * auto br = std::make_shared(...); + * 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 @@ -176,7 +176,7 @@ class BufferResource : public std::enable_shared_from_this { * To store the resource beyond the immediate call, promote the ref to an * owning `cuda::mr::any_resource`: * @code - * auto br = std::make_shared(...); + * auto br = BufferResource::create(...); * cuda::mr::any_resource mr = br->device_mr(); * br.reset(); // safe: `mr` keeps the BufferResource alive * mr.allocate(...); // safe @@ -185,14 +185,14 @@ class BufferResource : public std::enable_shared_from_this { * In the common case, no explicit promotion is needed because RMM and cuDF containers * that store a memory resource do this internally: * @code - * auto br = std::make_shared(...); + * auto br = BufferResource::create(...); * rmm::device_buffer buf{1024, stream, br->device_mr()}; * br.reset(); // safe: `buf` keeps the BufferResource alive internally * @endcode * * Returned objects from cuDF APIs typically behave the same way: * @code - * auto br = std::make_shared(...); + * auto br = BufferResource::create(...); * auto col = cudf::make_numeric_column( * cudf::data_type{cudf::type_id::INT32}, * 1000, From 19f50ef15eaf78c9c460a6bcb2882a184deff40d Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Mon, 1 Jun 2026 15:14:28 +0200 Subject: [PATCH 11/15] OwningDeviceMemoryResource --- .../rapidsmpf/memory/buffer_resource.pxd | 13 ++- .../rapidsmpf/memory/buffer_resource.pyi | 4 +- .../rapidsmpf/memory/buffer_resource.pyx | 89 ++++++++++++++----- .../rapidsmpf/tests/test_buffer_resource.py | 40 ++++++++- 4 files changed, 123 insertions(+), 23 deletions(-) diff --git a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pxd b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pxd index 7ca4a06ed..0dcb45c4a 100644 --- a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pxd +++ b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pxd @@ -8,7 +8,8 @@ from libcpp.memory cimport shared_ptr from libcpp.optional cimport optional from libcpp.unordered_map cimport unordered_map from rmm.librmm.cuda_stream_pool cimport cuda_stream_pool -from rmm.librmm.memory_resource cimport any_resource, device_accessible +from rmm.librmm.memory_resource cimport (any_resource, device_accessible, + device_async_resource_ref) from rmm.pylibrmm.cuda_stream_pool cimport CudaStreamPool from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource @@ -46,6 +47,7 @@ cdef extern from "" nogil: const cuda_stream_pool &stream_pool() except +ex_handler size_t release(cpp_MemoryReservation&, size_t) except +ex_handler shared_ptr[cpp_Statistics] statistics() except +ex_handler + device_async_resource_ref device_mr() noexcept cdef class BufferResource: cdef object __weakref__ @@ -57,3 +59,12 @@ cdef class BufferResource: cdef CudaStreamPool _stream_pool cdef Statistics _statistics cdef const cuda_stream_pool* stream_pool(self) + + +cdef class OwningDeviceMemoryResource(DeviceMemoryResource): + cdef any_resource[device_accessible] c_obj + + @staticmethod + cdef OwningDeviceMemoryResource _create( + any_resource[device_accessible] resource, + ) diff --git a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyi b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyi index ce21b941f..5f7ded0fa 100644 --- a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyi +++ b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyi @@ -33,7 +33,7 @@ class BufferResource: statistics: Statistics | None = None, ) -> Self: ... @property - def device_mr(self) -> DeviceMemoryResource: ... + def device_mr(self) -> OwningDeviceMemoryResource: ... @property def pinned_mr(self) -> PinnedMemoryResource | None: ... def memory_available(self, mem_type: MemoryType) -> int: ... @@ -55,6 +55,8 @@ class BufferResource: def release(self, reservation: MemoryReservation, size: int) -> int: ... def stream_pool_size(self) -> int: ... +class OwningDeviceMemoryResource(DeviceMemoryResource): ... + def device_limit_from_options(options: Options) -> int: ... def periodic_spill_check_from_options(options: Options) -> float | None: ... def stream_pool_from_options(options: Options) -> CudaStreamPool: ... diff --git a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx index f0b26d992..cf0b48ce8 100644 --- a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx +++ b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx @@ -15,10 +15,34 @@ from rmm.librmm.cuda_stream_pool cimport cuda_stream_pool from rmm.pylibrmm import CudaStreamFlags -from rmm.librmm.memory_resource cimport make_any_device_resource +from rmm.librmm.memory_resource cimport (any_resource, device_accessible, + device_async_resource_ref, + make_any_device_resource) from rmm.pylibrmm.cuda_stream_pool cimport CudaStreamPool from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource + +cdef extern from *: + """ + // Construct a device_async_resource_ref from an owning any_resource. + // The free template in RMM only declares overloads for concrete RMM + // types, this overload covers `cuda::mr::any_resource`. + std::optional + cpp_make_device_async_resource_ref_from_any( + cuda::mr::any_resource& mr + ) { + + // `cython_device_async_resource_ref` is declared inline in RMM's + // `librmm/memory_resource.pxd`. + return std::optional( + rmm::device_async_resource_ref(mr) + ); + } + """ + optional[device_async_resource_ref] cpp_make_device_async_resource_ref_from_any( + any_resource[device_accessible]& + ) except +ex_handler + from rapidsmpf._detail.exception_handling cimport ex_handler from rapidsmpf.memory.memory_reservation cimport MemoryReservation from rapidsmpf.memory.pinned_memory_resource cimport (PinnedMemoryResource, @@ -142,22 +166,15 @@ cdef class BufferResource: Notes ----- Allocation tracking only applies to allocations routed through this - ``BufferResource``. The constructor wraps the supplied ``device_mr`` in - an internal RMM adaptor that records all allocations and deallocations; - that adaptor is exposed via ``BufferResource.device_mr``. - - Allocations made through the original, unwrapped ``device_mr`` 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 original - ``device_mr`` passed to the constructor. - - Tracking allocations made outside ``BufferResource``, for example - allocations performed before construction or through code paths that use - a raw RMM memory resource directly, is a separate design concern and is - not handled by this class. + ``BufferResource``. + + To ensure allocations are included in memory-limit accounting and + statistics, use ``BufferResource.device_mr`` for all CUDA allocations + associated with this resource. + + Allocations performed through other memory resources, including the + original resource passed to the constructor or allocations made outside + ``BufferResource``, are not tracked by this class. """ def __cinit__( self, @@ -294,13 +311,18 @@ cdef class BufferResource: @property def device_mr(self): """ - The memory resource used for device memory allocations. + The tracked device memory resource. + + Allocations made through this resource count against the ``BufferResource`` + memory limits and appear in its statistics. Returns ------- - The device memory resource. + The tracked device memory resource. """ - return self._device_mr + return OwningDeviceMemoryResource._create( + make_any_device_resource(deref(self._handle).device_mr()) + ) @property def pinned_mr(self): @@ -606,3 +628,30 @@ def stream_pool_from_options(Options options not None): pool_size=pool_size, flags=CudaStreamFlags.NON_BLOCKING, ) + + +cdef class OwningDeviceMemoryResource(DeviceMemoryResource): + """ + Owning :class:`DeviceMemoryResource`. + + Useful for exposing device memory resources to Python in a form that is + compatible with cuDF/RMM APIs while preserving ownership semantics. + + Notes + ----- + RMM does not currently provide an equivalent owning wrapper. If one is + added in the future, this class can likely be replaced by the + RMM-provided implementation. + """ + @staticmethod + cdef OwningDeviceMemoryResource _create( + any_resource[device_accessible] resource, + ): + cdef OwningDeviceMemoryResource self = ( + OwningDeviceMemoryResource.__new__(OwningDeviceMemoryResource) + ) + # Owning storage for the underlying CCCL resource. The base class's + # `c_ref` is a non-owning view into `c_obj`. + self.c_obj = resource + self.c_ref = cpp_make_device_async_resource_ref_from_any(self.c_obj) + return self diff --git a/python/rapidsmpf/rapidsmpf/tests/test_buffer_resource.py b/python/rapidsmpf/rapidsmpf/tests/test_buffer_resource.py index 15801e171..95cabae14 100644 --- a/python/rapidsmpf/rapidsmpf/tests/test_buffer_resource.py +++ b/python/rapidsmpf/rapidsmpf/tests/test_buffer_resource.py @@ -2,13 +2,15 @@ # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations +import numpy as np import pytest +import rmm import rmm.mr from rapidsmpf.error import ReservationError from rapidsmpf.memory.buffer import MemoryType -from rapidsmpf.memory.buffer_resource import BufferResource +from rapidsmpf.memory.buffer_resource import BufferResource, OwningDeviceMemoryResource from rapidsmpf.memory.memory_reservation import opaque_memory_usage from rapidsmpf.statistics import Statistics @@ -193,6 +195,42 @@ def test_opaque_memory_usage_partial_consumption(mem_type: MemoryType) -> None: assert res.size == 0 +def test_device_mr_property_tracks_allocations() -> None: + mr = rmm.mr.CudaMemoryResource() + br = BufferResource(mr, memory_limits={MemoryType.DEVICE: KiB(1024)}) + assert isinstance(br.device_mr, OwningDeviceMemoryResource) + + avail_before = br.memory_available(MemoryType.DEVICE) + + # Allocation through the tracked view is visible to the BR. + buf_tracked = rmm.DeviceBuffer(size=KiB(64), mr=br.device_mr) + avail_after_tracked = br.memory_available(MemoryType.DEVICE) + assert avail_after_tracked == avail_before - KiB(64) + + # Allocation through the original (unwrapped) MR bypasses tracking. + buf_untracked = rmm.DeviceBuffer(size=KiB(64), mr=mr) + avail_after_untracked = br.memory_available(MemoryType.DEVICE) + assert avail_after_untracked == avail_after_tracked + + del buf_tracked, buf_untracked + + +def test_device_mr_outlives_buffer_resource() -> None: + expected = np.arange(KiB(64), dtype=int) + + # Write data and delete the buffer resource. + mr = rmm.mr.CudaMemoryResource() + br = BufferResource(mr) + buf = rmm.DeviceBuffer(size=expected.nbytes, mr=br.device_mr) + buf.copy_from_host(expected.tobytes()) + del br + + # Read the data back and verify it survived. + got = np.frombuffer(buf.copy_to_host(), dtype=int) + assert np.array_equal(got, expected) + del buf + + def test_reserve_or_fail() -> None: mr = rmm.mr.CudaMemoryResource() br = BufferResource(mr, memory_limits={MemoryType.DEVICE: KiB(100)}) From b0f7cc6a08e7662d727056510ee2a2227add17b5 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Mon, 1 Jun 2026 15:26:28 +0200 Subject: [PATCH 12/15] cleanup --- python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx | 1 - 1 file changed, 1 deletion(-) diff --git a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx index cf0b48ce8..1981c7709 100644 --- a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx +++ b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx @@ -31,7 +31,6 @@ cdef extern from *: cpp_make_device_async_resource_ref_from_any( cuda::mr::any_resource& mr ) { - // `cython_device_async_resource_ref` is declared inline in RMM's // `librmm/memory_resource.pxd`. return std::optional( From eb5f3bbce71b28deae3a3de1f7f25ee08eb516a0 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Mon, 1 Jun 2026 15:48:49 +0200 Subject: [PATCH 13/15] fix docs --- docs/source/conf.py | 2 ++ python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/docs/source/conf.py b/docs/source/conf.py index 57b177d28..9b0cb9345 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -145,4 +145,6 @@ def setup(app): ("py:obj", "rapidsmpf.rmm_resource_adaptor.RmmResourceAdaptor.allocate"), ("py:obj", "rapidsmpf.rmm_resource_adaptor.RmmResourceAdaptor.deallocate"), ("py:obj", "rapidsmpf.rmm_resource_adaptor.RmmResourceAdaptor.get_upstream"), + ("py:obj", "rapidsmpf.memory.buffer_resource.OwningDeviceMemoryResource.allocate"), + ("py:obj", "rapidsmpf.memory.buffer_resource.OwningDeviceMemoryResource.deallocate"), ] diff --git a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx index 1981c7709..bf00608e4 100644 --- a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx +++ b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx @@ -631,7 +631,7 @@ def stream_pool_from_options(Options options not None): cdef class OwningDeviceMemoryResource(DeviceMemoryResource): """ - Owning :class:`DeviceMemoryResource`. + Owning ``DeviceMemoryResource``. Useful for exposing device memory resources to Python in a form that is compatible with cuDF/RMM APIs while preserving ownership semantics. From c02c32fd21b5b614d3279f3f03ea03e29b1eecd4 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Mon, 1 Jun 2026 16:17:27 +0200 Subject: [PATCH 14/15] docs --- python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx index bf00608e4..cdaf9f8cb 100644 --- a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx +++ b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx @@ -225,8 +225,14 @@ cdef class BufferResource: # checked cast requires the GIL stats_handle = (statistics)._handle - # Stored for the Python device_mr/pinned_mr property accessors. - # The C++ BufferResource owns the resource via any_resource. + # Keep the original Python memory resources alive while the C++ + # BufferResource holds resources derived from them. These anchors are + # likely redundant: `make_any_device_resource(...)` deep-copies the + # device MR into a self-sufficient owning any_resource, and + # `cpp_PinnedMemoryResource` is copied by value into the C++ BR. Kept + # defensively for now. + # TODO: drop these once verified against pool/upstream-adaptor MRs. + # https://github.com/rapidsai/rapidsmpf/issues/1074 self._device_mr = device_mr self._pinned_mr = pinned_mr cdef optional[cpp_PinnedMemoryResource] cpp_pinned_mr From d259e94c72e9c097a2891ca1e064192d5f80b5d4 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Mon, 1 Jun 2026 17:36:24 +0200 Subject: [PATCH 15/15] remove the device_adaptor_ member --- cpp/include/rapidsmpf/memory/buffer_resource.hpp | 4 ---- cpp/src/memory/buffer_resource.cpp | 11 +++++------ 2 files changed, 5 insertions(+), 10 deletions(-) diff --git a/cpp/include/rapidsmpf/memory/buffer_resource.hpp b/cpp/include/rapidsmpf/memory/buffer_resource.hpp index 84f71b5f7..d477f903a 100644 --- a/cpp/include/rapidsmpf/memory/buffer_resource.hpp +++ b/cpp/include/rapidsmpf/memory/buffer_resource.hpp @@ -508,7 +508,6 @@ class BufferResource : public std::enable_shared_from_this { private: /** @brief Private constructor, use `create()` or `from_options()`. */ BufferResource( - cuda::mr::any_resource device_mr, std::optional pinned_mr, std::unordered_map memory_limits, std::optional periodic_spill_check, @@ -517,9 +516,6 @@ class BufferResource : public std::enable_shared_from_this { ); std::mutex mutex_; - // The internal RmmResourceAdaptor wraps the user's device MR so that - // allocations are tracked for the DEVICE memory_available calculation. - RmmResourceAdaptor device_adaptor_; // Stable storage for the device MR returned by `device_mr()`. See // `OwningResourceAdaptor` for the lifetime semantics. Installed by // `create()` after construction. diff --git a/cpp/src/memory/buffer_resource.cpp b/cpp/src/memory/buffer_resource.cpp index 4ce779505..0fb385a5f 100644 --- a/cpp/src/memory/buffer_resource.cpp +++ b/cpp/src/memory/buffer_resource.cpp @@ -23,15 +23,13 @@ namespace rapidsmpf { BufferResource::BufferResource( - cuda::mr::any_resource device_mr, std::optional pinned_mr, std::unordered_map memory_limits, std::optional periodic_spill_check, std::shared_ptr stream_pool, std::shared_ptr statistics ) - : device_adaptor_{std::move(device_mr)}, - pinned_mr_{std::move(pinned_mr)}, + : pinned_mr_{std::move(pinned_mr)}, host_mr_{}, stream_pool_{std::move(stream_pool)}, spill_manager_{this, periodic_spill_check}, @@ -58,7 +56,6 @@ std::shared_ptr BufferResource::create( std::shared_ptr statistics ) { std::shared_ptr br{new BufferResource{ - std::move(device_mr), std::move(pinned_mr), std::move(memory_limits), periodic_spill_check, @@ -75,7 +72,9 @@ std::shared_ptr BufferResource::create( // `rmm::device_buffer`), CCCL deep-copies the adaptor. The copy promotes the weak // ref to a `shared_ptr`, acquiring shared ownership so the // `BufferResource` stays alive for as long as the owning resource exists. - br->owning_mr_.emplace(br->device_adaptor_, br->weak_from_this()); + br->owning_mr_.emplace( + RmmResourceAdaptor{std::move(device_mr)}, br->weak_from_this() + ); return br; } @@ -104,7 +103,7 @@ std::int64_t BufferResource::memory_available(MemoryType mem_type) const noexcep ); switch (mem_type) { case MemoryType::DEVICE: - return limit - device_adaptor_.current_allocated(); + return limit - owning_mr_->get().current_allocated(); case MemoryType::PINNED_HOST: if (pinned_mr_ == PinnedMemoryResource::Disabled) { return 0;