Use cuda::stream_ref for compute APIs

cpp/include/cudf/aggregation/host_udf.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -11,9 +11,10 @@
 #include <cudf/utilities/export.hpp>
 #include <cudf/utilities/span.hpp>
 
-#include <rmm/cuda_stream_view.hpp>
 #include <rmm/resource_ref.hpp>
 
+#include <cuda/stream_ref>
+
 #include <functional>
 #include <optional>
 
@@ -96,7 +97,7 @@ class host_udf_base {
  *     column_view const& input,
  *     data_type output_dtype,
  *     std::optional<std::reference_wrapper<scalar const>> init,
- *     rmm::cuda_stream_view stream,
+ *     cuda::stream_ref stream,
  *     rmm::device_async_resource_ref mr) const override
  *   {
  *     // Perform reduction computation using the input data and return the reduction result.
@@ -132,7 +133,7 @@ struct reduce_host_udf : host_udf_base {
     column_view const& input,
     data_type output_dtype,
     std::optional<std::reference_wrapper<scalar const>> init,
-    rmm::cuda_stream_view stream,
+    cuda::stream_ref stream,
     rmm::device_async_resource_ref mr) const = 0;
 };
 
@@ -154,7 +155,7 @@ struct reduce_host_udf : host_udf_base {
  *     data_type output_dtype,
  *     null_policy null_handling,
  *     std::optional<std::reference_wrapper<scalar const>> init,
- *     rmm::cuda_stream_view stream,
+ *     cuda::stream_ref stream,
  *     rmm::device_async_resource_ref mr) const override
  *   {
  *     // Perform computation using the input data and return the result.
@@ -196,7 +197,7 @@ struct segmented_reduce_host_udf : host_udf_base {
     data_type output_dtype,
     null_policy null_handling,
     std::optional<std::reference_wrapper<scalar const>> init,
-    rmm::cuda_stream_view stream,
+    cuda::stream_ref stream,
     rmm::device_async_resource_ref mr) const = 0;
 };
 
@@ -226,14 +227,14 @@ struct aggregate_result_functor;
  *   my_udf_aggregation() = default;
  *
  *   [[nodiscard]] std::unique_ptr<column> get_empty_output(
- *     rmm::cuda_stream_view stream,
+ *     cuda::stream_ref stream,
  *     rmm::device_async_resource_ref mr) const override
  *   {
  *     // Return a column corresponding to the result when the input values column is empty.
  *   }
  *
  *   [[nodiscard]] std::unique_ptr<column> operator()(
- *     rmm::cuda_stream_view stream,
+ *     cuda::stream_ref stream,
  *     rmm::device_async_resource_ref mr) const override
  *   {
  *     // Perform UDF computation using the input data and return the result.
@@ -265,7 +266,7 @@ struct groupby_host_udf : host_udf_base {
    * @return The output result of the aggregation when the input values column is empty
    */
   [[nodiscard]] virtual std::unique_ptr<column> get_empty_output(
-    rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const = 0;
+    cuda::stream_ref stream, rmm::device_async_resource_ref mr) const = 0;
 
   /**
    * @brief Perform the main groupby computation for the host-based UDF.
@@ -275,7 +276,7 @@ struct groupby_host_udf : host_udf_base {
    * @return The output result of the aggregation
    */
   [[nodiscard]] virtual std::unique_ptr<column> operator()(
-    rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const = 0;
+    cuda::stream_ref stream, rmm::device_async_resource_ref mr) const = 0;
 
  private:
   // Allow the struct `aggregate_result_functor` to set its private callback variables.

cpp/include/cudf/binaryop.hpp

@@ -160,7 +160,7 @@ std::unique_ptr<column> binary_operation(
   column_view const& rhs,
   binary_operator op,
   data_type output_type,
-  rmm::cuda_stream_view 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());
 
 /**
@@ -191,7 +191,7 @@ std::unique_ptr<column> binary_operation(
   scalar const& rhs,
   binary_operator op,
   data_type output_type,
-  rmm::cuda_stream_view 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());
 
 /**
@@ -221,7 +221,7 @@ std::unique_ptr<column> binary_operation(
   column_view const& rhs,
   binary_operator op,
   data_type output_type,
-  rmm::cuda_stream_view 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());
 
 /**
@@ -252,7 +252,7 @@ std::unique_ptr<column> binary_operation(
   column_view const& rhs,
   std::string const& ptx,
   data_type output_type,
-  rmm::cuda_stream_view 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());
 
 /**
@@ -305,7 +305,7 @@ bool is_supported_operation(data_type out, data_type lhs, data_type rhs, binary_
 std::pair<rmm::device_buffer, size_type> scalar_col_valid_mask_and(
   column_view const& col,
   scalar const& s,
-  rmm::cuda_stream_view 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());
 
 }  // namespace binops
@@ -333,6 +333,6 @@ void apply_sorting_struct_binary_op(mutable_column_view& out,
                                     bool is_lhs_scalar,
                                     bool is_rhs_scalar,
                                     binary_operator op,
-                                    rmm::cuda_stream_view stream);
+                                    cuda::stream_ref stream);
 }  // namespace binops::compiled::detail
 }  // namespace CUDF_EXPORT cudf

cpp/include/cudf/contiguous_split.hpp

@@ -104,6 +104,7 @@ struct contiguous_split_state;
  * // data. In memory constrained cases, this can be used to set aside scratch memory
  * // for `chunked_pack` at the beginning of a program.
  * auto mr = cudf::get_current_device_resource_ref();
+ * cuda::stream_ref stream = cudf::get_default_stream();
  *
  * // Define a buffer size for each chunk: the larger the buffer is, the more SMs can be
  * // occupied by this algorithm.
@@ -116,7 +117,7 @@ struct contiguous_split_state;
  * //
  * std::size_t user_buffer_size = 128*1024*1024;
  *
- * auto chunked_packer = cudf::chunked_pack::create(tv, user_buffer_size, mr);
+ * auto chunked_packer = cudf::chunked_pack::create(tv, user_buffer_size, stream, mr);
  *
  * std::size_t host_offset = 0;
  * auto host_buffer = ...; // obtain a host buffer you would like to copy to
@@ -134,7 +135,7 @@ struct contiguous_split_state;
  *     user_buffer.data(),
  *     bytes_copied,
  *     cudaMemcpyDefault,
- *     stream);
+ *     stream.get());
  *
  *   host_offset += bytes_copied;
  * }

cpp/include/cudf/detail/binaryop.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2018-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2018-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -8,54 +8,54 @@
 #include <cudf/utilities/default_stream.hpp>
 #include <cudf/utilities/memory_resource.hpp>
 
-#include <rmm/cuda_stream_view.hpp>
+#include <cuda/stream_ref>
 
 namespace cudf {
 //! Inner interfaces and implementations
 namespace detail {
 
 /**
  * @copydoc cudf::binary_operation(column_view const&, column_view const&,
- * std::string const&, data_type, rmm::cuda_stream_view, rmm::device_async_resource_ref)
+ * std::string const&, data_type, cuda::stream_ref, rmm::device_async_resource_ref)
  */
 std::unique_ptr<column> binary_operation(column_view const& lhs,
                                          column_view const& rhs,
                                          std::string const& ptx,
                                          data_type output_type,
-                                         rmm::cuda_stream_view stream,
+                                         cuda::stream_ref stream,
                                          rmm::device_async_resource_ref mr);
 
 /**
  * @copydoc cudf::binary_operation(scalar const&, column_view const&, binary_operator,
- * data_type, rmm::cuda_stream_view, rmm::device_async_resource_ref)
+ * data_type, cuda::stream_ref, rmm::device_async_resource_ref)
  */
 std::unique_ptr<column> binary_operation(scalar const& lhs,
                                          column_view const& rhs,
                                          binary_operator op,
                                          data_type output_type,
-                                         rmm::cuda_stream_view stream,
+                                         cuda::stream_ref stream,
                                          rmm::device_async_resource_ref mr);
 
 /**
  * @copydoc cudf::binary_operation(column_view const&, scalar const&, binary_operator,
- * data_type, rmm::cuda_stream_view, rmm::device_async_resource_ref)
+ * data_type, cuda::stream_ref, rmm::device_async_resource_ref)
  */
 std::unique_ptr<column> binary_operation(column_view const& lhs,
                                          scalar const& rhs,
                                          binary_operator op,
                                          data_type output_type,
-                                         rmm::cuda_stream_view stream,
+                                         cuda::stream_ref stream,
                                          rmm::device_async_resource_ref mr);
 
 /**
  * @copydoc cudf::binary_operation(column_view const&, column_view const&,
- * binary_operator, data_type, rmm::cuda_stream_view, rmm::device_async_resource_ref)
+ * binary_operator, data_type, cuda::stream_ref, rmm::device_async_resource_ref)
  */
 std::unique_ptr<column> binary_operation(column_view const& lhs,
                                          column_view const& rhs,
                                          binary_operator op,
                                          data_type output_type,
-                                         rmm::cuda_stream_view stream,
+                                         cuda::stream_ref stream,
                                          rmm::device_async_resource_ref mr);
 }  // namespace detail
 }  // namespace cudf

cpp/include/cudf/detail/groupby.hpp

@@ -9,7 +9,7 @@
 #include <cudf/utilities/memory_resource.hpp>
 #include <cudf/utilities/span.hpp>
 
-#include <rmm/cuda_stream_view.hpp>
+#include <cuda/stream_ref>
 
 #include <memory>
 #include <utility>
@@ -32,7 +32,7 @@ std::pair<std::unique_ptr<table>, std::vector<aggregation_result>> groupby(
   table_view const& keys,
   std::span<aggregation_request const> requests,
   null_policy include_null_keys,
-  rmm::cuda_stream_view stream,
+  cuda::stream_ref stream,
   rmm::device_async_resource_ref mr);
 }  // namespace groupby::detail::hash
 }  // namespace cudf

cpp/include/cudf/detail/groupby/group_replace_nulls.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -30,7 +30,7 @@ namespace detail {
 std::unique_ptr<column> group_replace_nulls(cudf::column_view const& grouped_value,
                                             device_span<size_type const> group_labels,
                                             cudf::replace_policy replace_policy,
-                                            rmm::cuda_stream_view stream,
+                                            cuda::stream_ref stream,
                                             rmm::device_async_resource_ref mr);
 
 }  // namespace detail

cpp/include/cudf/detail/groupby/sort_helper.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -11,9 +11,10 @@
 #include <cudf/types.hpp>
 #include <cudf/utilities/memory_resource.hpp>
 
-#include <rmm/cuda_stream_view.hpp>
 #include <rmm/device_uvector.hpp>
 
+#include <cuda/stream_ref>
+
 namespace cudf {
 namespace groupby::detail::sort {
 /**
@@ -76,7 +77,7 @@ struct sort_groupby_helper {
    * @return the sorted and grouped column
    */
   std::unique_ptr<column> sorted_values(column_view const& values,
-                                        rmm::cuda_stream_view stream,
+                                        cuda::stream_ref stream,
                                         rmm::device_async_resource_ref mr);
 
   /**
@@ -92,29 +93,27 @@ struct sort_groupby_helper {
    * @return the grouped column
    */
   std::unique_ptr<column> grouped_values(column_view const& values,
-                                         rmm::cuda_stream_view stream,
+                                         cuda::stream_ref stream,
                                          rmm::device_async_resource_ref mr);
 
   /**
    * @brief Get a table of sorted unique keys
    *
    * @return a new table in which each row is a unique row in the sorted key table.
    */
-  std::unique_ptr<table> unique_keys(rmm::cuda_stream_view stream,
-                                     rmm::device_async_resource_ref mr);
+  std::unique_ptr<table> unique_keys(cuda::stream_ref stream, rmm::device_async_resource_ref mr);
 
   /**
    * @brief Get a table of sorted keys
    *
    * @return a new table containing the sorted keys.
    */
-  std::unique_ptr<table> sorted_keys(rmm::cuda_stream_view stream,
-                                     rmm::device_async_resource_ref mr);
+  std::unique_ptr<table> sorted_keys(cuda::stream_ref stream, rmm::device_async_resource_ref mr);
 
   /**
    * @brief Get the number of groups in `keys`
    */
-  size_type num_groups(rmm::cuda_stream_view stream) { return group_offsets(stream).size() - 1; }
+  size_type num_groups(cuda::stream_ref stream) { return group_offsets(stream).size() - 1; }
 
   /**
    * @brief check if the groupby keys are presorted
@@ -128,7 +127,7 @@ struct sort_groupby_helper {
    * When include_null_keys = NO, returned value is the number of rows in `keys`
    *  in which no element is null
    */
-  size_type num_keys(rmm::cuda_stream_view stream);
+  size_type num_keys(cuda::stream_ref stream);
 
   /**
    * @brief Get the sorted order of `keys`.
@@ -143,7 +142,7 @@ struct sort_groupby_helper {
    *
    * @return the sort order indices for `keys`.
    */
-  column_view key_sort_order(rmm::cuda_stream_view stream);
+  column_view key_sort_order(cuda::stream_ref stream);
 
   /**
    * @brief Get each group's offset into the sorted order of `keys`.
@@ -156,7 +155,7 @@ struct sort_groupby_helper {
    * @return vector of offsets of the starting point of each group in the sorted
    * key table
    */
-  index_vector const& group_offsets(rmm::cuda_stream_view stream);
+  index_vector const& group_offsets(cuda::stream_ref stream);
 
   /**
    * @brief Get the group labels corresponding to the sorted order of `keys`.
@@ -171,7 +170,7 @@ struct sort_groupby_helper {
    *
    * @return vector of group labels for each row in the sorted key column
    */
-  index_vector const& group_labels(rmm::cuda_stream_view stream);
+  index_vector const& group_labels(cuda::stream_ref stream);
 
  private:
   /**
@@ -188,7 +187,7 @@ struct sort_groupby_helper {
    * @return A nullable column of `INT32` containing group labels in the order
    *         of the unsorted key table
    */
-  column_view unsorted_keys_labels(rmm::cuda_stream_view stream);
+  column_view unsorted_keys_labels(cuda::stream_ref stream);
 
   /**
    * @brief Get the column representing the row bitmask for the `keys`
@@ -202,7 +201,7 @@ struct sort_groupby_helper {
    * Computes and stores bitmask on first invocation and returns stored column
    * on subsequent calls.
    */
-  column_view keys_bitmask_column(rmm::cuda_stream_view stream);
+  column_view keys_bitmask_column(cuda::stream_ref stream);
 
   column_ptr _key_sorted_order;      ///< Indices to produce _keys in sorted order
   column_ptr _unsorted_keys_labels;  ///< Group labels for unsorted _keys