diff --git a/cuda_pathfinder/cuda/pathfinder/_static_libs/find_static_lib.py b/cuda_pathfinder/cuda/pathfinder/_static_libs/find_static_lib.py index 804b1c04be7..ea5a740aec4 100644 --- a/cuda_pathfinder/cuda/pathfinder/_static_libs/find_static_lib.py +++ b/cuda_pathfinder/cuda/pathfinder/_static_libs/find_static_lib.py @@ -9,6 +9,7 @@ from cuda.pathfinder._utils.env_vars import get_cuda_path_or_home 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 class StaticLibNotFoundError(RuntimeError): @@ -32,17 +33,28 @@ class _StaticLibInfo(TypedDict): site_packages_dirs: tuple[str, ...] +def _cudadevrt_info() -> _StaticLibInfo: + if not IS_WINDOWS: + return { + "filename": "libcudadevrt.a", + "ctk_rel_paths": ("lib64", "lib"), + "conda_rel_paths": ("lib",), + "site_packages_dirs": ("nvidia/cu13/lib", "nvidia/cuda_runtime/lib"), + } + + arch_dir = windows_python_arch() + component_wheel_dirs = ("nvidia/cuda_runtime/lib/x64",) if arch_dir == "x64" else () + conda_fallback_dirs = ("lib",) if arch_dir == "x64" else () + return { + "filename": "cudadevrt.lib", + "ctk_rel_paths": (os.path.join("lib", arch_dir),), + "conda_rel_paths": (os.path.join("lib", arch_dir), *conda_fallback_dirs), + "site_packages_dirs": (f"nvidia/cu13/lib/{arch_dir}", *component_wheel_dirs), + } + + _SUPPORTED_STATIC_LIBS_INFO: dict[str, _StaticLibInfo] = { - "cudadevrt": { - "filename": "cudadevrt.lib" if IS_WINDOWS else "libcudadevrt.a", - "ctk_rel_paths": (os.path.join("lib", "x64"),) if IS_WINDOWS else ("lib64", "lib"), - "conda_rel_paths": ((os.path.join("lib", "x64"), "lib") if IS_WINDOWS else ("lib",)), - "site_packages_dirs": ( - ("nvidia/cu13/lib/x64", "nvidia/cuda_runtime/lib/x64") - if IS_WINDOWS - else ("nvidia/cu13/lib", "nvidia/cuda_runtime/lib") - ), - }, + "cudadevrt": _cudadevrt_info(), } SUPPORTED_STATIC_LIBS: tuple[str, ...] = tuple(sorted(_SUPPORTED_STATIC_LIBS_INFO.keys())) diff --git a/cuda_pathfinder/docs/source/release/1.6.1-notes.rst b/cuda_pathfinder/docs/source/release/1.6.1-notes.rst index 29a3106e0f1..919963802ff 100644 --- a/cuda_pathfinder/docs/source/release/1.6.1-notes.rst +++ b/cuda_pathfinder/docs/source/release/1.6.1-notes.rst @@ -24,6 +24,11 @@ Highlights * Add ``UnsupportedArchError`` for unsupported Windows Python platform tags. +* Make Windows static-library discovery architecture-aware. Searches now use + the current Python interpreter architecture to select the matching + ``lib/x64`` or ``lib/arm64`` CUDA Toolkit and wheel directories. CUDA 12 + component-wheel and legacy Conda fallbacks remain x64-only. + Internal maintenance -------------------- diff --git a/cuda_pathfinder/tests/test_find_static_lib.py b/cuda_pathfinder/tests/test_find_static_lib.py index e5560dcabbf..6d29a8def11 100644 --- a/cuda_pathfinder/tests/test_find_static_lib.py +++ b/cuda_pathfinder/tests/test_find_static_lib.py @@ -143,6 +143,41 @@ def test_locate_static_lib_conda_rel_path_fallback(monkeypatch, tmp_path): assert located_lib.found_via == "conda" +@pytest.mark.parametrize( + ("target_arch", "expected_ctk_dirs", "expected_conda_dirs", "expected_site_packages_dirs"), + ( + ( + "x64", + (os.path.join("lib", "x64"),), + (os.path.join("lib", "x64"), "lib"), + ("nvidia/cu13/lib/x64", "nvidia/cuda_runtime/lib/x64"), + ), + ( + "arm64", + (os.path.join("lib", "arm64"),), + (os.path.join("lib", "arm64"),), + ("nvidia/cu13/lib/arm64",), + ), + ), +) +@pytest.mark.agent_authored(model="gpt-5.6") +def test_cudadevrt_windows_paths_follow_python_arch( + monkeypatch, + target_arch, + expected_ctk_dirs, + expected_conda_dirs, + expected_site_packages_dirs, +): + monkeypatch.setattr(find_static_lib_module, "IS_WINDOWS", True) + monkeypatch.setattr(find_static_lib_module, "windows_python_arch", lambda: target_arch) + + info = find_static_lib_module._cudadevrt_info() + + assert info["ctk_rel_paths"] == expected_ctk_dirs + assert info["conda_rel_paths"] == expected_conda_dirs + assert info["site_packages_dirs"] == expected_site_packages_dirs + + @pytest.mark.usefixtures("clear_find_static_lib_cache") def test_find_static_lib_not_found_error_includes_cuda_home_directory_listing(monkeypatch, tmp_path): filename = CUDADEVRT_INFO["filename"]