diff --git a/cuda_core/cuda/core/_cpp/DESIGN.md b/cuda_core/cuda/core/_cpp/DESIGN.md index 3aafe68e2f..bcf1a7060f 100644 --- a/cuda_core/cuda/core/_cpp/DESIGN.md +++ b/cuda_core/cuda/core/_cpp/DESIGN.md @@ -34,12 +34,16 @@ responsibilities, we achieve: Resource handles provide **referentially transparent** wrappers around CUDA resources: - **No rebinding**: A handle always refers to the same resource. -- **No invalidation**: If a handle exists, its resource is valid. +- **Ownership-backed validity**: A handle keeps independently owned resources + alive. - **Structural dependencies**: If resource A depends on resource B, A's handle embeds B's handle, automatically extending B's lifetime. -This eliminates global lifetime analysis. Correctness is enforced structurally—if you -have a handle, you have a valid resource. +Borrowed graph and node handles are an unavoidable exception. CUDA owns child +graphs and their nodes, and explicitly removing or replacing an owner node +destroys those resources even if cuda-core wrappers remain. cuda-core +invalidates the affected handles while keeping their storage stable; see +[Graph Node Attachments](GRAPH_ATTACHMENTS.md). ## Handle Types @@ -205,6 +209,9 @@ struct StreamBox { The shared pointer's custom deleter captures any additional state needed for destruction. This ensures resources are always destroyed in the correct order. +Graph node parameters require a specialized ownership model built on +`OpaqueHandle`; see [Graph Node Attachments](GRAPH_ATTACHMENTS.md). + ### GIL Management Handle destructors may run from any thread. The implementation includes RAII guards @@ -217,21 +224,6 @@ Handle destructors may run from any thread. The implementation includes RAII gua The handle API functions are safe to call with or without the GIL held. They will release the GIL (if necessary) before calling CUDA driver API functions. -### CUDA User-Object Cleanup - -CUDA user-object destructors may run on an internal driver thread where CUDA -API calls are forbidden. Graph attachment payloads can release handles whose -deleters call CUDA or run Python finalizers, so their CUDA callback must not -delete the payload directly. - -The callback enqueues a preallocated `DeferredCleanupItem` on a -process-lifetime lock-free `DeferredCleanupQueue` and schedules one coalesced -`Py_AddPendingCall`. The pending callback drains all queued payloads from -Python's main thread. If scheduling fails, payloads remain intact for a later -retry; they are intentionally leaked once interpreter finalization begins. A -single pending call is shared by all queued payloads because CPython's -main-thread pending-call queue is bounded. - ### Static Initialization and Deadlock Hazards When writing C++ code that interacts with Python, a subtle deadlock can occur diff --git a/cuda_core/cuda/core/_cpp/GRAPH_ATTACHMENTS.md b/cuda_core/cuda/core/_cpp/GRAPH_ATTACHMENTS.md new file mode 100644 index 0000000000..2b3ef9f44a --- /dev/null +++ b/cuda_core/cuda/core/_cpp/GRAPH_ATTACHMENTS.md @@ -0,0 +1,263 @@ +# Graph Node Attachments + +## Purpose + +CUDA user objects tie externally managed resource lifetimes to CUDA graphs. +cuda-core uses them to keep resources referenced by graph node parameters +alive. It also keeps non-owning metadata that maps each node to the +user-object attachment backing its current parameters. This mapping lets +cuda-core update the correct user-object references during graph changes and +reproduce the node associations in graph clones. + +A `CUgraph` is a mutable graph definition made of nodes. Each node has a +complete parameter set that may refer to kernels, memory, events, callbacks, or +other external resources. CUDA does not normally take ownership of those +resources. + +A definition can be cloned, embedded as a child graph, or instantiated as a +`CUgraphExec`. Each operation copies the definition's current nodes and +parameters. The source and copy then have independent lifetimes: changing the +source does not change an existing clone or executable. An embedded clone is +owned by its parent node and is destroyed when that node is removed or +replaced. + +Definitions can add, delete, or replace nodes and their parameters. Executable +graphs can adopt a new definition through a whole-graph update: success +replaces their inherited state, while failure leaves the old state intact. An +in-flight launch may continue using an older parameter set even after its +source node or executable has changed or been destroyed. Direct +executable-node updates require separate executable-specific ownership and are +not tracked by this definition metadata. + +Several versions of one logical node's parameters may therefore remain live at +the same time. Without version-specific ownership, replacing or deleting a node +could release memory, kernels, callbacks, or other resources that an existing +clone, executable, or launch still uses. + +cuda-core prevents this by giving each parameter version an immutable +`NodeAttachment` owned through a graph-retained CUDA user object. Non-owning +metadata records the current attachment for each graph node so cuda-core can +update the correct user-object references as graphs change. + +This document explains that ownership model and provides a map to the data +structures and functions that implement it. + +## Ownership model + +For each resource-bearing node version, cuda-core: + +1. Collects its resources in an immutable `NodeAttachment`. +2. Creates a CUDA user object that owns the attachment. +3. Retains one user-object reference on the `CUgraph`. + +CUDA copies graph-owned user-object references when it clones or instantiates a +graph. It also keeps the references needed by in-flight launches. Replacing or +deleting a source node therefore releases only the source graph's reference. +Clones, executable graphs, and launches retain the old attachment for as long +as they need it. + +```text +CUgraph definition ─┐ +CUgraph clone ──────┼── retains ──> CUuserObject ── owns ──> NodeAttachment +CUgraphExec ────────┤ │ +in-flight launch ───┘ └── owns resources + +GraphAttachmentMap ── non-owning lookup ──> current NodeAttachment +``` + +The CUDA user-object reference count controls the attachment lifetime. +`GraphAttachmentMap` only lets cuda-core find the attachment currently +associated with a node. + +Each `NodeAttachment` contains two type-erased `OpaqueHandle` owners: + +- kernel: kernel and argument storage +- host callback: callback and copied user data +- memcpy: destination and source +- memset or event: destination or event in the first owner + +`OpaqueHandle` is `shared_ptr`. Existing cuda-core handles reuse +their shared ownership when converted to it. Python objects and copied callback +data use custom deleters. + +## Graph hierarchy state + +One `GraphHierarchy` represents a root graph and every CUDA-owned child or +conditional-body graph below it. Every graph in the hierarchy has one +canonical `GraphBox`. + +A `GraphHandle` points at `GraphBox::resource` while sharing ownership of the +whole `GraphHierarchy`. Holding any root or child handle therefore keeps every +box and the root CUDA graph alive. Only the root graph is destroyed with +`cuGraphDestroy`; CUDA owns and destroys embedded child graphs. + +`GraphBox` stores: + +- the `CUgraph` +- its parent box and owning node, when it is a child graph +- a `GraphAttachmentMap` from `CUgraphNode` to `NodeAttachment*` +- a per-graph weak registry of canonical `GraphNodeHandle` objects + +`GraphHierarchy::graphs` stores live boxes in parent-before-child order. When +CUDA destroys a child graph, cuda-core unregisters its raw handle, nulls the +resource, clears its non-owning metadata, and moves the box to +`GraphHierarchy::graveyard`. Keeping this tombstone at a stable address lets +existing aliasing handles safely observe that the graph is invalid. Removing +the registry entry also lets a future CUDA graph reuse 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. + +## Code map + +The C++ implementation is in +[resource_handles.cpp](resource_handles.cpp). Its public C++ declarations are +in [resource_handles.hpp](resource_handles.hpp), and the Cython declarations +are in [_resource_handles.pxd](../_resource_handles.pxd). + +The main operations are: + +- `create_graph_handle`: create the root box and hierarchy +- `create_child_graph_handle`: register a CUDA-owned child graph and return an + aliasing handle +- `graph_get_attachment`: copy either or both current owners +- `graph_prepare_attachment`: graph-retain a staged replacement before mutation +- `graph_commit_attachment`: publish or anonymously retain the staged attachment +- `graph_clone_attachments`: copy attachment metadata into an embedded clone +- `invalidate_child_graph_state`: invalidate boxes and node handles after CUDA + destroys child graphs + +Private helpers `rekey_attachments` and `copy_attachments` implement recursive +clone mapping. + +The Cython graph code creates CUDA nodes and then calls these operations to keep +the C++ metadata synchronized with the driver. + +## Common operations + +### Adding a node + +Cython keeps the parameter owners alive while it calls the CUDA node-creation +API. Before the call, `graph_prepare_attachment` creates and graph-retains one +user object for the complete owner bundle. It also preallocates the metadata +entry needed by commit. + +After CUDA returns the new node handle, `graph_commit_attachment` publishes the +preallocated entry without allocating. If preparation or node creation fails, +destruction of the `PreparedAttachment` automatically releases the staged graph +reference. CUDA therefore never receives node parameters whose owners could +not be retained. + +### Reading and replacing attachments + +`graph_get_attachment` copies the requested owners. Either output pointer may +be null when the caller needs only one owner. A missing attachment returns +empty handles. + +`graph_prepare_attachment` creates and graph-retains a complete replacement +bundle before the CUDA parameter update. After the update succeeds, +`graph_commit_attachment` publishes the replacement before releasing the +graph's reference to the old user object, so metadata never points at a +released payload. Preparing two empty owners stages removal of the current +attachment. + +A partial parameter update first gets the current owners, replaces the values +covered by the update, prepares the resulting complete attachment, asks CUDA +to apply the new complete parameter set, and then commits the attachment. + +### Embedding a child graph + +`cuGraphAddChildGraphNode` clones the source graph into a node owned by the +parent. CUDA copies all user-object references into that embedded clone. + +After obtaining the embedded `CUgraph`, cuda-core creates its canonical child +box and calls `graph_clone_attachments`. This operation: + +1. Copies the source graph's non-owning attachment map. +2. Uses `cuGraphNodeFindInClone` to replace source node keys with cloned node + keys. +3. Finds nested cloned child graphs from their mapped owner nodes. +4. Recursively creates staged boxes and copies their maps. +5. Publishes the root map with `swap` and the descendant boxes with + `list::splice`. +6. Registers canonical handles for the published descendant graphs. + +All CUDA mapping happens before publication. A mapping error leaves the live +metadata unchanged. + +### Deleting a node + +cuda-core asks CUDA to destroy the node before changing any wrapper or +attachment state. If CUDA rejects deletion, all cuda-core state remains +unchanged. An empty `PreparedAttachment` is created before the call and +committed only after deletion succeeds. + +After a successful deletion, cuda-core: + +1. Clears the node attachment and releases that graph's user-object reference. +2. Finds child graphs owned by the node. +3. Invalidates their graph and node handles. +4. Clears their non-owning metadata and moves their boxes to the graveyard. +5. Invalidates the deleted node handle. + +CUDA already destroyed those child graphs and released their user-object +references, so cuda-core does not release the child attachments again. + +### Captured host callbacks + +`cuLaunchHostFunc` can add a host node during stream capture but does not return +its `CUgraphNode`. cuda-core normally recovers the node from the stream's +capture dependencies and records the callback attachment there. + +If CUDA accepted the callback but its node cannot be identified, cuda-core +commits the prepared attachment with a null node. This retains the owners +anonymously on the graph without publishing node metadata. The owners then +remain alive until graph destruction, preventing a dangling callback pointer. + +### Deferred cleanup + +CUDA invokes a user-object destructor on an internal thread where CUDA API +calls are forbidden. Destroying an attachment there could release handles whose +deleters call CUDA or run Python finalizers. + +`NodeAttachment` therefore inherits from `DeferredCleanupItem`. The CUDA +destructor callback only adds the attachment to the process-lifetime +`DeferredCleanupQueue` and requests a `Py_AddPendingCall`. + +One pending call drains all queued attachments from Python's main thread. The +queue coalesces work because CPython's pending-call queue is bounded. If +scheduling fails, attachments stay queued and a later enqueue or safe cuda-core +entry retries. Graph and executable-graph destruction and explicit close paths +provide additional retry points. During Python finalization, scheduling stops +and unreclaimable attachments are intentionally leaked rather than destroyed +in an unsafe context. + +## Invariants + +1. A published `NodeAttachment` is immutable. +2. CUDA user-object references, not metadata pointers, own attachments. +3. Metadata is removed or replaced before its graph reference is released. +4. Fallible attachment setup and metadata allocation happen before the CUDA + graph mutation they support. +5. Every live cuda-core `CUgraph` has one canonical `GraphBox` and registry + entry. +6. Graph boxes remain in parent-before-child order. +7. Destroyed child boxes remain at stable addresses in the graveyard. +8. A raw graph handle is unregistered before its box becomes a tombstone. +9. CUDA callbacks only enqueue attachments; they never release owners or call + CUDA. +10. Graph mutations and their metadata updates require the same external + synchronization as the underlying CUDA graph. + +## Scope + +- Attachment metadata tracks graph mutations performed through cuda-core. +- Raw driver clones receive the CUDA user-object references needed for safe + execution, but cuda-core cannot reconstruct their node-to-attachment map. +- Executable graphs rely on CUDA's inherited user-object references; they do + not use `GraphAttachmentMap`. +- Direct executable-node updates require separate executable ownership and are + not tracked by definition attachment metadata. +- Stream capture explicitly retains host callbacks. Other captured operations + keep their documented caller-owned lifetime contract. diff --git a/cuda_core/cuda/core/_cpp/REGISTRY_DESIGN.md b/cuda_core/cuda/core/_cpp/REGISTRY_DESIGN.md index 089f98acd9..233129d9c1 100644 --- a/cuda_core/cuda/core/_cpp/REGISTRY_DESIGN.md +++ b/cuda_core/cuda/core/_cpp/REGISTRY_DESIGN.md @@ -20,9 +20,9 @@ expires. ## Level 1: Driver Handle -> Resource Handle (C++) `HandleRegistry` in `resource_handles.cpp` maps a raw CUDA handle -(e.g., `CUevent`, `CUkernel`, `CUgraphNode`) to the `weak_ptr` that -owns it. When a `_ref` constructor receives a raw handle, it -checks the registry first. If found, it returns the existing +(e.g., `CUevent`, `CUkernel`, `CUgraph`) to a `weak_ptr` +for its owning resource handle. When a `_ref` constructor receives a raw +handle, it checks the registry first. If found, it returns the existing `shared_ptr`, preserving the Box and its metadata (e.g., `EventBox` carries timing/IPC flags, `KernelBox` carries the library dependency). @@ -30,7 +30,13 @@ 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_node_registry`. +`kernel_registry`, `graph_registry`. + +Graph node identity is scoped to its owning graph. Each `GraphBox` therefore +stores its own node-handle `HandleRegistry` instead of using a process-wide +`HandleRegistry`. This is to simplify invalidating handles when a child graph +is destroyed. The registry's mutex protects this internal cache when +concurrent reads reconstruct nodes. ## Level 2: Resource Handle -> Python Object (Cython) diff --git a/cuda_core/cuda/core/_cpp/resource_handles.cpp b/cuda_core/cuda/core/_cpp/resource_handles.cpp index 56e20b9572..d16cfa6079 100644 --- a/cuda_core/cuda/core/_cpp/resource_handles.cpp +++ b/cuda_core/cuda/core/_cpp/resource_handles.cpp @@ -11,6 +11,7 @@ #include #include #include +#include #include #include #include @@ -77,6 +78,9 @@ decltype(&cuGraphExecDestroy) p_cuGraphExecDestroy = nullptr; decltype(&cuUserObjectCreate) p_cuUserObjectCreate = nullptr; decltype(&cuUserObjectRelease) p_cuUserObjectRelease = nullptr; decltype(&cuGraphRetainUserObject) p_cuGraphRetainUserObject = nullptr; +decltype(&cuGraphReleaseUserObject) p_cuGraphReleaseUserObject = nullptr; +decltype(&cuGraphNodeFindInClone) p_cuGraphNodeFindInClone = nullptr; +decltype(&cuGraphChildGraphNodeGetGraph) p_cuGraphChildGraphNodeGetGraph = nullptr; // Linker decltype(&cuLinkDestroy) p_cuLinkDestroy = nullptr; @@ -189,29 +193,21 @@ class GILAcquireGuard { // ============================================================================ // CUDA user-object deferred cleanup // -// CUDA may invoke a user-object destructor on an internal thread where CUDA +// CUDA invokes a user-object destructor on an internal thread where CUDA // calls are forbidden. Payload cleanup can release resource handles whose // deleters call CUDA, so the callback only transfers a preallocated intrusive // node to this process-lifetime queue. One coalesced pending call drains all // queued payloads from Python's main thread. // ============================================================================ -namespace { - -class DeferredCleanupQueue; - -using CleanupFn = void (*)(void*) noexcept; - -// Preallocated queue item that transfers one payload out of CUDA's callback. +// Intrusive base for payloads transferred out of CUDA's callback. struct DeferredCleanupItem { DeferredCleanupItem* next = nullptr; - void* payload; - CleanupFn cleanup; - - DeferredCleanupItem(void* payload_, CleanupFn cleanup_) noexcept - : payload(payload_), cleanup(cleanup_) {} + virtual ~DeferredCleanupItem() noexcept = default; }; +namespace { + // Process-lifetime MPSC queue that drains payloads from Python's main thread. class DeferredCleanupQueue { public: @@ -230,7 +226,8 @@ class DeferredCleanupQueue { accepting_.store(false, std::memory_order_release); } - // Reattempt scheduling for payloads left queued after an earlier failure. + // Reattempt scheduling after Py_AddPendingCall() found CPython's bounded + // pending-call queue full and left payloads queued for a later safe entry. void retry_schedule() noexcept { schedule(); } @@ -244,8 +241,14 @@ class DeferredCleanupQueue { // Coalesce all queued work behind at most one CPython pending call. void schedule() noexcept { - if (!accepting_.load(std::memory_order_acquire) || - !head_.load(std::memory_order_acquire)) { + if (!accepting_.load(std::memory_order_acquire)) { + return; + } + if (!Py_IsInitialized() || py_is_finalizing()) { + stop(); + return; + } + if (!head_.load(std::memory_order_acquire)) { return; } bool expected = false; @@ -273,7 +276,6 @@ class DeferredCleanupQueue { head_.exchange(nullptr, std::memory_order_acquire)) { while (list) { DeferredCleanupItem* next = list->next; - list->cleanup(list->payload); delete list; list = next; } @@ -296,28 +298,13 @@ class DeferredCleanupQueue { // Published once at module initialization and intentionally never freed. std::atomic deferred_cleanup_queue{nullptr}; -void stop_deferred_cleanup() { - if (DeferredCleanupQueue* queue = - deferred_cleanup_queue.load(std::memory_order_acquire)) { - queue->stop(); - } -} - -DeferredCleanupItem* make_cleanup_item( - void* payload, CleanupFn cleanup) { +void ensure_deferred_cleanup_ready() { DeferredCleanupQueue* queue = deferred_cleanup_queue.load(std::memory_order_acquire); if (!queue) { throw std::runtime_error("deferred cleanup is not initialized"); } queue->retry_schedule(); - return new DeferredCleanupItem(payload, cleanup); -} - -// Execute an item's cleanup synchronously before CUDA has acquired ownership. -void execute_cleanup_now(DeferredCleanupItem* item) noexcept { - item->cleanup(item->payload); - delete item; } // CUDA's CUhostFn ABI is void (*)(void*); recover and enqueue the cleanup item. @@ -331,14 +318,23 @@ void enqueue_cleanup(void* item) noexcept { } // namespace +// Module initialization calls this once with the GIL held, which serializes +// the check, allocation, and publication below. void initialize_deferred_cleanup() { if (deferred_cleanup_queue.load(std::memory_order_acquire)) { return; } auto* queue = new DeferredCleanupQueue(); deferred_cleanup_queue.store(queue, std::memory_order_release); - if (Py_AtExit(stop_deferred_cleanup) != 0) { - queue->stop(); +} + +void retry_deferred_cleanup() noexcept { + if (!Py_IsInitialized() || py_is_finalizing()) { + return; + } + if (DeferredCleanupQueue* queue = + deferred_cleanup_queue.load(std::memory_order_acquire)) { + queue->retry_schedule(); } } @@ -377,6 +373,31 @@ class HandleRegistry { return {}; } + template + Handle get_or_create(const Key& key, Factory&& create) { + std::lock_guard lock(mutex_); + auto it = map_.find(key); + if (it != map_.end()) { + if (Handle h = it->second.lock()) { + return h; + } + map_.erase(it); + } + + Handle h = create(); + if (h) { + map_[key] = h; + } + return h; + } + + MapType drain() noexcept { + std::lock_guard lock(mutex_); + MapType extracted; + extracted.swap(map_); + return extracted; + } + private: std::mutex mutex_; MapType map_; @@ -1278,15 +1299,57 @@ LibraryHandle get_kernel_library(const KernelHandle& h) noexcept { namespace { -// Slot table layout (internal). Each graph maps CUgraphNode -> a fixed-size -// array of type-erased owners. The width is the most any single node needs: a -// kernel node holds its kernel and its packed arguments; a host node holds its -// callback and the userData. The table is heap-allocated and retained on the -// graph as a user object, so the driver frees it -- and every owner in it -- -// when the graph is destroyed. -constexpr std::size_t SLOTS_PER_NODE = 2; -using NodeSlots = std::array; -using GraphSlotTable = std::map; +struct NodeAttachment; +using GraphAttachmentMap = std::map; + +struct GraphHierarchy; + +// Standard-layout alias target for GraphHandle. +struct GraphBoxBase { + CUgraph resource = nullptr; +}; + +// Canonical state for one CUgraph. Its GraphHandle aliases resource, whose +// address remains stable for the lifetime of the hierarchy. +struct GraphBox : GraphBoxBase { + GraphHierarchy* hierarchy = nullptr; // Non-owning back-reference. + GraphBox* parent = nullptr; // Null for the root graph. + CUgraphNode owner_node = nullptr; // Node in parent that owns this graph. + GraphAttachmentMap attachments; // Non-owning attachment index. + HandleRegistry node_handles; + + GraphBox( + CUgraph resource_, + GraphHierarchy* hierarchy_, + GraphBox* parent_ = nullptr, + CUgraphNode owner_node_ = nullptr) noexcept + : GraphBoxBase{resource_}, + hierarchy(hierarchy_), + parent(parent_), + owner_node(owner_node_) {} +}; + +// Shared owner of stable GraphBox storage. Every GraphHandle aliases the same +// control block, so any graph handle keeps the entire hierarchy alive. +struct GraphHierarchy { + std::list graphs; // Parent boxes precede their descendants. + std::list graveyard; // Retired child graph tombstones. + + GraphBox* root() noexcept { + return graphs.empty() ? nullptr : &graphs.front(); + } +}; + +// See REGISTRY_DESIGN.md (Level 1: Driver Handle -> Resource Handle) +static HandleRegistry graph_registry; + +struct NodeAttachment : DeferredCleanupItem { + CUuserObject object = nullptr; + std::array owners; + + NodeAttachment(OpaqueHandle owner0, OpaqueHandle owner1) + : owners{std::move(owner0), std::move(owner1)} {} +}; // shared_ptr deleters for the payloads that need one. Typed handles convert to // OpaqueHandle by assignment and reuse their own control block, so they need no @@ -1303,60 +1366,92 @@ void free_deleter(const void* p) noexcept { std::free(const_cast(p)); } -void cleanup_graph_slot_table(void* table) noexcept { - delete static_cast(table); +GraphBox* get_box(const GraphHandle& h) noexcept { + auto* value = reinterpret_cast(h.get()); + return const_cast( + static_cast(value)); } -struct GraphBox { - CUgraph resource; - GraphHandle h_parent; // Keeps parent alive for child/branch graphs - mutable GraphSlotTable* slot_table = nullptr; // Lazily created; owned by the graph's user object -}; +// Rekey a staged attachment map from source nodes to their cloned nodes. +// The caller must release the GIL before calling this function. +CUresult rekey_attachments( + GraphAttachmentMap& attachments, CUgraph cloned_graph) { + if (!cloned_graph) { + return CUDA_ERROR_INVALID_VALUE; + } + if (!p_cuGraphNodeFindInClone) { + return CUDA_ERROR_NOT_SUPPORTED; + } -const GraphBox* get_box(const GraphHandle& h) { - const CUgraph* p = h.get(); - return reinterpret_cast( - reinterpret_cast(p) - offsetof(GraphBox, resource) - ); + GraphAttachmentMap remapped; + while (!attachments.empty()) { + auto attachment = attachments.extract(attachments.begin()); + CUgraphNode cloned_node = nullptr; + CUresult status = p_cuGraphNodeFindInClone( + &cloned_node, attachment.key(), cloned_graph); + if (status != CUDA_SUCCESS) { + return status; + } + attachment.key() = cloned_node; + if (!remapped.insert(std::move(attachment)).inserted) { + return CUDA_ERROR_INVALID_VALUE; + } + } + attachments.swap(remapped); + return CUDA_SUCCESS; } -// Return box's slot table, creating it on first use. The table is retained on -// the graph as a user object (MOVE transfers our only reference into the -// graph), so it -- and every owner in it -- is freed when the graph is -// destroyed. Returns nullptr if the driver lacks user-object support or a -// driver call fails; the cached pointer is non-owning. -GraphSlotTable* ensure_slot_table(const GraphBox* box) { - if (box->slot_table) { - return box->slot_table; - } - if (!p_cuUserObjectCreate || !p_cuGraphRetainUserObject || !p_cuUserObjectRelease) { - return nullptr; +// Recursively copy and rekey attachments for a cloned graph hierarchy. +// The caller must release the GIL before calling this function. +CUresult copy_attachments( + const GraphBox& source, + GraphBox& clone, + GraphAttachmentMap& attachments, + std::list& subgraphs) { + if (!p_cuGraphNodeFindInClone || !p_cuGraphChildGraphNodeGetGraph) { + return CUDA_ERROR_NOT_SUPPORTED; } - auto* table = new GraphSlotTable(); - DeferredCleanupItem* cleanup_item = nullptr; - try { - cleanup_item = make_cleanup_item(table, cleanup_graph_slot_table); - } catch (...) { - delete table; - throw; + + attachments = source.attachments; + CUresult status = rekey_attachments(attachments, clone.resource); + if (status != CUDA_SUCCESS) { + return status; } - CUuserObject user_obj = nullptr; - { - GILReleaseGuard gil; - if (p_cuUserObjectCreate(&user_obj, cleanup_item, - reinterpret_cast(enqueue_cleanup), - 1, CU_USER_OBJECT_NO_DESTRUCTOR_SYNC) != CUDA_SUCCESS) { - execute_cleanup_now(cleanup_item); - return nullptr; + + for (const GraphBox& source_child : source.hierarchy->graphs) { + if (source_child.parent != &source || !source_child.resource) { + continue; } - if (p_cuGraphRetainUserObject(box->resource, user_obj, 1, - CU_GRAPH_USER_OBJECT_MOVE) != CUDA_SUCCESS) { - p_cuUserObjectRelease(user_obj, 1); // drops refcount to 0 -> frees table - return nullptr; + + CUgraphNode cloned_owner = nullptr; + status = p_cuGraphNodeFindInClone( + &cloned_owner, source_child.owner_node, clone.resource); + if (status != CUDA_SUCCESS) { + return status; + } + + CUgraph cloned_graph = nullptr; + status = p_cuGraphChildGraphNodeGetGraph( + cloned_owner, &cloned_graph); + if (status != CUDA_SUCCESS) { + return status; + } + + GraphBox& cloned_child = subgraphs.emplace_back( + cloned_graph, + clone.hierarchy, + &clone, + cloned_owner); + status = copy_attachments( + source_child, + cloned_child, + cloned_child.attachments, + subgraphs); + if (status != CUDA_SUCCESS) { + return status; } } - box->slot_table = table; // non-owning cache; the user object owns it - return table; + return CUDA_SUCCESS; } } // namespace @@ -1370,36 +1465,279 @@ OpaqueHandle make_opaque_malloc(void* buf) { return OpaqueHandle(static_cast(buf), free_deleter); } -GraphHandle create_graph_handle(CUgraph graph) { - auto box = std::shared_ptr( - new GraphBox{graph, {}}, - [](const GraphBox* b) { +struct PreparedAttachmentState { + GraphHandle h_graph; + NodeAttachment* replacement = nullptr; + GraphAttachmentMap::node_type replacement_entry; + + explicit PreparedAttachmentState(GraphHandle h_graph_) + : h_graph(std::move(h_graph_)) {} +}; + +void PreparedAttachmentDeleter::operator()( + PreparedAttachmentState* state) const noexcept { + if (!state) { + return; + } + if (state->replacement) { + GraphBox* box = get_box(state->h_graph); + if (box->resource) { GILReleaseGuard gil; - p_cuGraphDestroy(b->resource); - delete b; + p_cuGraphReleaseUserObject( + box->resource, state->replacement->object, 1); + } + } + delete state; +} + +GraphHandle create_graph_handle(CUgraph graph) { + if (!graph) { + return {}; + } + + 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; + p_cuGraphDestroy(root->resource); + } + retry_deferred_cleanup(); + delete hierarchy; } ); - return GraphHandle(box, &box->resource); + 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; } -GraphHandle create_graph_handle_ref(CUgraph graph, const GraphHandle& h_parent) { - auto box = std::make_shared(GraphBox{graph, h_parent}); - return GraphHandle(box, &box->resource); +GraphHandle create_child_graph_handle( + CUgraph child_graph, const GraphHandle& h_parent, + CUgraphNode owner_node) { + if (!child_graph || !h_parent || !owner_node) { + return {}; + } + if (GraphHandle h_graph = graph_registry.lookup(child_graph)) { + 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); + try { + graph_registry.register_handle(child_graph, h_child); + } catch (...) { + hierarchy->graphs.pop_back(); + throw; + } + return h_child; } -CUresult graph_set_slot(const GraphHandle& h_graph, CUgraphNode node, - unsigned int slot, OpaqueHandle owner) { - if (!h_graph || slot >= SLOTS_PER_NODE) { +CUresult graph_get_attachment( + const GraphHandle& h_graph, CUgraphNode node, + OpaqueHandle* owner0, OpaqueHandle* owner1) { + if (!h_graph || !node || (!owner0 && !owner1)) { return CUDA_ERROR_INVALID_VALUE; } - if (!owner) { - return CUDA_SUCCESS; // nothing to retain; don't force table creation + if (owner0) { + owner0->reset(); + } + if (owner1) { + owner1->reset(); + } + + GraphBox* box = get_box(h_graph); + if (!box->resource) { + return CUDA_ERROR_INVALID_VALUE; + } + auto it = box->attachments.find(node); + if (it != box->attachments.end()) { + if (owner0) { + *owner0 = it->second->owners[0]; + } + if (owner1) { + *owner1 = it->second->owners[1]; + } + } + return CUDA_SUCCESS; +} + +CUresult graph_prepare_attachment( + const GraphHandle& h_graph, + OpaqueHandle owner0, + OpaqueHandle owner1, + PreparedAttachment* out_prepared) { + if (!out_prepared) { + return CUDA_ERROR_INVALID_VALUE; } - GraphSlotTable* table = ensure_slot_table(get_box(h_graph)); - if (!table) { + out_prepared->reset(); + if (!h_graph) { + return CUDA_ERROR_INVALID_VALUE; + } + + GraphBox* box = get_box(h_graph); + if (!box->resource) { + return CUDA_ERROR_INVALID_VALUE; + } + if (!p_cuGraphReleaseUserObject) { return CUDA_ERROR_NOT_SUPPORTED; } - (*table)[node][slot] = std::move(owner); + + PreparedAttachment prepared( + new PreparedAttachmentState(h_graph)); + if (owner0 || owner1) { + if (!p_cuUserObjectCreate || !p_cuUserObjectRelease || + !p_cuGraphRetainUserObject) { + return CUDA_ERROR_NOT_SUPPORTED; + } + + ensure_deferred_cleanup_ready(); + prepared->replacement = new NodeAttachment( + std::move(owner0), std::move(owner1)); + GraphAttachmentMap staged; + try { + staged.emplace(nullptr, prepared->replacement); + prepared->replacement_entry = + staged.extract(staged.begin()); + } catch (...) { + delete prepared->replacement; + prepared->replacement = nullptr; + throw; + } + auto* cleanup_item = + static_cast( + prepared->replacement); + + CUuserObject object = nullptr; + CUresult status; + { + GILReleaseGuard gil; + status = p_cuUserObjectCreate( + &object, cleanup_item, + reinterpret_cast(enqueue_cleanup), + 1, CU_USER_OBJECT_NO_DESTRUCTOR_SYNC); + if (status != CUDA_SUCCESS) { + prepared->replacement_entry.mapped() = nullptr; + delete prepared->replacement; + prepared->replacement = nullptr; + return status; + } + prepared->replacement->object = object; + status = p_cuGraphRetainUserObject( + box->resource, object, 1, CU_GRAPH_USER_OBJECT_MOVE); + if (status != CUDA_SUCCESS) { + prepared->replacement_entry.mapped() = nullptr; + prepared->replacement = nullptr; + p_cuUserObjectRelease(object, 1); + return status; + } + } + } + + *out_prepared = std::move(prepared); + return CUDA_SUCCESS; +} + +CUresult graph_commit_attachment( + PreparedAttachment& prepared, + CUgraphNode node) { + if (!prepared) { + return CUDA_ERROR_INVALID_VALUE; + } + + GraphHandle h_graph = prepared->h_graph; + GraphBox* box = get_box(h_graph); + if (!box->resource || (!node && !prepared->replacement)) { + delete prepared.release(); + return CUDA_ERROR_INVALID_VALUE; + } + if (!node) { + delete prepared.release(); + return CUDA_SUCCESS; + } + + NodeAttachment* previous = nullptr; + auto it = box->attachments.find(node); + if (it == box->attachments.end()) { + if (prepared->replacement) { + prepared->replacement_entry.key() = node; + auto result = box->attachments.insert( + std::move(prepared->replacement_entry)); + if (!result.inserted) { + prepared->replacement_entry = + std::move(result.node); + delete prepared.release(); + return CUDA_ERROR_INVALID_VALUE; + } + } + } else { + previous = it->second; + if (prepared->replacement) { + it->second = prepared->replacement; + } else { + box->attachments.erase(it); + } + } + + delete prepared.release(); + if (!previous) { + return CUDA_SUCCESS; + } + GILReleaseGuard gil; + return p_cuGraphReleaseUserObject( + box->resource, previous->object, 1); +} + +CUresult graph_clone_attachments( + const GraphHandle& h_clone, + const GraphHandle& h_source) { + if (!h_clone || !h_source) { + return CUDA_ERROR_INVALID_VALUE; + } + + GraphBox* clone = get_box(h_clone); + GraphBox* source = get_box(h_source); + if (!clone->resource || !source->resource || + !clone->attachments.empty()) { + return CUDA_ERROR_INVALID_VALUE; + } + + GraphAttachmentMap attachments = source->attachments; + std::list subgraphs; + CUresult status; + { + GILReleaseGuard gil; + status = copy_attachments( + *source, *clone, attachments, subgraphs); + } + if (status != CUDA_SUCCESS) { + return status; + } + + clone->attachments.swap(attachments); + if (subgraphs.empty()) { + return CUDA_SUCCESS; + } + + auto first = subgraphs.begin(); + clone->hierarchy->graphs.splice( + clone->hierarchy->graphs.end(), subgraphs); + for (auto it = first; it != clone->hierarchy->graphs.end(); ++it) { + GraphHandle h_graph(h_clone, &it->resource); + graph_registry.register_handle(it->resource, h_graph); + } return CUDA_SUCCESS; } @@ -1417,8 +1755,11 @@ GraphExecHandle create_graph_exec_handle(CUgraphExec graph_exec) { auto box = std::shared_ptr( new GraphExecBox{graph_exec}, [](const GraphExecBox* b) { - GILReleaseGuard gil; - p_cuGraphExecDestroy(b->resource); + { + GILReleaseGuard gil; + p_cuGraphExecDestroy(b->resource); + } + retry_deferred_cleanup(); delete b; } ); @@ -1439,21 +1780,61 @@ static const GraphNodeBox* get_box(const GraphNodeHandle& h) { ); } -// See REGISTRY_DESIGN.md (Level 1: Driver Handle -> Resource Handle) -static HandleRegistry graph_node_registry; +// graphs is ordered parent-before-child. Nulling a selected box marks its +// later descendants, whose parent pointers remain valid after list splicing. +// This permits one allocation-free sweep of the hierarchy. +void invalidate_child_graph_state( + const GraphHandle& h_parent, + CUgraphNode owner_node) noexcept { + if (!h_parent || !owner_node) { + return; + } -GraphNodeHandle create_graph_node_handle(CUgraphNode node, const GraphHandle& h_graph) { - if (node) { - if (auto h = graph_node_registry.lookup(node)) { - return h; + GraphBox* parent = get_box(h_parent); + if (!parent->resource) { + return; + } + GraphHierarchy& hierarchy = *parent->hierarchy; + for (auto it = hierarchy.graphs.begin(); + it != hierarchy.graphs.end();) { + auto graph = it++; + bool is_owned_root = graph->parent == parent && + graph->owner_node == owner_node; + bool is_descendant = graph->parent && + !graph->parent->resource; + if (!is_owned_root && !is_descendant) { + continue; + } + + // Empty node_handles and invalidate each one. + for (auto& entry : graph->node_handles.drain()) { + if (GraphNodeHandle h_node = entry.second.lock()) { + get_box(h_node)->resource = nullptr; + } } + graph_registry.unregister_handle(graph->resource); + graph->resource = nullptr; + graph->attachments.clear(); + hierarchy.graveyard.splice( + hierarchy.graveyard.end(), hierarchy.graphs, graph); } - auto box = std::make_shared(GraphNodeBox{node, h_graph}); - GraphNodeHandle h(box, &box->resource); - if (node) { - graph_node_registry.register_handle(node, h); +} + +GraphNodeHandle create_graph_node_handle(CUgraphNode node, const GraphHandle& h_graph) { + if (!node) { + auto box = std::make_shared( + GraphNodeBox{nullptr, h_graph}); + return GraphNodeHandle(box, &box->resource); } - return h; + + GraphBox* graph = get_box(h_graph); + return graph->node_handles.get_or_create( + node, + [node, &h_graph] { + auto box = std::make_shared( + GraphNodeBox{node, h_graph}); + return GraphNodeHandle(box, &box->resource); + }); } GraphHandle graph_node_get_graph(const GraphNodeHandle& h) noexcept { @@ -1461,13 +1842,18 @@ GraphHandle graph_node_get_graph(const GraphNodeHandle& h) noexcept { } void invalidate_graph_node(const GraphNodeHandle& h) noexcept { - if (h) { - CUgraphNode node = get_box(h)->resource; - if (node) { - graph_node_registry.unregister_handle(node); - } - get_box(h)->resource = nullptr; + if (!h) { + return; + } + + const GraphNodeBox* node_box = get_box(h); + CUgraphNode node = node_box->resource; + if (!node) { + return; } + GraphBox* graph = get_box(node_box->h_graph); + graph->node_handles.unregister_handle(node); + node_box->resource = nullptr; } // ============================================================================ diff --git a/cuda_core/cuda/core/_cpp/resource_handles.hpp b/cuda_core/cuda/core/_cpp/resource_handles.hpp index bb6bb3aa84..f6fdeff06b 100644 --- a/cuda_core/cuda/core/_cpp/resource_handles.hpp +++ b/cuda_core/cuda/core/_cpp/resource_handles.hpp @@ -112,6 +112,9 @@ extern decltype(&cuGraphExecDestroy) p_cuGraphExecDestroy; extern decltype(&cuUserObjectCreate) p_cuUserObjectCreate; extern decltype(&cuUserObjectRelease) p_cuUserObjectRelease; extern decltype(&cuGraphRetainUserObject) p_cuGraphRetainUserObject; +extern decltype(&cuGraphReleaseUserObject) p_cuGraphReleaseUserObject; +extern decltype(&cuGraphNodeFindInClone) p_cuGraphNodeFindInClone; +extern decltype(&cuGraphChildGraphNodeGetGraph) p_cuGraphChildGraphNodeGetGraph; // Linker extern decltype(&cuLinkDestroy) p_cuLinkDestroy; @@ -259,6 +262,7 @@ void py_object_user_object_destroy(void* py_object) noexcept; // Initialize the process-lifetime CUDA user-object cleanup queue. Called once // from module initialization while Python is fully initialized. void initialize_deferred_cleanup(); +void retry_deferred_cleanup() noexcept; // Return the context dependency associated with a stream handle, if any. ContextHandle get_stream_context(const StreamHandle& h) noexcept; @@ -462,26 +466,19 @@ LibraryHandle get_kernel_library(const KernelHandle& h) noexcept; // Graph handle functions // ============================================================================ -// Wrap an externally-created CUgraph with RAII cleanup. -// When the last reference is released, cuGraphDestroy is called automatically. -// The caller must have already created the graph via cuGraphCreate. +// Create the owning handle for a root graph and its hierarchy. GraphHandle create_graph_handle(CUgraph graph); -// Create a non-owning graph handle that keeps h_parent alive. -// Use for graphs owned by a child/conditional node in a parent graph. -// The child graph will NOT be destroyed when this handle is released, -// but h_parent will be prevented from destruction while this handle exists. -GraphHandle create_graph_handle_ref(CUgraph graph, const GraphHandle& h_parent); +// Create the canonical handle for a graph whose CUDA lifetime is owned by a +// node in h_parent. +GraphHandle create_child_graph_handle( + CUgraph child_graph, const GraphHandle& h_parent, CUgraphNode owner_node); // ============================================================================ -// Graph slot attachments +// Graph node attachments // -// A graph carries a side table that keeps resources used by its nodes (kernel -// arguments, host callbacks, events, ...) alive for as long as the graph can -// execute. The table is created on first use and retained on the CUgraph as a -// user object, so the driver releases it -- and everything attached through it -// -- when the graph is destroyed. The table layout is an internal detail; -// callers use the abstract API below. +// Each resource-bearing node has one immutable attachment retained on its +// CUgraph as a CUDA user object. // ============================================================================ // Type-erased shared owner of an attached resource. Typed handles such as @@ -497,14 +494,44 @@ OpaqueHandle make_opaque_py(PyObject* obj); // Build an OpaqueHandle from a malloc'd buffer: std::free on release. OpaqueHandle make_opaque_malloc(void* buf); -// Attach owner to one of node's fixed slots on h_graph, replacing whatever was -// there. The graph's slot table is created on first use. A null owner is a -// no-op (returns CUDA_SUCCESS without creating the table), so callers need not -// guard optional owners. Returns CUDA_SUCCESS, or an error if slot is out of -// range or the graph cannot hold a table (e.g. the driver lacks user-object -// support). -CUresult graph_set_slot(const GraphHandle& h_graph, CUgraphNode node, - unsigned int slot, OpaqueHandle owner); +struct PreparedAttachmentState; +struct PreparedAttachmentDeleter { + void operator()(PreparedAttachmentState* state) const noexcept; +}; +using PreparedAttachment = + std::unique_ptr; + +// Copy requested owners from node's current attachment. Pass nullptr to ignore +// either owner; a missing attachment produces empty handles. +CUresult graph_get_attachment( + const GraphHandle& h_graph, + CUgraphNode node, + OpaqueHandle* owner0, + OpaqueHandle* owner1); + +// Create and graph-retain a replacement attachment before a CUDA mutation. +// Destruction rolls the prepared attachment back unless it is committed. +CUresult graph_prepare_attachment( + const GraphHandle& h_graph, + OpaqueHandle owner0, + OpaqueHandle owner1, + PreparedAttachment* out_prepared); + +// Publish a prepared attachment after the CUDA mutation succeeds. A null node +// retains the attachment anonymously without publishing node metadata. +CUresult graph_commit_attachment( + PreparedAttachment& prepared, + CUgraphNode node); + +// Copy attachment metadata from a source graph hierarchy into its CUDA clone. +CUresult graph_clone_attachments( + const GraphHandle& h_clone, + const GraphHandle& h_source); + +// Invalidate cuda-core state for child graphs CUDA destroyed with owner_node. +void invalidate_child_graph_state( + const GraphHandle& h_parent, + CUgraphNode owner_node) noexcept; // ============================================================================ // Graph exec handle functions diff --git a/cuda_core/cuda/core/_resource_handles.pxd b/cuda_core/cuda/core/_resource_handles.pxd index aaf3e75222..2f481fee4f 100644 --- a/cuda_core/cuda/core/_resource_handles.pxd +++ b/cuda_core/cuda/core/_resource_handles.pxd @@ -5,7 +5,7 @@ from libc.stddef cimport size_t from libc.stdint cimport intptr_t -from libcpp.memory cimport shared_ptr +from libcpp.memory cimport shared_ptr, unique_ptr from cuda.bindings cimport cydriver from cuda.bindings cimport cynvrtc @@ -63,6 +63,14 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core": # custom deleter. ctypedef shared_ptr[const void] OpaqueHandle + cppclass PreparedAttachmentState: + pass + cppclass PreparedAttachmentDeleter: + pass + ctypedef unique_ptr[ + PreparedAttachmentState, PreparedAttachmentDeleter + ] PreparedAttachment + # as_cu() - extract the raw CUDA handle (inline C++) cydriver.CUcontext as_cu(ContextHandle h) noexcept nogil cydriver.CUgreenCtx as_cu(GreenCtxHandle h) noexcept nogil @@ -160,6 +168,7 @@ cdef StreamHandle create_stream_handle( cdef StreamHandle create_stream_handle_ref(cydriver.CUstream stream) except+ nogil cdef StreamHandle create_stream_handle_with_owner(cydriver.CUstream stream, object owner) except+ nogil cdef void py_object_user_object_destroy(void* py_object) noexcept nogil +cdef void retry_deferred_cleanup() noexcept cdef ContextHandle get_stream_context(const StreamHandle& h) noexcept nogil cdef StreamHandle get_legacy_stream() except+ nogil cdef StreamHandle get_per_thread_stream() except+ nogil @@ -227,14 +236,25 @@ cdef LibraryHandle get_kernel_library(const KernelHandle& h) noexcept nogil # Graph handles cdef GraphHandle create_graph_handle(cydriver.CUgraph graph) except+ nogil -cdef GraphHandle create_graph_handle_ref(cydriver.CUgraph graph, const GraphHandle& h_parent) except+ nogil +cdef GraphHandle create_child_graph_handle( + cydriver.CUgraph child_graph, const GraphHandle& h_parent, + cydriver.CUgraphNode owner_node) except+ nogil -# Graph slot attachments +# Graph node attachments cdef OpaqueHandle make_opaque_py(object obj) except+ cdef OpaqueHandle make_opaque_malloc(void* buf) except+ -cdef cydriver.CUresult graph_set_slot( +cdef cydriver.CUresult graph_get_attachment( const GraphHandle& h_graph, cydriver.CUgraphNode node, - unsigned int slot, OpaqueHandle owner) except+ + OpaqueHandle* owner0, OpaqueHandle* owner1) except+ +cdef cydriver.CUresult graph_prepare_attachment( + const GraphHandle& h_graph, OpaqueHandle owner0, OpaqueHandle owner1, + PreparedAttachment* out_prepared) except+ +cdef cydriver.CUresult graph_commit_attachment( + PreparedAttachment& prepared, cydriver.CUgraphNode node) except+ +cdef cydriver.CUresult graph_clone_attachments( + const GraphHandle& h_clone, const GraphHandle& h_source) except+ +cdef void invalidate_child_graph_state( + const GraphHandle& h_parent, cydriver.CUgraphNode owner_node) noexcept # Graph exec handles cdef GraphExecHandle create_graph_exec_handle(cydriver.CUgraphExec graph_exec) except+ nogil diff --git a/cuda_core/cuda/core/_resource_handles.pyi b/cuda_core/cuda/core/_resource_handles.pyi index 4236df5d6e..f11f6f08e0 100644 --- a/cuda_core/cuda/core/_resource_handles.pyi +++ b/cuda_core/cuda/core/_resource_handles.pyi @@ -2,7 +2,7 @@ from __future__ import annotations -from libcpp.memory import shared_ptr +from libcpp.memory import shared_ptr, unique_ptr ContextHandle = shared_ptr GreenCtxHandle = shared_ptr @@ -25,4 +25,5 @@ OpaqueArrayHandle = shared_ptr MipmappedArrayHandle = shared_ptr TexObjectHandle = shared_ptr SurfObjectHandle = shared_ptr -OpaqueHandle = shared_ptr \ No newline at end of file +OpaqueHandle = shared_ptr +PreparedAttachment = unique_ptr \ No newline at end of file diff --git a/cuda_core/cuda/core/_resource_handles.pyx b/cuda_core/cuda/core/_resource_handles.pyx index 7407da4556..f6fb6ac4e2 100644 --- a/cuda_core/cuda/core/_resource_handles.pyx +++ b/cuda_core/cuda/core/_resource_handles.pyx @@ -62,6 +62,7 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core": void py_object_user_object_destroy "cuda_core::py_object_user_object_destroy" ( void* py_object) noexcept nogil void initialize_deferred_cleanup "cuda_core::initialize_deferred_cleanup" () except+ + void retry_deferred_cleanup "cuda_core::retry_deferred_cleanup" () noexcept ContextHandle get_stream_context "cuda_core::get_stream_context" ( const StreamHandle& h) noexcept nogil StreamHandle get_legacy_stream "cuda_core::get_legacy_stream" () except+ nogil @@ -149,15 +150,25 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core": # Graph handles GraphHandle create_graph_handle "cuda_core::create_graph_handle" ( cydriver.CUgraph graph) except+ nogil - GraphHandle create_graph_handle_ref "cuda_core::create_graph_handle_ref" ( - cydriver.CUgraph graph, const GraphHandle& h_parent) except+ nogil + GraphHandle create_child_graph_handle "cuda_core::create_child_graph_handle" ( + cydriver.CUgraph child_graph, const GraphHandle& h_parent, + cydriver.CUgraphNode owner_node) except+ nogil - # Graph slot attachments + # Graph node attachments OpaqueHandle make_opaque_py "cuda_core::make_opaque_py" (object obj) except+ OpaqueHandle make_opaque_malloc "cuda_core::make_opaque_malloc" (void* buf) except+ - cydriver.CUresult graph_set_slot "cuda_core::graph_set_slot" ( + cydriver.CUresult graph_get_attachment "cuda_core::graph_get_attachment" ( const GraphHandle& h_graph, cydriver.CUgraphNode node, - unsigned int slot, OpaqueHandle owner) except+ + OpaqueHandle* owner0, OpaqueHandle* owner1) except+ + cydriver.CUresult graph_prepare_attachment "cuda_core::graph_prepare_attachment" ( + const GraphHandle& h_graph, OpaqueHandle owner0, OpaqueHandle owner1, + PreparedAttachment* out_prepared) except+ + cydriver.CUresult graph_commit_attachment "cuda_core::graph_commit_attachment" ( + PreparedAttachment& prepared, cydriver.CUgraphNode node) except+ + cydriver.CUresult graph_clone_attachments "cuda_core::graph_clone_attachments" ( + const GraphHandle& h_clone, const GraphHandle& h_source) except+ + void invalidate_child_graph_state "cuda_core::invalidate_child_graph_state" ( + const GraphHandle& h_parent, cydriver.CUgraphNode owner_node) noexcept # Graph exec handles GraphExecHandle create_graph_exec_handle "cuda_core::create_graph_exec_handle" ( @@ -315,6 +326,9 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core": void* p_cuUserObjectCreate "reinterpret_cast(cuda_core::p_cuUserObjectCreate)" void* p_cuUserObjectRelease "reinterpret_cast(cuda_core::p_cuUserObjectRelease)" void* p_cuGraphRetainUserObject "reinterpret_cast(cuda_core::p_cuGraphRetainUserObject)" + void* p_cuGraphReleaseUserObject "reinterpret_cast(cuda_core::p_cuGraphReleaseUserObject)" + void* p_cuGraphNodeFindInClone "reinterpret_cast(cuda_core::p_cuGraphNodeFindInClone)" + void* p_cuGraphChildGraphNodeGetGraph "reinterpret_cast(cuda_core::p_cuGraphChildGraphNodeGetGraph)" # Linker void* p_cuLinkDestroy "reinterpret_cast(cuda_core::p_cuLinkDestroy)" @@ -375,7 +389,9 @@ cdef void _init_driver_fn_pointers() noexcept: global p_cuMemPoolImportPointer global p_cuLibraryLoadFromFile, p_cuLibraryLoadData, p_cuLibraryUnload, p_cuLibraryGetKernel global p_cuGraphDestroy, p_cuGraphExecDestroy - global p_cuUserObjectCreate, p_cuUserObjectRelease, p_cuGraphRetainUserObject + global p_cuUserObjectCreate, p_cuUserObjectRelease + global p_cuGraphRetainUserObject, p_cuGraphReleaseUserObject + global p_cuGraphNodeFindInClone, p_cuGraphChildGraphNodeGetGraph global p_cuLinkDestroy global p_cuGraphicsUnmapResources, p_cuGraphicsUnregisterResource global p_cuDevSmResourceSplit @@ -439,6 +455,9 @@ cdef void _init_driver_fn_pointers() noexcept: p_cuUserObjectCreate = _get_driver_fn("cuUserObjectCreate") p_cuUserObjectRelease = _get_driver_fn("cuUserObjectRelease") p_cuGraphRetainUserObject = _get_driver_fn("cuGraphRetainUserObject") + p_cuGraphReleaseUserObject = _get_driver_fn("cuGraphReleaseUserObject") + p_cuGraphNodeFindInClone = _get_driver_fn("cuGraphNodeFindInClone") + p_cuGraphChildGraphNodeGetGraph = _get_driver_fn("cuGraphChildGraphNodeGetGraph") # Linker p_cuLinkDestroy = _get_driver_fn("cuLinkDestroy") diff --git a/cuda_core/cuda/core/graph/_graph_builder.pyi b/cuda_core/cuda/core/graph/_graph_builder.pyi index f19d4067eb..6da55e71b4 100644 --- a/cuda_core/cuda/core/graph/_graph_builder.pyi +++ b/cuda_core/cuda/core/graph/_graph_builder.pyi @@ -490,4 +490,7 @@ class Graph: __all__ = ['Graph', 'GraphBuilder', 'GraphCompleteOptions', 'GraphDebugPrintOptions'] def _instantiate_graph(h_graph, options: GraphCompleteOptions | None=None) -> Graph: - ... \ No newline at end of file + ... + +def _capture_callback_with_tail_failure_for_testing(gb: GraphBuilder, fn, *, user_data=None): + """Exercise anonymous attachment retention after node discovery fails.""" \ No newline at end of file diff --git a/cuda_core/cuda/core/graph/_graph_builder.pyx b/cuda_core/cuda/core/graph/_graph_builder.pyx index b1e6859f54..cd1d785143 100644 --- a/cuda_core/cuda/core/graph/_graph_builder.pyx +++ b/cuda_core/cuda/core/graph/_graph_builder.pyx @@ -10,12 +10,18 @@ from libc.stdint cimport intptr_t from cuda.bindings cimport cydriver from cuda.core.graph._graph_definition cimport GraphCondition, GraphDefinition -from cuda.core.graph._host_callback cimport _attach_host_callback_owners, _resolve_host_callback +from cuda.core.graph._host_callback cimport _resolve_host_callback from cuda.core._resource_handles cimport ( GraphHandle, OpaqueHandle, + PreparedAttachment, as_cu, as_py, - create_graph_exec_handle, create_graph_handle, create_graph_handle_ref, + create_child_graph_handle, create_graph_exec_handle, create_graph_handle, + graph_clone_attachments, + graph_commit_attachment, + graph_prepare_attachment, + invalidate_child_graph_state, + retry_deferred_cleanup, ) from cuda.core._stream cimport Stream from cuda.core._utils.cuda_utils cimport HANDLE_RETURN @@ -282,6 +288,7 @@ cdef class GraphBuilder: GB_end_capture_if_needed(self, True) self._h_graph.reset() self._h_stream.reset() + retry_deferred_cleanup() self._state = CLOSED self._stream = None @@ -782,16 +789,36 @@ cdef class GraphBuilder: # See https://github.com/NVIDIA/cuda-python/pull/879#issuecomment-3211054159 # for rationale - deps_info_trimmed = deps_info_out[:num_dependencies_out] - deps_info_update = [ - [ - handle_return( - driver.cuGraphAddChildGraphNode( - graph_out, *deps_info_trimmed, num_dependencies_out, as_py(child._h_graph) - ) - ) - ] - ] + [None] * (len(deps_info_out) - 1) + dependencies_out = deps_info_out[0] + new_node = handle_return( + driver.cuGraphAddChildGraphNode( + graph_out, dependencies_out, + num_dependencies_out, as_py(child._h_graph) + ) + ) + cdef cydriver.CUgraphNode c_new_node = ( + int(new_node) + ) + cdef cydriver.CUgraph embedded_graph = NULL + cdef cydriver.CUresult rollback_status + cdef GraphHandle h_embedded + try: + with nogil: + HANDLE_RETURN(cydriver.cuGraphChildGraphNodeGetGraph( + c_new_node, &embedded_graph)) + h_embedded = create_child_graph_handle( + embedded_graph, self._h_graph, c_new_node) + HANDLE_RETURN(graph_clone_attachments( + h_embedded, child._h_graph)) + except: + with nogil: + rollback_status = cydriver.cuGraphDestroyNode(c_new_node) + if rollback_status == cydriver.CUDA_SUCCESS: + invalidate_child_graph_state( + self._h_graph, c_new_node) + raise + + deps_info_update = [[new_node]] + [None] * (len(deps_info_out) - 1) handle_return( driver.cuStreamUpdateCaptureDependencies( stream_handle, @@ -829,30 +856,56 @@ cdef class GraphBuilder: pointer (caller manages lifetime). If bytes-like, the data is copied and its lifetime is tied to the graph. """ - GB_check_open(self) - cdef Stream stream = self._stream - cdef cydriver.CUstream c_stream = as_cu(stream._h_stream) - cdef cydriver.CUstreamCaptureStatus capture_status + GB_callback(self, fn, user_data, False) - with nogil: - _get_capture_info(c_stream, &capture_status, NULL) - if capture_status != cydriver.CU_STREAM_CAPTURE_STATUS_ACTIVE: - raise RuntimeError("Cannot add callback when graph is not being built") +cdef inline void GB_callback( + GraphBuilder gb, object fn, object user_data, + bint fail_tail_discovery_for_testing) except *: + GB_check_open(gb) + cdef Stream stream = gb._stream + cdef cydriver.CUstream c_stream = as_cu(stream._h_stream) + cdef cydriver.CUstreamCaptureStatus capture_status + + with nogil: + _get_capture_info(c_stream, &capture_status, NULL) - cdef cydriver.CUhostFn c_fn - cdef void* c_user_data = NULL - cdef OpaqueHandle fn_owner, data_owner - _resolve_host_callback(fn, user_data, &c_fn, &c_user_data, &fn_owner, &data_owner) + if capture_status != cydriver.CU_STREAM_CAPTURE_STATUS_ACTIVE: + raise RuntimeError("Cannot add callback when graph is not being built") - with nogil: - HANDLE_RETURN(cydriver.cuLaunchHostFunc(c_stream, c_fn, c_user_data)) + cdef cydriver.CUhostFn c_fn + cdef void* c_user_data = NULL + cdef OpaqueHandle fn_owner, data_owner + cdef PreparedAttachment prepared + _resolve_host_callback( + fn, user_data, &c_fn, &c_user_data, &fn_owner, &data_owner) + HANDLE_RETURN(graph_prepare_attachment( + gb._h_graph, fn_owner, data_owner, &prepared)) - # Capturing the host function added a node to the graph; it is now the - # stream's sole capture dependency. Key the callback's owners to it so - # they live in the graph's slot table like any explicitly-added node. - cdef cydriver.CUgraphNode host_node = _capture_tail_node(c_stream) - _attach_host_callback_owners(self._h_graph, host_node, fn_owner, data_owner) + with nogil: + HANDLE_RETURN(cydriver.cuLaunchHostFunc(c_stream, c_fn, c_user_data)) + + # Capturing the host function added a node to the graph; it is now the + # stream's sole capture dependency. Attach the callback's owners to it. + cdef cydriver.CUgraphNode host_node + cdef cydriver.CUresult commit_status + try: + if fail_tail_discovery_for_testing: + raise RuntimeError("forced capture tail discovery failure") + host_node = _capture_tail_node(c_stream) + except: + # CUDA added the callback, but its node cannot be identified. + # Retain its owners anonymously to prevent dangling pointers. + commit_status = graph_commit_attachment(prepared, NULL) + HANDLE_RETURN(commit_status) + raise + HANDLE_RETURN(graph_commit_attachment(prepared, host_node)) + + +def _capture_callback_with_tail_failure_for_testing( + GraphBuilder gb, fn, *, user_data=None): + """Exercise anonymous attachment retention after node discovery fails.""" + GB_callback(gb, fn, user_data, True) cdef inline int GB_check_open(GraphBuilder gb) except -1: @@ -893,9 +946,7 @@ cdef inline int GB_end_capture_if_needed(GraphBuilder gb, bint check_status) exc cdef inline GraphBuilder GB_init_forked(Stream stream, GraphHandle h_primary_graph): cdef GraphBuilder gb = GraphBuilder.__new__(GraphBuilder) # A FORKED builder captures into the primary's CUgraph. It holds the - # primary's GraphHandle so conditional bodies created on it (via - # GB_init_conditional -> create_graph_handle_ref(cond_graph, parent._h_graph)) - # have a valid parent handle to pin. + # primary's GraphHandle so conditional bodies share its graph hierarchy. gb._h_graph = h_primary_graph gb._h_stream = stream._h_stream gb._kind = FORKED @@ -904,9 +955,12 @@ cdef inline GraphBuilder GB_init_forked(Stream stream, GraphHandle h_primary_gra return gb -cdef inline GraphBuilder GB_init_conditional(Stream stream, cydriver.CUgraph cond_graph, GraphBuilder parent): +cdef inline GraphBuilder GB_init_conditional( + Stream stream, cydriver.CUgraph cond_graph, + GraphBuilder parent, cydriver.CUgraphNode owner_node): cdef GraphBuilder gb = GraphBuilder.__new__(GraphBuilder) - gb._h_graph = create_graph_handle_ref(cond_graph, parent._h_graph) + gb._h_graph = create_child_graph_handle( + cond_graph, parent._h_graph, owner_node) gb._h_stream = stream._h_stream gb._kind = CONDITIONAL_BODY gb._state = CAPTURE_NOT_STARTED @@ -962,9 +1016,9 @@ cdef inline tuple GB_cond_with_params(GraphBuilder gb, node_params): if status != driver.CUstreamCaptureStatus.CU_STREAM_CAPTURE_STATUS_ACTIVE: raise RuntimeError("Cannot add conditional node when not actively capturing") - deps_info_update = [ - [handle_return(driver.cuGraphAddNode(graph, *deps_info, num_dependencies, node_params))] - ] + [None] * (len(deps_info) - 1) + new_node = handle_return( + driver.cuGraphAddNode(graph, *deps_info, num_dependencies, node_params)) + deps_info_update = [[new_node]] + [None] * (len(deps_info) - 1) handle_return( driver.cuStreamUpdateCaptureDependencies( @@ -980,6 +1034,7 @@ cdef inline tuple GB_cond_with_params(GraphBuilder gb, node_params): gb._stream.device.create_stream(), int(node_params.conditional.phGraph_out[i]), gb, + int(new_node), ) for i in range(node_params.conditional.size) ) @@ -1008,6 +1063,7 @@ cdef class Graph: def close(self) -> None: """Destroy the graph.""" self._h_graph_exec.reset() + retry_deferred_cleanup() @property def handle(self) -> driver.CUgraphExec: diff --git a/cuda_core/cuda/core/graph/_graph_definition.pyi b/cuda_core/cuda/core/graph/_graph_definition.pyi index 89652d879b..9780b53b58 100644 --- a/cuda_core/cuda/core/graph/_graph_definition.pyi +++ b/cuda_core/cuda/core/graph/_graph_definition.pyi @@ -56,6 +56,12 @@ class GraphDefinition: A GraphDefinition is used to construct a graph explicitly by adding nodes and specifying dependencies. Once construction is complete, call instantiate() to obtain an executable Graph. + + Notes + ----- + Definitions that view the same root, child, or conditional graph hierarchy + share underlying graph state. Mutations anywhere in that hierarchy must be + externally synchronized. """ def __init__(self): diff --git a/cuda_core/cuda/core/graph/_graph_definition.pyx b/cuda_core/cuda/core/graph/_graph_definition.pyx index 0bd05318ca..b2516034d6 100644 --- a/cuda_core/cuda/core/graph/_graph_definition.pyx +++ b/cuda_core/cuda/core/graph/_graph_definition.pyx @@ -7,6 +7,7 @@ from __future__ import annotations from libc.stddef cimport size_t +from libc.stdint cimport uintptr_t from libcpp.vector cimport vector @@ -105,6 +106,12 @@ cdef class GraphDefinition: A GraphDefinition is used to construct a graph explicitly by adding nodes and specifying dependencies. Once construction is complete, call instantiate() to obtain an executable Graph. + + Notes + ----- + Definitions that view the same root, child, or conditional graph hierarchy + share underlying graph state. Mutations anywhere in that hierarchy must be + externally synchronized. """ def __init__(self): @@ -127,10 +134,13 @@ cdef class GraphDefinition: def __eq__(self, other: object) -> bool: if not isinstance(other, GraphDefinition): return NotImplemented - return as_intptr(self._h_graph) == as_intptr((other)._h_graph) + return ( + self._h_graph.get() + == (other)._h_graph.get() + ) def __hash__(self) -> int: - return hash(as_intptr(self._h_graph)) + return hash(self._h_graph.get()) @property def _entry(self) -> GraphNode: diff --git a/cuda_core/cuda/core/graph/_graph_node.pyx b/cuda_core/cuda/core/graph/_graph_node.pyx index 0adf69ad4b..37d94607bd 100644 --- a/cuda_core/cuda/core/graph/_graph_node.pyx +++ b/cuda_core/cuda/core/graph/_graph_node.pyx @@ -44,20 +44,23 @@ from cuda.core._resource_handles cimport ( GraphHandle, GraphNodeHandle, OpaqueHandle, + PreparedAttachment, as_cu, as_intptr, as_py, - create_graph_handle_ref, + create_child_graph_handle, create_graph_node_handle, + graph_clone_attachments, + graph_commit_attachment, graph_node_get_graph, - graph_set_slot, + graph_prepare_attachment, + invalidate_child_graph_state, invalidate_graph_node, make_opaque_py, ) from cuda.core._utils.cuda_utils cimport HANDLE_RETURN, _parse_fill_value from cuda.core.graph._host_callback cimport ( - _attach_host_callback_owners, _resolve_host_callback, ) @@ -160,14 +163,24 @@ cdef class GraphNode: already-destroyed node (no-op). """ cdef cydriver.CUgraphNode node = as_cu(self._h_node) + cdef GraphHandle h_graph + cdef cydriver.CUresult cleanup_status + cdef PreparedAttachment prepared if node == NULL: return - _node_registry.pop(self._h_node.get(), None) - invalidate_graph_node(self._h_node) + h_graph = graph_node_get_graph(self._h_node) + HANDLE_RETURN(graph_prepare_attachment( + h_graph, OpaqueHandle(), OpaqueHandle(), &prepared)) with nogil: HANDLE_RETURN(cydriver.cuGraphDestroyNode(node)) + cleanup_status = graph_commit_attachment(prepared, node) + invalidate_child_graph_state(h_graph, node) + _node_registry.pop(self._h_node.get(), None) + invalidate_graph_node(self._h_node) + HANDLE_RETURN(cleanup_status) + @property def pred(self) -> AdjacencySetProxy: """A mutable set-like view of this node's predecessors.""" @@ -330,10 +343,12 @@ cdef class GraphNode: cdef cydriver.CUdeviceptr c_dst cdef unsigned int val cdef unsigned int elem_size - cdef OpaqueHandle dst_slot_owner - dst_slot_owner = _resolve_memcpy_operand(dst, dst_owner, "dst", &c_dst) + cdef OpaqueHandle dst_attachment_owner + dst_attachment_owner = _resolve_memcpy_operand(dst, dst_owner, "dst", &c_dst) val, elem_size = _parse_fill_value(value) - return GN_memset(self, c_dst, dst_slot_owner, val, elem_size, width, height, pitch) + return GN_memset( + self, c_dst, dst_attachment_owner, + val, elem_size, width, height, pitch) def memcpy( self, @@ -384,10 +399,12 @@ cdef class GraphNode: """ cdef cydriver.CUdeviceptr c_dst cdef cydriver.CUdeviceptr c_src - cdef OpaqueHandle dst_slot_owner, src_slot_owner - dst_slot_owner = _resolve_memcpy_operand(dst, dst_owner, "dst", &c_dst) - src_slot_owner = _resolve_memcpy_operand(src, src_owner, "src", &c_src) - return GN_memcpy(self, c_dst, dst_slot_owner, c_src, src_slot_owner, size) + cdef OpaqueHandle dst_attachment_owner, src_attachment_owner + dst_attachment_owner = _resolve_memcpy_operand(dst, dst_owner, "dst", &c_dst) + src_attachment_owner = _resolve_memcpy_operand(src, src_owner, "src", &c_src) + return GN_memcpy( + self, c_dst, dst_attachment_owner, + c_src, src_attachment_owner, size) def embed(self, child: GraphDefinition) -> ChildGraphNode: """Add a child graph node depending on this node. @@ -604,7 +621,7 @@ cdef inline ConditionalNode _make_conditional_node( cdef GraphHandle h_branch for i in range(size): bg = params.conditional.phGraph_out[i] - h_branch = create_graph_handle_ref(bg, h_graph) + h_branch = create_child_graph_handle(bg, h_graph, new_node) branch_list.append(GraphDefinition._from_handle(h_branch)) cdef tuple branches = tuple(branch_list) @@ -673,7 +690,8 @@ cdef inline KernelNode GN_launch(GraphNode self, LaunchConfig conf, Kernel ker, cdef cydriver.CUgraphNode pred_node = as_cu(self._h_node) cdef cydriver.CUgraphNode* deps = NULL cdef size_t num_deps = 0 - cdef OpaqueHandle owner + cdef OpaqueHandle kernel_owner, args_owner + cdef PreparedAttachment prepared if pred_node != NULL: deps = &pred_node @@ -692,24 +710,20 @@ cdef inline KernelNode GN_launch(GraphNode self, LaunchConfig conf, Kernel ker, node_params.extra = NULL node_params.ctx = NULL + # Keep the kernel and argument objects alive because CUDA copies argument + # values but does not retain the resources they reference. + kernel_owner = ker._h_kernel + kernel_args = ker_args.kernel_args + if kernel_args is not None: + args_owner = make_opaque_py(kernel_args) + HANDLE_RETURN(graph_prepare_attachment( + h_graph, kernel_owner, args_owner, &prepared)) + with nogil: HANDLE_RETURN(cydriver.cuGraphAddKernelNode( &new_node, as_cu(h_graph), deps, num_deps, &node_params)) - # Slot 0 keeps the kernel loaded; slot 1 keeps the Python kernel-argument - # objects (notably device Buffers) alive for the graph's lifetime. The - # driver copies argument values into the node at add time but does not own - # the device memory they reference. - owner = ker._h_kernel - try: - HANDLE_RETURN(graph_set_slot(h_graph, new_node, 0, owner)) - kernel_args = ker_args.kernel_args - if kernel_args is not None: - owner = make_opaque_py(kernel_args) - HANDLE_RETURN(graph_set_slot(h_graph, new_node, 1, owner)) - except: - cydriver.cuGraphDestroyNode(new_node) # best effort - raise + HANDLE_RETURN(graph_commit_attachment(prepared, new_node)) return _registered(KernelNode._create_with_params( create_graph_node_handle(new_node, h_graph), @@ -836,18 +850,18 @@ cdef inline FreeNode GN_free(GraphNode self, cydriver.CUdeviceptr c_dptr): return _registered(FreeNode._create_with_params(create_graph_node_handle(new_node, h_graph), c_dptr)) -cdef inline OpaqueHandle _buffer_slot_owner(Buffer buf, str label): - """Copy a Buffer's device-pointer handle into a graph slot owner.""" - cdef OpaqueHandle slot_owner +cdef inline OpaqueHandle _buffer_attachment_owner(Buffer buf, str label): + """Copy a Buffer's device-pointer handle into an attachment owner.""" + cdef OpaqueHandle attachment_owner if not buf._h_ptr: raise ValueError(f"{label} Buffer has no active allocation") - slot_owner = buf._h_ptr - return slot_owner + attachment_owner = buf._h_ptr + return attachment_owner cdef inline OpaqueHandle _resolve_memcpy_operand( object operand, object owner, str side, cydriver.CUdeviceptr* out_ptr): - """Resolve a memcpy/memset operand to a pointer and optional slot owner. + """Resolve an operand to a pointer and optional attachment owner. ``operand`` is a :class:`Buffer` or a raw integer address; its device pointer is written to ``out_ptr``. For a :class:`Buffer` operand, returns an @@ -862,6 +876,7 @@ cdef inline OpaqueHandle _resolve_memcpy_operand( If a :class:`Buffer` operand or ``*_owner`` has no active allocation. """ cdef Buffer buf + cdef OpaqueHandle attachment_owner if isinstance(operand, Buffer): if owner is not None: @@ -869,14 +884,14 @@ cdef inline OpaqueHandle _resolve_memcpy_operand( f"{side}_owner cannot be used when {side} is a Buffer" ) buf = operand - slot_owner = _buffer_slot_owner(buf, side) + attachment_owner = _buffer_attachment_owner(buf, side) out_ptr[0] = as_cu(buf._h_ptr) - return slot_owner + return attachment_owner out_ptr[0] = operand if owner is None: return OpaqueHandle() if isinstance(owner, Buffer): - return _buffer_slot_owner(owner, f"{side}_owner") + return _buffer_attachment_owner(owner, f"{side}_owner") return make_opaque_py(owner) @@ -890,6 +905,7 @@ cdef inline MemsetNode GN_memset( cdef cydriver.CUgraphNode pred_node = as_cu(self._h_node) cdef cydriver.CUgraphNode* deps = NULL cdef size_t num_deps = 0 + cdef PreparedAttachment prepared if pred_node != NULL: deps = &pred_node @@ -907,20 +923,17 @@ cdef inline MemsetNode GN_memset( memset_params.height = height memset_params.pitch = pitch + if dst_owner: + HANDLE_RETURN(graph_prepare_attachment( + h_graph, dst_owner, OpaqueHandle(), &prepared)) + with nogil: HANDLE_RETURN(cydriver.cuGraphAddMemsetNode( &new_node, as_cu(h_graph), deps, num_deps, &memset_params, ctx)) - # Retain the destination allocation for the graph's lifetime (slot 0). Roll - # the node back if slot attachment fails, so it cannot outlive an - # unretained allocation. if dst_owner: - try: - HANDLE_RETURN(graph_set_slot(h_graph, new_node, 0, dst_owner)) - except: - cydriver.cuGraphDestroyNode(new_node) # best effort - raise + HANDLE_RETURN(graph_commit_attachment(prepared, new_node)) return _registered(MemsetNode._create_with_params( create_graph_node_handle(new_node, h_graph), c_dst, @@ -972,27 +985,24 @@ cdef inline MemcpyNode GN_memcpy( cdef cydriver.CUgraphNode pred_node = as_cu(self._h_node) cdef cydriver.CUgraphNode* deps = NULL cdef size_t num_deps = 0 + cdef PreparedAttachment prepared if pred_node != NULL: deps = &pred_node num_deps = 1 cdef cydriver.CUcontext ctx = NULL + if dst_owner or src_owner: + HANDLE_RETURN(graph_prepare_attachment( + h_graph, dst_owner, src_owner, &prepared)) + with nogil: HANDLE_RETURN(cydriver.cuCtxGetCurrent(&ctx)) HANDLE_RETURN(cydriver.cuGraphAddMemcpyNode( &new_node, as_cu(h_graph), deps, num_deps, ¶ms, ctx)) - # Retain operand allocations for the graph's lifetime (dst -> slot 0, src -> - # slot 1). Roll the node back if slot attachment fails. - try: - if dst_owner: - HANDLE_RETURN(graph_set_slot(h_graph, new_node, 0, dst_owner)) - if src_owner: - HANDLE_RETURN(graph_set_slot(h_graph, new_node, 1, src_owner)) - except: - cydriver.cuGraphDestroyNode(new_node) # best effort - raise + if dst_owner or src_owner: + HANDLE_RETURN(graph_commit_attachment(prepared, new_node)) return _registered(MemcpyNode._create_with_params( create_graph_node_handle(new_node, h_graph), c_dst, c_src, size, @@ -1005,6 +1015,7 @@ cdef inline ChildGraphNode GN_embed(GraphNode self, GraphDefinition child_def): cdef cydriver.CUgraphNode pred_node = as_cu(self._h_node) cdef cydriver.CUgraphNode* deps = NULL cdef size_t num_deps = 0 + cdef cydriver.CUresult rollback_status if pred_node != NULL: deps = &pred_node @@ -1015,11 +1026,21 @@ cdef inline ChildGraphNode GN_embed(GraphNode self, GraphDefinition child_def): &new_node, as_cu(h_graph), deps, num_deps, as_cu(child_def._h_graph))) cdef cydriver.CUgraph embedded_graph = NULL - with nogil: - HANDLE_RETURN(cydriver.cuGraphChildGraphNodeGetGraph( - new_node, &embedded_graph)) - - cdef GraphHandle h_embedded = create_graph_handle_ref(embedded_graph, h_graph) + cdef GraphHandle h_embedded + try: + with nogil: + HANDLE_RETURN(cydriver.cuGraphChildGraphNodeGetGraph( + new_node, &embedded_graph)) + h_embedded = create_child_graph_handle( + embedded_graph, h_graph, new_node) + HANDLE_RETURN(graph_clone_attachments( + h_embedded, child_def._h_graph)) + except: + with nogil: + rollback_status = cydriver.cuGraphDestroyNode(new_node) + if rollback_status == cydriver.CUDA_SUCCESS: + invalidate_child_graph_state(h_graph, new_node) + raise return _registered(ChildGraphNode._create_with_params( create_graph_node_handle(new_node, h_graph), h_embedded)) @@ -1032,21 +1053,21 @@ cdef inline EventRecordNode GN_record_event(GraphNode self, Event ev): cdef cydriver.CUgraphNode* deps = NULL cdef size_t num_deps = 0 cdef OpaqueHandle owner + cdef PreparedAttachment prepared if pred_node != NULL: deps = &pred_node num_deps = 1 + owner = ev._h_event + HANDLE_RETURN(graph_prepare_attachment( + h_graph, owner, OpaqueHandle(), &prepared)) + with nogil: HANDLE_RETURN(cydriver.cuGraphAddEventRecordNode( &new_node, as_cu(h_graph), deps, num_deps, as_cu(ev._h_event))) - owner = ev._h_event - try: - HANDLE_RETURN(graph_set_slot(h_graph, new_node, 0, owner)) - except: - cydriver.cuGraphDestroyNode(new_node) # best effort - raise + HANDLE_RETURN(graph_commit_attachment(prepared, new_node)) return _registered(EventRecordNode._create_with_params( create_graph_node_handle(new_node, h_graph), ev._h_event)) @@ -1059,21 +1080,21 @@ cdef inline EventWaitNode GN_wait_event(GraphNode self, Event ev): cdef cydriver.CUgraphNode* deps = NULL cdef size_t num_deps = 0 cdef OpaqueHandle owner + cdef PreparedAttachment prepared if pred_node != NULL: deps = &pred_node num_deps = 1 + owner = ev._h_event + HANDLE_RETURN(graph_prepare_attachment( + h_graph, owner, OpaqueHandle(), &prepared)) + with nogil: HANDLE_RETURN(cydriver.cuGraphAddEventWaitNode( &new_node, as_cu(h_graph), deps, num_deps, as_cu(ev._h_event))) - owner = ev._h_event - try: - HANDLE_RETURN(graph_set_slot(h_graph, new_node, 0, owner)) - except: - cydriver.cuGraphDestroyNode(new_node) # best effort - raise + HANDLE_RETURN(graph_commit_attachment(prepared, new_node)) return _registered(EventWaitNode._create_with_params( create_graph_node_handle(new_node, h_graph), ev._h_event)) @@ -1087,6 +1108,7 @@ cdef inline HostCallbackNode GN_callback(GraphNode self, object fn, object user_ cdef cydriver.CUgraphNode* deps = NULL cdef size_t num_deps = 0 cdef OpaqueHandle fn_owner, data_owner + cdef PreparedAttachment prepared if pred_node != NULL: deps = &pred_node @@ -1095,18 +1117,14 @@ cdef inline HostCallbackNode GN_callback(GraphNode self, object fn, object user_ _resolve_host_callback( fn, user_data, &node_params.fn, &node_params.userData, &fn_owner, &data_owner) + HANDLE_RETURN(graph_prepare_attachment( + h_graph, fn_owner, data_owner, &prepared)) with nogil: HANDLE_RETURN(cydriver.cuGraphAddHostNode( &new_node, as_cu(h_graph), deps, num_deps, &node_params)) - # Roll the node back if owner attachment fails: an unretained host node - # leaves the driver holding a raw pointer into freed Python memory. - try: - _attach_host_callback_owners(h_graph, new_node, fn_owner, data_owner) - except: - cydriver.cuGraphDestroyNode(new_node) # best effort - raise + HANDLE_RETURN(graph_commit_attachment(prepared, new_node)) cdef object callable_obj = fn if not isinstance(fn, ct._CFuncPtr) else None return _registered(HostCallbackNode._create_with_params( diff --git a/cuda_core/cuda/core/graph/_host_callback.pxd b/cuda_core/cuda/core/graph/_host_callback.pxd index dac249c74e..fc77809c84 100644 --- a/cuda_core/cuda/core/graph/_host_callback.pxd +++ b/cuda_core/cuda/core/graph/_host_callback.pxd @@ -4,7 +4,7 @@ from cuda.bindings cimport cydriver -from cuda.core._resource_handles cimport GraphHandle, OpaqueHandle +from cuda.core._resource_handles cimport OpaqueHandle cdef bint _is_py_host_trampoline(cydriver.CUhostFn fn) noexcept nogil @@ -13,7 +13,3 @@ cdef void _resolve_host_callback( object fn, object user_data, cydriver.CUhostFn* out_fn, void** out_user_data, OpaqueHandle* out_fn_owner, OpaqueHandle* out_data_owner) except * - -cdef int _attach_host_callback_owners( - const GraphHandle& h_graph, cydriver.CUgraphNode node, - OpaqueHandle fn_owner, OpaqueHandle data_owner) except -1 diff --git a/cuda_core/cuda/core/graph/_host_callback.pyx b/cuda_core/cuda/core/graph/_host_callback.pyx index bed2d8152f..27f251abea 100644 --- a/cuda_core/cuda/core/graph/_host_callback.pyx +++ b/cuda_core/cuda/core/graph/_host_callback.pyx @@ -9,13 +9,10 @@ from libc.string cimport memcpy as c_memcpy from cuda.bindings cimport cydriver from cuda.core._resource_handles cimport ( - GraphHandle, OpaqueHandle, - graph_set_slot, make_opaque_malloc, make_opaque_py, ) -from cuda.core._utils.cuda_utils cimport HANDLE_RETURN import ctypes as ct @@ -38,8 +35,7 @@ cdef void _resolve_host_callback( On return ``*out_fn`` / ``*out_user_data`` are ready to pass to ``cuGraphAddHostNode`` or ``cuLaunchHostFunc``. ``*out_fn_owner`` owns the callback object; ``*out_data_owner`` owns a copied ``user_data`` buffer and - is left null otherwise. The caller attaches the owners to the node's graph - slots. + is left null otherwise. The caller attaches both owners to the graph node. """ if isinstance(fn, ct._CFuncPtr): out_fn[0] = ct.cast(fn, ct.c_void_p).value @@ -65,15 +61,3 @@ cdef void _resolve_host_callback( out_user_data[0] = fn out_fn_owner[0] = make_opaque_py(fn) - - -cdef int _attach_host_callback_owners( - const GraphHandle& h_graph, cydriver.CUgraphNode node, - OpaqueHandle fn_owner, OpaqueHandle data_owner) except -1: - """Attach a resolved host callback's owners to its node's graph slots: the - callback in slot 0 and any copied ``user_data`` buffer in slot 1. - """ - HANDLE_RETURN(graph_set_slot(h_graph, node, 0, fn_owner)) - if data_owner: - HANDLE_RETURN(graph_set_slot(h_graph, node, 1, data_owner)) - return 0 diff --git a/cuda_core/cuda/core/graph/_subclasses.pyx b/cuda_core/cuda/core/graph/_subclasses.pyx index a393dfff69..2fa08e2a6a 100644 --- a/cuda_core/cuda/core/graph/_subclasses.pyx +++ b/cuda_core/cuda/core/graph/_subclasses.pyx @@ -24,7 +24,7 @@ from cuda.core._resource_handles cimport ( as_cu, as_intptr, create_event_handle_ref, - create_graph_handle_ref, + create_child_graph_handle, create_kernel_handle_ref, graph_node_get_graph, ) @@ -470,7 +470,8 @@ cdef class ChildGraphNode(GraphNode): with nogil: HANDLE_RETURN(cydriver.cuGraphChildGraphNodeGetGraph(node, &child_graph)) cdef GraphHandle h_graph = graph_node_get_graph(h_node) - cdef GraphHandle h_child = create_graph_handle_ref(child_graph, h_graph) + cdef GraphHandle h_child = create_child_graph_handle( + child_graph, h_graph, node) return ChildGraphNode._create_with_params(h_node, h_child) def __repr__(self) -> str: @@ -656,9 +657,9 @@ cdef class ConditionalNode(GraphNode): cdef GraphHandle h_branch if cond_params.phGraph_out is not None: for i in range(size): - h_branch = create_graph_handle_ref( + h_branch = create_child_graph_handle( int(cond_params.phGraph_out[i]), - h_graph) + h_graph, node) branch_list.append(GraphDefinition._from_handle(h_branch)) cdef tuple branches = tuple(branch_list) diff --git a/cuda_core/docs/source/release/1.2.0-notes.rst b/cuda_core/docs/source/release/1.2.0-notes.rst index 3f90ba4b91..6a047c9cfe 100644 --- a/cuda_core/docs/source/release/1.2.0-notes.rst +++ b/cuda_core/docs/source/release/1.2.0-notes.rst @@ -6,9 +6,21 @@ ``cuda.core`` 1.2.0 Release Notes ================================== +Fixes and enhancements +---------------------- + +- Graph node resources are now retained independently across graph clones, + executable graphs, updates, node deletion, and in-flight launches. Previously, + modifying a graph definition could release resources still used by an + existing executable graph. Releasing graph-attached objects could also call + CUDA from a CUDA-invoked callback, where CUDA API calls are prohibited; these + objects are now released from a safe context. + (`#2357 `__, + `#2371 `__) + Deprecation Notices ------------------- -* Support for using ``cuda-core`` with Python 3.10 is deprecated and will be +- Support for using ``cuda-core`` with Python 3.10 is deprecated and will be removed in a future version. Python 3.10 reaches end of life in October 2026 per the `CPython support cycle `_. diff --git a/cuda_core/tests/graph/test_graph_builder.py b/cuda_core/tests/graph/test_graph_builder.py index 3247bbecef..6fcb540cb7 100644 --- a/cuda_core/tests/graph/test_graph_builder.py +++ b/cuda_core/tests/graph/test_graph_builder.py @@ -4,6 +4,8 @@ """GraphBuilder stream capture tests.""" import gc +import time +import weakref import numpy as np import pytest @@ -13,6 +15,18 @@ from cuda.core import Device, LaunchConfig, LegacyPinnedMemoryResource, launch from cuda.core.graph import GraphBuilder, GraphDefinition +from cuda.core.graph._graph_builder import ( + _capture_callback_with_tail_failure_for_testing, +) + + +def _wait_until(predicate, timeout=5.0): + deadline = time.monotonic() + timeout + while not predicate(): + if time.monotonic() >= deadline: + raise AssertionError(f"condition not satisfied within {timeout}s") + gc.collect() + time.sleep(0.02) def test_graph_is_building(init_cuda): @@ -294,7 +308,7 @@ def read_byte(data): @pytest.mark.agent_authored(model="claude-opus-4.8") def test_graph_capture_callback_python_survives_del(init_cuda): - """Captured host callback is retained in the graph slot table after del.""" + """Captured callback is retained by its graph-node user object after del.""" called = [False] def my_callback(): @@ -341,6 +355,34 @@ def read_byte(data): assert result[0] == 0xAB +@pytest.mark.agent_authored(model="gpt-5.6") +def test_capture_tail_discovery_failure_preserves_callback_attachment(init_cuda): + """A committed captured node keeps its owner when discovery fails.""" + called = [False] + + def callback(): + called[0] = True + + callback_weak = weakref.ref(callback) + stream = Device().create_stream() + gb = stream.create_graph_builder().begin_building() + + with pytest.raises(RuntimeError, match="forced capture tail discovery failure"): + _capture_callback_with_tail_failure_for_testing(gb, callback) + + del callback + gc.collect() + assert callback_weak() is not None + + graph = gb.end_building().complete() + graph.launch(stream) + stream.sync() + assert called[0] + + del graph, gb + _wait_until(lambda: callback_weak() is None) + + @pytest.mark.skipif(tuple(int(i) for i in np.__version__.split(".")[:2]) < (2, 1), reason="need numpy 2.1.0+") def test_graph_child_graph(init_cuda): mod = compile_common_kernels() diff --git a/cuda_core/tests/graph/test_graph_definition_lifetime.py b/cuda_core/tests/graph/test_graph_definition_lifetime.py index a01ff8162f..78ccb3008b 100644 --- a/cuda_core/tests/graph/test_graph_definition_lifetime.py +++ b/cuda_core/tests/graph/test_graph_definition_lifetime.py @@ -20,17 +20,18 @@ from cuda_python_test_helpers import under_compute_sanitizer # Resource finalization triggered by graph destruction is not synchronous. A -# CUDA user-object callback transfers the slot table to a pending-call -# cleanup queue, which releases each owner from Python's main thread. Release is -# deterministic at the reference-count level, so the predicate normally flips -# within milliseconds; this budget only bounds a slow/loaded runner. It stays a -# hard failure rather than a warning so a real leak still fails the suite. +# CUDA user-object callback transfers each node attachment bundle to a +# pending-call cleanup queue, which releases each owner from Python's main +# thread. Release is deterministic at the reference-count level, so the +# predicate normally flips within milliseconds; this budget only bounds a +# slow/loaded runner. It stays a hard failure rather than a warning so a real +# leak still fails the suite. # Compute-sanitizer slows everything down, hence the larger ceiling there. _FINALIZE_TIMEOUT = 30.0 if under_compute_sanitizer() else 5.0 class _Sentinel: - """Weak-referenceable stand-in for an owner attached to a graph slot. + """Weak-referenceable stand-in for a graph-node attachment owner. Bare ``object()`` instances do not support weak references, so tests that observe owner release through a :class:`weakref.ref` use this trivial @@ -80,6 +81,7 @@ def _wait_until(predicate, timeout=None, interval=0.02): ChildGraphNode, ConditionalNode, GraphDefinition, + HostCallbackNode, KernelNode, ) @@ -177,6 +179,52 @@ def test_reconstructed_body_survives_parent_deletion(init_cuda): assert body.nodes() == set() +@pytest.mark.parametrize("builder, expected_count", _COND_BUILDERS) +@pytest.mark.agent_authored(model="gpt-5.6") +def test_destroying_conditional_node_releases_branch_attachments(init_cuda, builder, expected_count): + """Destroying a conditional owner releases every branch attachment.""" + graph_def = GraphDefinition() + condition = try_create_condition(graph_def) + branches = builder(graph_def, condition) + callback_refs = [] + + for branch in branches: + + def callback(): + pass + + callback_refs.append(weakref.ref(callback)) + branch.callback(callback) + + del callback, branches + gc.collect() + assert len(callback_refs) == expected_count + assert all(ref() is not None for ref in callback_refs) + + owner = next(node for node in graph_def.nodes() if isinstance(node, ConditionalNode)) + owner.destroy() + del owner + + _wait_until(lambda: all(ref() is None for ref in callback_refs)) + + +@pytest.mark.parametrize("builder, expected_count", _COND_BUILDERS) +@pytest.mark.agent_authored(model="gpt-5.6") +def test_destroying_conditional_node_invalidates_branch_handles(init_cuda, builder, expected_count): + """Destroying a conditional owner invalidates every branch graph and node.""" + graph_def = GraphDefinition() + condition = try_create_condition(graph_def) + branches = builder(graph_def, condition) + branch_nodes = [branch.callback(lambda: None) for branch in branches] + + owner = next(node for node in graph_def.nodes() if isinstance(node, ConditionalNode)) + owner.destroy() + + assert len(branches) == expected_count + assert all(int(branch.handle) == 0 for branch in branches) + assert all(not node.is_valid for node in branch_nodes) + + # ============================================================================= # Child graph (embed) lifetime # ============================================================================= @@ -202,6 +250,28 @@ def test_child_graph_survives_parent_deletion(init_cuda): assert len(child_ref.nodes()) == 2 +@pytest.mark.agent_authored(model="gpt-5.6") +def test_destroyed_child_graph_preserves_python_identity(init_cuda): + """Tombstoning does not change graph equality or hashing.""" + parent = GraphDefinition() + first_owner = parent.embed(GraphDefinition()) + second_owner = parent.embed(GraphDefinition()) + first = first_owner.child_graph + second = second_owner.child_graph + first_hash = hash(first) + second_hash = hash(second) + values = {first: "first", second: "second"} + + first_owner.destroy() + second_owner.destroy() + + assert hash(first) == first_hash + assert hash(second) == second_hash + assert first != second + assert values[first] == "first" + assert values[second] == "second" + + def test_nested_child_graph_lifetime(init_cuda): """Grandchild graph keeps entire ancestor chain alive.""" mod = compile_common_kernels() @@ -228,6 +298,161 @@ def test_nested_child_graph_lifetime(init_cuda): assert len(grandchild.nodes()) == 1 +@pytest.mark.agent_authored(model="gpt-5.6") +def test_reconstructed_child_view_releases_node_attachment(init_cuda): + """Equivalent child-graph handles share attachment metadata.""" + parent = GraphDefinition() + child_node = parent.embed(GraphDefinition()) + child_view = child_node.child_graph + + def callback(): + pass + + callback_weak = weakref.ref(callback) + callback_node = child_view.callback(callback) + del callback, callback_node, child_view, child_node + gc.collect() + assert callback_weak() is not None + + reconstructed_child_node = next(node for node in parent.nodes() if isinstance(node, ChildGraphNode)) + reconstructed_child = reconstructed_child_node.child_graph + reconstructed_callback = next(node for node in reconstructed_child.nodes() if isinstance(node, HostCallbackNode)) + reconstructed_callback.destroy() + + del reconstructed_callback, reconstructed_child, reconstructed_child_node + _wait_until(lambda: callback_weak() is None) + + +@pytest.mark.agent_authored(model="gpt-5.6") +def test_populated_child_clone_releases_attachment_on_node_destroy(init_cuda): + """An embedded clone imports metadata for attachments created before embed.""" + + def callback(): + pass + + callback_weak = weakref.ref(callback) + child = GraphDefinition() + source_callback = child.callback(callback) + parent = GraphDefinition() + child_node = parent.embed(child) + embedded_child = child_node.child_graph + embedded_callback = next(node for node in embedded_child.nodes() if isinstance(node, HostCallbackNode)) + + del callback, source_callback, child + gc.collect() + assert callback_weak() is not None + + embedded_callback.destroy() + del embedded_callback, embedded_child + _wait_until(lambda: callback_weak() is None) + + +@pytest.mark.agent_authored(model="gpt-5.6") +def test_nested_child_clone_releases_attachment_on_node_destroy(init_cuda): + """Attachment metadata is imported recursively through nested child clones.""" + + def callback(): + pass + + callback_weak = weakref.ref(callback) + inner = GraphDefinition() + inner_callback = inner.callback(callback) + middle = GraphDefinition() + middle_child = middle.embed(inner) + outer = GraphDefinition() + outer_child = outer.embed(middle) + + embedded_middle = outer_child.child_graph + embedded_child_node = next(node for node in embedded_middle.nodes() if isinstance(node, ChildGraphNode)) + embedded_inner = embedded_child_node.child_graph + embedded_callback = next(node for node in embedded_inner.nodes() if isinstance(node, HostCallbackNode)) + + del callback, inner_callback, inner, middle_child, middle + gc.collect() + assert callback_weak() is not None + + embedded_callback.destroy() + del embedded_callback, embedded_inner, embedded_child_node, embedded_middle + _wait_until(lambda: callback_weak() is None) + + +@pytest.mark.agent_authored(model="gpt-5.6") +def test_destroying_child_node_recursively_releases_subgraph_attachments(init_cuda): + """Destroying a child owner releases attachments in all nested clones.""" + + def callback(): + pass + + callback_weak = weakref.ref(callback) + inner = GraphDefinition() + inner_callback = inner.callback(callback) + middle = GraphDefinition() + middle_child = middle.embed(inner) + outer = GraphDefinition() + outer_child = outer.embed(middle) + + del callback, inner_callback, inner, middle_child, middle + gc.collect() + assert callback_weak() is not None + + outer_child.destroy() + del outer_child + _wait_until(lambda: callback_weak() is None) + + +@pytest.mark.agent_authored(model="gpt-5.6") +def test_destroying_child_node_invalidates_embedded_handles(init_cuda): + """Destroying a child owner invalidates all embedded graph and node views.""" + inner = GraphDefinition() + inner.callback(lambda: None) + middle = GraphDefinition() + middle.embed(inner) + outer = GraphDefinition() + outer_child = outer.embed(middle) + + embedded_middle = outer_child.child_graph + embedded_child = next(node for node in embedded_middle.nodes() if isinstance(node, ChildGraphNode)) + embedded_inner = embedded_child.child_graph + embedded_callback = next(node for node in embedded_inner.nodes() if isinstance(node, HostCallbackNode)) + + outer_child.destroy() + + assert not outer_child.is_valid + assert int(embedded_middle.handle) == 0 + assert int(embedded_inner.handle) == 0 + assert not embedded_child.is_valid + assert not embedded_callback.is_valid + + +@pytest.mark.agent_authored(model="gpt-5.6") +def test_builder_embedded_clone_releases_attachment_on_node_destroy(init_cuda): + """GraphBuilder.embed imports metadata from the captured child graph.""" + + def callback(): + pass + + callback_weak = weakref.ref(callback) + child = Device().create_graph_builder().begin_building() + child.callback(callback) + child.end_building() + + parent = Device().create_graph_builder().begin_building() + parent.embed(child) + parent.end_building() + parent_def = parent.graph_definition + child_node = next(node for node in parent_def.nodes() if isinstance(node, ChildGraphNode)) + embedded_child = child_node.child_graph + embedded_callback = next(node for node in embedded_child.nodes() if isinstance(node, HostCallbackNode)) + + del callback, child + gc.collect() + assert callback_weak() is not None + + embedded_callback.destroy() + del embedded_callback, embedded_child + _wait_until(lambda: callback_weak() is None) + + # ============================================================================= # Event lifetime — event nodes should keep the Event alive # ============================================================================= @@ -358,6 +583,8 @@ def test_event_survives_graph_clone_and_execution(init_cuda): stream = dev.create_stream() handle_return(driver.cuGraphLaunch(graph_exec, driver.CUstream(int(stream.handle)))) stream.sync() + handle_return(driver.cuGraphExecDestroy(graph_exec)) + handle_return(driver.cuGraphDestroy(cloned_cu_graph)) # ============================================================================= @@ -419,15 +646,12 @@ def fill_queue_and_destroy(): worker.join() assert queue_was_full == [True] - # Preparing another attachment retries the first payload if its scheduling - # attempt observed the full queue. CUDA may also invoke its destructor - # later, after pending-call queue space is available. - second_callback = _ThreadRecordingCallback(finalized_threads) - second_graph = GraphDefinition() - second_graph.callback(second_callback) - del second_callback, second_graph + # A later safe cuda-core close retries after the main thread has had an + # opportunity to drain the foreign pending calls. + retry_builder = Device().create_graph_builder() + retry_builder.close() - _wait_until(lambda: len(finalized_threads) == 2) + _wait_until(lambda: len(finalized_threads) == 1) assert set(finalized_threads) == {main_thread} @@ -482,6 +706,200 @@ def my_callback(): assert called[0] +@pytest.mark.agent_authored(model="gpt-5.6") +def test_callback_survives_source_node_deletion(init_cuda): + """An instantiated graph independently retains a deleted source node's callback.""" + called = [False] + + def callback(): + called[0] = True + + callback_weak = weakref.ref(callback) + graph_def = GraphDefinition() + node = graph_def.callback(callback) + graph = graph_def.instantiate() + + del callback + node.destroy() + del node, graph_def + gc.collect() + assert callback_weak() is not None + + stream = Device().create_stream() + graph.launch(stream) + stream.sync() + assert called[0] + + del graph + _wait_until(lambda: callback_weak() is None) + + +@pytest.mark.agent_authored(model="gpt-5.6") +def test_inflight_launch_retains_attachments_until_completion(init_cuda): + """An in-flight launch retains the final allocation reference.""" + from cuda.core._utils._weak_handles import weak_handle + + _skip_if_no_mempool() + dev = Device() + mr = DeviceMemoryResource(dev) + buf = mr.allocate(8, stream=dev.default_stream) + dev.default_stream.sync() + dptr = int(buf.handle) + allocation_weak = weak_handle(buf) + callback_started = threading.Event() + release_callback = threading.Event() + + def blocking_callback(): + callback_started.set() + assert release_callback.wait(timeout=_FINALIZE_TIMEOUT) + + graph_def = GraphDefinition() + copy_node = graph_def.memcpy(dptr, dptr + 4, 4, dst_owner=buf, src_owner=buf) + callback_node = copy_node.callback(blocking_callback) + graph = graph_def.instantiate() + + buf.close() + del buf, blocking_callback, callback_node, copy_node, graph_def + gc.collect() + assert allocation_weak + + stream = dev.create_stream() + graph.launch(stream) + assert callback_started.wait(timeout=_FINALIZE_TIMEOUT) + + close_started = threading.Event() + close_errors = [] + + def close_graph(graph_to_close): + close_started.set() + try: + graph_to_close.close() + except BaseException as exc: # pragma: no cover - surfaced below + close_errors.append(exc) + + close_thread = threading.Thread(target=close_graph, args=(graph,)) + close_thread.start() + assert close_started.wait(timeout=_FINALIZE_TIMEOUT) + try: + time.sleep(0.05) + assert allocation_weak + finally: + release_callback.set() + stream.sync() + close_thread.join(timeout=_FINALIZE_TIMEOUT) + + assert not close_thread.is_alive() + assert close_errors == [] + del close_graph, close_thread, graph + _wait_until(lambda: not allocation_weak) + + +@pytest.mark.agent_authored(model="gpt-5.6") +def test_callback_survives_source_node_deletion_after_clone(init_cuda): + """A clone independently retains a callback removed from its source graph.""" + from cuda.core._utils.cuda_utils import driver, handle_return + + called = [False] + + def callback(): + called[0] = True + + callback_weak = weakref.ref(callback) + graph_def = GraphDefinition() + node = graph_def.callback(callback) + cloned_graph = handle_return(driver.cuGraphClone(driver.CUgraph(graph_def.handle))) + + node.destroy() + del callback, node, graph_def + gc.collect() + assert callback_weak() is not None + + graph_exec = handle_return(driver.cuGraphInstantiate(cloned_graph, 0)) + try: + stream = Device().create_stream() + handle_return(driver.cuGraphLaunch(graph_exec, driver.CUstream(int(stream.handle)))) + stream.sync() + assert called[0] + finally: + handle_return(driver.cuGraphExecDestroy(graph_exec)) + handle_return(driver.cuGraphDestroy(cloned_graph)) + + _wait_until(lambda: callback_weak() is None) + + +@pytest.mark.agent_authored(model="gpt-5.6") +def test_attachment_released_only_after_last_executable_is_destroyed(init_cuda): + """Each executable independently retains the source node's attachment.""" + called = [0] + + def callback(): + called[0] += 1 + + callback_weak = weakref.ref(callback) + graph_def = GraphDefinition() + graph_def.callback(callback) + first = graph_def.instantiate() + second = graph_def.instantiate() + + del callback, graph_def + gc.collect() + assert callback_weak() is not None + + del first + gc.collect() + assert callback_weak() is not None + + stream = Device().create_stream() + second.launch(stream) + stream.sync() + assert called == [1] + + del second + _wait_until(lambda: callback_weak() is None) + + +@pytest.mark.agent_authored(model="gpt-5.6") +def test_new_node_attachment_does_not_replace_deleted_node_exec_attachment( + init_cuda, +): + """A new source node cannot replace an older executable's attachment.""" + called = [] + + def old_callback(): + called.append("old") + + def new_callback(): + called.append("new") + + old_weak = weakref.ref(old_callback) + new_weak = weakref.ref(new_callback) + graph_def = GraphDefinition() + old_node = graph_def.callback(old_callback) + old_graph = graph_def.instantiate() + + old_node.destroy() + new_node = graph_def.callback(new_callback) + new_graph = graph_def.instantiate() + + del old_callback, new_callback, old_node, new_node, graph_def + gc.collect() + assert old_weak() is not None + assert new_weak() is not None + + stream = Device().create_stream() + old_graph.launch(stream) + new_graph.launch(stream) + stream.sync() + assert called == ["old", "new"] + + del old_graph + _wait_until(lambda: old_weak() is None) + assert new_weak() is not None + + del new_graph + _wait_until(lambda: new_weak() is None) + + def test_cfunc_callback_survives_del(init_cuda): """ctypes CFUNCTYPE wrapper is kept alive by the graph after Python ref is dropped.""" import ctypes @@ -599,6 +1017,8 @@ def test_kernel_survives_graph_clone_and_execution(init_cuda): stream = dev.create_stream() handle_return(driver.cuGraphLaunch(graph_exec, driver.CUstream(int(stream.handle)))) stream.sync() + handle_return(driver.cuGraphExecDestroy(graph_exec)) + handle_return(driver.cuGraphDestroy(cloned_cu_graph)) # ============================================================================= @@ -732,6 +1152,8 @@ def test_kernel_args_survive_graph_clone(init_cuda): out = (ctypes.c_int * 1)(0) handle_return(driver.cuMemcpyDtoH(out, dptr, ctypes.sizeof(ctypes.c_int))) assert out[0] == 1 + handle_return(driver.cuGraphExecDestroy(graph_exec)) + handle_return(driver.cuGraphDestroy(cloned_cu_graph)) # ============================================================================= @@ -855,8 +1277,8 @@ def test_memcpy_buffer_allocations_released_after_graph_destroyed(init_cuda): Each operand's device-pointer handle is observed via a weak handle (see ``cuda.core._utils._weak_handles``), so release is checked at the reference-count level rather than through a driver side effect. With both - Buffer wrappers closed, the graph's slots are the only remaining owners; - destroying the graph releases them and the weak handles expire. + Buffer wrappers closed, the node attachment is the only remaining owner; + destroying the graph releases it and the weak handles expire. """ from cuda.core._utils._weak_handles import weak_handle @@ -871,12 +1293,12 @@ def test_memcpy_buffer_allocations_released_after_graph_destroyed(init_cuda): g.memcpy(dst, src, 4) # Observe the allocations, then drop the wrappers' strong references; the - # graph slots remain the sole owners. + # graph attachment remains the sole owner. src_weak = weak_handle(src) dst_weak = weak_handle(dst) src.close() dst.close() - assert src_weak and dst_weak # graph slots still retain both allocations + assert src_weak and dst_weak # graph attachment still retains both allocations del g _wait_until(lambda: not src_weak and not dst_weak) @@ -911,6 +1333,8 @@ def test_memcpy_buffers_survive_graph_clone(init_cuda): out = (ctypes.c_uint8 * 4)(0) handle_return(driver.cuMemcpyDtoH(out, dst_dptr, 4)) assert list(out) == [0xCD] * 4 + handle_return(driver.cuGraphExecDestroy(graph_exec)) + handle_return(driver.cuGraphDestroy(cloned_cu_graph)) # ============================================================================= @@ -944,14 +1368,14 @@ def test_memset_raw_ptr_with_dst_owner(init_cuda): @pytest.mark.agent_authored(model="claude-opus-4.8") -def test_slot_owners_released_after_graph_destroyed(init_cuda): - """Destroying the graph releases every owner held in its slot table. +def test_attachment_owners_released_after_graph_destroyed(init_cuda): + """Destroying the graph releases every per-node attachment owner. Raw-pointer operands with explicit sentinel owners make release observable - in pure Python: the slot table holds a strong Python reference to each owner - (via ``make_opaque_py``), and graph destruction frees the table -- dropping - those references. This exercises the same teardown that releases a Buffer - operand's device-pointer handle (slot 0 for ``dst``, slot 1 for ``src``). + in pure Python: the node's CUDA user object holds strong Python references + (via ``make_opaque_py``), and graph destruction drops those references. + This exercises the same teardown that releases Buffer device-pointer + handles retained for ``dst`` and ``src``. """ _skip_if_no_mempool() dev = Device() @@ -979,6 +1403,35 @@ def test_slot_owners_released_after_graph_destroyed(init_cuda): buf.close() +@pytest.mark.agent_authored(model="gpt-5.6") +def test_deleted_memcpy_bundle_survives_in_existing_exec(init_cuda): + """Both memcpy owners remain together until an existing exec releases them.""" + _skip_if_no_mempool() + dev = Device() + mr = DeviceMemoryResource(dev) + buf = mr.allocate(8, stream=dev.default_stream) + dev.default_stream.sync() + dptr = int(buf.handle) + + dst_owner = _Sentinel() + src_owner = _Sentinel() + dst_weak = weakref.ref(dst_owner) + src_weak = weakref.ref(src_owner) + graph_def = GraphDefinition() + node = graph_def.memcpy(dptr, dptr + 4, 4, dst_owner=dst_owner, src_owner=src_owner) + graph = graph_def.instantiate() + + del dst_owner, src_owner + node.destroy() + del node, graph_def + gc.collect() + assert dst_weak() is not None and src_weak() is not None + + del graph + _wait_until(lambda: dst_weak() is None and src_weak() is None) + buf.close() + + @pytest.mark.agent_authored(model="claude-opus-4.8") def test_memcpy_raw_ptrs_with_owners(init_cuda): """Raw src/dst plus Buffer owners retain allocations after close.""" diff --git a/cuda_core/tests/graph/test_graph_definition_mutation.py b/cuda_core/tests/graph/test_graph_definition_mutation.py index 1db1089f82..e118c88c8d 100644 --- a/cuda_core/tests/graph/test_graph_definition_mutation.py +++ b/cuda_core/tests/graph/test_graph_definition_mutation.py @@ -3,6 +3,9 @@ """Tests for mutating a graph definition (edge changes, node removal).""" +import gc +import weakref + import numpy as np import pytest from helpers.collection_interface_testers import assert_mutable_set_interface @@ -307,6 +310,42 @@ def test_destroyed_node(init_cuda): assert not b.is_valid +@pytest.mark.agent_authored(model="gpt-5.6") +def test_failed_destroy_preserves_node_and_attachments(init_cuda): + """A graph-memory restriction must not invalidate a failed node deletion.""" + if not Device(0).properties.memory_pools_supported: + pytest.skip("graph memory nodes require memory pool support") + + called = [False] + + def callback(): + called[0] = True + + callback_weak = weakref.ref(callback) + graph_def = GraphDefinition() + alloc = graph_def.allocate(4) + node = alloc.callback(callback) + edge = (alloc, node) + + del callback + gc.collect() + assert callback_weak() is not None + + with pytest.raises(CUDAError): + node.destroy() + + assert node.is_valid + assert node in graph_def.nodes() + assert edge in graph_def.edges() + assert callback_weak() is not None + + graph = graph_def.instantiate() + stream = Device().create_stream() + graph.launch(stream) + stream.sync() + assert called[0] + + def test_add_wrong_type(init_cuda): """Adding a non-GraphNode raises TypeError.""" g = GraphDefinition() diff --git a/cuda_core/tests/graph/test_graph_update.py b/cuda_core/tests/graph/test_graph_update.py index 206fabcd58..448cebcf26 100644 --- a/cuda_core/tests/graph/test_graph_update.py +++ b/cuda_core/tests/graph/test_graph_update.py @@ -3,6 +3,10 @@ """Tests for whole-graph update (Graph.update).""" +import gc +import time +import weakref + import numpy as np import pytest from helpers.graph_kernels import compile_common_kernels, compile_conditional_kernels @@ -13,6 +17,15 @@ from cuda.core.graph import GraphDefinition +def _wait_until(predicate, timeout=5.0): + deadline = time.monotonic() + timeout + while not predicate(): + if time.monotonic() >= deadline: + raise AssertionError(f"condition not satisfied within {timeout}s") + gc.collect() + time.sleep(0.02) + + @pytest.mark.parametrize("builder", ["GraphBuilder", "GraphDefinition"]) @requires_module(np, "2.1") def test_graph_update_kernel_args(init_cuda, builder): @@ -150,6 +163,108 @@ def build_graph(condition_value): b.close() +@pytest.mark.agent_authored(model="gpt-5.6") +def test_graph_update_replaces_attachment_user_objects(init_cuda): + """A successful update adopts new owners and safely retires old ones.""" + called = [] + + def old_callback(): + called.append("old") + + def new_callback(): + called.append("new") + + old_weak = weakref.ref(old_callback) + new_weak = weakref.ref(new_callback) + old_def = GraphDefinition() + old_def.callback(old_callback) + graph = old_def.instantiate() + new_def = GraphDefinition() + new_def.callback(new_callback) + + del old_callback, new_callback + graph.update(new_def) + del old_def, new_def + gc.collect() + + _wait_until(lambda: old_weak() is None) + assert new_weak() is not None + + stream = Device().create_stream() + graph.launch(stream) + stream.sync() + assert called == ["new"] + + +@pytest.mark.agent_authored(model="gpt-5.6") +def test_failed_graph_update_does_not_adopt_attachments(init_cuda): + """A rejected source graph keeps ownership separate from the exec.""" + called = [] + + def active_callback(): + called.append("active") + + def rejected_callback(): + called.append("rejected") + + rejected_weak = weakref.ref(rejected_callback) + active_def = GraphDefinition() + active_def.callback(active_callback) + graph = active_def.instantiate() + + rejected_def = GraphDefinition() + rejected_def.callback(rejected_callback) + rejected_def.empty() # Force a topology mismatch. + del rejected_callback + + with pytest.raises(CUDAError): + graph.update(rejected_def) + + del rejected_def + _wait_until(lambda: rejected_weak() is None) + + stream = Device().create_stream() + graph.launch(stream) + stream.sync() + assert called == ["active"] + + +@pytest.mark.agent_authored(model="gpt-5.6") +def test_failed_graph_update_preserves_active_attachment_ownership(init_cuda): + """A rejected update leaves the executable's active owner unchanged.""" + called = [] + + def active_callback(): + called.append("active") + + def rejected_callback(): + called.append("rejected") + + active_weak = weakref.ref(active_callback) + active_def = GraphDefinition() + active_def.callback(active_callback) + graph = active_def.instantiate() + + rejected_def = GraphDefinition() + rejected_def.callback(rejected_callback) + rejected_def.empty() # Force a topology mismatch. + + with pytest.raises(CUDAError): + graph.update(rejected_def) + + del active_callback, active_def, rejected_callback, rejected_def + gc.collect() + assert active_weak() is not None + + stream = Device().create_stream() + graph.launch(stream) + stream.sync() + assert called == ["active"] + + del graph + _wait_until(lambda: active_weak() is None) + + # ============================================================================= # Error cases # =============================================================================