Repository: cupy/cupy
Pull request: add `cudaStreamCreateWithPriority`
This PR aims to close #8210 by adding the `cudaStreamCreateWithPriority` API. 

Problem statement:
Support Cuda Stream creation with Priority
### Description

Hi,

It appears that CuPy does not support the creation of custom CUDA streams with priority. Since the API for this functionality is already available in CUDA, it would be very helpful if CuPy provided this feature.

Thanks

### Additional Information

_No response_

Hints from issue discussion:
As a workaround, you can access the `cudaStreamCreateWithPriority` API through CUDA Python (need to install it yourself via `pip` or `conda`)
https://nvidia.github.io/cuda-python/module/cudart.html?highlight=cudastreamcreatewithpriority#cuda.cudart.cudaStreamCreateWithPriority
and wrap the returned pointer with `ExternalStream`
https://docs.cupy.dev/en/stable/reference/generated/cupy.cuda.ExternalStream.html#cupy.cuda.ExternalStream
so that you can use it as a CuPy stream.
Hello, I'm drafting a fix for this issue, and I have a question. I've added a cython wrapper for the runtime API [`cudaStreamCreateWithPriority`](https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__STREAM.html#group__CUDART__STREAM_1ge2be9e9858849bf62ba4a8b66d1c3540), by making changes to the` cupy_backends/cuda/api/runtime.pyx`, `cupy_backends/cuda/api/runtime.pxd` and `cupy_backends/cuda/api/_runtime_extern.pxi` files. 

Would I also have to modify the following class's `__init__` method in `cupy/cuda/stream.pyx` to take `priority` as a parameter and then make a call to `runtime.streamCreateWithPriority` ? 

https://github.com/cupy/cupy/blob/20ccd63c0acc40969c851b1917dedeb032209e8b/cupy/cuda/stream.pyx#L427


Please let me know if I should create a PR and move the discussion there instead of here.
Sure, that sounds reasonable. I'd suggest to set `priority=None` as default, meaning we don't care. In terms implementation, you can then fall back to the old path if it's `None`. Do pay attention to the interaction between `priority` and `blocking`.

Commit message: fixed typo

Patch to review:
diff --git a/cupy/cuda/stream.pyx b/cupy/cuda/stream.pyx
index 96e34e6618f..73cd7c6f416 100644
--- a/cupy/cuda/stream.pyx
+++ b/cupy/cuda/stream.pyx
@@ -423,6 +423,38 @@ class _BaseStream:
         except RuntimeError:  # can be RuntimeError or CUDARuntimeError
             raise
 
+    def get_stream_flags(self):
+        """Query the flags of a stream.
+
+        Returns:
+            bool:
+                The return value is True, if the stream is non_blocking,
+                else False is returned to indicate the default stream
+                creation flag.
+        """
+        cdef unsigned int flags
+        flags = runtime.streamGetFlags(self.ptr)
+        if flags == runtime.streamNonBlocking:
+            return True
+        return False
+
+    def get_stream_priority(self):
+        """Query the priority of a stream.
+
+        Returns:
+            signed int:
+                  Lower numbers represent higher priorities.
+
+        .. seealso:: `cudaStreamGetPriority()`_
+
+        .. _cudaStreamGetPriority():
+            https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__STREAM.html#group__CUDART__STREAM_1g192bb727d15c4407c119747de7d198a6
+
+        """
+        cdef int priority
+        priority = runtime.streamGetPriority(self.ptr)
+        return priority
+
 
 class Stream(_BaseStream):
 
@@ -445,6 +477,8 @@ class Stream(_BaseStream):
             per-thread default stream object.
         non_blocking (bool): If ``True`` and both ``null`` and ``ptds`` are
             ``False``, the stream does not synchronize with the NULL stream.
+        priority (int): Priority of the stream. Lower numbers represent higher
+            priorities.
 
     Attributes:
         ~Stream.ptr (intptr_t): Raw stream handle.
@@ -457,7 +491,8 @@ class Stream(_BaseStream):
     null = None
     ptds = None
 
-    def __init__(self, null=False, non_blocking=False, ptds=False):
+    def __init__(self, null=False, non_blocking=False, ptds=False,
+                 priority=None):
         if null:
             # TODO(pentschev): move to streamLegacy. This wasn't possible
             # because of a NCCL bug that should be fixed in the version
@@ -470,11 +505,14 @@ class Stream(_BaseStream):
                                  'default stream (ptds)')
             ptr = runtime.streamPerThread
             device_id = -1
-        elif non_blocking:
-            ptr = runtime.streamCreateWithFlags(runtime.streamNonBlocking)
-            device_id = runtime.getDevice()
         else:
-            ptr = runtime.streamCreate()
+            if priority is None:
+                priority = 0  # default
+            if non_blocking:
+                flag = runtime.streamNonBlocking
+            else:
+                flag = runtime.streamDefault
+            ptr = runtime.streamCreateWithPriority(flag, priority)
             device_id = runtime.getDevice()
         super().__init__(ptr, device_id)
 
diff --git a/cupy_backends/cuda/api/_runtime_extern.pxi b/cupy_backends/cuda/api/_runtime_extern.pxi
index d95e5629999..62fdac30f1e 100644
--- a/cupy_backends/cuda/api/_runtime_extern.pxi
+++ b/cupy_backends/cuda/api/_runtime_extern.pxi
@@ -110,6 +110,10 @@ cdef extern from '../../cupy_backend_runtime.h' nogil:
     int cudaStreamCreate(driver.Stream* pStream)
     int cudaStreamCreateWithFlags(driver.Stream* pStream,
                                   unsigned int flags)
+    int cudaStreamCreateWithPriority(driver.Stream* pStream,
+                                     unsigned int flags, int priority)
+    int cudaStreamGetFlags(driver.Stream pStream, unsigned int* flags)
+    int cudaStreamGetPriority(driver.Stream pStream, int* priority)
     int cudaStreamDestroy(driver.Stream stream)
     int cudaStreamSynchronize(driver.Stream stream)
     int cudaStreamAddCallback(driver.Stream stream, StreamCallback callback,
diff --git a/cupy_backends/cuda/api/runtime.pxd b/cupy_backends/cuda/api/runtime.pxd
index fccc6f6560a..603e17554a3 100644
--- a/cupy_backends/cuda/api/runtime.pxd
+++ b/cupy_backends/cuda/api/runtime.pxd
@@ -266,6 +266,10 @@ cpdef memPoolSetAttribute(intptr_t, int, object)
 
 cpdef intptr_t streamCreate() except? 0
 cpdef intptr_t streamCreateWithFlags(unsigned int flags) except? 0
+cpdef intptr_t streamCreateWithPriority(unsigned int flags,
+                                        int priority) except? 0
+cpdef unsigned int streamGetFlags(intptr_t stream) except? 0
+cpdef int streamGetPriority(intptr_t stream) except? 0
 cpdef streamDestroy(intptr_t stream)
 cpdef streamSynchronize(intptr_t stream)
 cpdef streamAddCallback(intptr_t stream, callback, intptr_t arg,
diff --git a/cupy_backends/cuda/api/runtime.pyx b/cupy_backends/cuda/api/runtime.pyx
index 2eff2004883..58cb5af54fe 100644
--- a/cupy_backends/cuda/api/runtime.pyx
+++ b/cupy_backends/cuda/api/runtime.pyx
@@ -839,6 +839,28 @@ cpdef intptr_t streamCreateWithFlags(unsigned int flags) except? 0:
     return <intptr_t>stream
 
 
+cpdef intptr_t streamCreateWithPriority(unsigned int flags,
+                                        int priority) except? 0:
+    cdef driver.Stream stream
+    status = cudaStreamCreateWithPriority(&stream, flags, priority)
+    check_status(status)
+    return <intptr_t>stream
+
+
+cpdef unsigned int streamGetFlags(intptr_t stream) except? 0:
+    cdef unsigned int flags
+    status = cudaStreamGetFlags(<driver.Stream>stream, &flags)
+    check_status(status)
+    return flags
+
+
+cpdef int streamGetPriority(intptr_t stream) except? 0:
+    cdef int priority
+    status = cudaStreamGetPriority(<driver.Stream>stream, &priority)
+    check_status(status)
+    return priority
+
+
 cpdef streamDestroy(intptr_t stream):
     status = cudaStreamDestroy(<driver.Stream>stream)
     check_status(status)
diff --git a/cupy_backends/hip/cupy_hip_runtime.h b/cupy_backends/hip/cupy_hip_runtime.h
index 51a54323ebc..3a1e4efd910 100644
--- a/cupy_backends/hip/cupy_hip_runtime.h
+++ b/cupy_backends/hip/cupy_hip_runtime.h
@@ -345,6 +345,20 @@ cudaError_t cudaStreamCreateWithFlags(cudaStream_t *stream,
     return hipStreamCreateWithFlags(stream, flags);
 }
 
+cudaError_t cudaStreamCreateWithPriority(cudaStream_t *stream,
+                                         unsigned int flags,
+                                         int priority) {
+    return hipStreamCreateWithPriority(stream, flags, priority);
+}
+
+cudaError_t cudaStreamGetFlags(cudaStream_t stream, unsigned int *flags) {
+    return hipStreamGetFlags(stream, flags);
+}
+
+cudaError_t cudaStreamGetPriority(cudaStream_t stream, int *priority) {
+    return hipStreamGetPriority(stream, priority);
+}
+
 cudaError_t cudaStreamDestroy(cudaStream_t stream) {
     return hipStreamDestroy(stream);
 }
diff --git a/cupy_backends/stub/cupy_cuda_runtime.h b/cupy_backends/stub/cupy_cuda_runtime.h
index 6cde83108b3..d42308b0490 100644
--- a/cupy_backends/stub/cupy_cuda_runtime.h
+++ b/cupy_backends/stub/cupy_cuda_runtime.h
@@ -292,6 +292,18 @@ cudaError_t cudaStreamCreateWithFlags(...) {
     return cudaSuccess;
 }
 
+cudaError_t cudaStreamCreateWithPriority(...) {
+    return cudaSuccess;
+}
+
+cudaError_t cudaStreamGetFlags(...) {
+    return cudaSuccess;
+}
+
+cudaError_t cudaStreamGetPriority(...) {
+    return cudaSuccess;
+}
+
 cudaError_t cudaStreamDestroy(...) {
     return cudaSuccess;
 }
diff --git a/docs/source/reference/cuda.rst b/docs/source/reference/cuda.rst
index fb12351b7cb..32a276df5f4 100644
--- a/docs/source/reference/cuda.rst
+++ b/docs/source/reference/cuda.rst
@@ -193,6 +193,7 @@ to use these functions.
    cupy.cuda.runtime.pointerGetAttributes
    cupy.cuda.runtime.streamCreate
    cupy.cuda.runtime.streamCreateWithFlags
+   cupy.cuda.runtime.streamCreateWithPriority
    cupy.cuda.runtime.streamDestroy
    cupy.cuda.runtime.streamSynchronize
    cupy.cuda.runtime.streamAddCallback
diff --git a/tests/cupy_tests/cuda_tests/test_stream.py b/tests/cupy_tests/cuda_tests/test_stream.py
index 93fce5ca8de..f31d75f2c93 100644
--- a/tests/cupy_tests/cuda_tests/test_stream.py
+++ b/tests/cupy_tests/cuda_tests/test_stream.py
@@ -243,6 +243,23 @@ def f2(barrier, errors):
         for err in errors:
             assert err is False
 
+    def test_create_with_flags(self):
+        s1 = cuda.Stream()
+        s2 = cuda.Stream(non_blocking=True)
+        assert s1.get_stream_flags() is False
+        assert s2.get_stream_flags() is True
+
+    def test_create_with_priority(self):
+        # parameterize wasn't used since priority gets
+        # clamped when it isn't initialized within a specific
+        # returned by `cudaDeviceGetStreamPriorityRange`.
+        s1 = cuda.Stream(priority=0)
+        s2 = cuda.Stream(priority=-1)
+        s3 = cuda.Stream(priority=-3)
+        assert s1.get_stream_priority() == 0
+        assert s2.get_stream_priority() == -1
+        assert s3.get_stream_priority() == -3
+
 
 class TestExternalStream(unittest.TestCase):
 


Code review comments:
diff_hunk: @@ -423,6 +423,38 @@ class _BaseStream:
         except RuntimeError:  # can be RuntimeError or CUDARuntimeError
             raise
 
+    def get_stream_flags(self):
+        """Query the flags of a stream.
+
+        Returns:
+            bool:
+                The return value is True, if the stream is non_blocking,
+                else False is returned to indicate the default stream
+                creation flag.
+        """
+        cdef unsigned int flags
+        flags = runtime.streamGetFlags(self.ptr)
+        if flags == runtime.streamNonBlocking:
+            return True
+        return False
path: cupy/cuda/stream.pyx
text: @user1:
```suggestion
        return flags == runtime.streamNonBlocking
```

diff_hunk: @@ -423,6 +423,38 @@ class _BaseStream:
         except RuntimeError:  # can be RuntimeError or CUDARuntimeError
             raise
 
+    def get_stream_flags(self):
+        """Query the flags of a stream.
+
+        Returns:
+            bool:
+                The return value is True, if the stream is non_blocking,
+                else False is returned to indicate the default stream
+                creation flag.
+        """
+        cdef unsigned int flags
+        flags = runtime.streamGetFlags(self.ptr)
+        if flags == runtime.streamNonBlocking:
+            return True
+        return False
+
+    def get_stream_priority(self):
path: cupy/cuda/stream.pyx
text: @user1:
ditto (making it a property rather than a method)

diff_hunk: @@ -423,6 +423,38 @@ class _BaseStream:
         except RuntimeError:  # can be RuntimeError or CUDARuntimeError
             raise
 
+    def get_stream_flags(self):
path: cupy/cuda/stream.pyx
text: @user1:
How about making it a property with name `flags` for simplicity? It could be accessed as `stream.flags` where `stream` is an `_BaseStream` instance.

@user1:
You may name this `is_non_blocking`, then add another property `is_default`. I suppose having properties for each flag value is acceptable here as a new flag is not added to CUDA for long.

@author:
Made the changes!