Skip to content

Commit 3d1eab9

Browse files
Add cuDNN discovery and NCCL header support
- cuDNN: added dynamic-library loading and header discovery. - NCCL: dynamic-library loading already existed; added the missing header discovery.
1 parent 819c586 commit 3d1eab9

7 files changed

Lines changed: 69 additions & 0 deletions

File tree

‎cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -411,6 +411,17 @@ class DescriptorSpec:
411411
dependencies=("nvshmem_host",),
412412
requires_rtld_deepbind=True,
413413
),
414+
DescriptorSpec(
415+
name="cudnn",
416+
packaged_with="other",
417+
linux_sonames=("libcudnn.so.9",),
418+
windows_dlls=("cudnn64_9.dll",),
419+
supported_windows_arch=("x64",),
420+
site_packages_linux=("nvidia/cudnn/lib",),
421+
site_packages_windows=WindowsSearchDirs.x64_only("nvidia/cudnn/bin"),
422+
anchor_rel_dirs_windows=WindowsSearchDirs.x64_only("bin/x64", "bin"),
423+
requires_add_dll_directory=True,
424+
),
414425
DescriptorSpec(
415426
name="cusolverMp",
416427
packaged_with="other",

‎cuda_pathfinder/cuda/pathfinder/_headers/header_descriptor_catalog.py‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -150,6 +150,14 @@ class HeaderDescriptorSpec:
150150
# -----------------------------------------------------------------------
151151
# Third-party / separately packaged headers
152152
# -----------------------------------------------------------------------
153+
HeaderDescriptorSpec(
154+
name="cudnn",
155+
packaged_with="other",
156+
header_basename="cudnn.h",
157+
site_packages_dirs=("nvidia/cudnn/include",),
158+
conda_targets_layout=False,
159+
use_ctk_root_canary=False,
160+
),
153161
HeaderDescriptorSpec(
154162
name="cusolverMp",
155163
packaged_with="other",
@@ -244,6 +252,15 @@ class HeaderDescriptorSpec:
244252
conda_targets_layout=False,
245253
use_ctk_root_canary=False,
246254
),
255+
HeaderDescriptorSpec(
256+
name="nccl",
257+
packaged_with="other",
258+
header_basename="nccl.h",
259+
site_packages_dirs=("nvidia/nccl/include",),
260+
available_on_windows=False,
261+
conda_targets_layout=False,
262+
use_ctk_root_canary=False,
263+
),
247264
HeaderDescriptorSpec(
248265
name="nvshmem",
249266
packaged_with="other",

‎cuda_pathfinder/pyproject.toml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ cu12 = [
2424
"cuquantum-cu12; sys_platform != 'win32'",
2525
"cutensor-cu12",
2626
"nvidia-cublasmp-cu12; sys_platform != 'win32'",
27+
"nvidia-cudnn-cu12",
2728
"nvidia-cudss-cu12",
2829
"nvidia-cufftmp-cu12; sys_platform != 'win32'",
2930
"nvidia-cusolvermp-cu12; sys_platform != 'win32'",
@@ -39,6 +40,7 @@ cu13 = [
3940
"cutensor-cu13",
4041
"nvidia-cublasmp-cu13; sys_platform != 'win32'",
4142
"nvidia-cudla; platform_system == 'Linux' and platform_machine == 'aarch64'",
43+
"nvidia-cudnn-cu13",
4244
"nvidia-cudss-cu13",
4345
"nvidia-cufftmp-cu13; sys_platform != 'win32'",
4446
"nvidia-cusolvermp-cu13; sys_platform != 'win32'",

‎cuda_pathfinder/tests/test_descriptor_catalog.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -115,6 +115,20 @@ def test_cusparselt_windows_metadata_matches_wheel_layouts():
115115
)
116116

117117

118+
@pytest.mark.agent_authored(model="gpt-5")
119+
def test_cudnn_metadata_matches_wheel_layouts():
120+
spec = _CATALOG_BY_NAME["cudnn"]
121+
assert spec.packaged_with == "other"
122+
assert spec.linux_sonames == ("libcudnn.so.9",)
123+
assert spec.windows_dlls == ("cudnn64_9.dll",)
124+
assert spec.supported_windows_arch == ("x64",)
125+
assert spec.site_packages_linux == ("nvidia/cudnn/lib",)
126+
assert spec.site_packages_windows == WindowsSearchDirs.x64_only("nvidia/cudnn/bin")
127+
assert spec.anchor_rel_dirs_windows == WindowsSearchDirs.x64_only("bin/x64", "bin")
128+
assert spec.dependencies == ()
129+
assert spec.requires_add_dll_directory
130+
131+
118132
@pytest.mark.parametrize("spec", DESCRIPTOR_CATALOG, ids=lambda s: s.name)
119133
def test_ctk_root_canary_anchors_reference_known_ctk_libs(spec: DescriptorSpec):
120134
for anchor in spec.ctk_root_canary_anchor_libnames:

‎cuda_pathfinder/tests/test_find_nvidia_headers.py‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
from cuda.pathfinder._dynamic_libs.load_nvidia_dynamic_lib import (
2828
_resolve_system_loaded_abs_path_in_subprocess,
2929
)
30+
from cuda.pathfinder._headers.header_descriptor import HEADER_DESCRIPTORS
3031
from cuda.pathfinder._headers.supported_nvidia_headers import (
3132
SUPPORTED_HEADERS_CTK,
3233
SUPPORTED_HEADERS_CTK_ALL,
@@ -43,6 +44,7 @@
4344

4445
NON_CTK_IMPORTLIB_METADATA_DISTRIBUTIONS_NAMES = {
4546
"cudensitymat": r"^cudensitymat-.*$",
47+
"cudnn": r"^nvidia-cudnn-.*$",
4648
"cupauliprop": r"^cupauliprop-.*$",
4749
"cusolverMp": r"^nvidia-cusolvermp-.*$",
4850
"cusparseLt": r"^nvidia-cusparselt-.*$",
@@ -53,6 +55,7 @@
5355
"custatevec": r"^custatevec-.*$",
5456
"cutlass": r"^nvidia-cutlass$",
5557
"mathdx": r"^nvidia-libmathdx-.*$",
58+
"nccl": r"^nvidia-nccl-.*$",
5659
"nvshmem": r"^nvidia-nvshmem-.*$",
5760
}
5861

@@ -78,6 +81,25 @@ def test_non_ctk_importlib_metadata_distributions_names():
7881
assert sorted(NON_CTK_IMPORTLIB_METADATA_DISTRIBUTIONS_NAMES) == sorted(SUPPORTED_HEADERS_NON_CTK_ALL)
7982

8083

84+
@pytest.mark.agent_authored(model="gpt-5")
85+
def test_cudnn_and_nccl_header_metadata_matches_wheel_layouts():
86+
cudnn = HEADER_DESCRIPTORS["cudnn"]
87+
assert cudnn.header_basename == "cudnn.h"
88+
assert cudnn.site_packages_dirs == ("nvidia/cudnn/include",)
89+
assert cudnn.available_on_linux
90+
assert cudnn.available_on_windows
91+
assert not cudnn.conda_targets_layout
92+
assert not cudnn.use_ctk_root_canary
93+
94+
nccl = HEADER_DESCRIPTORS["nccl"]
95+
assert nccl.header_basename == "nccl.h"
96+
assert nccl.site_packages_dirs == ("nvidia/nccl/include",)
97+
assert nccl.available_on_linux
98+
assert not nccl.available_on_windows
99+
assert not nccl.conda_targets_layout
100+
assert not nccl.use_ctk_root_canary
101+
102+
81103
@functools.cache
82104
def have_distribution_for(libname: str) -> bool:
83105
pattern = re.compile(NON_CTK_IMPORTLIB_METADATA_DISTRIBUTIONS_NAMES[libname])

‎toolshed/conda_create_for_pathfinder_testing.ps1‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ conda activate "pathfinder_testing_cu$CudaMajorMinorPatch"
1919
# Keep this list aligned with the Windows-installable subset of
2020
# cuda_pathfinder/pyproject.toml.
2121
$cpkgs = @(
22+
"cudnn",
2223
"cusparselt-dev",
2324
"cutensor",
2425
"cutlass",

‎toolshed/conda_create_for_pathfinder_testing.sh‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ set -u
2626
# cuda_pathfinder/pyproject.toml.
2727
cpkgs=(
2828
"cuquantum"
29+
"cudnn"
2930
"cusparselt-dev"
3031
"cutensor"
3132
"cutlass"
@@ -34,6 +35,7 @@ cpkgs=(
3435
"libcufftmp-dev"
3536
"libcusolvermp-dev"
3637
"libmathdx-dev"
38+
"nccl"
3739
"libnvshmem3"
3840
"libnvshmem-dev"
3941
)

0 commit comments

Comments
 (0)