Simplify Stream_accept. Default arguments can more easily be handled outside that function.

Author: Andy-Jost

cuda_core/cuda/core/experimental/_launcher.pyx

@@ -54,7 +54,7 @@ def launch(stream: Stream | GraphBuilder | IsStreamT, config: LaunchConfig, kern
         launching kernel.
 
     """
-    stream = Stream_accept(stream, allow_default=False, default_value=None, allow_stream_protocol=True)
+    stream = Stream_accept(stream, allow_stream_protocol=True)
     assert_type(kernel, Kernel)
     _lazy_init()
     config = check_or_create_options(LaunchConfig, config, "launch config")

cuda_core/cuda/core/experimental/_memory/_buffer.pyx

@@ -273,7 +273,7 @@ cdef class Buffer:
 cdef inline void Buffer_close(Buffer self, stream):
     cdef Stream s
     if self._ptr and self._memory_resource is not None:
-        s = Stream_accept(stream, allow_default=True, default_value=self._alloc_stream)
+        s = Stream_accept(stream) if stream is not None else self._alloc_stream
         self._memory_resource.deallocate(self._ptr, self._size, s)
         self._ptr = 0
         self._memory_resource = None

cuda_core/cuda/core/experimental/_memory/_device_memory_resource.pyx

@@ -12,7 +12,7 @@ from cuda.bindings cimport cydriver
 from cuda.core.experimental._memory._buffer cimport Buffer, MemoryResource
 from cuda.core.experimental._memory cimport _ipc
 from cuda.core.experimental._memory._ipc cimport IPCAllocationHandle, IPCData
-from cuda.core.experimental._stream cimport Stream_accept, Stream
+from cuda.core.experimental._stream cimport default_stream, Stream_accept, Stream
 from cuda.core.experimental._utils.cuda_utils cimport (
     check_or_create_options,
     HANDLE_RETURN,
@@ -334,7 +334,7 @@ cdef class DeviceMemoryResource(MemoryResource):
         """
         if self.is_mapped:
             raise TypeError("Cannot allocate from a mapped IPC-enabled memory resource")
-        stream = Stream_accept(stream, allow_default=True)
+        stream = Stream_accept(stream) if stream is not None else default_stream()
         return DMR_allocate(self, size, <Stream> stream)
 
     def deallocate(self, ptr: DevicePointerT, size_t size, stream: Stream | GraphBuilder | None = None):
@@ -351,7 +351,7 @@ cdef class DeviceMemoryResource(MemoryResource):
             If the buffer is deallocated without an explicit stream, the allocation stream
             is used.
         """
-        stream = Stream_accept(stream, allow_default=True)
+        stream = Stream_accept(stream) if stream is not None else default_stream()
         DMR_deallocate(self, <uintptr_t>ptr, size, <Stream> stream)
 
     @property

cuda_core/cuda/core/experimental/_memory/_graph_memory_resource.pyx

@@ -8,7 +8,7 @@ from libc.stdint cimport intptr_t
 
 from cuda.bindings cimport cydriver
 from cuda.core.experimental._memory._buffer cimport Buffer, MemoryResource
-from cuda.core.experimental._stream cimport Stream_accept, Stream
+from cuda.core.experimental._stream cimport default_stream, Stream_accept, Stream
 from cuda.core.experimental._utils.cuda_utils cimport HANDLE_RETURN
 
 from functools import cache
@@ -106,14 +106,14 @@ cdef class cyGraphMemoryResource(MemoryResource):
         """
         Allocate a buffer of the requested size. See documentation for :obj:`~_memory.MemoryResource`.
         """
-        stream = Stream_accept(stream, allow_default=True)
+        stream = Stream_accept(stream) if stream is not None else default_stream()
         return GMR_allocate(self, size, <Stream> stream)
 
     def deallocate(self, ptr: DevicePointerT, size_t size, stream: Stream | GraphBuilder | None = None):
         """
         Deallocate a buffer of the requested size. See documentation for :obj:`~_memory.MemoryResource`.
         """
-        stream = Stream_accept(stream, allow_default=True)
+        stream = Stream_accept(stream) if stream is not None else default_stream()
         return GMR_deallocate(ptr, size, <Stream> stream)
 
     def close(self):

cuda_core/cuda/core/experimental/_stream.pxd

@@ -22,5 +22,4 @@ cdef class Stream:
 
 
 cpdef Stream default_stream()
-cdef Stream Stream_accept(arg, bint allow_default=*, Stream default_value=*, bint allow_stream_protocol=*)
-# from cuda.core.experimental._stream cimport Stream_accept
+cdef Stream Stream_accept(arg, bint allow_stream_protocol=*)

cuda_core/cuda/core/experimental/_stream.pyx

@@ -451,14 +451,8 @@ cdef cydriver.CUstream _handle_from_stream_protocol(obj) except*:
 
 # Helper for API functions that accept either Stream or GraphBuilder. Performs
 # needed checks and returns the relevant stream.
-cdef Stream Stream_accept(arg, bint allow_default=False, Stream default_value=None, bint allow_stream_protocol=False):
-    if arg is None:
-        if allow_default:
-            if default_value is not None:
-                return default_value
-            else:
-                return default_stream()
-    elif isinstance(arg, Stream):
+cdef Stream Stream_accept(arg, bint allow_stream_protocol=False):
+    if isinstance(arg, Stream):
         return <Stream> arg
     elif isinstance(arg, GraphBuilder):
         return <Stream> arg.stream