Cythonize away some perf hot spots (#709)

* cythonize event

* cythonize context/event/util

* Cython 3.0+ supports __del__ for cdef classes

* inline precondition to reduce overhead

* centralize check_or_create_options

* add back error check

* cythonize stream + bug fixes

* make linter happy

* bug fix

* Cython mis-compiles Optional types

* Modified _get_context_device() helper routine

To fix test failures with CTK 11.8 and driver 535.247.01
only attempt to query _ctx_handle if _device_id is None.

Ensure that context handle is set in Stream.context property

* Verify that stream.context handle is not None in test_stream_context

* cache success enums

* nit: avoid cdef void

* In Event.close set handle to None before raising error

---------

Co-authored-by: Oleksandr Pavlyk <[email protected]>

Author: leofang

cuda_core/cuda/core/experimental/_context.py renamed to cuda_core/cuda/core/experimental/_context.pyx

@@ -13,16 +13,21 @@ class ContextOptions:
     pass  # TODO
 
 
-class Context:
-    __slots__ = ("_handle", "_id")
+cdef class Context:
 
-    def __new__(self, *args, **kwargs):
+    cdef:
+        readonly object _handle
+        int _device_id
+
+    def __init__(self, *args, **kwargs):
         raise RuntimeError("Context objects cannot be instantiated directly. Please use Device or Stream APIs.")
 
     @classmethod
-    def _from_ctx(cls, obj, dev_id):
-        assert_type(obj, driver.CUcontext)
-        ctx = super().__new__(cls)
-        ctx._handle = obj
-        ctx._id = dev_id
+    def _from_ctx(cls, handle: driver.CUcontext, int device_id):
+        cdef Context ctx = Context.__new__(Context)
+        ctx._handle = handle
+        ctx._device_id = device_id
         return ctx
+
+    def __eq__(self, other):
+        return int(self._handle) == int(other._handle)

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,8 +1222,7 @@ 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:
+    def create_stream(self, obj: Optional[IsStreamT] = None, options: Optional[StreamOptions] = None) -> Stream:
         """Create a Stream object.
 
         New stream objects can be created in two different ways:
@@ -1235,9 +1250,9 @@ def create_stream(self, obj: Optional[IsStreamT] = None, options: StreamOptions
             Newly created stream object.
 
         """
-        return Stream._init(obj=obj, options=options)
+        self._check_context_initialized()
+        return Stream._init(obj=obj, options=options, device_id=self._id)
 
-    @precondition(_check_context_initialized)
     def create_event(self, options: Optional[EventOptions] = None) -> Event:
         """Create an Event object without recording it to a Stream.
 
@@ -1256,9 +1271,10 @@ def create_event(self, options: Optional[EventOptions] = None) -> Event:
             Newly created event object.
 
         """
-        return Event._init(self._id, self.context._handle, options)
+        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.
 
@@ -1285,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.
 
@@ -1298,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.
 
@@ -1310,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/_event.py renamed to cuda_core/cuda/core/experimental/_event.pyx

@@ -4,28 +4,28 @@
 
 from __future__ import annotations
 
-import weakref
+from cuda.core.experimental._utils.cuda_utils cimport (
+    _check_driver_error as raise_if_driver_error,
+    check_or_create_options,
+)
+
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Optional
 
 from cuda.core.experimental._context import Context
 from cuda.core.experimental._utils.cuda_utils import (
     CUDAError,
-    check_or_create_options,
     driver,
     handle_return,
 )
-from cuda.core.experimental._utils.cuda_utils import (
-    _check_driver_error as raise_if_driver_error,
-)
 
 if TYPE_CHECKING:
     import cuda.bindings
     from cuda.core.experimental._device import Device
 
 
 @dataclass
-class EventOptions:
+cdef class EventOptions:
     """Customizable :obj:`~_event.Event` options.
 
     Attributes
@@ -49,7 +49,7 @@ class EventOptions:
     support_ipc: Optional[bool] = False
 
 
-class Event:
+cdef class Event:
     """Represent a record at a specific point of execution within a CUDA stream.
 
     Applications can asynchronously record events at any point in
@@ -77,49 +77,46 @@ class Event:
     and they should instead be created through a :obj:`~_stream.Stream` object.
 
     """
-
-    class _MembersNeededForFinalize:
-        __slots__ = ("handle",)
-
-        def __init__(self, event_obj, handle):
-            self.handle = handle
-            weakref.finalize(event_obj, self.close)
-
-        def close(self):
-            if self.handle is not None:
-                handle_return(driver.cuEventDestroy(self.handle))
-                self.handle = None
-
-    def __new__(self, *args, **kwargs):
+    cdef:
+        object _handle
+        bint _timing_disabled
+        bint _busy_waited
+        int _device_id
+        object _ctx_handle
+
+    def __init__(self, *args, **kwargs):
         raise RuntimeError("Event objects cannot be instantiated directly. Please use Stream APIs (record).")
 
-    __slots__ = ("__weakref__", "_mnff", "_timing_disabled", "_busy_waited", "_device_id", "_ctx_handle")
-
     @classmethod
-    def _init(cls, device_id: int, ctx_handle: Context, options: Optional[EventOptions] = None):
-        self = super().__new__(cls)
-        self._mnff = Event._MembersNeededForFinalize(self, None)
-
-        options = check_or_create_options(EventOptions, options, "Event options")
+    def _init(cls, device_id: int, ctx_handle: Context, options=None):
+        cdef Event self = Event.__new__(Event)
+        cdef EventOptions opts = check_or_create_options(EventOptions, options, "Event options")
         flags = 0x0
         self._timing_disabled = False
         self._busy_waited = False
-        if not options.enable_timing:
+        if not opts.enable_timing:
             flags |= driver.CUevent_flags.CU_EVENT_DISABLE_TIMING
             self._timing_disabled = True
-        if options.busy_waited_sync:
+        if opts.busy_waited_sync:
             flags |= driver.CUevent_flags.CU_EVENT_BLOCKING_SYNC
             self._busy_waited = True
-        if options.support_ipc:
+        if opts.support_ipc:
             raise NotImplementedError("WIP: https://github.com/NVIDIA/cuda-python/issues/103")
-        self._mnff.handle = handle_return(driver.cuEventCreate(flags))
+        err, self._handle = driver.cuEventCreate(flags)
+        raise_if_driver_error(err)
         self._device_id = device_id
         self._ctx_handle = ctx_handle
         return self
 
-    def close(self):
+    cpdef close(self):
         """Destroy the event."""
-        self._mnff.close()
+        if self._handle is not None:
+            err, = driver.cuEventDestroy(self._handle)
+            self._handle = None
+            raise_if_driver_error(err)
+
+    def __del__(self):
+        self.close()
 
     def __isub__(self, other):
         return NotImplemented
@@ -129,7 +126,7 @@ def __rsub__(self, other):
 
     def __sub__(self, other):
         # return self - other (in milliseconds)
-        err, timing = driver.cuEventElapsedTime(other.handle, self.handle)
+        err, timing = driver.cuEventElapsedTime(other.handle, self._handle)
         try:
             raise_if_driver_error(err)
             return timing
@@ -180,12 +177,12 @@ def sync(self):
         has been completed.
 
         """
-        handle_return(driver.cuEventSynchronize(self._mnff.handle))
+        handle_return(driver.cuEventSynchronize(self._handle))
 
     @property
     def is_done(self) -> bool:
         """Return True if all captured works have been completed, otherwise False."""
-        (result,) = driver.cuEventQuery(self._mnff.handle)
+        result, = driver.cuEventQuery(self._handle)
         if result == driver.CUresult.CUDA_SUCCESS:
             return True
         if result == driver.CUresult.CUDA_ERROR_NOT_READY:
@@ -201,7 +198,7 @@ def handle(self) -> cuda.bindings.driver.CUevent:
             This handle is a Python object. To get the memory address of the underlying C
             handle, call ``int(Event.handle)``.
         """
-        return self._mnff.handle
+        return self._handle
 
     @property
     def device(self) -> Device: