Fix Windows binary utility discovery on Arm64

cuda_pathfinder/cuda/pathfinder/_binaries/find_nvidia_binary_utility.py

@@ -2,13 +2,29 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import functools
+import importlib
 import os
+from collections.abc import Iterable
+from typing import Any
 
 from cuda.pathfinder._binaries import supported_nvidia_binaries
 from cuda.pathfinder._utils.ctk_root_canary import CTK_ROOT_CANARY_ANCHOR_LIBNAMES
 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_machine_arch
+
+_NSIGHT_REGISTRY_ROOT = r"SOFTWARE\NVIDIA Corporation\Installed Products\Nsight"
+
+_NSYS_TARGET_DIR_BY_ARCH = {
+    "x64": "target-windows-x64",
+    "arm64": "target-windows-armv8",
+}
+
+_NCU_TARGET_DIR_BY_ARCH = {
+    "x64": os.path.join("target", "windows-desktop-win7-x64"),
+    "arm64": os.path.join("target", "windows-desktop-win10-t23x-a64"),
+}
 
 
 class UnsupportedBinaryError(Exception):
@@ -46,6 +62,82 @@ def _ctk_bin_subdirs(root: str) -> list[str]:
     return [os.path.join(root, "bin")]
 
 
+def _resolve_candidate_paths(candidates: Iterable[str]) -> str | None:
+    """Return the first executable candidate, preserving candidate order."""
+    seen: set[str] = set()
+    for candidate in candidates:
+        if candidate in seen:
+            continue
+        seen.add(candidate)
+        if _is_executable_candidate(candidate):
+            return os.path.abspath(candidate)
+    return None
+
+
+def _find_windows_compute_sanitizer(ctk_root: str) -> str | None:
+    return _resolve_candidate_paths(
+        (
+            os.path.join(ctk_root, "bin", "compute-sanitizer.bat"),
+            os.path.join(ctk_root, "compute-sanitizer", "compute-sanitizer.exe"),
+        )
+    )
+
+
+def _windows_installed_nsight_root(product: str) -> str | None:
+    """Return the active Nsight product installation recorded by its MSI."""
+    # ``winreg`` attributes are absent from the type stubs on non-Windows hosts.
+    winreg: Any = importlib.import_module("winreg")
+
+    access = winreg.KEY_READ | winreg.KEY_WOW64_64KEY
+    product_key_path = rf"{_NSIGHT_REGISTRY_ROOT}\{product}"
+    try:
+        product_context = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, product_key_path, 0, access)
+    except FileNotFoundError:
+        return None
+
+    try:
+        with product_context as product_key:
+            current_version, _ = winreg.QueryValueEx(product_key, "CurrentVersion")
+            if not isinstance(current_version, str) or not current_version.strip():
+                raise RuntimeError(
+                    f"Invalid CurrentVersion value {current_version!r} in "
+                    f"Nsight {product!r} registry registration at {product_key_path!r}"
+                )
+            with winreg.OpenKey(product_key, current_version, 0, access) as version_key:
+                install_root, _ = winreg.QueryValueEx(version_key, None)
+    except FileNotFoundError as exc:
+        raise RuntimeError(f"Incomplete Nsight {product!r} registry registration at {product_key_path!r}") from exc
+
+    if not isinstance(install_root, str) or not install_root.strip():
+        raise RuntimeError(
+            f"Invalid installation directory {install_root!r} in Nsight {product!r} "
+            f"registry registration at {product_key_path!r} version {current_version!r}"
+        )
+    return install_root
+
+
+def _find_windows_nsys() -> str | None:
+    install_root = _windows_installed_nsight_root("Systems")
+    if install_root is None:
+        return None
+
+    target_dir = _NSYS_TARGET_DIR_BY_ARCH[windows_machine_arch()]
+    return _resolve_candidate_paths((os.path.join(install_root, target_dir, "nsys.exe"),))
+
+
+def _find_windows_ncu() -> str | None:
+    install_root = _windows_installed_nsight_root("Compute")
+    if install_root is None:
+        return None
+
+    launcher = os.path.join(install_root, "ncu.bat")
+    if (found := _resolve_candidate_paths((launcher,))) is not None:
+        return found
+
+    target_dir = _NCU_TARGET_DIR_BY_ARCH[windows_machine_arch()]
+    return _resolve_candidate_paths((os.path.join(install_root, target_dir, "ncu.exe"),))
+
+
 def _resolve_ctk_root_via_canary() -> str | None:
     from cuda.pathfinder._dynamic_libs.load_nvidia_dynamic_lib import resolve_ctk_root_via_canary
 
@@ -69,6 +161,20 @@ def _resolve_in_trusted_dirs(normalized_name: str, dirs: list[str]) -> str | Non
     return None
 
 
+def _resolve_names_in_trusted_dirs(candidate_names: tuple[str, ...], dirs: list[str]) -> str | None:
+    """Resolve ordered candidate names within each trusted directory."""
+    seen: set[str] = set()
+    for directory in dirs:
+        if directory in seen:
+            continue
+        assert directory
+        seen.add(directory)
+        found = _resolve_candidate_paths(os.path.join(directory, name) for name in candidate_names)
+        if found is not None:
+            return found
+    return None
+
+
 @functools.cache
 def find_nvidia_binary_utility(utility_name: str) -> str | None:
     """Locate a CUDA binary utility executable.
@@ -87,6 +193,19 @@ def find_nvidia_binary_utility(utility_name: str) -> str | None:
     Raises:
         UnsupportedBinaryError: If ``utility_name`` is not in the supported set
             (see ``SUPPORTED_BINARY_UTILITIES``).
+        RuntimeError: If a native Windows architecture needed for an
+            architecture-specific utility layout cannot be determined, or an
+            installed Nsight product has incomplete or invalid registry data.
+
+    Windows on ARM (WoA) Note:
+        Binary utilities execute in separate processes and do not need to match
+        the Python process architecture. When choosing among architecture-specific
+        Windows layouts, this API deliberately targets the native machine
+        architecture rather than the Python interpreter architecture. For
+        example, standalone ``nsys`` and ``ncu`` discovery under x64 Python on an
+        Arm64 machine selects the Arm64 target. This differs from
+        ``load_nvidia_dynamic_lib`` and ``find_static_lib``, which target the
+        Python interpreter architecture.
 
     Search order:
         1. **NVIDIA Python wheels**
@@ -100,17 +219,27 @@ def find_nvidia_binary_utility(utility_name: str) -> str | None:
              environment variable, which use platform-specific bin directory
              layouts (``Library/bin`` on Windows, ``bin`` on Linux).
 
-        3. **CUDA Toolkit environment variables**
+        3. **Library-specific standalone installations**
+
+           - Search the installation paths for the CUDA Toolkit, Nsight Systems,
+             and Nsight Compute.
 
-           - Use ``CUDA_HOME`` or ``CUDA_PATH`` (in that order), searching
-             ``bin/x64``, ``bin/x86_64``, and ``bin`` subdirectories on Windows,
-             or just ``bin`` on Linux.
+           3.1. **Nsight installations**: On Windows, locate Nsight Systems and
+                Nsight Compute from their installer registry entries. Select
+                architecture-specific binaries using the native machine
+                architecture, independent of Python. Lookup of the standalone
+                ``nsys`` and ``ncu`` CLIs is terminal; a miss does not fall
+                through to CUDA Toolkit locations.
+
+           3.2. **CUDA Toolkit installation**: Use ``CUDA_HOME`` or ``CUDA_PATH``
+                (in that order), searching ``bin/x64``, ``bin/x86_64``, and
+                ``bin`` subdirectories on Windows, or just ``bin`` on Linux.
 
         4. **CTK-root canary fallback**
 
-           - Only when steps 1-3 miss: resolve the ``cudart`` library through the
-             OS dynamic loader, derive the CUDA Toolkit root from it, and search
-             that root's bin layout.
+           - For utilities that reach this step after the earlier searches miss,
+             resolve the ``cudart`` library through the OS dynamic loader, derive
+             the CUDA Toolkit root from it, and search that root's bin layout.
 
     Note:
         Results are cached using ``@functools.cache`` for performance. The cache
@@ -146,17 +275,35 @@ def find_nvidia_binary_utility(utility_name: str) -> str | None:
         else:
             dirs.append(os.path.join(conda_prefix, "bin"))
 
-    # 3. Search in CUDA Toolkit (CUDA_HOME/CUDA_PATH)
-    if (cuda_home := get_cuda_path_or_home()) is not None:
-        dirs.extend(_ctk_bin_subdirs(cuda_home))
-
     normalized_name = _normalize_utility_name(utility_name)
-    found = _resolve_in_trusted_dirs(normalized_name, dirs)
+    if IS_WINDOWS and utility_name in ("compute-sanitizer", "ncu"):
+        candidate_names = (f"{utility_name}.bat", normalized_name)
+        found = _resolve_names_in_trusted_dirs(candidate_names, dirs)
+    else:
+        found = _resolve_in_trusted_dirs(normalized_name, dirs)
     if found is not None:
         return found
 
+    # 3. Search library-specific standalone installations.
+    # 3.1. Standalone Nsight CLI lookup is terminal; CTK does not contain nsys/ncu.
+    if IS_WINDOWS and utility_name == "nsys":
+        return _find_windows_nsys()
+    if IS_WINDOWS and utility_name == "ncu":
+        return _find_windows_ncu()
+
+    # 3.2. Search in CUDA Toolkit (CUDA_HOME/CUDA_PATH).
+    if (cuda_home := get_cuda_path_or_home()) is not None:
+        if IS_WINDOWS and utility_name == "compute-sanitizer":
+            found = _find_windows_compute_sanitizer(cuda_home)
+        else:
+            found = _resolve_in_trusted_dirs(normalized_name, _ctk_bin_subdirs(cuda_home))
+        if found is not None:
+            return found
+
     # 4. CTK-root canary fallback.
     ctk_root = _resolve_ctk_root_via_canary()
     if ctk_root is not None:
+        if IS_WINDOWS and utility_name == "compute-sanitizer":
+            return _find_windows_compute_sanitizer(ctk_root)
         return _resolve_in_trusted_dirs(normalized_name, _ctk_bin_subdirs(ctk_root))
     return None

cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_nvidia_dynamic_lib.py

@@ -231,6 +231,15 @@ def load_nvidia_dynamic_lib(libname: str) -> LoadedDL:
         DynamicLibNotFoundError: If the library cannot be found or loaded.
         RuntimeError: If Python is not 64-bit.
 
+    Windows on ARM (WoA) Note:
+        On Windows, this API aims to load a dynamic library whose architecture
+        matches the Python interpreter architecture. For example, x64 Python
+        running on an Arm64 machine targets an x64 DLL, while native Arm64 Python
+        targets an Arm64 DLL. A library loaded into the Python process must be
+        compatible with that process. This differs from
+        ``find_nvidia_binary_utility``, which targets the native machine
+        architecture when selecting architecture-specific executables.
+
     Search order:
         0. **Already loaded in the current process**
 

cuda_pathfinder/cuda/pathfinder/_static_libs/find_static_lib.py

@@ -175,5 +175,13 @@ def find_static_lib(name: str) -> str:
     Raises:
         ValueError: If ``name`` is not a supported static library.
         StaticLibNotFoundError: If the static library cannot be found.
+
+    Windows on ARM (WoA) Note:
+        On Windows, this API aims to return the path to a static library whose
+        architecture matches the Python interpreter architecture. For example,
+        x64 Python running on an Arm64 machine targets the x64 library, while
+        native Arm64 Python targets the Arm64 library. This differs from
+        ``find_nvidia_binary_utility``, which targets the native machine
+        architecture when selecting architecture-specific executables.
     """
     return locate_static_lib(name).abs_path

cuda_pathfinder/cuda/pathfinder/_utils/windows_arch.py

@@ -3,13 +3,16 @@
 
 from __future__ import annotations
 
+import platform
 import sysconfig
 
 WINDOWS_PE_MACHINE_BY_ARCH = {
     "x64": 0x8664,
     "arm64": 0xAA64,
 }
 
+_WINDOWS_ARCH_BY_PE_MACHINE = {machine: arch for arch, machine in WINDOWS_PE_MACHINE_BY_ARCH.items()}
+
 
 class UnsupportedArchError(RuntimeError):
     """Raised when Python reports an unsupported Windows architecture."""
@@ -35,6 +38,76 @@ def windows_python_arch() -> str:
     raise UnsupportedArchError(raw_platform_tag)
 
 
+def _windows_machine_arch_from_platform() -> str:
+    """Return the Windows architecture reported by Python's platform module."""
+    raw_machine = platform.machine()
+    machine = raw_machine.lower().replace("_", "-")
+
+    if machine in ("amd64", "x86-64"):
+        return "x64"
+
+    if machine in ("arm64", "aarch64"):
+        return "arm64"
+
+    raise RuntimeError(f"Unsupported Windows machine architecture: {raw_machine!r}")
+
+
+def _windows_native_machine() -> int | None:
+    """Return the native Windows PE machine type, or None on older Windows."""
+    import ctypes
+    from ctypes import wintypes
+
+    try:
+        # These ctypes attributes are absent from the type stubs on non-Windows hosts.
+        kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)  # type: ignore[attr-defined, unused-ignore]
+    except OSError as exc:
+        raise RuntimeError("Failed to load kernel32 while detecting the native Windows architecture") from exc
+
+    get_current_process = kernel32.GetCurrentProcess
+    try:
+        is_wow64_process2 = kernel32.IsWow64Process2
+    except AttributeError:
+        return None
+
+    get_current_process.argtypes = ()
+    get_current_process.restype = wintypes.HANDLE
+    is_wow64_process2.argtypes = (
+        wintypes.HANDLE,
+        ctypes.POINTER(wintypes.USHORT),
+        ctypes.POINTER(wintypes.USHORT),
+    )
+    is_wow64_process2.restype = wintypes.BOOL
+
+    process_machine = wintypes.USHORT()
+    native_machine = wintypes.USHORT()
+    if not is_wow64_process2(
+        get_current_process(),
+        ctypes.byref(process_machine),
+        ctypes.byref(native_machine),
+    ):
+        error_code = ctypes.get_last_error()  # type: ignore[attr-defined, unused-ignore]
+        error = ctypes.WinError(error_code)  # type: ignore[attr-defined, unused-ignore]
+        raise RuntimeError(
+            f"IsWow64Process2 failed while detecting the native Windows architecture "
+            f"(Windows error {error_code}): {error}"
+        ) from error
+    return native_machine.value
+
+
+def windows_machine_arch() -> str:
+    """Return the native Windows machine architecture, ignoring process emulation."""
+    native_machine = _windows_native_machine()
+    if native_machine is None:
+        # IsWow64Process2 predates x64-on-Arm emulation, so this fallback is only
+        # needed on older Windows versions where platform.machine() is sufficient.
+        return _windows_machine_arch_from_platform()
+
+    try:
+        return _WINDOWS_ARCH_BY_PE_MACHINE[native_machine]
+    except KeyError:
+        raise RuntimeError(f"Unsupported native Windows PE machine type: 0x{native_machine:04x}") from None
+
+
 def windows_pe_matches_arch(path: str, target_arch: str) -> bool:
     """Return whether a Windows Portable Executable (PE) targets the requested architecture.
 

cuda_pathfinder/docs/source/install.rst

@@ -9,7 +9,7 @@ Runtime Requirements
 
 ``cuda.pathfinder`` is a pure-Python package with no runtime dependencies:
 
-* Linux (x86-64, arm64) and Windows (x86-64)
+* Linux (x86-64, arm64) and Windows (x86-64, arm64)
 * Python 3.10 - 3.14
 
 Installing from PyPI