cuda.core: preserve graph identity when CUDA reuses handles

cuda_core/cuda/core/_cpp/rt/GRAPH_ATTACHMENTS.md

@@ -167,18 +167,19 @@ metadata for the root graph and subgraphs. Each `GraphBox` stores:
 - the `CUgraph`
 - its parent box and owning node, when it is a child graph
 - a `GraphAttachmentMap` from nodes to node attachments (`NodeAttachment*`)
-- a per-graph weak registry to canonicalize node handles
-
-When CUDA destroys a child graph, the associated `GraphBox` is cleared and
-moved to `GraphHierarchy::graveyard`. That invalidates existing handles and
-preserves the storage until the hierarchy's last shared owner is released.
-Removing the registry entry also prevents corruption if a future CUDA graph
-reuses the raw handle value.
-
-The process-wide graph registry maps a live `CUgraph` to its canonical
-`GraphHandle`. Node handles are scoped to their `GraphBox`, where they can all
-be invalidated when CUDA destroys that graph. They use separate
-`GraphNodeBox::resource` fields and are nulled before the graph box is retired.
+- weak registries to canonicalize direct child graph and node handles
+
+A retired `GraphBox` is cleared and moved to `GraphHierarchy::graveyard`.
+That invalidates existing handles and preserves the storage until the
+hierarchy's last shared owner is released.
+
+Each `GraphBox` holds two registries: its direct child graphs keyed by
+`CUgraph` and owner node, and its nodes keyed by `CUgraphNode`. Both are
+scoped to the box, so a raw handle recycled by CUDA in another hierarchy
+cannot resolve to this hierarchy's metadata; see
+[registry identity](REGISTRY_DESIGN.md#reused-graph-handles). Node handles use
+separate `GraphNodeBox::resource` fields and are nulled before the graph box
+is retired.
 
 ## Invariants
 

cuda_core/cuda/core/_cpp/rt/REGISTRY_DESIGN.md

@@ -30,7 +30,26 @@ Without this level, a round-tripped handle would produce a new Box
 with default metadata, losing information that was set at creation.
 
 Instances: `context_registry`, `stream_registry`, `event_registry`,
-`kernel_registry`, `graph_registry`, and one node registry per `GraphBox`.
+`kernel_registry`, and, per `GraphBox`, one registry for its direct child
+graphs and one for its nodes.
+
+### Reused graph handles
+
+A `GraphHandle` weak reference locks whenever any handle of its hierarchy is
+alive, so a lockable entry does not prove that a child graph is alive. CUDA
+destroys child graphs inside node operations and capture invalidation, before
+cuda.core retires their boxes. A process-wide map keyed by raw `CUgraph` could
+therefore resolve a recycled address to another hierarchy's metadata.
+
+Scoping the child registry to its parent `GraphBox` removes that case by
+construction: `create_child_graph_handle` always knows the parent, and roots
+are never looked up. Within the parent, an entry is keyed by the child's
+`CUgraph` and its owner node, which together identify a child graph. If CUDA
+destroys a child before cuda.core retires its box, the owner node is still
+alive, so a newer child at the recycled address has a different owner and its
+own entry. Retiring the stale box erases only the stale entry. Within one
+owner, external synchronization and the retire-before-register order in
+`graph_commit_child_graph_update` keep one live entry per key.
 
 ## Level 2: Resource Handle -> Python Object (Cython)
 

cuda_core/cuda/core/_cpp/rt/graph.cpp

@@ -11,6 +11,7 @@
 #include <array>
 #include <cstddef>
 #include <cstdlib>
+#include <functional>
 #include <list>
 #include <map>
 #include <memory>
@@ -69,6 +70,29 @@ using GraphAttachmentMap = std::map<CUgraphNode, NodeAttachment *>;
 
 struct GraphHierarchy;
 
+// A child graph is identified by its address and its owner node in the parent.
+// A box that CUDA destroyed before cuda.core retired it keeps a key distinct
+// from a newer child that reuses its address under another owner node.
+struct ChildGraphKey {
+    CUgraph graph;
+    CUgraphNode owner_node;
+
+    bool
+    operator==(const ChildGraphKey &other) const noexcept
+    {
+        return graph == other.graph && owner_node == other.owner_node;
+    }
+};
+
+struct ChildGraphKeyHash {
+    std::size_t
+    operator()(const ChildGraphKey &key) const noexcept
+    {
+        std::size_t h = std::hash<const void *>{}(key.graph);
+        return h ^ (std::hash<const void *>{}(key.owner_node) + 0x9e3779b9 + (h << 6) + (h >> 2));
+    }
+};
+
 // Standard-layout alias target for GraphHandle.
 struct GraphBoxBase {
     CUgraph resource = nullptr;
@@ -82,6 +106,8 @@ struct GraphBox : GraphBoxBase {
     CUgraphNode owner_node = nullptr;     // Node in parent that owns this graph.
     GraphAttachmentMap attachments;       // Non-owning attachment index.
     HandleRegistry<CUgraphNode, GraphNodeHandle> node_handles;
+    // See REGISTRY_DESIGN.md (Level 1: Driver Handle -> Resource Handle)
+    HandleRegistry<ChildGraphKey, GraphHandle, ChildGraphKeyHash> children;  // Direct child graphs.
 
     GraphBox(
         CUgraph resource_, GraphHierarchy *hierarchy_, GraphBox *parent_ = nullptr, CUgraphNode owner_node_ = nullptr
@@ -104,10 +130,6 @@ struct GraphHierarchy {
     }
 };
 
-// See REGISTRY_DESIGN.md (Level 1: Driver Handle -> Resource Handle)
-using GraphRegistry = HandleRegistry<CUgraph, GraphHandle>;
-static GraphRegistry graph_registry;
-
 // Immutable resource owners for one version of a graph node's parameters.
 // Inheriting DeferredCleanupItem lets CUDA's user-object destructor enqueue
 // the payload without destroying owners on the callback thread.
@@ -144,6 +166,17 @@ get_box(const GraphHandle &h) noexcept
     return const_cast<GraphBox *>(static_cast<const GraphBox *>(value));
 }
 
+// Register child handles with their immediate parents. Each handle aliases a
+// GraphBox whose parent is set; a hierarchy's root is never registered.
+void
+register_child_handles(const std::vector<GraphHandle> &handles)
+{
+    for (const GraphHandle &h : handles) {
+        GraphBox *box = get_box(h);
+        box->parent->children.register_handle({box->resource, box->owner_node}, h);
+    }
+}
+
 // Rekey a staged attachment map from source nodes to their cloned nodes.
 // The caller must release the GIL before calling this function.
 CUresult
@@ -302,11 +335,6 @@ create_graph_handle(CUgraph graph)
     auto hierarchy = std::shared_ptr<GraphHierarchy>(
         new GraphHierarchy{},
         [](GraphHierarchy *hierarchy) {
-            for (const GraphBox& box : hierarchy->graphs) {
-                if (box.resource) {
-                    graph_registry.unregister_handle(box.resource);
-                }
-            }
             GraphBox *root = hierarchy->root();
             if (root && root->resource) {
                 GILReleaseGuard gil;
@@ -320,7 +348,6 @@ create_graph_handle(CUgraph graph)
     GraphBox &root = hierarchy->graphs.emplace_back(graph, hierarchy.get());
 
     GraphHandle h_graph(hierarchy, &root.resource);
-    graph_registry.register_handle(graph, h_graph);
     return h_graph;
 }
 
@@ -330,16 +357,16 @@ create_child_graph_handle(CUgraph child_graph, const GraphHandle &h_parent, CUgr
     if (!child_graph || !h_parent || !owner_node) {
         return {};
     }
-    if (GraphHandle h_graph = graph_registry.lookup(child_graph)) {
+    GraphBox *parent = get_box(h_parent);
+    if (GraphHandle h_graph = parent->children.lookup({child_graph, owner_node})) {
         return h_graph;
     }
 
-    GraphBox *parent = get_box(h_parent);
     GraphHierarchy *hierarchy = parent->hierarchy;
     GraphBox &child = hierarchy->graphs.emplace_back(child_graph, hierarchy, parent, owner_node);
 
     GraphHandle h_child(h_parent, &child.resource);
-    graph_registry.register_handle(child_graph, h_child);
+    parent->children.register_handle({child_graph, owner_node}, h_child);
     return h_child;
 }
 
@@ -395,7 +422,7 @@ publish_child_graph_update(PreparedChildGraphUpdateState &state, GraphHandle *ou
     GraphBox *parent = get_box(state.h_parent);
     parent->hierarchy->graphs.splice(parent->hierarchy->graphs.end(), state.replacement);
     *out_child = state.handles.front();
-    graph_registry.register_handles(state.handles);
+    register_child_handles(state.handles);
 }
 
 CUresult
@@ -629,7 +656,7 @@ graph_clone_attachments(const GraphHandle &h_clone, const GraphHandle &h_source)
     }
 
     clone->hierarchy->graphs.splice(clone->hierarchy->graphs.end(), subgraphs);
-    graph_registry.register_handles(handles);
+    register_child_handles(handles);
     return CUDA_SUCCESS;
 }
 
@@ -659,7 +686,9 @@ retire_graph(GraphHierarchy &hierarchy, std::list<GraphBox>::iterator graph) noe
         }
     }
     if (graph->resource) {
-        graph_registry.unregister_handle(graph->resource);
+        if (graph->parent) {
+            graph->parent->children.unregister_handle({graph->resource, graph->owner_node});
+        }
         graph->resource = nullptr;
     }
     graph->attachments.clear();

cuda_core/cuda/core/_cpp/rt/internal.hpp

@@ -16,7 +16,6 @@
 #include <type_traits>
 #include <unordered_map>
 #include <utility>
-#include <vector>
 
 namespace cuda_core::rt::detail {
 
@@ -193,17 +192,6 @@ class HandleRegistry {
         map_.erase(key);
     }
 
-    void
-    register_handles(const std::vector<Handle> &handles)
-    {
-        std::lock_guard<std::mutex> lock(mutex_);
-        for (const Handle &h : handles) {
-            if (h) {
-                map_[*h] = h;
-            }
-        }
-    }
-
     Handle
     lookup(const Key &key)
     {

cuda_core/tests/cython/build_tests.py

@@ -8,6 +8,8 @@
 so `cimport cuda.bindings.*` fails to locate the .pxd files. We resolve
 the namespace package's source root from `cuda.bindings.__file__` and pass
 it via `include_path=` so cythonize finds the .pxd tree on every platform.
+Tests that cimport cuda.core also need this checkout's Core source root on
+that search path; the editable install's Python finder does not supply it.
 
 When CUDA_PYTHON_CYTHON_CACHE_DIR is set, cythonize uses the same cache
 namespacing and include-path aliasing as the package build (``_build_shared.py``).
@@ -54,6 +56,7 @@ def _bindings_source_root() -> Path:
 
 
 def _cythonize_tests(pyx_files):
+    core_source_root = Path(__file__).resolve().parents[2]
     cache_path = build_shared._cython_cache_path(
         "cuda-core-cython-tests",
         compiler_directives=_COMPILER_DIRECTIVES,
@@ -68,7 +71,7 @@ def _cythonize_tests(pyx_files):
     if cache_path is None:
         return cythonize(
             pyx_files,
-            include_path=[str(_bindings_source_root())],
+            include_path=[str(core_source_root), str(_bindings_source_root())],
             **cythonize_kwargs,
         )
 
@@ -78,11 +81,12 @@ def _cythonize_tests(pyx_files):
     stdlib_target = Path(Cython.__file__).parent / "Includes"
     with (
         build_shared._stable_cython_alias(stdlib_target, Path(".cython-stdlib-tests")) as rel_stdlib,
+        build_shared._stable_cython_alias(core_source_root, Path(".cython-core-tests")) as rel_core,
         build_shared._stable_cython_alias(_bindings_source_root(), Path(".cython-bindings-tests")) as rel_bindings,
     ):
         return cythonize(
             pyx_files,
-            include_path=[".", rel_bindings, rel_stdlib],
+            include_path=[".", rel_core, rel_bindings, rel_stdlib],
             **cythonize_kwargs,
         )
 
@@ -101,6 +105,9 @@ def main() -> None:
     # are stable regardless of the invoking cwd.
     os.chdir(script_dir)
     ext_modules = _cythonize_tests(pyx_files)
+    # Core's runtime consumer headers require C++17 with either toolchain.
+    for extension in ext_modules:
+        extension.extra_compile_args.append("/std:c++17" if os.name == "nt" else "-std=c++17")
 
     sys.argv = [sys.argv[0], "build_ext", "--inplace"]
     setup(name="cuda_core_cython_tests", ext_modules=ext_modules)

cuda_core/tests/cython/test_cython.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2021-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2021-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # SPDX-License-Identifier: Apache-2.0
 
@@ -21,6 +21,7 @@ def wrapped(*args, **kwargs):
 
 cython_test_modules = [
     "test_get_cuda_native_handle",
+    "test_graph_registry",
 ]
 
 

cuda_core/tests/cython/test_get_cuda_native_handle.pyx

@@ -1,9 +1,8 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # SPDX-License-Identifier: Apache-2.0
 
 # distutils: language = c++
-# distutils: extra_compile_args = -std=c++17
 
 from libc.stdint cimport intptr_t