Implement device and context properties for Event (#618)

* Implement device and context properties for Event

* Fix documentation typo

* Address comments

* Add new attributes to slots

* Add missing import

* Change order of arguments to preserve defaults

* Always pass ctx_handle to event constructor

* Fix order of passed arguments to constructor

* Pass context handle instead of object

* Update release notes

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, self.context._handle, options)
 
     @precondition(_check_context_initialized)
     def allocate(self, size, stream=None) -> Buffer:

cuda_core/cuda/core/experimental/_event.py

@@ -8,6 +8,7 @@
 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,
@@ -20,6 +21,7 @@
 
 if TYPE_CHECKING:
     import cuda.bindings
+    from cuda.core.experimental._device import Device
 
 
 @dataclass
@@ -91,10 +93,10 @@ def close(self):
     def __new__(self, *args, **kwargs):
         raise RuntimeError("Event objects cannot be instantiated directly. Please use Stream APIs (record).")
 
-    __slots__ = ("__weakref__", "_mnff", "_timing_disabled", "_busy_waited")
+    __slots__ = ("__weakref__", "_mnff", "_timing_disabled", "_busy_waited", "_device_id", "_ctx_handle")
 
     @classmethod
-    def _init(cls, options: Optional[EventOptions] = None):
+    def _init(cls, device_id: int, ctx_handle: Context, options: Optional[EventOptions] = None):
         self = super().__new__(cls)
         self._mnff = Event._MembersNeededForFinalize(self, None)
 
@@ -111,6 +113,8 @@ 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 = device_id
+        self._ctx_handle = ctx_handle
         return self
 
     def close(self):
@@ -198,3 +202,24 @@ def handle(self) -> cuda.bindings.driver.CUevent:
             handle, call ``int(Event.handle)``.
         """
         return self._mnff.handle
+
+    @property
+    def device(self) -> Device:
+        """Return the :obj:`~_device.Device` singleton associated with this event.
+
+        Note
+        ----
+        The current context on the device may differ from this
+        event's context. This case occurs when a different CUDA
+        context is set current after a event is created.
+
+        """
+
+        from cuda.core.experimental._device import Device  # avoid circular import
+
+        return Device(self._device_id)
+
+    @property
+    def context(self) -> Context:
+        """Return the :obj:`~_context.Context` associated with this event."""
+        return Context._from_ctx(self._ctx_handle, self._device_id)

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/docs/source/release/0.3.0-notes.rst

@@ -28,3 +28,5 @@ New examples
 
 Fixes and enhancements
 ----------------------
+
+- An :class:`Event` can now be used to look up its corresponding device and context using the ``.device`` and ``.context`` attributes respectively.

cuda_core/tests/test_event.py

@@ -169,3 +169,15 @@ def test_error_timing_incomplete():
     arr[0] = 1
     event3.sync()
     event3 - event1  # this should work
+
+
+def test_event_device(init_cuda):
+    device = Device()
+    event = device.create_event(options=EventOptions())
+    assert event.device is device
+
+
+def test_event_context(init_cuda):
+    event = Device().create_event(options=EventOptions())
+    context = event.context
+    assert context is not None