fix error handling

Author: leofang

cuda_core/cuda/core/experimental/_event.pyx

@@ -10,8 +10,8 @@ from libc.stdint cimport uintptr_t
 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,
+    HANDLE_RETURN
 )
 
 from dataclasses import dataclass
@@ -110,8 +110,7 @@ cdef class Event:
             self._busy_waited = True
         if opts.support_ipc:
             raise NotImplementedError("WIP: https://github.com/NVIDIA/cuda-python/issues/103")
-        # TODO: use HANDLE_RETURN
-        err = cydriver.cuEventCreate(&self._handle, flags)
+        HANDLE_RETURN(cydriver.cuEventCreate(&self._handle, flags))
         self._device_id = device_id
         self._ctx_handle = ctx_handle
         return self
@@ -120,8 +119,7 @@ cdef class Event:
         if is_shutting_down and is_shutting_down():
             return
         if self._handle != NULL:
-            # TODO: use HANDLE_RETURN
-            err = cydriver.cuEventDestroy(self._handle)
+            HANDLE_RETURN(cydriver.cuEventDestroy(self._handle))
             self._handle = <cydriver.CUevent>(NULL)
 
     cpdef close(self):
@@ -190,8 +188,7 @@ cdef class Event:
         has been completed.
 
         """
-        # TODO: use HANDLE_RETURN
-        err = cydriver.cuEventSynchronize(self._handle)
+        HANDLE_RETURN(cydriver.cuEventSynchronize(self._handle))
 
     @property
     def is_done(self) -> bool:
@@ -201,7 +198,7 @@ cdef class Event:
             return True
         if result == cydriver.CUresult.CUDA_ERROR_NOT_READY:
             return False
-        # TODO: use HANDLE_RETURN
+        HANDLE_RETURN(result)
 
     @property
     def handle(self) -> cuda.bindings.driver.CUevent:

cuda_core/cuda/core/experimental/_stream.pyx

@@ -10,8 +10,8 @@ from libc.stdint cimport uintptr_t
 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,
+    HANDLE_RETURN,
 )
 
 import sys
@@ -178,18 +178,16 @@ cdef class Stream:
         priority = opts.priority
 
         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)
+        HANDLE_RETURN(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
 
         cdef cydriver.CUstream s
-        # TODO: add HANDLE_RETURN macro to check driver error code?
-        err = cydriver.cuStreamCreateWithPriority(&s, flags, priority)
+        HANDLE_RETURN(cydriver.cuStreamCreateWithPriority(&s, flags, priority))
         self._handle = s
         self._owner = None
         self._nonblocking = nonblocking
@@ -207,8 +205,7 @@ cdef class Stream:
 
         if self._owner is None:
             if self._handle and not self._builtin:
-                # TODO: use HANDLE_RETURN
-                err = cydriver.cuStreamDestroy(self._handle)
+                HANDLE_RETURN(cydriver.cuStreamDestroy(self._handle))
         else:
             self._owner = None
         self._handle = <cydriver.CUstream>(NULL)
@@ -242,8 +239,7 @@ cdef class Stream:
         """Return True if this is a nonblocking stream, otherwise False."""
         cdef unsigned int flags
         if self._nonblocking is None:
-            # TODO: switch to HANDLE_RETURN
-            err = cydriver.cuStreamGetFlags(self._handle, &flags)
+            HANDLE_RETURN(cydriver.cuStreamGetFlags(self._handle, &flags))
             if flags & cydriver.CUstream_flags.CU_STREAM_NON_BLOCKING:
                 self._nonblocking = True
             else:
@@ -255,15 +251,13 @@ cdef class Stream:
         """Return the stream priority."""
         cdef int prio
         if self._priority is None:
-            # TODO: switch to HANDLE_RETURN
-            err = cydriver.cuStreamGetPriority(self._handle, &prio)
+            HANDLE_RETURN(cydriver.cuStreamGetPriority(self._handle, &prio))
             self._priority = prio
         return self._priority
 
     def sync(self):
         """Synchronize the stream."""
-        # TODO: switch to HANDLE_RETURN
-        err = cydriver.cuStreamSynchronize(self._handle)
+        HANDLE_RETURN(cydriver.cuStreamSynchronize(self._handle))
 
     def record(self, event: Event = None, options: EventOptions = None) -> Event:
         """Record an event onto the stream.
@@ -290,9 +284,8 @@ cdef class Stream:
         if event is None:
             self._get_device_and_context()
             event = Event._init(self._device_id, self._ctx_handle, options)
-        # TODO: switch to HANDLE_RETURN
         # TODO: revisit after Event is cythonized
-        err = cydriver.cuEventRecord(<cydriver.CUevent><uintptr_t>(event.handle), self._handle)
+        HANDLE_RETURN(cydriver.cuEventRecord(<cydriver.CUevent><uintptr_t>(event.handle), self._handle))
         return event
 
     def wait(self, event_or_stream: Union[Event, Stream]):
@@ -324,16 +317,14 @@ cdef class Stream:
                         f" got {type(event_or_stream)}"
                     ) from e
                 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)
+            HANDLE_RETURN(cydriver.cuEventCreate(&event, cydriver.CUevent_flags.CU_EVENT_DISABLE_TIMING))
+            HANDLE_RETURN(cydriver.cuEventRecord(event, stream))
             discard_event = True
 
         # TODO: support flags other than 0?
-        # TODO: switch to HANDLE_RETURN
-        err = cydriver.cuStreamWaitEvent(self._handle, event, 0)
+        HANDLE_RETURN(cydriver.cuStreamWaitEvent(self._handle, event, 0))
         if discard_event:
-            err = cydriver.cuEventDestroy(event)
+            HANDLE_RETURN(cydriver.cuEventDestroy(event))
 
     @property
     def device(self) -> Device:
@@ -354,8 +345,7 @@ cdef class Stream:
         # TODO: consider making self._ctx_handle typed?
         cdef cydriver.CUcontext ctx
         if self._ctx_handle is None:
-            # TODO: switch to HANDLE_RETURN
-            err = cydriver.cuStreamGetCtx(self._handle, &ctx)
+            HANDLE_RETURN(cydriver.cuStreamGetCtx(self._handle, &ctx))
             self._ctx_handle = driver.CUcontext(<uintptr_t>ctx)
         return 0
 

cuda_core/cuda/core/experimental/_utils/cuda_utils.pxd

@@ -2,18 +2,30 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-
 cimport cpython
-cimport libc.stdint
+from libc.stdint cimport int64_t
+
+# TODO: how about cuda.bindings < 12.6.2?
+from cuda.bindings cimport cydriver
+
+
+ctypedef fused supported_error_type:
+    cydriver.CUresult
 
 
+cdef int HANDLE_RETURN(supported_error_type err) except?-1
+
+
+# TODO: stop exposing these within the codebase?
 cpdef int _check_driver_error(error) except?-1
 cpdef int _check_runtime_error(error) except?-1
 cpdef int _check_nvrtc_error(error) except?-1
+
+
 cpdef check_or_create_options(type cls, options, str options_description=*, bint keep_none=*)
 
 
-cdef inline tuple carray_int64_t_to_tuple(libc.stdint.int64_t *ptr, int length):
+cdef inline tuple carray_int64_t_to_tuple(int64_t *ptr, int length):
     # Construct shape and strides tuples using the Python/C API for speed
     result = cpython.PyTuple_New(length)
     for i in range(length):

cuda_core/cuda/core/experimental/_utils/cuda_utils.pyx

@@ -52,6 +52,12 @@ def _reduce_3_tuple(t: tuple):
     return t[0] * t[1] * t[2]
 
 
+cdef int HANDLE_RETURN(supported_error_type err) except?-1:
+    if supported_error_type is cydriver.CUresult:
+        if err != cydriver.CUresult.CUDA_SUCCESS:
+            return _check_driver_error(err)
+
+
 cdef object _DRIVER_SUCCESS = driver.CUresult.CUDA_SUCCESS
 cdef object _RUNTIME_SUCCESS = runtime.cudaError_t.cudaSuccess
 cdef object _NVRTC_SUCCESS = nvrtc.nvrtcResult.NVRTC_SUCCESS