Make Windows pathfinder dynamic library searches architecture-aware

cuda_pathfinder/cuda/pathfinder/__init__.py

@@ -20,9 +20,7 @@
 )
 from cuda.pathfinder._dynamic_libs.load_dl_common import LoadedDL as LoadedDL
 from cuda.pathfinder._dynamic_libs.load_nvidia_dynamic_lib import load_nvidia_dynamic_lib as load_nvidia_dynamic_lib
-from cuda.pathfinder._dynamic_libs.supported_nvidia_libs import (
-    SUPPORTED_LIBNAMES as SUPPORTED_NVIDIA_LIBNAMES,
-)
+from cuda.pathfinder._dynamic_libs.supported_nvidia_libs import SUPPORTED_LIBNAMES as _SUPPORTED_NVIDIA_LIBNAMES
 from cuda.pathfinder._headers.find_nvidia_headers import LocatedHeaderDir as LocatedHeaderDir
 from cuda.pathfinder._headers.find_nvidia_headers import find_nvidia_header_directory as find_nvidia_header_directory
 from cuda.pathfinder._headers.find_nvidia_headers import (
@@ -60,6 +58,7 @@
     locate_static_lib as locate_static_lib,
 )
 from cuda.pathfinder._utils.env_vars import get_cuda_path_or_home as get_cuda_path_or_home
+from cuda.pathfinder._utils.windows_arch import UnsupportedArchError as UnsupportedArchError
 
 from cuda.pathfinder._version import __version__  # isort: skip
 
@@ -76,6 +75,11 @@
 #: Example utilities: ``"nvdisasm"``, ``"cuobjdump"``, ``"nvcc"``.
 SUPPORTED_BINARY_UTILITIES = _SUPPORTED_BINARIES
 
+#: Tuple of CUDA Toolkit dynamic library names supported by
+#: :func:`load_nvidia_dynamic_lib` for the current operating system and
+#: interpreter architecture.
+SUPPORTED_NVIDIA_LIBNAMES = _SUPPORTED_NVIDIA_LIBNAMES
+
 #: Tuple of supported bitcode library names that can be resolved
 #: via ``locate_bitcode_lib()`` and ``find_bitcode_lib()``.
 #: Example value: ``"device"``.

cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_nvidia_dynamic_lib.py

@@ -34,6 +34,7 @@
     build_dynamic_lib_subprocess_command,
     parse_dynamic_lib_subprocess_payload,
 )
+from cuda.pathfinder._dynamic_libs.supported_nvidia_libs import ALL_AVAILABLE_LIBNAMES
 from cuda.pathfinder._utils.platform_aware import IS_WINDOWS
 
 if TYPE_CHECKING:
@@ -42,9 +43,6 @@
 # All libnames recognized by load_nvidia_dynamic_lib, across all categories
 # (CTK, third-party, driver).
 _ALL_KNOWN_LIBNAMES: frozenset[str] = frozenset(LIB_DESCRIPTORS)
-_ALL_SUPPORTED_LIBNAMES: frozenset[str] = frozenset(
-    name for name, desc in LIB_DESCRIPTORS.items() if (desc.windows_dlls if IS_WINDOWS else desc.linux_sonames)
-)
 _PLATFORM_NAME = "Windows" if IS_WINDOWS else "Linux"
 _CANARY_PROBE_TIMEOUT_SECONDS = 10.0
 
@@ -308,9 +306,9 @@ def load_nvidia_dynamic_lib(libname: str) -> LoadedDL:
         )
     if libname not in _ALL_KNOWN_LIBNAMES:
         raise DynamicLibUnknownError(f"Unknown library name: {libname!r}. Known names: {sorted(_ALL_KNOWN_LIBNAMES)}")
-    if libname not in _ALL_SUPPORTED_LIBNAMES:
+    if libname not in ALL_AVAILABLE_LIBNAMES:
         raise DynamicLibNotAvailableError(
             f"Library name {libname!r} is known but not available on {_PLATFORM_NAME}. "
-            f"Supported names on {_PLATFORM_NAME}: {sorted(_ALL_SUPPORTED_LIBNAMES)}"
+            f"Supported names on {_PLATFORM_NAME}: {sorted(ALL_AVAILABLE_LIBNAMES)}"
         )
     return _load_lib_no_cache(libname)

cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_platform.py

@@ -14,12 +14,14 @@
 import os
 from collections.abc import Sequence
 from dataclasses import dataclass
+from pathlib import PurePath
 from typing import Protocol, cast
 
 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.platform_aware import IS_WINDOWS
+from cuda.pathfinder._utils.windows_arch import windows_pe_matches_arch, windows_python_arch
 
 
 def _no_such_file_in_sub_dirs(
@@ -41,7 +43,7 @@ def _find_so_in_rel_dirs(
     sub_dirs_searched: list[tuple[str, ...]] = []
     file_wild = so_basename + "*"
     for rel_dir in rel_dirs:
-        sub_dir = tuple(rel_dir.split(os.path.sep))
+        sub_dir = PurePath(rel_dir).parts
         for abs_dir in find_sub_dirs_all_sitepackages(sub_dir):
             # Exact unversioned match first; fall back to versioned names because some
             # distros only ship lib<name>.so.<major> (e.g. conda libcupti). Only one match
@@ -61,12 +63,15 @@ def _find_so_in_rel_dirs(
     return None
 
 
-def _find_dll_under_dir(dirpath: str, file_wild: str) -> str | 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))):
         if not os.path.isfile(path):
             continue
-        if not is_suppressed_dll_file(os.path.basename(path)):
-            return path
+        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 None
 
 
@@ -78,7 +83,7 @@ def _find_dll_in_rel_dirs(
 ) -> str | None:
     sub_dirs_searched: list[tuple[str, ...]] = []
     for rel_dir in rel_dirs:
-        sub_dir = tuple(rel_dir.split(os.path.sep))
+        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)
             if dll_name is not None:
@@ -109,7 +114,7 @@ def find_in_site_packages(
     def find_in_lib_dir(
         self,
         lib_dir: str,
-        libname: str,
+        desc: LibDescriptor,
         lib_searched_for: str,
         error_messages: list[str],
         attachments: list[str],
@@ -142,7 +147,7 @@ def find_in_site_packages(
     def find_in_lib_dir(
         self,
         lib_dir: str,
-        _libname: str,
+        _desc: LibDescriptor,
         lib_searched_for: str,
         error_messages: list[str],
         attachments: list[str],
@@ -173,17 +178,19 @@ def find_in_lib_dir(
 
 @dataclass(frozen=True, slots=True)
 class WindowsSearchPlatform:
+    target_arch: str
+
     def lib_searched_for(self, libname: str) -> str:
         return f"{libname}*.dll"
 
     def site_packages_rel_dirs(self, desc: LibDescriptor) -> tuple[str, ...]:
-        return cast(tuple[str, ...], desc.site_packages_windows)
+        return cast(tuple[str, ...], desc.site_packages_windows.for_arch(self.target_arch))
 
     def conda_anchor_point(self, conda_prefix: str) -> str:
         return os.path.join(conda_prefix, "Library")
 
     def anchor_rel_dirs(self, desc: LibDescriptor) -> tuple[str, ...]:
-        return cast(tuple[str, ...], desc.anchor_rel_dirs_windows)
+        return cast(tuple[str, ...], desc.anchor_rel_dirs_windows.for_arch(self.target_arch))
 
     def find_in_site_packages(
         self,
@@ -197,16 +204,20 @@ def find_in_site_packages(
     def find_in_lib_dir(
         self,
         lib_dir: str,
-        libname: str,
+        desc: LibDescriptor,
         _lib_searched_for: str,
         error_messages: list[str],
         attachments: list[str],
     ) -> str | None:
-        file_wild = libname + "*.dll"
-        dll_name = _find_dll_under_dir(lib_dir, file_wild)
+        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)
         if dll_name is not None:
             return dll_name
-        error_messages.append(f"No such file: {file_wild}")
+        if target_arch is None:
+            error_messages.append(f"No such file: {file_wild}")
+        else:
+            error_messages.append(f"No {target_arch}-compatible PE file: {file_wild}")
         attachments.append(f'  listdir("{lib_dir}"):')
         if not os.path.isdir(lib_dir):
             attachments.append("    DIRECTORY DOES NOT EXIST")
@@ -216,4 +227,10 @@ def find_in_lib_dir(
         return None
 
 
-PLATFORM: SearchPlatform = WindowsSearchPlatform() if IS_WINDOWS else LinuxSearchPlatform()
+def _platform_for_current_system() -> SearchPlatform:
+    if IS_WINDOWS:
+        return WindowsSearchPlatform(target_arch=windows_python_arch())
+    return LinuxSearchPlatform()
+
+
+PLATFORM = _platform_for_current_system()

cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_steps.py

@@ -88,7 +88,7 @@ def _find_using_lib_dir(ctx: SearchContext, lib_dir: str | None) -> str | None:
         str | None,
         ctx.platform.find_in_lib_dir(
             lib_dir,
-            ctx.libname,
+            ctx.desc,
             ctx.lib_searched_for,
             ctx.error_messages,
             ctx.attachments,
@@ -121,13 +121,14 @@ def _derive_ctk_root_windows(resolved_lib_path: str) -> str | None:
 
     Supports:
     - ``$CTK_ROOT/bin/x64/foo.dll`` (CTK 13 style)
+    - ``$CTK_ROOT/bin/arm64/foo.dll`` (Windows on Arm CTK 13 style)
     - ``$CTK_ROOT/bin/foo.dll`` (CTK 12 style)
     """
     import ntpath
 
     lib_dir = ntpath.dirname(resolved_lib_path)
     basename = ntpath.basename(lib_dir).lower()
-    if basename == "x64":
+    if basename in ("x64", "arm64"):
         parent = ntpath.dirname(lib_dir)
         if ntpath.basename(parent).lower() == "bin":
             return ntpath.dirname(parent)

cuda_pathfinder/cuda/pathfinder/_dynamic_libs/supported_nvidia_libs.py

@@ -6,12 +6,18 @@
 The canonical data entry point is :mod:`descriptor_catalog`. This module keeps
 historical constant names for backward compatibility by deriving them from the
 catalog.
+
+The unsuffixed ``SUPPORTED_LIBNAMES_WINDOWS`` and
+``SITE_PACKAGES_LIBDIRS_WINDOWS*`` constants retain their historical x64
+meaning for compatibility, but are not recommended for new code. Use the
+explicit ``*_X64`` or ``*_ARM64`` projection instead. Never combine the two
+architecture projections.
 """
 
 from __future__ import annotations
 
 from cuda.pathfinder._dynamic_libs.descriptor_catalog import DESCRIPTOR_CATALOG
-from cuda.pathfinder._utils.platform_aware import IS_WINDOWS
+from cuda.pathfinder._utils.platform_aware import IS_WINDOWS, IS_WINDOWS_ARM64, IS_WINDOWS_X64
 
 _CTK_DESCRIPTORS = tuple(desc for desc in DESCRIPTOR_CATALOG if desc.packaged_with == "ctk")
 _OTHER_DESCRIPTORS = tuple(desc for desc in DESCRIPTOR_CATALOG if desc.packaged_with == "other")
@@ -26,10 +32,30 @@
     desc.name for desc in _CTK_DESCRIPTORS if desc.windows_dlls and not desc.linux_sonames
 )
 
+if not IS_WINDOWS:
+    ALL_AVAILABLE_LIBNAMES = frozenset(desc.name for desc in DESCRIPTOR_CATALOG if desc.linux_sonames)
+else:
+    assert IS_WINDOWS_X64 != IS_WINDOWS_ARM64
+    _current_windows_arch = "x64" if IS_WINDOWS_X64 else "arm64"
+    ALL_AVAILABLE_LIBNAMES = frozenset(
+        desc.name for desc in DESCRIPTOR_CATALOG if _current_windows_arch in desc.supported_windows_arch
+    )
+
 SUPPORTED_LIBNAMES_LINUX = SUPPORTED_LIBNAMES_COMMON + SUPPORTED_LIBNAMES_LINUX_ONLY
-SUPPORTED_LIBNAMES_WINDOWS = SUPPORTED_LIBNAMES_COMMON + SUPPORTED_LIBNAMES_WINDOWS_ONLY
+SUPPORTED_LIBNAMES_WINDOWS_X64 = tuple(desc.name for desc in _CTK_DESCRIPTORS if "x64" in desc.supported_windows_arch)
+SUPPORTED_LIBNAMES_WINDOWS_ARM64 = tuple(
+    desc.name for desc in _CTK_DESCRIPTORS if "arm64" in desc.supported_windows_arch
+)
+# Backward-compatible alias preserves the historical x64 meaning.
+SUPPORTED_LIBNAMES_WINDOWS = SUPPORTED_LIBNAMES_WINDOWS_X64
 SUPPORTED_LIBNAMES_ALL = SUPPORTED_LIBNAMES_COMMON + SUPPORTED_LIBNAMES_LINUX_ONLY + SUPPORTED_LIBNAMES_WINDOWS_ONLY
-SUPPORTED_LIBNAMES = SUPPORTED_LIBNAMES_WINDOWS if IS_WINDOWS else SUPPORTED_LIBNAMES_LINUX
+if not IS_WINDOWS:
+    SUPPORTED_LIBNAMES = SUPPORTED_LIBNAMES_LINUX
+elif IS_WINDOWS_X64:
+    SUPPORTED_LIBNAMES = SUPPORTED_LIBNAMES_WINDOWS_X64
+else:
+    assert IS_WINDOWS_ARM64
+    SUPPORTED_LIBNAMES = SUPPORTED_LIBNAMES_WINDOWS_ARM64
 
 DIRECT_DEPENDENCIES_CTK = {desc.name: desc.dependencies for desc in _CTK_DESCRIPTORS if desc.dependencies}
 DIRECT_DEPENDENCIES = {desc.name: desc.dependencies for desc in DESCRIPTOR_CATALOG if desc.dependencies}
@@ -51,7 +77,6 @@
     desc.name for desc in DESCRIPTOR_CATALOG if desc.requires_rtld_deepbind and desc.linux_sonames
 )
 
-# Based on output of toolshed/make_site_packages_libdirs_linux.py
 SITE_PACKAGES_LIBDIRS_LINUX_CTK = {
     desc.name: desc.site_packages_linux for desc in _CTK_DESCRIPTORS if desc.site_packages_linux
 }
@@ -60,13 +85,29 @@
 }
 SITE_PACKAGES_LIBDIRS_LINUX = SITE_PACKAGES_LIBDIRS_LINUX_CTK | SITE_PACKAGES_LIBDIRS_LINUX_OTHER
 
-SITE_PACKAGES_LIBDIRS_WINDOWS_CTK = {
-    desc.name: desc.site_packages_windows for desc in _CTK_DESCRIPTORS if desc.site_packages_windows
+# Architecture-specific Windows projections. Keep these separate: combining
+# them would make the table unsafe to consume for either process ABI.
+SITE_PACKAGES_LIBDIRS_WINDOWS_CTK_X64 = {
+    desc.name: desc.site_packages_windows.x64 for desc in _CTK_DESCRIPTORS if desc.site_packages_windows.x64
+}
+SITE_PACKAGES_LIBDIRS_WINDOWS_CTK_ARM64 = {
+    desc.name: desc.site_packages_windows.arm64 for desc in _CTK_DESCRIPTORS if desc.site_packages_windows.arm64
 }
-SITE_PACKAGES_LIBDIRS_WINDOWS_OTHER = {
-    desc.name: desc.site_packages_windows for desc in _NON_CTK_DESCRIPTORS if desc.site_packages_windows
+SITE_PACKAGES_LIBDIRS_WINDOWS_OTHER_X64 = {
+    desc.name: desc.site_packages_windows.x64 for desc in _NON_CTK_DESCRIPTORS if desc.site_packages_windows.x64
 }
-SITE_PACKAGES_LIBDIRS_WINDOWS = SITE_PACKAGES_LIBDIRS_WINDOWS_CTK | SITE_PACKAGES_LIBDIRS_WINDOWS_OTHER
+SITE_PACKAGES_LIBDIRS_WINDOWS_OTHER_ARM64 = {
+    desc.name: desc.site_packages_windows.arm64 for desc in _NON_CTK_DESCRIPTORS if desc.site_packages_windows.arm64
+}
+SITE_PACKAGES_LIBDIRS_WINDOWS_X64 = SITE_PACKAGES_LIBDIRS_WINDOWS_CTK_X64 | SITE_PACKAGES_LIBDIRS_WINDOWS_OTHER_X64
+SITE_PACKAGES_LIBDIRS_WINDOWS_ARM64 = (
+    SITE_PACKAGES_LIBDIRS_WINDOWS_CTK_ARM64 | SITE_PACKAGES_LIBDIRS_WINDOWS_OTHER_ARM64
+)
+
+# Backward-compatible aliases preserve the historical x64 meaning.
+SITE_PACKAGES_LIBDIRS_WINDOWS_CTK = SITE_PACKAGES_LIBDIRS_WINDOWS_CTK_X64
+SITE_PACKAGES_LIBDIRS_WINDOWS_OTHER = SITE_PACKAGES_LIBDIRS_WINDOWS_OTHER_X64
+SITE_PACKAGES_LIBDIRS_WINDOWS = SITE_PACKAGES_LIBDIRS_WINDOWS_X64
 
 
 def is_suppressed_dll_file(path_basename: str) -> bool:

cuda_pathfinder/cuda/pathfinder/_utils/platform_aware.py

@@ -4,6 +4,18 @@
 import sys
 
 IS_WINDOWS = sys.platform == "win32"
+_WINDOWS_PYTHON_ARCH: str | None
+
+if IS_WINDOWS:
+    from cuda.pathfinder._utils.windows_arch import windows_python_arch
+
+    _WINDOWS_PYTHON_ARCH = windows_python_arch()
+else:
+    _WINDOWS_PYTHON_ARCH = None
+
+# These describe the Python process ABI, not the Windows host architecture.
+IS_WINDOWS_X64 = _WINDOWS_PYTHON_ARCH == "x64"
+IS_WINDOWS_ARM64 = _WINDOWS_PYTHON_ARCH == "arm64"
 
 
 def quote_for_shell(s: str) -> str:

cuda_pathfinder/cuda/pathfinder/_utils/windows_arch.py

@@ -0,0 +1,65 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import sysconfig
+
+WINDOWS_PE_MACHINE_BY_ARCH = {
+    "x64": 0x8664,
+    "arm64": 0xAA64,
+}
+
+
+class UnsupportedArchError(RuntimeError):
+    """Raised when Python reports an unsupported Windows architecture."""
+
+    def __init__(self, platform_tag: str) -> None:
+        self.platform_tag = platform_tag
+        super().__init__(
+            f"Unsupported Windows Python platform tag: {platform_tag!r}; expected 'win-amd64' or 'win-arm64'"
+        )
+
+
+def windows_python_arch() -> str:
+    """Return the current Windows Python interpreter architecture."""
+    raw_platform_tag = sysconfig.get_platform()
+    platform_tag = raw_platform_tag.lower().replace("_", "-")
+
+    if platform_tag == "win-arm64":
+        return "arm64"
+
+    if platform_tag == "win-amd64":
+        return "x64"
+
+    raise UnsupportedArchError(raw_platform_tag)
+
+
+def windows_pe_matches_arch(path: str, target_arch: str) -> bool:
+    """Return whether a Windows Portable Executable (PE) targets the requested architecture.
+
+    PE is the file format used for Windows executables and DLLs. This reads the
+    PE/COFF header's machine field to distinguish x64 images from Arm64 images.
+    """
+    expected_machine = WINDOWS_PE_MACHINE_BY_ARCH.get(target_arch)
+    if expected_machine is None:
+        raise ValueError(f"Unsupported Windows target architecture: {target_arch!r}")
+
+    try:
+        with open(path, "rb") as stream:
+            if stream.read(2) != b"MZ":
+                return False
+            stream.seek(0x3C)
+            pe_offset_bytes = stream.read(4)
+            if len(pe_offset_bytes) != 4:
+                return False
+            stream.seek(int.from_bytes(pe_offset_bytes, "little"))
+            if stream.read(4) != b"PE\0\0":
+                return False
+            machine_bytes = stream.read(2)
+            if len(machine_bytes) != 2:
+                return False
+    except OSError:
+        return False
+
+    return int.from_bytes(machine_bytes, "little") == expected_machine