Adding context to the stream hash implementation

Author: rparolin

cuda_core/cuda/core/experimental/_event.pyx

@@ -233,16 +233,20 @@ cdef class Event:
     def from_ipc_descriptor(cls, ipc_descriptor: IPCEventDescriptor) -> Event:
         """Import an event that was exported from another process."""
         cdef cydriver.CUipcEventHandle data
+        cdef cydriver.CUcontext curr_ctx
+        cdef cydriver.CUdevice dev
         memcpy(data.reserved, <const void*><const char*>(ipc_descriptor._reserved), sizeof(data.reserved))
         cdef Event self = Event.__new__(cls)
         with nogil:
             HANDLE_RETURN(cydriver.cuIpcOpenEventHandle(&self._handle, data))
+            HANDLE_RETURN(cydriver.cuCtxGetCurrent(&curr_ctx))
+            HANDLE_RETURN(cydriver.cuCtxGetDevice(&dev))
         self._timing_disabled = True
         self._busy_waited = ipc_descriptor._busy_waited
         self._ipc_enabled = True
         self._ipc_descriptor = ipc_descriptor
-        self._device_id = -1  # ??
-        self._ctx_handle = None  # ??
+        self._device_id = <int>dev
+        self._ctx_handle = driver.CUcontext(<uintptr_t>curr_ctx)
         return self
 
     @property

cuda_core/cuda/core/experimental/_stream.pyx

@@ -126,15 +126,23 @@ cdef class Stream:
     @classmethod
     def _legacy_default(cls):
         cdef Stream self = Stream.__new__(cls)
+        cdef cydriver.CUcontext ctx
         self._handle = <cydriver.CUstream>(cydriver.CU_STREAM_LEGACY)
         self._builtin = True
+        with nogil:
+            HANDLE_RETURN(cydriver.cuCtxGetCurrent(&ctx))
+        self._ctx_handle = ctx
         return self
 
     @classmethod
     def _per_thread_default(cls):
         cdef Stream self = Stream.__new__(cls)
+        cdef cydriver.CUcontext ctx
         self._handle = <cydriver.CUstream>(cydriver.CU_STREAM_PER_THREAD)
         self._builtin = True
+        with nogil:
+            HANDLE_RETURN(cydriver.cuCtxGetCurrent(&ctx))
+        self._ctx_handle = ctx
         return self
 
     @classmethod
@@ -144,9 +152,16 @@ cdef class Stream:
         if obj is not None and options is not None:
             raise ValueError("obj and options cannot be both specified")
         if obj is not None:
+            cdef cydriver.CUcontext ctx
+            cdef cydriver.CUresult err
             self._handle = _try_to_get_stream_ptr(obj)
             # TODO: check if obj is created under the current context/device
             self._owner = obj
+            with nogil:
+                err = cydriver.cuStreamGetCtx(self._handle, &ctx)
+                if err != cydriver.CUresult.CUDA_SUCCESS:
+                    HANDLE_RETURN(cydriver.cuCtxGetCurrent(&ctx))
+            self._ctx_handle = ctx
             return self
 
         cdef StreamOptions opts = check_or_create_options(StreamOptions, options, "Stream options")
@@ -167,12 +182,15 @@ cdef class Stream:
             prio = high
 
         cdef cydriver.CUstream s
+        cdef cydriver.CUcontext ctx
         with nogil:
             HANDLE_RETURN(cydriver.cuStreamCreateWithPriority(&s, flags, prio))
+            HANDLE_RETURN(cydriver.cuStreamGetCtx(s, &ctx))
         self._handle = s
         self._nonblocking = int(nonblocking)
         self._priority = prio
         self._device_id = device_id if device_id is not None else self._device_id
+        self._ctx_handle = ctx
         return self
 
     def __dealloc__(self):
@@ -198,7 +216,7 @@ cdef class Stream:
         return (0, <uintptr_t>(self._handle))
 
     def __hash__(self) -> int:
-        """Return hash based on the underlying CUstream handle address.
+        """Return hash based on the underlying CUstream handle address and context.
 
         This enables Stream objects to be used as dictionary keys and in sets.
         Two Stream objects wrapping the same underlying CUDA stream will hash
@@ -207,15 +225,24 @@ cdef class Stream:
         Returns
         -------
         int
-            Hash value based on the stream handle address.
+            Hash value based on the stream handle address and context handle.
+
+        Notes
+        -----
+        Includes the context handle in the hash to prevent collisions when
+        handles are reused across different contexts. While handles are
+        context-scoped and typically not reused across contexts, including
+        the context provides defense-in-depth against hash collisions.
+        The context is fetched and cached during Stream construction.
 
         Warning
         -------
         Using a closed or destroyed stream as a dictionary key or in a set
         results in undefined behavior. The stream handle may be reused by
         the CUDA driver for new streams.
         """
-        return hash((type(self), <uintptr_t>(self._handle)))
+        # Context should always be set as a post-condition of Stream construction
+        return hash((type(self), <uintptr_t>(self._ctx_handle), <uintptr_t>(self._handle)))
 
     def __eq__(self, other) -> bool:
         """Check equality based on the underlying CUstream handle address.