Extract shared graph instantiation helper

Move instantiation logic from GraphBuilder.complete() into a shared
_instantiate_graph() helper so GraphDef.instantiate() can reuse it
with full GraphCompleteOptions support and error handling.

Made-with: Cursor

Author: Andy-Jost

cuda_core/cuda/core/_graph/__init__.py

@@ -155,6 +155,44 @@ class GraphCompleteOptions:
     use_node_priority: bool = False
 
 
+def _instantiate_graph(h_graph, options: GraphCompleteOptions | None = None) -> Graph:
+    params = driver.CUDA_GRAPH_INSTANTIATE_PARAMS()
+    if options:
+        flags = 0
+        if options.auto_free_on_launch:
+            flags |= driver.CUgraphInstantiate_flags.CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH
+        if options.upload_stream:
+            flags |= driver.CUgraphInstantiate_flags.CUDA_GRAPH_INSTANTIATE_FLAG_UPLOAD
+            params.hUploadStream = options.upload_stream.handle
+        if options.device_launch:
+            flags |= driver.CUgraphInstantiate_flags.CUDA_GRAPH_INSTANTIATE_FLAG_DEVICE_LAUNCH
+        if options.use_node_priority:
+            flags |= driver.CUgraphInstantiate_flags.CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
+        params.flags = flags
+
+    graph = Graph._init(handle_return(driver.cuGraphInstantiateWithParams(h_graph, params)))
+    if params.result_out == driver.CUgraphInstantiateResult.CUDA_GRAPH_INSTANTIATE_ERROR:
+        raise RuntimeError(
+            "Instantiation failed for an unexpected reason which is described in the return value of the function."
+        )
+    elif params.result_out == driver.CUgraphInstantiateResult.CUDA_GRAPH_INSTANTIATE_INVALID_STRUCTURE:
+        raise RuntimeError("Instantiation failed due to invalid structure, such as cycles.")
+    elif params.result_out == driver.CUgraphInstantiateResult.CUDA_GRAPH_INSTANTIATE_NODE_OPERATION_NOT_SUPPORTED:
+        raise RuntimeError(
+            "Instantiation for device launch failed because the graph contained an unsupported operation."
+        )
+    elif params.result_out == driver.CUgraphInstantiateResult.CUDA_GRAPH_INSTANTIATE_MULTIPLE_CTXS_NOT_SUPPORTED:
+        raise RuntimeError("Instantiation for device launch failed due to the nodes belonging to different contexts.")
+    elif (
+        _py_major_minor >= (12, 8)
+        and params.result_out == driver.CUgraphInstantiateResult.CUDA_GRAPH_INSTANTIATE_CONDITIONAL_HANDLE_UNUSED
+    ):
+        raise RuntimeError("One or more conditional handles are not associated with conditional builders.")
+    elif params.result_out != driver.CUgraphInstantiateResult.CUDA_GRAPH_INSTANTIATE_SUCCESS:
+        raise RuntimeError(f"Graph instantiation failed with unexpected error code: {params.result_out}")
+    return graph
+
+
 class GraphBuilder:
     """Represents a graph under construction.
 
@@ -317,53 +355,7 @@ def complete(self, options: GraphCompleteOptions | None = None) -> Graph:
         if not self._building_ended:
             raise RuntimeError("Graph has not finished building.")
 
-        if (_driver_ver < 12000) or (_py_major_minor < (12, 0)):
-            flags = 0
-            if options:
-                if options.auto_free_on_launch:
-                    flags |= driver.CUgraphInstantiate_flags.CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH
-                if options.use_node_priority:
-                    flags |= driver.CUgraphInstantiate_flags.CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
-            return Graph._init(handle_return(driver.cuGraphInstantiateWithFlags(self._mnff.graph, flags)))
-
-        params = driver.CUDA_GRAPH_INSTANTIATE_PARAMS()
-        if options:
-            flags = 0
-            if options.auto_free_on_launch:
-                flags |= driver.CUgraphInstantiate_flags.CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH
-            if options.upload_stream:
-                flags |= driver.CUgraphInstantiate_flags.CUDA_GRAPH_INSTANTIATE_FLAG_UPLOAD
-                params.hUploadStream = options.upload_stream.handle
-            if options.device_launch:
-                flags |= driver.CUgraphInstantiate_flags.CUDA_GRAPH_INSTANTIATE_FLAG_DEVICE_LAUNCH
-            if options.use_node_priority:
-                flags |= driver.CUgraphInstantiate_flags.CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
-            params.flags = flags
-
-        graph = Graph._init(handle_return(driver.cuGraphInstantiateWithParams(self._mnff.graph, params)))
-        if params.result_out == driver.CUgraphInstantiateResult.CUDA_GRAPH_INSTANTIATE_ERROR:
-            # NOTE: Should never get here since the handle_return should have caught this case
-            raise RuntimeError(
-                "Instantiation failed for an unexpected reason which is described in the return value of the function."
-            )
-        elif params.result_out == driver.CUgraphInstantiateResult.CUDA_GRAPH_INSTANTIATE_INVALID_STRUCTURE:
-            raise RuntimeError("Instantiation failed due to invalid structure, such as cycles.")
-        elif params.result_out == driver.CUgraphInstantiateResult.CUDA_GRAPH_INSTANTIATE_NODE_OPERATION_NOT_SUPPORTED:
-            raise RuntimeError(
-                "Instantiation for device launch failed because the graph contained an unsupported operation."
-            )
-        elif params.result_out == driver.CUgraphInstantiateResult.CUDA_GRAPH_INSTANTIATE_MULTIPLE_CTXS_NOT_SUPPORTED:
-            raise RuntimeError(
-                "Instantiation for device launch failed due to the nodes belonging to different contexts."
-            )
-        elif (
-            _py_major_minor >= (12, 8)
-            and params.result_out == driver.CUgraphInstantiateResult.CUDA_GRAPH_INSTANTIATE_CONDITIONAL_HANDLE_UNUSED
-        ):
-            raise RuntimeError("One or more conditional handles are not associated with conditional builders.")
-        elif params.result_out != driver.CUgraphInstantiateResult.CUDA_GRAPH_INSTANTIATE_SUCCESS:
-            raise RuntimeError(f"Graph instantiation failed with unexpected error code: {params.result_out}")
-        return graph
+        return _instantiate_graph(self._mnff.graph, options)
 
     def debug_dot_print(self, path, options: GraphDebugPrintOptions | None = None):
         """Generates a DOT debug file for the graph builder.

cuda_core/cuda/core/_graph/_graphdef.pyx

@@ -457,20 +457,23 @@ cdef class GraphDef:
         """
         return self._entry.switch(condition, count)
 
-    def instantiate(self):
+    def instantiate(self, options=None):
         """Instantiate the graph definition into an executable Graph.
 
+        Parameters
+        ----------
+        options : :obj:`~_graph.GraphCompleteOptions`, optional
+            Customizable dataclass for graph instantiation options.
+
         Returns
         -------
         Graph
             An executable graph that can be launched on a stream.
         """
-        from cuda.core._graph import Graph
-        from cuda.core._utils.cuda_utils import handle_return
+        from cuda.core._graph import _instantiate_graph
 
-        graph_exec = handle_return(driver.cuGraphInstantiate(
-            driver.CUgraph(as_intptr(self._h_graph)), 0))
-        return Graph._init(graph_exec)
+        return _instantiate_graph(
+            driver.CUgraph(as_intptr(self._h_graph)), options)
 
     def debug_dot_print(self, path: str, options=None) -> None:
         """Write a GraphViz DOT representation of the graph to a file.