Add phase 1 of CUDA Graphs support (#455)

* Draft

Draft

Initial weakref work

Use stream as name

stash

* Align with design

* Finish graphs APIs and add tests

* Revert removed line in pyproject

* Self review

* Run pre-commit

* Return clean errors when conditional is unsupported

* Add numpy version check to tests

* Add build modes

* Check for both old versions of NVRTC and driver

* Fix recursive is building check

* Run pre-commit

* Close MR allocs so that they don't GC during future Graph building

* Check binding version in addition to driver version

* Check NVRTC error correctly

* Fix options check

* Change default build mode to RELAXED

* Resolve some review comments

* Resolve review comments about naming

* Extend test and fix spelling

* Pre-commit and revert fix

* use CUDA context memoized by stream

* avoid erasing the type CUgraphConditionalHandle

* launch_graph -> gb.add_child

* add docs

* fix typo

---------

Co-authored-by: Leo Fang <[email protected]>
Co-authored-by: Leo Fang <[email protected]>

Author: vzhurba01

cuda_core/cuda/core/experimental/__init__.py

@@ -5,6 +5,12 @@
 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,
+    GraphCompleteOptions,
+    GraphDebugPrintOptions,
+)
 from cuda.core.experimental._launch_config import LaunchConfig
 from cuda.core.experimental._launcher import launch
 from cuda.core.experimental._linker import Linker, LinkerOptions

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
@@ -1298,3 +1299,15 @@ def sync(self):
 
         """
         handle_return(runtime.cudaDeviceSynchronize())
+
+    @precondition(_check_context_initialized)
+    def create_graph_builder(self) -> GraphBuilder:
+        """Create a new :obj:`~_graph.GraphBuilder` object.
+
+        Returns
+        -------
+        :obj:`~_graph.GraphBuilder`
+            Newly created graph builder object.
+
+        """
+        return GraphBuilder._init(stream=self.create_stream(), is_stream_owner=True)

cuda_core/cuda/core/experimental/_graph.py

@@ -16,6 +16,7 @@ import ctypes
 import numpy
 
 from cuda.core.experimental._memory import Buffer
+from cuda.core.experimental._utils.cuda_utils import driver
 
 
 ctypedef cpp_complex.complex[float] cpp_single_complex
@@ -235,6 +236,10 @@ cdef class ParamHolder:
             if not_prepared:
                 not_prepared = prepare_ctypes_arg(self.data, self.data_addresses, arg, i)
             if not_prepared:
+                # TODO: revisit this treatment if we decide to cythonize cuda.core
+                if isinstance(arg, driver.CUgraphConditionalHandle):
+                    prepare_arg[intptr_t](self.data, self.data_addresses, <intptr_t>int(arg), i)
+                    continue
                 # TODO: support ctypes/numpy struct
                 raise TypeError("the argument is of unsupported type: " + str(type(arg)))
 

cuda_core/cuda/core/experimental/_kernel_arg_handler.pyx

@@ -15,6 +15,7 @@
     from cuda.core.experimental._device import Device
 from cuda.core.experimental._context import Context
 from cuda.core.experimental._event import Event, EventOptions
+from cuda.core.experimental._graph import GraphBuilder
 from cuda.core.experimental._utils.clear_error_support import assert_type
 from cuda.core.experimental._utils.cuda_utils import (
     check_or_create_options,
@@ -342,6 +343,19 @@ def __cuda_stream__(self):
 
         return Stream._init(obj=_stream_holder())
 
+    def create_graph_builder(self) -> GraphBuilder:
+        """Create a new :obj:`~_graph.GraphBuilder` object.
+
+        The new graph builder will be associated with this stream.
+
+        Returns
+        -------
+        :obj:`~_graph.GraphBuilder`
+            Newly created graph builder object.
+
+        """
+        return GraphBuilder._init(stream=self, is_stream_owner=False)
+
 
 LEGACY_DEFAULT_STREAM = Stream._legacy_default()
 PER_THREAD_DEFAULT_STREAM = Stream._per_thread_default()

cuda_core/cuda/core/experimental/_stream.py

@@ -17,11 +17,15 @@ CUDA runtime
    :toctree: generated/
 
    Device
+   Graph
+   GraphBuilder
    launch
 
    :template: dataclass.rst
 
    EventOptions
+   GraphCompleteOptions
+   GraphDebugPrintOptions
    StreamOptions
    LaunchConfig
 

cuda_core/docs/source/api.rst

@@ -11,6 +11,9 @@ Highlights
 ----------
 
 - Starting this release ``cuda.core`` is licensed under Apache 2.0.
+- Initial support for CUDA graphs (phase 1).
+   - In this release, we support building a CUDA graph that captures kernel launches. The captured graph can be replayed to reduce
+     latency. Graph split/join and conditional nodes are supported.
 
 
 Breaking Changes

cuda_core/docs/source/release/0.3.0-notes.rst

@@ -165,6 +165,7 @@ def test_error_timing_incomplete():
     arr[0] = 1
     event3.sync()
     event3 - event1  # this should work
+    b.close()
 
 
 def test_event_device(init_cuda):