Merge branch 'main' into default_mr

Author: leofang

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, DeviceMemoryResource, MemoryResource, _SynchronousMemoryResource
 from cuda.core.experimental._stream import IsStreamT, 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,
@@ -355,6 +356,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,6 +17,8 @@ CUDA runtime
    :toctree: generated/
 
    Device
+   Graph
+   GraphBuilder
    launch
    Buffer
    MemoryResource
@@ -26,6 +28,8 @@ CUDA runtime
    :template: dataclass.rst
 
    EventOptions
+   GraphCompleteOptions
+   GraphDebugPrintOptions
    StreamOptions
    LaunchConfig
 

cuda_core/docs/source/api.rst

@@ -12,6 +12,9 @@ Highlights
 ----------
 
 - Starting this release ``cuda.core`` is licensed under Apache 2.0. The biggest implication of this change is that we are open to external contribution now! Please kindly follow the :ref:`Contributor Guide <contributor_guide>` for detailed instructions.
+- 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

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