launch_graph -> gb.add_child

Author: leofang

cuda_core/cuda/core/experimental/__init__.py

@@ -10,7 +10,6 @@
     GraphBuilder,
     GraphCompleteOptions,
     GraphDebugPrintOptions,
-    launch_graph,
 )
 from cuda.core.experimental._launch_config import LaunchConfig
 from cuda.core.experimental._launcher import launch

cuda_core/cuda/core/experimental/_graph.py

@@ -654,6 +654,45 @@ def close(self):
         """
         self._mnff.close()
 
+    def add_child(self, child_graph: GraphBuilder):
+        """Adds the child :obj:`~_graph.GraphBuilder` builder into self.
+
+        The child graph builder will be added as a child node to the parent graph builder.
+
+        Parameters
+        ----------
+        child_graph : :obj:`~_graph.GraphBuilder`
+            The child graph builder. Must have finished building.
+        """
+        if (_driver_ver < 12000) or (_py_major_minor < (12, 0)):
+            raise NotImplementedError(
+                f"Launching child graphs is not implemented for versions older than CUDA 12."
+                f"Found driver version is {_driver_ver} and binding version is {_py_major_minor}"
+            )
+
+        if not child_graph._building_ended:
+            raise ValueError("Child graph has not finished building.")
+
+        if not self.is_building:
+            raise ValueError("Parent graph is being built.")
+
+        stream_handle = self._mnff.stream.handle
+        _, _, graph_out, dependencies_out, num_dependencies_out = handle_return(
+            driver.cuStreamGetCaptureInfo(stream_handle)
+        )
+
+        child_node = handle_return(
+            driver.cuGraphAddChildGraphNode(graph_out, dependencies_out, num_dependencies_out, child_graph._mnff.graph)
+        )
+        handle_return(
+            driver.cuStreamUpdateCaptureDependencies(
+                stream_handle,
+                [child_node],
+                1,
+                driver.CUstreamUpdateCaptureDependencies_flags.CU_STREAM_SET_CAPTURE_DEPENDENCIES,
+            )
+        )
+
 
 class Graph:
     """Represents an executable graph.
@@ -733,46 +772,3 @@ def launch(self, stream: Stream):
 
         """
         handle_return(driver.cuGraphLaunch(self._mnff.graph, stream.handle))
-
-
-def launch_graph(parent_graph: GraphBuilder, child_graph: GraphBuilder):
-    """Adds the child :obj:`~_graph.GraphBuilder` builder into parent graph builder.
-
-    The child graph builder will be added as a child node to the parent graph builder.
-
-    Parameters
-    ----------
-    parent_graph : :obj:`~_graph.GraphBuilder`
-        The parent graph builder. Must be in building state.
-    child_graph : :obj:`~_graph.GraphBuilder`
-        The child graph builder. Must have finished building.
-
-    """
-
-    if (_driver_ver < 12000) or (_py_major_minor < (12, 0)):
-        raise NotImplementedError(
-            f"Launching child graphs is not implemented for versions older than CUDA 12."
-            f"Found driver version is {_driver_ver} and binding version is {_py_major_minor}"
-        )
-
-    if not child_graph._building_ended:
-        raise ValueError("Child graph has not finished building.")
-
-    if not parent_graph.is_building:
-        raise ValueError("Parent graph is being built.")
-
-    _, _, graph_out, dependencies_out, num_dependencies_out = handle_return(
-        driver.cuStreamGetCaptureInfo(parent_graph.stream.handle)
-    )
-
-    child_node = handle_return(
-        driver.cuGraphAddChildGraphNode(graph_out, dependencies_out, num_dependencies_out, child_graph._mnff.graph)
-    )
-    handle_return(
-        driver.cuStreamUpdateCaptureDependencies(
-            parent_graph.stream.handle,
-            [child_node],
-            1,
-            driver.CUstreamUpdateCaptureDependencies_flags.CU_STREAM_SET_CAPTURE_DEPENDENCIES,
-        )
-    )

cuda_core/tests/test_graph.py

@@ -18,7 +18,6 @@
     Program,
     ProgramOptions,
     launch,
-    launch_graph,
 )
 from cuda.core.experimental._memory import _DefaultPinnedMemorySource
 from cuda.core.experimental._utils.cuda_utils import NVRTCError, handle_return
@@ -517,7 +516,7 @@ def test_graph_child_graph(init_cuda):
 
     ## Add child
     try:
-        launch_graph(gb_parent, gb_child)
+        gb_parent.add_child(gb_child)
     except NotImplementedError as e:
         with pytest.raises(
             NotImplementedError,