diff --git a/cuda_pathfinder/cuda/pathfinder/__init__.py b/cuda_pathfinder/cuda/pathfinder/__init__.py index dc818dfd08f..a64b5ef64c7 100644 --- a/cuda_pathfinder/cuda/pathfinder/__init__.py +++ b/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"``. diff --git a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py index a26862c5435..32f7146eee4 100644 --- a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py +++ b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py @@ -6,11 +6,51 @@ 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 PackagedWith = Literal["ctk", "other", "driver"] +WindowsArch = Literal["x64", "arm64"] + + +@dataclass(frozen=True, slots=True) +class WindowsSearchDirs: + """Ordered Windows search locations grouped by process architecture.""" + + x64: tuple[str, ...] = () + arm64: tuple[str, ...] = () + + @classmethod + def x64_only(cls, *paths: str) -> WindowsSearchDirs: + return cls(x64=paths) + + @classmethod + def arm64_only(cls, *paths: str) -> WindowsSearchDirs: + return cls(arm64=paths) + + def for_arch(self, target_arch: str) -> tuple[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}") + + +DEFAULT_WINDOWS_CTK_ANCHOR_DIRS = WindowsSearchDirs( + x64=("bin/x64", "bin"), + arm64=("bin/arm64", "bin"), +) + + +def _ctk_windows_wheel_dirs(cuda13_bin_dir: str, cuda12_dir: str) -> WindowsSearchDirs: + """Search CUDA 13 first, with the x64-only CUDA 12 wheel as an x64 fallback.""" + cuda13_bin_path = PurePosixPath(cuda13_bin_dir) + return WindowsSearchDirs( + x64=((cuda13_bin_path / "x86_64").as_posix(), cuda12_dir), + arm64=((cuda13_bin_path / "arm64").as_posix(),), + ) @dataclass(frozen=True, slots=True) @@ -19,11 +59,12 @@ class DescriptorSpec: packaged_with: PackagedWith linux_sonames: tuple[str, ...] = () windows_dlls: tuple[str, ...] = () + supported_windows_arch: tuple[WindowsArch, ...] = () 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 @@ -38,32 +79,36 @@ class DescriptorSpec: packaged_with="ctk", linux_sonames=("libcudart.so.12", "libcudart.so.13"), windows_dlls=("cudart64_12.dll", "cudart64_13.dll"), + supported_windows_arch=("x64", "arm64"), 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", packaged_with="ctk", linux_sonames=("libnvfatbin.so.12", "libnvfatbin.so.13"), windows_dlls=("nvfatbin_120_0.dll", "nvfatbin_130_0.dll"), + supported_windows_arch=("x64", "arm64"), 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", packaged_with="ctk", linux_sonames=("libnvJitLink.so.12", "libnvJitLink.so.13"), windows_dlls=("nvJitLink_120_0.dll", "nvJitLink_130_0.dll"), + supported_windows_arch=("x64", "arm64"), 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", packaged_with="ctk", linux_sonames=("libnvrtc.so.12", "libnvrtc.so.13"), windows_dlls=("nvrtc64_120_0.dll", "nvrtc64_130_0.dll"), + supported_windows_arch=("x64", "arm64"), 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( @@ -71,10 +116,14 @@ class DescriptorSpec: packaged_with="ctk", linux_sonames=("libnvvm.so.4",), windows_dlls=("nvvm64.dll", "nvvm64_40_0.dll", "nvvm70.dll"), + supported_windows_arch=("x64", "arm64"), 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( + x64=("nvvm/bin/x64", "nvvm/bin"), + arm64=("nvvm/bin/arm64", "nvvm/bin"), + ), ctk_root_canary_anchor_libnames=CTK_ROOT_CANARY_ANCHOR_LIBNAMES, ), DescriptorSpec( @@ -82,8 +131,9 @@ class DescriptorSpec: packaged_with="ctk", linux_sonames=("libcublas.so.12", "libcublas.so.13"), windows_dlls=("cublas64_12.dll", "cublas64_13.dll"), + supported_windows_arch=("x64", "arm64"), 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( @@ -91,16 +141,18 @@ class DescriptorSpec: packaged_with="ctk", linux_sonames=("libcublasLt.so.12", "libcublasLt.so.13"), windows_dlls=("cublasLt64_12.dll", "cublasLt64_13.dll"), + supported_windows_arch=("x64", "arm64"), 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", packaged_with="ctk", linux_sonames=("libcufft.so.11", "libcufft.so.12"), windows_dlls=("cufft64_11.dll", "cufft64_12.dll"), + supported_windows_arch=("x64", "arm64"), 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( @@ -108,8 +160,9 @@ class DescriptorSpec: packaged_with="ctk", linux_sonames=("libcufftw.so.11", "libcufftw.so.12"), windows_dlls=("cufftw64_11.dll", "cufftw64_12.dll"), + supported_windows_arch=("x64", "arm64"), 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( @@ -117,16 +170,18 @@ class DescriptorSpec: packaged_with="ctk", linux_sonames=("libcurand.so.10",), windows_dlls=("curand64_10.dll",), + supported_windows_arch=("x64", "arm64"), 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", packaged_with="ctk", linux_sonames=("libcusolver.so.11", "libcusolver.so.12"), windows_dlls=("cusolver64_11.dll", "cusolver64_12.dll"), + supported_windows_arch=("x64", "arm64"), 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( @@ -134,8 +189,9 @@ class DescriptorSpec: packaged_with="ctk", linux_sonames=("libcusolverMg.so.11", "libcusolverMg.so.12"), windows_dlls=("cusolverMg64_11.dll", "cusolverMg64_12.dll"), + supported_windows_arch=("x64", "arm64"), 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( @@ -143,8 +199,9 @@ class DescriptorSpec: packaged_with="ctk", linux_sonames=("libcusparse.so.12",), windows_dlls=("cusparse64_12.dll",), + supported_windows_arch=("x64", "arm64"), 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( @@ -152,16 +209,18 @@ class DescriptorSpec: packaged_with="ctk", linux_sonames=("libnppc.so.12", "libnppc.so.13"), windows_dlls=("nppc64_12.dll", "nppc64_13.dll"), + supported_windows_arch=("x64", "arm64"), 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", packaged_with="ctk", linux_sonames=("libnppial.so.12", "libnppial.so.13"), windows_dlls=("nppial64_12.dll", "nppial64_13.dll"), + supported_windows_arch=("x64", "arm64"), 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( @@ -169,8 +228,9 @@ class DescriptorSpec: packaged_with="ctk", linux_sonames=("libnppicc.so.12", "libnppicc.so.13"), windows_dlls=("nppicc64_12.dll", "nppicc64_13.dll"), + supported_windows_arch=("x64", "arm64"), 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( @@ -178,8 +238,9 @@ class DescriptorSpec: packaged_with="ctk", linux_sonames=("libnppidei.so.12", "libnppidei.so.13"), windows_dlls=("nppidei64_12.dll", "nppidei64_13.dll"), + supported_windows_arch=("x64", "arm64"), 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( @@ -187,8 +248,9 @@ class DescriptorSpec: packaged_with="ctk", linux_sonames=("libnppif.so.12", "libnppif.so.13"), windows_dlls=("nppif64_12.dll", "nppif64_13.dll"), + supported_windows_arch=("x64", "arm64"), 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( @@ -196,8 +258,9 @@ class DescriptorSpec: packaged_with="ctk", linux_sonames=("libnppig.so.12", "libnppig.so.13"), windows_dlls=("nppig64_12.dll", "nppig64_13.dll"), + supported_windows_arch=("x64", "arm64"), 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( @@ -205,8 +268,9 @@ class DescriptorSpec: packaged_with="ctk", linux_sonames=("libnppim.so.12", "libnppim.so.13"), windows_dlls=("nppim64_12.dll", "nppim64_13.dll"), + supported_windows_arch=("x64", "arm64"), 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( @@ -214,8 +278,9 @@ class DescriptorSpec: packaged_with="ctk", linux_sonames=("libnppist.so.12", "libnppist.so.13"), windows_dlls=("nppist64_12.dll", "nppist64_13.dll"), + supported_windows_arch=("x64", "arm64"), 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( @@ -223,8 +288,9 @@ class DescriptorSpec: packaged_with="ctk", linux_sonames=("libnppisu.so.12", "libnppisu.so.13"), windows_dlls=("nppisu64_12.dll", "nppisu64_13.dll"), + supported_windows_arch=("x64", "arm64"), 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( @@ -232,8 +298,9 @@ class DescriptorSpec: packaged_with="ctk", linux_sonames=("libnppitc.so.12", "libnppitc.so.13"), windows_dlls=("nppitc64_12.dll", "nppitc64_13.dll"), + supported_windows_arch=("x64", "arm64"), 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( @@ -241,8 +308,9 @@ class DescriptorSpec: packaged_with="ctk", linux_sonames=("libnpps.so.12", "libnpps.so.13"), windows_dlls=("npps64_12.dll", "npps64_13.dll"), + supported_windows_arch=("x64", "arm64"), 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( @@ -250,8 +318,9 @@ class DescriptorSpec: packaged_with="ctk", linux_sonames=("libnvblas.so.12", "libnvblas.so.13"), windows_dlls=("nvblas64_12.dll", "nvblas64_13.dll"), + supported_windows_arch=("x64", "arm64"), 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( @@ -259,8 +328,9 @@ class DescriptorSpec: packaged_with="ctk", linux_sonames=("libnvjpeg.so.12", "libnvjpeg.so.13"), windows_dlls=("nvjpeg64_12.dll", "nvjpeg64_13.dll"), + supported_windows_arch=("x64", "arm64"), 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", @@ -290,10 +360,14 @@ class DescriptorSpec: "cupti64_2023.1.1.dll", "cupti64_2022.4.1.dll", ), + supported_windows_arch=("x64", "arm64"), 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( + x64=("extras/CUPTI/lib64", "bin/x64", "bin"), + arm64=("extras/CUPTI/lib64", "bin/arm64", "bin"), + ), ctk_root_canary_anchor_libnames=CTK_ROOT_CANARY_ANCHOR_LIBNAMES, ), DescriptorSpec( @@ -301,12 +375,11 @@ class DescriptorSpec: packaged_with="ctk", linux_sonames=("libcudla.so.1",), windows_dlls=("cudla.dll",), + supported_windows_arch=("arm64",), 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"), + anchor_rel_dirs_windows=WindowsSearchDirs.arm64_only("bin/arm64"), ), # ----------------------------------------------------------------------- # Third-party / separately packaged libraries @@ -338,8 +411,12 @@ class DescriptorSpec: packaged_with="other", linux_sonames=("libmathdx.so.0",), windows_dlls=("mathdx64_0.dll",), + supported_windows_arch=("x64", "arm64"), site_packages_linux=("nvidia/cu13/lib", "nvidia/cu12/lib"), - site_packages_windows=("nvidia/cu13/bin", "nvidia/cu12/bin"), + site_packages_windows=WindowsSearchDirs( + x64=("nvidia/cu13/bin", "nvidia/cu12/bin"), + arm64=("nvidia/cu13/bin", "nvidia/cu12/bin"), + ), dependencies=("nvrtc",), ), DescriptorSpec( @@ -347,8 +424,12 @@ class DescriptorSpec: packaged_with="other", linux_sonames=("libcudss.so.0",), windows_dlls=("cudss64_0.dll",), + supported_windows_arch=("x64", "arm64"), site_packages_linux=("nvidia/cu13/lib", "nvidia/cu12/lib"), - site_packages_windows=("nvidia/cu13/bin", "nvidia/cu12/bin"), + site_packages_windows=WindowsSearchDirs( + x64=("nvidia/cu13/bin", "nvidia/cu12/bin"), + arm64=("nvidia/cu13/bin", "nvidia/cu12/bin"), + ), dependencies=("cublas", "cublasLt"), ), DescriptorSpec( @@ -356,16 +437,21 @@ class DescriptorSpec: packaged_with="other", linux_sonames=("libcusparseLt.so.0",), windows_dlls=("cusparseLt.dll",), + supported_windows_arch=("x64", "arm64"), site_packages_linux=("nvidia/cu13/lib", "nvidia/cusparselt/lib"), - site_packages_windows=("nvidia/cu13/bin/x64", "nvidia/cusparselt/bin"), + site_packages_windows=WindowsSearchDirs( + x64=("nvidia/cu13/bin/x64", "nvidia/cusparselt/bin"), + arm64=("nvidia/cu13/bin/x64", "nvidia/cusparselt/bin"), + ), ), DescriptorSpec( name="cutensor", packaged_with="other", linux_sonames=("libcutensor.so.2",), windows_dlls=("cutensor.dll",), + supported_windows_arch=("x64", "arm64"), site_packages_linux=("cutensor/lib",), - site_packages_windows=("cutensor/bin",), + site_packages_windows=WindowsSearchDirs(x64=("cutensor/bin",), arm64=("cutensor/bin",)), dependencies=("cublasLt",), ), DescriptorSpec( @@ -373,8 +459,9 @@ class DescriptorSpec: packaged_with="other", linux_sonames=("libcutensorMg.so.2",), windows_dlls=("cutensorMg.dll",), + supported_windows_arch=("x64", "arm64"), site_packages_linux=("cutensor/lib",), - site_packages_windows=("cutensor/bin",), + site_packages_windows=WindowsSearchDirs(x64=("cutensor/bin",), arm64=("cutensor/bin",)), dependencies=("cutensor", "cublasLt"), ), DescriptorSpec( @@ -452,6 +539,7 @@ class DescriptorSpec: packaged_with="driver", linux_sonames=("libcuda.so.1",), windows_dlls=("nvcuda.dll",), + supported_windows_arch=("x64", "arm64"), ), DescriptorSpec( name="nvcudla", @@ -463,5 +551,6 @@ class DescriptorSpec: packaged_with="driver", linux_sonames=("libnvidia-ml.so.1",), windows_dlls=("nvml.dll",), + supported_windows_arch=("x64", "arm64"), ), ) diff --git a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_platform.py b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_platform.py index 37fd6eb1700..c84e0c5a2f0 100644 --- a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_platform.py +++ b/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.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 + 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() 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..47417bcd357 100644 --- a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/supported_nvidia_libs.py +++ b/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: diff --git a/cuda_pathfinder/cuda/pathfinder/_utils/platform_aware.py b/cuda_pathfinder/cuda/pathfinder/_utils/platform_aware.py index 72ecbc53593..af0610a6cdb 100644 --- a/cuda_pathfinder/cuda/pathfinder/_utils/platform_aware.py +++ b/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: 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..7123b04face --- /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 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) diff --git a/cuda_pathfinder/docs/source/api.rst b/cuda_pathfinder/docs/source/api.rst index e49478c09ec..f65014923f9 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 + UnsupportedArchError SUPPORTED_HEADERS_CTK find_nvidia_header_directory diff --git a/cuda_pathfinder/docs/source/release/1.6.0-notes.rst b/cuda_pathfinder/docs/source/release/1.6.0-notes.rst index a02ac297dbd..86363f1843e 100644 --- a/cuda_pathfinder/docs/source/release/1.6.0-notes.rst +++ b/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 `_) diff --git a/cuda_pathfinder/tests/test_catalog_writer.py b/cuda_pathfinder/tests/test_catalog_writer.py new file mode 100644 index 00000000000..020baccb70a --- /dev/null +++ b/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 + ) 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..788fc18ecf6 100644 --- a/cuda_pathfinder/tests/test_descriptor_catalog.py +++ b/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: diff --git a/cuda_pathfinder/tests/test_lib_descriptor.py b/cuda_pathfinder/tests/test_lib_descriptor.py index cda96131e13..8715f9181cc 100644 --- a/cuda_pathfinder/tests/test_lib_descriptor.py +++ b/cuda_pathfinder/tests/test_lib_descriptor.py @@ -13,10 +13,21 @@ LIBNAMES_REQUIRING_RTLD_DEEPBIND, SITE_PACKAGES_LIBDIRS_LINUX, SITE_PACKAGES_LIBDIRS_WINDOWS, + SITE_PACKAGES_LIBDIRS_WINDOWS_ARM64, + 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_X64, SUPPORTED_LIBNAMES, + SUPPORTED_LIBNAMES_LINUX, + SUPPORTED_LIBNAMES_WINDOWS, + SUPPORTED_LIBNAMES_WINDOWS_ARM64, + SUPPORTED_LIBNAMES_WINDOWS_X64, SUPPORTED_LINUX_SONAMES, SUPPORTED_WINDOWS_DLLS, ) +from cuda.pathfinder._utils.platform_aware import IS_WINDOWS, IS_WINDOWS_ARM64, IS_WINDOWS_X64 # --------------------------------------------------------------------------- # Registry completeness @@ -56,9 +67,49 @@ def test_site_packages_linux_match(name): assert LIB_DESCRIPTORS[name].site_packages_linux == SITE_PACKAGES_LIBDIRS_LINUX.get(name, ()) +@pytest.mark.parametrize( + ("target_arch", "site_packages_libdirs"), + [ + ("x64", SITE_PACKAGES_LIBDIRS_WINDOWS_X64), + ("arm64", SITE_PACKAGES_LIBDIRS_WINDOWS_ARM64), + ], +) @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, ()) +@pytest.mark.agent_authored(model="gpt-5") +def test_site_packages_windows_match(name, target_arch, site_packages_libdirs): + assert LIB_DESCRIPTORS[name].site_packages_windows.for_arch(target_arch) == site_packages_libdirs.get(name, ()) + + +@pytest.mark.agent_authored(model="gpt-5") +def test_legacy_site_packages_windows_tables_are_x64_aliases(): + assert SITE_PACKAGES_LIBDIRS_WINDOWS_CTK is SITE_PACKAGES_LIBDIRS_WINDOWS_CTK_X64 + assert SITE_PACKAGES_LIBDIRS_WINDOWS_OTHER is SITE_PACKAGES_LIBDIRS_WINDOWS_OTHER_X64 + assert SITE_PACKAGES_LIBDIRS_WINDOWS is SITE_PACKAGES_LIBDIRS_WINDOWS_X64 + + +@pytest.mark.agent_authored(model="gpt-5") +def test_legacy_supported_libnames_windows_is_x64_alias(): + assert SUPPORTED_LIBNAMES_WINDOWS is SUPPORTED_LIBNAMES_WINDOWS_X64 + + +@pytest.mark.agent_authored(model="gpt-5") +def test_supported_libnames_selects_current_platform_and_arch(): + if not IS_WINDOWS: + expected = SUPPORTED_LIBNAMES_LINUX + elif IS_WINDOWS_X64: + expected = SUPPORTED_LIBNAMES_WINDOWS_X64 + else: + assert IS_WINDOWS_ARM64 + expected = SUPPORTED_LIBNAMES_WINDOWS_ARM64 + assert SUPPORTED_LIBNAMES is expected + + +@pytest.mark.agent_authored(model="gpt-5") +def test_arch_specific_ctk_libname_projections(): + assert "cudla" not in SUPPORTED_LIBNAMES_WINDOWS_X64 + assert "cudla" in SUPPORTED_LIBNAMES_WINDOWS_ARM64 + assert "nvjpeg" in SUPPORTED_LIBNAMES_WINDOWS_X64 + assert "nvjpeg" in SUPPORTED_LIBNAMES_WINDOWS_ARM64 @pytest.mark.parametrize("name", sorted(LIB_DESCRIPTORS)) diff --git a/cuda_pathfinder/tests/test_load_nvidia_dynamic_lib.py b/cuda_pathfinder/tests/test_load_nvidia_dynamic_lib.py index 3e240dcf468..8c35663ebbe 100644 --- a/cuda_pathfinder/tests/test_load_nvidia_dynamic_lib.py +++ b/cuda_pathfinder/tests/test_load_nvidia_dynamic_lib.py @@ -31,10 +31,22 @@ def test_supported_libnames_linux_sonames_consistency(): ) -def test_supported_libnames_windows_dlls_consistency(): - assert tuple(sorted(supported_nvidia_libs.SUPPORTED_LIBNAMES_WINDOWS)) == tuple( - sorted(supported_nvidia_libs.SUPPORTED_WINDOWS_DLLS_CTK.keys()) - ) +@pytest.mark.parametrize( + ("target_arch", "supported_libnames"), + [ + ("x64", supported_nvidia_libs.SUPPORTED_LIBNAMES_WINDOWS_X64), + ("arm64", supported_nvidia_libs.SUPPORTED_LIBNAMES_WINDOWS_ARM64), + ], +) +@pytest.mark.agent_authored(model="gpt-5") +def test_supported_libnames_windows_dlls_consistency(target_arch, supported_libnames): + expected = { + desc.name + for desc in load_nvidia_dynamic_lib_module.LIB_DESCRIPTORS.values() + if desc.packaged_with == "ctk" and target_arch in desc.supported_windows_arch + } + assert set(supported_libnames) == expected + assert expected <= set(supported_nvidia_libs.SUPPORTED_WINDOWS_DLLS_CTK) def test_supported_libnames_linux_site_packages_libdirs_ctk_consistency(): @@ -43,12 +55,20 @@ def test_supported_libnames_linux_site_packages_libdirs_ctk_consistency(): ) -def test_supported_libnames_windows_site_packages_libdirs_ctk_consistency(): +@pytest.mark.parametrize( + "site_packages_libdirs", + [ + supported_nvidia_libs.SITE_PACKAGES_LIBDIRS_WINDOWS_CTK_X64, + supported_nvidia_libs.SITE_PACKAGES_LIBDIRS_WINDOWS_CTK_ARM64, + ], +) +@pytest.mark.agent_authored(model="gpt-5") +def test_supported_libnames_windows_site_packages_libdirs_ctk_consistency(site_packages_libdirs): # Not every Windows CTK library ships in a pip wheel (e.g. cudla is loaded # from the local CUDA Toolkit only), so a library may legitimately omit # site_packages_windows. Only assert that every site-packages entry maps to # a supported Windows libname, not the other way around. - site_packages_libnames = set(supported_nvidia_libs.SITE_PACKAGES_LIBDIRS_WINDOWS_CTK.keys()) + site_packages_libnames = set(site_packages_libdirs) supported_libnames = set(supported_nvidia_libs.SUPPORTED_LIBNAMES_WINDOWS) assert site_packages_libnames <= supported_libnames diff --git a/cuda_pathfinder/tests/test_search_steps.py b/cuda_pathfinder/tests/test_search_steps.py index 1b881707dfb..d223c2f3354 100644 --- a/cuda_pathfinder/tests/test_search_steps.py +++ b/cuda_pathfinder/tests/test_search_steps.py @@ -9,6 +9,9 @@ import pytest +from cuda.pathfinder import UnsupportedArchError +from cuda.pathfinder._dynamic_libs import search_platform as search_platform_mod +from cuda.pathfinder._dynamic_libs.descriptor_catalog import 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 +26,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 +44,10 @@ 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( + x64=(os.path.join("nvidia", "cuda_runtime", "bin"),), + arm64=(os.path.join("nvidia", "cuda_runtime", "bin"),), + ), } defaults.update(overrides) return LibDescriptor(**defaults) @@ -52,6 +59,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 # --------------------------------------------------------------------------- @@ -67,7 +82,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): @@ -83,6 +98,43 @@ def test_raise_not_found_empty_messages(self): ctx.raise_not_found() +# --------------------------------------------------------------------------- +# Windows Python architecture detection +# --------------------------------------------------------------------------- + + +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") + + 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( + UnsupportedArchError, + 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 +142,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 +179,91 @@ 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( + x64=(os.path.join("nvidia", "cuda_runtime", "bin"),), + arm64=(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" + 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.agent_authored(model="gpt-5") + def test_found_windows_x64_uses_cuda12_when_cuda13_is_absent(self, mocker, tmp_path): + 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="x64"))) + assert result is not None + assert result.abs_path == str(cuda12_dll) + assert result.found_via == "site-packages" + + @pytest.mark.agent_authored(model="gpt-5") + def test_found_windows_arm64_skips_cuda12_when_cuda13_is_absent(self, mocker, tmp_path): + cuda12_dir = tmp_path / "nvidia" / "cuda_runtime" / "bin" + cuda12_dir.mkdir(parents=True) + (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"] + platform = WindowsSearchPlatform(target_arch="arm64") + assert platform.site_packages_rel_dirs(desc) == ("nvidia/cu13/bin/arm64",) + + result = find_in_site_packages(_ctx(desc, platform=platform)) + + assert result is None + def test_not_found_appends_error(self, mocker, tmp_path): empty_dir = tmp_path / "nvidia" / "cuda_runtime" / "lib" empty_dir.mkdir(parents=True) @@ -241,11 +371,27 @@ 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" + 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 @@ -326,11 +472,27 @@ 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" + 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 @@ -388,19 +550,51 @@ def test_early_and_late_are_disjoint(self): class TestAnchorRelDirs: """Verify that descriptor anchor paths drive directory resolution.""" + @pytest.mark.agent_authored(model="gpt-5") + def test_windows_search_dirs_arch_only_constructors(self): + assert WindowsSearchDirs.x64_only("first", "second") == WindowsSearchDirs(x64=("first", "second")) + assert WindowsSearchDirs.arm64_only("first", "second") == WindowsSearchDirs(arm64=("first", "second")) + def test_nvvm_has_custom_linux_paths(self): desc = LIB_DESCRIPTORS["nvvm"] assert desc.anchor_rel_dirs_linux == ("nvvm/lib64",) 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") + + @pytest.mark.agent_authored(model="gpt-5") + def test_cudla_uses_arm64_only_windows_anchor(self): + desc = LIB_DESCRIPTORS["cudla"] + + assert desc.anchor_rel_dirs_windows.for_arch("x64") == () + assert desc.anchor_rel_dirs_windows.for_arch("arm64") == ("bin/arm64",) + + def test_windows_anchor_dirs_select_arm64(self): + desc = _make_desc( + anchor_rel_dirs_windows=WindowsSearchDirs( + x64=("bin/x64", "bin"), + arm64=("bin/arm64", "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( + x64=("bin/x64", "bin"), + arm64=("bin/arm64", "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 +607,32 @@ 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")) - result = _find_lib_dir_using_anchor(desc, WindowsSearchPlatform(), str(tmp_path)) + desc = _make_desc( + name="nvvm", + anchor_rel_dirs_windows=WindowsSearchDirs( + x64=("nvvm/bin/x64", "nvvm/bin"), + arm64=("nvvm/bin/arm64", "nvvm/bin"), + ), + ) + 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")) + 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( + x64=("nvvm/bin/x64", "nvvm/bin"), + arm64=("nvvm/bin/arm64", "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 diff --git a/toolshed/_catalog_writer.py b/toolshed/_catalog_writer.py index b41fb5838dd..ebaa33cfab5 100644 --- a/toolshed/_catalog_writer.py +++ b/toolshed/_catalog_writer.py @@ -26,6 +26,7 @@ from cuda.pathfinder._dynamic_libs.descriptor_catalog import ( # noqa: E402 DESCRIPTOR_CATALOG, DescriptorSpec, + WindowsSearchDirs, ) CATALOG_PATH = _PATHFINDER_ROOT / "cuda" / "pathfinder" / "_dynamic_libs" / "descriptor_catalog.py" @@ -63,6 +64,15 @@ def _render_tuple(values: tuple[str, ...]) -> str: return "(" + ", ".join(_quote(v) for v in values) + ")" +def _render_windows_search_dirs(value: WindowsSearchDirs) -> str: + fields: list[str] = [] + if value.x64: + fields.append(f"x64={_render_tuple(value.x64)}") + if value.arm64: + fields.append(f"arm64={_render_tuple(value.arm64)}") + return f"WindowsSearchDirs({', '.join(fields)})" + + def _render_spec(spec: DescriptorSpec) -> str: """Render a single DescriptorSpec constructor call, omitting default-valued fields.""" lines = [ @@ -74,13 +84,16 @@ def _render_spec(spec: DescriptorSpec) -> str: tuple_fields = [ "linux_sonames", "windows_dlls", + "supported_windows_arch", "site_packages_linux", - "site_packages_windows", "dependencies", "anchor_rel_dirs_linux", - "anchor_rel_dirs_windows", "ctk_root_canary_anchor_libnames", ] + windows_search_dir_fields = [ + "site_packages_windows", + "anchor_rel_dirs_windows", + ] bool_fields = [ "requires_add_dll_directory", "requires_rtld_deepbind", @@ -92,6 +105,12 @@ def _render_spec(spec: DescriptorSpec) -> str: if value != default: lines.append(f" {field}={_render_tuple(value)},") + for field in windows_search_dir_fields: + value = getattr(spec, field) + default = getattr(_DEFAULTS, field) + if value != default: + lines.append(f" {field}={_render_windows_search_dirs(value)},") + for field in bool_fields: value = getattr(spec, field) default = getattr(_DEFAULTS, field) @@ -116,6 +135,36 @@ def render_catalog(specs: tuple[DescriptorSpec, ...]) -> str: from typing import Literal PackagedWith = Literal["ctk", "other", "driver"] +WindowsArch = Literal["x64", "arm64"] + + +@dataclass(frozen=True, slots=True) +class WindowsSearchDirs: + """Ordered Windows search locations grouped by process architecture.""" + + x64: tuple[str, ...] = () + arm64: tuple[str, ...] = () + + @classmethod + def x64_only(cls, *paths: str) -> WindowsSearchDirs: + return cls(x64=paths) + + @classmethod + def arm64_only(cls, *paths: str) -> WindowsSearchDirs: + return cls(arm64=paths) + + def for_arch(self, target_arch: str) -> tuple[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}") + + +DEFAULT_WINDOWS_CTK_ANCHOR_DIRS = WindowsSearchDirs( + x64=("bin/x64", "bin"), + arm64=("bin/arm64", "bin"), +) @dataclass(frozen=True, slots=True) @@ -124,11 +173,12 @@ class DescriptorSpec: packaged_with: PackagedWith linux_sonames: tuple[str, ...] = () windows_dlls: tuple[str, ...] = () + supported_windows_arch: tuple[WindowsArch, ...] = () 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 diff --git a/toolshed/make_site_packages_libdirs.py b/toolshed/make_site_packages_libdirs.py index e1cbcb28825..2a3eb6d3438 100755 --- a/toolshed/make_site_packages_libdirs.py +++ b/toolshed/make_site_packages_libdirs.py @@ -13,6 +13,7 @@ from __future__ import annotations import argparse +import dataclasses import os import re import sys @@ -85,6 +86,12 @@ def main() -> None: ) ap.add_argument("platform", choices=["linux", "windows"]) ap.add_argument("path", help="Text file with one library path per line") + ap.add_argument( + "--windows-arch", + choices=["x64", "arm64"], + default="x64", + help="Windows collection architecture (default: x64)", + ) args = ap.parse_args() with open(args.path, encoding="utf-8") as f: @@ -98,12 +105,18 @@ def main() -> None: field = "site_packages_windows" catalog = load_catalog() - catalog_names = {spec.name for spec in catalog} + catalog_by_name = {spec.name: spec for spec in catalog} updates: dict[str, dict[str, object]] = {} for name, dirs in parsed.items(): - if name in catalog_names: - updates[name] = {field: tuple(sorted(dirs))} + if name not in catalog_by_name: + continue + paths = tuple(sorted(dirs)) + if args.platform == "windows": + windows_dirs = catalog_by_name[name].site_packages_windows + updates[name] = {field: dataclasses.replace(windows_dirs, **{args.windows_arch: paths})} + else: + updates[name] = {field: paths} if updates: write_catalog(update_specs(catalog, updates)) @@ -112,7 +125,7 @@ def main() -> None: else: print("No matching libraries found.") - unmatched = set(parsed.keys()) - catalog_names + unmatched = set(parsed) - set(catalog_by_name) if unmatched: print(f"\nLibraries not in catalog ({len(unmatched)}):") for name in sorted(unmatched):