cythonize stream module

leofang · leofang · commit 67db25e9ea4c · 2025-09-28T02:17:58.000Z
diff --git a/cuda_core/cuda/core/experimental/_launcher.pyx b/cuda_core/cuda/core/experimental/_launcher.pyx
@@ -2,12 +2,16 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+from libc.stdint cimport uintptr_t
+
+from cuda.core.experimental._stream cimport _try_to_get_stream_ptr
+
 from typing import Union
 
 from cuda.core.experimental._kernel_arg_handler import ParamHolder
 from cuda.core.experimental._launch_config import LaunchConfig, _to_native_launch_config
 from cuda.core.experimental._module import Kernel
-from cuda.core.experimental._stream import IsStreamT, Stream, _try_to_get_stream_ptr
+from cuda.core.experimental._stream import IsStreamT, Stream
 from cuda.core.experimental._utils.clear_error_support import assert_type
 from cuda.core.experimental._utils.cuda_utils import (
     _reduce_3_tuple,
@@ -60,7 +64,7 @@ def launch(stream: Union[Stream, IsStreamT], config: LaunchConfig, kernel: Kerne
         stream_handle = stream.handle
     except AttributeError:
         try:
-            stream_handle = _try_to_get_stream_ptr(stream)
+            stream_handle = driver.CUstream(<uintptr_t>(_try_to_get_stream_ptr(stream)))
         except Exception:
             raise ValueError(
                 f"stream must either be a Stream object or support __cuda_stream__ (got {type(stream)})"
diff --git a/cuda_core/cuda/core/experimental/_stream.pxd b/cuda_core/cuda/core/experimental/_stream.pxd
@@ -0,0 +1,9 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# TODO: how about cuda.bindings < 12.6.2?
+from cuda.bindings cimport cydriver
+
+
+cdef cydriver.CUstream _try_to_get_stream_ptr(obj: IsStreamT) except*
diff --git a/cuda_core/cuda/core/experimental/_stream.pyx b/cuda_core/cuda/core/experimental/_stream.pyx
@@ -4,10 +4,16 @@
 
 from __future__ import annotations
 
+from libc.stdint cimport uintptr_t
+
+# TODO: how about cuda.bindings < 12.6.2?
+from cuda.bindings cimport cydriver
+
 from cuda.core.experimental._utils.cuda_utils cimport (
     _check_driver_error as raise_if_driver_error,
     check_or_create_options,
 )
+
 import sys
 
 import cython
@@ -59,7 +65,7 @@ class IsStreamT(Protocol):
         ...
 
 
-def _try_to_get_stream_ptr(obj: IsStreamT):
+cdef cydriver.CUstream _try_to_get_stream_ptr(obj: IsStreamT) except*:
     try:
         cuda_stream_attr = obj.__cuda_stream__
     except AttributeError:
@@ -86,7 +92,7 @@ def _try_to_get_stream_ptr(obj: IsStreamT):
         raise RuntimeError(
             f"The first element of the sequence returned by obj.__cuda_stream__ must be 0, got {repr(info[0])}"
         )
-    return driver.CUstream(info[1])
+    return <cydriver.CUstream><uintptr_t>(info[1])
 
 
 cdef class Stream:
@@ -108,14 +114,17 @@ cdef class Stream:
     """
 
     cdef:
-        object _handle
+        cydriver.CUstream _handle
         object _owner
         object _builtin
         object _nonblocking
         object _priority
         object _device_id
         object _ctx_handle
 
+    def __cinit__(self, *args, **kwargs):
+        self._handle = <cydriver.CUstream>(NULL)
+
     def __init__(self, *args, **kwargs):
         raise RuntimeError(
             "Stream objects cannot be instantiated directly. "
@@ -125,7 +134,7 @@ cdef class Stream:
     @classmethod
     def _legacy_default(cls):
         cdef Stream self = Stream.__new__(cls)
-        self._handle = driver.CUstream(driver.CU_STREAM_LEGACY)
+        self._handle = <cydriver.CUstream>(cydriver.CU_STREAM_LEGACY)
         self._owner = None
         self._builtin = True
         self._nonblocking = None  # delayed
@@ -137,7 +146,7 @@ cdef class Stream:
     @classmethod
     def _per_thread_default(cls):
         cdef Stream self = Stream.__new__(cls)
-        self._handle = driver.CUstream(driver.CU_STREAM_PER_THREAD)
+        self._handle = <cydriver.CUstream>(cydriver.CU_STREAM_PER_THREAD)
         self._owner = None
         self._builtin = True
         self._nonblocking = None  # delayed
@@ -149,7 +158,6 @@ cdef class Stream:
     @classmethod
     def _init(cls, obj: Optional[IsStreamT] = None, options=None, device_id: int = None):
         cdef Stream self = Stream.__new__(cls)
-        self._handle = None
         self._owner = None
         self._builtin = False
 
@@ -169,16 +177,20 @@ cdef class Stream:
         nonblocking = opts.nonblocking
         priority = opts.priority
 
-        flags = driver.CUstream_flags.CU_STREAM_NON_BLOCKING if nonblocking else driver.CUstream_flags.CU_STREAM_DEFAULT
-        err, high, low = driver.cuCtxGetStreamPriorityRange()
-        raise_if_driver_error(err)
+        flags = cydriver.CUstream_flags.CU_STREAM_NON_BLOCKING if nonblocking else cydriver.CUstream_flags.CU_STREAM_DEFAULT
+        # TODO: use HANDLE_RETURN
+        cdef int high, low
+        err = cydriver.cuCtxGetStreamPriorityRange(&high, &low)
         if priority is not None:
             if not (low <= priority <= high):
                 raise ValueError(f"{priority=} is out of range {[low, high]}")
         else:
             priority = high
 
-        self._handle = handle_return(driver.cuStreamCreateWithPriority(flags, priority))
+        cdef cydriver.CUstream s
+        # TODO: add HANDLE_RETURN macro to check driver error code?
+        err = cydriver.cuStreamCreateWithPriority(&s, flags, priority)
+        self._handle = s
         self._owner = None
         self._nonblocking = nonblocking
         self._priority = priority
@@ -195,10 +207,11 @@ cdef class Stream:
 
         if self._owner is None:
             if self._handle and not self._builtin:
-                handle_return(driver.cuStreamDestroy(self._handle))
+                # TODO: use HANDLE_RETURN
+                err = cydriver.cuStreamDestroy(self._handle)
         else:
             self._owner = None
-        self._handle = None
+        self._handle = <cydriver.CUstream>(NULL)
 
     cpdef close(self):
         """Destroy the stream.
@@ -222,14 +235,16 @@ cdef class Stream:
             This handle is a Python object. To get the memory address of the underlying C
             handle, call ``int(Stream.handle)``.
         """
-        return self._handle
+        return driver.CUstream(<uintptr_t><void*>(self._handle))
 
     @property
     def is_nonblocking(self) -> bool:
         """Return True if this is a nonblocking stream, otherwise False."""
+        cdef unsigned int flags
         if self._nonblocking is None:
-            flag = handle_return(driver.cuStreamGetFlags(self._handle))
-            if flag == driver.CUstream_flags.CU_STREAM_NON_BLOCKING:
+            # TODO: switch to HANDLE_RETURN
+            err = cydriver.cuStreamGetFlags(self._handle, &flags)
+            if flags & cydriver.CUstream_flags.CU_STREAM_NON_BLOCKING:
                 self._nonblocking = True
             else:
                 self._nonblocking = False
@@ -238,14 +253,17 @@ cdef class Stream:
     @property
     def priority(self) -> int:
         """Return the stream priority."""
+        cdef int prio
         if self._priority is None:
-            prio = handle_return(driver.cuStreamGetPriority(self._handle))
+            # TODO: switch to HANDLE_RETURN
+            err = cydriver.cuStreamGetPriority(self._handle, &prio)
             self._priority = prio
         return self._priority
 
     def sync(self):
         """Synchronize the stream."""
-        handle_return(driver.cuStreamSynchronize(self._handle))
+        # TODO: switch to HANDLE_RETURN
+        err = cydriver.cuStreamSynchronize(self._handle)
 
     def record(self, event: Event = None, options: EventOptions = None) -> Event:
         """Record an event onto the stream.
@@ -272,8 +290,9 @@ cdef class Stream:
         if event is None:
             self._get_device_and_context()
             event = Event._init(self._device_id, self._ctx_handle, options)
-        err, = driver.cuEventRecord(event.handle, self._handle)
-        raise_if_driver_error(err)
+        # TODO: switch to HANDLE_RETURN
+        # TODO: revisit after Event is cythonized
+        err = cydriver.cuEventRecord(<cydriver.CUevent><uintptr_t>(event.handle), self._handle)
         return event
 
     def wait(self, event_or_stream: Union[Event, Stream]):
@@ -286,28 +305,35 @@ cdef class Stream:
         on the stream and then waiting on it.
 
         """
+        cdef cydriver.CUevent event
+        cdef cydriver.CUstream stream
+        cdef bint discard_event
+
         if isinstance(event_or_stream, Event):
-            event = event_or_stream.handle
+            event = <cydriver.CUevent><uintptr_t>(event_or_stream.handle)
             discard_event = False
         else:
             if isinstance(event_or_stream, Stream):
-                stream = event_or_stream
+                stream = <cydriver.CUstream><uintptr_t>(event_or_stream.handle)
             else:
                 try:
-                    stream = Stream._init(obj=event_or_stream)
+                    s = Stream._init(obj=event_or_stream)
                 except Exception as e:
                     raise ValueError(
                         "only an Event, Stream, or object supporting __cuda_stream__ can be waited,"
                         f" got {type(event_or_stream)}"
                     ) from e
-            event = handle_return(driver.cuEventCreate(driver.CUevent_flags.CU_EVENT_DISABLE_TIMING))
-            handle_return(driver.cuEventRecord(event, stream.handle))
+                stream = <cydriver.CUstream><uintptr_t>(s.handle)
+            # TODO: switch to HANDLE_RETURN
+            err = cydriver.cuEventCreate(&event, cydriver.CUevent_flags.CU_EVENT_DISABLE_TIMING)
+            err = cydriver.cuEventRecord(event, stream)
             discard_event = True
 
         # TODO: support flags other than 0?
-        handle_return(driver.cuStreamWaitEvent(self._handle, event, 0))
+        # TODO: switch to HANDLE_RETURN
+        err = cydriver.cuStreamWaitEvent(self._handle, event, 0)
         if discard_event:
-            handle_return(driver.cuEventDestroy(event))
+            err = cydriver.cuEventDestroy(event)
 
     @property
     def device(self) -> Device:
@@ -325,9 +351,12 @@ cdef class Stream:
         return Device(self._device_id)
 
     cdef int _get_context(Stream self) except?-1:
+        # TODO: consider making self._ctx_handle typed?
+        cdef cydriver.CUcontext ctx
         if self._ctx_handle is None:
-            err, self._ctx_handle = driver.cuStreamGetCtx(self._handle)
-            raise_if_driver_error(err)
+            # TODO: switch to HANDLE_RETURN
+            err = cydriver.cuStreamGetCtx(self._handle, &ctx)
+            self._ctx_handle = driver.CUcontext(<uintptr_t>ctx)
         return 0
 
     cdef int _get_device_and_context(Stream self) except?-1:
