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.
@@ -100,17 +195,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 +251,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/_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.