Replace libcudf cuda stream owners

cpp/benchmarks/io/fst.cu

@@ -14,18 +14,25 @@
 #include <cudf/types.hpp>
 #include <cudf/utilities/error.hpp>
 
-#include <rmm/cuda_stream.hpp>
 #include <rmm/device_buffer.hpp>
 #include <rmm/device_uvector.hpp>
 
 #include <cuda/iterator>
 #include <cuda/stream>
+#include <cuda_runtime_api.h>
 
 #include <nvbench/nvbench.cuh>
 
 #include <cstdlib>
 
 namespace {
+cuda::stream make_stream()
+{
+  int device{};
+  CUDF_CUDA_TRY(cudaGetDevice(&device));
+  return cuda::stream{cuda::device_ref{device}};
+}
+
 auto make_test_json_data(nvbench::state& state)
 {
   auto const string_size{cudf::size_type(state.get_int64("string_size"))};
@@ -67,7 +74,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 +96,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 +105,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 +115,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 +137,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 +146,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 +156,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 +176,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 +185,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 +195,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 +216,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 +225,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());
   });
 }
 

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 `rmm::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.

cpp/tests/ast/transform_tests.cpp

@@ -20,10 +20,11 @@
 #include <cudf/scalar/scalar.hpp>
 #include <cudf/table/table_view.hpp>
 #include <cudf/transform.hpp>
-
-#include <rmm/cuda_stream.hpp>
+#include <cudf/utilities/error.hpp>
 
 #include <cuda/iterator>
+#include <cuda/stream>
+#include <cuda_runtime_api.h>
 
 #include <algorithm>
 #include <array>
@@ -1629,7 +1630,9 @@ TYPED_TEST(TransformTest, NonDefaultStream)
 
   using Executor = TypeParam;
 
-  rmm::cuda_stream stream;
+  int device{};
+  CUDF_CUDA_TRY(cudaGetDevice(&device));
+  cuda::stream stream{cuda::device_ref{device}};
 
   auto c_0   = column_wrapper<int32_t>{3, 20, 1, 50};
   auto c_1   = column_wrapper<int32_t>{10, 7, 20, 0};
@@ -1641,7 +1644,7 @@ TYPED_TEST(TransformTest, NonDefaultStream)
 
   auto expected = column_wrapper<int32_t>{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);
 }

cpp/tests/column/column_test.cpp

@@ -25,7 +25,8 @@
 #include <cudf/utilities/error.hpp>
 #include <cudf/utilities/memory_resource.hpp>
 
-#include <rmm/cuda_stream.hpp>
+#include <cuda/stream>
+#include <cuda_runtime_api.h>
 
 #include <numeric>
 #include <random>
@@ -642,8 +643,11 @@ struct RebindStreamColumnTest : public cudf::test::BaseFixture {};
 
 TEST_F(RebindStreamColumnTest, RebindStreamPreservesNestedStructData)
 {
-  rmm::cuda_stream stream_a{};
-  rmm::cuda_stream stream_b{};
+  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};
 
   constexpr cudf::size_type num_rows{4};
   std::vector<int32_t> h_ints(static_cast<std::size_t>(num_rows));
@@ -654,7 +658,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<std::unique_ptr<cudf::column>> children;
   children.push_back(std::make_unique<cudf::column>(std::move(d_ints), std::move(null_mask), 0));

cpp/tests/groupby/streaming_groupby_test.cpp

@@ -17,12 +17,14 @@
 #include <cudf/table/table.hpp>
 #include <cudf/table/table_view.hpp>
 #include <cudf/unary.hpp>
+#include <cudf/utilities/error.hpp>
 #include <cudf/utilities/traits.hpp>
 
-#include <rmm/cuda_device.hpp>
-#include <rmm/cuda_stream.hpp>
 #include <rmm/mr/statistics_resource_adaptor.hpp>
 
+#include <cuda/stream>
+#include <cuda_runtime_api.h>
+
 #include <atomic>
 #include <thread>
 #include <vector>
@@ -353,16 +355,18 @@ TEST_F(StreamingGroupbyTest, ConcurrentAggregate)
     batches.push_back(cudf::table_view{{keys[i], vals[i]}});
   }
 
-  std::vector<std::unique_ptr<rmm::cuda_stream>> streams;
+  int device{};
+  CUDF_CUDA_TRY(cudaGetDevice(&device));
+  auto const stream_device = cuda::device_ref{device};
+  std::vector<std::unique_ptr<cuda::stream>> streams;
   streams.reserve(num_batches);
   for (int i = 0; i < num_batches; ++i) {
-    streams.push_back(std::make_unique<rmm::cuda_stream>());
+    streams.push_back(std::make_unique<cuda::stream>(stream_device));
   }
 
   auto reqs = single_agg_req(1, cudf::make_sum_aggregation<cudf::groupby_aggregation>());
   cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS);
 
-  auto const device = rmm::get_current_cuda_device();
   std::vector<std::thread> threads;
   std::vector<std::exception_ptr> errors(num_batches);
   // `ready` lets the main thread wait until every worker is spinning, and `start` then releases
@@ -372,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();
@@ -395,7 +399,7 @@ TEST_F(StreamingGroupbyTest, ConcurrentAggregate)
     EXPECT_FALSE(error);
   }
   for (auto const& stream : streams) {
-    stream->synchronize();
+    stream->sync();
   }
 
   auto [out_keys, results] = streaming_agg.finalize();

cpp/tests/io/fst/fst_test.cu

@@ -15,18 +15,26 @@
 #include <cudf/scalar/scalar_factories.hpp>
 #include <cudf/strings/repeat_strings.hpp>
 #include <cudf/types.hpp>
+#include <cudf/utilities/error.hpp>
 
-#include <rmm/cuda_stream.hpp>
 #include <rmm/device_buffer.hpp>
 #include <rmm/device_uvector.hpp>
 
 #include <cuda/stream>
+#include <cuda_runtime_api.h>
 
 #include <cstdlib>
 #include <vector>
 
 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,8 +130,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()};
+  auto stream = make_stream();
+  cuda::stream_ref stream_view{stream.get()};
 
   // Test input
   std::string input = R"(  {)"
@@ -171,7 +179,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 +203,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());

cpp/tests/io/fst/logical_stack_test.cu

@@ -9,11 +9,12 @@
 #include <cudf_test/testing_main.hpp>
 
 #include <cudf/types.hpp>
+#include <cudf/utilities/error.hpp>
 
-#include <rmm/cuda_stream.hpp>
 #include <rmm/device_uvector.hpp>
 
 #include <cuda/stream>
+#include <cuda_runtime_api.h>
 
 #include <src/io/fst/logical_stack.cuh>
 
@@ -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,8 +161,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()};
+  auto stream = make_stream();
+  cuda::stream_ref stream_view{stream.get()};
 
   // Test input,
   std::string input = R"(  {)"
@@ -200,13 +208,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<fst::stack_op_support::NO_RESET_SUPPORT, StackLevelT>(
@@ -217,7 +225,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 +240,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());