release gil

Author: leofang

cuda_core/cuda/core/experimental/_device.pyx

@@ -55,7 +55,8 @@ cdef class DeviceProperties:
     cdef inline _get_attribute(self, cydriver.CUdevice_attribute attr):
         """Retrieve the attribute value directly from the driver."""
         cdef int val
-        HANDLE_RETURN(cydriver.cuDeviceGetAttribute(&val, attr, self._handle))
+        with nogil:
+            HANDLE_RETURN(cydriver.cuDeviceGetAttribute(&val, attr, self._handle))
         return val
 
     cdef _get_cached_attribute(self, attr):
@@ -949,7 +950,8 @@ cdef cydriver.CUcontext _get_primary_context(int dev_id) except?NULL:
         primary_ctxs = _tls.primary_ctxs = [0] * total
     cdef cydriver.CUcontext ctx = <cydriver.CUcontext><uintptr_t>(primary_ctxs[dev_id])
     if ctx == NULL:
-        HANDLE_RETURN(cydriver.cuDevicePrimaryCtxRetain(&ctx, dev_id))
+        with nogil:
+            HANDLE_RETURN(cydriver.cuDevicePrimaryCtxRetain(&ctx, dev_id))
         primary_ctxs[dev_id] = <uintptr_t>(ctx)
     return ctx
 
@@ -985,19 +987,21 @@ class Device:
     def __new__(cls, device_id: Optional[int] = None):
         global _is_cuInit
         if _is_cuInit is False:
-            with _lock:
+            with _lock, nogil:
                 HANDLE_RETURN(cydriver.cuInit(0))
                 _is_cuInit = True
 
         # important: creating a Device instance does not initialize the GPU!
         cdef cydriver.CUdevice dev
         cdef cydriver.CUcontext ctx
         if device_id is None:
-            err = cydriver.cuCtxGetDevice(&dev)
+            with nogil:
+                err = cydriver.cuCtxGetDevice(&dev)
             if err == cydriver.CUresult.CUDA_SUCCESS:
                 device_id = int(dev)
             elif err == cydriver.CUresult.CUDA_ERROR_INVALID_CONTEXT:
-                HANDLE_RETURN(cydriver.cuCtxGetCurrent(&ctx))
+                with nogil:
+                    HANDLE_RETURN(cydriver.cuCtxGetCurrent(&ctx))
                 assert <void*>(ctx) == NULL
                 device_id = 0  # cudart behavior
             else:
@@ -1010,18 +1014,20 @@ class Device:
         try:
             devices = _tls.devices
         except AttributeError:
-            HANDLE_RETURN(cydriver.cuDeviceGetCount(&total))
+            with nogil:
+                HANDLE_RETURN(cydriver.cuDeviceGetCount(&total))
             devices = _tls.devices = []
             for dev_id in range(total):
                 device = super().__new__(cls)
                 device._id = dev_id
                 # If the device is in TCC mode, or does not support memory pools for some other reason,
                 # use the SynchronousMemoryResource which does not use memory pools.
-                HANDLE_RETURN(
-                    cydriver.cuDeviceGetAttribute(
-                        &attr, cydriver.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MEMORY_POOLS_SUPPORTED, dev_id
+                with nogil:
+                    HANDLE_RETURN(
+                        cydriver.cuDeviceGetAttribute(
+                            &attr, cydriver.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MEMORY_POOLS_SUPPORTED, dev_id
+                        )
                     )
-                )
                 if attr == 1:
                     device._mr = DeviceMemoryResource(dev_id)
                 else:
@@ -1042,16 +1048,18 @@ class Device:
                 f"Device {self._id} is not yet initialized, perhaps you forgot to call .set_current() first?"
             )
 
-    def _get_current_context(self, check_consistency=False) -> driver.CUcontext:
+    def _get_current_context(self, bint check_consistency=False) -> driver.CUcontext:
         cdef cydriver.CUcontext ctx
-        HANDLE_RETURN(cydriver.cuCtxGetCurrent(&ctx))
-        if ctx == NULL:
-            raise CUDAError("No context is bound to the calling CPU thread.")
         cdef cydriver.CUdevice dev
-        if check_consistency:
-            HANDLE_RETURN(cydriver.cuCtxGetDevice(&dev))
-            if <int>(dev) != self._id:
-                raise CUDAError("Internal error (current device is not equal to Device.device_id)")
+        cdef cydriver.CUdevice this_dev = self._id
+        with nogil:
+            HANDLE_RETURN(cydriver.cuCtxGetCurrent(&ctx))
+            if ctx == NULL:
+                raise CUDAError("No context is bound to the calling CPU thread.")
+            if check_consistency:
+                HANDLE_RETURN(cydriver.cuCtxGetDevice(&dev))
+                if dev != this_dev:
+                    raise CUDAError("Internal error (current device is not equal to Device.device_id)")
         return driver.CUcontext(<uintptr_t>ctx)
 
     @property
@@ -1080,10 +1088,12 @@ class Device:
 
         """
         cdef cydriver.CUuuid uuid
-        IF CUDA_CORE_BUILD_MAJOR == "12":
-            HANDLE_RETURN(cydriver.cuDeviceGetUuid_v2(&uuid, self._id))
-        ELSE:  # 13.0+
-            HANDLE_RETURN(cydriver.cuDeviceGetUuid(&uuid, self._id))
+        cdef cydriver.CUdevice this_dev = self._id
+        with nogil:
+            IF CUDA_CORE_BUILD_MAJOR == "12":
+                HANDLE_RETURN(cydriver.cuDeviceGetUuid_v2(&uuid, this_dev))
+            ELSE:  # 13.0+
+                HANDLE_RETURN(cydriver.cuDeviceGetUuid(&uuid, this_dev))
         cdef bytes uuid_b = cpython.PyBytes_FromStringAndSize(uuid.bytes, sizeof(uuid.bytes))
         cdef str uuid_hex = uuid_b.hex()
         # 8-4-4-4-12
@@ -1095,7 +1105,10 @@ class Device:
         # Use 256 characters to be consistent with CUDA Runtime
         cdef int LENGTH = 256
         cdef bytes name = bytes(LENGTH)
-        HANDLE_RETURN(cydriver.cuDeviceGetName(<char*>name, LENGTH, self._id))
+        cdef char* name_ptr = name
+        cdef cydriver.CUdevice this_dev = self._id
+        with nogil:
+            HANDLE_RETURN(cydriver.cuDeviceGetName(name_ptr, LENGTH, this_dev))
         name = name.split(b"\0")[0]
         return name.decode()
 
@@ -1198,7 +1211,8 @@ class Device:
         >>> # ... do work on device 0 ...
 
         """
-        cdef cydriver.CUcontext _ctx
+        cdef cydriver.CUcontext prev_ctx
+        cdef cydriver.CUcontext curr_ctx
         if ctx is not None:
             # TODO: revisit once Context is cythonized
             assert_type(ctx, Context)
@@ -1207,16 +1221,19 @@ class Device:
                     "the provided context was created on the device with"
                     f" id={ctx._id}, which is different from the target id={self._id}"
                 )
-            # _ctx is the previous context
-            HANDLE_RETURN(cydriver.cuCtxPopCurrent(&_ctx))
-            HANDLE_RETURN(cydriver.cuCtxPushCurrent(<cydriver.CUcontext>(ctx._handle)))
+            # prev_ctx is the previous context
+            curr_ctx = <cydriver.CUcontext>(ctx._handle)
+            with nogil:
+                HANDLE_RETURN(cydriver.cuCtxPopCurrent(&prev_ctx))
+                HANDLE_RETURN(cydriver.cuCtxPushCurrent(curr_ctx))
             self._has_inited = True
-            if _ctx != NULL:
-                return Context._from_ctx(<uintptr_t>(_ctx), self._id)
+            if prev_ctx != NULL:
+                return Context._from_ctx(<uintptr_t>(prev_ctx), self._id)
         else:
             # use primary ctx
-            _ctx = _get_primary_context(self._id)
-            HANDLE_RETURN(cydriver.cuCtxSetCurrent(_ctx))
+            curr_ctx = _get_primary_context(self._id)
+            with nogil:
+                HANDLE_RETURN(cydriver.cuCtxSetCurrent(curr_ctx))
             self._has_inited = True
 
     def create_context(self, options: ContextOptions = None) -> Context:

cuda_core/cuda/core/experimental/_event.pyx

@@ -109,7 +109,8 @@ cdef class Event:
             self._busy_waited = True
         if opts.support_ipc:
             raise NotImplementedError("WIP: https://github.com/NVIDIA/cuda-python/issues/103")
-        HANDLE_RETURN(cydriver.cuEventCreate(&self._handle, flags))
+        with nogil:
+            HANDLE_RETURN(cydriver.cuEventCreate(&self._handle, flags))
         self._device_id = device_id
         self._ctx_handle = ctx_handle
         return self
@@ -118,7 +119,8 @@ cdef class Event:
         if is_shutting_down and is_shutting_down():
             return
         if self._handle != NULL:
-            HANDLE_RETURN(cydriver.cuEventDestroy(self._handle))
+            with nogil:
+                HANDLE_RETURN(cydriver.cuEventDestroy(self._handle))
             self._handle = <cydriver.CUevent>(NULL)
 
     cpdef close(self):
@@ -137,7 +139,8 @@ cdef class Event:
     def __sub__(self, other: Event):
         # return self - other (in milliseconds)
         cdef float timing
-        err = cydriver.cuEventElapsedTime(&timing, other._handle, self._handle)
+        with nogil:
+            err = cydriver.cuEventElapsedTime(&timing, other._handle, self._handle)
         if err == 0:
             return timing
         else:
@@ -187,12 +190,14 @@ cdef class Event:
         has been completed.
 
         """
-        HANDLE_RETURN(cydriver.cuEventSynchronize(self._handle))
+        with nogil:
+            HANDLE_RETURN(cydriver.cuEventSynchronize(self._handle))
 
     @property
     def is_done(self) -> bool:
         """Return True if all captured works have been completed, otherwise False."""
-        result = cydriver.cuEventQuery(self._handle)
+        with nogil:
+            result = cydriver.cuEventQuery(self._handle)
         if result == cydriver.CUresult.CUDA_SUCCESS:
             return True
         if result == cydriver.CUresult.CUDA_ERROR_NOT_READY:

cuda_core/cuda/core/experimental/_stream.pyx

@@ -177,16 +177,21 @@ cdef class Stream:
         priority = opts.priority
 
         flags = cydriver.CUstream_flags.CU_STREAM_NON_BLOCKING if nonblocking else cydriver.CUstream_flags.CU_STREAM_DEFAULT
+        # TODO: we might want to consider memoizing high/low per CUDA context and avoid this call
         cdef int high, low
-        HANDLE_RETURN(cydriver.cuCtxGetStreamPriorityRange(&high, &low))
+        with nogil:
+            HANDLE_RETURN(cydriver.cuCtxGetStreamPriorityRange(&high, &low))
+        cdef int prio
         if priority is not None:
-            if not (low <= priority <= high):
+            prio = priority
+            if not (low <= prio <= high):
                 raise ValueError(f"{priority=} is out of range {[low, high]}")
         else:
-            priority = high
+            prio = high
 
         cdef cydriver.CUstream s
-        HANDLE_RETURN(cydriver.cuStreamCreateWithPriority(&s, flags, priority))
+        with nogil:
+            HANDLE_RETURN(cydriver.cuStreamCreateWithPriority(&s, flags, prio))
         self._handle = s
         self._owner = None
         self._nonblocking = nonblocking
@@ -204,7 +209,8 @@ cdef class Stream:
 
         if self._owner is None:
             if self._handle and not self._builtin:
-                HANDLE_RETURN(cydriver.cuStreamDestroy(self._handle))
+                with nogil:
+                    HANDLE_RETURN(cydriver.cuStreamDestroy(self._handle))
         else:
             self._owner = None
         self._handle = <cydriver.CUstream>(NULL)
@@ -238,7 +244,8 @@ cdef class Stream:
         """Return True if this is a nonblocking stream, otherwise False."""
         cdef unsigned int flags
         if self._nonblocking is None:
-            HANDLE_RETURN(cydriver.cuStreamGetFlags(self._handle, &flags))
+            with nogil:
+                HANDLE_RETURN(cydriver.cuStreamGetFlags(self._handle, &flags))
             if flags & cydriver.CUstream_flags.CU_STREAM_NON_BLOCKING:
                 self._nonblocking = True
             else:
@@ -250,13 +257,15 @@ cdef class Stream:
         """Return the stream priority."""
         cdef int prio
         if self._priority is None:
-            HANDLE_RETURN(cydriver.cuStreamGetPriority(self._handle, &prio))
+            with nogil:
+                HANDLE_RETURN(cydriver.cuStreamGetPriority(self._handle, &prio))
             self._priority = prio
         return self._priority
 
     def sync(self):
         """Synchronize the stream."""
-        HANDLE_RETURN(cydriver.cuStreamSynchronize(self._handle))
+        with nogil:
+            HANDLE_RETURN(cydriver.cuStreamSynchronize(self._handle))
 
     def record(self, event: Event = None, options: EventOptions = None) -> Event:
         """Record an event onto the stream.
@@ -299,11 +308,12 @@ cdef class Stream:
         """
         cdef cydriver.CUevent event
         cdef cydriver.CUstream stream
-        cdef bint discard_event
 
         if isinstance(event_or_stream, Event):
             event = <cydriver.CUevent><uintptr_t>(event_or_stream.handle)
-            discard_event = False
+            with nogil:
+                # TODO: support flags other than 0?
+                HANDLE_RETURN(cydriver.cuStreamWaitEvent(self._handle, event, 0))
         else:
             if isinstance(event_or_stream, Stream):
                 stream = <cydriver.CUstream><uintptr_t>(event_or_stream.handle)
@@ -316,14 +326,12 @@ cdef class Stream:
                         f" got {type(event_or_stream)}"
                     ) from e
                 stream = <cydriver.CUstream><uintptr_t>(s.handle)
-            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?
-        HANDLE_RETURN(cydriver.cuStreamWaitEvent(self._handle, event, 0))
-        if discard_event:
-            HANDLE_RETURN(cydriver.cuEventDestroy(event))
+            with nogil:
+                HANDLE_RETURN(cydriver.cuEventCreate(&event, cydriver.CUevent_flags.CU_EVENT_DISABLE_TIMING))
+                HANDLE_RETURN(cydriver.cuEventRecord(event, stream))
+                # TODO: support flags other than 0?
+                HANDLE_RETURN(cydriver.cuStreamWaitEvent(self._handle, event, 0))
+                HANDLE_RETURN(cydriver.cuEventDestroy(event))
 
     @property
     def device(self) -> Device:
@@ -344,7 +352,8 @@ cdef class Stream:
         # TODO: consider making self._ctx_handle typed?
         cdef cydriver.CUcontext ctx
         if self._ctx_handle is None:
-            HANDLE_RETURN(cydriver.cuStreamGetCtx(self._handle, &ctx))
+            with nogil:
+                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

@@ -12,11 +12,11 @@ ctypedef fused supported_error_type:
     cydriver.CUresult
 
 
-cdef int HANDLE_RETURN(supported_error_type err) except?-1
+cdef int HANDLE_RETURN(supported_error_type err) except?-1 nogil
 
 
 # TODO: stop exposing these within the codebase?
-cpdef int _check_driver_error(error) except?-1
+cpdef int _check_driver_error(cydriver.CUresult error) except?-1 nogil
 cpdef int _check_runtime_error(error) except?-1
 cpdef int _check_nvrtc_error(error) except?-1
 

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

@@ -52,32 +52,33 @@ def _reduce_3_tuple(t: tuple):
     return t[0] * t[1] * t[2]
 
 
-cdef int HANDLE_RETURN(supported_error_type err) except?-1:
+cdef int HANDLE_RETURN(supported_error_type err) except?-1 nogil:
     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
 
 
-cpdef inline int _check_driver_error(error) except?-1:
-    if error == _DRIVER_SUCCESS:
+cpdef inline int _check_driver_error(cydriver.CUresult error) except?-1 nogil:
+    if error == cydriver.CUresult.CUDA_SUCCESS:
         return 0
-    name_err, name = driver.cuGetErrorName(error)
-    if name_err != _DRIVER_SUCCESS:
+    cdef const char* name
+    name_err = cydriver.cuGetErrorName(error, &name)
+    if name_err != cydriver.CUresult.CUDA_SUCCESS:
         raise CUDAError(f"UNEXPECTED ERROR CODE: {error}")
-    name = name.decode()
-    expl = DRIVER_CU_RESULT_EXPLANATIONS.get(int(error))
-    if expl is not None:
-        raise CUDAError(f"{name}: {expl}")
-    desc_err, desc = driver.cuGetErrorString(error)
-    if desc_err != _DRIVER_SUCCESS:
-        raise CUDAError(f"{name}")
-    desc = desc.decode()
-    raise CUDAError(f"{name}: {desc}")
+    with gil:
+        # TODO: consider lower this to Cython
+        expl = DRIVER_CU_RESULT_EXPLANATIONS.get(int(error))
+        if expl is not None:
+            raise CUDAError(f"{name.decode()}: {expl}")
+    cdef const char* desc
+    desc_err = cydriver.cuGetErrorString(error, &desc)
+    if desc_err != cydriver.CUresult.CUDA_SUCCESS:
+        raise CUDAError(f"{name.decode()}")
+    raise CUDAError(f"{name.decode()}: {desc.decode()}")
 
 
 cpdef inline int _check_runtime_error(error) except?-1: