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,71 @@ 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:
+        with winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, product_key_path, 0, access) as product_key:
+            current_version, _ = winreg.QueryValueEx(product_key, "CurrentVersion")
+            if not isinstance(current_version, str) or not current_version:
+                raise RuntimeError(f"Invalid CurrentVersion value in {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:
+        return None
+
+    if not isinstance(install_root, str) or not install_root:
+        raise RuntimeError(f"Invalid installation directory for {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 +150,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.
@@ -88,6 +183,16 @@ def find_nvidia_binary_utility(utility_name: str) -> str | None:
         UnsupportedBinaryError: If ``utility_name`` is not in the supported set
             (see ``SUPPORTED_BINARY_UTILITIES``).
 
+    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 +205,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 +261,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,6 +3,7 @@
 
 from __future__ import annotations
 
+import platform
 import sysconfig
 
 WINDOWS_PE_MACHINE_BY_ARCH = {
@@ -35,6 +36,20 @@ def windows_python_arch() -> str:
     raise UnsupportedArchError(raw_platform_tag)
 
 
+def windows_machine_arch() -> str:
+    """Return the native Windows machine architecture."""
+    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_pe_matches_arch(path: str, target_arch: str) -> bool:
     """Return whether a Windows Portable Executable (PE) targets the requested architecture.