Add GraphNode identity cache for stable Python object round-trips

Nodes retrieved via GraphDef.nodes(), edges(), or pred/succ traversal
now return the same Python object that was originally created, enabling
identity checks with `is`. A C++ HandleRegistry deduplicates
CUgraphNode handles, and a Cython WeakValueDictionary caches the
Python wrapper objects.

Made-with: Cursor

Author: Andy-Jost

cuda_core/cuda/core/_cpp/resource_handles.cpp

@@ -969,9 +969,16 @@ static const GraphNodeBox* get_box(const GraphNodeHandle& h) {
     );
 }
 
+static HandleRegistry<CUgraphNode, GraphNodeHandle> graph_node_registry;
+
 GraphNodeHandle create_graph_node_handle(CUgraphNode node, const GraphHandle& h_graph) {
+    if (auto h = graph_node_registry.lookup(node)) {
+        return h;
+    }
     auto box = std::make_shared<const GraphNodeBox>(GraphNodeBox{node, h_graph});
-    return GraphNodeHandle(box, &box->resource);
+    GraphNodeHandle h(box, &box->resource);
+    graph_node_registry.register_handle(node, h);
+    return h;
 }
 
 GraphHandle graph_node_get_graph(const GraphNodeHandle& h) noexcept {

cuda_core/cuda/core/_graph/_graph_def/_graph_node.pyx

@@ -56,10 +56,19 @@ from cuda.core._graph._utils cimport (
     _attach_user_object,
 )
 
+import weakref
+
 from cuda.core import Device
 from cuda.core._graph._graph_def._adjacency_set_proxy import AdjacencySetProxy
 from cuda.core._utils.cuda_utils import driver, handle_return
 
+_node_cache = weakref.WeakValueDictionary()
+
+
+cdef inline GraphNode _cached(GraphNode n):
+    _node_cache[<uintptr_t>n._h_node.get()] = n
+    return n
+
 
 cdef class GraphNode:
     """Base class for all graph nodes.
@@ -510,18 +519,30 @@ cdef inline ConditionalNode _make_conditional_node(
     n._cond_type = cond_type
     n._branches = branches
 
-    return n
+    return _cached(n)
+
 
 cdef inline GraphNode GN_create(GraphHandle h_graph, cydriver.CUgraphNode node):
+    cdef GraphNodeHandle h_node = create_graph_node_handle(node, h_graph)
+
+    # Sentinel: virtual node to represent the graph entry point.
     if node == NULL:
         n = GraphNode.__new__(GraphNode)
-        (<GraphNode>n)._h_node = create_graph_node_handle(node, h_graph)
+        (<GraphNode>n)._h_node = h_node
         return n
 
-    cdef GraphNodeHandle h_node = create_graph_node_handle(node, h_graph)
+    # Return a cached object or create and cache a new one.
+    cached = _node_cache.get(<uintptr_t>h_node.get())
+    if cached is not None:
+        return <GraphNode>cached
+    else:
+        return _cached(GN_create_impl(h_node))
+
+
+cdef inline GraphNode GN_create_impl(GraphNodeHandle h_node):
     cdef cydriver.CUgraphNodeType node_type
     with nogil:
-        HANDLE_RETURN(cydriver.cuGraphNodeGetType(node, &node_type))
+        HANDLE_RETURN(cydriver.cuGraphNodeGetType(as_cu(h_node), &node_type))
 
     if node_type == cydriver.CU_GRAPH_NODE_TYPE_EMPTY:
         return EmptyNode._create_impl(h_node)
@@ -583,10 +604,10 @@ cdef inline KernelNode GN_launch(GraphNode self, LaunchConfig conf, Kernel ker,
     _attach_user_object(as_cu(h_graph), <void*>new KernelHandle(ker._h_kernel),
                         <cydriver.CUhostFn>_destroy_kernel_handle_copy)
 
-    return KernelNode._create_with_params(
+    return _cached(KernelNode._create_with_params(
         create_graph_node_handle(new_node, h_graph),
         conf.grid, conf.block, conf.shmem_size,
-        ker._h_kernel)
+        ker._h_kernel))
 
 
 cdef inline EmptyNode GN_join(GraphNode self, tuple nodes):
@@ -612,7 +633,7 @@ cdef inline EmptyNode GN_join(GraphNode self, tuple nodes):
         HANDLE_RETURN(cydriver.cuGraphAddEmptyNode(
             &new_node, as_cu(h_graph), deps_ptr, num_deps))
 
-    return EmptyNode._create_impl(create_graph_node_handle(new_node, h_graph))
+    return _cached(EmptyNode._create_impl(create_graph_node_handle(new_node, h_graph)))
 
 
 cdef inline AllocNode GN_alloc(GraphNode self, size_t size, object options):
@@ -688,9 +709,9 @@ cdef inline AllocNode GN_alloc(GraphNode self, size_t size, object options):
         HANDLE_RETURN(cydriver.cuGraphAddMemAllocNode(
             &new_node, as_cu(h_graph), deps, num_deps, &alloc_params))
 
-    return AllocNode._create_with_params(
+    return _cached(AllocNode._create_with_params(
         create_graph_node_handle(new_node, h_graph), alloc_params.dptr, size,
-        device_id, memory_type, tuple(peer_ids))
+        device_id, memory_type, tuple(peer_ids)))
 
 
 cdef inline FreeNode GN_free(GraphNode self, cydriver.CUdeviceptr c_dptr):
@@ -708,7 +729,7 @@ cdef inline FreeNode GN_free(GraphNode self, cydriver.CUdeviceptr c_dptr):
         HANDLE_RETURN(cydriver.cuGraphAddMemFreeNode(
             &new_node, as_cu(h_graph), deps, num_deps, c_dptr))
 
-    return FreeNode._create_with_params(create_graph_node_handle(new_node, h_graph), c_dptr)
+    return _cached(FreeNode._create_with_params(create_graph_node_handle(new_node, h_graph), c_dptr))
 
 
 cdef inline MemsetNode GN_memset(
@@ -743,9 +764,9 @@ cdef inline MemsetNode GN_memset(
             &new_node, as_cu(h_graph), deps, num_deps,
             &memset_params, ctx))
 
-    return MemsetNode._create_with_params(
+    return _cached(MemsetNode._create_with_params(
         create_graph_node_handle(new_node, h_graph), c_dst,
-        val, elem_size, width, height, pitch)
+        val, elem_size, width, height, pitch))
 
 
 cdef inline MemcpyNode GN_memcpy(
@@ -804,9 +825,9 @@ cdef inline MemcpyNode GN_memcpy(
         HANDLE_RETURN(cydriver.cuGraphAddMemcpyNode(
             &new_node, as_cu(h_graph), deps, num_deps, &params, ctx))
 
-    return MemcpyNode._create_with_params(
+    return _cached(MemcpyNode._create_with_params(
         create_graph_node_handle(new_node, h_graph), c_dst, c_src, size,
-        c_dst_type, c_src_type)
+        c_dst_type, c_src_type))
 
 
 cdef inline ChildGraphNode GN_embed(GraphNode self, GraphDef child_def):
@@ -831,8 +852,8 @@ cdef inline ChildGraphNode GN_embed(GraphNode self, GraphDef child_def):
 
     cdef GraphHandle h_embedded = create_graph_handle_ref(embedded_graph, h_graph)
 
-    return ChildGraphNode._create_with_params(
-        create_graph_node_handle(new_node, h_graph), h_embedded)
+    return _cached(ChildGraphNode._create_with_params(
+        create_graph_node_handle(new_node, h_graph), h_embedded))
 
 
 cdef inline EventRecordNode GN_record_event(GraphNode self, Event ev):
@@ -853,8 +874,8 @@ cdef inline EventRecordNode GN_record_event(GraphNode self, Event ev):
     _attach_user_object(as_cu(h_graph), <void*>new EventHandle(ev._h_event),
                         <cydriver.CUhostFn>_destroy_event_handle_copy)
 
-    return EventRecordNode._create_with_params(
-        create_graph_node_handle(new_node, h_graph), ev._h_event)
+    return _cached(EventRecordNode._create_with_params(
+        create_graph_node_handle(new_node, h_graph), ev._h_event))
 
 
 cdef inline EventWaitNode GN_wait_event(GraphNode self, Event ev):
@@ -875,8 +896,8 @@ cdef inline EventWaitNode GN_wait_event(GraphNode self, Event ev):
     _attach_user_object(as_cu(h_graph), <void*>new EventHandle(ev._h_event),
                         <cydriver.CUhostFn>_destroy_event_handle_copy)
 
-    return EventWaitNode._create_with_params(
-        create_graph_node_handle(new_node, h_graph), ev._h_event)
+    return _cached(EventWaitNode._create_with_params(
+        create_graph_node_handle(new_node, h_graph), ev._h_event))
 
 
 cdef inline HostCallbackNode GN_callback(GraphNode self, object fn, object user_data):
@@ -902,6 +923,6 @@ cdef inline HostCallbackNode GN_callback(GraphNode self, object fn, object user_
             &new_node, as_cu(h_graph), deps, num_deps, &node_params))
 
     cdef object callable_obj = fn if not isinstance(fn, ct._CFuncPtr) else None
-    return HostCallbackNode._create_with_params(
+    return _cached(HostCallbackNode._create_with_params(
         create_graph_node_handle(new_node, h_graph), callable_obj,
-        node_params.fn, node_params.userData)
+        node_params.fn, node_params.userData))

cuda_core/tests/graph/test_graphdef.py

@@ -661,6 +661,7 @@ def test_node_type_preserved_by_nodes(node_spec):
     matched = [n for n in all_nodes if n == node]
     assert len(matched) == 1
     assert isinstance(matched[0], spec.roundtrip_class)
+    assert matched[0] is node
 
 
 def test_node_type_preserved_by_pred_succ(node_spec):
@@ -670,6 +671,7 @@ def test_node_type_preserved_by_pred_succ(node_spec):
         matched = [s for s in predecessor.succ if s == node]
         assert len(matched) == 1
         assert isinstance(matched[0], spec.roundtrip_class)
+        assert matched[0] is node
 
 
 def test_node_attrs(node_spec):

cuda_core/tests/graph/test_graphdef_mutation.py

@@ -372,7 +372,7 @@ def test_convert_linear_to_fan_in(init_cuda):
     for node in g.nodes():
         if isinstance(node, MemsetNode):
             node.pred.clear()
-        elif isinstance(node, KernelNode) and node != reduce_node:
+        elif isinstance(node, KernelNode) and node is not reduce_node:
             node.succ.add(reduce_node)
 
     assert len(g.edges()) == 8