Draft

Author: vzhurba01

cuda_core/cuda/core/experimental/__init__.py

@@ -5,7 +5,7 @@
 from cuda.core.experimental import utils
 from cuda.core.experimental._device import Device
 from cuda.core.experimental._event import EventOptions
-from cuda.core.experimental._graph import GraphBuilder, Graph
+from cuda.core.experimental._graph import Graph, GraphBuilder
 from cuda.core.experimental._launcher import LaunchConfig, launch
 from cuda.core.experimental._linker import Linker, LinkerOptions
 from cuda.core.experimental._program import Program, ProgramOptions

cuda_core/cuda/core/experimental/_device.py

@@ -340,6 +340,6 @@ def sync(self):
         handle_return(runtime.cudaDeviceSynchronize())
 
     @precondition(_check_context_initialized)
-    def create_graph(self) -> GraphBuilder:
+    def build_graph(self) -> GraphBuilder:
         private_stream = self.create_stream()
         return GraphBuilder._init(stream=private_stream)

cuda_core/cuda/core/experimental/_graph.py

@@ -4,11 +4,8 @@
 
 from __future__ import annotations
 
-import os
-import warnings
-import weakref
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Optional, Tuple
 
 if TYPE_CHECKING:
     from cuda.core.experimental._stream import Stream
@@ -17,8 +14,7 @@
 
 @dataclass
 class DebugPrintOptions:
-    """
-    """
+    """ """
 
     VERBOSE: bool = False
     RUNTIME_TYPES: bool = False
@@ -60,7 +56,7 @@ def _init(stream, is_root=True):
         #       If from Stream, then we can't
         self._stream = stream
         self._capturing = False
-        self._is_root = is_root # TODO: Is this info needed?
+        self._is_root = is_root  # TODO: Is this info needed?
         return self
 
     def _check_capture_stream_provided(self, *args, **kwargs):
@@ -101,8 +97,10 @@ def is_capture_active(self) -> bool:
         elif capture_status == driver.CUstreamCaptureStatus.CU_STREAM_CAPTURE_STATUS_ACTIVE:
             return True
         elif capture_status == driver.CUstreamCaptureStatus.CU_STREAM_CAPTURE_STATUS_INVALIDATED:
-            raise RuntimeError("Stream is part of a capture sequence that has been invalidated, but "
-                               "not terminated. The capture sequence must be terminated with self.`().")
+            raise RuntimeError(
+                "Stream is part of a capture sequence that has been invalidated, but "
+                "not terminated. The capture sequence must be terminated with self.`()."
+            )
         else:
             raise NotImplementedError(f"Unsupported capture stuse type received: {capture_status}")
 
@@ -124,30 +122,30 @@ def debug_dot_print(self, path, options: Optional[DebugPrintOptions] = None):
 
         handle_return(driver.cuGraphDebugDotPrint(self._graph, path, options_value))
 
-    def split(self, count) -> Tuple[Graph, ...]:
+    def fork(self, count) -> Tuple[Graph, ...]:
         if count <= 1:
-            raise ValueError(f"Invalid split count: expecting >= 2, got {count}")
+            raise ValueError(f"Invalid fork count: expecting >= 2, got {count}")
 
         # 1. Record an event on our stream
         event = self._stream.record()
 
         # TODO: Steps 2,3,4 can be combined under a single loop
 
-        # 2. Create a streams for each of the new splits
-        # TODO: Optimization where one of the split stream is allowed to use
+        # 2. Create a streams for each of the new forks
+        # TODO: Optimization where one of the fork stream is allowed to use
         # TODO: Should use the same stream options as initial stream??
-        split_stream = [self._stream.device.create_stream() for i in range(count)]
+        fork_stream = [self._stream.device.create_stream() for i in range(count)]
 
         # 3. Have each new stream wait on our singular event
-        for stream in split_stream:
+        for stream in fork_stream:
             stream.wait(event)
 
         # 4. Discard the event
         # TODO: Is this actually allowed when using with a graph? Surely, since it just needs to create an edge for us... right?
         event.close()
 
-        # 5. Create new graph builders for each new stream split
-        return [GraphBuilder._init(stream=stream, is_root=False) for stream in split_stream]
+        # 5. Create new graph builders for each new stream fork
+        return [GraphBuilder._init(stream=stream, is_root=False) for stream in fork_stream]
 
     def join(self, *graph_builders):
         if len(graph_builders) < 1:
@@ -158,6 +156,18 @@ def join(self, *graph_builders):
             # TODO: Can we close each of those new streams? Do we need weakref?
             graph.close()
 
+    def create_conditional_handle(self, default_value=None):
+        pass
+
+    def if_cond(self, handle):
+        pass
+
+    def if_else(self, handle):
+        pass
+
+    def switch(self, handle, count):
+        pass
+
     def close(self):
         if self._capturing:
             raise RuntimeError("Trying to close a graph builder who is still capturing")
@@ -166,8 +176,7 @@ def close(self):
 
 
 class Graph:
-    """
-    """
+    """ """
 
     def __init__(self):
         raise RuntimeError("directly constructing a Graph instance is not supported")

cuda_core/cuda/core/experimental/_stream.py

@@ -294,7 +294,7 @@ def __cuda_stream__(self):
 
         return Stream._init(obj=_stream_holder())
 
-    def create_graph(self) -> GraphBuilder:
+    def build_graph(self) -> GraphBuilder:
         return GraphBuilder._init(stream=self)
 
 

cuda_core/tests/test_graph.py

@@ -6,16 +6,16 @@
 # this software and related documentation outside the terms of the EULA
 # is strictly prohibited.
 
-import pytest
 
 from cuda.core.experimental import Device, LaunchConfig, Program, ProgramOptions, launch
+
 # from cuda.core.experimental import Device, Stream, StreamOptions
 # from cuda.core.experimental._event import Event
 # from cuda.core.experimental._stream import LEGACY_DEFAULT_STREAM, PER_THREAD_DEFAULT_STREAM, default_stream
 
 
 def test_graph_is_capture_alive(init_cuda):
-    graph_builder = Device().create_graph()
+    graph_builder = Device().build_graph()
     assert graph_builder.is_capture_active() is False
     graph_builder.begin_capture()
     assert graph_builder.is_capture_active() is True
@@ -24,7 +24,6 @@ def test_graph_is_capture_alive(init_cuda):
 
 
 def test_graph_straight(init_cuda):
-
     # TODO: Maybe share these between tests?
     code = """
     __global__ void empty_kernel() {}
@@ -35,9 +34,8 @@ def test_graph_straight(init_cuda):
     mod = prog.compile("cubin", name_expressions=("empty_kernel",))
     empty_kernel = mod.get_kernel("empty_kernel")
 
-
     # Test start
-    graph_builder = Device().create_graph()
+    graph_builder = Device().build_graph()
     config = LaunchConfig(grid=1, block=1, stream=graph_builder.legacy_stream_capture)
 
     assert graph_builder.is_capture_active() is False
@@ -48,8 +46,7 @@ def test_graph_straight(init_cuda):
     graph_builder.end_capture()
 
 
-def test_graph_split_join(init_cuda):
-
+def test_graph_fork_join(init_cuda):
     # TODO: Maybe share these between tests?
     code = """
     __global__ void empty_kernel() {}
@@ -60,14 +57,13 @@ def test_graph_split_join(init_cuda):
     mod = prog.compile("cubin", name_expressions=("empty_kernel",))
     empty_kernel = mod.get_kernel("empty_kernel")
 
-
     # Test start
-    graph_builder = Device().create_graph()
+    graph_builder = Device().build_graph()
     assert graph_builder.is_capture_active() is False
     graph_builder.begin_capture()
     launch(empty_kernel, {"grid": 1, "block": 1, "stream": graph_builder.legacy_stream_capture})
 
-    left, right = graph_builder.split(2)
+    left, right = graph_builder.fork(2)
     launch(empty_kernel, {"grid": 1, "block": 1, "stream": left.legacy_stream_capture})
     launch(empty_kernel, {"grid": 1, "block": 1, "stream": left.legacy_stream_capture})
     launch(empty_kernel, {"grid": 1, "block": 1, "stream": right.legacy_stream_capture})