Use cuda::stream_ref in benchmarks

bdice · bdice · commit 73ad4eb9d898 · 2026-04-30T22:16:12.000-05:00
diff --git a/cpp/benchmarks/cuda_stream_pool/cuda_stream_pool_bench.cpp b/cpp/benchmarks/cuda_stream_pool/cuda_stream_pool_bench.cpp
@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2020-2021, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -18,7 +18,7 @@ static void BM_StreamPoolGetStream(benchmark::State& state)
 
   for (auto _ : state) {  // NOLINT(clang-analyzer-deadcode.DeadStores)
     auto stream = stream_pool.get_stream();
-    cudaStreamQuery(stream.value());
+    cudaStreamQuery(cuda::stream_ref{stream}.get());
   }
 
   state.SetItemsProcessed(static_cast<int64_t>(state.iterations()));
@@ -29,7 +29,7 @@ static void BM_CudaStreamClass(benchmark::State& state)
 {
   for (auto _ : state) {  // NOLINT(clang-analyzer-deadcode.DeadStores)
     auto stream = rmm::cuda_stream{};
-    cudaStreamQuery(stream.view().value());
+    cudaStreamQuery(cuda::stream_ref{stream}.get());
   }
 
   state.SetItemsProcessed(static_cast<int64_t>(state.iterations()));
diff --git a/cpp/benchmarks/device_uvector/device_uvector_bench.cu b/cpp/benchmarks/device_uvector/device_uvector_bench.cu
@@ -15,6 +15,7 @@
 #include <rmm/mr/per_device_resource.hpp>
 #include <rmm/mr/pool_memory_resource.hpp>
 
+#include <cuda/stream_ref>
 #include <cuda_runtime_api.h>
 #include <thrust/device_vector.h>
 #include <thrust/memory.h>
@@ -31,7 +32,7 @@ void BM_UvectorSizeConstruction(benchmark::State& state)
 
   for (auto _ : state) {  // NOLINT(clang-analyzer-deadcode.DeadStores)
     rmm::device_uvector<std::int32_t> vec(static_cast<std::size_t>(state.range(0)),
-                                          rmm::cuda_stream_view{});
+                                          cuda::stream_ref{cudaStream_t{nullptr}});
     cudaDeviceSynchronize();
   }
 
@@ -78,7 +79,7 @@ using rmm_vector    = rmm::device_vector<int32_t>;
 using rmm_uvector   = rmm::device_uvector<int32_t>;
 
 template <typename Vector>
-Vector make_vector(std::size_t num_elements, rmm::cuda_stream_view stream, bool zero_init = false)
+Vector make_vector(std::size_t num_elements, cuda::stream_ref stream, bool zero_init = false)
 {
   static_assert(std::is_same_v<Vector, thrust_vector> or std::is_same_v<Vector, rmm_vector> or
                   std::is_same_v<Vector, rmm_uvector>,
@@ -90,7 +91,7 @@ Vector make_vector(std::size_t num_elements, rmm::cuda_stream_view stream, bool
   } else if constexpr (std::is_same_v<Vector, rmm_uvector>) {
     auto vec = Vector(num_elements, stream);
     if (zero_init) {
-      cudaMemsetAsync(vec.data(), 0, num_elements * sizeof(std::int32_t), stream.value());
+      cudaMemsetAsync(vec.data(), 0, num_elements * sizeof(std::int32_t), stream.get());
     }
     return vec;
   }
@@ -111,14 +112,14 @@ void vector_workflow(std::size_t num_elements,
 {
   auto input = make_vector<Vector>(num_elements, input_stream, true);
   input_stream.synchronize();
-  for (rmm::cuda_stream_view stream : streams) {
+  for (cuda::stream_ref stream : streams) {
     auto output = make_vector<Vector>(num_elements, stream);
-    kernel<<<num_blocks, block_size, 0, stream.value()>>>(
+    kernel<<<num_blocks, block_size, 0, stream.get()>>>(
       vector_data(input), vector_data(output), num_elements);
   }
 
-  for (rmm::cuda_stream_view stream : streams) {
-    stream.synchronize();
+  for (cuda::stream_ref stream : streams) {
+    RMM_CUDA_TRY(cudaStreamSynchronize(stream.get()));
   }
 }
 
diff --git a/cpp/benchmarks/multi_stream_allocations/multi_stream_allocations_bench.cu b/cpp/benchmarks/multi_stream_allocations/multi_stream_allocations_bench.cu
@@ -15,6 +15,7 @@
 #include <rmm/mr/pool_memory_resource.hpp>
 #include <rmm/resource_ref.hpp>
 
+#include <cuda/stream_ref>
 #include <cuda_runtime_api.h>
 
 #include <benchmark/benchmark.h>
@@ -54,9 +55,9 @@ static void run_test(std::size_t num_kernels,
                      rmm::device_async_resource_ref mr)
 {
   for (std::size_t i = 0; i < num_kernels; i++) {
-    auto stream = stream_pool.get_stream(i);
+    auto stream = cuda::stream_ref{stream_pool.get_stream(i)};
     auto buffer = rmm::device_uvector<int64_t>(1, stream, mr);
-    compute_bound_kernel<<<1, 1, 0, stream.value()>>>(buffer.data());
+    compute_bound_kernel<<<1, 1, 0, stream.get()>>>(buffer.data());
   }
 }
 
diff --git a/cpp/benchmarks/random_allocations/random_allocations.cpp b/cpp/benchmarks/random_allocations/random_allocations.cpp
@@ -13,6 +13,9 @@
 #include <rmm/mr/pool_memory_resource.hpp>
 #include <rmm/resource_ref.hpp>
 
+#include <cuda/stream_ref>
+#include <cuda_runtime_api.h>
+
 #include <benchmark/benchmark.h>
 #include <benchmarks/utilities/cxxopts.hpp>
 
@@ -54,7 +57,7 @@ void random_allocation_free(rmm::device_async_resource_ref mr,
                             SizeDistribution size_distribution,
                             std::size_t num_allocations,
                             std::size_t max_usage,  // in MiB
-                            rmm::cuda_stream_view stream = {})
+                            cuda::stream_ref stream = cuda::stream_ref{cudaStream_t{nullptr}})
 {
   std::default_random_engine generator;
 
@@ -132,7 +135,7 @@ void uniform_random_allocations(
   std::size_t num_allocations,      // NOLINT(bugprone-easily-swappable-parameters)
   std::size_t max_allocation_size,  // size in MiB
   std::size_t max_usage,
-  rmm::cuda_stream_view stream = {})
+  cuda::stream_ref stream = cuda::stream_ref{cudaStream_t{nullptr}})
 {
   std::uniform_int_distribution<std::size_t> size_distribution(1, max_allocation_size * size_mb);
   random_allocation_free(mr, size_distribution, num_allocations, max_usage, stream);
@@ -144,7 +147,7 @@ void uniform_random_allocations(
                                 std::size_t mean_allocation_size = 500, // in MiB
                                 std::size_t stddev_allocation_size = 500, // in MiB
                                 std::size_t max_usage = 8 << 20,
-                                cuda_stream_view stream) {
+                                cuda::stream_ref stream) {
   std::normal_distribution<std::size_t> size_distribution(, max_allocation_size * size_mb);
 }*/
 
diff --git a/cpp/benchmarks/replay/replay.cpp b/cpp/benchmarks/replay/replay.cpp
@@ -4,7 +4,6 @@
  */
 
 #include <rmm/aligned.hpp>
-#include <rmm/cuda_stream_view.hpp>
 #include <rmm/detail/error.hpp>
 #include <rmm/logger.hpp>
 #include <rmm/mr/arena_memory_resource.hpp>
@@ -16,6 +15,7 @@
 #include <rmm/resource_ref.hpp>
 
 #include <cuda/iterator>
+#include <cuda/stream_ref>
 #include <thrust/execution_policy.h>
 #include <thrust/reduce.h>
 
@@ -253,8 +253,14 @@ std::vector<std::vector<rmm::detail::event>> parse_per_thread_events(std::string
                           [](auto const& event) {
                             cudaStream_t custream;
                             memcpy(&custream, &event.stream, sizeof(cudaStream_t));
-                            auto stream = rmm::cuda_stream_view{custream};
-                            return stream.is_default() or stream.is_per_thread_default();
+                            auto stream = cuda::stream_ref{custream};
+#ifdef CUDA_API_PER_THREAD_DEFAULT_STREAM
+                            return stream.get() == cudaStreamLegacy or
+                                   stream.get() == cudaStreamPerThread or stream.get() == nullptr;
+#else
+                            return stream.get() == cudaStreamLegacy or stream.get() == nullptr or
+                                   stream.get() == cudaStreamPerThread;
+#endif
                           }),
               "Non-default streams not currently supported.");
 
diff --git a/cpp/benchmarks/synchronization/synchronization.cpp b/cpp/benchmarks/synchronization/synchronization.cpp
@@ -19,7 +19,7 @@
 
 cuda_event_timer::cuda_event_timer(benchmark::State& state,
                                    bool flush_l2_cache,
-                                   rmm::cuda_stream_view stream)
+                                   cuda::stream_ref stream)
   : stream(stream), p_state(&state)
 {
   // flush all of L2$
@@ -36,18 +36,18 @@ cuda_event_timer::cuda_event_timer(benchmark::State& state,
       RMM_CUDA_TRY(cudaMemsetAsync(l2_cache_buffer.data(),
                                    memset_value,
                                    static_cast<std::size_t>(l2_cache_bytes),
-                                   stream.value()));
+                                   stream.get()));
     }
   }
 
   RMM_CUDA_TRY(cudaEventCreate(&start));
   RMM_CUDA_TRY(cudaEventCreate(&stop));
-  RMM_CUDA_TRY(cudaEventRecord(start, stream.value()));
+  RMM_CUDA_TRY(cudaEventRecord(start, stream.get()));
 }
 
 cuda_event_timer::~cuda_event_timer()
 {
-  RMM_CUDA_ASSERT_OK(cudaEventRecord(stop, stream.value()));
+  RMM_CUDA_ASSERT_OK(cudaEventRecord(stop, stream.get()));
   RMM_CUDA_ASSERT_OK(cudaEventSynchronize(stop));
 
   float milliseconds = 0.0F;
diff --git a/cpp/benchmarks/synchronization/synchronization.hpp b/cpp/benchmarks/synchronization/synchronization.hpp
@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2019-2021, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -47,9 +47,7 @@
 
 #pragma once
 
-#include <rmm/cuda_stream_view.hpp>
-
-// Google Benchmark library
+#include <cuda/stream_ref>
 #include <cuda_runtime_api.h>
 
 #include <benchmark/benchmark.h>
@@ -68,7 +66,7 @@ class cuda_event_timer {
    */
   cuda_event_timer(benchmark::State& state,
                    bool flush_l2_cache,
-                   rmm::cuda_stream_view stream = rmm::cuda_stream_default);
+                   cuda::stream_ref stream = cuda::stream_ref{cudaStream_t{nullptr}});
 
   // The user will HAVE to provide a benchmark::State object to set
   // the timer so we disable the default c'tor.
@@ -88,6 +86,6 @@ class cuda_event_timer {
  private:
   cudaEvent_t start{};
   cudaEvent_t stop{};
-  rmm::cuda_stream_view stream{};
+  cuda::stream_ref stream{cudaStream_t{nullptr}};
   benchmark::State* p_state{};
 };

Original file line number	Diff line number	Diff line change
`@@ -1,5 +1,5 @@`
`1`	`1`	`/*`
`2`		`- * SPDX-FileCopyrightText: Copyright (c) 2020-2021, NVIDIA CORPORATION.`
	`2`	`+ * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION.`
`3`	`3`	`* SPDX-License-Identifier: Apache-2.0`
`4`	`4`	`*/`
`5`	`5`
`@@ -18,7 +18,7 @@ static void BM_StreamPoolGetStream(benchmark::State& state)`
`18`	`18`
`19`	`19`	`for (auto _ : state) { // NOLINT(clang-analyzer-deadcode.DeadStores)`
`20`	`20`	`auto stream = stream_pool.get_stream();`
`21`		`- cudaStreamQuery(stream.value());`
	`21`	`+ cudaStreamQuery(cuda::stream_ref{stream}.get());`
`22`	`22`	`}`
`23`	`23`
`24`	`24`	`state.SetItemsProcessed(static_cast<int64_t>(state.iterations()));`
`@@ -29,7 +29,7 @@ static void BM_CudaStreamClass(benchmark::State& state)`
`29`	`29`	`{`
`30`	`30`	`for (auto _ : state) { // NOLINT(clang-analyzer-deadcode.DeadStores)`
`31`	`31`	`auto stream = rmm::cuda_stream{};`
`32`		`- cudaStreamQuery(stream.view().value());`
	`32`	`+ cudaStreamQuery(cuda::stream_ref{stream}.get());`
`33`	`33`	`}`
`34`	`34`
`35`	`35`	`state.SetItemsProcessed(static_cast<int64_t>(state.iterations()));`
Original file line number	Diff line number	Diff line change
`@@ -15,6 +15,7 @@`
`15`	`15`	`#include <rmm/mr/per_device_resource.hpp>`
`16`	`16`	`#include <rmm/mr/pool_memory_resource.hpp>`
`17`	`17`
	`18`	`+#include <cuda/stream_ref>`
`18`	`19`	`#include <cuda_runtime_api.h>`
`19`	`20`	`#include <thrust/device_vector.h>`
`20`	`21`	`#include <thrust/memory.h>`
`@@ -31,7 +32,7 @@ void BM_UvectorSizeConstruction(benchmark::State& state)`
`31`	`32`
`32`	`33`	`for (auto _ : state) { // NOLINT(clang-analyzer-deadcode.DeadStores)`
`33`	`34`	`rmm::device_uvector<std::int32_t> vec(static_cast<std::size_t>(state.range(0)),`
`34`		`- rmm::cuda_stream_view{});`
	`35`	`+ cuda::stream_ref{cudaStream_t{nullptr}});`
`35`	`36`	`cudaDeviceSynchronize();`
`36`	`37`	`}`
`37`	`38`
`@@ -78,7 +79,7 @@ using rmm_vector = rmm::device_vector<int32_t>;`
`78`	`79`	`using rmm_uvector = rmm::device_uvector<int32_t>;`
`79`	`80`
`80`	`81`	`template <typename Vector>`
`81`		`-Vector make_vector(std::size_t num_elements, rmm::cuda_stream_view stream, bool zero_init = false)`
	`82`	`+Vector make_vector(std::size_t num_elements, cuda::stream_ref stream, bool zero_init = false)`
`82`	`83`	`{`
`83`	`84`	`static_assert(std::is_same_v<Vector, thrust_vector> or std::is_same_v<Vector, rmm_vector> or`
`84`	`85`	`std::is_same_v<Vector, rmm_uvector>,`
`@@ -90,7 +91,7 @@ Vector make_vector(std::size_t num_elements, rmm::cuda_stream_view stream, bool`
`90`	`91`	`} else if constexpr (std::is_same_v<Vector, rmm_uvector>) {`
`91`	`92`	`auto vec = Vector(num_elements, stream);`
`92`	`93`	`if (zero_init) {`
`93`		`- cudaMemsetAsync(vec.data(), 0, num_elements * sizeof(std::int32_t), stream.value());`
	`94`	`+ cudaMemsetAsync(vec.data(), 0, num_elements * sizeof(std::int32_t), stream.get());`
`94`	`95`	`}`
`95`	`96`	`return vec;`
`96`	`97`	`}`
`@@ -111,14 +112,14 @@ void vector_workflow(std::size_t num_elements,`
`111`	`112`	`{`
`112`	`113`	`auto input = make_vector<Vector>(num_elements, input_stream, true);`
`113`	`114`	`input_stream.synchronize();`
`114`		`- for (rmm::cuda_stream_view stream : streams) {`
	`115`	`+ for (cuda::stream_ref stream : streams) {`
`115`	`116`	`auto output = make_vector<Vector>(num_elements, stream);`
`116`		`- kernel<<<num_blocks, block_size, 0, stream.value()>>>(`
	`117`	`+ kernel<<<num_blocks, block_size, 0, stream.get()>>>(`
`117`	`118`	`vector_data(input), vector_data(output), num_elements);`
`118`	`119`	`}`
`119`	`120`
`120`		`- for (rmm::cuda_stream_view stream : streams) {`
`121`		`- stream.synchronize();`
	`121`	`+ for (cuda::stream_ref stream : streams) {`
	`122`	`+ RMM_CUDA_TRY(cudaStreamSynchronize(stream.get()));`
`122`	`123`	`}`
`123`	`124`	`}`
`124`	`125`
Original file line number	Diff line number	Diff line change
`@@ -19,7 +19,7 @@`
`19`	`19`
`20`	`20`	`cuda_event_timer::cuda_event_timer(benchmark::State& state,`
`21`	`21`	`bool flush_l2_cache,`
`22`		`- rmm::cuda_stream_view stream)`
	`22`	`+ cuda::stream_ref stream)`
`23`	`23`	`: stream(stream), p_state(&state)`
`24`	`24`	`{`
`25`	`25`	`// flush all of L2$`
`@@ -36,18 +36,18 @@ cuda_event_timer::cuda_event_timer(benchmark::State& state,`
`36`	`36`	`RMM_CUDA_TRY(cudaMemsetAsync(l2_cache_buffer.data(),`
`37`	`37`	`memset_value,`
`38`	`38`	`static_cast<std::size_t>(l2_cache_bytes),`
`39`		`- stream.value()));`
	`39`	`+ stream.get()));`
`40`	`40`	`}`
`41`	`41`	`}`
`42`	`42`
`43`	`43`	`RMM_CUDA_TRY(cudaEventCreate(&start));`
`44`	`44`	`RMM_CUDA_TRY(cudaEventCreate(&stop));`
`45`		`- RMM_CUDA_TRY(cudaEventRecord(start, stream.value()));`
	`45`	`+ RMM_CUDA_TRY(cudaEventRecord(start, stream.get()));`
`46`	`46`	`}`
`47`	`47`
`48`	`48`	`cuda_event_timer::~cuda_event_timer()`
`49`	`49`	`{`
`50`		`- RMM_CUDA_ASSERT_OK(cudaEventRecord(stop, stream.value()));`
	`50`	`+ RMM_CUDA_ASSERT_OK(cudaEventRecord(stop, stream.get()));`
`51`	`51`	`RMM_CUDA_ASSERT_OK(cudaEventSynchronize(stop));`
`52`	`52`
`53`	`53`	`float milliseconds = 0.0F;`