diff --git a/cuda_core/cuda/core/_cpp/rt/GRAPH_ATTACHMENTS.md b/cuda_core/cuda/core/_cpp/rt/GRAPH_ATTACHMENTS.md index fdaf77785b2..02897f727cd 100644 --- a/cuda_core/cuda/core/_cpp/rt/GRAPH_ATTACHMENTS.md +++ b/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 diff --git a/cuda_core/cuda/core/_cpp/rt/REGISTRY_DESIGN.md b/cuda_core/cuda/core/_cpp/rt/REGISTRY_DESIGN.md index f4c760163e6..003887f296c 100644 --- a/cuda_core/cuda/core/_cpp/rt/REGISTRY_DESIGN.md +++ b/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) diff --git a/cuda_core/cuda/core/_cpp/rt/graph.cpp b/cuda_core/cuda/core/_cpp/rt/graph.cpp index 13720e6a9b1..09261623154 100644 --- a/cuda_core/cuda/core/_cpp/rt/graph.cpp +++ b/cuda_core/cuda/core/_cpp/rt/graph.cpp @@ -11,6 +11,7 @@ #include #include #include +#include #include #include #include @@ -69,6 +70,29 @@ using GraphAttachmentMap = std::map; 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{}(key.graph); + return h ^ (std::hash{}(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 node_handles; + // See REGISTRY_DESIGN.md (Level 1: Driver Handle -> Resource Handle) + HandleRegistry 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; -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(static_cast(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 &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( 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::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(); diff --git a/cuda_core/cuda/core/_cpp/rt/internal.hpp b/cuda_core/cuda/core/_cpp/rt/internal.hpp index 270ee486aa2..382c4901f11 100644 --- a/cuda_core/cuda/core/_cpp/rt/internal.hpp +++ b/cuda_core/cuda/core/_cpp/rt/internal.hpp @@ -16,7 +16,6 @@ #include #include #include -#include namespace cuda_core::rt::detail { @@ -193,17 +192,6 @@ class HandleRegistry { map_.erase(key); } - void - register_handles(const std::vector &handles) - { - std::lock_guard lock(mutex_); - for (const Handle &h : handles) { - if (h) { - map_[*h] = h; - } - } - } - Handle lookup(const Key &key) { diff --git a/cuda_core/tests/cython/build_tests.py b/cuda_core/tests/cython/build_tests.py index 39539e3100c..0ae2d3531e6 100644 --- a/cuda_core/tests/cython/build_tests.py +++ b/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) diff --git a/cuda_core/tests/cython/test_cython.py b/cuda_core/tests/cython/test_cython.py index 87b97e8a2a4..ff40e9d841c 100644 --- a/cuda_core/tests/cython/test_cython.py +++ b/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", ] diff --git a/cuda_core/tests/cython/test_get_cuda_native_handle.pyx b/cuda_core/tests/cython/test_get_cuda_native_handle.pyx index 1480a3b96bd..fa56c586a61 100644 --- a/cuda_core/tests/cython/test_get_cuda_native_handle.pyx +++ b/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 diff --git a/cuda_core/tests/cython/test_graph_registry.pyx b/cuda_core/tests/cython/test_graph_registry.pyx new file mode 100644 index 00000000000..09ddca2f706 --- /dev/null +++ b/cuda_core/tests/cython/test_graph_registry.pyx @@ -0,0 +1,110 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# distutils: language = c++ +# distutils: include_dirs = ../../cuda/core + +"""Deterministic raw-handle reuse tests against the per-parent child registries. + +Parents are real empty CUDA graphs. Stack-local addresses stand in for reused +child and owner handles only inside the metadata API; they are never passed to +CUDA. Distinct addresses per invocation also keep parallel workers independent. +The parents retain ownership of their real graphs and destroy them normally. +""" + +from libc.stdint cimport uintptr_t +from cuda.bindings.cydriver cimport CUgraph, CUgraphNode +from cuda.core._rt cimport ( + GraphHandle, + as_cu, + create_child_graph_handle, + invalidate_child_graph_state, +) +from cuda.core.graph._graph_definition cimport GraphDefinition + +import pytest + + +@pytest.mark.agent_authored(model="gpt-6") +def test_child_graph_registry_rejects_foreign_parent(init_cuda): + """Reusing both raw addresses must not alias an independent parent tree.""" + cdef uintptr_t tokens[2] + cdef CUgraph raw_child = &tokens[0] + cdef CUgraphNode owner = &tokens[1] + cdef GraphDefinition old_parent = GraphDefinition() + cdef GraphDefinition new_parent = GraphDefinition() + cdef GraphHandle old_child = create_child_graph_handle( + raw_child, old_parent._h_graph, owner + ) + cdef GraphHandle same_child = create_child_graph_handle( + raw_child, old_parent._h_graph, owner + ) + assert same_child.get() == old_child.get() + + cdef GraphHandle new_child = create_child_graph_handle( + raw_child, new_parent._h_graph, owner + ) + assert new_child.get() != old_child.get() + assert as_cu(new_child) == raw_child + + invalidate_child_graph_state(old_parent._h_graph, owner) + assert as_cu(old_child) == NULL + assert as_cu(same_child) == NULL + assert as_cu(new_child) == raw_child + + same_child = create_child_graph_handle(raw_child, new_parent._h_graph, owner) + assert same_child.get() == new_child.get() + + +@pytest.mark.agent_authored(model="gpt-6") +def test_child_graph_registry_rejects_foreign_owner_in_same_parent(init_cuda): + """Retiring an old owner preserves the replacement within one hierarchy.""" + cdef uintptr_t tokens[3] + cdef CUgraph raw_child = &tokens[0] + cdef CUgraphNode old_owner = &tokens[1] + cdef CUgraphNode new_owner = &tokens[2] + cdef GraphDefinition parent = GraphDefinition() + cdef GraphHandle old_child = create_child_graph_handle( + raw_child, parent._h_graph, old_owner + ) + cdef GraphHandle new_child = create_child_graph_handle( + raw_child, parent._h_graph, new_owner + ) + assert new_child.get() != old_child.get() + + invalidate_child_graph_state(parent._h_graph, old_owner) + assert as_cu(old_child) == NULL + assert as_cu(new_child) == raw_child + + cdef GraphHandle same_child = create_child_graph_handle( + raw_child, parent._h_graph, new_owner + ) + assert same_child.get() == new_child.get() + + +@pytest.mark.agent_authored(model="gpt-6") +def test_graph_hierarchy_deletion_preserves_replacement_mapping(init_cuda): + """An old tree's destructor must not erase another tree's newer entry.""" + cdef uintptr_t tokens[3] + cdef CUgraph raw_child = &tokens[0] + cdef CUgraphNode old_owner = &tokens[1] + cdef CUgraphNode new_owner = &tokens[2] + cdef GraphDefinition old_parent = GraphDefinition() + cdef GraphDefinition new_parent = GraphDefinition() + cdef GraphHandle old_child = create_child_graph_handle( + raw_child, old_parent._h_graph, old_owner + ) + cdef GraphHandle new_child = create_child_graph_handle( + raw_child, new_parent._h_graph, new_owner + ) + assert new_child.get() != old_child.get() + + # Drop every owner of the old hierarchy so its real CUDA root is destroyed. + old_child.reset() + old_parent = None + assert as_cu(new_child) == raw_child + + cdef GraphHandle same_child = create_child_graph_handle( + raw_child, new_parent._h_graph, new_owner + ) + assert same_child.get() == new_child.get()