inline @precondition to reduce overhead

Author: leofang

cuda_core/cuda/core/experimental/_device.py

@@ -17,7 +17,6 @@
     _check_driver_error,
     driver,
     handle_return,
-    precondition,
     runtime,
 )
 
@@ -1017,12 +1016,31 @@ def __new__(cls, device_id: Optional[int] = None):
         except IndexError:
             raise ValueError(f"device_id must be within [0, {len(devices)}), got {device_id}") from None
 
-    def _check_context_initialized(self, *args, **kwargs):
+    def _check_context_initialized(self):
         if not self._has_inited:
             raise CUDAError(
                 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:
+        err, ctx = driver.cuCtxGetCurrent()
+
+        # TODO: We want to just call this:
+        #_check_driver_error(err)
+        # but even the simplest success check causes 50-100 ns. Wait until we cythonize this file...
+        if ctx is None:
+            _check_driver_error(err)
+
+        if int(ctx) == 0:
+            raise CUDAError("No context is bound to the calling CPU thread.")
+        if check_consistency:
+            err, dev = driver.cuCtxGetDevice()
+            if err != _SUCCESS:
+                handle_return((err,))
+            if int(dev) != self._id:
+                raise CUDAError("Internal error (current device is not equal to Device.device_id)")
+        return ctx
+
     @property
     def device_id(self) -> int:
         """Return device ordinal."""
@@ -1083,7 +1101,6 @@ def compute_capability(self) -> ComputeCapability:
         return cc
 
     @property
-    @precondition(_check_context_initialized)
     def context(self) -> Context:
         """Return the current :obj:`~_context.Context` associated with this device.
 
@@ -1092,9 +1109,8 @@ def context(self) -> Context:
         Device must be initialized.
 
         """
-        ctx = handle_return(driver.cuCtxGetCurrent())
-        if int(ctx) == 0:
-            raise CUDAError("No context is bound to the calling CPU thread.")
+        self._check_context_initialized()
+        ctx = self._get_current_context(check_consistency=True)
         return Context._from_ctx(ctx, self._id)
 
     @property
@@ -1206,7 +1222,6 @@ def create_context(self, options: ContextOptions = None) -> Context:
         """
         raise NotImplementedError("WIP: https://github.com/NVIDIA/cuda-python/issues/189")
 
-    @precondition(_check_context_initialized)
     def create_stream(self, obj: Optional[IsStreamT] = None, options: StreamOptions = None) -> Stream:
         """Create a Stream object.
 
@@ -1235,6 +1250,7 @@ def create_stream(self, obj: Optional[IsStreamT] = None, options: StreamOptions
             Newly created stream object.
 
         """
+        self._check_context_initialized()
         return Stream._init(obj=obj, options=options)
 
     def create_event(self, options: Optional[EventOptions] = None) -> Event:
@@ -1255,12 +1271,10 @@ def create_event(self, options: Optional[EventOptions] = None) -> Event:
             Newly created event object.
 
         """
-        ctx = driver.cuCtxGetCurrent()[1]
-        if int(ctx) == 0:
-            raise CUDAError("No context is bound to the calling CPU thread.")
+        self._check_context_initialized()
+        ctx = self._get_current_context()
         return Event._init(self._id, ctx, options)
 
-    @precondition(_check_context_initialized)
     def allocate(self, size, stream: Optional[Stream] = None) -> Buffer:
         """Allocate device memory from a specified stream.
 
@@ -1287,11 +1301,11 @@ def allocate(self, size, stream: Optional[Stream] = None) -> Buffer:
             Newly created buffer object.
 
         """
+        self._check_context_initialized()
         if stream is None:
             stream = default_stream()
         return self._mr.allocate(size, stream)
 
-    @precondition(_check_context_initialized)
     def sync(self):
         """Synchronize the device.
 
@@ -1300,9 +1314,9 @@ def sync(self):
         Device must be initialized.
 
         """
+        self._check_context_initialized()
         handle_return(runtime.cudaDeviceSynchronize())
 
-    @precondition(_check_context_initialized)
     def create_graph_builder(self) -> GraphBuilder:
         """Create a new :obj:`~_graph.GraphBuilder` object.
 
@@ -1312,4 +1326,5 @@ def create_graph_builder(self) -> GraphBuilder:
             Newly created graph builder object.
 
         """
+        self._check_context_initialized()
         return GraphBuilder._init(stream=self.create_stream(), is_stream_owner=True)

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

@@ -52,26 +52,29 @@ def _reduce_3_tuple(t: tuple):
     return t[0] * t[1] * t[2]
 
 
-cpdef inline void _check_driver_error(error) except*:
-    if error == driver.CUresult.CUDA_SUCCESS:
-        return
+cdef object _DRIVER_SUCCESS = driver.CUresult.CUDA_SUCCESS
+
+
+cpdef inline int _check_driver_error(error) except?-1:
+    if error == _DRIVER_SUCCESS:
+        return 0
     name_err, name = driver.cuGetErrorName(error)
-    if name_err != driver.CUresult.CUDA_SUCCESS:
+    if name_err != _DRIVER_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.CUresult.CUDA_SUCCESS:
+    if desc_err != _DRIVER_SUCCESS:
         raise CUDAError(f"{name}")
     desc = desc.decode()
     raise CUDAError(f"{name}: {desc}")
 
 
-cpdef inline void _check_runtime_error(error) except*:
+cpdef inline int _check_runtime_error(error) except?-1:
     if error == runtime.cudaError_t.cudaSuccess:
-        return
+        return 0
     name_err, name = runtime.cudaGetErrorName(error)
     if name_err != runtime.cudaError_t.cudaSuccess:
         raise CUDAError(f"UNEXPECTED ERROR CODE: {error}")
@@ -86,30 +89,35 @@ cpdef inline void _check_runtime_error(error) except*:
     raise CUDAError(f"{name}: {desc}")
 
 
-cdef inline void _check_error(error, handle=None) except*:
+cpdef inline int _check_nvrtc_error(error, handle=None) except?-1:
+    if error == nvrtc.nvrtcResult.NVRTC_SUCCESS:
+        return 0
+    err = f"{error}: {nvrtc.nvrtcGetErrorString(error)[1].decode()}"
+    if handle is not None:
+        _, logsize = nvrtc.nvrtcGetProgramLogSize(handle)
+        log = b" " * logsize
+        _ = nvrtc.nvrtcGetProgramLog(handle, log)
+        err += f", compilation log:\n\n{log.decode('utf-8', errors='backslashreplace')}"
+    raise NVRTCError(err)
+
+
+cdef inline int _check_error(error, handle=None) except?-1:
     if isinstance(error, driver.CUresult):
-        _check_driver_error(error)
+        return _check_driver_error(error)
     elif isinstance(error, runtime.cudaError_t):
-        _check_runtime_error(error)
+        return _check_runtime_error(error)
     elif isinstance(error, nvrtc.nvrtcResult):
-        if error == nvrtc.nvrtcResult.NVRTC_SUCCESS:
-            return
-        err = f"{error}: {nvrtc.nvrtcGetErrorString(error)[1].decode()}"
-        if handle is not None:
-            _, logsize = nvrtc.nvrtcGetProgramLogSize(handle)
-            log = b" " * logsize
-            _ = nvrtc.nvrtcGetProgramLog(handle, log)
-            err += f", compilation log:\n\n{log.decode('utf-8', errors='backslashreplace')}"
-        raise NVRTCError(err)
+        return _check_nvrtc_error(error, handle=handle)
     else:
         raise RuntimeError(f"Unknown error type: {error}")
 
 
 def handle_return(tuple result, handle=None):
     _check_error(result[0], handle=handle)
-    if len(result) == 1:
+    cdef int out_len = len(result)
+    if out_len == 1:
         return
-    elif len(result) == 2:
+    elif out_len == 2:
         return result[1]
     else:
         return result[1:]
@@ -144,7 +152,7 @@ def _handle_boolean_option(option: bool) -> str:
     return "true" if bool(option) else "false"
 
 
-def precondition(checker: Callable[..., None], what: str = "") -> Callable:
+def precondition(checker: Callable[..., None], str what="") -> Callable:
     """
     A decorator that adds checks to ensure any preconditions are met.