Reorganize graph tests and add device-side launch support #1512

Andy-Jost · 2026-01-16T21:55:37Z

Summary

This PR reorganizes the graph tests into a new tests/graph/ subdirectory, adds a public Graph.handle property, and introduces tests for device-side graph launch.

Closes #1331

Changes

Test Reorganization

Create new tests/graph/ subdirectory following the pattern of tests/memory_ipc/
Break apart test_graph.py into logical groupings:
- test_basic.py: basic graph construction and topology tests
- test_conditional.py: conditional node tests (if, if-else, switch, while)
- test_advanced.py: child graphs, update, stream lifetime
- test_options.py: debug print, complete options, build mode
Incorporate test_graph_mem.py (renamed to test_capture_alloc.py)
Add shared kernel compilation helpers in tests/helpers/graph_kernels.py

API Addition

Add Graph.handle property to expose the underlying CUgraphExec handle
Follows the existing pattern of Stream.handle and Event.handle
Forward-compatible with future RAII handle refactoring

Device-Side Graph Launch Tests

Add test_device_launch.py with tests for device-side graph launch
Uses cudaGraphLaunch() from device code via cudaStreamGraphTailLaunch
Requires Hopper (sm_90+) architecture
Uses cuda.pathfinder.find_nvidia_header_directory() to locate libcudadevrt.a

Test Coverage

58 tests total in tests/graph/
All tests pass on Hopper (H200)
Device launch tests skip on pre-Hopper architectures

copy-pr-bot · 2026-01-16T21:55:40Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

Andy-Jost · 2026-01-16T21:55:49Z

/ok to test 6e30c53

Andy-Jost · 2026-01-16T21:58:23Z

cuda_core/tests/graph/test_advanced.py

Moved from test_graph.py

Andy-Jost · 2026-01-16T21:58:35Z

cuda_core/tests/graph/test_basic.py

Moved from test_graph.py

Andy-Jost · 2026-01-16T21:59:00Z

cuda_core/tests/graph/test_conditional.py

Moved from test_graph.py

Andy-Jost · 2026-01-16T21:59:12Z

cuda_core/tests/graph/test_device_launch.py

Andy-Jost · 2026-01-16T21:59:25Z

cuda_core/tests/graph/test_options.py

Moved from test_graph.py

Andy-Jost · 2026-01-16T21:59:38Z

cuda_core/tests/helpers/graph_kernels.py

Moved from test_graph.py

Andy-Jost · 2026-01-16T22:02:33Z

/ok to test 5c63733

- Reorganize graph tests into tests/graph/ subdirectory: - test_basic.py: basic graph construction tests - test_conditional.py: conditional node tests (if, if-else, switch, while) - test_advanced.py: child graphs, update, stream lifetime - test_options.py: debug print, complete options, build mode - test_capture_alloc.py: graph memory resource tests (moved from test_graph_mem.py) - test_device_launch.py: new device-side graph launch tests - Add Graph.handle property to expose CUgraphExec handle, consistent with Stream.handle and Event.handle patterns - Add tests/helpers/graph_kernels.py with shared kernel compilation helpers - Device-side graph launch tests require Hopper (sm_90+) architecture

Andy-Jost · 2026-01-16T22:06:10Z

/ok to test da6205c

Andy-Jost · 2026-01-16T22:07:30Z

cuda_core/tests/graph/test_device_launch.py

+def _find_cudadevrt_library():
+    """Find the CUDA device runtime static library using pathfinder.
+
+    Uses load_nvidia_dynamic_lib("cudart") to locate the CUDA runtime,
+    then derives the static library path from that location.
+
+    Returns:
+        Path to libcudadevrt.a (Linux) or cudadevrt.lib (Windows), or None if not found.
+    """
+    try:
+        loaded = load_nvidia_dynamic_lib("cudart")
+    except Exception:
+        return None
+
+    lib_dir = os.path.dirname(loaded.abs_path)
+
+    if IS_WINDOWS:
+        # Windows: .dll is in bin/, .lib is in lib/x64/
+        base = os.path.dirname(lib_dir)
+        path = os.path.join(base, "lib", "x64", "cudadevrt.lib")
+        if os.path.isfile(path):
+            return path
+    else:
+        # Linux: .so and .a are both in lib/
+        path = os.path.join(lib_dir, "libcudadevrt.a")
+        if os.path.isfile(path):
+            return path
+
+    return None


I'm not sure of the best way to locate the cudadevrt library (needed for device launch). I'm open to suggestions here.

We have the existing requirement that CUDA_HOME needs to be set when running the cuda-core tests: did you consider finding the .lib/.a files relative to CUDA_HOME? I'm guessing that's probably pretty simple.

In any case, could you please reference #716 from the helper function, and maybe add a comment under 716 to point back to the helper function, after this PR is merged?

If you want to keep the current implementation:

try: loaded = load_nvidia_dynamic_lib("cudart") except DynamicLibNotFoundError: return None

That's far less likely to mask bugs.

github-actions · 2026-01-16T22:16:13Z

Doc Preview CI
🚀 View preview at https://nvidia.github.io/cuda-python/pr-preview/pr-1512/
https://nvidia.github.io/cuda-python/pr-preview/pr-1512/cuda-core/
https://nvidia.github.io/cuda-python/pr-preview/pr-1512/cuda-bindings/
https://nvidia.github.io/cuda-python/pr-preview/pr-1512/cuda-pathfinder/
Preview will be ready when the GitHub Pages deployment is complete.

leofang · 2026-01-17T07:16:40Z

cuda_core/cuda/core/_graph.py

        self._mnff.close()

+    @property
+    def handle(self) -> driver.CUgraphExec:


How do we decide if we want to expose CUgraph or CUgraphExec? Both can find their use cases.

It's a great question that I tried to sidestep in this change. (I'm working on another change to add resource handles to the graph module.)

My understanding is that class GraphBuilder corresponds to CUgraph and class Graph corresponds to CUgraphExec.

Following the existing pattern, after moving to resource handles, we would update the internal names to clarify:

GraphBuilder._mnff.graph → GraphBuilder._mnff._h_graph (CUgraph)

Graph._mnff.graph → Graph._mnff._h_graph_exec (CUgraphExec)

And the property types would be:

GraphBuilder.handle -> driver.CUgraph

Graph.handle -> driver.CUgraphExec

FYI, I added this property to expose CUgraphExec for the following code in test_device_launch.py:

inner_graph_handle = int(inner_graph.handle) ... launch(gb_outer, LaunchConfig(grid=1, block=1), launcher_kernel, inner_graph_handle)

Andy-Jost added the cuda.core Everything related to the cuda.core module label Jan 16, 2026

Andy-Jost self-assigned this Jan 16, 2026

Andy-Jost commented Jan 16, 2026

View reviewed changes

Andy-Jost force-pushed the graph-tests-reorg branch from 6e30c53 to 5c63733 Compare January 16, 2026 22:02

Andy-Jost force-pushed the graph-tests-reorg branch from 5c63733 to da6205c Compare January 16, 2026 22:06

Andy-Jost commented Jan 16, 2026

View reviewed changes

leofang reviewed Jan 17, 2026

View reviewed changes

Reorganize graph tests and add device-side launch support #1512

Are you sure you want to change the base?

Reorganize graph tests and add device-side launch support #1512

Uh oh!

Conversation

Andy-Jost commented Jan 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Changes

Test Reorganization

API Addition

Device-Side Graph Launch Tests

Test Coverage

Uh oh!

copy-pr-bot bot commented Jan 16, 2026

Uh oh!

Andy-Jost commented Jan 16, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Andy-Jost commented Jan 16, 2026

Uh oh!

Andy-Jost commented Jan 16, 2026

Uh oh!

Andy-Jost Jan 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented Jan 16, 2026

Preview will be ready when the GitHub Pages deployment is complete.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Andy-Jost commented Jan 16, 2026 •

edited

Loading

Andy-Jost Jan 16, 2026 •

edited

Loading