cuda.core: Cythonize GraphBuilder and Graph with handle-layer cleanup

cuda_core/cuda/core/_cpp/resource_handles.cpp

@@ -69,6 +69,7 @@ decltype(&cuLibraryGetKernel) p_cuLibraryGetKernel = nullptr;
 
 // Graph
 decltype(&cuGraphDestroy) p_cuGraphDestroy = nullptr;
+decltype(&cuGraphExecDestroy) p_cuGraphExecDestroy = nullptr;
 
 // Linker
 decltype(&cuLinkDestroy) p_cuLinkDestroy = nullptr;
@@ -1126,6 +1127,28 @@ GraphHandle create_graph_handle_ref(CUgraph graph, const GraphHandle& h_parent)
     return GraphHandle(box, &box->resource);
 }
 
+// ============================================================================
+// Graph Exec Handles
+// ============================================================================
+
+namespace {
+struct GraphExecBox {
+    CUgraphExec resource;
+};
+}  // namespace
+
+GraphExecHandle create_graph_exec_handle(CUgraphExec graph_exec) {
+    auto box = std::shared_ptr<const GraphExecBox>(
+        new GraphExecBox{graph_exec},
+        [](const GraphExecBox* b) {
+            GILReleaseGuard gil;
+            p_cuGraphExecDestroy(b->resource);
+            delete b;
+        }
+    );
+    return GraphExecHandle(box, &box->resource);
+}
+
 namespace {
 struct GraphNodeBox {
     mutable CUgraphNode resource;

cuda_core/cuda/core/_cpp/resource_handles.hpp

@@ -100,6 +100,7 @@ extern decltype(&cuLibraryGetKernel) p_cuLibraryGetKernel;
 
 // Graph
 extern decltype(&cuGraphDestroy) p_cuGraphDestroy;
+extern decltype(&cuGraphExecDestroy) p_cuGraphExecDestroy;
 
 // Linker
 extern decltype(&cuLinkDestroy) p_cuLinkDestroy;
@@ -164,6 +165,7 @@ using MemoryPoolHandle = std::shared_ptr<const CUmemoryPool>;
 using LibraryHandle = std::shared_ptr<const CUlibrary>;
 using KernelHandle = std::shared_ptr<const CUkernel>;
 using GraphHandle = std::shared_ptr<const CUgraph>;
+using GraphExecHandle = std::shared_ptr<const CUgraphExec>;
 using GraphNodeHandle = std::shared_ptr<const CUgraphNode>;
 using GraphicsResourceHandle = std::shared_ptr<const CUgraphicsResource>;
 using NvrtcProgramHandle = std::shared_ptr<const nvrtcProgram>;
@@ -441,6 +443,14 @@ 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 exec handle functions
+// ============================================================================
+
+// Wrap an externally-created CUgraphExec with RAII cleanup.
+// When the last reference is released, cuGraphExecDestroy is called automatically.
+GraphExecHandle create_graph_exec_handle(CUgraphExec graph_exec);
+
 // ============================================================================
 // Graph node handle functions
 // ============================================================================
@@ -571,6 +581,10 @@ inline CUgraph as_cu(const GraphHandle& h) noexcept {
     return h ? *h : nullptr;
 }
 
+inline CUgraphExec as_cu(const GraphExecHandle& h) noexcept {
+    return h ? *h : nullptr;
+}
+
 inline CUgraphNode as_cu(const GraphNodeHandle& h) noexcept {
     return h ? *h : nullptr;
 }
@@ -633,6 +647,10 @@ inline std::intptr_t as_intptr(const GraphHandle& h) noexcept {
     return reinterpret_cast<std::intptr_t>(as_cu(h));
 }
 
+inline std::intptr_t as_intptr(const GraphExecHandle& h) noexcept {
+    return reinterpret_cast<std::intptr_t>(as_cu(h));
+}
+
 inline std::intptr_t as_intptr(const GraphNodeHandle& h) noexcept {
     return reinterpret_cast<std::intptr_t>(as_cu(h));
 }
@@ -743,6 +761,10 @@ inline PyObject* as_py(const GraphHandle& h) noexcept {
     return detail::make_py("cuda.bindings.driver", "CUgraph", as_intptr(h));
 }
 
+inline PyObject* as_py(const GraphExecHandle& h) noexcept {
+    return detail::make_py("cuda.bindings.driver", "CUgraphExec", as_intptr(h));
+}
+
 inline PyObject* as_py(const GraphNodeHandle& h) noexcept {
     if (!as_intptr(h)) {
         Py_RETURN_NONE;

cuda_core/cuda/core/_device.pyx

@@ -1454,7 +1454,7 @@ class Device:
         from cuda.core.graph._graph_builder import GraphBuilder
 
         self._check_context_initialized()
-        return GraphBuilder._init(stream=self.create_stream(), is_stream_owner=True)
+        return GraphBuilder._init(self.create_stream())
 
 
 cdef inline int Device_ensure_cuda_initialized() except? -1:

cuda_core/cuda/core/_resource_handles.pxd

@@ -28,6 +28,7 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core":
     ctypedef shared_ptr[const cydriver.CUlibrary] LibraryHandle
     ctypedef shared_ptr[const cydriver.CUkernel] KernelHandle
     ctypedef shared_ptr[const cydriver.CUgraph] GraphHandle
+    ctypedef shared_ptr[const cydriver.CUgraphExec] GraphExecHandle
     ctypedef shared_ptr[const cydriver.CUgraphNode] GraphNodeHandle
     ctypedef shared_ptr[const cydriver.CUgraphicsResource] GraphicsResourceHandle
     ctypedef shared_ptr[const cynvrtc.nvrtcProgram] NvrtcProgramHandle
@@ -54,6 +55,7 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core":
     cydriver.CUlibrary as_cu(LibraryHandle h) noexcept nogil
     cydriver.CUkernel as_cu(KernelHandle h) noexcept nogil
     cydriver.CUgraph as_cu(GraphHandle h) noexcept nogil
+    cydriver.CUgraphExec as_cu(GraphExecHandle h) noexcept nogil
     cydriver.CUgraphNode as_cu(GraphNodeHandle h) noexcept nogil
     cydriver.CUgraphicsResource as_cu(GraphicsResourceHandle h) noexcept nogil
     cynvrtc.nvrtcProgram as_cu(NvrtcProgramHandle h) noexcept nogil
@@ -71,6 +73,7 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core":
     intptr_t as_intptr(LibraryHandle h) noexcept nogil
     intptr_t as_intptr(KernelHandle h) noexcept nogil
     intptr_t as_intptr(GraphHandle h) noexcept nogil
+    intptr_t as_intptr(GraphExecHandle h) noexcept nogil
     intptr_t as_intptr(GraphNodeHandle h) noexcept nogil
     intptr_t as_intptr(GraphicsResourceHandle h) noexcept nogil
     intptr_t as_intptr(NvrtcProgramHandle h) noexcept nogil
@@ -89,6 +92,7 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core":
     object as_py(LibraryHandle h)
     object as_py(KernelHandle h)
     object as_py(GraphHandle h)
+    object as_py(GraphExecHandle h)
     object as_py(GraphNodeHandle h)
     object as_py(GraphicsResourceHandle h)
     object as_py(NvrtcProgramHandle h)
@@ -195,6 +199,9 @@ 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 exec handles
+cdef GraphExecHandle create_graph_exec_handle(cydriver.CUgraphExec graph_exec) except+ nogil
+
 # Graph node handles
 cdef GraphNodeHandle create_graph_node_handle(cydriver.CUgraphNode node, const GraphHandle& h_graph) except+ nogil
 cdef GraphHandle graph_node_get_graph(const GraphNodeHandle& h) noexcept nogil

cuda_core/cuda/core/_resource_handles.pyi

@@ -13,6 +13,7 @@ DevicePtrHandle = shared_ptr
 LibraryHandle = shared_ptr
 KernelHandle = shared_ptr
 GraphHandle = shared_ptr
+GraphExecHandle = shared_ptr
 GraphNodeHandle = shared_ptr
 GraphicsResourceHandle = shared_ptr
 NvrtcProgramHandle = shared_ptr

cuda_core/cuda/core/_resource_handles.pyx

@@ -151,6 +151,10 @@ 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 exec handles
+    GraphExecHandle create_graph_exec_handle "cuda_core::create_graph_exec_handle" (
+        cydriver.CUgraphExec graph_exec) except+ nogil
+
     # Graph node handles
     GraphNodeHandle create_graph_node_handle "cuda_core::create_graph_node_handle" (
         cydriver.CUgraphNode node, const GraphHandle& h_graph) except+ nogil
@@ -276,6 +280,7 @@ 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)"
 
     # Linker
     void* p_cuLinkDestroy "reinterpret_cast<void*&>(cuda_core::p_cuLinkDestroy)"
@@ -324,7 +329,7 @@ cdef void _init_driver_fn_pointers() noexcept:
     global p_cuMemFreeAsync, p_cuMemFree, p_cuMemFreeHost
     global p_cuMemPoolImportPointer
     global p_cuLibraryLoadFromFile, p_cuLibraryLoadData, p_cuLibraryUnload, p_cuLibraryGetKernel
-    global p_cuGraphDestroy
+    global p_cuGraphDestroy, p_cuGraphExecDestroy
     global p_cuLinkDestroy
     global p_cuGraphicsUnmapResources, p_cuGraphicsUnregisterResource
     global p_cuDevSmResourceSplit
@@ -380,6 +385,7 @@ cdef void _init_driver_fn_pointers() noexcept:
 
     # Graph
     p_cuGraphDestroy = _get_driver_fn("cuGraphDestroy")
+    p_cuGraphExecDestroy = _get_driver_fn("cuGraphExecDestroy")
 
     # Linker
     p_cuLinkDestroy = _get_driver_fn("cuLinkDestroy")

cuda_core/cuda/core/_stream.pyx

@@ -414,7 +414,7 @@ cdef class Stream:
         """
         from cuda.core.graph._graph_builder import GraphBuilder
 
-        return GraphBuilder._init(stream=self, is_stream_owner=False)
+        return GraphBuilder._init(self)
 
 
 LEGACY_DEFAULT_STREAM: Stream = Stream._legacy_default()

cuda_core/cuda/core/graph/_graph_builder.pxd

@@ -0,0 +1,27 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from cuda.bindings cimport cydriver
+
+from cuda.core._resource_handles cimport GraphExecHandle, GraphHandle, StreamHandle
+from cuda.core._stream cimport Stream
+
+
+cdef class GraphBuilder:
+    cdef:
+        GraphHandle _h_graph
+        StreamHandle _h_stream
+        int _kind
+        int _state
+        Stream _stream  # cached to avoid reconstruction from _h_stream handle
+        object __weakref__
+
+
+cdef class Graph:
+    cdef:
+        GraphExecHandle _h_graph_exec
+        object __weakref__
+
+    @staticmethod
+    cdef Graph _init(cydriver.CUgraphExec graph_exec)

cuda_core/cuda/core/graph/_graph_builder.pyi

@@ -8,6 +8,8 @@ from cuda.core._stream import Stream
 from cuda.core._utils.cuda_utils import driver
 from cuda.core.graph._graph_definition import GraphCondition, GraphDefinition
 
+_BuilderKind = int
+_CaptureState = int
 
 @dataclass
 class GraphDebugPrintOptions:
@@ -106,23 +108,19 @@ class GraphBuilder:
 
     """
 
-    class _MembersNeededForFinalize:
-        __slots__ = ('conditional_graph', 'graph', 'is_join_required', 'is_stream_owner', 'stream')
-
-        def __init__(self, graph_builder_obj: GraphBuilder, stream_obj: Stream | None, is_stream_owner: bool, conditional_graph, is_join_required: bool) -> None:
-            ...
-
-        def close(self) -> None:
-            ...
-    __slots__ = ('__weakref__', '_building_ended', '_mnff')
+    def __init__(self):
+        ...
 
-    def __init__(self) -> None:
+    def __dealloc__(self):
         ...
 
-    @classmethod
-    def _init(cls, stream: Stream | None, is_stream_owner: bool, conditional_graph: object=None, is_join_required: bool=False) -> GraphBuilder:
+    @staticmethod
+    def _init(stream: Stream):
         ...
 
+    def close(self):
+        """Destroy the graph builder."""
+
     @property
     def stream(self) -> Stream:
         """Returns the stream associated with the graph builder."""
@@ -155,7 +153,7 @@ class GraphBuilder:
     def end_building(self) -> GraphBuilder:
         """Ends the building process."""
 
-    def complete(self, options: GraphCompleteOptions | None=None) -> 'Graph':
+    def complete(self, options: GraphCompleteOptions | None=None) -> Graph:
         """Completes the graph builder and returns the built :obj:`~graph.Graph` object.
 
         Parameters
@@ -245,9 +243,6 @@ class GraphBuilder:
             A condition variable for controlling conditional execution.
         """
 
-    def _cond_with_params(self, node_params: object) -> tuple[GraphBuilder, ...]:
-        ...
-
     def if_then(self, condition: GraphCondition) -> GraphBuilder:
         """Adds an if condition branch and returns a new graph builder for it.
 
@@ -335,15 +330,7 @@ class GraphBuilder:
 
         """
 
-    def close(self) -> None:
-        """Destroy the graph builder.
-
-        Closes the associated stream if we own it. Borrowed stream
-        object will instead have their references released.
-
-        """
-
-    def embed(self, child: GraphBuilder) -> None:
+    def embed(self, child: GraphBuilder):
         """Embed a previously-built :obj:`~graph.GraphBuilder` as a child node.
 
         Parameters
@@ -392,21 +379,7 @@ class Graph:
 
     """
 
-    class _MembersNeededForFinalize:
-        __slots__ = 'graph'
-
-        def __init__(self, graph_obj: Graph, graph: driver.CUgraphExec) -> None:
-            ...
-
-        def close(self) -> None:
-            ...
-    __slots__ = ('__weakref__', '_mnff')
-
-    def __init__(self) -> None:
-        ...
-
-    @classmethod
-    def _init(cls, graph: driver.CUgraphExec) -> Graph:
+    def __init__(self):
         ...
 
     def close(self) -> None:
@@ -457,5 +430,5 @@ class Graph:
         """
 __all__ = ['Graph', 'GraphBuilder', 'GraphCompleteOptions', 'GraphDebugPrintOptions']
 
-def _instantiate_graph(h_graph, options: GraphCompleteOptions | None=None) -> 'Graph':
+def _instantiate_graph(h_graph, options: GraphCompleteOptions | None=None) -> Graph:
     ...