Add cuDNN discovery and NCCL header support

.pre-commit-config.yaml

@@ -167,6 +167,8 @@ repos:
       - id: cython-lint
         args: [--no-pycodestyle]
         exclude: ^cuda_bindings/
+        additional_dependencies:
+          - Cython==3.2.9
 
 
 default_language_version:

cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py

@@ -61,12 +61,19 @@ class DescriptorSpec:
     packaged_with: PackagedWith
     linux_sonames: tuple[str, ...] = ()
     windows_dlls: tuple[str, ...] = ()
+    windows_dll_fallback_globs: tuple[str, ...] = ()
     supported_windows_arch: tuple[WindowsArch, ...] = ()
     site_packages_linux: tuple[str, ...] = ()
     site_packages_windows: WindowsSearchDirs = WindowsSearchDirs()
     dependencies: tuple[str, ...] = ()
+    optional_dependencies: tuple[str, ...] = ()
     anchor_rel_dirs_linux: tuple[str, ...] = ("lib64", "lib")
     anchor_rel_dirs_windows: WindowsSearchDirs = DEFAULT_WINDOWS_CTK_ANCHOR_DIRS
+    install_root_env_vars_linux: tuple[str, ...] = ()
+    install_root_env_rel_dirs_linux: tuple[str, ...] = ()
+    install_root_env_vars_windows: tuple[str, ...] = ()
+    install_root_env_rel_dirs_windows: WindowsSearchDirs = WindowsSearchDirs()
+    program_files_root_globs_windows: WindowsSearchDirs = WindowsSearchDirs()
     ctk_root_canary_anchor_libnames: tuple[str, ...] = ()
     requires_add_dll_directory: bool = False
     requires_rtld_deepbind: bool = False
@@ -353,23 +360,24 @@ class DescriptorSpec:
         packaged_with="ctk",
         linux_sonames=("libcupti.so.12", "libcupti.so.13"),
         windows_dlls=(
-            "cupti64_2026.3.0.dll",
-            "cupti64_2026.2.1.dll",
-            "cupti64_2026.2.0.dll",
-            "cupti64_2026.1.1.dll",
-            "cupti64_2026.1.0.dll",
-            "cupti64_2025.4.1.dll",
-            "cupti64_2025.3.1.dll",
-            "cupti64_2025.2.1.dll",
-            "cupti64_2025.1.1.dll",
-            "cupti64_2024.3.2.dll",
-            "cupti64_2024.2.1.dll",
-            "cupti64_2024.1.1.dll",
-            "cupti64_2023.3.1.dll",
-            "cupti64_2023.2.2.dll",
-            "cupti64_2023.1.1.dll",
             "cupti64_2022.4.1.dll",
+            "cupti64_2023.1.1.dll",
+            "cupti64_2023.2.2.dll",
+            "cupti64_2023.3.1.dll",
+            "cupti64_2024.1.1.dll",
+            "cupti64_2024.2.1.dll",
+            "cupti64_2024.3.2.dll",
+            "cupti64_2025.1.1.dll",
+            "cupti64_2025.2.1.dll",
+            "cupti64_2025.3.1.dll",
+            "cupti64_2025.4.1.dll",
+            "cupti64_2026.1.0.dll",
+            "cupti64_2026.1.1.dll",
+            "cupti64_2026.2.0.dll",
+            "cupti64_2026.2.1.dll",
+            "cupti64_2026.3.0.dll",
         ),
+        windows_dll_fallback_globs=("cupti64_*.dll",),
         supported_windows_arch=("x64", "arm64"),
         site_packages_linux=("nvidia/cu13/lib", "nvidia/cuda_cupti/lib"),
         site_packages_windows=_ctk_windows_wheel_dirs("nvidia/cu13/bin", "nvidia/cuda_cupti/bin"),
@@ -411,6 +419,29 @@ class DescriptorSpec:
         dependencies=("nvshmem_host",),
         requires_rtld_deepbind=True,
     ),
+    DescriptorSpec(
+        name="cudnn",
+        packaged_with="other",
+        linux_sonames=("libcudnn.so.9",),
+        windows_dlls=("cudnn64_9.dll",),
+        supported_windows_arch=("x64", "arm64"),
+        site_packages_linux=("nvidia/cudnn/lib",),
+        site_packages_windows=WindowsSearchDirs.x64_only("nvidia/cudnn/bin"),
+        dependencies=("cublasLt",),
+        optional_dependencies=("nvrtc",),
+        install_root_env_vars_linux=("CUDNN_PATH",),
+        install_root_env_rel_dirs_linux=("lib", "lib64"),
+        # The ARM64 layout is verified only for the standalone archive rooted
+        # at CUDNN_PATH, not for conda, CUDA_PATH, or Program Files installs.
+        anchor_rel_dirs_windows=WindowsSearchDirs.x64_only("bin/x64", "bin"),
+        install_root_env_vars_windows=("CUDNN_PATH",),
+        install_root_env_rel_dirs_windows=WindowsSearchDirs(
+            x64=("bin/x64", "bin"),
+            arm64=("bin/arm64",),
+        ),
+        program_files_root_globs_windows=WindowsSearchDirs.x64_only("NVIDIA/CUDNN/v9.*"),
+        requires_add_dll_directory=True,
+    ),
     DescriptorSpec(
         name="cusolverMp",
         packaged_with="other",
@@ -530,6 +561,8 @@ class DescriptorSpec:
         packaged_with="other",
         linux_sonames=("libnccl.so.2",),
         site_packages_linux=("nvidia/nccl/lib",),
+        install_root_env_vars_linux=("NCCL_HOME",),
+        install_root_env_rel_dirs_linux=("lib", "lib64", "build/lib"),
     ),
     DescriptorSpec(
         name="nvpl_fftw",

cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_dl_common.py

@@ -32,5 +32,20 @@ class LoadedDL:
 
 
 def load_dependencies(desc: LibDescriptor, load_func: Callable[[str], LoadedDL]) -> None:
+    """Load required dependencies, then best-effort runtime dependencies.
+
+    A plain ``DynamicLibNotFoundError`` from an optional dependency is
+    suppressed. More specific contract errors and failures while loading a
+    dependency that was found remain errors.
+    """
     for dep in desc.dependencies:
         load_func(dep)
+    for dep in desc.optional_dependencies:
+        try:
+            load_func(dep)
+        except DynamicLibNotFoundError as exc:
+            # Both public contract errors inherit DynamicLibNotFoundError, but
+            # neither an unknown descriptor nor platform incompatibility means
+            # that an optional runtime component is simply absent.
+            if type(exc) is not DynamicLibNotFoundError:
+                raise

cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_dl_linux.py

@@ -138,7 +138,7 @@ def _candidate_sonames(desc: LibDescriptor) -> list[str]:
 
 if sys.platform == "linux":
 
-    def check_if_already_loaded_from_elsewhere(desc: LibDescriptor, _have_abs_path: bool) -> LoadedDL | None:
+    def check_if_already_loaded_from_elsewhere(desc: LibDescriptor) -> LoadedDL | None:
         for soname in _candidate_sonames(desc):
             try:
                 handle = ctypes.CDLL(soname, mode=os.RTLD_NOLOAD)
@@ -160,7 +160,7 @@ def _load_lib(desc: LibDescriptor, filename: str) -> ctypes.CDLL:
         return ctypes.CDLL(filename, cdll_mode)
 else:
 
-    def check_if_already_loaded_from_elsewhere(_desc: LibDescriptor, _have_abs_path: bool) -> LoadedDL | None:
+    def check_if_already_loaded_from_elsewhere(_desc: LibDescriptor) -> LoadedDL | None:
         raise RuntimeError(f"check_if_already_loaded_from_elsewhere() is not supported on platform {sys.platform!r}")
 
     def _load_lib(_desc: LibDescriptor, _filename: str) -> ctypes.CDLL:

cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_dl_windows.py

@@ -120,15 +120,14 @@ def abs_path_for_dynamic_library(libname: str, handle: ctypes.wintypes.HMODULE)
     return buffer.value
 
 
-def check_if_already_loaded_from_elsewhere(desc: LibDescriptor, have_abs_path: bool) -> LoadedDL | None:
+def check_if_already_loaded_from_elsewhere(desc: LibDescriptor) -> LoadedDL | None:
     for dll_name in desc.windows_dlls:
         handle = kernel32.GetModuleHandleW(dll_name)
         if handle:
             abs_path = abs_path_for_dynamic_library(desc.name, handle)
-            if have_abs_path and desc.requires_add_dll_directory:
-                # This is a side-effect if the pathfinder loads the library via
-                # load_with_abs_path(). To make the side-effect more deterministic,
-                # activate it even if the library was already loaded from elsewhere.
+            if desc.requires_add_dll_directory:
+                # Match load_with_abs_path(): lazy component DLLs need the directory
+                # of the module that is actually loaded, regardless of how it arrived.
                 add_dll_directory(abs_path)
             return LoadedDL(abs_path, True, ctypes_handle_to_unsigned_int(handle), "was-already-loaded-from-elsewhere")
     return None

cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_nvidia_dynamic_lib.py

@@ -60,7 +60,7 @@ def _load_driver_lib_no_cache(desc: LibDescriptor) -> LoadedDL:
     native loader mechanisms, so the full CTK search cascade (site-packages,
     conda, CUDA_PATH, canary) is unnecessary.
     """
-    loaded = LOADER.check_if_already_loaded_from_elsewhere(desc, False)
+    loaded = LOADER.check_if_already_loaded_from_elsewhere(desc)
     if loaded is not None:
         return loaded
     loaded = LOADER.load_with_system_search(desc)
@@ -174,9 +174,7 @@ def _load_lib_no_cache(libname: str) -> LoadedDL:
     find = run_find_steps(ctx, EARLY_FIND_STEPS)
 
     # Phase 2: Cross-cutting — already-loaded check and dependency loading.
-    # The already-loaded check on Windows uses the "have we found a path?"
-    # flag to decide whether to apply AddDllDirectory side-effects.
-    loaded = LOADER.check_if_already_loaded_from_elsewhere(desc, find is not None)
+    loaded = LOADER.check_if_already_loaded_from_elsewhere(desc)
     load_dependencies(desc, load_nvidia_dynamic_lib)
     if loaded is not None:
         return loaded
@@ -275,12 +273,21 @@ def load_nvidia_dynamic_lib(libname: str) -> LoadedDL:
 
         4. **Environment variables**
 
-           - If set, use ``CUDA_PATH`` or ``CUDA_HOME`` (in that order).
-             On Windows, this is the typical way system-installed CTK DLLs are
-             located. Note that the NVIDIA CTK installer automatically
+           - First search library-specific roots declared by the descriptor,
+             such as ``CUDNN_PATH`` and ``NCCL_HOME``, using their
+             platform-specific product layouts. Then use ``CUDA_PATH`` or
+             ``CUDA_HOME`` (in that order).
+             On Windows, ``CUDA_PATH`` is the typical way system-installed CTK
+             DLLs are located. Note that the NVIDIA CTK installer automatically
              adds ``CUDA_PATH`` to the system-wide environment.
 
-        5. **CTK root canary probe (discoverable libs only)**
+        5. **Windows Program Files (configured libraries only)**
+
+           - Search descriptor-configured standalone installation roots, such
+             as versioned x64 cuDNN directories under ``ProgramFiles``, using
+             the general per-library anchor layout.
+
+        6. **CTK root canary probe (discoverable libs only)**
 
            - For selected libraries whose shared object doesn't reside on the
              standard linker path (currently ``nvvm``), attempt to derive CTK
@@ -298,8 +305,8 @@ def load_nvidia_dynamic_lib(libname: str) -> LoadedDL:
         0. Already loaded in the current process
         1. OS default mechanisms (``dlopen`` / ``LoadLibraryExW``)
 
-        The CTK-specific steps (site-packages, conda, ``CUDA_PATH``, canary
-        probe) are skipped entirely.
+        The non-driver steps (site-packages, conda, environment roots,
+        ``ProgramFiles``, and canary probe) are skipped entirely.
 
     Notes:
         The search is performed **per library**. There is currently no mechanism to

cuda_pathfinder/cuda/pathfinder/_dynamic_libs/platform_loader.py

@@ -24,7 +24,7 @@
 
 
 class PlatformLoader(Protocol):
-    def check_if_already_loaded_from_elsewhere(self, desc: LibDescriptor, have_abs_path: bool) -> LoadedDL | None: ...
+    def check_if_already_loaded_from_elsewhere(self, desc: LibDescriptor) -> LoadedDL | None: ...
 
     def load_with_system_search(self, desc: LibDescriptor) -> LoadedDL | None: ...
 

cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_platform.py

@@ -20,6 +20,7 @@
 from cuda.pathfinder._dynamic_libs.lib_descriptor import LibDescriptor
 from cuda.pathfinder._dynamic_libs.supported_nvidia_libs import is_suppressed_dll_file
 from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages
+from cuda.pathfinder._utils.path_sort import natural_path_sort_key
 from cuda.pathfinder._utils.platform_aware import IS_WINDOWS
 from cuda.pathfinder._utils.windows_arch import windows_pe_matches_arch, windows_python_arch
 
@@ -63,29 +64,48 @@ def _find_so_in_rel_dirs(
     return None
 
 
-def _find_dll_under_dir(dirpath: str, file_wild: str, target_arch: str | None = None) -> str | None:
-    for path in sorted(glob.glob(os.path.join(dirpath, file_wild))):
+def _find_descriptor_dll_under_dir(
+    dirpath: str,
+    desc: LibDescriptor,
+    target_arch: str | None = None,
+) -> str | None:
+    def candidate_is_usable(path: str) -> bool:
         if not os.path.isfile(path):
-            continue
+            return False
         if is_suppressed_dll_file(os.path.basename(path)):
-            continue
-        if target_arch is not None and not windows_pe_matches_arch(path, target_arch):
-            continue
-        return path
+            return False
+        return target_arch is None or windows_pe_matches_arch(path, target_arch)
+
+    # Prefer the descriptor's known DLL names in its established search order.
+    # Explicit globs provide a collision-safe forward-compatible fallback for
+    # libraries whose full version is encoded in the filename (for example CUPTI).
+    for dll_basename in reversed(cast(tuple[str, ...], desc.windows_dlls)):
+        path = os.path.join(dirpath, dll_basename)
+        if candidate_is_usable(path):
+            return path
+
+    for dll_glob in desc.windows_dll_fallback_globs:
+        file_wild = os.path.join(dirpath, dll_glob)
+        for path in sorted(glob.glob(file_wild), key=natural_path_sort_key, reverse=True):
+            if candidate_is_usable(path):
+                return path
     return None
 
 
 def _find_dll_in_rel_dirs(
     rel_dirs: tuple[str, ...],
+    desc: LibDescriptor,
+    target_arch: str,
     lib_searched_for: str,
     error_messages: list[str],
     attachments: list[str],
 ) -> str | None:
     sub_dirs_searched: list[tuple[str, ...]] = []
+    checked_arch = target_arch if desc.requires_windows_binary_arch_check else None
     for rel_dir in rel_dirs:
         sub_dir = PurePath(rel_dir).parts
         for abs_dir in find_sub_dirs_all_sitepackages(sub_dir):
-            dll_name = _find_dll_under_dir(abs_dir, lib_searched_for)
+            dll_name = _find_descriptor_dll_under_dir(abs_dir, desc, checked_arch)
             if dll_name is not None:
                 return dll_name
         sub_dirs_searched.append(sub_dir)
@@ -103,9 +123,16 @@ def conda_anchor_point(self, conda_prefix: str) -> str: ...
 
     def anchor_rel_dirs(self, desc: LibDescriptor) -> tuple[str, ...]: ...
 
+    def install_root_env_vars(self, desc: LibDescriptor) -> tuple[str, ...]: ...
+
+    def install_root_env_rel_dirs(self, desc: LibDescriptor) -> tuple[str, ...]: ...
+
+    def program_files_root_globs(self, desc: LibDescriptor) -> tuple[str, ...]: ...
+
     def find_in_site_packages(
         self,
         rel_dirs: tuple[str, ...],
+        desc: LibDescriptor,
         lib_searched_for: str,
         error_messages: list[str],
         attachments: list[str],
@@ -135,9 +162,19 @@ def conda_anchor_point(self, conda_prefix: str) -> str:
     def anchor_rel_dirs(self, desc: LibDescriptor) -> tuple[str, ...]:
         return cast(tuple[str, ...], desc.anchor_rel_dirs_linux)
 
+    def install_root_env_vars(self, desc: LibDescriptor) -> tuple[str, ...]:
+        return cast(tuple[str, ...], desc.install_root_env_vars_linux)
+
+    def install_root_env_rel_dirs(self, desc: LibDescriptor) -> tuple[str, ...]:
+        return cast(tuple[str, ...], desc.install_root_env_rel_dirs_linux)
+
+    def program_files_root_globs(self, _desc: LibDescriptor) -> tuple[str, ...]:
+        return ()
+
     def find_in_site_packages(
         self,
         rel_dirs: tuple[str, ...],
+        _desc: LibDescriptor,
         lib_searched_for: str,
         error_messages: list[str],
         attachments: list[str],
@@ -192,14 +229,39 @@ def conda_anchor_point(self, conda_prefix: str) -> str:
     def anchor_rel_dirs(self, desc: LibDescriptor) -> tuple[str, ...]:
         return cast(tuple[str, ...], desc.anchor_rel_dirs_windows.for_arch(self.target_arch))
 
+    def install_root_env_vars(self, desc: LibDescriptor) -> tuple[str, ...]:
+        if self.target_arch not in desc.supported_windows_arch:
+            return ()
+        return cast(tuple[str, ...], desc.install_root_env_vars_windows)
+
+    def install_root_env_rel_dirs(self, desc: LibDescriptor) -> tuple[str, ...]:
+        if self.target_arch not in desc.supported_windows_arch:
+            return ()
+        return cast(tuple[str, ...], desc.install_root_env_rel_dirs_windows.for_arch(self.target_arch))
+
+    def program_files_root_globs(self, desc: LibDescriptor) -> tuple[str, ...]:
+        program_files = os.environ.get("PROGRAMW6432") or os.environ.get("PROGRAMFILES")
+        if not program_files:
+            return ()
+        rel_globs = desc.program_files_root_globs_windows.for_arch(self.target_arch)
+        return tuple(os.path.join(program_files, rel_glob) for rel_glob in rel_globs)
+
     def find_in_site_packages(
         self,
         rel_dirs: tuple[str, ...],
+        desc: LibDescriptor,
         lib_searched_for: str,
         error_messages: list[str],
         attachments: list[str],
     ) -> str | None:
-        return _find_dll_in_rel_dirs(rel_dirs, lib_searched_for, error_messages, attachments)
+        return _find_dll_in_rel_dirs(
+            rel_dirs,
+            desc,
+            self.target_arch,
+            lib_searched_for,
+            error_messages,
+            attachments,
+        )
 
     def find_in_lib_dir(
         self,
@@ -211,7 +273,7 @@ def find_in_lib_dir(
     ) -> str | None:
         file_wild = desc.name + "*.dll"
         target_arch = self.target_arch if desc.requires_windows_binary_arch_check else None
-        dll_name = _find_dll_under_dir(lib_dir, file_wild, target_arch)
+        dll_name = _find_descriptor_dll_under_dir(lib_dir, desc, target_arch)
         if dll_name is not None:
             return dll_name
         if target_arch is None: