Add build modes

Author: vzhurba01

cuda_core/cuda/core/experimental/_graph.py

@@ -189,10 +189,32 @@ def is_join_required(self) -> bool:
         """Returns True if this graph builder must be joined before building is ended."""
         return self._mnff.is_join_required
 
-    def begin_building(self) -> GraphBuilder:
-        """Begins the building process."""
+    def begin_building(self, mode="global") -> GraphBuilder:
+        """Begins the building process.
+
+        Build `mode` for controlling interaction with other API calls must be one of the following:
+
+        - `global` : Prohibit potentially unsafe operations across all streams in the process.
+        - `thread_local` : Prohibit potentially unsafe operations in streams created by the current thread.
+        - `relaxed` : The local thread is not prohibited from potentially unsafe operations.
+
+        Parameters
+        ----------
+        mode : str, optional
+            Build mode to control the interaction with other API calls that are porentially unsafe.
+
+        """
         if self._building_ended:
             raise RuntimeError("Cannot resume building after building has ended.")
+        if mode not in ("global", "thread_local", "relaxed"):
+            raise ValueError(f"Unsupported build mode: {mode}")
+        if mode == "global":
+            capture_mode = driver.CUstreamCaptureMode.CU_STREAM_CAPTURE_MODE_GLOBAL
+        elif mode == "thread_local":
+            capture_mode = driver.CUstreamCaptureMode.CU_STREAM_CAPTURE_MODE_THREAD_LOCAL
+        elif mode == "relaxed":
+            capture_mode = driver.CUstreamCaptureMode.CU_STREAM_CAPTURE_MODE_RELAXED
+
         if self._mnff.conditional_graph:
             handle_return(
                 driver.cuStreamBeginCaptureToGraph(
@@ -201,15 +223,11 @@ def begin_building(self) -> GraphBuilder:
                     None,  # dependencies
                     None,  # dependencyData
                     0,  # numDependencies
-                    driver.CUstreamCaptureMode.CU_STREAM_CAPTURE_MODE_GLOBAL,
+                    capture_mode,
                 )
             )
         else:
-            handle_return(
-                driver.cuStreamBeginCapture(
-                    self._mnff.stream.handle, driver.CUstreamCaptureMode.CU_STREAM_CAPTURE_MODE_GLOBAL
-                )
-            )
+            handle_return(driver.cuStreamBeginCapture(self._mnff.stream.handle, capture_mode))
         return self
 
     @property

cuda_core/tests/test_graph.py

@@ -638,3 +638,30 @@ def test_graph_complete_options(init_cuda):
     gb.complete(options).close()
     options = CompleteOptions(use_node_priority=True)
     gb.complete(options).close()
+
+
+def test_graph_build_mode(init_cuda):
+    mod = _common_kernels()
+    empty_kernel = mod.get_kernel("empty_kernel")
+
+    # Simple linear topology
+    gb = Device().create_graph_builder().begin_building(mode="global")
+    launch(gb, LaunchConfig(grid=1, block=1), empty_kernel)
+    launch(gb, LaunchConfig(grid=1, block=1), empty_kernel)
+    launch(gb, LaunchConfig(grid=1, block=1), empty_kernel)
+    gb.end_building()
+
+    gb = Device().create_graph_builder().begin_building(mode="thread_local")
+    launch(gb, LaunchConfig(grid=1, block=1), empty_kernel)
+    launch(gb, LaunchConfig(grid=1, block=1), empty_kernel)
+    launch(gb, LaunchConfig(grid=1, block=1), empty_kernel)
+    gb.end_building()
+
+    gb = Device().create_graph_builder().begin_building(mode="relaxed")
+    launch(gb, LaunchConfig(grid=1, block=1), empty_kernel)
+    launch(gb, LaunchConfig(grid=1, block=1), empty_kernel)
+    launch(gb, LaunchConfig(grid=1, block=1), empty_kernel)
+    gb.end_building()
+
+    with pytest.raises(ValueError, match="^Unsupported build mode:"):
+        gb = Device().create_graph_builder().begin_building(mode=None)