cuda.core: preserve graph identity when CUDA reuses handles (#3294)

* cuda.core: preserve graph identity across native handle reuse

Signed-off-by: Ralf W. Grosse-Kunstleve <rgrossekunst@nvidia.com>

* cuda.core: scope the child-graph registry to the parent GraphBox

The process-wide graph registry mapped a raw CUgraph to its canonical
GraphHandle. Every handle of a hierarchy aliases the hierarchy's single
control block, so a lockable registry entry proved only that the hierarchy
was alive, not that the child graph was. CUDA destroys a child graph inside
cuGraphDestroyNode, cuGraphNodeSetParams, or capture invalidation, and the
box is retired only after the driver call returns. In that window the map
could hand one hierarchy's recycled address to another.

The only reader, create_child_graph_handle, always knows the parent, and
roots are never looked up. Each GraphBox now holds the registry of its
direct child graphs, next to its node registry. A recycled address can only
collide inside one hierarchy, where mutations are externally synchronized
and every retire precedes the matching registration. The conditional
unregister_handle overload, the bulk register_handles, the static registry,
and the hierarchy deleter's unregister loop are gone. The same-parent
registry test modeled two live generations at one address under one parent,
which the synchronization contract excludes, and is removed.

Written by Codex (gpt-6.1-sol); reviewed by Claude.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* cuda.core: key child graphs by address and owner node

A child graph is identified within its parent by its CUgraph and the node
that owns it. Keying the per-parent registry by that pair makes the owner
check part of the lookup. A box whose graph CUDA destroyed before cuda.core
retired it keeps a key distinct from a newer child that reuses its address
under another owner, so retiring the stale box cannot erase the newer entry.

The same-parent registry test is restored; it passes under the pair key.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

---------

Signed-off-by: Ralf W. Grosse-Kunstleve <rgrossekunst@nvidia.com>
Co-authored-by: Andy Jost <ajost@nvidia.com>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

Author: 3 people

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

@@ -171,18 +171,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