cuda.core: graph slot table for node attachment lifetimes (#2280)

* cuda.core: add GraphBuilder.graph_definition property

Completes step 3 of #1330 by exposing the captured graph as an explicit
`GraphDefinition` view that shares ownership of the underlying `CUgraph`.
The handle-layer plumbing landed in PR #2008; this commit wires up the
user-facing surface and locks in the state-guard rules.

State semantics:

- PRIMARY builder: only valid after `end_building()`. Before
  `begin_building()` no graph exists; during capture the driver is the
  sole writer, so explicit access is unsafe.
- CONDITIONAL_BODY builder: valid both before `begin_building()` (the
  body graph is allocated at conditional-node creation time) and after
  `end_building()`. This enables a hybrid flow where a conditional body
  is populated entirely via the explicit API, with no capture at all.
- FORKED builder: never valid. Forked builders share the primary's
  graph; access through the primary instead.

Tests cover the happy path, both hybrid flows on conditional bodies
(populate-via-explicit-API and capture-then-augment), the three error
states (forked, capturing, primary pre-capture), and the
shared-ownership guarantee (the `GraphDefinition` survives the
builder's `close()`).

Co-authored-by: Cursor <cursoragent@cursor.com>

* cuda.core: add graph slot table infrastructure (phase 1)

Introduce OpaqueHandle and a per-graph slot table retained on the CUgraph
as a user object, preparing to replace ad-hoc per-resource user objects when
wiring graph node attachments in a follow-up change.

* cuda.core: wire graph node attachments to the slot table (phase 2)

Replace the per-resource CUDA user objects attached at each graph node
with the per-graph slot table from phase 1. Kernel, event-record,
event-wait, and host-callback nodes now store their owning handles in
node slots via graph_set_slot. Stream-captured callbacks map the
just-captured host node from cuStreamGetCaptureInfo and use the same
path; forked builders share the primary's graph handle so their
attachments reach the same table.

Refine the phase 1 surface to support this: the slot table is created
lazily on first attachment, so conditional-branch bodies (ref handles)
get one too, and graph_set_slot returns CUresult for HANDLE_RETURN-style
error checking. Removes _attach_user_object and the per-type heap-copy
deleters.

* cuda.core: rename graph host-callback module and retain kernel args

Rename graph/_utils to graph/_host_callback now that it holds only
host-callback machinery (the trampoline, _is_py_host_trampoline, and
_resolve_host_callback), matching the concept-named files around it, and
update the three cimport sites. Add _attach_host_callback_owners to share
the "callback -> slot 0, user_data -> slot 1" attachment between the eager
(GN_callback) and capture (add_callback) paths. Guard a zero-length
user_data copy against malloc(0) and hoist the per-call ctypes import.

Attach the kernel-argument tuple to the kernel node's slot 1 so the Python
objects backing the arguments -- notably device Buffers -- outlive the
graph. The driver copies argument values into the node at add time but does
not keep the referenced device memory alive, so without this a kernel node
could be left with a stale device pointer. This is the slot-table port of
the user-object fix from #2041 (currently only on main).

* cuda.core: accept Buffer in graph memcpy/memset and retain operands

GraphNode.memcpy/memset (and the GraphDefinition pass-throughs) now accept a
Buffer or a raw int for each address. A new _resolve_ptr helper reads the
device pointer from a Buffer and returns it as an owner; a raw int casts
through with no owner. GN_memcpy attaches a Buffer dst to slot 0 and src to
slot 1, and GN_memset attaches dst to slot 0, so buffers passed by value
outlive the graph. Raw ints behave exactly as before (caller owns the
lifetime), so this is backward compatible.

Document the stream-capture lifetime contract on GraphBuilder: operations
recorded during capture reference caller-owned memory and are not retained,
unlike explicit GraphDefinition construction. Host callbacks are the one
exception, retained on both the capture and explicit paths.

* cuda.core: add slot-table lifetime tests for Buffer memcpy/memset and capture callbacks

Cover GraphDefinition memset/memcpy with Buffer operands (including clone),
and GraphBuilder capture host callbacks retained after dropping Python refs.

* cuda.core: add explicit dst/src_owner for graph memcpy/memset

Keyword-only *_owner args retain arbitrary objects for raw pointer
operands; Buffer+owner combinations are rejected. Strengthen owner tests
with weakref retention checks and add src_owner rejection test.

* cuda.core: retain device allocations in graph memcpy/memset slots

Store DevicePtrHandle in slot table instead of Buffer wrappers so
reset/close cannot release memory while a graph still references it.
Add test-only weak_handle() for deterministic allocation lifetime
checks and extend graph lifetime tests accordingly.

* cuda.core: keep memset height/pitch positional; mark new graph tests

Address PR #2280 review feedback:

- Move the keyword-only "*" marker in GraphNode.memset and
  GraphDefinition.memset to after height/pitch, so pre-existing positional
  calls memset(dst, value, width, height, pitch) keep working. The new
  dst_owner argument remains keyword-only. This avoids a public API break
  across 1.x. memcpy is unchanged (its dst_owner/src_owner args are new,
  so the existing "*" placement is non-breaking).
- Add @pytest.mark.agent_authored markers to the new graph tests in
  test_graph_builder.py and test_graph_definition_lifetime.py.

* cuda.core: roll back graph node when owner-slot attachment fails

A node added via cuGraphAdd*Node is committed to the graph before its
owner slots are attached. If graph_set_slot fails (e.g. the driver lacks
cuUserObjectCreate, or a transient error), the node would remain in the
graph referencing Python-owned memory with nothing keeping it alive,
risking a later launch dereferencing freed memory.

Guard the slot-attachment at each explicit-add site (kernel, memset,
memcpy, event record/wait, host callback) with a try/except that destroys
the node (best effort) and re-raises. The capture-path callback in
_graph_builder is intentionally left alone: its node is created by
cuLaunchHostFunc during active capture, where destroying a capture
dependency would corrupt capture state.

* cuda.core: use if/elif chain in graph_definition guard

Convert the sequential guard checks in GraphBuilder.graph_definition to an
if/elif chain (splitting the final compound condition into a nested if).
Behavior is unchanged since each leading branch raises; the chain lets
Cython generate tighter branch code. Addresses a review nit on PR #2280.

* cuda.core: make graph_set_slot a no-op for null owners

Centralize null-owner handling in graph_set_slot: a null OpaqueHandle now
returns CUDA_SUCCESS without forcing slot-table (and user-object) creation.
This resolves the reviewer question about the asymmetric per-call-site NULL
checks -- optional owners are uniformly safe at the source, so callers no
longer need to guard them. Update the header doc accordingly.

---------

Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Leo Fang <leof@nvidia.com>

Author: 3 people

cuda_core/cuda/core/_cpp/resource_handles.cpp

@@ -6,8 +6,11 @@
 
 #include "resource_handles.hpp"
 #include <cuda.h>
+#include <array>
 #include <cstdint>
+#include <cstdlib>
 #include <cstring>
+#include <map>
 #include <mutex>
 #include <stdexcept>
 #include <unordered_map>
@@ -70,6 +73,9 @@ decltype(&cuLibraryGetKernel) p_cuLibraryGetKernel = nullptr;
 // Graph
 decltype(&cuGraphDestroy) p_cuGraphDestroy = nullptr;
 decltype(&cuGraphExecDestroy) p_cuGraphExecDestroy = nullptr;
+decltype(&cuUserObjectCreate) p_cuUserObjectCreate = nullptr;
+decltype(&cuUserObjectRelease) p_cuUserObjectRelease = nullptr;
+decltype(&cuGraphRetainUserObject) p_cuGraphRetainUserObject = nullptr;
 
 // Linker
 decltype(&cuLinkDestroy) p_cuLinkDestroy = nullptr;
@@ -1114,12 +1120,92 @@ 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<OpaqueHandle, SLOTS_PER_NODE>;
+using GraphSlotTable = std::map<CUgraphNode, NodeSlots>;
+
+// 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
+// deleter here. The Python deleter follows the owner-release pattern used by
+// the stream/deviceptr handles above.
+void py_deleter(const void* p) noexcept {
+    GILAcquireGuard gil;
+    if (gil.acquired()) {
+        Py_DECREF(const_cast<PyObject*>(static_cast<const PyObject*>(p)));
+    }
+}
+
+void free_deleter(const void* p) noexcept {
+    std::free(const_cast<void*>(p));
+}
+
+void destroy_graph_slot_table(void* table) noexcept {
+    delete static_cast<GraphSlotTable*>(table);
+}
+
 struct GraphBox {
     CUgraph resource;
-    GraphHandle h_parent;  // Keeps parent alive for child/branch graphs
+    GraphHandle h_parent;                  // Keeps parent alive for child/branch graphs
+    mutable GraphSlotTable* slot_table = nullptr;  // Lazily created; owned by the graph's user object
 };
+
+const GraphBox* get_box(const GraphHandle& h) {
+    const CUgraph* p = h.get();
+    return reinterpret_cast<const GraphBox*>(
+        reinterpret_cast<const char*>(p) - offsetof(GraphBox, resource)
+    );
+}
+
+// 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;
+    }
+    auto* table = new GraphSlotTable();
+    CUuserObject user_obj = nullptr;
+    {
+        GILReleaseGuard gil;
+        if (p_cuUserObjectCreate(&user_obj, table,
+                                 reinterpret_cast<CUhostFn>(destroy_graph_slot_table),
+                                 1, CU_USER_OBJECT_NO_DESTRUCTOR_SYNC) != CUDA_SUCCESS) {
+            delete table;  // no user object created; nothing else owns the table
+            return nullptr;
+        }
+        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;
+        }
+    }
+    box->slot_table = table;  // non-owning cache; the user object owns it
+    return table;
+}
+
 }  // namespace
 
+OpaqueHandle make_opaque_py(PyObject* obj) {
+    Py_INCREF(obj);
+    return OpaqueHandle(static_cast<const void*>(obj), py_deleter);
+}
+
+OpaqueHandle make_opaque_malloc(void* buf) {
+    return OpaqueHandle(static_cast<const void*>(buf), free_deleter);
+}
+
 GraphHandle create_graph_handle(CUgraph graph) {
     auto box = std::shared_ptr<const GraphBox>(
         new GraphBox{graph, {}},
@@ -1137,6 +1223,22 @@ GraphHandle create_graph_handle_ref(CUgraph graph, const GraphHandle& h_parent)
     return GraphHandle(box, &box->resource);
 }
 
+CUresult graph_set_slot(const GraphHandle& h_graph, CUgraphNode node,
+                        unsigned int slot, OpaqueHandle owner) {
+    if (!h_graph || slot >= SLOTS_PER_NODE) {
+        return CUDA_ERROR_INVALID_VALUE;
+    }
+    if (!owner) {
+        return CUDA_SUCCESS;  // nothing to retain; don't force table creation
+    }
+    GraphSlotTable* table = ensure_slot_table(get_box(h_graph));
+    if (!table) {
+        return CUDA_ERROR_NOT_SUPPORTED;
+    }
+    (*table)[node][slot] = std::move(owner);
+    return CUDA_SUCCESS;
+}
+
 // ============================================================================
 // Graph Exec Handles
 // ============================================================================

cuda_core/cuda/core/_cpp/resource_handles.hpp

@@ -109,6 +109,9 @@ extern decltype(&cuLibraryGetKernel) p_cuLibraryGetKernel;
 // Graph
 extern decltype(&cuGraphDestroy) p_cuGraphDestroy;
 extern decltype(&cuGraphExecDestroy) p_cuGraphExecDestroy;
+extern decltype(&cuUserObjectCreate) p_cuUserObjectCreate;
+extern decltype(&cuUserObjectRelease) p_cuUserObjectRelease;
+extern decltype(&cuGraphRetainUserObject) p_cuGraphRetainUserObject;
 
 // Linker
 extern decltype(&cuLinkDestroy) p_cuLinkDestroy;
@@ -466,6 +469,39 @@ GraphHandle create_graph_handle(CUgraph graph);
 // but h_parent will be prevented from destruction while this handle exists.
 GraphHandle create_graph_handle_ref(CUgraph graph, const GraphHandle& h_parent);
 
+// ============================================================================
+// Graph slot 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.
+// ============================================================================
+
+// Type-erased shared owner of an attached resource. Typed handles such as
+// EventHandle and KernelHandle convert to OpaqueHandle by assignment, reusing
+// their existing control block; the helpers below build OpaqueHandles for the
+// two cases that need a custom deleter.
+using OpaqueHandle = std::shared_ptr<const void>;
+
+// Build an OpaqueHandle from a Python object: increments its refcount now and
+// decrements it (under the GIL) on release. The caller must hold the GIL.
+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);
+
 // ============================================================================
 // Graph exec handle functions
 // ============================================================================

cuda_core/cuda/core/_resource_handles.pxd

@@ -57,6 +57,12 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core":
     ctypedef shared_ptr[const TexObjectValue] TexObjectHandle
     ctypedef shared_ptr[const SurfObjectValue] SurfObjectHandle
 
+    # Type-erased shared owner for resources attached to graph node slots.
+    # Typed handles above assign directly to an OpaqueHandle (shared control
+    # block); make_opaque_py / make_opaque_malloc cover the two cases needing a
+    # custom deleter.
+    ctypedef shared_ptr[const void] OpaqueHandle
+
     # 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
@@ -223,6 +229,13 @@ cdef LibraryHandle get_kernel_library(const KernelHandle& h) noexcept nogil
 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
 
+# Graph slot attachments
+cdef OpaqueHandle make_opaque_py(object obj) except+
+cdef OpaqueHandle make_opaque_malloc(void* buf) except+
+cdef cydriver.CUresult graph_set_slot(
+    const GraphHandle& h_graph, cydriver.CUgraphNode node,
+    unsigned int slot, OpaqueHandle owner) except+
+
 # Graph exec handles
 cdef GraphExecHandle create_graph_exec_handle(cydriver.CUgraphExec graph_exec) except+ nogil
 

cuda_core/cuda/core/_resource_handles.pyi

@@ -24,4 +24,5 @@ FileDescriptorHandle = shared_ptr
 OpaqueArrayHandle = shared_ptr
 MipmappedArrayHandle = shared_ptr
 TexObjectHandle = shared_ptr
-SurfObjectHandle = shared_ptr
+SurfObjectHandle = shared_ptr
+OpaqueHandle = shared_ptr

cuda_core/cuda/core/_resource_handles.pyx

@@ -151,6 +151,13 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core":
     GraphHandle create_graph_handle_ref "cuda_core::create_graph_handle_ref" (
         cydriver.CUgraph graph, const GraphHandle& h_parent) except+ nogil
 
+    # Graph slot 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" (
+        const GraphHandle& h_graph, cydriver.CUgraphNode node,
+        unsigned int slot, OpaqueHandle owner) except+
+
     # Graph exec handles
     GraphExecHandle create_graph_exec_handle "cuda_core::create_graph_exec_handle" (
         cydriver.CUgraphExec graph_exec) except+ nogil
@@ -304,6 +311,9 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core":
     # Graph
     void* p_cuGraphDestroy "reinterpret_cast<void*&>(cuda_core::p_cuGraphDestroy)"
     void* p_cuGraphExecDestroy "reinterpret_cast<void*&>(cuda_core::p_cuGraphExecDestroy)"
+    void* p_cuUserObjectCreate "reinterpret_cast<void*&>(cuda_core::p_cuUserObjectCreate)"
+    void* p_cuUserObjectRelease "reinterpret_cast<void*&>(cuda_core::p_cuUserObjectRelease)"
+    void* p_cuGraphRetainUserObject "reinterpret_cast<void*&>(cuda_core::p_cuGraphRetainUserObject)"
 
     # Linker
     void* p_cuLinkDestroy "reinterpret_cast<void*&>(cuda_core::p_cuLinkDestroy)"
@@ -364,6 +374,7 @@ 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_cuLinkDestroy
     global p_cuGraphicsUnmapResources, p_cuGraphicsUnregisterResource
     global p_cuDevSmResourceSplit
@@ -424,6 +435,9 @@ cdef void _init_driver_fn_pointers() noexcept:
     # Graph
     p_cuGraphDestroy = _get_driver_fn("cuGraphDestroy")
     p_cuGraphExecDestroy = _get_driver_fn("cuGraphExecDestroy")
+    p_cuUserObjectCreate = _get_driver_fn("cuUserObjectCreate")
+    p_cuUserObjectRelease = _get_driver_fn("cuUserObjectRelease")
+    p_cuGraphRetainUserObject = _get_driver_fn("cuGraphRetainUserObject")
 
     # Linker
     p_cuLinkDestroy = _get_driver_fn("cuLinkDestroy")

cuda_core/cuda/core/_utils/_weak_handles.pyi

@@ -0,0 +1,55 @@
+# This file was generated by stubgen-pyx v0.2.6 from cuda_core/cuda/core/_utils/_weak_handles.pyx
+
+"""Test-only weak handles for resource-handle lifetime checks.
+
+This module is **not** part of the public ``cuda.core`` API. It is built into
+the package (like other private ``_utils`` modules) purely so the test suite can
+observe, deterministically, when the strong references that keep a CUDA resource
+alive have all been released -- without relying on driver- or hardware-specific
+side effects (for example, whether freed device memory happens to remain
+readable).
+
+Every resource handle is owned by a C++ ``std::shared_ptr``. A **weak handle**
+is a non-owning ``std::weak_ptr`` observer of that control block: truthy while
+some strong owner remains, falsy once the last one is gone. Use :func:`weak_handle`
+to obtain a weak handle from a supported front-end object.
+
+To support another type, add a ``cdef _weak_from_<type>`` that reads its ``cdef``
+handle field (see ``*.pxd``), assigns to :ctype:`OpaqueHandle`, and extend the
+``isinstance`` chain in :func:`weak_handle`. Types whose slots hold arbitrary
+Python owners via ``make_opaque_py`` are not covered here -- use
+:class:`weakref.ref` on a weak-referenceable owner object in tests instead.
+"""
+from __future__ import annotations
+
+
+class WeakHandle:
+    """Non-owning weak handle for a resource's shared control block.
+
+    Truthy while some strong owner of the underlying resource handle remains,
+    falsy once the last strong reference is released. Obtain instances via
+    :func:`weak_handle` rather than constructing directly.
+    """
+
+    def __bool__(self):
+        ...
+
+    def expired(self):
+        """Return ``True`` once every strong owner of the handle is gone."""
+
+    def use_count(self):
+        """Number of strong owners currently sharing the handle."""
+
+def weak_handle(obj):
+    """Return a :class:`WeakHandle` observing the resource behind ``obj``.
+
+    Currently supports :class:`~cuda.core.Buffer` (device allocation handle).
+    See the module docstring for how to add more types.
+
+    Raises
+    ------
+    ValueError
+        If ``obj`` is a :class:`~cuda.core.Buffer` with no active allocation.
+    TypeError
+        If ``obj`` is not a supported type.
+    """