From 08ab6912f879de200e47f7aa22e2e39ba499f0b1 Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Thu, 20 Aug 2026 09:45:17 -0700 Subject: [PATCH] Use cuda::stream_ref for central stream helpers --- cpp/include/cudf/detail/utilities/default_stream.hpp | 7 +++---- cpp/include/cudf/utilities/default_stream.hpp | 4 ++-- .../tests/streaming/test_bloom_filter.cu | 4 ++-- cpp/src/utilities/default_stream.cpp | 10 +++++----- java/src/main/native/src/CompiledExpression.cpp | 2 +- .../pylibcudf/libcudf/detail/utilities/stream_pool.pxd | 10 ++++++---- python/pylibcudf/pylibcudf/libcudf/interop.pxd | 2 +- 7 files changed, 20 insertions(+), 19 deletions(-) diff --git a/cpp/include/cudf/detail/utilities/default_stream.hpp b/cpp/include/cudf/detail/utilities/default_stream.hpp index 11fb1aae6dd2..9e17fd8e4818 100644 --- a/cpp/include/cudf/detail/utilities/default_stream.hpp +++ b/cpp/include/cudf/detail/utilities/default_stream.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -7,8 +7,7 @@ #include -#include -#include +#include namespace CUDF_EXPORT cudf { @@ -20,7 +19,7 @@ namespace detail { * Use this value to ensure the correct stream is used when compiled with per * thread default stream. */ -extern rmm::cuda_stream_view const default_stream_value; +extern cuda::stream_ref const default_stream_value; } // namespace detail diff --git a/cpp/include/cudf/utilities/default_stream.hpp b/cpp/include/cudf/utilities/default_stream.hpp index 8d56848ca0be..822dad21fcda 100644 --- a/cpp/include/cudf/utilities/default_stream.hpp +++ b/cpp/include/cudf/utilities/default_stream.hpp @@ -7,7 +7,7 @@ #include -#include +#include /** * @file @@ -25,7 +25,7 @@ namespace CUDF_EXPORT cudf { * * @return The current default stream. */ -rmm::cuda_stream_view const get_default_stream(); +cuda::stream_ref const get_default_stream(); /** * @brief Check if per-thread default stream is enabled. diff --git a/cpp/libcudf_streaming/tests/streaming/test_bloom_filter.cu b/cpp/libcudf_streaming/tests/streaming/test_bloom_filter.cu index b8fc5b2e978a..418540526f25 100644 --- a/cpp/libcudf_streaming/tests/streaming/test_bloom_filter.cu +++ b/cpp/libcudf_streaming/tests/streaming/test_bloom_filter.cu @@ -42,8 +42,8 @@ TEST(BloomFilterPolicyTest, UsesBlocksBeyondFormerArrowLimit) auto const stream = cudf::get_default_stream(); rmm::device_scalar index{0, stream}; - block_index_kernel<<<1, 1, 0, stream.value()>>>(upper_hash, num_blocks, index.data()); - CUDF_CHECK_CUDA(stream.value()); + block_index_kernel<<<1, 1, 0, stream.get()>>>(upper_hash, num_blocks, index.data()); + CUDF_CHECK_CUDA(stream.get()); EXPECT_EQ(index.value(stream), arrow_max_blocks); } diff --git a/cpp/src/utilities/default_stream.cpp b/cpp/src/utilities/default_stream.cpp index fa8ab67c53f1..7de1ca987b6e 100644 --- a/cpp/src/utilities/default_stream.cpp +++ b/cpp/src/utilities/default_stream.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2020-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -12,9 +12,9 @@ namespace cudf { namespace detail { #if defined(CUDF_USE_PER_THREAD_DEFAULT_STREAM) -rmm::cuda_stream_view const default_stream_value{rmm::cuda_stream_per_thread}; +cuda::stream_ref const default_stream_value{cuda::stream_ref{cudaStreamPerThread}}; #else -rmm::cuda_stream_view const default_stream_value{}; +cuda::stream_ref const default_stream_value{cudaStream_t{nullptr}}; #endif } // namespace detail @@ -33,11 +33,11 @@ bool is_ptds_enabled() #endif } -rmm::cuda_stream_view const get_default_stream() +cuda::stream_ref const get_default_stream() { static auto const default_stream = []() { if (std::getenv("CUDF_PER_THREAD_STREAM") != nullptr) { - return rmm::cuda_stream_per_thread; + return cuda::stream_ref{cudaStreamPerThread}; } else { return detail::default_stream_value; } diff --git a/java/src/main/native/src/CompiledExpression.cpp b/java/src/main/native/src/CompiledExpression.cpp index 44748ef6b9b9..dcfb0989218f 100644 --- a/java/src/main/native/src/CompiledExpression.cpp +++ b/java/src/main/native/src/CompiledExpression.cpp @@ -538,7 +538,7 @@ std::unique_ptr compile_serialized_ast(jni_serial if (!jni_ast.at_eof()) { throw std::invalid_argument("Extra bytes at end of serialized AST"); } // The expression may be handed to a thread with a different default stream. - if (jni_expr_ptr->has_literals()) { cudf::get_default_stream().synchronize(); } + if (jni_expr_ptr->has_literals()) { cudf::get_default_stream().sync(); } return jni_expr_ptr; } diff --git a/python/pylibcudf/pylibcudf/libcudf/detail/utilities/stream_pool.pxd b/python/pylibcudf/pylibcudf/libcudf/detail/utilities/stream_pool.pxd index b05edceb3abc..7930d731b224 100644 --- a/python/pylibcudf/pylibcudf/libcudf/detail/utilities/stream_pool.pxd +++ b/python/pylibcudf/pylibcudf/libcudf/detail/utilities/stream_pool.pxd @@ -12,9 +12,7 @@ cdef extern from * nogil: """ #include #include - #include - #include namespace { @@ -22,8 +20,12 @@ cdef extern from * nogil: cudf::host_span streams, cudaStream_t stream ) { - std::vector stream_refs(streams.begin(), streams.end()); - cudf::detail::join_streams(stream_refs, stream); + std::vector stream_refs; + stream_refs.reserve(streams.size()); + for (auto const s : streams) { + stream_refs.emplace_back(s); + } + cudf::detail::join_streams(stream_refs, cuda::stream_ref{stream}); } } """ diff --git a/python/pylibcudf/pylibcudf/libcudf/interop.pxd b/python/pylibcudf/pylibcudf/libcudf/interop.pxd index 78fc455dd351..176c72d7f7ac 100644 --- a/python/pylibcudf/pylibcudf/libcudf/interop.pxd +++ b/python/pylibcudf/pylibcudf/libcudf/interop.pxd @@ -175,7 +175,7 @@ cdef extern from *: ArrowDeviceArray* to_arrow_device_raw( ViewType const& obj, PyObject* owner, - cudaStream_t stream = cudf::get_default_stream(), + cuda::stream_ref stream = cudf::get_default_stream(), rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()) { auto tmp = cudf::to_arrow_device(obj, stream, mr);