Skip to content

Commit 1d0a111

Browse files
isVoidisVoid
committed
Make Windows pathfinder searches architecture-aware
1 parent 270be80 commit 1d0a111

11 files changed

Lines changed: 362 additions & 55 deletions

File tree

‎cuda_pathfinder/cuda/pathfinder/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,7 @@
6060
locate_static_lib as locate_static_lib,
6161
)
6262
from cuda.pathfinder._utils.env_vars import get_cuda_path_or_home as get_cuda_path_or_home
63+
from cuda.pathfinder._utils.windows_arch import UnsupportedWindowsArchError as UnsupportedWindowsArchError
6364

6465
from cuda.pathfinder._version import __version__ # isort: skip
6566

‎cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py‎

Lines changed: 103 additions & 39 deletions
Large diffs are not rendered by default.

‎cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_platform.py‎

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -13,13 +13,15 @@
1313
import glob
1414
import os
1515
from collections.abc import Sequence
16-
from dataclasses import dataclass
16+
from dataclasses import dataclass, field
17+
from pathlib import PurePath
1718
from typing import Protocol, cast
1819

1920
from cuda.pathfinder._dynamic_libs.lib_descriptor import LibDescriptor
2021
from cuda.pathfinder._dynamic_libs.supported_nvidia_libs import is_suppressed_dll_file
2122
from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages
2223
from cuda.pathfinder._utils.platform_aware import IS_WINDOWS
24+
from cuda.pathfinder._utils.windows_arch import windows_python_arch
2325

2426

2527
def _no_such_file_in_sub_dirs(
@@ -41,7 +43,7 @@ def _find_so_in_rel_dirs(
4143
sub_dirs_searched: list[tuple[str, ...]] = []
4244
file_wild = so_basename + "*"
4345
for rel_dir in rel_dirs:
44-
sub_dir = tuple(rel_dir.split(os.path.sep))
46+
sub_dir = PurePath(rel_dir).parts
4547
for abs_dir in find_sub_dirs_all_sitepackages(sub_dir):
4648
# Exact unversioned match first; fall back to versioned names because some
4749
# distros only ship lib<name>.so.<major> (e.g. conda libcupti). Only one match
@@ -78,7 +80,7 @@ def _find_dll_in_rel_dirs(
7880
) -> str | None:
7981
sub_dirs_searched: list[tuple[str, ...]] = []
8082
for rel_dir in rel_dirs:
81-
sub_dir = tuple(rel_dir.split(os.path.sep))
83+
sub_dir = PurePath(rel_dir).parts
8284
for abs_dir in find_sub_dirs_all_sitepackages(sub_dir):
8385
dll_name = _find_dll_under_dir(abs_dir, lib_searched_for)
8486
if dll_name is not None:
@@ -173,17 +175,19 @@ def find_in_lib_dir(
173175

174176
@dataclass(frozen=True, slots=True)
175177
class WindowsSearchPlatform:
178+
target_arch: str = field(default_factory=windows_python_arch)
179+
176180
def lib_searched_for(self, libname: str) -> str:
177181
return f"{libname}*.dll"
178182

179183
def site_packages_rel_dirs(self, desc: LibDescriptor) -> tuple[str, ...]:
180-
return cast(tuple[str, ...], desc.site_packages_windows)
184+
return desc.site_packages_windows.for_arch(self.target_arch)
181185

182186
def conda_anchor_point(self, conda_prefix: str) -> str:
183187
return os.path.join(conda_prefix, "Library")
184188

185189
def anchor_rel_dirs(self, desc: LibDescriptor) -> tuple[str, ...]:
186-
return cast(tuple[str, ...], desc.anchor_rel_dirs_windows)
190+
return desc.anchor_rel_dirs_windows.for_arch(self.target_arch)
187191

188192
def find_in_site_packages(
189193
self,

‎cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_steps.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -121,13 +121,14 @@ def _derive_ctk_root_windows(resolved_lib_path: str) -> str | None:
121121
122122
Supports:
123123
- ``$CTK_ROOT/bin/x64/foo.dll`` (CTK 13 style)
124+
- ``$CTK_ROOT/bin/arm64/foo.dll`` (Windows on Arm CTK 13 style)
124125
- ``$CTK_ROOT/bin/foo.dll`` (CTK 12 style)
125126
"""
126127
import ntpath
127128

128129
lib_dir = ntpath.dirname(resolved_lib_path)
129130
basename = ntpath.basename(lib_dir).lower()
130-
if basename == "x64":
131+
if basename in ("x64", "arm64"):
131132
parent = ntpath.dirname(lib_dir)
132133
if ntpath.basename(parent).lower() == "bin":
133134
return ntpath.dirname(parent)

‎cuda_pathfinder/cuda/pathfinder/_dynamic_libs/supported_nvidia_libs.py‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -60,11 +60,16 @@
6060
}
6161
SITE_PACKAGES_LIBDIRS_LINUX = SITE_PACKAGES_LIBDIRS_LINUX_CTK | SITE_PACKAGES_LIBDIRS_LINUX_OTHER
6262

63+
# Historical table exports represent the original x64 catalog.
6364
SITE_PACKAGES_LIBDIRS_WINDOWS_CTK = {
64-
desc.name: desc.site_packages_windows for desc in _CTK_DESCRIPTORS if desc.site_packages_windows
65+
desc.name: desc.site_packages_windows.for_arch("x64")
66+
for desc in _CTK_DESCRIPTORS
67+
if desc.site_packages_windows.paths
6568
}
6669
SITE_PACKAGES_LIBDIRS_WINDOWS_OTHER = {
67-
desc.name: desc.site_packages_windows for desc in _NON_CTK_DESCRIPTORS if desc.site_packages_windows
70+
desc.name: desc.site_packages_windows.for_arch("x64")
71+
for desc in _NON_CTK_DESCRIPTORS
72+
if desc.site_packages_windows.paths
6873
}
6974
SITE_PACKAGES_LIBDIRS_WINDOWS = SITE_PACKAGES_LIBDIRS_WINDOWS_CTK | SITE_PACKAGES_LIBDIRS_WINDOWS_OTHER
7075

‎cuda_pathfinder/cuda/pathfinder/_utils/windows_arch.py‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
4+
from __future__ import annotations
5+
6+
import sysconfig
7+
8+
9+
class UnsupportedWindowsArchError(RuntimeError):
10+
"""Raised when Python reports an unsupported Windows architecture."""
11+
12+
def __init__(self, platform_tag: str) -> None:
13+
self.platform_tag = platform_tag
14+
super().__init__(
15+
f"Unsupported Windows Python platform tag: {platform_tag!r}; expected 'win-amd64' or 'win-arm64'"
16+
)
17+
18+
19+
def windows_python_arch() -> str:
20+
"""Return the current Windows Python interpreter architecture."""
21+
raw_platform_tag = sysconfig.get_platform()
22+
platform_tag = raw_platform_tag.lower().replace("_", "-")
23+
24+
if platform_tag == "win-arm64":
25+
return "arm64"
26+
27+
if platform_tag == "win-amd64":
28+
return "x64"
29+
30+
raise UnsupportedWindowsArchError(raw_platform_tag)

‎cuda_pathfinder/docs/source/api.rst‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ CUDA bitcode and static libraries.
2424
DynamicLibNotFoundError
2525
DynamicLibUnknownError
2626
DynamicLibNotAvailableError
27+
UnsupportedWindowsArchError
2728

2829
SUPPORTED_HEADERS_CTK
2930
find_nvidia_header_directory

‎cuda_pathfinder/tests/test_ctk_root_discovery.py‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
_resolve_system_loaded_abs_path_in_subprocess,
1919
_try_ctk_root_canary,
2020
)
21+
from cuda.pathfinder._dynamic_libs.search_platform import WindowsSearchPlatform
2122
from cuda.pathfinder._dynamic_libs.search_steps import (
2223
SearchContext,
2324
_derive_ctk_root_linux,
@@ -125,6 +126,11 @@ def test_derive_ctk_root_windows_ctk13():
125126
assert _derive_ctk_root_windows(path) == r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.0"
126127

127128

129+
def test_derive_ctk_root_windows_ctk13_arm64():
130+
path = r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4\bin\arm64\cudart64_13.dll"
131+
assert _derive_ctk_root_windows(path) == r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4"
132+
133+
128134
def test_derive_ctk_root_windows_ctk12():
129135
path = r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.8\bin\cudart64_12.dll"
130136
assert _derive_ctk_root_windows(path) == r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.8"
@@ -189,6 +195,23 @@ def test_try_via_ctk_root_regular_lib(tmp_path):
189195
assert result.found_via == "system-ctk-root"
190196

191197

198+
def test_try_via_ctk_root_windows_arm64_prefers_arch_dir(tmp_path):
199+
ctk_root = tmp_path / "cuda-13"
200+
x64_dir = ctk_root / "bin" / "x64"
201+
arm64_dir = ctk_root / "bin" / "arm64"
202+
x64_dir.mkdir(parents=True)
203+
arm64_dir.mkdir(parents=True)
204+
(x64_dir / "cudart64_13.dll").write_bytes(b"fake")
205+
arm64_lib = arm64_dir / "cudart64_13.dll"
206+
arm64_lib.write_bytes(b"fake")
207+
208+
ctx = SearchContext(LIB_DESCRIPTORS["cudart"], platform=WindowsSearchPlatform(target_arch="arm64"))
209+
result = find_via_ctk_root(ctx, str(ctk_root))
210+
assert result is not None
211+
assert result.abs_path == str(arm64_lib)
212+
assert result.found_via == "system-ctk-root"
213+
214+
192215
# ---------------------------------------------------------------------------
193216
# _resolve_system_loaded_abs_path_in_subprocess
194217
# ---------------------------------------------------------------------------

‎cuda_pathfinder/tests/test_descriptor_catalog.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,7 @@ def test_no_self_dependency(spec: DescriptorSpec):
5959
def test_driver_libs_have_no_site_packages(spec: DescriptorSpec):
6060
"""Driver libs are system-search-only; site-packages paths would be unused."""
6161
assert not spec.site_packages_linux, f"driver lib {spec.name} has site_packages_linux"
62-
assert not spec.site_packages_windows, f"driver lib {spec.name} has site_packages_windows"
62+
assert not spec.site_packages_windows.paths, f"driver lib {spec.name} has site_packages_windows"
6363

6464

6565
@pytest.mark.parametrize(

‎cuda_pathfinder/tests/test_lib_descriptor.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@ def test_site_packages_linux_match(name):
5858

5959
@pytest.mark.parametrize("name", sorted(LIB_DESCRIPTORS))
6060
def test_site_packages_windows_match(name):
61-
assert LIB_DESCRIPTORS[name].site_packages_windows == SITE_PACKAGES_LIBDIRS_WINDOWS.get(name, ())
61+
assert LIB_DESCRIPTORS[name].site_packages_windows.for_arch("x64") == SITE_PACKAGES_LIBDIRS_WINDOWS.get(name, ())
6262

6363

6464
@pytest.mark.parametrize("name", sorted(LIB_DESCRIPTORS))

0 commit comments

Comments
 (0)