From 8c55ee10ef50d5f483881027911f135bdd82522b Mon Sep 17 00:00:00 2001 From: isvoid <13521008+isVoid@users.noreply.github.com> Date: Mon, 20 Jul 2026 13:27:23 -0700 Subject: [PATCH 1/2] Make Windows pathfinder searches architecture-aware --- cuda_pathfinder/cuda/pathfinder/__init__.py | 1 + .../_dynamic_libs/descriptor_catalog.py | 142 +++++++++---- .../_dynamic_libs/search_platform.py | 14 +- .../pathfinder/_dynamic_libs/search_steps.py | 3 +- .../_dynamic_libs/supported_nvidia_libs.py | 9 +- .../cuda/pathfinder/_utils/windows_arch.py | 30 +++ cuda_pathfinder/docs/source/api.rst | 1 + .../tests/test_ctk_root_discovery.py | 23 +++ .../tests/test_descriptor_catalog.py | 2 +- cuda_pathfinder/tests/test_lib_descriptor.py | 2 +- cuda_pathfinder/tests/test_search_steps.py | 190 +++++++++++++++++- 11 files changed, 362 insertions(+), 55 deletions(-) create mode 100644 cuda_pathfinder/cuda/pathfinder/_utils/windows_arch.py diff --git a/cuda_pathfinder/cuda/pathfinder/__init__.py b/cuda_pathfinder/cuda/pathfinder/__init__.py index dc818dfd08f..69ae2e08efa 100644 --- a/cuda_pathfinder/cuda/pathfinder/__init__.py +++ b/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 diff --git a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py index a26862c5435..bc1155f8491 100644 --- a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py +++ b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py @@ -6,6 +6,7 @@ from __future__ import annotations from dataclasses import dataclass +from pathlib import PurePosixPath from typing import Literal from cuda.pathfinder._utils.ctk_root_canary import CTK_ROOT_CANARY_ANCHOR_LIBNAMES @@ -13,6 +14,61 @@ PackagedWith = Literal["ctk", "other", "driver"] +@dataclass(frozen=True, slots=True) +class WindowsSearchDir: + """One Windows search location with architecture-specific spellings.""" + + x64: str + arm64: str + + @classmethod + def common(cls, path: str) -> WindowsSearchDir: + return cls(x64=path, arm64=path) + + def for_arch(self, target_arch: str) -> str: + if target_arch == "x64": + return self.x64 + if target_arch == "arm64": + return self.arm64 + raise ValueError(f"Unsupported Windows target architecture: {target_arch!r}") + + +@dataclass(frozen=True, slots=True) +class WindowsSearchDirs: + """Ordered Windows search locations selected for one architecture.""" + + paths: tuple[WindowsSearchDir, ...] = () + + @classmethod + def common(cls, *paths: str) -> WindowsSearchDirs: + return cls(tuple(WindowsSearchDir.common(path) for path in paths)) + + def for_arch(self, target_arch: str) -> tuple[str, ...]: + return tuple(path.for_arch(target_arch) for path in self.paths) + + +DEFAULT_WINDOWS_CTK_ANCHOR_DIRS = WindowsSearchDirs( + paths=( + WindowsSearchDir(x64="bin/x64", arm64="bin/arm64"), + WindowsSearchDir.common("bin"), + ) +) + + +def _ctk_windows_wheel_dirs(cuda13_bin_dir: str, cuda12_dir: str) -> WindowsSearchDirs: + """Search the CUDA 13 aggregate wheel before its CUDA 12 component wheel.""" + cuda13_bin_path = PurePosixPath(cuda13_bin_dir) + return WindowsSearchDirs( + paths=( + WindowsSearchDir( + x64=(cuda13_bin_path / "x86_64").as_posix(), + arm64=(cuda13_bin_path / "arm64").as_posix(), + ), + WindowsSearchDir.common(cuda12_dir), + ) + ) + + @dataclass(frozen=True, slots=True) class DescriptorSpec: name: str @@ -20,10 +76,10 @@ class DescriptorSpec: linux_sonames: tuple[str, ...] = () windows_dlls: tuple[str, ...] = () site_packages_linux: tuple[str, ...] = () - site_packages_windows: tuple[str, ...] = () + site_packages_windows: WindowsSearchDirs = WindowsSearchDirs() dependencies: tuple[str, ...] = () anchor_rel_dirs_linux: tuple[str, ...] = ("lib64", "lib") - anchor_rel_dirs_windows: tuple[str, ...] = ("bin/x64", "bin") + anchor_rel_dirs_windows: WindowsSearchDirs = DEFAULT_WINDOWS_CTK_ANCHOR_DIRS ctk_root_canary_anchor_libnames: tuple[str, ...] = () requires_add_dll_directory: bool = False requires_rtld_deepbind: bool = False @@ -39,7 +95,7 @@ class DescriptorSpec: linux_sonames=("libcudart.so.12", "libcudart.so.13"), windows_dlls=("cudart64_12.dll", "cudart64_13.dll"), site_packages_linux=("nvidia/cu13/lib", "nvidia/cuda_runtime/lib"), - site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/cuda_runtime/bin"), + site_packages_windows=_ctk_windows_wheel_dirs("nvidia/cu13/bin", "nvidia/cuda_runtime/bin"), ), DescriptorSpec( name="nvfatbin", @@ -47,7 +103,7 @@ class DescriptorSpec: linux_sonames=("libnvfatbin.so.12", "libnvfatbin.so.13"), windows_dlls=("nvfatbin_120_0.dll", "nvfatbin_130_0.dll"), site_packages_linux=("nvidia/cu13/lib", "nvidia/nvfatbin/lib"), - site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/nvfatbin/bin"), + site_packages_windows=_ctk_windows_wheel_dirs("nvidia/cu13/bin", "nvidia/nvfatbin/bin"), ), DescriptorSpec( name="nvJitLink", @@ -55,7 +111,7 @@ class DescriptorSpec: linux_sonames=("libnvJitLink.so.12", "libnvJitLink.so.13"), windows_dlls=("nvJitLink_120_0.dll", "nvJitLink_130_0.dll"), site_packages_linux=("nvidia/cu13/lib", "nvidia/nvjitlink/lib"), - site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/nvjitlink/bin"), + site_packages_windows=_ctk_windows_wheel_dirs("nvidia/cu13/bin", "nvidia/nvjitlink/bin"), ), DescriptorSpec( name="nvrtc", @@ -63,7 +119,7 @@ class DescriptorSpec: linux_sonames=("libnvrtc.so.12", "libnvrtc.so.13"), windows_dlls=("nvrtc64_120_0.dll", "nvrtc64_130_0.dll"), site_packages_linux=("nvidia/cu13/lib", "nvidia/cuda_nvrtc/lib"), - site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/cuda_nvrtc/bin"), + site_packages_windows=_ctk_windows_wheel_dirs("nvidia/cu13/bin", "nvidia/cuda_nvrtc/bin"), requires_add_dll_directory=True, ), DescriptorSpec( @@ -72,9 +128,14 @@ class DescriptorSpec: linux_sonames=("libnvvm.so.4",), windows_dlls=("nvvm64.dll", "nvvm64_40_0.dll", "nvvm70.dll"), site_packages_linux=("nvidia/cu13/lib", "nvidia/cuda_nvcc/nvvm/lib64"), - site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/cuda_nvcc/nvvm/bin"), + site_packages_windows=_ctk_windows_wheel_dirs("nvidia/cu13/bin", "nvidia/cuda_nvcc/nvvm/bin"), anchor_rel_dirs_linux=("nvvm/lib64",), - anchor_rel_dirs_windows=("nvvm/bin/*", "nvvm/bin"), + anchor_rel_dirs_windows=WindowsSearchDirs( + paths=( + WindowsSearchDir(x64="nvvm/bin/x64", arm64="nvvm/bin/arm64"), + WindowsSearchDir.common("nvvm/bin"), + ) + ), ctk_root_canary_anchor_libnames=CTK_ROOT_CANARY_ANCHOR_LIBNAMES, ), DescriptorSpec( @@ -83,7 +144,7 @@ class DescriptorSpec: linux_sonames=("libcublas.so.12", "libcublas.so.13"), windows_dlls=("cublas64_12.dll", "cublas64_13.dll"), site_packages_linux=("nvidia/cu13/lib", "nvidia/cublas/lib"), - site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/cublas/bin"), + site_packages_windows=_ctk_windows_wheel_dirs("nvidia/cu13/bin", "nvidia/cublas/bin"), dependencies=("cublasLt",), ), DescriptorSpec( @@ -92,7 +153,7 @@ class DescriptorSpec: linux_sonames=("libcublasLt.so.12", "libcublasLt.so.13"), windows_dlls=("cublasLt64_12.dll", "cublasLt64_13.dll"), site_packages_linux=("nvidia/cu13/lib", "nvidia/cublas/lib"), - site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/cublas/bin"), + site_packages_windows=_ctk_windows_wheel_dirs("nvidia/cu13/bin", "nvidia/cublas/bin"), ), DescriptorSpec( name="cufft", @@ -100,7 +161,7 @@ class DescriptorSpec: linux_sonames=("libcufft.so.11", "libcufft.so.12"), windows_dlls=("cufft64_11.dll", "cufft64_12.dll"), site_packages_linux=("nvidia/cu13/lib", "nvidia/cufft/lib"), - site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/cufft/bin"), + site_packages_windows=_ctk_windows_wheel_dirs("nvidia/cu13/bin", "nvidia/cufft/bin"), requires_add_dll_directory=True, ), DescriptorSpec( @@ -109,7 +170,7 @@ class DescriptorSpec: linux_sonames=("libcufftw.so.11", "libcufftw.so.12"), windows_dlls=("cufftw64_11.dll", "cufftw64_12.dll"), site_packages_linux=("nvidia/cu13/lib", "nvidia/cufft/lib"), - site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/cufft/bin"), + site_packages_windows=_ctk_windows_wheel_dirs("nvidia/cu13/bin", "nvidia/cufft/bin"), dependencies=("cufft",), ), DescriptorSpec( @@ -118,7 +179,7 @@ class DescriptorSpec: linux_sonames=("libcurand.so.10",), windows_dlls=("curand64_10.dll",), site_packages_linux=("nvidia/cu13/lib", "nvidia/curand/lib"), - site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/curand/bin"), + site_packages_windows=_ctk_windows_wheel_dirs("nvidia/cu13/bin", "nvidia/curand/bin"), ), DescriptorSpec( name="cusolver", @@ -126,7 +187,7 @@ class DescriptorSpec: linux_sonames=("libcusolver.so.11", "libcusolver.so.12"), windows_dlls=("cusolver64_11.dll", "cusolver64_12.dll"), site_packages_linux=("nvidia/cu13/lib", "nvidia/cusolver/lib"), - site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/cusolver/bin"), + site_packages_windows=_ctk_windows_wheel_dirs("nvidia/cu13/bin", "nvidia/cusolver/bin"), dependencies=("nvJitLink", "cusparse", "cublasLt", "cublas"), ), DescriptorSpec( @@ -135,7 +196,7 @@ class DescriptorSpec: linux_sonames=("libcusolverMg.so.11", "libcusolverMg.so.12"), windows_dlls=("cusolverMg64_11.dll", "cusolverMg64_12.dll"), site_packages_linux=("nvidia/cu13/lib", "nvidia/cusolver/lib"), - site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/cusolver/bin"), + site_packages_windows=_ctk_windows_wheel_dirs("nvidia/cu13/bin", "nvidia/cusolver/bin"), dependencies=("nvJitLink", "cublasLt", "cublas"), ), DescriptorSpec( @@ -144,7 +205,7 @@ class DescriptorSpec: linux_sonames=("libcusparse.so.12",), windows_dlls=("cusparse64_12.dll",), site_packages_linux=("nvidia/cu13/lib", "nvidia/cusparse/lib"), - site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/cusparse/bin"), + site_packages_windows=_ctk_windows_wheel_dirs("nvidia/cu13/bin", "nvidia/cusparse/bin"), dependencies=("nvJitLink",), ), DescriptorSpec( @@ -153,7 +214,7 @@ class DescriptorSpec: linux_sonames=("libnppc.so.12", "libnppc.so.13"), windows_dlls=("nppc64_12.dll", "nppc64_13.dll"), site_packages_linux=("nvidia/cu13/lib", "nvidia/npp/lib"), - site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"), + site_packages_windows=_ctk_windows_wheel_dirs("nvidia/cu13/bin", "nvidia/npp/bin"), ), DescriptorSpec( name="nppial", @@ -161,7 +222,7 @@ class DescriptorSpec: linux_sonames=("libnppial.so.12", "libnppial.so.13"), windows_dlls=("nppial64_12.dll", "nppial64_13.dll"), site_packages_linux=("nvidia/cu13/lib", "nvidia/npp/lib"), - site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"), + site_packages_windows=_ctk_windows_wheel_dirs("nvidia/cu13/bin", "nvidia/npp/bin"), dependencies=("nppc",), ), DescriptorSpec( @@ -170,7 +231,7 @@ class DescriptorSpec: linux_sonames=("libnppicc.so.12", "libnppicc.so.13"), windows_dlls=("nppicc64_12.dll", "nppicc64_13.dll"), site_packages_linux=("nvidia/cu13/lib", "nvidia/npp/lib"), - site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"), + site_packages_windows=_ctk_windows_wheel_dirs("nvidia/cu13/bin", "nvidia/npp/bin"), dependencies=("nppc",), ), DescriptorSpec( @@ -179,7 +240,7 @@ class DescriptorSpec: linux_sonames=("libnppidei.so.12", "libnppidei.so.13"), windows_dlls=("nppidei64_12.dll", "nppidei64_13.dll"), site_packages_linux=("nvidia/cu13/lib", "nvidia/npp/lib"), - site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"), + site_packages_windows=_ctk_windows_wheel_dirs("nvidia/cu13/bin", "nvidia/npp/bin"), dependencies=("nppc",), ), DescriptorSpec( @@ -188,7 +249,7 @@ class DescriptorSpec: linux_sonames=("libnppif.so.12", "libnppif.so.13"), windows_dlls=("nppif64_12.dll", "nppif64_13.dll"), site_packages_linux=("nvidia/cu13/lib", "nvidia/npp/lib"), - site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"), + site_packages_windows=_ctk_windows_wheel_dirs("nvidia/cu13/bin", "nvidia/npp/bin"), dependencies=("nppc",), ), DescriptorSpec( @@ -197,7 +258,7 @@ class DescriptorSpec: linux_sonames=("libnppig.so.12", "libnppig.so.13"), windows_dlls=("nppig64_12.dll", "nppig64_13.dll"), site_packages_linux=("nvidia/cu13/lib", "nvidia/npp/lib"), - site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"), + site_packages_windows=_ctk_windows_wheel_dirs("nvidia/cu13/bin", "nvidia/npp/bin"), dependencies=("nppc",), ), DescriptorSpec( @@ -206,7 +267,7 @@ class DescriptorSpec: linux_sonames=("libnppim.so.12", "libnppim.so.13"), windows_dlls=("nppim64_12.dll", "nppim64_13.dll"), site_packages_linux=("nvidia/cu13/lib", "nvidia/npp/lib"), - site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"), + site_packages_windows=_ctk_windows_wheel_dirs("nvidia/cu13/bin", "nvidia/npp/bin"), dependencies=("nppc",), ), DescriptorSpec( @@ -215,7 +276,7 @@ class DescriptorSpec: linux_sonames=("libnppist.so.12", "libnppist.so.13"), windows_dlls=("nppist64_12.dll", "nppist64_13.dll"), site_packages_linux=("nvidia/cu13/lib", "nvidia/npp/lib"), - site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"), + site_packages_windows=_ctk_windows_wheel_dirs("nvidia/cu13/bin", "nvidia/npp/bin"), dependencies=("nppc",), ), DescriptorSpec( @@ -224,7 +285,7 @@ class DescriptorSpec: linux_sonames=("libnppisu.so.12", "libnppisu.so.13"), windows_dlls=("nppisu64_12.dll", "nppisu64_13.dll"), site_packages_linux=("nvidia/cu13/lib", "nvidia/npp/lib"), - site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"), + site_packages_windows=_ctk_windows_wheel_dirs("nvidia/cu13/bin", "nvidia/npp/bin"), dependencies=("nppc",), ), DescriptorSpec( @@ -233,7 +294,7 @@ class DescriptorSpec: linux_sonames=("libnppitc.so.12", "libnppitc.so.13"), windows_dlls=("nppitc64_12.dll", "nppitc64_13.dll"), site_packages_linux=("nvidia/cu13/lib", "nvidia/npp/lib"), - site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"), + site_packages_windows=_ctk_windows_wheel_dirs("nvidia/cu13/bin", "nvidia/npp/bin"), dependencies=("nppc",), ), DescriptorSpec( @@ -242,7 +303,7 @@ class DescriptorSpec: linux_sonames=("libnpps.so.12", "libnpps.so.13"), windows_dlls=("npps64_12.dll", "npps64_13.dll"), site_packages_linux=("nvidia/cu13/lib", "nvidia/npp/lib"), - site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"), + site_packages_windows=_ctk_windows_wheel_dirs("nvidia/cu13/bin", "nvidia/npp/bin"), dependencies=("nppc",), ), DescriptorSpec( @@ -251,7 +312,7 @@ class DescriptorSpec: linux_sonames=("libnvblas.so.12", "libnvblas.so.13"), windows_dlls=("nvblas64_12.dll", "nvblas64_13.dll"), site_packages_linux=("nvidia/cu13/lib", "nvidia/cublas/lib"), - site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/cublas/bin"), + site_packages_windows=_ctk_windows_wheel_dirs("nvidia/cu13/bin", "nvidia/cublas/bin"), dependencies=("cublas", "cublasLt"), ), DescriptorSpec( @@ -260,7 +321,7 @@ class DescriptorSpec: linux_sonames=("libnvjpeg.so.12", "libnvjpeg.so.13"), windows_dlls=("nvjpeg64_12.dll", "nvjpeg64_13.dll"), site_packages_linux=("nvidia/cu13/lib", "nvidia/nvjpeg/lib"), - site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/nvjpeg/bin"), + site_packages_windows=_ctk_windows_wheel_dirs("nvidia/cu13/bin", "nvidia/nvjpeg/bin"), ), DescriptorSpec( name="cufile", @@ -291,9 +352,15 @@ class DescriptorSpec: "cupti64_2022.4.1.dll", ), site_packages_linux=("nvidia/cu13/lib", "nvidia/cuda_cupti/lib"), - site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/cuda_cupti/bin"), + site_packages_windows=_ctk_windows_wheel_dirs("nvidia/cu13/bin", "nvidia/cuda_cupti/bin"), anchor_rel_dirs_linux=("extras/CUPTI/lib64", "lib"), - anchor_rel_dirs_windows=("extras/CUPTI/lib64", "bin"), + anchor_rel_dirs_windows=WindowsSearchDirs( + paths=( + WindowsSearchDir.common("extras/CUPTI/lib64"), + WindowsSearchDir(x64="bin/x64", arm64="bin/arm64"), + WindowsSearchDir.common("bin"), + ) + ), ctk_root_canary_anchor_libnames=CTK_ROOT_CANARY_ANCHOR_LIBNAMES, ), DescriptorSpec( @@ -304,9 +371,6 @@ class DescriptorSpec: site_packages_linux=("nvidia/cu13/lib",), # No Windows pip wheel ships cudla.dll today; it is loaded from the local # CUDA Toolkit only, so site_packages_windows is intentionally left empty. - # The Windows CUDA Toolkit ships cudla.dll under per-architecture bin - # subdirs (e.g. bin/arm64 on N1X); search those ahead of the defaults. - anchor_rel_dirs_windows=("bin/arm64", "bin/x64", "bin"), ), # ----------------------------------------------------------------------- # Third-party / separately packaged libraries @@ -339,7 +403,7 @@ class DescriptorSpec: linux_sonames=("libmathdx.so.0",), windows_dlls=("mathdx64_0.dll",), site_packages_linux=("nvidia/cu13/lib", "nvidia/cu12/lib"), - site_packages_windows=("nvidia/cu13/bin", "nvidia/cu12/bin"), + site_packages_windows=WindowsSearchDirs.common("nvidia/cu13/bin", "nvidia/cu12/bin"), dependencies=("nvrtc",), ), DescriptorSpec( @@ -348,7 +412,7 @@ class DescriptorSpec: linux_sonames=("libcudss.so.0",), windows_dlls=("cudss64_0.dll",), site_packages_linux=("nvidia/cu13/lib", "nvidia/cu12/lib"), - site_packages_windows=("nvidia/cu13/bin", "nvidia/cu12/bin"), + site_packages_windows=WindowsSearchDirs.common("nvidia/cu13/bin", "nvidia/cu12/bin"), dependencies=("cublas", "cublasLt"), ), DescriptorSpec( @@ -357,7 +421,7 @@ class DescriptorSpec: linux_sonames=("libcusparseLt.so.0",), windows_dlls=("cusparseLt.dll",), site_packages_linux=("nvidia/cu13/lib", "nvidia/cusparselt/lib"), - site_packages_windows=("nvidia/cu13/bin/x64", "nvidia/cusparselt/bin"), + site_packages_windows=WindowsSearchDirs.common("nvidia/cu13/bin/x64", "nvidia/cusparselt/bin"), ), DescriptorSpec( name="cutensor", @@ -365,7 +429,7 @@ class DescriptorSpec: linux_sonames=("libcutensor.so.2",), windows_dlls=("cutensor.dll",), site_packages_linux=("cutensor/lib",), - site_packages_windows=("cutensor/bin",), + site_packages_windows=WindowsSearchDirs.common("cutensor/bin"), dependencies=("cublasLt",), ), DescriptorSpec( @@ -374,7 +438,7 @@ class DescriptorSpec: linux_sonames=("libcutensorMg.so.2",), windows_dlls=("cutensorMg.dll",), site_packages_linux=("cutensor/lib",), - site_packages_windows=("cutensor/bin",), + site_packages_windows=WindowsSearchDirs.common("cutensor/bin"), dependencies=("cutensor", "cublasLt"), ), DescriptorSpec( diff --git a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_platform.py b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_platform.py index 37fd6eb1700..ef93eae1d4d 100644 --- a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_platform.py +++ b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_platform.py @@ -13,13 +13,15 @@ import glob import os from collections.abc import Sequence -from dataclasses import dataclass +from dataclasses import dataclass, field +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.so. (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 = field(default_factory=windows_python_arch) + 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 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 desc.anchor_rel_dirs_windows.for_arch(self.target_arch) def find_in_site_packages( self, diff --git a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_steps.py b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_steps.py index 55d8a8aa674..56d47a69aac 100644 --- a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_steps.py +++ b/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) diff --git a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/supported_nvidia_libs.py b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/supported_nvidia_libs.py index db06411c6d0..e5909313f35 100644 --- a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/supported_nvidia_libs.py +++ b/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 diff --git a/cuda_pathfinder/cuda/pathfinder/_utils/windows_arch.py b/cuda_pathfinder/cuda/pathfinder/_utils/windows_arch.py new file mode 100644 index 00000000000..ea3bbdf1c64 --- /dev/null +++ b/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) diff --git a/cuda_pathfinder/docs/source/api.rst b/cuda_pathfinder/docs/source/api.rst index e49478c09ec..c532ba43bdc 100644 --- a/cuda_pathfinder/docs/source/api.rst +++ b/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 diff --git a/cuda_pathfinder/tests/test_ctk_root_discovery.py b/cuda_pathfinder/tests/test_ctk_root_discovery.py index 9ad148dccd2..98f1eb9e608 100644 --- a/cuda_pathfinder/tests/test_ctk_root_discovery.py +++ b/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 # --------------------------------------------------------------------------- diff --git a/cuda_pathfinder/tests/test_descriptor_catalog.py b/cuda_pathfinder/tests/test_descriptor_catalog.py index b2c8eece4bb..ee820f584c2 100644 --- a/cuda_pathfinder/tests/test_descriptor_catalog.py +++ b/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( diff --git a/cuda_pathfinder/tests/test_lib_descriptor.py b/cuda_pathfinder/tests/test_lib_descriptor.py index cda96131e13..57be934e47d 100644 --- a/cuda_pathfinder/tests/test_lib_descriptor.py +++ b/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)) diff --git a/cuda_pathfinder/tests/test_search_steps.py b/cuda_pathfinder/tests/test_search_steps.py index 1b881707dfb..c5f935062c7 100644 --- a/cuda_pathfinder/tests/test_search_steps.py +++ b/cuda_pathfinder/tests/test_search_steps.py @@ -9,6 +9,8 @@ import pytest +from cuda.pathfinder import UnsupportedWindowsArchError +from cuda.pathfinder._dynamic_libs.descriptor_catalog import WindowsSearchDir, WindowsSearchDirs from cuda.pathfinder._dynamic_libs.lib_descriptor import LIB_DESCRIPTORS, LibDescriptor from cuda.pathfinder._dynamic_libs.load_dl_common import DynamicLibNotFoundError from cuda.pathfinder._dynamic_libs.search_platform import LinuxSearchPlatform, WindowsSearchPlatform @@ -23,6 +25,7 @@ find_in_site_packages, run_find_steps, ) +from cuda.pathfinder._utils import windows_arch as windows_arch_mod _STEPS_MOD = "cuda.pathfinder._dynamic_libs.search_steps" _PLAT_MOD = "cuda.pathfinder._dynamic_libs.search_platform" @@ -40,7 +43,7 @@ def _make_desc(name: str = "cudart", **overrides) -> LibDescriptor: "linux_sonames": ("libcudart.so",), "windows_dlls": ("cudart64_12.dll",), "site_packages_linux": (os.path.join("nvidia", "cuda_runtime", "lib"),), - "site_packages_windows": (os.path.join("nvidia", "cuda_runtime", "bin"),), + "site_packages_windows": WindowsSearchDirs.common(os.path.join("nvidia", "cuda_runtime", "bin")), } defaults.update(overrides) return LibDescriptor(**defaults) @@ -52,6 +55,14 @@ def _ctx(desc: LibDescriptor | None = None, *, platform=None) -> SearchContext: return SearchContext(desc or _make_desc(), platform=platform) +def _patch_site_packages_search(mocker, root): + def _find_sub_dirs(sub_dirs): + path = root.joinpath(*sub_dirs) + return [str(path)] if path.is_dir() else [] + + return mocker.patch(f"{_PLAT_MOD}.find_sub_dirs_all_sitepackages", side_effect=_find_sub_dirs) + + # --------------------------------------------------------------------------- # SearchContext # --------------------------------------------------------------------------- @@ -83,6 +94,33 @@ def test_raise_not_found_empty_messages(self): ctx.raise_not_found() +# --------------------------------------------------------------------------- +# Windows Python architecture detection +# --------------------------------------------------------------------------- + + +class TestWindowsPythonArch: + def test_detects_sysconfig_x64(self, mocker): + mocker.patch.object(windows_arch_mod.sysconfig, "get_platform", return_value="win-amd64") + + assert windows_arch_mod.windows_python_arch() == "x64" + + def test_detects_sysconfig_arm64(self, mocker): + mocker.patch.object(windows_arch_mod.sysconfig, "get_platform", return_value="win-arm64") + + assert windows_arch_mod.windows_python_arch() == "arm64" + + def test_rejects_unknown_sysconfig_tag(self, mocker): + mocker.patch.object(windows_arch_mod.sysconfig, "get_platform", return_value="custom-win") + + with pytest.raises( + UnsupportedWindowsArchError, + match=r"Unsupported Windows Python platform tag: 'custom-win'.*win-amd64.*win-arm64", + ) as exc_info: + windows_arch_mod.windows_python_arch() + assert exc_info.value.platform_tag == "custom-win" + + # --------------------------------------------------------------------------- # find_in_site_packages # --------------------------------------------------------------------------- @@ -90,7 +128,7 @@ def test_raise_not_found_empty_messages(self): class TestFindInSitePackages: def test_returns_none_when_no_rel_dirs(self): - desc = _make_desc(site_packages_linux=(), site_packages_windows=()) + desc = _make_desc(site_packages_linux=(), site_packages_windows=WindowsSearchDirs()) result = find_in_site_packages(_ctx(desc)) assert result is None @@ -127,13 +165,72 @@ def test_found_windows(self, mocker, tmp_path): desc = _make_desc( name="cudart", - site_packages_windows=(os.path.join("nvidia", "cuda_runtime", "bin"),), + site_packages_windows=WindowsSearchDirs.common(os.path.join("nvidia", "cuda_runtime", "bin")), ) result = find_in_site_packages(_ctx(desc, platform=WindowsSearchPlatform())) assert result is not None assert result.abs_path == str(dll) assert result.found_via == "site-packages" + def test_found_windows_arm64_prefers_cuda13_arch_dir_to_cuda12(self, mocker, tmp_path): + x86_64_dir = tmp_path / "nvidia" / "cu13" / "bin" / "x86_64" + arm64_dir = tmp_path / "nvidia" / "cu13" / "bin" / "arm64" + cuda12_dir = tmp_path / "nvidia" / "cuda_runtime" / "bin" + x86_64_dir.mkdir(parents=True) + arm64_dir.mkdir(parents=True) + cuda12_dir.mkdir(parents=True) + (x86_64_dir / "cudart64_12.dll").touch() + (cuda12_dir / "cudart64_12.dll").touch() + arm64_dll = arm64_dir / "cudart64_12.dll" + arm64_dll.touch() + + _patch_site_packages_search(mocker, tmp_path) + mocker.patch(f"{_PLAT_MOD}.is_suppressed_dll_file", return_value=False) + + desc = LIB_DESCRIPTORS["cudart"] + result = find_in_site_packages(_ctx(desc, platform=WindowsSearchPlatform(target_arch="arm64"))) + assert result is not None + assert result.abs_path == str(arm64_dll) + assert result.found_via == "site-packages" + + def test_found_windows_x64_prefers_cuda13_arch_dir_to_cuda12(self, mocker, tmp_path): + x86_64_dir = tmp_path / "nvidia" / "cu13" / "bin" / "x86_64" + arm64_dir = tmp_path / "nvidia" / "cu13" / "bin" / "arm64" + cuda12_dir = tmp_path / "nvidia" / "cuda_runtime" / "bin" + x86_64_dir.mkdir(parents=True) + arm64_dir.mkdir(parents=True) + cuda12_dir.mkdir(parents=True) + x86_64_dll = x86_64_dir / "cudart64_12.dll" + x86_64_dll.touch() + (arm64_dir / "cudart64_12.dll").touch() + (cuda12_dir / "cudart64_12.dll").touch() + + _patch_site_packages_search(mocker, tmp_path) + mocker.patch(f"{_PLAT_MOD}.is_suppressed_dll_file", return_value=False) + + desc = LIB_DESCRIPTORS["cudart"] + result = find_in_site_packages(_ctx(desc, platform=WindowsSearchPlatform(target_arch="x64"))) + assert result is not None + assert result.abs_path == str(x86_64_dll) + assert result.found_via == "site-packages" + + @pytest.mark.parametrize("target_arch", ["x64", "arm64"]) + @pytest.mark.agent_authored(model="gpt-5") + def test_found_windows_uses_cuda12_when_cuda13_is_absent(self, mocker, tmp_path, target_arch): + cuda12_dir = tmp_path / "nvidia" / "cuda_runtime" / "bin" + cuda12_dir.mkdir(parents=True) + cuda12_dll = cuda12_dir / "cudart64_12.dll" + cuda12_dll.touch() + + _patch_site_packages_search(mocker, tmp_path) + mocker.patch(f"{_PLAT_MOD}.is_suppressed_dll_file", return_value=False) + + desc = LIB_DESCRIPTORS["cudart"] + result = find_in_site_packages(_ctx(desc, platform=WindowsSearchPlatform(target_arch=target_arch))) + assert result is not None + assert result.abs_path == str(cuda12_dll) + assert result.found_via == "site-packages" + def test_not_found_appends_error(self, mocker, tmp_path): empty_dir = tmp_path / "nvidia" / "cuda_runtime" / "lib" empty_dir.mkdir(parents=True) @@ -246,6 +343,22 @@ def test_found_windows(self, mocker, tmp_path): assert result.abs_path == str(dll) assert result.found_via == "conda" + def test_found_windows_arm64_prefers_arch_dir(self, mocker, tmp_path): + x64_dir = tmp_path / "Library" / "bin" / "x64" + arm64_dir = tmp_path / "Library" / "bin" / "arm64" + x64_dir.mkdir(parents=True) + arm64_dir.mkdir(parents=True) + (x64_dir / "cudart64_12.dll").touch() + arm64_dll = arm64_dir / "cudart64_12.dll" + arm64_dll.touch() + + mocker.patch.dict(os.environ, {"CONDA_PREFIX": str(tmp_path)}) + + result = find_in_conda(_ctx(platform=WindowsSearchPlatform(target_arch="arm64"))) + assert result is not None + assert result.abs_path == str(arm64_dll) + assert result.found_via == "conda" + # The next three tests cover the Linux glob fallback in # cuda.pathfinder._dynamic_libs.search_platform.LinuxSearchPlatform.find_in_lib_dir, # which is exercised by find_in_conda (and find_in_cuda_path) when the @@ -331,6 +444,22 @@ def test_found_windows(self, mocker, tmp_path): assert result.abs_path == str(dll) assert result.found_via == "CUDA_PATH" + def test_found_windows_arm64_prefers_arch_dir(self, mocker, tmp_path): + x64_dir = tmp_path / "bin" / "x64" + arm64_dir = tmp_path / "bin" / "arm64" + x64_dir.mkdir(parents=True) + arm64_dir.mkdir(parents=True) + (x64_dir / "cudart64_12.dll").touch() + arm64_dll = arm64_dir / "cudart64_12.dll" + arm64_dll.touch() + + mocker.patch(f"{_STEPS_MOD}.get_cuda_path_or_home", return_value=str(tmp_path)) + + result = find_in_cuda_path(_ctx(platform=WindowsSearchPlatform(target_arch="arm64"))) + assert result is not None + assert result.abs_path == str(arm64_dll) + assert result.found_via == "CUDA_PATH" + # --------------------------------------------------------------------------- # run_find_steps @@ -394,13 +523,37 @@ def test_nvvm_has_custom_linux_paths(self): def test_nvvm_has_custom_windows_paths(self): desc = LIB_DESCRIPTORS["nvvm"] - assert desc.anchor_rel_dirs_windows == ("nvvm/bin/*", "nvvm/bin") + assert desc.anchor_rel_dirs_windows.for_arch("x64") == ("nvvm/bin/x64", "nvvm/bin") + assert desc.anchor_rel_dirs_windows.for_arch("arm64") == ("nvvm/bin/arm64", "nvvm/bin") @pytest.mark.parametrize("libname", ["cudart", "cublas", "nvrtc"]) def test_regular_ctk_libs_use_defaults(self, libname): desc = LIB_DESCRIPTORS[libname] assert desc.anchor_rel_dirs_linux == ("lib64", "lib") - assert desc.anchor_rel_dirs_windows == ("bin/x64", "bin") + assert desc.anchor_rel_dirs_windows.for_arch("x64") == ("bin/x64", "bin") + assert desc.anchor_rel_dirs_windows.for_arch("arm64") == ("bin/arm64", "bin") + + def test_windows_anchor_dirs_select_arm64(self): + desc = _make_desc( + anchor_rel_dirs_windows=WindowsSearchDirs( + paths=( + WindowsSearchDir(x64="bin/x64", arm64="bin/arm64"), + WindowsSearchDir.common("bin"), + ) + ) + ) + assert WindowsSearchPlatform(target_arch="arm64").anchor_rel_dirs(desc) == ("bin/arm64", "bin") + + def test_windows_anchor_dirs_select_x64(self): + desc = _make_desc( + anchor_rel_dirs_windows=WindowsSearchDirs( + paths=( + WindowsSearchDir(x64="bin/x64", arm64="bin/arm64"), + WindowsSearchDir.common("bin"), + ) + ) + ) + assert WindowsSearchPlatform(target_arch="x64").anchor_rel_dirs(desc) == ("bin/x64", "bin") def test_find_lib_dir_uses_descriptor_linux(self, tmp_path): (tmp_path / "nvvm" / "lib64").mkdir(parents=True) @@ -413,11 +566,36 @@ def test_find_lib_dir_uses_descriptor_linux(self, tmp_path): def test_find_lib_dir_uses_descriptor_windows(self, tmp_path): (tmp_path / "nvvm" / "bin").mkdir(parents=True) - desc = _make_desc(name="nvvm", anchor_rel_dirs_windows=("nvvm/bin/*", "nvvm/bin")) + desc = _make_desc( + name="nvvm", + anchor_rel_dirs_windows=WindowsSearchDirs( + paths=( + WindowsSearchDir(x64="nvvm/bin/x64", arm64="nvvm/bin/arm64"), + WindowsSearchDir.common("nvvm/bin"), + ) + ), + ) result = _find_lib_dir_using_anchor(desc, WindowsSearchPlatform(), str(tmp_path)) assert result is not None assert result.endswith(os.path.join("nvvm", "bin")) + def test_find_lib_dir_windows_arm64_uses_arm64_anchor(self, tmp_path): + (tmp_path / "nvvm" / "bin" / "x64").mkdir(parents=True) + (tmp_path / "nvvm" / "bin" / "arm64").mkdir(parents=True) + + desc = _make_desc( + name="nvvm", + anchor_rel_dirs_windows=WindowsSearchDirs( + paths=( + WindowsSearchDir(x64="nvvm/bin/x64", arm64="nvvm/bin/arm64"), + WindowsSearchDir.common("nvvm/bin"), + ) + ), + ) + result = _find_lib_dir_using_anchor(desc, WindowsSearchPlatform(target_arch="arm64"), str(tmp_path)) + assert result is not None + assert result.endswith(os.path.join("nvvm", "bin", "arm64")) + def test_find_lib_dir_returns_none_when_no_match(self, tmp_path): desc = _make_desc(anchor_rel_dirs_linux=("nonexistent",)) assert _find_lib_dir_using_anchor(desc, LinuxSearchPlatform(), str(tmp_path)) is None From 32470721c982f41234504e3b9bb58b793e9abd9d Mon Sep 17 00:00:00 2001 From: isvoid <13521008+isVoid@users.noreply.github.com> Date: Tue, 21 Jul 2026 09:36:56 -0700 Subject: [PATCH 2/2] Avoid Windows architecture detection on Linux --- .../_dynamic_libs/search_platform.py | 16 +++++++++----- cuda_pathfinder/tests/test_search_steps.py | 21 ++++++++++++++----- 2 files changed, 27 insertions(+), 10 deletions(-) diff --git a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_platform.py b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_platform.py index ef93eae1d4d..c84e0c5a2f0 100644 --- a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_platform.py +++ b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_platform.py @@ -13,7 +13,7 @@ import glob import os from collections.abc import Sequence -from dataclasses import dataclass, field +from dataclasses import dataclass from pathlib import PurePath from typing import Protocol, cast @@ -175,19 +175,19 @@ def find_in_lib_dir( @dataclass(frozen=True, slots=True) class WindowsSearchPlatform: - target_arch: str = field(default_factory=windows_python_arch) + 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 desc.site_packages_windows.for_arch(self.target_arch) + 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 desc.anchor_rel_dirs_windows.for_arch(self.target_arch) + return cast(tuple[str, ...], desc.anchor_rel_dirs_windows.for_arch(self.target_arch)) def find_in_site_packages( self, @@ -220,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() diff --git a/cuda_pathfinder/tests/test_search_steps.py b/cuda_pathfinder/tests/test_search_steps.py index c5f935062c7..1c34fb076ee 100644 --- a/cuda_pathfinder/tests/test_search_steps.py +++ b/cuda_pathfinder/tests/test_search_steps.py @@ -10,6 +10,7 @@ import pytest from cuda.pathfinder import UnsupportedWindowsArchError +from cuda.pathfinder._dynamic_libs import search_platform as search_platform_mod from cuda.pathfinder._dynamic_libs.descriptor_catalog import WindowsSearchDir, WindowsSearchDirs from cuda.pathfinder._dynamic_libs.lib_descriptor import LIB_DESCRIPTORS, LibDescriptor from cuda.pathfinder._dynamic_libs.load_dl_common import DynamicLibNotFoundError @@ -78,7 +79,7 @@ def test_lib_searched_for_linux(self): assert ctx.lib_searched_for == "libcublas.so" def test_lib_searched_for_windows(self): - ctx = SearchContext(_make_desc(name="cublas"), platform=WindowsSearchPlatform()) + ctx = SearchContext(_make_desc(name="cublas"), platform=WindowsSearchPlatform(target_arch="x64")) assert ctx.lib_searched_for == "cublas*.dll" def test_raise_not_found_includes_messages(self): @@ -100,6 +101,16 @@ def test_raise_not_found_empty_messages(self): class TestWindowsPythonArch: + @pytest.mark.agent_authored(model="gpt-5") + def test_linux_platform_does_not_detect_windows_arch(self, mocker): + mocker.patch.object(search_platform_mod, "IS_WINDOWS", False) + get_windows_arch = mocker.patch.object(search_platform_mod, "windows_python_arch") + + platform = search_platform_mod._platform_for_current_system() + + assert isinstance(platform, LinuxSearchPlatform) + get_windows_arch.assert_not_called() + def test_detects_sysconfig_x64(self, mocker): mocker.patch.object(windows_arch_mod.sysconfig, "get_platform", return_value="win-amd64") @@ -167,7 +178,7 @@ def test_found_windows(self, mocker, tmp_path): name="cudart", site_packages_windows=WindowsSearchDirs.common(os.path.join("nvidia", "cuda_runtime", "bin")), ) - result = find_in_site_packages(_ctx(desc, platform=WindowsSearchPlatform())) + result = find_in_site_packages(_ctx(desc, platform=WindowsSearchPlatform(target_arch="x64"))) assert result is not None assert result.abs_path == str(dll) assert result.found_via == "site-packages" @@ -338,7 +349,7 @@ def test_found_windows(self, mocker, tmp_path): mocker.patch.dict(os.environ, {"CONDA_PREFIX": str(tmp_path)}) - result = find_in_conda(_ctx(platform=WindowsSearchPlatform())) + result = find_in_conda(_ctx(platform=WindowsSearchPlatform(target_arch="x64"))) assert result is not None assert result.abs_path == str(dll) assert result.found_via == "conda" @@ -439,7 +450,7 @@ def test_found_windows(self, mocker, tmp_path): mocker.patch(f"{_STEPS_MOD}.get_cuda_path_or_home", return_value=str(tmp_path)) - result = find_in_cuda_path(_ctx(platform=WindowsSearchPlatform())) + result = find_in_cuda_path(_ctx(platform=WindowsSearchPlatform(target_arch="x64"))) assert result is not None assert result.abs_path == str(dll) assert result.found_via == "CUDA_PATH" @@ -575,7 +586,7 @@ def test_find_lib_dir_uses_descriptor_windows(self, tmp_path): ) ), ) - result = _find_lib_dir_using_anchor(desc, WindowsSearchPlatform(), str(tmp_path)) + result = _find_lib_dir_using_anchor(desc, WindowsSearchPlatform(target_arch="x64"), str(tmp_path)) assert result is not None assert result.endswith(os.path.join("nvvm", "bin"))