[FEA] Migrate RAPIDS to CCCL `cuda::stream_ref` / `cuda::stream` in C++ and `cuda.core` streams in Python

[bdice]

RMM currently provides the "vocabulary type" for CUDA streams in RAPIDS. This has been used to provide C++ and Cython/Python wrappers for cudaStream_t for some time. However, now we have CCCL's cuda::stream_ref and cuda::stream, and we want to rely on CCCL to provide this vocabulary type. Moreover, we would like for cuda.core to be the sole source of this vocabulary type for the Python ecosystem, providing Cython/Python wrappers for CCCL stream types. See prior discussion in #2040 on why parity/alignment between CCCL cuda::stream and cuda.core streams is important.

I discussed this with @brandon-b-miller today. It seems like the C++ side of this should be a straightforward find-replace in many parts of RAPIDS, while the Cython/Python side may take a little more effort to integrate with @Andy-Jost's work on resource handles in cuda.core.

This migration would have additional benefits: starting with streams will prepare us to use CCCL and cuda.core as a source of memory resource definitions/bindings. This aims towards the ultimate goal of memory resources being a core part of the CUDA platform (both C++ and Python) with RMM providing a stripped-down library of memory resource "extras" like adaptors, specialized resources, and functionality specific to RAPIDS.

References:

• https://nvidia.github.io/cccl/libcudacxx/api/classstream__ref.html#_CPPv410stream_ref
• https://nvidia.github.io/cccl/libcudacxx/api/structstream.html#structstream
• [FEA]: Migrate to using CCCL-Runtime NVIDIA/cuda-python#1598
• Turn all cdef class objects in the Python layer to cdef public class NVIDIA/cuda-python#464

Proposed Tasks

C++

Concretely, we want to replace rmm::cuda_stream with cuda::stream and rmm::cuda_stream_view with cuda::stream_ref. Other stream-related classes like rmm::cuda_stream_pool can stay in RMM for now. Note: There may be subtle differences in handling of the legacy default stream that we need to consider.

I think there are no blockers, this is ready to work on.

Python

First, we should work on #1782 so we only have one module to deprecate/replace. Then we can start replacing rmm.pylibrmm.cuda_stream.CudaStream with the cuda.core equivalent.

If we succeed at migrating the C++ side of things first and don't have bindings in cuda.core yet, we can add Cython bindings for the CCCL stream types in RMM just to keep things moving forward towards deprecation of the RMM stream types in C++.

[wence-]

Concretely, we want to replace rmm::cuda_stream with cuda::stream and rmm::cuda_stream_view with cuda::stream_ref. Other stream-related classes like rmm::cuda_stream_pool can stay in RMM for now. Note: There may be subtle differences in handling of the legacy default stream that we need to consider.

I think there are no blockers, this is ready to work on.

It is not quite that trivial, I think. The libcudacxx cuda:: types are device-explicit https://nvidia.github.io/cccl/libcudacxx/runtime/cudart_interactions.html#device-selection and the ctor needs a device_ref

In addition, use of the legacy default stream via a cuda::stream_ref requires that the calling code call cudaSetDevice to set an active device first.

That is, one cannot blindly do:

// These don't work, because there is no active device.
auto ptds = cuda::stream_ref{cudaStreamPerThread};
auto default_stream = cuda::stream_ref{cudaStream_t{nullptr}};

We can polyfill with a stream creation function rmm::create_stream() that does cudaGetDevice under the hood, but I don't think that addresses the broader question.

If we want to move to CCCL and/or core.core owning vocabulary types we need to decide how to manage device state. Both of the libraries want you to be explicit in your use of devices (e.g. https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Stream.html#cuda.core.Stream).

I agree that migrating rmm::cuda_stream_view to cuda::stream_ref ought to be trivial. We can do it first by just accepting cuda::stream_ref in the APIs instead (since there's an implicit conversion from stream_view)