88from dataclasses import dataclass
99from typing import TYPE_CHECKING , Optional
1010
11+ from cuda .core .experimental ._context import Context
1112from cuda .core .experimental ._utils .cuda_utils import (
1213 CUDAError ,
1314 check_or_create_options ,
2021
2122if TYPE_CHECKING :
2223 import cuda .bindings
24+ from cuda .core .experimental ._device import Device
2325
2426
2527@dataclass
@@ -91,10 +93,10 @@ def close(self):
9193 def __new__ (self , * args , ** kwargs ):
9294 raise RuntimeError ("Event objects cannot be instantiated directly. Please use Stream APIs (record)." )
9395
94- __slots__ = ("__weakref__" , "_mnff" , "_timing_disabled" , "_busy_waited" )
96+ __slots__ = ("__weakref__" , "_mnff" , "_timing_disabled" , "_busy_waited" , "_device_id" , "_ctx_handle" )
9597
9698 @classmethod
97- def _init (cls , options : Optional [EventOptions ] = None ):
99+ def _init (cls , device_id : int , ctx_handle : Context , options : Optional [EventOptions ] = None ):
98100 self = super ().__new__ (cls )
99101 self ._mnff = Event ._MembersNeededForFinalize (self , None )
100102
@@ -111,6 +113,8 @@ def _init(cls, options: Optional[EventOptions] = None):
111113 if options .support_ipc :
112114 raise NotImplementedError ("WIP: https://github.com/NVIDIA/cuda-python/issues/103" )
113115 self ._mnff .handle = handle_return (driver .cuEventCreate (flags ))
116+ self ._device_id = device_id
117+ self ._ctx_handle = ctx_handle
114118 return self
115119
116120 def close (self ):
@@ -198,3 +202,24 @@ def handle(self) -> cuda.bindings.driver.CUevent:
198202 handle, call ``int(Event.handle)``.
199203 """
200204 return self ._mnff .handle
205+
206+ @property
207+ def device (self ) -> Device :
208+ """Return the :obj:`~_device.Device` singleton associated with this event.
209+
210+ Note
211+ ----
212+ The current context on the device may differ from this
213+ event's context. This case occurs when a different CUDA
214+ context is set current after a event is created.
215+
216+ """
217+
218+ from cuda .core .experimental ._device import Device # avoid circular import
219+
220+ return Device (self ._device_id )
221+
222+ @property
223+ def context (self ) -> Context :
224+ """Return the :obj:`~_context.Context` associated with this event."""
225+ return Context ._from_ctx (self ._ctx_handle , self ._device_id )
0 commit comments