cythonize stream + bug fixes

Author: leofang

cuda_core/cuda/core/experimental/_context.pyx

@@ -16,7 +16,7 @@ class ContextOptions:
 cdef class Context:
 
     cdef:
-        object _handle
+        readonly object _handle
         int _device_id
 
     def __init__(self, *args, **kwargs):

cuda_core/cuda/core/experimental/_device.py

@@ -1251,7 +1251,7 @@ def create_stream(self, obj: Optional[IsStreamT] = None, options: StreamOptions
 
         """
         self._check_context_initialized()
-        return Stream._init(obj=obj, options=options)
+        return Stream._init(obj=obj, options=options, device_id=self._id)
 
     def create_event(self, options: Optional[EventOptions] = None) -> Event:
         """Create an Event object without recording it to a Stream.

cuda_core/cuda/core/experimental/_event.pyx

@@ -88,19 +88,19 @@ cdef class Event:
         raise RuntimeError("Event objects cannot be instantiated directly. Please use Stream APIs (record).")
 
     @classmethod
-    def _init(cls, device_id: int, ctx_handle: Context, opts=None):
+    def _init(cls, device_id: int, ctx_handle: Context, options: Optional[EventOptions] = None):
         cdef Event self = Event.__new__(Event)
-        cdef EventOptions options = check_or_create_options(EventOptions, opts, "Event options")
+        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")
         err, self._handle = driver.cuEventCreate(flags)
         raise_if_driver_error(err)

…a_core/cuda/core/experimental/_stream.py …_core/cuda/core/experimental/_stream.pyxcuda_core/cuda/core/experimental/_stream.py renamed to cuda_core/cuda/core/experimental/_stream.pyx cuda_core/cuda/core/experimental/_stream.py renamed to cuda_core/cuda/core/experimental/_stream.pyx

@@ -4,9 +4,13 @@
 
 from __future__ import annotations
 
+from cuda.core.experimental._utils.cuda_utils cimport (
+    _check_driver_error as raise_if_driver_error,
+    check_or_create_options,
+)
+
 import os
 import warnings
-import weakref
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Optional, Protocol, Tuple, Union
 
@@ -18,16 +22,14 @@
 from cuda.core.experimental._graph import GraphBuilder
 from cuda.core.experimental._utils.clear_error_support import assert_type
 from cuda.core.experimental._utils.cuda_utils import (
-    check_or_create_options,
     driver,
     get_device_from_ctx,
     handle_return,
-    runtime,
 )
 
 
 @dataclass
-class StreamOptions:
+cdef class StreamOptions:
     """Customizable :obj:`~_stream.Stream` options.
 
     Attributes
@@ -85,7 +87,7 @@ def _try_to_get_stream_ptr(obj: IsStreamT):
     return driver.CUstream(info[1])
 
 
-class Stream:
+cdef class Stream:
     """Represent a queue of GPU operations that are executed in a specific order.
 
     Applications use streams to control the order of execution for
@@ -103,35 +105,27 @@ class Stream:
 
     """
 
-    class _MembersNeededForFinalize:
-        __slots__ = ("handle", "owner", "builtin")
-
-        def __init__(self, stream_obj, handle, owner, builtin):
-            self.handle = handle
-            self.owner = owner
-            self.builtin = builtin
-            weakref.finalize(stream_obj, self.close)
+    cdef:
+        object _handle
+        object _owner
+        object _builtin
+        object _nonblocking
+        object _priority
+        object _device_id
+        object _ctx_handle
 
-        def close(self):
-            if self.owner is None:
-                if self.handle and not self.builtin:
-                    handle_return(driver.cuStreamDestroy(self.handle))
-            else:
-                self.owner = None
-            self.handle = None
-
-    def __new__(self, *args, **kwargs):
+    def __init__(self, *args, **kwargs):
         raise RuntimeError(
             "Stream objects cannot be instantiated directly. "
             "Please use Device APIs (create_stream) or other Stream APIs (from_handle)."
         )
 
-    __slots__ = ("__weakref__", "_mnff", "_nonblocking", "_priority", "_device_id", "_ctx_handle")
-
     @classmethod
     def _legacy_default(cls):
-        self = super().__new__(cls)
-        self._mnff = Stream._MembersNeededForFinalize(self, driver.CUstream(driver.CU_STREAM_LEGACY), None, True)
+        cdef Stream self = Stream.__new__(Stream)
+        self._handle = driver.CUstream(driver.CU_STREAM_LEGACY)
+        self._owner = None
+        self._builtin = True
         self._nonblocking = None  # delayed
         self._priority = None  # delayed
         self._device_id = None  # delayed
@@ -140,66 +134,76 @@ def _legacy_default(cls):
 
     @classmethod
     def _per_thread_default(cls):
-        self = super().__new__(cls)
-        self._mnff = Stream._MembersNeededForFinalize(self, driver.CUstream(driver.CU_STREAM_PER_THREAD), None, True)
+        cdef Stream self = Stream.__new__(Stream)
+        self._handle = driver.CUstream(driver.CU_STREAM_PER_THREAD)
+        self._owner = None
+        self._builtin = True
         self._nonblocking = None  # delayed
         self._priority = None  # delayed
         self._device_id = None  # delayed
         self._ctx_handle = None  # delayed
         return self
 
     @classmethod
-    def _init(cls, obj: Optional[IsStreamT] = None, *, options: Optional[StreamOptions] = None):
-        self = super().__new__(cls)
-        self._mnff = Stream._MembersNeededForFinalize(self, None, None, False)
+    def _init(cls, obj: Optional[IsStreamT] = None, *, options: Optional[StreamOptions] = None, device_id: int = None):
+        cdef Stream self = Stream.__new__(Stream)
+        self._handle = None
+        self._owner = None
+        self._builtin = False
 
         if obj is not None and options is not None:
             raise ValueError("obj and options cannot be both specified")
         if obj is not None:
-            self._mnff.handle = _try_to_get_stream_ptr(obj)
+            self._handle = _try_to_get_stream_ptr(obj)
             # TODO: check if obj is created under the current context/device
-            self._mnff.owner = obj
+            self._owner = obj
             self._nonblocking = None  # delayed
             self._priority = None  # delayed
             self._device_id = None  # delayed
             self._ctx_handle = None  # delayed
             return self
 
-        options = check_or_create_options(StreamOptions, options, "Stream options")
-        nonblocking = options.nonblocking
-        priority = options.priority
+        cdef StreamOptions opts = check_or_create_options(StreamOptions, options, "Stream options")
+        nonblocking = opts.nonblocking
+        priority = opts.priority
 
         flags = driver.CUstream_flags.CU_STREAM_NON_BLOCKING if nonblocking else driver.CUstream_flags.CU_STREAM_DEFAULT
-
-        high, low = handle_return(runtime.cudaDeviceGetStreamPriorityRange())
+        err, high, low = driver.cuCtxGetStreamPriorityRange()
+        raise_if_driver_error(err)
         if priority is not None:
             if not (low <= priority <= high):
                 raise ValueError(f"{priority=} is out of range {[low, high]}")
         else:
             priority = high
 
-        self._mnff.handle = handle_return(driver.cuStreamCreateWithPriority(flags, priority))
-        self._mnff.owner = None
+        self._handle = handle_return(driver.cuStreamCreateWithPriority(flags, priority))
+        self._owner = None
         self._nonblocking = nonblocking
         self._priority = priority
-        # don't defer this because we will have to pay a cost for context
-        # switch later
-        self._device_id = int(handle_return(driver.cuCtxGetDevice()))
+        self._device_id = device_id
         self._ctx_handle = None  # delayed
         return self
 
-    def close(self):
+    def __del__(self):
+        self.close()
+
+    cpdef close(self):
         """Destroy the stream.
 
         Destroy the stream if we own it. Borrowed foreign stream
         object will instead have their references released.
 
         """
-        self._mnff.close()
+        if self._owner is None:
+            if self._handle and not self._builtin:
+                handle_return(driver.cuStreamDestroy(self._handle))
+        else:
+            self._owner = None
+        self._handle = None
 
     def __cuda_stream__(self) -> Tuple[int, int]:
         """Return an instance of a __cuda_stream__ protocol."""
-        return (0, self.handle)
+        return (0, int(self.handle))
 
     @property
     def handle(self) -> cuda.bindings.driver.CUstream:
@@ -210,13 +214,13 @@ def handle(self) -> cuda.bindings.driver.CUstream:
             This handle is a Python object. To get the memory address of the underlying C
             handle, call ``int(Stream.handle)``.
         """
-        return self._mnff.handle
+        return self._handle
 
     @property
     def is_nonblocking(self) -> bool:
         """Return True if this is a nonblocking stream, otherwise False."""
         if self._nonblocking is None:
-            flag = handle_return(driver.cuStreamGetFlags(self._mnff.handle))
+            flag = handle_return(driver.cuStreamGetFlags(self._handle))
             if flag == driver.CUstream_flags.CU_STREAM_NON_BLOCKING:
                 self._nonblocking = True
             else:
@@ -227,13 +231,13 @@ def is_nonblocking(self) -> bool:
     def priority(self) -> int:
         """Return the stream priority."""
         if self._priority is None:
-            prio = handle_return(driver.cuStreamGetPriority(self._mnff.handle))
+            prio = handle_return(driver.cuStreamGetPriority(self._handle))
             self._priority = prio
         return self._priority
 
     def sync(self):
         """Synchronize the stream."""
-        handle_return(driver.cuStreamSynchronize(self._mnff.handle))
+        handle_return(driver.cuStreamSynchronize(self._handle))
 
     def record(self, event: Event = None, options: EventOptions = None) -> Event:
         """Record an event onto the stream.
@@ -259,8 +263,8 @@ def record(self, event: Event = None, options: EventOptions = None) -> Event:
         # and CU_EVENT_RECORD_EXTERNAL, can be set in EventOptions.
         if event is None:
             event = Event._init(self._device_id, self._ctx_handle, options)
-        assert_type(event, Event)
-        handle_return(driver.cuEventRecord(event.handle, self._mnff.handle))
+        err, = driver.cuEventRecord(event.handle, self._handle)
+        raise_if_driver_error(err)
         return event
 
     def wait(self, event_or_stream: Union[Event, Stream]):
@@ -281,7 +285,7 @@ def wait(self, event_or_stream: Union[Event, Stream]):
                 stream = event_or_stream
             else:
                 try:
-                    stream = Stream._init(event_or_stream)
+                    stream = Stream._init(obj=event_or_stream)
                 except Exception as e:
                     raise ValueError(
                         "only an Event, Stream, or object supporting __cuda_stream__ can be waited,"
@@ -292,7 +296,7 @@ def wait(self, event_or_stream: Union[Event, Stream]):
             discard_event = True
 
         # TODO: support flags other than 0?
-        handle_return(driver.cuStreamWaitEvent(self._mnff.handle, event, 0))
+        handle_return(driver.cuStreamWaitEvent(self._handle, event, 0))
         if discard_event:
             handle_return(driver.cuEventDestroy(event))
 
@@ -308,21 +312,22 @@ def device(self) -> Device:
 
         """
         from cuda.core.experimental._device import Device  # avoid circular import
+        self._get_device_and_context()
+        return Device(self._device_id)
 
+    cdef _get_device_and_context(self):
+        # Get the stream context first
+        if self._ctx_handle is None:
+            err, self._ctx_handle = driver.cuStreamGetCtx(self._handle)
+            raise_if_driver_error(err)
         if self._device_id is None:
-            # Get the stream context first
-            if self._ctx_handle is None:
-                self._ctx_handle = handle_return(driver.cuStreamGetCtx(self._mnff.handle))
             self._device_id = get_device_from_ctx(self._ctx_handle)
-        return Device(self._device_id)
+            raise_if_driver_error(err)
 
     @property
     def context(self) -> Context:
         """Return the :obj:`~_context.Context` associated with this stream."""
-        if self._ctx_handle is None:
-            self._ctx_handle = handle_return(driver.cuStreamGetCtx(self._mnff.handle))
-        if self._device_id is None:
-            self._device_id = get_device_from_ctx(self._ctx_handle)
+        self._get_device_and_context()
         return Context._from_ctx(self._ctx_handle, self._device_id)
 
     @staticmethod

cuda_core/tests/test_cuda_utils.py

@@ -44,7 +44,7 @@ def test_check_driver_error():
     num_unexpected = 0
     for error in driver.CUresult:
         if error == driver.CUresult.CUDA_SUCCESS:
-            assert cuda_utils._check_driver_error(error) is None
+            assert cuda_utils._check_driver_error(error) == 0
         else:
             with pytest.raises(cuda_utils.CUDAError) as e:
                 cuda_utils._check_driver_error(error)
@@ -63,7 +63,7 @@ def test_check_runtime_error():
     num_unexpected = 0
     for error in runtime.cudaError_t:
         if error == runtime.cudaError_t.cudaSuccess:
-            assert cuda_utils._check_runtime_error(error) is None
+            assert cuda_utils._check_runtime_error(error) == 0
         else:
             with pytest.raises(cuda_utils.CUDAError) as e:
                 cuda_utils._check_runtime_error(error)

cuda_core/tests/test_stream.py

@@ -52,7 +52,7 @@ def test_stream_record(init_cuda):
 
 def test_stream_record_invalid_event(init_cuda):
     stream = Device().create_stream(options=StreamOptions())
-    with pytest.raises(TypeError):
+    with pytest.raises(AttributeError):
         stream.record(event="invalid_event")