Add hash and equality support to Stream, Event, Context, and Device classes, enabling their use as dictionary keys and in sets.

Author: rparolin

cuda_core/cuda/core/experimental/_context.pyx

@@ -29,4 +29,35 @@ cdef class Context:
         return ctx
 
     def __eq__(self, other):
+        """Check equality based on the underlying CUcontext handle address.
+        
+        Two Context objects are considered equal if they wrap the same
+        underlying CUDA context.
+        
+        Parameters
+        ----------
+        other : object
+            Another object to compare with.
+        
+        Returns
+        -------
+        bool
+            True if other is a Context wrapping the same handle, False otherwise.
+        """
+        if not isinstance(other, Context):
+            return NotImplemented
         return int(self._handle) == int(other._handle)
+
+    def __hash__(self) -> int:
+        """Return hash based on the underlying CUcontext handle address.
+        
+        This enables Context objects to be used as dictionary keys and in sets.
+        Two Context objects wrapping the same underlying CUDA context will hash
+        to the same value and be considered equal.
+        
+        Returns
+        -------
+        int
+            Hash value based on the context handle address.
+        """
+        return hash(int(self._handle))

cuda_core/cuda/core/experimental/_device.pyx

@@ -1145,6 +1145,49 @@ class Device:
     def __repr__(self):
         return f"<Device {self._id} ({self.name})>"
 
+    def __hash__(self) -> int:
+        """Return hash based on the device ordinal.
+        
+        This enables Device objects to be used as dictionary keys and in sets.
+        Device objects with the same device_id will hash to the same value
+        and be considered equal, even if they are different Python objects or
+        exist on different threads.
+        
+        Returns
+        -------
+        int
+            Hash value based on the device ordinal (device_id).
+        
+        Notes
+        -----
+        Device is a per-thread singleton, but equality is based on logical
+        device identity (device_id), not Python object identity. This means
+        Device(0) on thread A equals Device(0) on thread B.
+        """
+        return hash(self._id)
+
+    def __eq__(self, other) -> bool:
+        """Check equality based on the device ordinal.
+        
+        Two Device objects are considered equal if they have the same device_id,
+        regardless of whether they are the same Python object or exist on
+        different threads.
+        
+        Parameters
+        ----------
+        other : object
+            Another object to compare with.
+        
+        Returns
+        -------
+        bool
+            True if other is a Device with the same device_id, False otherwise.
+            Returns NotImplemented if other is not a Device.
+        """
+        if not isinstance(other, Device):
+            return NotImplemented
+        return self._id == other._id
+
     def __reduce__(self):
         return Device, (self.device_id,)
 

cuda_core/cuda/core/experimental/_event.pyx

@@ -165,6 +165,48 @@ cdef class Event:
                 raise CUDAError(err)
             raise RuntimeError(explanation)
 
+    def __hash__(self) -> int:
+        """Return hash based on the underlying CUevent handle address.
+        
+        This enables Event objects to be used as dictionary keys and in sets.
+        Two Event objects wrapping the same underlying CUDA event will hash
+        to the same value and be considered equal.
+        
+        Returns
+        -------
+        int
+            Hash value based on the event handle address.
+        
+        Warning
+        -------
+        Using a closed or destroyed event as a dictionary key or in a set
+        results in undefined behavior. The event handle may be reused by
+        the CUDA driver for new events.
+        """
+        return hash(<uintptr_t>(self._handle))
+
+    def __eq__(self, other) -> bool:
+        """Check equality based on the underlying CUevent handle address.
+        
+        Two Event objects are considered equal if they wrap the same
+        underlying CUDA event, regardless of whether they are the same
+        Python object.
+        
+        Parameters
+        ----------
+        other : object
+            Another object to compare with.
+        
+        Returns
+        -------
+        bool
+            True if other is an Event wrapping the same handle, False otherwise.
+            Returns NotImplemented if other is not an Event.
+        """
+        if not isinstance(other, Event):
+            return NotImplemented
+        return <uintptr_t>(self._handle) == <uintptr_t>((<Event>other)._handle)
+
     def get_ipc_descriptor(self) -> IPCEventDescriptor:
         """Export an event allocated for sharing between processes."""
         if self._ipc_descriptor is not None:

cuda_core/cuda/core/experimental/_stream.pyx

@@ -197,6 +197,48 @@ cdef class Stream:
         """Return an instance of a __cuda_stream__ protocol."""
         return (0, <uintptr_t>(self._handle))
 
+    def __hash__(self) -> int:
+        """Return hash based on the underlying CUstream handle address.
+        
+        This enables Stream objects to be used as dictionary keys and in sets.
+        Two Stream objects wrapping the same underlying CUDA stream will hash
+        to the same value and be considered equal.
+        
+        Returns
+        -------
+        int
+            Hash value based on the stream handle address.
+        
+        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(<uintptr_t>(self._handle))
+
+    def __eq__(self, other) -> bool:
+        """Check equality based on the underlying CUstream handle address.
+        
+        Two Stream objects are considered equal if they wrap the same
+        underlying CUDA stream, regardless of whether they are the same
+        Python object.
+        
+        Parameters
+        ----------
+        other : object
+            Another object to compare with.
+        
+        Returns
+        -------
+        bool
+            True if other is a Stream wrapping the same handle, False otherwise.
+            Returns NotImplemented if other is not a Stream.
+        """
+        if not isinstance(other, Stream):
+            return NotImplemented
+        return <uintptr_t>(self._handle) == <uintptr_t>((<Stream>other)._handle)
+
     @property
     def handle(self) -> cuda.bindings.driver.CUstream:
         """Return the underlying ``CUstream`` object.