Address comments

Author: NaderAlAwar

cuda_core/cuda/core/experimental/_device.py

@@ -1228,7 +1228,7 @@ def create_event(self, options: Optional[EventOptions] = None) -> Event:
             Newly created event object.
 
         """
-        return Event._init(options)
+        return Event._init(self._id, options)
 
     @precondition(_check_context_initialized)
     def allocate(self, size, stream=None) -> Buffer:

cuda_core/cuda/core/experimental/_event.py

@@ -9,7 +9,6 @@
 from typing import TYPE_CHECKING, Optional
 
 from cuda.core.experimental._context import Context
-from cuda.core.experimental._device import Device
 from cuda.core.experimental._utils.cuda_utils import (
     CUDAError,
     check_or_create_options,
@@ -22,6 +21,7 @@
 
 if TYPE_CHECKING:
     import cuda.bindings
+    from cuda.core.experimental._device import Device
 
 
 @dataclass
@@ -96,7 +96,7 @@ def __new__(self, *args, **kwargs):
     __slots__ = ("__weakref__", "_mnff", "_timing_disabled", "_busy_waited")
 
     @classmethod
-    def _init(cls, options: Optional[EventOptions] = None):
+    def _init(cls, device_id: int, ctx_handle=None, options: Optional[EventOptions] = None):
         self = super().__new__(cls)
         self._mnff = Event._MembersNeededForFinalize(self, None)
 
@@ -113,8 +113,11 @@ def _init(cls, options: Optional[EventOptions] = None):
         if options.support_ipc:
             raise NotImplementedError("WIP: https://github.com/NVIDIA/cuda-python/issues/103")
         self._mnff.handle = handle_return(driver.cuEventCreate(flags))
-        self._device_id = int(handle_return(driver.cuCtxGetDevice()))
-        self._ctx_handle = handle_return(driver.cuStreamGetCtx(self._mnff.handle))
+        self._device_id = device_id
+        if ctx_handle is not None:
+            self._ctx_handle = ctx_handle
+        else:
+            self._ctx_handle = handle_return(driver.cuCtxGetCurrent())
         return self
 
     def close(self):

cuda_core/cuda/core/experimental/_stream.py

@@ -244,7 +244,7 @@ def record(self, event: Event = None, options: EventOptions = None) -> Event:
         # on the stream. Event flags such as disabling timing, nonblocking,
         # and CU_EVENT_RECORD_EXTERNAL, can be set in EventOptions.
         if event is None:
-            event = Event._init(options)
+            event = Event._init(self._device_id, self._ctx_handle, options)
         assert_type(event, Event)
         handle_return(driver.cuEventRecord(event.handle, self._mnff.handle))
         return event

cuda_core/tests/test_event.py

@@ -177,9 +177,9 @@ def test_error_timing_incomplete():
 
 
 def test_event_device(init_cuda):
-    event = Device().create_event(options=EventOptions())
-    device = event.device
-    assert isinstance(device, Device)
+    device = Device()
+    event = device.create_event(options=EventOptions())
+    assert event.device is device
 
 
 def test_event_context(init_cuda):