Skip to content

Commit 644d69e

Browse files
cpcloudcursoragent
cpcloud
and
cursoragent
committed
feat(pathfinder): support finding static libraries like libcudadevrt.a (NVIDIA#716)
Add locate_static_lib() and find_static_lib() APIs to cuda.pathfinder for discovering native static libraries (e.g. libcudadevrt.a on Linux, cudadevrt.lib on Windows). Searches site-packages, conda, and CUDA_HOME/CUDA_PATH in that order, matching the existing bitcode lib pattern. Update cuda_core's device launch test to use the new API. Co-authored-by: Cursor <[email protected]>
1 parent 8d0ccdd commit 644d69e

4 files changed

Lines changed: 360 additions & 30 deletions

File tree

‎cuda_core/tests/graph/test_device_launch.py‎

Lines changed: 6 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,6 @@
1010
- The kernel calling cudaGraphLaunch() must itself be launched from within a graph
1111
"""
1212

13-
import os
14-
import sys
15-
1613
import numpy as np
1714
import pytest
1815
from cuda.core import (
@@ -29,30 +26,6 @@
2926
)
3027

3128

32-
def _find_cudadevrt_library():
33-
"""Find the CUDA device runtime static library using CUDA_HOME.
34-
35-
See https://github.com/NVIDIA/cuda-python/issues/716 for future improvements
36-
to make this discovery more robust via cuda.pathfinder.
37-
38-
Returns:
39-
Path to libcudadevrt.a (Linux) or cudadevrt.lib (Windows), or None if not found.
40-
"""
41-
cuda_home = os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH")
42-
if not cuda_home:
43-
return None
44-
45-
if sys.platform == "win32":
46-
path = os.path.join(cuda_home, "lib", "x64", "cudadevrt.lib")
47-
else:
48-
# Try lib64 first (common on Linux), fall back to lib
49-
path = os.path.join(cuda_home, "lib64", "libcudadevrt.a")
50-
if not os.path.isfile(path):
51-
path = os.path.join(cuda_home, "lib", "libcudadevrt.a")
52-
53-
return path if os.path.isfile(path) else None
54-
55-
5629
def _get_device_arch():
5730
"""Get the current device's architecture string."""
5831
return "".join(f"{i}" for i in Device().compute_capability)
@@ -80,9 +53,12 @@ def _compile_device_launcher_kernel():
8053
8154
Raises pytest.skip if libcudadevrt.a cannot be found.
8255
"""
83-
cudadevrt_path = _find_cudadevrt_library()
84-
if cudadevrt_path is None:
85-
pytest.skip("cudadevrt library not found (set CUDA_HOME or CUDA_PATH)")
56+
try:
57+
from cuda.pathfinder import StaticLibNotFoundError, find_static_lib
58+
59+
cudadevrt_path = find_static_lib("cudadevrt")
60+
except (ImportError, StaticLibNotFoundError):
61+
pytest.skip("cudadevrt library not found")
8662

8763
code = """
8864
extern "C" __global__ void launch_graph_from_device(cudaGraphExec_t graph) {

‎cuda_pathfinder/cuda/pathfinder/__init__.py‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,21 @@
4040
from cuda.pathfinder._static_libs.find_bitcode_lib import (
4141
locate_bitcode_lib as locate_bitcode_lib,
4242
)
43+
from cuda.pathfinder._static_libs.find_static_lib import (
44+
SUPPORTED_STATIC_LIBS as _SUPPORTED_STATIC_LIBS,
45+
)
46+
from cuda.pathfinder._static_libs.find_static_lib import (
47+
LocatedStaticLib as LocatedStaticLib,
48+
)
49+
from cuda.pathfinder._static_libs.find_static_lib import (
50+
StaticLibNotFoundError as StaticLibNotFoundError,
51+
)
52+
from cuda.pathfinder._static_libs.find_static_lib import (
53+
find_static_lib as find_static_lib,
54+
)
55+
from cuda.pathfinder._static_libs.find_static_lib import (
56+
locate_static_lib as locate_static_lib,
57+
)
4358

4459
from cuda.pathfinder._version import __version__ # isort: skip # noqa: F401
4560

@@ -61,6 +76,11 @@
6176
#: Example value: ``"device"``.
6277
SUPPORTED_BITCODE_LIBS = _SUPPORTED_BITCODE_LIBS
6378

79+
#: Tuple of supported static library names that can be resolved
80+
#: via ``locate_static_lib()`` and ``find_static_lib()``.
81+
#: Example value: ``"cudadevrt"``.
82+
SUPPORTED_STATIC_LIBS = _SUPPORTED_STATIC_LIBS
83+
6484
# Backward compatibility: _find_nvidia_header_directory was added in release 1.2.2.
6585
# It will be removed in release 1.2.4.
6686
_find_nvidia_header_directory = find_nvidia_header_directory

‎cuda_pathfinder/cuda/pathfinder/_static_libs/find_static_lib.py‎

Lines changed: 166 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,166 @@
1+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
4+
import functools
5+
import os
6+
from dataclasses import dataclass
7+
from typing import NoReturn, TypedDict
8+
9+
from cuda.pathfinder._utils.env_vars import get_cuda_home_or_path
10+
from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages
11+
from cuda.pathfinder._utils.platform_aware import IS_WINDOWS
12+
13+
14+
class StaticLibNotFoundError(RuntimeError):
15+
"""Raised when a static library cannot be found."""
16+
17+
18+
@dataclass(frozen=True)
19+
class LocatedStaticLib:
20+
"""Information about a located static library."""
21+
22+
name: str
23+
abs_path: str
24+
filename: str
25+
found_via: str
26+
27+
28+
class _StaticLibInfo(TypedDict):
29+
filename: str
30+
ctk_rel_paths: tuple[str, ...]
31+
conda_rel_path: str
32+
site_packages_dirs: tuple[str, ...]
33+
34+
35+
_SUPPORTED_STATIC_LIBS_INFO: dict[str, _StaticLibInfo] = {
36+
"cudadevrt": {
37+
"filename": "cudadevrt.lib" if IS_WINDOWS else "libcudadevrt.a",
38+
"ctk_rel_paths": (os.path.join("lib", "x64"),) if IS_WINDOWS else ("lib64", "lib"),
39+
"conda_rel_path": os.path.join("lib", "x64") if IS_WINDOWS else "lib",
40+
"site_packages_dirs": (
41+
("nvidia/cu13/lib/x64", "nvidia/cuda_runtime/lib/x64")
42+
if IS_WINDOWS
43+
else ("nvidia/cu13/lib", "nvidia/cuda_runtime/lib")
44+
),
45+
},
46+
}
47+
48+
SUPPORTED_STATIC_LIBS: tuple[str, ...] = tuple(sorted(_SUPPORTED_STATIC_LIBS_INFO.keys()))
49+
50+
51+
def _no_such_file_in_dir(dir_path: str, filename: str, error_messages: list[str], attachments: list[str]) -> None:
52+
error_messages.append(f"No such file: {os.path.join(dir_path, filename)}")
53+
if os.path.isdir(dir_path):
54+
attachments.append(f' listdir("{dir_path}"):')
55+
for node in sorted(os.listdir(dir_path)):
56+
attachments.append(f" {node}")
57+
else:
58+
attachments.append(f' Directory does not exist: "{dir_path}"')
59+
60+
61+
class _FindStaticLib:
62+
def __init__(self, name: str) -> None:
63+
if name not in _SUPPORTED_STATIC_LIBS_INFO:
64+
raise ValueError(f"Unknown static library: '{name}'. Supported: {', '.join(SUPPORTED_STATIC_LIBS)}")
65+
self.name: str = name
66+
self.config: _StaticLibInfo = _SUPPORTED_STATIC_LIBS_INFO[name]
67+
self.filename: str = self.config["filename"]
68+
self.ctk_rel_paths: tuple[str, ...] = self.config["ctk_rel_paths"]
69+
self.conda_rel_path: str = self.config["conda_rel_path"]
70+
self.site_packages_dirs: tuple[str, ...] = self.config["site_packages_dirs"]
71+
self.error_messages: list[str] = []
72+
self.attachments: list[str] = []
73+
74+
def try_site_packages(self) -> str | None:
75+
for rel_dir in self.site_packages_dirs:
76+
sub_dir = tuple(rel_dir.split("/"))
77+
for abs_dir in find_sub_dirs_all_sitepackages(sub_dir):
78+
file_path = os.path.join(abs_dir, self.filename)
79+
if os.path.isfile(file_path):
80+
return file_path
81+
return None
82+
83+
def try_with_conda_prefix(self) -> str | None:
84+
conda_prefix = os.environ.get("CONDA_PREFIX")
85+
if not conda_prefix:
86+
return None
87+
88+
anchor = os.path.join(conda_prefix, "Library") if IS_WINDOWS else conda_prefix
89+
file_path = os.path.join(anchor, self.conda_rel_path, self.filename)
90+
if os.path.isfile(file_path):
91+
return file_path
92+
return None
93+
94+
def try_with_cuda_home(self) -> str | None:
95+
cuda_home = get_cuda_home_or_path()
96+
if cuda_home is None:
97+
self.error_messages.append("CUDA_HOME/CUDA_PATH not set")
98+
return None
99+
100+
for rel_path in self.ctk_rel_paths:
101+
file_path = os.path.join(cuda_home, rel_path, self.filename)
102+
if os.path.isfile(file_path):
103+
return file_path
104+
105+
_no_such_file_in_dir(
106+
os.path.join(cuda_home, self.ctk_rel_paths[0]),
107+
self.filename,
108+
self.error_messages,
109+
self.attachments,
110+
)
111+
return None
112+
113+
def raise_not_found_error(self) -> NoReturn:
114+
err = ", ".join(self.error_messages) if self.error_messages else "No search paths available"
115+
att = "\n".join(self.attachments) if self.attachments else ""
116+
raise StaticLibNotFoundError(f'Failure finding "{self.filename}": {err}\n{att}')
117+
118+
119+
def locate_static_lib(name: str) -> LocatedStaticLib:
120+
"""Locate a static library by name.
121+
122+
Raises:
123+
ValueError: If ``name`` is not a supported static library.
124+
StaticLibNotFoundError: If the static library cannot be found.
125+
"""
126+
finder = _FindStaticLib(name)
127+
128+
abs_path = finder.try_site_packages()
129+
if abs_path is not None:
130+
return LocatedStaticLib(
131+
name=name,
132+
abs_path=abs_path,
133+
filename=finder.filename,
134+
found_via="site-packages",
135+
)
136+
137+
abs_path = finder.try_with_conda_prefix()
138+
if abs_path is not None:
139+
return LocatedStaticLib(
140+
name=name,
141+
abs_path=abs_path,
142+
filename=finder.filename,
143+
found_via="conda",
144+
)
145+
146+
abs_path = finder.try_with_cuda_home()
147+
if abs_path is not None:
148+
return LocatedStaticLib(
149+
name=name,
150+
abs_path=abs_path,
151+
filename=finder.filename,
152+
found_via="CUDA_HOME",
153+
)
154+
155+
finder.raise_not_found_error()
156+
157+
158+
@functools.cache
159+
def find_static_lib(name: str) -> str:
160+
"""Find the absolute path to a static library.
161+
162+
Raises:
163+
ValueError: If ``name`` is not a supported static library.
164+
StaticLibNotFoundError: If the static library cannot be found.
165+
"""
166+
return locate_static_lib(name).abs_path

0 commit comments

Comments
 (0)