Implement __hash__ and __eq__ for cuda.core classes (NVIDIA#1198)

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

* pre-commit fixes

* pre-commit fixes

* Update cuda_core/cuda/core/experimental/_event.pyx

Co-authored-by: Oleksandr Pavlyk <21087696+oleksandr-pavlyk@users.noreply.github.com>

* return False on type error

* reverting int conversion fix due to cython compilation error

* Addressing performance nit feedback

* Reverting back to returning NotImplemented on type error in __eq__ based on feedback.

* Adding  to the __hash__ function for all core types.

* Updating __eq__: Adding comment to explain perf optimization and ensuring we are doing a direct c-cast instead of isinstance

* Test organization: moving type specific tests to their corresponding files

* Device __hash__ function not uses uuid instead of device ordinal

* Adding context to the stream hash implementation

* forgot to include event changes

* cython compilation fixes

* Checking for context initalization

* Handle builtin streams for fetching context

* Fixing seg fault

* memoize uuid property

* memoize the uuid

* Removing unnecessary casts

* Making docstrings consistent with other dunder methods

* Address _event.pyx potential seg fault feedback when casting Event objects

* reverting change that prevented lazy init

* Adding an isinstance check for Stream.__eq__ to avoid seg faults

* Fixes equality and hash to make these consistent.  Both include context handle.

* Moving the uuid string generation code from a global function to an inline

* Pushing missing updates to Device.__eq__

* Addressing feedback and adding more unit tests that validate equality and hash behaviour is consistent.

* formatting

* pre-commit fixes

---------

Co-authored-by: Oleksandr Pavlyk <21087696+oleksandr-pavlyk@users.noreply.github.com>

Author: rparolin

cuda_core/cuda/core/experimental/_context.pyx

@@ -29,4 +29,10 @@ cdef class Context:
         return ctx
 
     def __eq__(self, other):
-        return int(self._handle) == int(other._handle)
+        if not isinstance(other, Context):
+            return NotImplemented
+        cdef Context _other = <Context>other
+        return int(self._handle) == int(_other._handle)
+
+    def __hash__(self) -> int:
+        return hash(int(self._handle))

cuda_core/cuda/core/experimental/_device.pyx

@@ -949,7 +949,7 @@ class Device:
         Default value of `None` return the currently used device.
 
     """
-    __slots__ = ("_id", "_mr", "_has_inited", "_properties")
+    __slots__ = ("_id", "_mr", "_has_inited", "_properties", "_uuid")
 
     def __new__(cls, device_id: int | None = None):
         global _is_cuInit
@@ -1002,6 +1002,7 @@ class Device:
 
                 device._has_inited = False
                 device._properties = None
+                device._uuid = None
                 devices.append(device)
 
         try:
@@ -1053,18 +1054,26 @@ class Device:
         MIG UUID is only returned when device is in MIG mode and the
         driver is older than CUDA 11.4.
 
+        The UUID is cached after first access to avoid repeated CUDA API calls.
+
         """
         cdef cydriver.CUuuid uuid
-        cdef cydriver.CUdevice this_dev = self._id
-        with nogil:
-            IF CUDA_CORE_BUILD_MAJOR == "12":
-                HANDLE_RETURN(cydriver.cuDeviceGetUuid_v2(&uuid, this_dev))
-            ELSE:  # 13.0+
-                HANDLE_RETURN(cydriver.cuDeviceGetUuid(&uuid, this_dev))
-        cdef bytes uuid_b = cpython.PyBytes_FromStringAndSize(uuid.bytes, sizeof(uuid.bytes))
-        cdef str uuid_hex = uuid_b.hex()
-        # 8-4-4-4-12
-        return f"{uuid_hex[:8]}-{uuid_hex[8:12]}-{uuid_hex[12:16]}-{uuid_hex[16:20]}-{uuid_hex[20:]}"
+        cdef cydriver.CUdevice dev
+        cdef bytes uuid_b
+        cdef str uuid_hex
+
+        if self._uuid is None:
+            dev = self._id
+            with nogil:
+                IF CUDA_CORE_BUILD_MAJOR == "12":
+                    HANDLE_RETURN(cydriver.cuDeviceGetUuid_v2(&uuid, dev))
+                ELSE:  # 13.0+
+                    HANDLE_RETURN(cydriver.cuDeviceGetUuid(&uuid, dev))
+            uuid_b = cpython.PyBytes_FromStringAndSize(uuid.bytes, sizeof(uuid.bytes))
+            uuid_hex = uuid_b.hex()
+            # 8-4-4-4-12
+            self._uuid = f"{uuid_hex[:8]}-{uuid_hex[8:12]}-{uuid_hex[12:16]}-{uuid_hex[16:20]}-{uuid_hex[20:]}"
+        return self._uuid
 
     @property
     def name(self) -> str:
@@ -1145,6 +1154,14 @@ class Device:
     def __repr__(self):
         return f"<Device {self._id} ({self.name})>"
 
+    def __hash__(self) -> int:
+        return hash(self.uuid)
+
+    def __eq__(self, other) -> bool:
+        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,16 @@ cdef class Event:
                 raise CUDAError(err)
             raise RuntimeError(explanation)
 
+    def __hash__(self) -> int:
+        return hash((self._ctx_handle, <uintptr_t>(self._handle)))
+
+    def __eq__(self, other) -> bool:
+        # Note: using isinstance because `Event` can be subclassed.
+        if not isinstance(other, Event):
+            return NotImplemented
+        cdef Event _other = <Event>other
+        return <uintptr_t>(self._handle) == <uintptr_t>(_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,27 @@ cdef class Stream:
         """Return an instance of a __cuda_stream__ protocol."""
         return (0, <uintptr_t>(self._handle))
 
+    def __hash__(self) -> int:
+        # Ensure context is initialized for hash consistency
+        if self._ctx_handle == CU_CONTEXT_INVALID:
+            self._get_context()
+        return hash((<uintptr_t>(self._ctx_handle), <uintptr_t>(self._handle)))
+
+    def __eq__(self, other) -> bool:
+        if not isinstance(other, Stream):
+            return NotImplemented
+        cdef Stream _other = <Stream>other
+        # Fast path: compare handles first
+        if <uintptr_t>(self._handle) != <uintptr_t>((_other)._handle):
+            return False
+        # Ensure contexts are initialized for both streams
+        if self._ctx_handle == CU_CONTEXT_INVALID:
+            self._get_context()
+        if _other._ctx_handle == CU_CONTEXT_INVALID:
+            _other._get_context()
+        # Compare contexts as well
+        return <uintptr_t>(self._ctx_handle) == <uintptr_t>((_other)._ctx_handle)
+
     @property
     def handle(self) -> cuda.bindings.driver.CUstream:
         """Return the underlying ``CUstream`` object.

cuda_core/tests/test_comparable.py

@@ -0,0 +1,175 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Tests for __eq__ and __ne__ implementations in cuda.core classes.
+
+These tests verify multi-type equality behavior and subclassing equality behavior
+across Device, Stream, Event, and Context objects.
+"""
+
+from cuda.core.experimental import Device, Stream
+from cuda.core.experimental._context import Context
+from cuda.core.experimental._event import Event, EventOptions
+from cuda.core.experimental._stream import StreamOptions
+
+# ============================================================================
+# Equality Contract Tests
+# ============================================================================
+
+
+def test_equality_is_not_identity():
+    """Test that equality (==) is different from identity (is)."""
+    device = Device(0)
+    device.set_current()
+
+    # Streams: Different objects can be equal
+    s1 = device.create_stream()
+    s2 = Stream.from_handle(int(s1.handle))
+
+    assert s1 == s2, "Streams with same handle are equal"
+    assert s1 is not s2, "But they are not the same object"
+
+    # Device: Same object due to singleton (special case)
+    d1 = Device(0)
+    d2 = Device(0)
+
+    assert d1 == d2, "Devices with same ID are equal"
+    assert d1 is d2, "And they ARE the same object (singleton)"
+
+
+# ============================================================================
+# Subclassing Equality Tests
+# ============================================================================
+
+
+def test_device_subclass_equality(init_cuda):
+    """Test Device subclass equality behavior.
+
+    Device uses a singleton pattern where Device(0) always returns the same
+    cached instance. This means subclassing Device doesn't create new instances;
+    MyDevice(0) returns the original Device(0) instance from the cache.
+    """
+
+    class MyDevice(Device):
+        pass
+
+    device = Device(0)
+    device.set_current()
+    my_device = MyDevice(0)
+
+    # Due to singleton pattern, both return the exact same instance
+    assert device is my_device, "Device singleton returns same instance for same device_id"
+    assert type(device) is Device, "Singleton returns original Device type, not subclass"
+    assert type(my_device) is Device, "Even MyDevice(0) returns Device instance due to singleton"
+
+    # Since they're the same object, they're equal
+    assert device == my_device
+
+
+def test_stream_subclass_equality(init_cuda):
+    """Test Stream subclass equality behavior.
+
+    Stream uses isinstance() for equality checking, which means a Stream instance
+    and a MyStream subclass instance wrapping the same handle will compare equal.
+    """
+
+    class MyStream(Stream):
+        pass
+
+    device = Device(0)
+    device.set_current()
+
+    # Create base Stream instance
+    stream = Stream._init(options=StreamOptions(), device_id=device.device_id)
+
+    # Create another Stream wrapping same handle
+    stream2 = Stream.from_handle(int(stream.handle))
+    assert stream == stream2, "Streams wrapping same handle are equal"
+
+    # Create subclass instance with different handle
+    my_stream = MyStream._init(options=StreamOptions(), device_id=device.device_id)
+
+    # Different handles -> not equal
+    assert stream != my_stream, "Streams with different handles are not equal"
+    assert stream.handle != my_stream.handle
+
+    # sanity check: base and subclass compare equal (and hash equal)
+    stream_from_handle = MyStream.from_handle(int(my_stream.handle))
+    assert my_stream == stream_from_handle, "MyStream and Stream wrapping same handle compare equal"
+    assert hash(my_stream) == hash(stream_from_handle)
+
+
+def test_event_subclass_equality(init_cuda):
+    """Test Event subclass equality behavior.
+
+    Event uses isinstance() for equality checking, similar to Stream.
+    """
+
+    class MyEvent(Event):
+        pass
+
+    device = Device(0)
+    device.set_current()
+
+    # Create two different events
+    event = Event._init(device.device_id, device.context, options=EventOptions())
+    my_event = MyEvent._init(device.device_id, device.context, options=EventOptions())
+
+    # Different events should not be equal (different handles)
+    assert event != my_event, "Different Event instances are not equal"
+
+    # Same subclass type with different handles
+    my_event2 = MyEvent._init(device.device_id, device.context, options=EventOptions())
+    assert my_event != my_event2, "Different MyEvent instances are not equal"
+
+
+def test_context_subclass_equality(init_cuda):
+    """Test Context subclass equality behavior."""
+
+    class MyContext(Context):
+        pass
+
+    device = Device(0)
+    device.set_current()
+    stream = device.create_stream()
+    context = stream.context
+
+    # MyContext._from_ctx() returns a Context instance, not MyContext
+    my_context = MyContext._from_ctx(context._handle, device.device_id)
+    assert type(my_context) is Context, "_from_ctx returns Context, not subclass"
+    assert type(my_context) is not MyContext
+
+    # Since both are Context instances with same handle, they're equal
+    assert context == my_context, "Context instances with same handle are equal"
+
+    # Create another context from different stream
+    stream2 = device.create_stream()
+    context2 = stream2.context
+
+    # Same device, same primary context, should be equal
+    assert context == context2, "Contexts from same device are equal"
+
+
+def test_subclass_type_safety(init_cuda):
+    """Test that equality checks with incompatible types return False or NotImplemented."""
+    device = Device(0)
+    device.set_current()
+
+    stream = device.create_stream()
+    event = stream.record()
+    context = stream.context
+
+    # None of these should be equal to each other
+    assert device != stream
+    assert device != event
+    assert device != context
+    assert stream != event
+    assert stream != context
+    assert event != context
+
+    # None should be equal to arbitrary types
+    assert device != "device"
+    assert stream != 123
+    assert event != []
+    assert context != {"key": "value"}

cuda_core/tests/test_context.py

@@ -3,8 +3,88 @@
 
 import cuda.core.experimental
 import pytest
+from cuda.core.experimental import Device
 
 
 def test_context_init_disabled():
     with pytest.raises(RuntimeError, match=r"^Context objects cannot be instantiated directly\."):
         cuda.core.experimental._context.Context()  # Ensure back door is locked.
+
+
+# ============================================================================
+# Context Equality Tests
+# ============================================================================
+
+
+def test_context_equality_same_context(init_cuda):
+    """Contexts from same device should be equal."""
+    device = Device()
+
+    s1 = device.create_stream()
+    s2 = device.create_stream()
+
+    ctx1 = s1.context
+    ctx2 = s2.context
+
+    # Same device, should have same context
+    assert ctx1 == ctx2, "Streams on same device should share context"
+
+
+def test_context_equality_reflexive(init_cuda):
+    """Context should equal itself (reflexive property)."""
+    device = Device()
+    stream = device.create_stream()
+    context = stream.context
+
+    assert context == context, "Context should equal itself"
+
+
+def test_context_type_safety(init_cuda):
+    """Comparing Context with wrong type should return False."""
+    device = Device()
+    context = device.create_stream().context
+
+    assert (context == "not a context") is False
+    assert (context == 123) is False
+    assert (context is None) is False
+
+
+# ============================================================================
+# Context Hash Tests
+# ============================================================================
+
+
+def test_context_hash_consistency(init_cuda):
+    """Hash of same Context object should be consistent."""
+    device = Device()
+    stream = device.create_stream()
+    context = stream.context
+
+    hash1 = hash(context)
+    hash2 = hash(context)
+    assert hash1 == hash2, "Hash should be consistent for same object"
+
+
+def test_context_hash_equality(init_cuda):
+    """Contexts from same device should hash equal."""
+    device = Device()
+
+    s1 = device.create_stream()
+    s2 = device.create_stream()
+
+    ctx1 = s1.context
+    ctx2 = s2.context
+
+    # Same device, should have same context
+    assert ctx1 == ctx2, "Streams on same device should share context"
+    assert hash(ctx1) == hash(ctx2), "Same context should hash equal"
+
+
+def test_context_dict_key(init_cuda):
+    """Contexts should be usable as dictionary keys."""
+    device = Device()
+    stream = device.create_stream()
+    context = stream.context
+
+    ctx_cache = {context: "context_data"}
+    assert ctx_cache[context] == "context_data"