Draft

Draft

Initial weakref work

Use stream as name

stash

Author: vzhurba01

cuda_core/cuda/core/experimental/__init__.py

@@ -5,6 +5,7 @@
 from cuda.core.experimental import utils
 from cuda.core.experimental._device import Device
 from cuda.core.experimental._event import Event, EventOptions
+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._module import ObjectCode

cuda_core/cuda/core/experimental/_device.py

@@ -7,6 +7,7 @@
 
 from cuda.core.experimental._context import Context, ContextOptions
 from cuda.core.experimental._event import Event, EventOptions
+from cuda.core.experimental._graph import GraphBuilder
 from cuda.core.experimental._memory import Buffer, MemoryResource, _DefaultAsyncMempool, _SynchronousMemoryResource
 from cuda.core.experimental._stream import Stream, StreamOptions, default_stream
 from cuda.core.experimental._utils.clear_error_support import assert_type
@@ -1271,3 +1272,8 @@ def sync(self):
 
         """
         handle_return(runtime.cudaDeviceSynchronize())
+
+    @precondition(_check_context_initialized)
+    def build_graph(self) -> GraphBuilder:
+        private_stream = self.create_stream()
+        return GraphBuilder._init(stream=private_stream)

cuda_core/cuda/core/experimental/_graph.py

@@ -0,0 +1,215 @@
+# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+import weakref
+from typing import TYPE_CHECKING, Optional, Tuple
+
+if TYPE_CHECKING:
+    from cuda.core.experimental._stream import Stream
+from cuda.core.experimental._utils import driver, handle_return, precondition
+
+
+@dataclass
+class DebugPrintOptions:
+    """ """
+
+    VERBOSE: bool = False
+    RUNTIME_TYPES: bool = False
+    KERNEL_NODE_PARAMS: bool = False
+    MEMCPY_NODE_PARAMS: bool = False
+    MEMSET_NODE_PARAMS: bool = False
+    HOST_NODE_PARAMS: bool = False
+    EVENT_NODE_PARAMS: bool = False
+    EXT_SEMAS_SIGNAL_NODE_PARAMS: bool = False
+    EXT_SEMAS_WAIT_NODE_PARAMS: bool = False
+    KERNEL_NODE_ATTRIBUTES: bool = False
+    HANDLES: bool = False
+    MEM_ALLOC_NODE_PARAMS: bool = False
+    MEM_FREE_NODE_PARAMS: bool = False
+    BATCH_MEM_OP_NODE_PARAMS: bool = False
+    EXTRA_TOPO_INFO: bool = False
+    CONDITIONAL_NODE_PARAMS: bool = False
+
+
+class GraphBuilder:
+    """TBD
+
+    Directly creating a :obj:`~_graph.GraphBuilder` is not supported due
+    to ambiguity. New graph builders should instead be created through a
+    :obj:`~_device.Device`, or a :obj:`~_stream.stream` object
+    """
+
+    class _MembersNeededForFinalize:
+        __slots__ = ("stream", "graph")
+
+        def __init__(self, graph_builder_obj, stream_obj):
+            self.stream = stream_obj
+            self.graph = None
+            weakref.finalize(graph_builder_obj, self.close)
+
+        def close(self):
+            # FIXME: Are the stream and graph builder racing for the weakref callback?
+            #        If so, maybe we need to enforce that all capture is completed
+            status = handle_return(driver.cuStreamGetCaptureInfo(self.stream.handle))[0]
+            if status != driver.CUstreamCaptureStatus.CU_STREAM_CAPTURE_STATUS_NONE:
+                # Callback routine needs to end capture for error free handling
+                handle_return(driver.cuStreamEndCapture(self.stream.handle))
+            if self.graph:
+                handle_return(driver.cuGraphDestroy(self.graph))
+            self.graph = None
+
+    __slots__ = ("__weakref__", "_mnff", "_is_primary", "_capturing")
+
+    def __init__(self):
+        raise NotImplementedError(
+            "directly creating a Graph object can be ambiguous. Please either "
+            "call Device.create_graph() or stream.creating_graph()"
+        )
+
+    @staticmethod
+    def _init(stream, _is_primary=True):
+        self = GraphBuilder.__new__(GraphBuilder)
+        # TODO: I need to know if we own this stream object.
+        #       If from Device(), then we can destroy it on close
+        #       If from Stream, then we can't
+        self._capturing = False
+        self._is_primary = _is_primary
+        self._mnff = GraphBuilder._MembersNeededForFinalize(self, stream)
+        return self
+
+    def _check_capture_stream_provided(self, *args, **kwargs):
+        if self._mnff.stream == None:
+            raise RuntimeError("Tried to use a stream capture operation on a graph builder without a stream")
+
+    @property
+    def stream(self) -> Stream:
+        return self._mnff.stream
+
+    @property
+    def is_primary(self) -> bool:
+        return self._is_primary
+
+    @precondition(_check_capture_stream_provided)
+    def begin_capture(self, mode="global"):
+        # Supports "global", "local" or "relaxed"
+        # TODO; Test case for each mode and fail
+        if mode == "global":
+            capture_mode = driver.CUstreamCaptureMode.CU_STREAM_CAPTURE_MODE_GLOBAL
+        elif mode == "local":
+            capture_mode = driver.CUstreamCaptureMode.CU_STREAM_CAPTURE_MODE_THREAD_LOCAL
+        elif mode == "relaxed":
+            capture_mode = driver.CUstreamCaptureMode.CU_STREAM_CAPTURE_MODE_RELAXED
+        else:
+            raise ValueError(f"Only 'global', 'local' or 'relaxed' capture mode are supported, got {capture_mode}")
+
+        handle_return(driver.cuStreamBeginCapture(self._mnff.stream.handle, capture_mode))
+        self._capturing = True
+
+    @precondition(_check_capture_stream_provided)
+    def is_capture_active(self) -> bool:
+        result = handle_return(driver.cuStreamGetCaptureInfo(self._mnff.stream.handle))
+
+        capture_status = result[0]
+        if capture_status == driver.CUstreamCaptureStatus.CU_STREAM_CAPTURE_STATUS_NONE:
+            return False
+        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.`()."
+            )
+        else:
+            raise NotImplementedError(f"Unsupported capture stuse type received: {capture_status}")
+
+    @precondition(_check_capture_stream_provided)
+    def end_capture(self):
+        if not self._capturing:
+            raise RuntimeError("Stream is not capturing. Did you forget to call begin_capture()?")
+        self._mnff.graph = handle_return(driver.cuStreamEndCapture(self.stream.handle))
+        self._capturing = False
+
+    def debug_dot_print(self, path, options: Optional[DebugPrintOptions] = None):
+        # TODO: We should be able to print one while the capture is happening right? Just need to make sure driver version is new enough.
+        if self._mnff.graph == None:
+            raise RuntimeError("Graph needs to be built before generating a DOT debug file")
+
+        # TODO: Apply each option to the value
+        options_value = 0
+
+        handle_return(driver.cuGraphDebugDotPrint(self._mnff.graph, path, options_value))
+
+    def fork(self, count) -> Tuple[GraphBuilder, ...]:
+        if count <= 1:
+            raise ValueError(f"Invalid fork count: expecting >= 2, got {count}")
+
+        # 1. Record an event on our stream
+        event = self._mnff.stream.record()
+
+        # TODO: Steps 2,3,4 can be combined under a single loop
+
+        # 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??
+        fork_stream = [self._mnff.stream.device.create_stream() for i in range(count)]
+
+        # 3. Have each new stream wait on our singular event
+        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 fork
+        return [GraphBuilder._init(stream=stream, is_primary=False) for stream in fork_stream]
+
+    def join(self, *graph_builders):
+        if len(graph_builders) < 1:
+            raise ValueError("Must specify which graphs should join but none were given")
+
+        # Assert that none of the graph_builders are primary
+        for graph in graph_builders:
+            if graph.is_primary:
+                raise ValueError("The primary graph builder should not be joined. Others builders should instead be joined onto it.")
+
+        for graph in graph_builders:
+            self._mnff.stream.wait(graph.stream)
+            # TODO: Do we close them now or let weakref handle it during garbage collection?
+            #       This is a perf question, is there a good default?
+            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._mnff.capturing:
+            # Explicitly trying to close a graph builder who is still capturing is not allowed
+            raise RuntimeError("Trying to close a graph builder who is still capturing. Did you forget to call end_capture()?")
+        self._mnff.close()
+
+
+class Graph:
+    """ """
+
+    def __init__(self):
+        raise RuntimeError("directly constructing a Graph instance is not supported")
+
+    @staticmethod
+    def _init(graph):
+        self = Graph.__new__(Graph)
+        self._graph = graph
+        return self

cuda_core/cuda/core/experimental/_stream.py

@@ -23,6 +23,7 @@
     handle_return,
     runtime,
 )
+from cuda.core.experimental._graph import GraphBuilder
 
 
 @dataclass
@@ -342,6 +343,9 @@ def __cuda_stream__(self):
 
         return Stream._init(obj=_stream_holder())
 
+    def build_graph(self) -> GraphBuilder:
+        return GraphBuilder._init(stream=self)
+
 
 LEGACY_DEFAULT_STREAM = Stream._legacy_default()
 PER_THREAD_DEFAULT_STREAM = Stream._per_thread_default()

cuda_core/tests/test_graph.py

@@ -0,0 +1,76 @@
+# Copyright 2025 NVIDIA Corporation.  All rights reserved.
+#
+# Please refer to the NVIDIA end user license agreement (EULA) associated
+# with this source code for terms and conditions that govern your use of
+# this software. Any use, reproduction, disclosure, or distribution of
+# this software and related documentation outside the terms of the EULA
+# is strictly prohibited.
+
+
+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().build_graph()
+    assert graph_builder.is_capture_active() is False
+    graph_builder.begin_capture()
+    assert graph_builder.is_capture_active() is True
+    graph_builder.end_capture()
+    assert graph_builder.is_capture_active() is False
+
+
+def test_graph_straight(init_cuda):
+    # TODO: Maybe share these between tests?
+    code = """
+    __global__ void empty_kernel() {}
+    """
+    arch = "".join(f"{i}" for i in Device().compute_capability)
+    program_options = ProgramOptions(std="c++17", arch=f"sm_{arch}")
+    prog = Program(code, code_type="c++", options=program_options)
+    mod = prog.compile("cubin", name_expressions=("empty_kernel",))
+    empty_kernel = mod.get_kernel("empty_kernel")
+
+    # Test start
+    graph_builder = Device().build_graph()
+    config = LaunchConfig(grid=1, block=1, stream=graph_builder.stream)
+
+    assert graph_builder.is_capture_active() is False
+    graph_builder.begin_capture()
+    launch(empty_kernel, {"grid": 1, "block": 1, "stream": graph_builder.stream})
+    launch(empty_kernel, {"grid": 1, "block": 1, "stream": graph_builder.stream})
+    launch(empty_kernel, {"grid": 1, "block": 1, "stream": graph_builder.stream})
+    graph_builder.end_capture()
+
+
+def test_graph_fork_join(init_cuda):
+    # TODO: Maybe share these between tests?
+    code = """
+    __global__ void empty_kernel() {}
+    """
+    arch = "".join(f"{i}" for i in Device().compute_capability)
+    program_options = ProgramOptions(std="c++17", arch=f"sm_{arch}")
+    prog = Program(code, code_type="c++", options=program_options)
+    mod = prog.compile("cubin", name_expressions=("empty_kernel",))
+    empty_kernel = mod.get_kernel("empty_kernel")
+
+    # Test start
+    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.stream})
+
+    left, right = graph_builder.fork(2)
+    launch(empty_kernel, {"grid": 1, "block": 1, "stream": left.stream})
+    launch(empty_kernel, {"grid": 1, "block": 1, "stream": left.stream})
+    launch(empty_kernel, {"grid": 1, "block": 1, "stream": right.stream})
+    launch(empty_kernel, {"grid": 1, "block": 1, "stream": right.stream})
+    graph_builder.join(left, right)
+
+    launch(empty_kernel, {"grid": 1, "block": 1, "stream": graph_builder.stream})
+    graph_builder.end_capture()
+
+    graph_builder.debug_dot_print(b"vlad.dot")