Skip to content

Commit 7ccf852

Browse files
authored
cuda_core forward compatibility changes. (#722)
* cuda_core forward compatibility changes. * Add missing file. * Fix logic error computing CCCL_INCLUDE_PATHS in cuda_core/tests/helpers.py
1 parent 79e4bf7 commit 7ccf852

File tree

5 files changed

+43
-23
lines changed

5 files changed

+43
-23
lines changed

cuda_core/cuda/core/experimental/_graph.py

Lines changed: 17 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -476,7 +476,7 @@ def create_conditional_handle(self, default_value=None) -> driver.CUgraphConditi
476476
default_value = 0
477477
flags = 0
478478

479-
status, _, graph, _, _ = handle_return(driver.cuStreamGetCaptureInfo(self._mnff.stream.handle))
479+
status, _, graph, *_, _ = handle_return(driver.cuStreamGetCaptureInfo(self._mnff.stream.handle))
480480
if status != driver.CUstreamCaptureStatus.CU_STREAM_CAPTURE_STATUS_ACTIVE:
481481
raise RuntimeError("Cannot create a conditional handle when graph is not being built")
482482

@@ -486,20 +486,22 @@ def create_conditional_handle(self, default_value=None) -> driver.CUgraphConditi
486486

487487
def _cond_with_params(self, node_params) -> GraphBuilder:
488488
# Get current capture info to ensure we're in a valid state
489-
status, _, graph, dependencies, num_dependencies = handle_return(
489+
status, _, graph, *deps_info, num_dependencies = handle_return(
490490
driver.cuStreamGetCaptureInfo(self._mnff.stream.handle)
491491
)
492492
if status != driver.CUstreamCaptureStatus.CU_STREAM_CAPTURE_STATUS_ACTIVE:
493493
raise RuntimeError("Cannot add conditional node when not actively capturing")
494494

495495
# Add the conditional node to the graph
496-
node = handle_return(driver.cuGraphAddNode(graph, dependencies, num_dependencies, node_params))
496+
deps_info_update = [
497+
[handle_return(driver.cuGraphAddNode(graph, *deps_info, num_dependencies, node_params))]
498+
] + [None] * (len(deps_info) - 1)
497499

498500
# Update the stream's capture dependencies
499501
handle_return(
500502
driver.cuStreamUpdateCaptureDependencies(
501503
self._mnff.stream.handle,
502-
[node], # dependencies
504+
*deps_info_update, # dependencies, edgeData
503505
1, # numDependencies
504506
driver.CUstreamUpdateCaptureDependencies_flags.CU_STREAM_SET_CAPTURE_DEPENDENCIES,
505507
)
@@ -677,17 +679,23 @@ def add_child(self, child_graph: GraphBuilder):
677679
raise ValueError("Parent graph is not being built.")
678680

679681
stream_handle = self._mnff.stream.handle
680-
_, _, graph_out, dependencies_out, num_dependencies_out = handle_return(
682+
_, _, graph_out, *deps_info_out, num_dependencies_out = handle_return(
681683
driver.cuStreamGetCaptureInfo(stream_handle)
682684
)
683685

684-
child_node = handle_return(
685-
driver.cuGraphAddChildGraphNode(graph_out, dependencies_out, num_dependencies_out, child_graph._mnff.graph)
686-
)
686+
deps_info_update = [
687+
[
688+
handle_return(
689+
driver.cuGraphAddChildGraphNode(
690+
graph_out, deps_info_out[0], num_dependencies_out, child_graph._mnff.graph
691+
)
692+
)
693+
]
694+
] + [None] * (len(deps_info_out) - 1)
687695
handle_return(
688696
driver.cuStreamUpdateCaptureDependencies(
689697
stream_handle,
690-
[child_node],
698+
*deps_info_update, # dependencies, edgeData
691699
1,
692700
driver.CUstreamUpdateCaptureDependencies_flags.CU_STREAM_SET_CAPTURE_DEPENDENCIES,
693701
)

cuda_core/tests/conftest.py

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
# Copyright 2024 NVIDIA Corporation. All rights reserved.
22
# SPDX-License-Identifier: Apache-2.0
33

4-
import os
4+
import helpers
55

66
try:
77
from cuda.bindings import driver
@@ -65,5 +65,4 @@ def pop_all_contexts():
6565
return pop_all_contexts
6666

6767

68-
# TODO: make the fixture more sophisticated using path finder
69-
skipif_need_cuda_headers = pytest.mark.skipif(os.environ.get("CUDA_PATH") is None, reason="need CUDA header")
68+
skipif_need_cuda_headers = pytest.mark.skipif(helpers.CUDA_INCLUDE_PATH is None, reason="need CUDA header")

cuda_core/tests/helpers.py

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
# Copyright 2025 NVIDIA Corporation. All rights reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
4+
import os
5+
6+
CUDA_PATH = os.environ.get("CUDA_PATH")
7+
CUDA_INCLUDE_PATH = None
8+
CCCL_INCLUDE_PATHS = None
9+
if CUDA_PATH is not None:
10+
path = os.path.join(CUDA_PATH, "include")
11+
if os.path.isdir(path):
12+
CUDA_INCLUDE_PATH = path
13+
CCCL_INCLUDE_PATHS = (path,)
14+
path = os.path.join(path, "cccl")
15+
if os.path.isdir(path):
16+
CCCL_INCLUDE_PATHS = (path,) + CCCL_INCLUDE_PATHS

cuda_core/tests/test_event.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import platform
77
import time
88

9+
import helpers
910
import numpy as np
1011
import pytest
1112
from conftest import skipif_need_cuda_headers
@@ -149,7 +150,7 @@ def test_error_timing_incomplete():
149150
program_options = ProgramOptions(
150151
std="c++17",
151152
arch=f"sm_{arch}",
152-
include_path=str(pathlib.Path(os.environ["CUDA_PATH"]) / pathlib.Path("include")),
153+
include_path=helpers.CCCL_INCLUDE_PATHS,
153154
)
154155
prog = Program(code, code_type="c++", options=program_options)
155156
mod = prog.compile(target_type="cubin")

cuda_core/tests/test_launcher.py

Lines changed: 6 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,8 @@
22
# SPDX-License-Identifier: Apache-2.0
33

44
import ctypes
5-
import os
6-
import pathlib
5+
6+
import helpers
77

88
try:
99
import cupy as cp
@@ -107,7 +107,7 @@ def test_launch_invalid_values(init_cuda):
107107
(ctypes.c_float, "float", 3.14),
108108
(ctypes.c_double, "double", 2.718),
109109
)
110-
if os.environ.get("CUDA_PATH"):
110+
if helpers.CCCL_INCLUDE_PATHS is not None:
111111
PARAMS += (
112112
(np.float16, "half", 0.78),
113113
(np.complex64, "cuda::std::complex<float>", 1 + 2j),
@@ -141,18 +141,15 @@ def test_launch_scalar_argument(python_type, cpp_type, init_value):
141141

142142
# Compile and force instantiation for this type
143143
arch = "".join(f"{i}" for i in dev.compute_capability)
144-
if os.environ.get("CUDA_PATH"):
145-
include_path = str(pathlib.Path(os.environ["CUDA_PATH"]) / pathlib.Path("include"))
144+
if helpers.CCCL_INCLUDE_PATHS is not None:
146145
code = (
147146
r"""
148147
#include <cuda_fp16.h>
149148
#include <cuda/std/complex>
150149
"""
151150
+ code
152151
)
153-
else:
154-
include_path = None
155-
pro_opts = ProgramOptions(std="c++11", arch=f"sm_{arch}", include_path=include_path)
152+
pro_opts = ProgramOptions(std="c++17", arch=f"sm_{arch}", include_path=helpers.CCCL_INCLUDE_PATHS)
156153
prog = Program(code, code_type="c++", options=pro_opts)
157154
ker_name = f"write_scalar<{cpp_type}>"
158155
mod = prog.compile("cubin", name_expressions=(ker_name,))
@@ -186,8 +183,7 @@ def test_cooperative_launch():
186183

187184
# Compile and force instantiation for this type
188185
arch = "".join(f"{i}" for i in dev.compute_capability)
189-
include_path = str(pathlib.Path(os.environ["CUDA_PATH"]) / pathlib.Path("include"))
190-
pro_opts = ProgramOptions(std="c++17", arch=f"sm_{arch}", include_path=include_path)
186+
pro_opts = ProgramOptions(std="c++17", arch=f"sm_{arch}", include_path=helpers.CCCL_INCLUDE_PATHS)
191187
prog = Program(code, code_type="c++", options=pro_opts)
192188
ker = prog.compile("cubin").get_kernel("test_grid_sync")
193189

0 commit comments

Comments
 (0)