Add join_streams to pylibcudf API (#20316)

This adds a `plc.experimental.join_streams` API to pylibcudf. It wraps `cudf::detail::join_streams`. We've used the `.experimental` namespace in part to reflect the fact that the C++ implementation we're wrapping is in a `detail` namespace.

This change will support our work on making cudf-polars use non-default streams for operations, which is important for performance with the in-progress rapidsmpf streaming network runtime. We only need the ability to join some stream to a sequence of upstream streams, so only `join_streams` is wrapped, nothing to do with `stream_pools` that are also provided in that module.

Closes #20315

Authors:
  - Tom Augspurger (https://github.com/TomAugspurger)

Approvers:
  - Lawrence Mitchell (https://github.com/wence-)
  - Matthew Murray (https://github.com/Matt711)
  - Bradley Dice (https://github.com/bdice)

URL: #20316

Author: TomAugspurger

docs/cudf/source/pylibcudf/api_docs/experimental/index.rst

@@ -0,0 +1,8 @@
+============
+Experimental
+============
+
+APIs in this namespace are experimental and may change without warning in the future.
+
+.. automodule:: pylibcudf.experimental
+   :members:

docs/cudf/source/pylibcudf/api_docs/index.rst

@@ -51,6 +51,7 @@ This page provides API documentation for pylibcudf.
     :maxdepth: 2
     :caption: Subpackages
 
+    experimental/index.rst
     io/index.rst
     strings/index.rst
     nvtext/index.rst

python/pylibcudf/pylibcudf/CMakeLists.txt

@@ -67,3 +67,4 @@ add_subdirectory(libcudf)
 add_subdirectory(strings)
 add_subdirectory(io)
 add_subdirectory(nvtext)
+add_subdirectory(experimental)

python/pylibcudf/pylibcudf/__init__.py

@@ -19,6 +19,7 @@
     contiguous_split,
     copying,
     datetime,
+    experimental,
     expressions,
     filling,
     groupby,
@@ -72,6 +73,7 @@
     "contiguous_split",
     "copying",
     "datetime",
+    "experimental",
     "expressions",
     "filling",
     "gpumemoryview",

python/pylibcudf/pylibcudf/experimental/CMakeLists.txt

@@ -0,0 +1,16 @@
+# =============================================================================
+# cmake-format: off
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION.
+# SPDX-License-Identifier: Apache-2.0
+# cmake-format: on
+# =============================================================================
+
+set(cython_sources _join_streams.pyx)
+
+set(linked_libraries cudf::cudf)
+rapids_cython_create_modules(
+  CXX
+  SOURCE_FILES "${cython_sources}"
+  LINKED_LIBRARIES "${linked_libraries}" MODULE_PREFIX pylibcudf_experimental_ ASSOCIATED_TARGETS
+                   cudf
+)

python/pylibcudf/pylibcudf/experimental/__init__.pxd

@@ -0,0 +1,4 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
+# SPDX-License-Identifier: Apache-2.0
+
+from ._join_streams cimport join_streams

python/pylibcudf/pylibcudf/experimental/__init__.py

@@ -0,0 +1,8 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
+# SPDX-License-Identifier: Apache-2.0
+
+from ._join_streams import join_streams
+
+__all__ = [
+    "join_streams",
+]

python/pylibcudf/pylibcudf/experimental/_join_streams.pxd

@@ -0,0 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
+# SPDX-License-Identifier: Apache-2.0
+
+from rmm.pylibrmm.stream cimport Stream
+
+cpdef void join_streams(list streams, Stream stream)

python/pylibcudf/pylibcudf/experimental/_join_streams.pyi

@@ -0,0 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
+# SPDX-License-Identifier: Apache-2.0
+
+from rmm.pylibrmm.stream import Stream
+
+def join_streams(streams: list[Stream], stream: Stream) -> None: ...

python/pylibcudf/pylibcudf/experimental/_join_streams.pyx

@@ -0,0 +1,56 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
+# SPDX-License-Identifier: Apache-2.0
+
+from libcpp.vector cimport vector
+
+from pylibcudf.libcudf.detail.utilities cimport stream_pool as cpp_stream_pool
+from pylibcudf.libcudf.utilities.span cimport host_span
+
+from rmm.librmm.cuda_stream_view cimport cuda_stream_view
+from rmm.pylibrmm.stream cimport Stream
+
+ctypedef const cuda_stream_view const_cuda_stream_view
+
+
+__all__ = ["join_streams"]
+
+
+cpdef void join_streams(list streams, Stream stream):
+    """Synchronize a stream to an event on a set of streams.
+
+    This function synchronizes the joined stream with the waited-on streams
+    by placing events on each of the waited-on streams and having the joined
+    stream wait on those events.
+
+    Parameters
+    ----------
+    streams : list
+        A list of Stream objects to wait on.
+    stream : Stream
+        The joined stream that synchronizes with the waited-on streams.
+
+    Examples
+    --------
+    >>> import pylibcudf as plc
+    >>> from rmm.pylibrmm.stream import Stream
+    >>> # Create streams
+    >>> stream1 = Stream()
+    >>> stream2 = Stream()
+    >>> join_stream = Stream()
+    >>> # ... do work on stream1 and stream2 ...
+    >>> # Wait for both streams before continuing work on join_stream
+    >>> plc.experimental.join_streams([stream1, stream2], join_stream)
+    >>> # ... continue work on join_stream ...
+    """
+    cdef Stream c_stream = <Stream?>stream
+    cdef vector[cuda_stream_view] c_streams
+
+    c_streams.reserve(len(streams))
+    for s in streams:
+        c_streams.push_back((<Stream?>s).view())
+
+    with nogil:
+        cpp_stream_pool.join_streams(
+            host_span[const_cuda_stream_view](c_streams.data(), c_streams.size()),
+            c_stream.view()
+        )