Use cudaStream_t instead of cuda_stream_view in pylibcudf Cython (#22368)

Contributes to rapidsai/rmm#2359

Authors:
  - Vyas Ramasubramani (https://github.com/vyasr)

Approvers:
  - Bradley Dice (https://github.com/bdice)

URL: #22368

Author: vyasr

python/cudf_polars/cudf_polars/utils/cuda_stream.py

@@ -13,6 +13,7 @@
 if TYPE_CHECKING:
     from collections.abc import Callable, Sequence
 
+    from pylibcudf.utils import CudaStreamLike
     from rmm.pylibrmm.stream import Stream
 
 
@@ -27,7 +28,7 @@ def get_cuda_stream() -> Stream:
 
 
 def join_cuda_streams(
-    *, downstreams: Sequence[Stream], upstreams: Sequence[Stream]
+    *, downstreams: Sequence[CudaStreamLike], upstreams: Sequence[CudaStreamLike]
 ) -> None:
     """
     Join multiple CUDA streams.
@@ -46,7 +47,7 @@ def join_cuda_streams(
 
 
 def get_joined_cuda_stream(
-    get_cuda_stream: Callable[[], Stream], *, upstreams: Sequence[Stream]
+    get_cuda_stream: Callable[[], Stream], *, upstreams: Sequence[CudaStreamLike]
 ) -> Stream:
     """
     Return a CUDA stream that is joined to the given streams.

python/pylibcudf/pylibcudf/binaryop.pxd

@@ -1,10 +1,9 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 
 from libcpp cimport bool
 from pylibcudf.libcudf.binaryop cimport binary_operator
 from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource
-from rmm.pylibrmm.stream cimport Stream
 
 from .column cimport Column
 from .scalar cimport Scalar
@@ -25,7 +24,7 @@ cpdef Column binary_operation(
     RightBinaryOperand rhs,
     binary_operator op,
     DataType output_type,
-    Stream stream=*,
+    object stream = *,
     DeviceMemoryResource mr=*,
 )
 

python/pylibcudf/pylibcudf/binaryop.pyi

@@ -1,14 +1,14 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 
 from enum import IntEnum
 
 from rmm.pylibrmm.memory_resource import DeviceMemoryResource
-from rmm.pylibrmm.stream import Stream
 
 from pylibcudf.column import Column
 from pylibcudf.scalar import Scalar
 from pylibcudf.types import DataType
+from pylibcudf.utils import CudaStreamLike
 
 class BinaryOperator(IntEnum):
     ADD = ...
@@ -52,7 +52,7 @@ def binary_operation(
     rhs: Column | Scalar,
     op: BinaryOperator,
     output_type: DataType,
-    stream: Stream | None = None,
+    stream: CudaStreamLike | None = None,
     mr: DeviceMemoryResource | None = None,
 ) -> Column: ...
 def is_supported_operation(

python/pylibcudf/pylibcudf/binaryop.pyx

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 
 from cython.operator import dereference
@@ -20,6 +20,7 @@ from .column cimport Column
 from .scalar cimport Scalar
 from .types cimport DataType
 from .utils cimport _get_stream, _get_memory_resource
+from cuda.bindings.cyruntime cimport cudaStream_t
 
 __all__ = ["BinaryOperator", "binary_operation", "is_supported_operation"]
 
@@ -28,7 +29,7 @@ cpdef Column binary_operation(
     RightBinaryOperand rhs,
     binary_operator op,
     DataType output_type,
-    Stream stream=None,
+    object stream=None,
     DeviceMemoryResource mr=None,
 ):
     """Perform a binary operation between a column and another column or scalar.
@@ -61,7 +62,8 @@ cpdef Column binary_operation(
         The result of the binary operation
     """
     cdef unique_ptr[column] result
-    stream = _get_stream(stream)
+    cdef Stream _stream = _get_stream(stream)
+    cdef cudaStream_t _cs = _stream.view().value()
     mr = _get_memory_resource(mr)
 
     if LeftBinaryOperand is Column and RightBinaryOperand is Column:
@@ -71,7 +73,7 @@ cpdef Column binary_operation(
                 rhs.view(),
                 op,
                 output_type.c_obj,
-                stream.view(),
+                _cs,
                 mr.get_mr()
             )
     elif LeftBinaryOperand is Column and RightBinaryOperand is Scalar:
@@ -81,7 +83,7 @@ cpdef Column binary_operation(
                 dereference(rhs.c_obj),
                 op,
                 output_type.c_obj,
-                stream.view(),
+                _cs,
                 mr.get_mr()
             )
     elif LeftBinaryOperand is Scalar and RightBinaryOperand is Column:
@@ -91,13 +93,13 @@ cpdef Column binary_operation(
                 rhs.view(),
                 op,
                 output_type.c_obj,
-                stream.view(),
+                _cs,
                 mr.get_mr()
             )
     else:
         raise ValueError(f"Invalid arguments {lhs} and {rhs}")
 
-    return Column.from_libcudf(move(result), stream, mr)
+    return Column.from_libcudf(move(result), _stream, mr)
 
 
 cpdef bool is_supported_operation(

python/pylibcudf/pylibcudf/column.pxd

@@ -6,7 +6,6 @@ from libcpp.vector cimport vector
 from libc.stdint cimport uint64_t
 
 from rmm.librmm.device_buffer cimport device_buffer
-from rmm.pylibrmm.stream cimport Stream
 from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource
 from pylibcudf.libcudf.column.column cimport column
 from pylibcudf.libcudf.column.column_view cimport (
@@ -27,7 +26,7 @@ cdef class OwnerWithCAI:
     cdef dict cai
 
     @staticmethod
-    cdef create(column_view cv, object owner, Stream stream)
+    cdef create(column_view cv, object owner, object stream)
 
 
 cdef class OwnerMaskWithCAI:
@@ -38,7 +37,7 @@ cdef class OwnerMaskWithCAI:
     cdef create(column_view cv, object owner)
 
 
-cdef gpumemoryview _copy_array_to_device(object buf, Stream stream=*)
+cdef gpumemoryview _copy_array_to_device(object buf, object stream=*)
 
 
 cdef class Column:
@@ -61,7 +60,7 @@ cdef class Column:
     @staticmethod
     cdef Column from_libcudf(
         unique_ptr[column] libcudf_col,
-        Stream stream,
+        object stream,
         DeviceMemoryResource mr
     )
 
@@ -72,7 +71,7 @@ cdef class Column:
     cdef Column from_column_view_of_arbitrary(
         const column_view& cv,
         object owner,
-        Stream stream,
+        object stream,
     )
 
     @staticmethod
@@ -81,10 +80,10 @@ cdef class Column:
         tuple shape,
         DataType dtype,
         Column base=*,
-        Stream stream=*,
+        object stream=*,
     )
 
-    cpdef Scalar to_scalar(self, Stream stream=*, DeviceMemoryResource mr=*)
+    cpdef Scalar to_scalar(self, object stream=*, DeviceMemoryResource mr=*)
     cpdef DataType type(self)
     cpdef Column child(self, size_type index)
     cpdef size_type num_children(self)
@@ -95,7 +94,7 @@ cdef class Column:
     cpdef object data(self)
     cpdef object null_mask(self)
     cpdef list children(self)
-    cpdef Column copy(self, Stream stream=*, DeviceMemoryResource mr=*)
+    cpdef Column copy(self, object stream=*, DeviceMemoryResource mr=*)
     cpdef uint64_t device_buffer_size(self)
     cpdef Column with_mask(self, object, size_type, bint validate=*)
 
@@ -108,10 +107,10 @@ cdef class ListsColumnView:
     cpdef child(self)
     cpdef offsets(self)
     cdef lists_column_view view(self) nogil
-    cpdef Column get_sliced_child(self, Stream stream=*)
+    cpdef Column get_sliced_child(self, object stream=*)
 
 
 cdef class StructsColumnView:
     cdef Column _column
     cdef structs_column_view view(self) nogil
-    cpdef Column get_sliced_child(self, int index, Stream stream=*)
+    cpdef Column get_sliced_child(self, int index, object stream=*)

python/pylibcudf/pylibcudf/column.pyi

@@ -6,12 +6,12 @@ from typing import Any, Protocol, TypedDict
 
 from rmm.pylibrmm.device_buffer import DeviceBuffer
 from rmm.pylibrmm.memory_resource import DeviceMemoryResource
-from rmm.pylibrmm.stream import Stream
 
 from pylibcudf._interop_helpers import ArrowLike, ColumnMetadata
 from pylibcudf.scalar import Scalar
 from pylibcudf.span import Span
 from pylibcudf.types import DataType
+from pylibcudf.utils import CudaStreamLike
 
 class ArrayInterfaceBase(TypedDict):
     shape: tuple[int, ...]
@@ -64,7 +64,7 @@ class Column:
     def num_children(self) -> int: ...
     def copy(
         self,
-        stream: Stream | None = None,
+        stream: CudaStreamLike | None = None,
         mr: DeviceMemoryResource | None = None,
     ) -> Column: ...
     def device_buffer_size(self) -> int: ...
@@ -77,19 +77,19 @@ class Column:
     def from_scalar(
         scalar: Scalar,
         size: int,
-        stream: Stream | None = None,
+        stream: CudaStreamLike | None = None,
         mr: DeviceMemoryResource | None = None,
     ) -> Column: ...
     def to_scalar(
         self,
-        stream: Stream | None = None,
+        stream: CudaStreamLike | None = None,
         mr: DeviceMemoryResource | None = None,
     ) -> Scalar: ...
     @staticmethod
     def all_null_like(
         like: Column,
         size: int,
-        stream: Stream | None = None,
+        stream: CudaStreamLike | None = None,
         mr: DeviceMemoryResource | None = None,
     ) -> Column: ...
     @staticmethod
@@ -99,54 +99,58 @@ class Column:
     def to_arrow(
         self,
         metadata: ColumnMetadata | str | None = None,
-        stream: Stream | None = None,
+        stream: CudaStreamLike | None = None,
     ) -> ArrowLike: ...
     # Private methods below are included because polars is currently using them,
     # but we want to remove stubs for these private methods eventually
     def _to_schema(self, metadata: Any = None) -> Any: ...
-    def _to_host_array(self, stream: Stream) -> Any: ...
+    def _to_host_array(self, stream: CudaStreamLike) -> Any: ...
     @staticmethod
     def from_arrow(
         obj: ArrowLike,
         dtype: DataType | None = None,
-        stream: Stream | None = None,
+        stream: CudaStreamLike | None = None,
         mr: DeviceMemoryResource | None = None,
     ) -> Column: ...
     @classmethod
     def from_cuda_array_interface(
-        cls, obj: SupportsCudaArrayInterface, stream: Stream | None = None
+        cls,
+        obj: SupportsCudaArrayInterface,
+        stream: CudaStreamLike | None = None,
     ) -> Column: ...
     @classmethod
     def from_array_interface(
-        cls, obj: SupportsArrayInterface, stream: Stream | None = None
+        cls, obj: SupportsArrayInterface, stream: CudaStreamLike | None = None
     ) -> Column: ...
     @classmethod
     def from_array(
         cls,
         obj: SupportsCudaArrayInterface | SupportsArrayInterface,
-        stream: Stream | None = None,
+        stream: CudaStreamLike | None = None,
     ) -> Column: ...
     @staticmethod
     def struct_from_children(children: Sequence[Column]) -> Column: ...
     @staticmethod
     def from_iterable_of_py(
         obj: Iterable,
         dtype: DataType | None = None,
-        stream: Stream | None = None,
+        stream: CudaStreamLike | None = None,
     ) -> Column: ...
 
 class ListsColumnView:
     def __init__(self, column: Column): ...
     def child(self) -> Column: ...
     def offsets(self) -> Column: ...
-    def get_sliced_child(self, stream: Stream | None = None) -> Column: ...
+    def get_sliced_child(
+        self, stream: CudaStreamLike | None = None
+    ) -> Column: ...
 
 class StructsColumnView:
     def __init__(self, column: Column): ...
     def child(self) -> Column: ...
     def offsets(self) -> Column: ...
     def get_sliced_child(
-        self, index: int, stream: Stream | None = None
+        self, index: int, stream: CudaStreamLike | None = None
     ) -> Column: ...
 
 def is_c_contiguous(