[FEA] Make `cuda.core.Stream` recognized by numba-cuda by supporting the `__cuda_stream__` protocol

[leofang]

Is your feature request related to a problem? Please describe.

I want to use cuda.core.Stream for

• kernel launch
• data copy
• ...

xref: https://nvidia.github.io/cuda-python/cuda-core/latest/interoperability.html#cuda-stream-protocol

Describe the solution you'd like

Describe alternatives you've considered

N/A

Additional context

N/A

[brandon-b-miller]

If all we want is what's described in the OP, the following code

from numba import cuda
from cuda.core.experimental import Stream
from cuda.core.experimental import Device

dev = Device()
dev.set_current()
stream = dev.create_stream()
stream2 = dev.create_stream()

@cuda.jit
def kernel():
    pass


kernel[1,1, stream]()
kernel[1,1, stream2]()

devary = cuda.to_device([1, 2, 3, 4, 5], stream=stream2)
print(devary.copy_to_host(stream=stream2))

[1 2 3 4 5] 

Works with just this patch:

Details

diff --git a/numba_cuda/numba/cuda/cudadrv/driver.py b/numba_cuda/numba/cuda/cudadrv/driver.py
index 6d3f307..b5a7c34 100644
--- a/numba_cuda/numba/cuda/cudadrv/driver.py
+++ b/numba_cuda/numba/cuda/cudadrv/driver.py
@@ -71,6 +71,7 @@ USE_NV_BINDING = config.CUDA_USE_NVIDIA_BINDING
 if USE_NV_BINDING:
     from cuda.bindings import driver as binding
 
+
     # There is no definition of the default stream in the Nvidia bindings (nor
     # is there at the C/C++ level), so we define it here so we don't need to
     # use a magic number 0 in places where we want the default stream.
@@ -3504,7 +3505,8 @@ def host_to_device(dst, src, size, stream=0):
     varargs = []
 
     if stream:
-        assert isinstance(stream, Stream)
+        if not USE_NV_BINDING:
+            assert isinstance(stream, Stream)
         fn = driver.cuMemcpyHtoDAsync
         varargs.append(stream.handle)
     else:
@@ -3522,7 +3524,8 @@ def device_to_host(dst, src, size, stream=0):
     varargs = []
 
     if stream:
-        assert isinstance(stream, Stream)
+        if not USE_NV_BINDING:
+            assert isinstance(stream, Stream)
         fn = driver.cuMemcpyDtoHAsync
         varargs.append(stream.handle)
     else:
@@ -3540,7 +3543,8 @@ def device_to_device(dst, src, size, stream=0):
     varargs = []
 
     if stream:
-        assert isinstance(stream, Stream)
+        if not USE_NV_BINDING:
+            assert isinstance(stream, Stream)
         fn = driver.cuMemcpyDtoDAsync
         varargs.append(stream.handle)
     else:
@@ -3561,7 +3565,8 @@ def device_memset(dst, val, size, stream=0):
     varargs = []
 
     if stream:
-        assert isinstance(stream, Stream)
+        if not USE_NV_BINDING:
+            assert isinstance(stream, Stream)
         fn = driver.cuMemsetD8Async
         varargs.append(stream.handle)
     else:

But I think that like with the Linker, the larger goal here would be for numba-cuda to use cuda.core.experimental.Stream entirely for its internals. The best thing would be so that this worked and passed all tests:

if config.CUDA_USE_NVIDIA_BINDING
    from cuda.core.experimental import Stream
else:
    class Stream...

Later on we could hopefully drop our own stream class entirely and there'd be one place (cuda-python) for python users to look for a stream object.

[rparolin]

blocked by #401