Make Windows pathfinder dynamic library searches architecture-aware

cuda_pathfinder/cuda/pathfinder/__init__.py

@@ -60,6 +60,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 UnsupportedWindowsArchError as UnsupportedWindowsArchError
 
 from cuda.pathfinder._version import __version__  # isort: skip
 

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

@@ -60,11 +60,16 @@
 }
 SITE_PACKAGES_LIBDIRS_LINUX = SITE_PACKAGES_LIBDIRS_LINUX_CTK | SITE_PACKAGES_LIBDIRS_LINUX_OTHER
 
+# Historical table exports represent the original x64 catalog.
 SITE_PACKAGES_LIBDIRS_WINDOWS_CTK = {
-    desc.name: desc.site_packages_windows for desc in _CTK_DESCRIPTORS if desc.site_packages_windows
+    desc.name: desc.site_packages_windows.for_arch("x64")
+    for desc in _CTK_DESCRIPTORS
+    if desc.site_packages_windows.paths
 }
 SITE_PACKAGES_LIBDIRS_WINDOWS_OTHER = {
-    desc.name: desc.site_packages_windows for desc in _NON_CTK_DESCRIPTORS if desc.site_packages_windows
+    desc.name: desc.site_packages_windows.for_arch("x64")
+    for desc in _NON_CTK_DESCRIPTORS
+    if desc.site_packages_windows.paths
 }
 SITE_PACKAGES_LIBDIRS_WINDOWS = SITE_PACKAGES_LIBDIRS_WINDOWS_CTK | SITE_PACKAGES_LIBDIRS_WINDOWS_OTHER
 

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 UnsupportedWindowsArchError(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 UnsupportedWindowsArchError(raw_platform_tag)

cuda_pathfinder/docs/source/api.rst

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

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

@@ -59,7 +59,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 not spec.site_packages_windows.paths, f"driver lib {spec.name} has site_packages_windows"
 
 
 @pytest.mark.parametrize(

cuda_pathfinder/tests/test_lib_descriptor.py

@@ -58,7 +58,7 @@ def test_site_packages_linux_match(name):
 
 @pytest.mark.parametrize("name", sorted(LIB_DESCRIPTORS))
 def test_site_packages_windows_match(name):
-    assert LIB_DESCRIPTORS[name].site_packages_windows == SITE_PACKAGES_LIBDIRS_WINDOWS.get(name, ())
+    assert LIB_DESCRIPTORS[name].site_packages_windows.for_arch("x64") == SITE_PACKAGES_LIBDIRS_WINDOWS.get(name, ())
 
 
 @pytest.mark.parametrize("name", sorted(LIB_DESCRIPTORS))