Skip to content

Commit 1633628

Browse files
Andy-JostAndy-Jost
committed
Reorganize graph test files for clarity
Rename test files to reflect what they actually test: - test_basic -> test_graph_builder (stream capture tests) - test_conditional -> test_graph_builder_conditional - test_advanced -> test_graph_update (moved child_graph and stream_lifetime tests into test_graph_builder) - test_capture_alloc -> test_graph_memory_resource - test_explicit* -> test_graphdef* Made-with: Cursor
1 parent d427c9b commit 1633628

9 files changed

Lines changed: 93 additions & 138 deletions

‎cuda_core/tests/graph/test_device_launch.py‎

Lines changed: 1 addition & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,7 @@
11
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
22
# SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
33

4-
"""Device-side graph launch tests.
5-
6-
Device-side graph launch allows a kernel running on the GPU to launch a CUDA graph.
7-
This feature requires:
8-
- CUDA 12.0+
9-
- Hopper architecture (sm_90+)
10-
- The kernel calling cudaGraphLaunch() must itself be launched from within a graph
11-
"""
4+
"""Tests for device-side graph launch (GPU kernel launching a CUDA graph)."""
125

136
import numpy as np
147
import pytest

cuda_core/tests/graph/test_basic.py renamed to cuda_core/tests/graph/test_graph_builder.py

Lines changed: 83 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
22
# SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
33

4-
"""Basic graph construction and topology tests."""
4+
"""GraphBuilder stream capture tests."""
55

66
import numpy as np
77
import pytest
@@ -205,3 +205,85 @@ def read_byte(data):
205205
launch_stream.sync()
206206

207207
assert result[0] == 0xAB
208+
209+
210+
@pytest.mark.skipif(tuple(int(i) for i in np.__version__.split(".")[:2]) < (2, 1), reason="need numpy 2.1.0+")
211+
def test_graph_child_graph(init_cuda):
212+
mod = compile_common_kernels()
213+
add_one = mod.get_kernel("add_one")
214+
215+
# Allocate memory
216+
launch_stream = Device().create_stream()
217+
mr = LegacyPinnedMemoryResource()
218+
b = mr.allocate(8)
219+
arr = np.from_dlpack(b).view(np.int32)
220+
arr[0] = 0
221+
arr[1] = 0
222+
223+
# Capture the child graph
224+
gb_child = Device().create_graph_builder().begin_building()
225+
launch(gb_child, LaunchConfig(grid=1, block=1), add_one, arr[1:].ctypes.data)
226+
launch(gb_child, LaunchConfig(grid=1, block=1), add_one, arr[1:].ctypes.data)
227+
launch(gb_child, LaunchConfig(grid=1, block=1), add_one, arr[1:].ctypes.data)
228+
gb_child.end_building()
229+
230+
# Capture the parent graph
231+
gb_parent = Device().create_graph_builder().begin_building()
232+
launch(gb_parent, LaunchConfig(grid=1, block=1), add_one, arr.ctypes.data)
233+
234+
## Add child
235+
try:
236+
gb_parent.add_child(gb_child)
237+
except NotImplementedError as e:
238+
with pytest.raises(
239+
NotImplementedError,
240+
match="^Launching child graphs is not implemented for versions older than CUDA 12",
241+
):
242+
raise e
243+
gb_parent.end_building()
244+
b.close()
245+
pytest.skip("Launching child graphs is not implemented for versions older than CUDA 12")
246+
247+
launch(gb_parent, LaunchConfig(grid=1, block=1), add_one, arr.ctypes.data)
248+
graph = gb_parent.end_building().complete()
249+
250+
# Parent updates first value, child updates second value
251+
assert arr[0] == 0
252+
assert arr[1] == 0
253+
graph.launch(launch_stream)
254+
launch_stream.sync()
255+
assert arr[0] == 2
256+
assert arr[1] == 3
257+
258+
b.close()
259+
260+
261+
def test_graph_stream_lifetime(init_cuda):
262+
mod = compile_common_kernels()
263+
empty_kernel = mod.get_kernel("empty_kernel")
264+
265+
# Create simple graph from device
266+
gb = Device().create_graph_builder().begin_building()
267+
launch(gb, LaunchConfig(grid=1, block=1), empty_kernel)
268+
graph = gb.end_building().complete()
269+
270+
# Destroy simple graph and builder
271+
gb.close()
272+
graph.close()
273+
274+
# Create simple graph from stream
275+
stream = Device().create_stream()
276+
gb = stream.create_graph_builder().begin_building()
277+
launch(gb, LaunchConfig(grid=1, block=1), empty_kernel)
278+
graph = gb.end_building().complete()
279+
280+
# Destroy simple graph and builder
281+
gb.close()
282+
graph.close()
283+
284+
# Verify the stream can still launch work
285+
launch(stream, LaunchConfig(grid=1, block=1), empty_kernel)
286+
stream.sync()
287+
288+
# Destroy the stream
289+
stream.close()

cuda_core/tests/graph/test_conditional.py renamed to cuda_core/tests/graph/test_graph_builder_conditional.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
22
# SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
33

4-
"""Conditional graph node tests (if, if-else, switch, while)."""
4+
"""Tests for GraphBuilder conditional node capture (if, if-else, switch, while)."""
55

66
import ctypes
77

cuda_core/tests/graph/test_capture_alloc.py renamed to cuda_core/tests/graph/test_graph_memory_resource.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
#
33
# SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
44

5-
"""Graph memory resource tests."""
5+
"""Tests for GraphMemoryResource allocation and attributes during graph capture."""
66

77
import pytest
88
from helpers import IS_WINDOWS, IS_WSL

cuda_core/tests/graph/test_advanced.py renamed to cuda_core/tests/graph/test_graph_update.py

Lines changed: 3 additions & 89 deletions
Original file line numberDiff line numberDiff line change
@@ -1,68 +1,15 @@
1-
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
1+
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
22
# SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
33

4-
"""Advanced graph feature tests (child graphs, update, stream lifetime)."""
4+
"""Tests for whole-graph update (Graph.update)."""
55

66
import numpy as np
77
import pytest
8-
from helpers.graph_kernels import compile_common_kernels, compile_conditional_kernels
8+
from helpers.graph_kernels import compile_conditional_kernels
99

1010
from cuda.core import Device, LaunchConfig, LegacyPinnedMemoryResource, launch
1111

1212

13-
@pytest.mark.skipif(tuple(int(i) for i in np.__version__.split(".")[:2]) < (2, 1), reason="need numpy 2.1.0+")
14-
def test_graph_child_graph(init_cuda):
15-
mod = compile_common_kernels()
16-
add_one = mod.get_kernel("add_one")
17-
18-
# Allocate memory
19-
launch_stream = Device().create_stream()
20-
mr = LegacyPinnedMemoryResource()
21-
b = mr.allocate(8)
22-
arr = np.from_dlpack(b).view(np.int32)
23-
arr[0] = 0
24-
arr[1] = 0
25-
26-
# Capture the child graph
27-
gb_child = Device().create_graph_builder().begin_building()
28-
launch(gb_child, LaunchConfig(grid=1, block=1), add_one, arr[1:].ctypes.data)
29-
launch(gb_child, LaunchConfig(grid=1, block=1), add_one, arr[1:].ctypes.data)
30-
launch(gb_child, LaunchConfig(grid=1, block=1), add_one, arr[1:].ctypes.data)
31-
gb_child.end_building()
32-
33-
# Capture the parent graph
34-
gb_parent = Device().create_graph_builder().begin_building()
35-
launch(gb_parent, LaunchConfig(grid=1, block=1), add_one, arr.ctypes.data)
36-
37-
## Add child
38-
try:
39-
gb_parent.add_child(gb_child)
40-
except NotImplementedError as e:
41-
with pytest.raises(
42-
NotImplementedError,
43-
match="^Launching child graphs is not implemented for versions older than CUDA 12",
44-
):
45-
raise e
46-
gb_parent.end_building()
47-
b.close()
48-
pytest.skip("Launching child graphs is not implemented for versions older than CUDA 12")
49-
50-
launch(gb_parent, LaunchConfig(grid=1, block=1), add_one, arr.ctypes.data)
51-
graph = gb_parent.end_building().complete()
52-
53-
# Parent updates first value, child updates second value
54-
assert arr[0] == 0
55-
assert arr[1] == 0
56-
graph.launch(launch_stream)
57-
launch_stream.sync()
58-
assert arr[0] == 2
59-
assert arr[1] == 3
60-
61-
# Close the memory resource now because the garbage collected might
62-
# de-allocate it during the next graph builder process
63-
b.close()
64-
65-
6613
@pytest.mark.skipif(tuple(int(i) for i in np.__version__.split(".")[:2]) < (2, 1), reason="need numpy 2.1.0+")
6714
def test_graph_update(init_cuda):
6815
mod = compile_conditional_kernels(int)
@@ -151,37 +98,4 @@ def build_graph(condition_value):
15198
assert arr[1] == 3
15299
assert arr[2] == 3
153100

154-
# Close the memory resource now because the garbage collected might
155-
# de-allocate it during the next graph builder process
156101
b.close()
157-
158-
159-
def test_graph_stream_lifetime(init_cuda):
160-
mod = compile_common_kernels()
161-
empty_kernel = mod.get_kernel("empty_kernel")
162-
163-
# Create simple graph from device
164-
gb = Device().create_graph_builder().begin_building()
165-
launch(gb, LaunchConfig(grid=1, block=1), empty_kernel)
166-
graph = gb.end_building().complete()
167-
168-
# Destroy simple graph and builder
169-
gb.close()
170-
graph.close()
171-
172-
# Create simple graph from stream
173-
stream = Device().create_stream()
174-
gb = stream.create_graph_builder().begin_building()
175-
launch(gb, LaunchConfig(grid=1, block=1), empty_kernel)
176-
graph = gb.end_building().complete()
177-
178-
# Destroy simple graph and builder
179-
gb.close()
180-
graph.close()
181-
182-
# Verify the stream can still launch work
183-
launch(stream, LaunchConfig(grid=1, block=1), empty_kernel)
184-
stream.sync()
185-
186-
# Destroy the stream
187-
stream.close()

cuda_core/tests/graph/test_explicit.py renamed to cuda_core/tests/graph/test_graphdef.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
22
# SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
33

4-
"""Tests for explicit CUDA graph construction (GraphDef and GraphNode)."""
4+
"""Tests for GraphDef topology, node types, instantiation, and execution."""
55

66
from collections.abc import Callable
77
from dataclasses import dataclass, field

cuda_core/tests/graph/test_explicit_errors.py renamed to cuda_core/tests/graph/test_graphdef_errors.py

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,7 @@
11
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
22
# SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
33

4-
"""Tests for error handling, input validation, and edge cases in explicit graphs.
5-
6-
These tests verify that the explicit graph API properly validates inputs,
7-
raises appropriate exceptions for misuse, and handles boundary conditions
8-
correctly.
9-
"""
4+
"""Tests for GraphDef input validation, error handling, and edge cases."""
105

116
import ctypes
127

cuda_core/tests/graph/test_explicit_integration.py renamed to cuda_core/tests/graph/test_graphdef_integration.py

Lines changed: 1 addition & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,31 +1,7 @@
11
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
22
# SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
33

4-
"""Integration tests for explicit CUDA graph construction.
5-
6-
Three test scenarios exercise complementary subsets of node types:
7-
8-
test_heat_diffusion
9-
1D heat bar evolving toward steady state via finite differences.
10-
Exercises: AllocNode, FreeNode, MemsetNode, ChildGraphNode,
11-
EmptyNode, EventRecordNode, EventWaitNode, WhileNode, KernelNode,
12-
MemcpyNode, HostCallbackNode.
13-
14-
test_bisection_root
15-
Find sqrt(2) by bisecting f(x) = x^2 - 2 on [0, 2], with an
16-
optional Newton polish step.
17-
Exercises: IfElseNode (interval halving), IfNode (refinement
18-
guard), WhileNode, KernelNode, AllocNode, MemsetNode, MemcpyNode,
19-
HostCallbackNode, FreeNode, EmptyNode.
20-
21-
test_switch_dispatch
22-
Apply one of four element-wise transforms selected at graph
23-
creation time via a switch condition.
24-
Exercises: SwitchNode, KernelNode, AllocNode, MemsetNode,
25-
MemcpyNode, FreeNode.
26-
27-
Together the three tests cover all 14 explicit-graph node types.
28-
"""
4+
"""End-to-end integration tests exercising all GraphDef node types in realistic scenarios."""
295

306
import ctypes
317

cuda_core/tests/graph/test_explicit_lifetime.py renamed to cuda_core/tests/graph/test_graphdef_lifetime.py

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,7 @@
11
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
22
# SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
33

4-
"""Tests for resource lifetime management in explicit CUDA graphs.
5-
6-
These tests verify that the RAII mechanism in GraphHandle correctly
7-
prevents dangling references when parent Python objects are deleted
8-
while child/body graph references remain alive.
9-
"""
4+
"""Tests for GraphDef resource lifetime management and RAII correctness."""
105

116
import gc
127

0 commit comments

Comments
 (0)