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/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_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
@@ -78,7 +80,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:
@@ -173,17 +175,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,
@@ -216,4 +220,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

@@ -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")
@@ -27,9 +33,20 @@
 )
 
 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}
@@ -60,13 +77,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,30 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import sysconfig
+
+
+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)

cuda_pathfinder/docs/source/api.rst

@@ -24,6 +24,7 @@ CUDA bitcode and static libraries.
    DynamicLibNotFoundError
    DynamicLibUnknownError
    DynamicLibNotAvailableError
+   UnsupportedArchError
 
    SUPPORTED_HEADERS_CTK
    find_nvidia_header_directory

cuda_pathfinder/docs/source/release/1.6.0-notes.rst

@@ -35,6 +35,17 @@ Highlights
 Internal maintenance
 --------------------
 
+* Add explicit ``_X64`` and ``_ARM64`` variants of the internal
+  ``SITE_PACKAGES_LIBDIRS_WINDOWS*`` tables. The unsuffixed names remain x64
+  aliases for backward compatibility, but are not recommended for new code.
+  Consumers should select the projection matching the current process
+  architecture and must not combine the two.
+
+* Record supported Windows architectures explicitly in each dynamic-library
+  descriptor. ``SUPPORTED_NVIDIA_LIBNAMES`` now selects the CTK library names
+  supported by the current Windows interpreter architecture; the legacy
+  ``SUPPORTED_LIBNAMES_WINDOWS`` table remains an x64 compatibility alias.
+
 * Clean up dead and misleading error-handling code in the Linux and Windows
   dynamic-library loaders; runtime behavior is unchanged.
   (`PR #2239 <https://github.com/NVIDIA/cuda-python/pull/2239>`_)

cuda_pathfinder/tests/test_catalog_writer.py

@@ -0,0 +1,24 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import dataclasses
+import runpy
+
+import pytest
+
+from cuda.pathfinder._dynamic_libs.descriptor_catalog import DESCRIPTOR_CATALOG
+from toolshed._catalog_writer import render_catalog
+
+
+@pytest.mark.agent_authored(model="gpt-5")
+def test_catalog_writer_round_trips_descriptor_metadata(tmp_path):
+    generated_catalog = tmp_path / "descriptor_catalog.py"
+    generated_catalog.write_text(render_catalog(DESCRIPTOR_CATALOG), encoding="utf-8")
+
+    rendered_specs = runpy.run_path(str(generated_catalog))["DESCRIPTOR_CATALOG"]
+
+    assert tuple(dataclasses.asdict(spec) for spec in rendered_specs) == tuple(
+        dataclasses.asdict(spec) for spec in DESCRIPTOR_CATALOG
+    )

cuda_pathfinder/tests/test_ctk_root_discovery.py

@@ -19,6 +19,7 @@
     _try_ctk_root_canary,
     resolve_ctk_root_via_canary,
 )
+from cuda.pathfinder._dynamic_libs.search_platform import WindowsSearchPlatform
 from cuda.pathfinder._dynamic_libs.search_steps import (
     SearchContext,
     _derive_ctk_root_linux,
@@ -126,6 +127,11 @@ def test_derive_ctk_root_windows_ctk13():
     assert _derive_ctk_root_windows(path) == r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.0"
 
 
+def test_derive_ctk_root_windows_ctk13_arm64():
+    path = r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4\bin\arm64\cudart64_13.dll"
+    assert _derive_ctk_root_windows(path) == r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4"
+
+
 def test_derive_ctk_root_windows_ctk12():
     path = r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.8\bin\cudart64_12.dll"
     assert _derive_ctk_root_windows(path) == r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.8"
@@ -190,6 +196,23 @@ def test_try_via_ctk_root_regular_lib(tmp_path):
     assert result.found_via == "system-ctk-root"
 
 
+def test_try_via_ctk_root_windows_arm64_prefers_arch_dir(tmp_path):
+    ctk_root = tmp_path / "cuda-13"
+    x64_dir = ctk_root / "bin" / "x64"
+    arm64_dir = ctk_root / "bin" / "arm64"
+    x64_dir.mkdir(parents=True)
+    arm64_dir.mkdir(parents=True)
+    (x64_dir / "cudart64_13.dll").write_bytes(b"fake")
+    arm64_lib = arm64_dir / "cudart64_13.dll"
+    arm64_lib.write_bytes(b"fake")
+
+    ctx = SearchContext(LIB_DESCRIPTORS["cudart"], platform=WindowsSearchPlatform(target_arch="arm64"))
+    result = find_via_ctk_root(ctx, str(ctk_root))
+    assert result is not None
+    assert result.abs_path == str(arm64_lib)
+    assert result.found_via == "system-ctk-root"
+
+
 # ---------------------------------------------------------------------------
 # _resolve_system_loaded_abs_path_in_subprocess
 # ---------------------------------------------------------------------------

cuda_pathfinder/tests/test_descriptor_catalog.py

@@ -13,10 +13,11 @@
 
 import pytest
 
-from cuda.pathfinder._dynamic_libs.descriptor_catalog import DESCRIPTOR_CATALOG, DescriptorSpec
+from cuda.pathfinder._dynamic_libs.descriptor_catalog import DESCRIPTOR_CATALOG, DescriptorSpec, WindowsSearchDirs
 
 _VALID_NAME_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
 _VALID_PACKAGED_WITH_VALUES = {"ctk", "other", "driver"}
+_VALID_WINDOWS_ARCHES = ("x64", "arm64")
 _CATALOG_BY_NAME = {spec.name: spec for spec in DESCRIPTOR_CATALOG}
 
 
@@ -59,7 +60,7 @@ def test_no_self_dependency(spec: DescriptorSpec):
 def test_driver_libs_have_no_site_packages(spec: DescriptorSpec):
     """Driver libs are system-search-only; site-packages paths would be unused."""
     assert not spec.site_packages_linux, f"driver lib {spec.name} has site_packages_linux"
-    assert not spec.site_packages_windows, f"driver lib {spec.name} has site_packages_windows"
+    assert spec.site_packages_windows == WindowsSearchDirs(), f"driver lib {spec.name} has site_packages_windows"
 
 
 @pytest.mark.parametrize(
@@ -85,6 +86,25 @@ def test_windows_dlls_look_like_dlls(spec: DescriptorSpec):
         assert dll.endswith(".dll"), f"Unexpected Windows DLL format: {dll}"
 
 
+@pytest.mark.parametrize("spec", DESCRIPTOR_CATALOG, ids=lambda s: s.name)
+@pytest.mark.agent_authored(model="gpt-5")
+def test_supported_windows_arch_is_explicit_and_canonical(spec: DescriptorSpec):
+    expected = tuple(arch for arch in _VALID_WINDOWS_ARCHES if arch in spec.supported_windows_arch)
+    assert spec.supported_windows_arch == expected
+    assert bool(spec.supported_windows_arch) == bool(spec.windows_dlls)
+
+
+@pytest.mark.parametrize("spec", DESCRIPTOR_CATALOG, ids=lambda s: s.name)
+@pytest.mark.agent_authored(model="gpt-5")
+def test_windows_search_dirs_do_not_include_unsupported_arches(spec: DescriptorSpec):
+    if not spec.windows_dlls:
+        return
+    for arch in _VALID_WINDOWS_ARCHES:
+        if arch not in spec.supported_windows_arch:
+            assert not spec.site_packages_windows.for_arch(arch)
+            assert not spec.anchor_rel_dirs_windows.for_arch(arch)
+
+
 @pytest.mark.parametrize("spec", DESCRIPTOR_CATALOG, ids=lambda s: s.name)
 def test_ctk_root_canary_anchors_reference_known_ctk_libs(spec: DescriptorSpec):
     for anchor in spec.ctk_root_canary_anchor_libnames: