Publish the graph API as cuda.core.graph

cuda_core/cuda/core/__init__.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # SPDX-License-Identifier: Apache-2.0
 
@@ -31,12 +31,6 @@
 from cuda.core import system, utils
 from cuda.core._device import Device
 from cuda.core._event import Event, EventOptions
-from cuda.core._graph import (
-    Graph,
-    GraphBuilder,
-    GraphCompleteOptions,
-    GraphDebugPrintOptions,
-)
 from cuda.core._graphics import GraphicsResource
 from cuda.core._launch_config import LaunchConfig
 from cuda.core._launcher import launch
@@ -69,3 +63,12 @@
     StreamOptions,
 )
 from cuda.core._tensor_map import TensorMapDescriptor, TensorMapDescriptorOptions
+from cuda.core.graph import (
+    Condition,
+    Graph,
+    GraphAllocOptions,
+    GraphBuilder,
+    GraphCompleteOptions,
+    GraphDebugPrintOptions,
+    GraphDef,
+)

cuda_core/cuda/core/_device.pyx

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # SPDX-License-Identifier: Apache-2.0
 
@@ -24,7 +24,6 @@ from cuda.core._resource_handles cimport (
     as_cu,
 )
 
-from cuda.core._graph import GraphBuilder
 from cuda.core._stream import IsStreamT, Stream, StreamOptions
 from cuda.core._utils.clear_error_support import assert_type
 from cuda.core._utils.cuda_utils import (
@@ -1363,15 +1362,17 @@ class Device:
         self._check_context_initialized()
         handle_return(runtime.cudaDeviceSynchronize())
 
-    def create_graph_builder(self) -> GraphBuilder:
-        """Create a new :obj:`~_graph.GraphBuilder` object.
+    def create_graph_builder(self) -> "GraphBuilder":
+        """Create a new :obj:`~graph.GraphBuilder` object.
 
         Returns
         -------
-        :obj:`~_graph.GraphBuilder`
+        :obj:`~graph.GraphBuilder`
             Newly created graph builder object.
 
         """
+        from cuda.core.graph._graph_builder import GraphBuilder
+
         self._check_context_initialized()
         return GraphBuilder._init(stream=self.create_stream(), is_stream_owner=True)
 

cuda_core/cuda/core/_launcher.pyx

@@ -26,7 +26,7 @@ def launch(stream: Stream | GraphBuilder | IsStreamT, config: LaunchConfig, kern
 
     Parameters
     ----------
-    stream : :obj:`~_stream.Stream` | :obj:`~_graph.GraphBuilder`
+    stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`
         The stream establishing the stream ordering semantic of a
         launch.
     config : :obj:`LaunchConfig`

cuda_core/cuda/core/_memory/_buffer.pyx

@@ -182,7 +182,7 @@ cdef class Buffer:
 
         Parameters
         ----------
-        stream : :obj:`~_stream.Stream` | :obj:`~_graph.GraphBuilder`, optional
+        stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`, optional
             The stream object to use for asynchronous deallocation. If None,
             the deallocation stream stored in the handle is used.
         """
@@ -206,7 +206,7 @@ cdef class Buffer:
         ----------
         dst : :obj:`~_memory.Buffer`
             Source buffer to copy data from
-        stream : :obj:`~_stream.Stream` | :obj:`~_graph.GraphBuilder`
+        stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`
             Keyword argument specifying the stream for the
             asynchronous copy
 
@@ -237,7 +237,7 @@ cdef class Buffer:
         ----------
         src : :obj:`~_memory.Buffer`
             Source buffer to copy data from
-        stream : :obj:`~_stream.Stream` | :obj:`~_graph.GraphBuilder`
+        stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`
             Keyword argument specifying the stream for the
             asynchronous copy
 
@@ -262,7 +262,7 @@ cdef class Buffer:
         value : int | :obj:`collections.abc.Buffer`
             - int: Must be in range [0, 256). Converted to 1 byte.
             - :obj:`collections.abc.Buffer`: Must be 1, 2, or 4 bytes.
-        stream : :obj:`~_stream.Stream` | :obj:`~_graph.GraphBuilder`
+        stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`
             Stream for the asynchronous fill operation.
 
         Raises
@@ -496,7 +496,7 @@ cdef class MemoryResource:
         ----------
         size : int
             The size of the buffer to allocate, in bytes.
-        stream : :obj:`~_stream.Stream` | :obj:`~_graph.GraphBuilder`, optional
+        stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`, optional
             The stream on which to perform the allocation asynchronously.
             If None, it is up to each memory resource implementation to decide
             and document the behavior.
@@ -518,7 +518,7 @@ cdef class MemoryResource:
             The pointer or handle to the buffer to deallocate.
         size : int
             The size of the buffer to deallocate, in bytes.
-        stream : :obj:`~_stream.Stream` | :obj:`~_graph.GraphBuilder`, optional
+        stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`, optional
             The stream on which to perform the deallocation asynchronously.
             If None, it is up to each memory resource implementation to decide
             and document the behavior.

cuda_core/cuda/core/_memory/_memory_pool.pyx

@@ -129,7 +129,7 @@ cdef class _MemPool(MemoryResource):
         ----------
         size : int
             The size of the buffer to allocate, in bytes.
-        stream : :obj:`~_stream.Stream` | :obj:`~_graph.GraphBuilder`, optional
+        stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`, optional
             The stream on which to perform the allocation asynchronously.
             If None, an internal stream is used.
 
@@ -153,7 +153,7 @@ cdef class _MemPool(MemoryResource):
             The pointer or handle to the buffer to deallocate.
         size : int
             The size of the buffer to deallocate, in bytes.
-        stream : :obj:`~_stream.Stream` | :obj:`~_graph.GraphBuilder`, optional
+        stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`, optional
             The stream on which to perform the deallocation asynchronously.
             If the buffer is deallocated without an explicit stream, the allocation stream
             is used.

cuda_core/cuda/core/_stream.pyx

@@ -38,7 +38,6 @@ from cuda.core._resource_handles cimport (
     as_py,
 )
 
-from cuda.core._graph import GraphBuilder
 
 
 @dataclass
@@ -354,17 +353,19 @@ cdef class Stream:
 
         return Stream._init(obj=_stream_holder())
 
-    def create_graph_builder(self) -> GraphBuilder:
-        """Create a new :obj:`~_graph.GraphBuilder` object.
+    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`
+        :obj:`~graph.GraphBuilder`
             Newly created graph builder object.
 
         """
+        from cuda.core.graph._graph_builder import GraphBuilder
+
         return GraphBuilder._init(stream=self, is_stream_owner=False)
 
 
@@ -466,6 +467,8 @@ cdef cydriver.CUstream _handle_from_stream_protocol(obj) except*:
 # Helper for API functions that accept either Stream or GraphBuilder. Performs
 # needed checks and returns the relevant stream.
 cdef Stream Stream_accept(arg, bint allow_stream_protocol=False):
+    from cuda.core.graph._graph_builder import GraphBuilder
+
     if isinstance(arg, Stream):
         return <Stream>(arg)
     elif isinstance(arg, GraphBuilder):

cuda_core/cuda/core/experimental/__init__.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # SPDX-License-Identifier: Apache-2.0
 
@@ -46,12 +46,6 @@ def _warn_deprecated():
 
 from cuda.core._device import Device
 from cuda.core._event import Event, EventOptions
-from cuda.core._graph import (
-    Graph,
-    GraphBuilder,
-    GraphCompleteOptions,
-    GraphDebugPrintOptions,
-)
 from cuda.core._launch_config import LaunchConfig
 from cuda.core._launcher import launch
 from cuda.core._layout import _StridedLayout
@@ -73,3 +67,9 @@ def _warn_deprecated():
 from cuda.core._module import Kernel, ObjectCode
 from cuda.core._program import Program, ProgramOptions
 from cuda.core._stream import Stream, StreamOptions
+from cuda.core.graph import (
+    Graph,
+    GraphBuilder,
+    GraphCompleteOptions,
+    GraphDebugPrintOptions,
+)

cuda_core/cuda/core/_graph/_graph_def/__init__.pxd → cuda_core/cuda/core/graph/__init__.pxd

@@ -2,9 +2,9 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from cuda.core._graph._graph_def._graph_def cimport Condition, GraphDef
-from cuda.core._graph._graph_def._graph_node cimport GraphNode
-from cuda.core._graph._graph_def._subclasses cimport (
+from cuda.core.graph._graph_def cimport Condition, GraphDef
+from cuda.core.graph._graph_node cimport GraphNode
+from cuda.core.graph._subclasses cimport (
     AllocNode,
     ChildGraphNode,
     ConditionalNode,

cuda_core/cuda/core/graph/__init__.py

@@ -0,0 +1,33 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from cuda.core.graph._graph_builder import (
+    Graph,
+    GraphBuilder,
+    GraphCompleteOptions,
+    GraphDebugPrintOptions,
+)
+from cuda.core.graph._graph_def import (
+    Condition,
+    GraphAllocOptions,
+    GraphDef,
+)
+from cuda.core.graph._graph_node import GraphNode
+from cuda.core.graph._subclasses import (
+    AllocNode,
+    ChildGraphNode,
+    ConditionalNode,
+    EmptyNode,
+    EventRecordNode,
+    EventWaitNode,
+    FreeNode,
+    HostCallbackNode,
+    IfElseNode,
+    IfNode,
+    KernelNode,
+    MemcpyNode,
+    MemsetNode,
+    SwitchNode,
+    WhileNode,
+)

cuda_core/cuda/core/_graph/_graph_def/_adjacency_set_proxy.pyx → cuda_core/cuda/core/graph/_adjacency_set_proxy.pyx

@@ -7,7 +7,7 @@
 from libc.stddef cimport size_t
 from libcpp.vector cimport vector
 from cuda.bindings cimport cydriver
-from cuda.core._graph._graph_def._graph_node cimport GraphNode
+from cuda.core.graph._graph_node cimport GraphNode
 from cuda.core._resource_handles cimport (
     GraphHandle,
     GraphNodeHandle,