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..e39046eec70 100644 --- a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py +++ b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py @@ -6,11 +6,53 @@ 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}") + + +# Windows CTK before 13.4 was x64-only and used the common bin directory. +# Native ARM64 support starts with the architecture-qualified 13.4 layout. +DEFAULT_WINDOWS_CTK_ANCHOR_DIRS = WindowsSearchDirs( + x64=("bin/x64", "bin"), + arm64=("bin/arm64",), +) + + +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,14 +61,16 @@ 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 + requires_windows_binary_arch_check: bool = False DESCRIPTOR_CATALOG: tuple[DescriptorSpec, ...] = ( @@ -38,32 +82,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,19 +119,31 @@ 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"), + # CTK 13.4 installs the ARM64 DLL directly in nvvm/bin, while x64 + # uses nvvm/bin/x64. Older x64 toolkits also used nvvm/bin, so the + # binary in the unqualified directory must be checked at runtime. + anchor_rel_dirs_windows=WindowsSearchDirs( + x64=("nvvm/bin/x64", "nvvm/bin"), + arm64=("nvvm/bin",), + ), ctk_root_canary_anchor_libnames=CTK_ROOT_CANARY_ANCHOR_LIBNAMES, + # requires_windows_binary_arch_check disambiguates pre-13.4 x64 DLLs + # from 13.4+ Arm64 DLLs in nvvm/bin; see + # _utils/windows_arch.py for the validation. + requires_windows_binary_arch_check=True, ), DescriptorSpec( name="cublas", 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 +151,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 +170,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 +180,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 +199,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 +209,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 +219,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 +238,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 +248,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 +258,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 +268,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 +278,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 +288,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 +298,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 +308,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 +318,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 +328,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 +338,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 +370,16 @@ 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"), + # CTK 13.4 uses architecture-qualified CUPTI directories. Older + # Windows CUPTI toolkits were x64-only and used extras/CUPTI/lib64. + anchor_rel_dirs_windows=WindowsSearchDirs( + x64=("extras/CUPTI/lib/x64", "extras/CUPTI/lib64", "bin"), + arm64=("extras/CUPTI/lib/arm64",), + ), ctk_root_canary_anchor_libnames=CTK_ROOT_CANARY_ANCHOR_LIBNAMES, ), DescriptorSpec( @@ -301,12 +387,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 +423,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 +436,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 +449,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/arm64",), + ), ), 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 +471,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 +551,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 +563,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/load_nvidia_dynamic_lib.py b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_nvidia_dynamic_lib.py index 95a71825793..53446107da3 100644 --- a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_nvidia_dynamic_lib.py +++ b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_nvidia_dynamic_lib.py @@ -34,6 +34,7 @@ build_dynamic_lib_subprocess_command, parse_dynamic_lib_subprocess_payload, ) +from cuda.pathfinder._dynamic_libs.supported_nvidia_libs import ALL_AVAILABLE_LIBNAMES from cuda.pathfinder._utils.platform_aware import IS_WINDOWS if TYPE_CHECKING: @@ -42,9 +43,6 @@ # All libnames recognized by load_nvidia_dynamic_lib, across all categories # (CTK, third-party, driver). _ALL_KNOWN_LIBNAMES: frozenset[str] = frozenset(LIB_DESCRIPTORS) -_ALL_SUPPORTED_LIBNAMES: frozenset[str] = frozenset( - name for name, desc in LIB_DESCRIPTORS.items() if (desc.windows_dlls if IS_WINDOWS else desc.linux_sonames) -) _PLATFORM_NAME = "Windows" if IS_WINDOWS else "Linux" _CANARY_PROBE_TIMEOUT_SECONDS = 10.0 @@ -308,9 +306,9 @@ def load_nvidia_dynamic_lib(libname: str) -> LoadedDL: ) if libname not in _ALL_KNOWN_LIBNAMES: raise DynamicLibUnknownError(f"Unknown library name: {libname!r}. Known names: {sorted(_ALL_KNOWN_LIBNAMES)}") - if libname not in _ALL_SUPPORTED_LIBNAMES: + if libname not in ALL_AVAILABLE_LIBNAMES: raise DynamicLibNotAvailableError( f"Library name {libname!r} is known but not available on {_PLATFORM_NAME}. " - f"Supported names on {_PLATFORM_NAME}: {sorted(_ALL_SUPPORTED_LIBNAMES)}" + f"Supported names on {_PLATFORM_NAME}: {sorted(ALL_AVAILABLE_LIBNAMES)}" ) return _load_lib_no_cache(libname) diff --git a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_platform.py b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_platform.py index 37fd6eb1700..2d6a5f016a7 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_pe_matches_arch, windows_python_arch def _no_such_file_in_sub_dirs( @@ -41,7 +43,7 @@ def _find_so_in_rel_dirs( sub_dirs_searched: list[tuple[str, ...]] = [] file_wild = so_basename + "*" for rel_dir in rel_dirs: - sub_dir = tuple(rel_dir.split(os.path.sep)) + sub_dir = PurePath(rel_dir).parts for abs_dir in find_sub_dirs_all_sitepackages(sub_dir): # Exact unversioned match first; fall back to versioned names because some # distros only ship lib.so. (e.g. conda libcupti). Only one match @@ -61,12 +63,15 @@ def _find_so_in_rel_dirs( return None -def _find_dll_under_dir(dirpath: str, file_wild: str) -> str | None: +def _find_dll_under_dir(dirpath: str, file_wild: str, target_arch: str | None = None) -> str | None: for path in sorted(glob.glob(os.path.join(dirpath, file_wild))): if not os.path.isfile(path): continue - if not is_suppressed_dll_file(os.path.basename(path)): - return path + if is_suppressed_dll_file(os.path.basename(path)): + continue + if target_arch is not None and not windows_pe_matches_arch(path, target_arch): + continue + return path return None @@ -78,7 +83,7 @@ def _find_dll_in_rel_dirs( ) -> str | None: sub_dirs_searched: list[tuple[str, ...]] = [] for rel_dir in rel_dirs: - sub_dir = tuple(rel_dir.split(os.path.sep)) + sub_dir = PurePath(rel_dir).parts for abs_dir in find_sub_dirs_all_sitepackages(sub_dir): dll_name = _find_dll_under_dir(abs_dir, lib_searched_for) if dll_name is not None: @@ -109,7 +114,7 @@ def find_in_site_packages( def find_in_lib_dir( self, lib_dir: str, - libname: str, + desc: LibDescriptor, lib_searched_for: str, error_messages: list[str], attachments: list[str], @@ -142,7 +147,7 @@ def find_in_site_packages( def find_in_lib_dir( self, lib_dir: str, - _libname: str, + _desc: LibDescriptor, lib_searched_for: str, error_messages: list[str], attachments: list[str], @@ -173,17 +178,19 @@ def find_in_lib_dir( @dataclass(frozen=True, slots=True) class WindowsSearchPlatform: + target_arch: str + def lib_searched_for(self, libname: str) -> str: return f"{libname}*.dll" def site_packages_rel_dirs(self, desc: LibDescriptor) -> tuple[str, ...]: - return cast(tuple[str, ...], desc.site_packages_windows) + return cast(tuple[str, ...], desc.site_packages_windows.for_arch(self.target_arch)) def conda_anchor_point(self, conda_prefix: str) -> str: return os.path.join(conda_prefix, "Library") def anchor_rel_dirs(self, desc: LibDescriptor) -> tuple[str, ...]: - return cast(tuple[str, ...], desc.anchor_rel_dirs_windows) + return cast(tuple[str, ...], desc.anchor_rel_dirs_windows.for_arch(self.target_arch)) def find_in_site_packages( self, @@ -197,16 +204,20 @@ def find_in_site_packages( def find_in_lib_dir( self, lib_dir: str, - libname: str, + desc: LibDescriptor, _lib_searched_for: str, error_messages: list[str], attachments: list[str], ) -> str | None: - file_wild = libname + "*.dll" - dll_name = _find_dll_under_dir(lib_dir, file_wild) + file_wild = desc.name + "*.dll" + target_arch = self.target_arch if desc.requires_windows_binary_arch_check else None + dll_name = _find_dll_under_dir(lib_dir, file_wild, target_arch) if dll_name is not None: return dll_name - error_messages.append(f"No such file: {file_wild}") + if target_arch is None: + error_messages.append(f"No such file: {file_wild}") + else: + error_messages.append(f"No {target_arch}-compatible PE file: {file_wild}") attachments.append(f' listdir("{lib_dir}"):') if not os.path.isdir(lib_dir): attachments.append(" DIRECTORY DOES NOT EXIST") @@ -216,4 +227,10 @@ def find_in_lib_dir( return None -PLATFORM: SearchPlatform = WindowsSearchPlatform() if IS_WINDOWS else LinuxSearchPlatform() +def _platform_for_current_system() -> SearchPlatform: + if IS_WINDOWS: + return WindowsSearchPlatform(target_arch=windows_python_arch()) + return LinuxSearchPlatform() + + +PLATFORM = _platform_for_current_system() diff --git a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_steps.py b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_steps.py index 55d8a8aa674..5901094fcaa 100644 --- a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_steps.py +++ b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_steps.py @@ -88,7 +88,7 @@ def _find_using_lib_dir(ctx: SearchContext, lib_dir: str | None) -> str | None: str | None, ctx.platform.find_in_lib_dir( lib_dir, - ctx.libname, + ctx.desc, ctx.lib_searched_for, ctx.error_messages, ctx.attachments, @@ -121,13 +121,14 @@ def _derive_ctk_root_windows(resolved_lib_path: str) -> str | None: Supports: - ``$CTK_ROOT/bin/x64/foo.dll`` (CTK 13 style) + - ``$CTK_ROOT/bin/arm64/foo.dll`` (Windows on Arm CTK 13 style) - ``$CTK_ROOT/bin/foo.dll`` (CTK 12 style) """ import ntpath lib_dir = ntpath.dirname(resolved_lib_path) basename = ntpath.basename(lib_dir).lower() - if basename == "x64": + if basename in ("x64", "arm64"): parent = ntpath.dirname(lib_dir) if ntpath.basename(parent).lower() == "bin": return ntpath.dirname(parent) 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..daf696b638e 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") @@ -26,10 +32,30 @@ desc.name for desc in _CTK_DESCRIPTORS if desc.windows_dlls and not desc.linux_sonames ) +if not IS_WINDOWS: + ALL_AVAILABLE_LIBNAMES = frozenset(desc.name for desc in DESCRIPTOR_CATALOG if desc.linux_sonames) +else: + assert IS_WINDOWS_X64 != IS_WINDOWS_ARM64 + _current_windows_arch = "x64" if IS_WINDOWS_X64 else "arm64" + ALL_AVAILABLE_LIBNAMES = frozenset( + desc.name for desc in DESCRIPTOR_CATALOG if _current_windows_arch in desc.supported_windows_arch + ) + SUPPORTED_LIBNAMES_LINUX = SUPPORTED_LIBNAMES_COMMON + SUPPORTED_LIBNAMES_LINUX_ONLY -SUPPORTED_LIBNAMES_WINDOWS = SUPPORTED_LIBNAMES_COMMON + SUPPORTED_LIBNAMES_WINDOWS_ONLY +SUPPORTED_LIBNAMES_WINDOWS_X64 = tuple(desc.name for desc in _CTK_DESCRIPTORS if "x64" in desc.supported_windows_arch) +SUPPORTED_LIBNAMES_WINDOWS_ARM64 = tuple( + desc.name for desc in _CTK_DESCRIPTORS if "arm64" in desc.supported_windows_arch +) +# Backward-compatible alias preserves the historical x64 meaning. +SUPPORTED_LIBNAMES_WINDOWS = SUPPORTED_LIBNAMES_WINDOWS_X64 SUPPORTED_LIBNAMES_ALL = SUPPORTED_LIBNAMES_COMMON + SUPPORTED_LIBNAMES_LINUX_ONLY + SUPPORTED_LIBNAMES_WINDOWS_ONLY -SUPPORTED_LIBNAMES = SUPPORTED_LIBNAMES_WINDOWS if IS_WINDOWS else SUPPORTED_LIBNAMES_LINUX +if not IS_WINDOWS: + SUPPORTED_LIBNAMES = SUPPORTED_LIBNAMES_LINUX +elif IS_WINDOWS_X64: + SUPPORTED_LIBNAMES = SUPPORTED_LIBNAMES_WINDOWS_X64 +else: + assert IS_WINDOWS_ARM64 + SUPPORTED_LIBNAMES = SUPPORTED_LIBNAMES_WINDOWS_ARM64 DIRECT_DEPENDENCIES_CTK = {desc.name: desc.dependencies for desc in _CTK_DESCRIPTORS if desc.dependencies} DIRECT_DEPENDENCIES = {desc.name: desc.dependencies for desc in DESCRIPTOR_CATALOG if desc.dependencies} @@ -51,7 +77,6 @@ desc.name for desc in DESCRIPTOR_CATALOG if desc.requires_rtld_deepbind and desc.linux_sonames ) -# Based on output of toolshed/make_site_packages_libdirs_linux.py SITE_PACKAGES_LIBDIRS_LINUX_CTK = { desc.name: desc.site_packages_linux for desc in _CTK_DESCRIPTORS if desc.site_packages_linux } @@ -60,13 +85,29 @@ } SITE_PACKAGES_LIBDIRS_LINUX = SITE_PACKAGES_LIBDIRS_LINUX_CTK | SITE_PACKAGES_LIBDIRS_LINUX_OTHER -SITE_PACKAGES_LIBDIRS_WINDOWS_CTK = { - desc.name: desc.site_packages_windows for desc in _CTK_DESCRIPTORS if desc.site_packages_windows +# Architecture-specific Windows projections. Keep these separate: combining +# them would make the table unsafe to consume for either process ABI. +SITE_PACKAGES_LIBDIRS_WINDOWS_CTK_X64 = { + desc.name: desc.site_packages_windows.x64 for desc in _CTK_DESCRIPTORS if desc.site_packages_windows.x64 +} +SITE_PACKAGES_LIBDIRS_WINDOWS_CTK_ARM64 = { + desc.name: desc.site_packages_windows.arm64 for desc in _CTK_DESCRIPTORS if desc.site_packages_windows.arm64 } -SITE_PACKAGES_LIBDIRS_WINDOWS_OTHER = { - desc.name: desc.site_packages_windows for desc in _NON_CTK_DESCRIPTORS if desc.site_packages_windows +SITE_PACKAGES_LIBDIRS_WINDOWS_OTHER_X64 = { + desc.name: desc.site_packages_windows.x64 for desc in _NON_CTK_DESCRIPTORS if desc.site_packages_windows.x64 } -SITE_PACKAGES_LIBDIRS_WINDOWS = SITE_PACKAGES_LIBDIRS_WINDOWS_CTK | SITE_PACKAGES_LIBDIRS_WINDOWS_OTHER +SITE_PACKAGES_LIBDIRS_WINDOWS_OTHER_ARM64 = { + desc.name: desc.site_packages_windows.arm64 for desc in _NON_CTK_DESCRIPTORS if desc.site_packages_windows.arm64 +} +SITE_PACKAGES_LIBDIRS_WINDOWS_X64 = SITE_PACKAGES_LIBDIRS_WINDOWS_CTK_X64 | SITE_PACKAGES_LIBDIRS_WINDOWS_OTHER_X64 +SITE_PACKAGES_LIBDIRS_WINDOWS_ARM64 = ( + SITE_PACKAGES_LIBDIRS_WINDOWS_CTK_ARM64 | SITE_PACKAGES_LIBDIRS_WINDOWS_OTHER_ARM64 +) + +# Backward-compatible aliases preserve the historical x64 meaning. +SITE_PACKAGES_LIBDIRS_WINDOWS_CTK = SITE_PACKAGES_LIBDIRS_WINDOWS_CTK_X64 +SITE_PACKAGES_LIBDIRS_WINDOWS_OTHER = SITE_PACKAGES_LIBDIRS_WINDOWS_OTHER_X64 +SITE_PACKAGES_LIBDIRS_WINDOWS = SITE_PACKAGES_LIBDIRS_WINDOWS_X64 def is_suppressed_dll_file(path_basename: str) -> bool: 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..9313f3a9f17 --- /dev/null +++ b/cuda_pathfinder/cuda/pathfinder/_utils/windows_arch.py @@ -0,0 +1,65 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import sysconfig + +WINDOWS_PE_MACHINE_BY_ARCH = { + "x64": 0x8664, + "arm64": 0xAA64, +} + + +class UnsupportedArchError(RuntimeError): + """Raised when Python reports an unsupported Windows architecture.""" + + def __init__(self, platform_tag: str) -> None: + self.platform_tag = platform_tag + super().__init__( + f"Unsupported Windows Python platform tag: {platform_tag!r}; expected 'win-amd64' or 'win-arm64'" + ) + + +def windows_python_arch() -> str: + """Return the current Windows Python interpreter architecture.""" + raw_platform_tag = sysconfig.get_platform() + platform_tag = raw_platform_tag.lower().replace("_", "-") + + if platform_tag == "win-arm64": + return "arm64" + + if platform_tag == "win-amd64": + return "x64" + + raise UnsupportedArchError(raw_platform_tag) + + +def windows_pe_matches_arch(path: str, target_arch: str) -> bool: + """Return whether a Windows Portable Executable (PE) targets the requested architecture. + + PE is the file format used for Windows executables and DLLs. This reads the + PE/COFF header's machine field to distinguish x64 images from Arm64 images. + """ + expected_machine = WINDOWS_PE_MACHINE_BY_ARCH.get(target_arch) + if expected_machine is None: + raise ValueError(f"Unsupported Windows target architecture: {target_arch!r}") + + try: + with open(path, "rb") as stream: + if stream.read(2) != b"MZ": + return False + stream.seek(0x3C) + pe_offset_bytes = stream.read(4) + if len(pe_offset_bytes) != 4: + return False + stream.seek(int.from_bytes(pe_offset_bytes, "little")) + if stream.read(4) != b"PE\0\0": + return False + machine_bytes = stream.read(2) + if len(machine_bytes) != 2: + return False + except OSError: + return False + + return int.from_bytes(machine_bytes, "little") == expected_machine 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.1-notes.rst b/cuda_pathfinder/docs/source/release/1.6.1-notes.rst new file mode 100644 index 00000000000..29a3106e0f1 --- /dev/null +++ b/cuda_pathfinder/docs/source/release/1.6.1-notes.rst @@ -0,0 +1,38 @@ +.. SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +.. SPDX-License-Identifier: Apache-2.0 + +.. py:currentmodule:: cuda.pathfinder + +``cuda-pathfinder`` 1.6.1 Release notes +======================================= + +Highlights +---------- + +* Make Windows dynamic-library discovery architecture-aware. Pathfinder now + detects whether the current Python interpreter is x64 or Arm64, searches + only the matching CUDA Toolkit and wheel directories, and reports only the + CTK libraries available for that architecture through + ``SUPPORTED_NVIDIA_LIBNAMES``. A known library unavailable for the current + architecture raises ``DynamicLibNotAvailableError``. + +* Add Windows Arm64 discovery for CUDA 13.4 layouts while retaining legacy + CUDA 12 wheel directories as x64-only fallbacks. This includes corrected + architecture-specific locations for cuDLA, NVVM, CUPTI, and cuSPARSELt. + NVVM binaries found in an unqualified legacy directory are checked for a + matching PE machine architecture before loading. + +* Add ``UnsupportedArchError`` for unsupported Windows Python platform tags. + +Internal maintenance +-------------------- + +* Group Windows search locations by architecture in the dynamic-library + descriptor catalog. Add explicit ``_X64`` and ``_ARM64`` variants of the + internal ``SUPPORTED_LIBNAMES_WINDOWS*`` and + ``SITE_PACKAGES_LIBDIRS_WINDOWS*`` tables. Unsuffixed names remain x64 + aliases for backward compatibility. + +* Remove the obsolete descriptor-catalog writer and its catalog-update tools. + The site-packages collection scripts remain available for gathering library + paths. diff --git a/cuda_pathfinder/tests/test_ctk_root_discovery.py b/cuda_pathfinder/tests/test_ctk_root_discovery.py index 9ad148dccd2..b4aa33d6a74 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, @@ -32,6 +33,7 @@ MODE_CANARY, ) from cuda.pathfinder._utils.platform_aware import IS_WINDOWS +from cuda.pathfinder._utils.windows_arch import windows_python_arch _MODULE = "cuda.pathfinder._dynamic_libs.load_nvidia_dynamic_lib" _STEPS_MODULE = "cuda.pathfinder._dynamic_libs.search_steps" @@ -60,11 +62,18 @@ def _create_nvvm_in_ctk(ctk_root): nvvm_dir = ctk_root / "nvvm" / "bin" nvvm_dir.mkdir(parents=True) nvvm_lib = nvvm_dir / "nvvm64.dll" + machine = {"x64": 0x8664, "arm64": 0xAA64}[windows_python_arch()] + image = bytearray(0x86) + image[:2] = b"MZ" + image[0x3C:0x40] = (0x80).to_bytes(4, "little") + image[0x80:0x84] = b"PE\0\0" + image[0x84:0x86] = machine.to_bytes(2, "little") + nvvm_lib.write_bytes(image) else: nvvm_dir = ctk_root / "nvvm" / "lib64" nvvm_dir.mkdir(parents=True) nvvm_lib = nvvm_dir / "libnvvm.so" - nvvm_lib.write_bytes(b"fake") + nvvm_lib.write_bytes(b"fake") return nvvm_lib @@ -126,6 +135,12 @@ def test_derive_ctk_root_windows_ctk13(): assert _derive_ctk_root_windows(path) == r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.0" +@pytest.mark.agent_authored(model="gpt-5") +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 +205,24 @@ def test_try_via_ctk_root_regular_lib(tmp_path): assert result.found_via == "system-ctk-root" +@pytest.mark.agent_authored(model="gpt-5") +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..3b643aa2e77 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,35 @@ 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.agent_authored(model="gpt-5") +def test_cusparselt_windows_metadata_matches_wheel_layouts(): + spec = _CATALOG_BY_NAME["cusparseLt"] + assert spec.supported_windows_arch == ("x64", "arm64") + assert spec.site_packages_windows == WindowsSearchDirs( + x64=("nvidia/cu13/bin/x64", "nvidia/cusparselt/bin"), + arm64=("nvidia/cu13/bin/arm64",), + ) + + @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: @@ -96,3 +126,10 @@ def test_ctk_root_canary_anchors_reference_known_ctk_libs(spec: DescriptorSpec): def test_only_ctk_libs_define_ctk_root_canary_anchors(spec: DescriptorSpec): if spec.ctk_root_canary_anchor_libnames: assert spec.packaged_with == "ctk", f"{spec.name} defines canary anchors but is not a CTK lib" + + +@pytest.mark.agent_authored(model="gpt-5") +def test_only_nvvm_requires_windows_binary_arch_check(): + checked_libs = {spec.name for spec in DESCRIPTOR_CATALOG if spec.requires_windows_binary_arch_check} + + assert checked_libs == {"nvvm"} 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..ad42e24e63e 100644 --- a/cuda_pathfinder/tests/test_load_nvidia_dynamic_lib.py +++ b/cuda_pathfinder/tests/test_load_nvidia_dynamic_lib.py @@ -25,32 +25,49 @@ assert STRICTNESS in ("see_what_works", "all_must_work") +@pytest.mark.agent_authored(model="gpt-5") +def test_loader_uses_all_available_libnames(): + assert supported_nvidia_libs.ALL_AVAILABLE_LIBNAMES == load_nvidia_dynamic_lib_module.ALL_AVAILABLE_LIBNAMES + + def test_supported_libnames_linux_sonames_consistency(): assert tuple(sorted(supported_nvidia_libs.SUPPORTED_LIBNAMES_LINUX)) == tuple( sorted(supported_nvidia_libs.SUPPORTED_LINUX_SONAMES_CTK.keys()) ) -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()) - ) - - def test_supported_libnames_linux_site_packages_libdirs_ctk_consistency(): assert tuple(sorted(supported_nvidia_libs.SUPPORTED_LIBNAMES_LINUX)) == tuple( sorted(supported_nvidia_libs.SITE_PACKAGES_LIBDIRS_LINUX_CTK.keys()) ) -def test_supported_libnames_windows_site_packages_libdirs_ctk_consistency(): +@pytest.mark.parametrize( + ("site_packages_libdirs", "supported_libnames"), + [ + pytest.param( + supported_nvidia_libs.SITE_PACKAGES_LIBDIRS_WINDOWS_CTK_X64, + supported_nvidia_libs.SUPPORTED_LIBNAMES_WINDOWS_X64, + id="x64", + ), + pytest.param( + supported_nvidia_libs.SITE_PACKAGES_LIBDIRS_WINDOWS_CTK_ARM64, + supported_nvidia_libs.SUPPORTED_LIBNAMES_WINDOWS_ARM64, + id="arm64", + ), + ], +) +@pytest.mark.human_reviewed +def test_supported_libnames_windows_site_packages_libdirs_ctk_consistency( + site_packages_libdirs, + supported_libnames, +): # 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()) - supported_libnames = set(supported_nvidia_libs.SUPPORTED_LIBNAMES_WINDOWS) - assert site_packages_libnames <= supported_libnames + site_packages_libnames = set(site_packages_libdirs) + assert site_packages_libnames <= set(supported_libnames) @pytest.mark.parametrize("dict_name", ["SUPPORTED_LINUX_SONAMES", "SUPPORTED_WINDOWS_DLLS"]) @@ -88,7 +105,7 @@ def test_unknown_libname_raises_dynamic_lib_unknown_error(): def test_known_but_platform_unavailable_libname_raises_dynamic_lib_not_available_error(monkeypatch): load_nvidia_dynamic_lib.cache_clear() monkeypatch.setattr(load_nvidia_dynamic_lib_module, "_ALL_KNOWN_LIBNAMES", frozenset(("known_but_unavailable",))) - monkeypatch.setattr(load_nvidia_dynamic_lib_module, "_ALL_SUPPORTED_LIBNAMES", frozenset()) + monkeypatch.setattr(load_nvidia_dynamic_lib_module, "ALL_AVAILABLE_LIBNAMES", frozenset()) monkeypatch.setattr(load_nvidia_dynamic_lib_module, "_PLATFORM_NAME", "TestOS") with pytest.raises( DynamicLibNotAvailableError, diff --git a/cuda_pathfinder/tests/test_search_steps.py b/cuda_pathfinder/tests/test_search_steps.py index 1b881707dfb..54136dc34e1 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,23 @@ 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) + + +def _write_pe(path, machine): + image = bytearray(0x86) + image[:2] = b"MZ" + image[0x3C:0x40] = (0x80).to_bytes(4, "little") + image[0x80:0x84] = b"PE\0\0" + image[0x84:0x86] = machine.to_bytes(2, "little") + path.write_bytes(image) + + # --------------------------------------------------------------------------- # SearchContext # --------------------------------------------------------------------------- @@ -67,7 +91,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 +107,71 @@ 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() + + @pytest.mark.agent_authored(model="gpt-5") + 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" + + @pytest.mark.agent_authored(model="gpt-5") + 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" + + @pytest.mark.agent_authored(model="gpt-5") + 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" + + +@pytest.mark.parametrize( + ("machine", "target_arch", "expected"), + ( + (0x8664, "x64", True), + (0x8664, "arm64", False), + (0xAA64, "x64", False), + (0xAA64, "arm64", True), + ), +) +@pytest.mark.agent_authored(model="gpt-5") +def test_windows_pe_matches_arch(tmp_path, machine, target_arch, expected): + dll = tmp_path / "test.dll" + _write_pe(dll, machine) + + assert windows_arch_mod.windows_pe_matches_arch(str(dll), target_arch) is expected + + +@pytest.mark.agent_authored(model="gpt-5") +def test_windows_pe_matches_arch_rejects_malformed_file(tmp_path): + dll = tmp_path / "test.dll" + dll.write_bytes(b"not a PE file") + + assert windows_arch_mod.windows_pe_matches_arch(str(dll), "x64") is False + + # --------------------------------------------------------------------------- # find_in_site_packages # --------------------------------------------------------------------------- @@ -90,7 +179,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 +216,93 @@ 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" + @pytest.mark.agent_authored(model="gpt-5") + 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" + + @pytest.mark.agent_authored(model="gpt-5") + 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 +410,28 @@ 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" + @pytest.mark.agent_authored(model="gpt-5") + 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 +512,57 @@ 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" + @pytest.mark.agent_authored(model="gpt-5") + 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" + + @pytest.mark.parametrize( + ("target_arch", "machine", "expected_found"), + ( + ("x64", 0x8664, True), + ("x64", 0xAA64, False), + ("arm64", 0x8664, False), + ("arm64", 0xAA64, True), + ), + ) + @pytest.mark.agent_authored(model="gpt-5") + def test_nvvm_windows_checks_binary_arch(self, mocker, tmp_path, target_arch, machine, expected_found): + nvvm_dir = tmp_path / "nvvm" / "bin" + nvvm_dir.mkdir(parents=True) + dll = nvvm_dir / "nvvm64_40_0.dll" + _write_pe(dll, machine) + + mocker.patch(f"{_STEPS_MOD}.get_cuda_path_or_home", return_value=str(tmp_path)) + + ctx = _ctx(LIB_DESCRIPTORS["nvvm"], platform=WindowsSearchPlatform(target_arch=target_arch)) + result = find_in_cuda_path(ctx) + + assert (result is not None) is expected_found + if expected_found: + assert result is not None + assert result.abs_path == str(dll) + assert result.found_via == "CUDA_PATH" + else: + assert any(f"No {target_arch}-compatible PE file" in message for message in ctx.error_messages) + # --------------------------------------------------------------------------- # run_find_steps @@ -388,19 +620,63 @@ 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",) + + @pytest.mark.agent_authored(model="gpt-5") + def test_cupti_has_custom_windows_paths(self): + desc = LIB_DESCRIPTORS["cupti"] + assert desc.anchor_rel_dirs_windows.for_arch("x64") == ( + "extras/CUPTI/lib/x64", + "extras/CUPTI/lib64", + "bin", + ) + assert desc.anchor_rel_dirs_windows.for_arch("arm64") == ("extras/CUPTI/lib/arm64",) @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",) + + @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",) + + @pytest.mark.agent_authored(model="gpt-5") + 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") + + @pytest.mark.agent_authored(model="gpt-5") + 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 +689,33 @@ 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")) + @pytest.mark.agent_authored(model="gpt-5") + def test_find_lib_dir_windows_arm64_uses_arm64_anchor(self, tmp_path): + (tmp_path / "bin" / "x64").mkdir(parents=True) + (tmp_path / "bin" / "arm64").mkdir(parents=True) + + desc = _make_desc( + name="cudart", + anchor_rel_dirs_windows=WindowsSearchDirs( + x64=("bin/x64", "bin"), + arm64=("bin/arm64",), + ), + ) + 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("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 deleted file mode 100644 index b41fb5838dd..00000000000 --- a/toolshed/_catalog_writer.py +++ /dev/null @@ -1,182 +0,0 @@ -#!/usr/bin/env python3 - -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Shared helper for reading, updating, and rewriting descriptor_catalog.py. - -Each toolshed script that extracts data from CTK distributions or wheel -layouts uses this module to merge its findings into the authored catalog -without touching fields it doesn't own. -""" - -from __future__ import annotations - -import dataclasses -import json -import sys -from pathlib import Path - -# Ensure the cuda_pathfinder package is importable. -_REPO_ROOT = Path(__file__).resolve().parents[1] -_PATHFINDER_ROOT = _REPO_ROOT / "cuda_pathfinder" -if str(_PATHFINDER_ROOT) not in sys.path: - sys.path.insert(0, str(_PATHFINDER_ROOT)) - -from cuda.pathfinder._dynamic_libs.descriptor_catalog import ( # noqa: E402 - DESCRIPTOR_CATALOG, - DescriptorSpec, -) - -CATALOG_PATH = _PATHFINDER_ROOT / "cuda" / "pathfinder" / "_dynamic_libs" / "descriptor_catalog.py" - -_DEFAULTS = DescriptorSpec(name="", packaged_with="ctk") - -_SECTION_COMMENTS = { - "ctk": ( - " # -----------------------------------------------------------------------\n" - " # CTK (CUDA Toolkit) libraries\n" - " # -----------------------------------------------------------------------" - ), - "other": ( - " # -----------------------------------------------------------------------\n" - " # Third-party / separately packaged libraries\n" - " # -----------------------------------------------------------------------" - ), - "driver": ( - " # -----------------------------------------------------------------------\n" - " # Driver libraries (system-search only, no CTK cascade)\n" - " # -----------------------------------------------------------------------" - ), -} - - -def _quote(s: str) -> str: - return json.dumps(s) - - -def _render_tuple(values: tuple[str, ...]) -> str: - if not values: - return "()" - if len(values) == 1: - return f"({_quote(values[0])},)" - return "(" + ", ".join(_quote(v) for v in values) + ")" - - -def _render_spec(spec: DescriptorSpec) -> str: - """Render a single DescriptorSpec constructor call, omitting default-valued fields.""" - lines = [ - " DescriptorSpec(", - f" name={_quote(spec.name)},", - f' packaged_with="{spec.packaged_with}",', - ] - - tuple_fields = [ - "linux_sonames", - "windows_dlls", - "site_packages_linux", - "site_packages_windows", - "dependencies", - "anchor_rel_dirs_linux", - "anchor_rel_dirs_windows", - "ctk_root_canary_anchor_libnames", - ] - bool_fields = [ - "requires_add_dll_directory", - "requires_rtld_deepbind", - ] - - for field in tuple_fields: - value = getattr(spec, field) - default = getattr(_DEFAULTS, field) - if value != default: - lines.append(f" {field}={_render_tuple(value)},") - - for field in bool_fields: - value = getattr(spec, field) - default = getattr(_DEFAULTS, field) - if value != default: - lines.append(f" {field}={value},") - - lines.append(" ),") - return "\n".join(lines) - - -def render_catalog(specs: tuple[DescriptorSpec, ...]) -> str: - """Render the full descriptor_catalog.py file content.""" - header = '''\ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Canonical authored descriptor catalog for dynamic libraries.""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import Literal - -PackagedWith = Literal["ctk", "other", "driver"] - - -@dataclass(frozen=True, slots=True) -class DescriptorSpec: - name: str - packaged_with: PackagedWith - linux_sonames: tuple[str, ...] = () - windows_dlls: tuple[str, ...] = () - site_packages_linux: tuple[str, ...] = () - site_packages_windows: tuple[str, ...] = () - dependencies: tuple[str, ...] = () - anchor_rel_dirs_linux: tuple[str, ...] = ("lib64", "lib") - anchor_rel_dirs_windows: tuple[str, ...] = ("bin/x64", "bin") - ctk_root_canary_anchor_libnames: tuple[str, ...] = () - requires_add_dll_directory: bool = False - requires_rtld_deepbind: bool = False - - -DESCRIPTOR_CATALOG: tuple[DescriptorSpec, ...] = ( -''' - - body_parts: list[str] = [] - prev_packaged_with = None - for spec in specs: - if spec.packaged_with != prev_packaged_with: - comment = _SECTION_COMMENTS.get(spec.packaged_with) - if comment is not None: - body_parts.append(comment) - prev_packaged_with = spec.packaged_with - body_parts.append(_render_spec(spec)) - - footer = ")\n" - return header + "\n".join(body_parts) + "\n" + footer - - -def load_catalog() -> tuple[DescriptorSpec, ...]: - """Return the current DESCRIPTOR_CATALOG from disk.""" - return DESCRIPTOR_CATALOG - - -def load_catalog_as_dict() -> dict[str, DescriptorSpec]: - """Return the current catalog keyed by name.""" - return {spec.name: spec for spec in DESCRIPTOR_CATALOG} - - -def update_specs( - catalog: tuple[DescriptorSpec, ...], - updates: dict[str, dict[str, object]], -) -> tuple[DescriptorSpec, ...]: - """Apply field updates to matching specs by name, preserving order.""" - result = [] - for spec in catalog: - if spec.name in updates: - result.append(dataclasses.replace(spec, **updates[spec.name])) - else: - result.append(spec) - return tuple(result) - - -def write_catalog(specs: tuple[DescriptorSpec, ...], path: Path | None = None) -> None: - """Render and write the catalog to disk.""" - if path is None: - path = CATALOG_PATH - path.write_text(render_catalog(specs), encoding="utf-8") diff --git a/toolshed/build_pathfinder_dlls.py b/toolshed/build_pathfinder_dlls.py deleted file mode 100755 index 63abba52386..00000000000 --- a/toolshed/build_pathfinder_dlls.py +++ /dev/null @@ -1,118 +0,0 @@ -#!/usr/bin/env python3 - -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Scan 7z listing files for .dll names, update descriptor_catalog.py. - -Usage: - # First generate listings from CTK .exe installers: - # for exe in *.exe; do 7z l "$exe" > "${exe%.exe}.txt"; done - python toolshed/build_pathfinder_dlls.py listing1.txt [listing2.txt ...] -""" - -from __future__ import annotations - -import collections -import sys -from pathlib import Path - -sys.path.insert(0, str(Path(__file__).resolve().parent)) -from _catalog_writer import load_catalog, update_specs, write_catalog - - -def _is_suppressed_dll(libname: str, dll: str) -> bool: - if libname == "cudart": - if dll.startswith("cudart32_"): - return True - if dll == "cudart64_65.dll": - # PhysX/files/Common/cudart64_65.dll from CTK 6.5, but shipped with CTK 12.0-12.9 - return True - if dll == "cudart64_101.dll": - # GFExperience.NvStreamSrv/amd64/server/cudart64_101.dll from CTK 10.1, but shipped with CTK 12.0-12.6 - return True - elif libname == "nvrtc": - if dll.endswith(".alt.dll"): - return True - if dll.startswith("nvrtc-builtins"): - return True - elif libname == "nvvm" and dll == "nvvm32.dll": - return True - return False - - -def _parse_listings(paths: list[str]) -> set[str]: - dlls: set[str] = set() - for filename in paths: - lines_iter = iter(Path(filename).read_text().splitlines()) - for line in lines_iter: - if line.startswith("-------------------"): - break - else: - raise RuntimeError(f"------------------- NOT FOUND in {filename}") - for line in lines_iter: - if line.startswith("-------------------"): - break - assert line[52] == " ", line - assert line[53] != " ", line - path = line[53:] - if path.endswith(".dll"): - dll = path.rsplit("/", 1)[1] - dlls.add(dll) - else: - raise RuntimeError(f"------------------- NOT FOUND in {filename}") - return dlls - - -def run(listing_files: list[str]) -> None: - dlls_from_files = _parse_listings(listing_files) - catalog = load_catalog() - - # Longest-prefix-first to avoid ambiguous matches (e.g. "cufftw" before "cufft"). - ctk_names = sorted( - (spec.name for spec in catalog if spec.packaged_with == "ctk"), - key=lambda n: (-len(n), n), - ) - - dlls_in_scope: set[str] = set() - dlls_by_name: dict[str, list[str]] = collections.defaultdict(list) - suppressed: set[str] = set() - - for name in ctk_names: - for dll in sorted(dlls_from_files): - if dll not in dlls_in_scope and dll.startswith(name): - if _is_suppressed_dll(name, dll): - suppressed.add(dll) - else: - dlls_by_name[name].append(dll) - dlls_in_scope.add(dll) - - updates: dict[str, dict[str, object]] = {} - for name, dlls in dlls_by_name.items(): - updates[name] = {"windows_dlls": tuple(dlls)} - - if updates: - write_catalog(update_specs(catalog, updates)) - for name in sorted(updates): - print(f" updated {name}: windows_dlls={updates[name]['windows_dlls']}") - else: - print("No matching DLLs found.") - - if suppressed: - print(f"\nSuppressed DLLs ({len(suppressed)}):") - for dll in sorted(suppressed): - print(f" {dll}") - - out_of_scope = dlls_from_files - dlls_in_scope - if out_of_scope: - print(f"\nDLLs out of scope ({len(out_of_scope)}):") - for dll in sorted(out_of_scope): - print(f" {dll}") - - -if __name__ == "__main__": - if len(sys.argv) < 2: - print("Usage: build_pathfinder_dlls.py <7z-listing.txt> ...", file=sys.stderr) - sys.exit(1) - run(listing_files=sys.argv[1:]) diff --git a/toolshed/build_pathfinder_sonames.py b/toolshed/build_pathfinder_sonames.py deleted file mode 100755 index b3fa6c2efc9..00000000000 --- a/toolshed/build_pathfinder_sonames.py +++ /dev/null @@ -1,93 +0,0 @@ -#!/usr/bin/env python3 - -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Scan directories for .so files, extract SONAMEs, update descriptor_catalog.py. - -Usage: - python toolshed/build_pathfinder_sonames.py /path/to/cuda [/more/paths ...] -""" - -from __future__ import annotations - -import os -import subprocess -import sys -from pathlib import Path - -sys.path.insert(0, str(Path(__file__).resolve().parent)) -from _catalog_writer import load_catalog, update_specs, write_catalog - - -def _extract_soname(path: str) -> str | None: - try: - out = subprocess.run( # noqa: S603 - ["readelf", "-d", path], # noqa: S607 - capture_output=True, - text=True, - timeout=10, - ) - except (FileNotFoundError, subprocess.TimeoutExpired): - return None - for line in out.stdout.splitlines(): - if "SONAME" in line: - # Format: 0x000000000000000e (SONAME) Library soname: [libfoo.so.1] - start = line.find("[") - end = line.find("]") - if start != -1 and end != -1: - return line[start + 1 : end] - return None - - -def _find_sonames(roots: list[str]) -> set[str]: - sonames: set[str] = set() - for root in roots: - for dirpath, _dirnames, filenames in os.walk(root): - for fname in filenames: - if ".so" not in fname: - continue - full = os.path.join(dirpath, fname) - if os.path.islink(full): - continue - soname = _extract_soname(full) - if soname is not None: - sonames.add(soname) - return sonames - - -def run(roots: list[str]) -> None: - sonames_found = _find_sonames(roots) - catalog = load_catalog() - - updates: dict[str, dict[str, object]] = {} - matched: set[str] = set() - for spec in catalog: - if spec.packaged_with != "ctk": - continue - prefix = "lib" + spec.name + ".so" - found = tuple(sorted(s for s in sonames_found if s.startswith(prefix))) - if found: - updates[spec.name] = {"linux_sonames": found} - matched.update(found) - - if updates: - write_catalog(update_specs(catalog, updates)) - for name, upd in sorted(updates.items()): - print(f" updated {name}: linux_sonames={upd['linux_sonames']}") - else: - print("No matching sonames found.") - - unmatched = sonames_found - matched - if unmatched: - print(f"\nSONAMEs not matched to any CTK descriptor ({len(unmatched)}):") - for s in sorted(unmatched): - print(f" {s}") - - -if __name__ == "__main__": - if len(sys.argv) < 2: - print("Usage: build_pathfinder_sonames.py [ ...]", file=sys.stderr) - sys.exit(1) - run(roots=sys.argv[1:]) diff --git a/toolshed/collect_site_packages_dll_files.ps1 b/toolshed/collect_site_packages_dll_files.ps1 index f0a6f799242..4efebbf3aab 100644 --- a/toolshed/collect_site_packages_dll_files.ps1 +++ b/toolshed/collect_site_packages_dll_files.ps1 @@ -1,12 +1,11 @@ # collect_site_packages_dll_files.ps1 -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # Usage: # cd cuda-python # powershell -File toolshed\collect_site_packages_dll_files.ps1 -# python .\toolshed\make_site_packages_libdirs.py windows site_packages_dll.txt $ErrorActionPreference = 'Stop' diff --git a/toolshed/collect_site_packages_so_files.sh b/toolshed/collect_site_packages_so_files.sh index 974f6eeae86..a88652bdfe0 100755 --- a/toolshed/collect_site_packages_so_files.sh +++ b/toolshed/collect_site_packages_so_files.sh @@ -1,12 +1,11 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # Usage: # cd cuda-python # ./toolshed/collect_site_packages_so_files.sh -# ./toolshed/make_site_packages_libdirs.py linux site_packages_so.txt set -euo pipefail fresh_venv() { diff --git a/toolshed/make_site_packages_libdirs.py b/toolshed/make_site_packages_libdirs.py deleted file mode 100755 index e1cbcb28825..00000000000 --- a/toolshed/make_site_packages_libdirs.py +++ /dev/null @@ -1,123 +0,0 @@ -#!/usr/bin/env python3 - -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Parse collected site-packages library paths, update descriptor_catalog.py. - -Usage: - python toolshed/make_site_packages_libdirs.py linux collected_linux.txt - python toolshed/make_site_packages_libdirs.py windows collected_windows.txt -""" - -from __future__ import annotations - -import argparse -import os -import re -import sys -from pathlib import Path -from typing import Dict, Set - -sys.path.insert(0, str(Path(__file__).resolve().parent)) -from _catalog_writer import load_catalog, update_specs, write_catalog - -_SITE_PACKAGES_RE = re.compile(r"(?i)^.*?/site-packages/") - - -def _strip_site_packages_prefix(p: str) -> str: - """Remove any leading '.../site-packages/' (handles '\\' or '/', case-insensitive).""" - p = p.replace("\\", "/") - return _SITE_PACKAGES_RE.sub("", p) - - -def _parse_lines_linux(lines: list[str]) -> Dict[str, Set[str]]: - d: Dict[str, Set[str]] = {} - for raw in lines: - line = raw.strip() - if not line or line.startswith("#"): - continue - line = _strip_site_packages_prefix(line) - dirpath, fname = os.path.split(line) - # Require something like libNAME.so, libNAME.so.12, libNAME.so.12.1, etc. - i = fname.find(".so") - if not fname.startswith("lib") or i == -1: - continue - name = fname[3:i] # e.g. "libnvrtc" -> "nvrtc" - d.setdefault(name, set()).add(dirpath) - return d - - -def _extract_libname_from_dll(fname: str) -> str | None: - """Return base libname per the heuristic, or None if not a .dll.""" - base = os.path.basename(fname) - if not base.lower().endswith(".dll"): - return None - stem = base[:-4] # drop ".dll" - out = [] - for ch in stem: - if ch == "_" or ch.isdigit(): - break - out.append(ch) - name = "".join(out) - return name or None - - -def _parse_lines_windows(lines: list[str]) -> Dict[str, Set[str]]: - """Collect {libname: set(dirnames)} with deduped directories.""" - m: Dict[str, Set[str]] = {} - for raw in lines: - line = raw.strip() - if not line or line.startswith("#"): - continue - line = _strip_site_packages_prefix(line) - dirpath, fname = os.path.split(line) - libname = _extract_libname_from_dll(fname) - if not libname: - continue - m.setdefault(libname, set()).add(dirpath) - return m - - -def main() -> None: - ap = argparse.ArgumentParser( - description="Update site_packages_* in descriptor_catalog.py from collected library paths" - ) - ap.add_argument("platform", choices=["linux", "windows"]) - ap.add_argument("path", help="Text file with one library path per line") - args = ap.parse_args() - - with open(args.path, encoding="utf-8") as f: - lines = f.read().splitlines() - - if args.platform == "linux": - parsed = _parse_lines_linux(lines) - field = "site_packages_linux" - else: - parsed = _parse_lines_windows(lines) - field = "site_packages_windows" - - catalog = load_catalog() - catalog_names = {spec.name 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 updates: - write_catalog(update_specs(catalog, updates)) - for name in sorted(updates): - print(f" updated {name}: {field}={updates[name][field]}") - else: - print("No matching libraries found.") - - unmatched = set(parsed.keys()) - catalog_names - if unmatched: - print(f"\nLibraries not in catalog ({len(unmatched)}):") - for name in sorted(unmatched): - print(f" {name}") - - -if __name__ == "__main__": - main() diff --git a/toolshed/update_catalog.py b/toolshed/update_catalog.py deleted file mode 100644 index 800451ca45d..00000000000 --- a/toolshed/update_catalog.py +++ /dev/null @@ -1,47 +0,0 @@ -#!/usr/bin/env python3 - -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Update descriptor_catalog.py from CTK installations. - -On Linux, scans directories for .so files and extracts SONAMEs via readelf. -On Windows, parses 7z listing files generated from CTK .exe installers. - -Usage: - # Linux — pass one or more CTK lib directories: - python toolshed/update_catalog.py /path/to/ctk12/lib64 /path/to/ctk13/lib64 - - # Windows — pass 7z listing .txt files: - # for exe in *.exe; do 7z l "$exe" > "${exe%.exe}.txt"; done - python toolshed/update_catalog.py listing12.txt listing13.txt -""" - -from __future__ import annotations - -import sys -from pathlib import Path - -sys.path.insert(0, str(Path(__file__).resolve().parent)) - - -def main() -> None: - if len(sys.argv) < 2: - print(__doc__, file=sys.stderr) - sys.exit(1) - - args = sys.argv[1:] - - if sys.platform == "win32": - from build_pathfinder_dlls import run as run_dlls - - run_dlls(listing_files=args) - else: - from build_pathfinder_sonames import run as run_sonames - - run_sonames(roots=args) - - -if __name__ == "__main__": - main()