From 801e7f901fb0806701ffa7d58a43e92922ac5bd0 Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Fri, 18 Sep 2026 17:24:51 -0700 Subject: [PATCH 1/2] Replace libcudf cuda stream owners --- cpp/benchmarks/io/fst.cu | 31 +++++++++++-------- cpp/doxygen/developer_guide/TESTING.md | 2 +- cpp/tests/ast/transform_tests.cpp | 7 +++-- cpp/tests/column/column_test.cpp | 11 ++++--- cpp/tests/groupby/streaming_groupby_test.cpp | 12 ++++--- cpp/tests/io/fst/fst_test.cu | 10 +++--- cpp/tests/io/fst/logical_stack_test.cu | 14 ++++----- cpp/tests/join/streaming_hash_join_tests.cpp | 15 +++++---- cpp/tests/scalar/scalar_test.cpp | 13 ++++---- cpp/tests/utilities/identify_stream_usage.cpp | 4 +-- 10 files changed, 65 insertions(+), 54 deletions(-) diff --git a/cpp/benchmarks/io/fst.cu b/cpp/benchmarks/io/fst.cu index 0dbc4603b792..d3ee3853764c 100644 --- a/cpp/benchmarks/io/fst.cu +++ b/cpp/benchmarks/io/fst.cu @@ -14,7 +14,7 @@ #include #include -#include +#include #include #include @@ -26,6 +26,11 @@ #include namespace { +cuda::stream make_stream() +{ + return cuda::stream{cuda::device_ref{rmm::get_current_cuda_device().value()}}; +} + auto make_test_json_data(nvbench::state& state) { auto const string_size{cudf::size_type(state.get_int64("string_size"))}; @@ -67,7 +72,7 @@ void BM_FST_JSON(nvbench::state& state) "Benchmarks only support up to size_type's maximum number of items"); auto const string_size{cudf::size_type(state.get_int64("string_size"))}; // Prepare cuda stream for data transfers & kernels - rmm::cuda_stream stream{}; + auto stream = make_stream(); cuda::stream_ref stream_view(stream); auto input_string = make_test_json_data(state); @@ -89,7 +94,7 @@ void BM_FST_JSON(nvbench::state& state) max_translated_out>(pda_out_tt), stream); - state.set_cuda_stream(nvbench::make_cuda_stream_view(stream.value())); + state.set_cuda_stream(nvbench::make_cuda_stream_view(stream.get())); state.exec(nvbench::exec_tag::sync, [&](nvbench::launch& launch) { // Allocate device-side temporary storage & run algorithm parser.Transduce(d_input.data(), @@ -98,7 +103,7 @@ void BM_FST_JSON(nvbench::state& state) out_indexes_gpu.device_ptr(), output_gpu_size.device_ptr(), start_state, - stream.value()); + stream.get()); }); } @@ -108,7 +113,7 @@ void BM_FST_JSON_no_outidx(nvbench::state& state) "Benchmarks only support up to size_type's maximum number of items"); auto const string_size{cudf::size_type(state.get_int64("string_size"))}; // Prepare cuda stream for data transfers & kernels - rmm::cuda_stream stream{}; + auto stream = make_stream(); cuda::stream_ref stream_view(stream); auto input_string = make_test_json_data(state); @@ -130,7 +135,7 @@ void BM_FST_JSON_no_outidx(nvbench::state& state) max_translated_out>(pda_out_tt), stream); - state.set_cuda_stream(nvbench::make_cuda_stream_view(stream.value())); + state.set_cuda_stream(nvbench::make_cuda_stream_view(stream.get())); state.exec(nvbench::exec_tag::sync, [&](nvbench::launch& launch) { // Allocate device-side temporary storage & run algorithm parser.Transduce(d_input.data(), @@ -139,7 +144,7 @@ void BM_FST_JSON_no_outidx(nvbench::state& state) cuda::make_discard_iterator(), output_gpu_size.device_ptr(), start_state, - stream.value()); + stream.get()); }); } @@ -149,7 +154,7 @@ void BM_FST_JSON_no_out(nvbench::state& state) "Benchmarks only support up to size_type's maximum number of items"); auto const string_size{cudf::size_type(state.get_int64("string_size"))}; // Prepare cuda stream for data transfers & kernels - rmm::cuda_stream stream{}; + auto stream = make_stream(); cuda::stream_ref stream_view(stream); auto input_string = make_test_json_data(state); @@ -169,7 +174,7 @@ void BM_FST_JSON_no_out(nvbench::state& state) max_translated_out>(pda_out_tt), stream); - state.set_cuda_stream(nvbench::make_cuda_stream_view(stream.value())); + state.set_cuda_stream(nvbench::make_cuda_stream_view(stream.get())); state.exec(nvbench::exec_tag::sync, [&](nvbench::launch& launch) { // Allocate device-side temporary storage & run algorithm parser.Transduce(d_input.data(), @@ -178,7 +183,7 @@ void BM_FST_JSON_no_out(nvbench::state& state) cuda::make_discard_iterator(), output_gpu_size.device_ptr(), start_state, - stream.value()); + stream.get()); }); } @@ -188,7 +193,7 @@ void BM_FST_JSON_no_str(nvbench::state& state) "Benchmarks only support up to size_type's maximum number of items"); auto const string_size{cudf::size_type(state.get_int64("string_size"))}; // Prepare cuda stream for data transfers & kernels - rmm::cuda_stream stream{}; + auto stream = make_stream(); cuda::stream_ref stream_view(stream); auto input_string = make_test_json_data(state); @@ -209,7 +214,7 @@ void BM_FST_JSON_no_str(nvbench::state& state) max_translated_out>(pda_out_tt), stream); - state.set_cuda_stream(nvbench::make_cuda_stream_view(stream.value())); + state.set_cuda_stream(nvbench::make_cuda_stream_view(stream.get())); state.exec(nvbench::exec_tag::sync, [&](nvbench::launch& launch) { // Allocate device-side temporary storage & run algorithm parser.Transduce(d_input.data(), @@ -218,7 +223,7 @@ void BM_FST_JSON_no_str(nvbench::state& state) out_indexes_gpu.device_ptr(), output_gpu_size.device_ptr(), start_state, - stream.value()); + stream.get()); }); } diff --git a/cpp/doxygen/developer_guide/TESTING.md b/cpp/doxygen/developer_guide/TESTING.md index 87d4108ec6da..20c94a820ba0 100644 --- a/cpp/doxygen/developer_guide/TESTING.md +++ b/cpp/doxygen/developer_guide/TESTING.md @@ -486,7 +486,7 @@ Stream validity is determined by overloading the definition of libcudf's default libcudf `cudf::get_default_stream` returns one of `rmm`'s default stream values (depending on whether or not libcudf is compiled with per thread default stream enabled). In the preload library, this function is redefined to instead return a new user-created stream managed using a -function-local static `rmm::cuda_stream`. An invalid stream in this situation is defined as any of +function-local static `cuda::stream`. An invalid stream in this situation is defined as any of CUDA's default stream values (cudaStreamLegacy, cudaStreamDefault, or cudaStreamPerThread), since any kernel that properly uses `cudf::get_default_stream` will now instead be using the custom stream created by the preload library. diff --git a/cpp/tests/ast/transform_tests.cpp b/cpp/tests/ast/transform_tests.cpp index 59bd36eda7f1..c0c2c20db46c 100644 --- a/cpp/tests/ast/transform_tests.cpp +++ b/cpp/tests/ast/transform_tests.cpp @@ -21,9 +21,10 @@ #include #include -#include +#include #include +#include #include #include @@ -1629,7 +1630,7 @@ TYPED_TEST(TransformTest, NonDefaultStream) using Executor = TypeParam; - rmm::cuda_stream stream; + cuda::stream stream{cuda::device_ref{rmm::get_current_cuda_device().value()}}; auto c_0 = column_wrapper{3, 20, 1, 50}; auto c_1 = column_wrapper{10, 7, 20, 0}; @@ -1641,7 +1642,7 @@ TYPED_TEST(TransformTest, NonDefaultStream) auto expected = column_wrapper{13, 27, 21, 50}; auto result = Executor::compute_column(table, expression, stream); - stream.synchronize(); + stream.sync(); CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view(), verbosity); } diff --git a/cpp/tests/column/column_test.cpp b/cpp/tests/column/column_test.cpp index d7aade39aef4..fe0114e47bc3 100644 --- a/cpp/tests/column/column_test.cpp +++ b/cpp/tests/column/column_test.cpp @@ -25,7 +25,9 @@ #include #include -#include +#include + +#include #include #include @@ -642,8 +644,9 @@ struct RebindStreamColumnTest : public cudf::test::BaseFixture {}; TEST_F(RebindStreamColumnTest, RebindStreamPreservesNestedStructData) { - rmm::cuda_stream stream_a{}; - rmm::cuda_stream stream_b{}; + auto const device = cuda::device_ref{rmm::get_current_cuda_device().value()}; + cuda::stream stream_a{device}; + cuda::stream stream_b{device}; constexpr cudf::size_type num_rows{4}; std::vector h_ints(static_cast(num_rows)); @@ -654,7 +657,7 @@ TEST_F(RebindStreamColumnTest, RebindStreamPreservesNestedStructData) auto null_mask = cudf::create_null_mask( num_rows, cudf::mask_state::ALL_VALID, stream_a, cudf::get_current_device_resource_ref()); - stream_a.synchronize(); + stream_a.sync(); std::vector> children; children.push_back(std::make_unique(std::move(d_ints), std::move(null_mask), 0)); diff --git a/cpp/tests/groupby/streaming_groupby_test.cpp b/cpp/tests/groupby/streaming_groupby_test.cpp index ebef2d70c3c7..5e4c96bf1cab 100644 --- a/cpp/tests/groupby/streaming_groupby_test.cpp +++ b/cpp/tests/groupby/streaming_groupby_test.cpp @@ -20,9 +20,10 @@ #include #include -#include #include +#include + #include #include #include @@ -353,16 +354,17 @@ TEST_F(StreamingGroupbyTest, ConcurrentAggregate) batches.push_back(cudf::table_view{{keys[i], vals[i]}}); } - std::vector> streams; + auto const device = rmm::get_current_cuda_device(); + auto const stream_device = cuda::device_ref{device.value()}; + std::vector> streams; streams.reserve(num_batches); for (int i = 0; i < num_batches; ++i) { - streams.push_back(std::make_unique()); + streams.push_back(std::make_unique(stream_device)); } auto reqs = single_agg_req(1, cudf::make_sum_aggregation()); cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); - auto const device = rmm::get_current_cuda_device(); std::vector threads; std::vector errors(num_batches); // `ready` lets the main thread wait until every worker is spinning, and `start` then releases @@ -395,7 +397,7 @@ TEST_F(StreamingGroupbyTest, ConcurrentAggregate) EXPECT_FALSE(error); } for (auto const& stream : streams) { - stream->synchronize(); + stream->sync(); } auto [out_keys, results] = streaming_agg.finalize(); diff --git a/cpp/tests/io/fst/fst_test.cu b/cpp/tests/io/fst/fst_test.cu index 4efd2958ad4e..5363ea43ebcf 100644 --- a/cpp/tests/io/fst/fst_test.cu +++ b/cpp/tests/io/fst/fst_test.cu @@ -16,7 +16,7 @@ #include #include -#include +#include #include #include @@ -122,8 +122,8 @@ TEST_F(FstTest, GroundTruth) using SymbolOffsetT = uint32_t; // Prepare cuda stream for data transfers & kernels - rmm::cuda_stream stream{}; - cuda::stream_ref stream_view{stream.value()}; + cuda::stream stream{cuda::device_ref{rmm::get_current_cuda_device().value()}}; + cuda::stream_ref stream_view{stream.get()}; // Test input std::string input = R"( {)" @@ -171,7 +171,7 @@ TEST_F(FstTest, GroundTruth) out_indexes_gpu.device_ptr(), output_gpu_size.device_ptr(), start_state, - stream.value()); + stream.get()); // Async copy results from device to host output_gpu.device_to_host_async(stream_view); @@ -195,7 +195,7 @@ TEST_F(FstTest, GroundTruth) std::back_inserter(out_index_cpu)); // Make sure results have been copied back to host - stream.synchronize(); + stream.sync(); // Verify results ASSERT_EQ(output_gpu_size[0], output_cpu.size()); diff --git a/cpp/tests/io/fst/logical_stack_test.cu b/cpp/tests/io/fst/logical_stack_test.cu index 7f56989fbee1..11b257ec480c 100644 --- a/cpp/tests/io/fst/logical_stack_test.cu +++ b/cpp/tests/io/fst/logical_stack_test.cu @@ -10,7 +10,7 @@ #include -#include +#include #include #include @@ -153,8 +153,8 @@ TEST_F(LogicalStackTest, GroundTruth) constexpr SymbolT read_symbol = 'x'; // Prepare cuda stream for data transfers & kernels - rmm::cuda_stream stream{}; - cuda::stream_ref stream_view{stream.value()}; + cuda::stream stream{cuda::device_ref{rmm::get_current_cuda_device().value()}}; + cuda::stream_ref stream_view{stream.get()}; // Test input, std::string input = R"( {)" @@ -200,13 +200,13 @@ TEST_F(LogicalStackTest, GroundTruth) stack_symbols.data(), stack_symbols.size() * sizeof(SymbolT), cudaMemcpyDefault, - stream.value())); + stream.get())); CUDF_CUDA_TRY(cudaMemcpyAsync(d_stack_op_indexes.data(), stack_op_indexes.data(), stack_op_indexes.size() * sizeof(SymbolOffsetT), cudaMemcpyDefault, - stream.value())); + stream.get())); // Run algorithm fst::sparse_stack_op_to_top_of_stack( @@ -217,7 +217,7 @@ TEST_F(LogicalStackTest, GroundTruth) empty_stack_symbol, read_symbol, string_size, - stream.value()); + stream.get()); // Async copy results from device to host top_of_stack_gpu.device_to_host_async(stream_view); @@ -232,7 +232,7 @@ TEST_F(LogicalStackTest, GroundTruth) std::back_inserter(top_of_stack_cpu)); // Make sure results have been copied back to host - stream.synchronize(); + stream.sync(); // Verify results ASSERT_EQ(string_size, top_of_stack_cpu.size()); diff --git a/cpp/tests/join/streaming_hash_join_tests.cpp b/cpp/tests/join/streaming_hash_join_tests.cpp index e70a994e763a..11ca58939894 100644 --- a/cpp/tests/join/streaming_hash_join_tests.cpp +++ b/cpp/tests/join/streaming_hash_join_tests.cpp @@ -25,7 +25,6 @@ #include #include -#include #include #include @@ -120,16 +119,18 @@ TEST_F(StreamingHashJoinTest, ConcurrentInsert) auto const right_partitions = cudf::slice(right_view, slice_indices); column_wrapper left(values.begin(), values.end()); - std::vector> streams; + auto const device = rmm::get_current_cuda_device(); + auto const stream_device = cuda::device_ref{device.value()}; + std::vector> streams; streams.reserve(num_batches); for (size_type i = 0; i < num_batches; ++i) { - streams.push_back(std::make_unique()); + streams.push_back(std::make_unique(stream_device)); } std::vector const keys{0}; // Construct on a stream of its own so the inserts below run on different streams than the // hash table was built on, then synchronize it as the `insert()` docs require. - rmm::cuda_stream const build_stream; + cuda::stream const build_stream{stream_device}; cudf::streaming_hash_join joiner{right_view, keys, /*total_right_rows=*/num_batches, @@ -138,9 +139,7 @@ TEST_F(StreamingHashJoinTest, ConcurrentInsert) cudf::null_equality::EQUAL, /*load_factor=*/0.5, build_stream}; - build_stream.synchronize(); - - auto const device = rmm::get_current_cuda_device(); + build_stream.sync(); std::vector threads; std::vector errors(num_batches); // `ready` lets the main thread wait until every worker is spawned and spinning, and `start` @@ -174,7 +173,7 @@ TEST_F(StreamingHashJoinTest, ConcurrentInsert) EXPECT_FALSE(error); } for (auto const& insert_stream : streams) { - insert_stream->synchronize(); + insert_stream->sync(); } auto [left_indices, right_indices] = joiner.inner_join(cudf::table_view{{left}}, {}, stream); diff --git a/cpp/tests/scalar/scalar_test.cpp b/cpp/tests/scalar/scalar_test.cpp index df5c8e35e4ef..daa3ddd3dcca 100644 --- a/cpp/tests/scalar/scalar_test.cpp +++ b/cpp/tests/scalar/scalar_test.cpp @@ -13,9 +13,10 @@ #include #include -#include +#include #include +#include #include #include @@ -110,13 +111,13 @@ TYPED_TEST(TypedScalarTestWithoutFixedPoint, SetValue) TEST_F(ScalarTest, AsyncSetValueOwnsHostSource) { - rmm::cuda_stream stream; - auto const stream_ref = cuda::stream_ref{stream.value()}; + cuda::stream stream{cuda::device_ref{rmm::get_current_cuda_device().value()}}; + auto const stream_ref = cuda::stream_ref{stream.get()}; int32_t source = 42; lifetime_test_scalar scalar{0, true, stream_ref}; host_func_gate gate; CUDF_CUDA_TRY(cudaLaunchHostFunc( - stream.value(), [](void* data) { static_cast(data)->wait(); }, &gate)); + stream.get(), [](void* data) { static_cast(data)->wait(); }, &gate)); scalar.set_data_async(source, stream_ref); source = -1; @@ -128,8 +129,8 @@ TEST_F(ScalarTest, AsyncSetValueOwnsHostSource) TEST_F(ScalarTest, AsyncStringConstructionOwnsHostSource) { - rmm::cuda_stream stream; - auto const stream_ref = cuda::stream_ref{stream.value()}; + cuda::stream stream{cuda::device_ref{rmm::get_current_cuda_device().value()}}; + auto const stream_ref = cuda::stream_ref{stream.get()}; host_func_gate gate; auto upstream = cudf::get_current_device_resource_ref(); int allocations{0}; diff --git a/cpp/tests/utilities/identify_stream_usage.cpp b/cpp/tests/utilities/identify_stream_usage.cpp index 76dcdba61a0f..ff9ec7b9f635 100644 --- a/cpp/tests/utilities/identify_stream_usage.cpp +++ b/cpp/tests/utilities/identify_stream_usage.cpp @@ -5,7 +5,7 @@ #include -#include +#include #include #include @@ -50,7 +50,7 @@ namespace test { cuda::stream_ref const get_default_stream() { - static rmm::cuda_stream stream{}; + static cuda::stream stream{cuda::device_ref{rmm::get_current_cuda_device().value()}}; return stream; } From 3e647147efd1c4adbdc12eab51109b596e5c775a Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Tue, 22 Sep 2026 10:54:28 -0700 Subject: [PATCH 2/2] Use CUDA Runtime for current device queries --- cpp/benchmarks/io/fst.cu | 6 ++++-- cpp/doxygen/developer_guide/TESTING.md | 3 ++- cpp/tests/ast/transform_tests.cpp | 8 +++++--- cpp/tests/column/column_test.cpp | 7 ++++--- cpp/tests/groupby/streaming_groupby_test.cpp | 10 ++++++---- cpp/tests/io/fst/fst_test.cu | 12 ++++++++++-- cpp/tests/io/fst/logical_stack_test.cu | 12 ++++++++++-- cpp/tests/join/streaming_hash_join_tests.cpp | 10 ++++++---- cpp/tests/scalar/scalar_test.cpp | 12 +++++++++--- cpp/tests/utilities/identify_stream_usage.cpp | 11 +++++++---- 10 files changed, 63 insertions(+), 28 deletions(-) diff --git a/cpp/benchmarks/io/fst.cu b/cpp/benchmarks/io/fst.cu index d3ee3853764c..35538211ad61 100644 --- a/cpp/benchmarks/io/fst.cu +++ b/cpp/benchmarks/io/fst.cu @@ -14,12 +14,12 @@ #include #include -#include #include #include #include #include +#include #include @@ -28,7 +28,9 @@ namespace { cuda::stream make_stream() { - return cuda::stream{cuda::device_ref{rmm::get_current_cuda_device().value()}}; + int device{}; + CUDF_CUDA_TRY(cudaGetDevice(&device)); + return cuda::stream{cuda::device_ref{device}}; } auto make_test_json_data(nvbench::state& state) diff --git a/cpp/doxygen/developer_guide/TESTING.md b/cpp/doxygen/developer_guide/TESTING.md index 20c94a820ba0..3d28716fdd6d 100644 --- a/cpp/doxygen/developer_guide/TESTING.md +++ b/cpp/doxygen/developer_guide/TESTING.md @@ -486,7 +486,8 @@ Stream validity is determined by overloading the definition of libcudf's default libcudf `cudf::get_default_stream` returns one of `rmm`'s default stream values (depending on whether or not libcudf is compiled with per thread default stream enabled). In the preload library, this function is redefined to instead return a new user-created stream managed using a -function-local static `cuda::stream`. An invalid stream in this situation is defined as any of +function-local static pointer to an intentionally leaked `cuda::stream`. An invalid stream in this +situation is defined as any of CUDA's default stream values (cudaStreamLegacy, cudaStreamDefault, or cudaStreamPerThread), since any kernel that properly uses `cudf::get_default_stream` will now instead be using the custom stream created by the preload library. diff --git a/cpp/tests/ast/transform_tests.cpp b/cpp/tests/ast/transform_tests.cpp index c0c2c20db46c..7238057341ac 100644 --- a/cpp/tests/ast/transform_tests.cpp +++ b/cpp/tests/ast/transform_tests.cpp @@ -20,11 +20,11 @@ #include #include #include - -#include +#include #include #include +#include #include #include @@ -1630,7 +1630,9 @@ TYPED_TEST(TransformTest, NonDefaultStream) using Executor = TypeParam; - cuda::stream stream{cuda::device_ref{rmm::get_current_cuda_device().value()}}; + int device{}; + CUDF_CUDA_TRY(cudaGetDevice(&device)); + cuda::stream stream{cuda::device_ref{device}}; auto c_0 = column_wrapper{3, 20, 1, 50}; auto c_1 = column_wrapper{10, 7, 20, 0}; diff --git a/cpp/tests/column/column_test.cpp b/cpp/tests/column/column_test.cpp index fe0114e47bc3..7f68c268f248 100644 --- a/cpp/tests/column/column_test.cpp +++ b/cpp/tests/column/column_test.cpp @@ -25,9 +25,8 @@ #include #include -#include - #include +#include #include #include @@ -644,7 +643,9 @@ struct RebindStreamColumnTest : public cudf::test::BaseFixture {}; TEST_F(RebindStreamColumnTest, RebindStreamPreservesNestedStructData) { - auto const device = cuda::device_ref{rmm::get_current_cuda_device().value()}; + int device_id{}; + CUDF_CUDA_TRY(cudaGetDevice(&device_id)); + auto const device = cuda::device_ref{device_id}; cuda::stream stream_a{device}; cuda::stream stream_b{device}; diff --git a/cpp/tests/groupby/streaming_groupby_test.cpp b/cpp/tests/groupby/streaming_groupby_test.cpp index 5e4c96bf1cab..8d0ad1044570 100644 --- a/cpp/tests/groupby/streaming_groupby_test.cpp +++ b/cpp/tests/groupby/streaming_groupby_test.cpp @@ -17,12 +17,13 @@ #include #include #include +#include #include -#include #include #include +#include #include #include @@ -354,8 +355,9 @@ TEST_F(StreamingGroupbyTest, ConcurrentAggregate) batches.push_back(cudf::table_view{{keys[i], vals[i]}}); } - auto const device = rmm::get_current_cuda_device(); - auto const stream_device = cuda::device_ref{device.value()}; + int device{}; + CUDF_CUDA_TRY(cudaGetDevice(&device)); + auto const stream_device = cuda::device_ref{device}; std::vector> streams; streams.reserve(num_batches); for (int i = 0; i < num_batches; ++i) { @@ -374,7 +376,7 @@ TEST_F(StreamingGroupbyTest, ConcurrentAggregate) threads.reserve(num_batches); for (int i = 0; i < num_batches; ++i) { threads.emplace_back([&, i] { - rmm::cuda_set_device_raii const device_guard{device}; + CUDF_CUDA_TRY(cudaSetDevice(device)); ready.fetch_add(1, std::memory_order_relaxed); while (!start.load(std::memory_order_acquire)) { std::this_thread::yield(); diff --git a/cpp/tests/io/fst/fst_test.cu b/cpp/tests/io/fst/fst_test.cu index 5363ea43ebcf..6a0de25e5739 100644 --- a/cpp/tests/io/fst/fst_test.cu +++ b/cpp/tests/io/fst/fst_test.cu @@ -15,18 +15,26 @@ #include #include #include +#include -#include #include #include #include +#include #include #include namespace { +cuda::stream make_stream() +{ + int device{}; + CUDF_CUDA_TRY(cudaGetDevice(&device)); + return cuda::stream{cuda::device_ref{device}}; +} + //------------------------------------------------------------------------------ // CPU-BASED IMPLEMENTATIONS FOR VERIFICATION //------------------------------------------------------------------------------ @@ -122,7 +130,7 @@ TEST_F(FstTest, GroundTruth) using SymbolOffsetT = uint32_t; // Prepare cuda stream for data transfers & kernels - cuda::stream stream{cuda::device_ref{rmm::get_current_cuda_device().value()}}; + auto stream = make_stream(); cuda::stream_ref stream_view{stream.get()}; // Test input diff --git a/cpp/tests/io/fst/logical_stack_test.cu b/cpp/tests/io/fst/logical_stack_test.cu index 11b257ec480c..63fa0e3136a0 100644 --- a/cpp/tests/io/fst/logical_stack_test.cu +++ b/cpp/tests/io/fst/logical_stack_test.cu @@ -9,11 +9,12 @@ #include #include +#include -#include #include #include +#include #include @@ -26,6 +27,13 @@ namespace { namespace fst = cudf::io::fst; +cuda::stream make_stream() +{ + int device{}; + CUDF_CUDA_TRY(cudaGetDevice(&device)); + return cuda::stream{cuda::device_ref{device}}; +} + /** * @brief Generates the sparse representation of stack operations to feed into the logical * stack @@ -153,7 +161,7 @@ TEST_F(LogicalStackTest, GroundTruth) constexpr SymbolT read_symbol = 'x'; // Prepare cuda stream for data transfers & kernels - cuda::stream stream{cuda::device_ref{rmm::get_current_cuda_device().value()}}; + auto stream = make_stream(); cuda::stream_ref stream_view{stream.get()}; // Test input, diff --git a/cpp/tests/join/streaming_hash_join_tests.cpp b/cpp/tests/join/streaming_hash_join_tests.cpp index 11ca58939894..a2a82e292064 100644 --- a/cpp/tests/join/streaming_hash_join_tests.cpp +++ b/cpp/tests/join/streaming_hash_join_tests.cpp @@ -22,12 +22,13 @@ #include #include #include +#include #include -#include #include #include +#include #include #include @@ -119,8 +120,9 @@ TEST_F(StreamingHashJoinTest, ConcurrentInsert) auto const right_partitions = cudf::slice(right_view, slice_indices); column_wrapper left(values.begin(), values.end()); - auto const device = rmm::get_current_cuda_device(); - auto const stream_device = cuda::device_ref{device.value()}; + int device{}; + CUDF_CUDA_TRY(cudaGetDevice(&device)); + auto const stream_device = cuda::device_ref{device}; std::vector> streams; streams.reserve(num_batches); for (size_type i = 0; i < num_batches; ++i) { @@ -150,7 +152,7 @@ TEST_F(StreamingHashJoinTest, ConcurrentInsert) threads.reserve(num_batches); for (size_type i = 0; i < num_batches; ++i) { threads.emplace_back([&, i] { - rmm::cuda_set_device_raii const device_guard{device}; + CUDF_CUDA_TRY(cudaSetDevice(device)); ready.fetch_add(1, std::memory_order_relaxed); while (!start.load(std::memory_order_acquire)) { std::this_thread::yield(); diff --git a/cpp/tests/scalar/scalar_test.cpp b/cpp/tests/scalar/scalar_test.cpp index daa3ddd3dcca..b8adb4f30c4a 100644 --- a/cpp/tests/scalar/scalar_test.cpp +++ b/cpp/tests/scalar/scalar_test.cpp @@ -13,7 +13,6 @@ #include #include -#include #include #include @@ -67,6 +66,13 @@ class lifetime_test_scalar : public cudf::numeric_scalar { } }; +cuda::stream make_stream() +{ + int device{}; + CUDF_CUDA_TRY(cudaGetDevice(&device)); + return cuda::stream{cuda::device_ref{device}}; +} + } // namespace template @@ -111,7 +117,7 @@ TYPED_TEST(TypedScalarTestWithoutFixedPoint, SetValue) TEST_F(ScalarTest, AsyncSetValueOwnsHostSource) { - cuda::stream stream{cuda::device_ref{rmm::get_current_cuda_device().value()}}; + auto stream = make_stream(); auto const stream_ref = cuda::stream_ref{stream.get()}; int32_t source = 42; lifetime_test_scalar scalar{0, true, stream_ref}; @@ -129,7 +135,7 @@ TEST_F(ScalarTest, AsyncSetValueOwnsHostSource) TEST_F(ScalarTest, AsyncStringConstructionOwnsHostSource) { - cuda::stream stream{cuda::device_ref{rmm::get_current_cuda_device().value()}}; + auto stream = make_stream(); auto const stream_ref = cuda::stream_ref{stream.get()}; host_func_gate gate; auto upstream = cudf::get_current_device_resource_ref(); diff --git a/cpp/tests/utilities/identify_stream_usage.cpp b/cpp/tests/utilities/identify_stream_usage.cpp index ff9ec7b9f635..25777324f450 100644 --- a/cpp/tests/utilities/identify_stream_usage.cpp +++ b/cpp/tests/utilities/identify_stream_usage.cpp @@ -4,8 +4,7 @@ */ #include - -#include +#include #include #include @@ -50,8 +49,12 @@ namespace test { cuda::stream_ref const get_default_stream() { - static cuda::stream stream{cuda::device_ref{rmm::get_current_cuda_device().value()}}; - return stream; + static auto* stream = new cuda::stream{[] { + int device{}; + CUDF_CUDA_TRY(cudaGetDevice(&device)); + return cuda::device_ref{device}; + }()}; + return *stream; } #ifdef STREAM_MODE_TESTING