initial commit

cuda_pathfinder/cuda/pathfinder/__init__.py

@@ -5,6 +5,8 @@
 
 from cuda.pathfinder._version import __version__  # noqa: F401
 
+from cuda.pathfinder._binaries.find_nvidia_binaries import find_nvidia_binary as find_nvidia_binary
+from cuda.pathfinder._binaries.supported_nvidia_binaries import SUPPORTED_BINARIES as SUPPORTED_NVIDIA_BINARIES
 from cuda.pathfinder._dynamic_libs.load_dl_common import DynamicLibNotFoundError as DynamicLibNotFoundError
 from cuda.pathfinder._dynamic_libs.load_dl_common import LoadedDL as LoadedDL
 from cuda.pathfinder._dynamic_libs.load_nvidia_dynamic_lib import load_nvidia_dynamic_lib as load_nvidia_dynamic_lib
@@ -13,6 +15,15 @@
 )
 from cuda.pathfinder._headers.find_nvidia_headers import find_nvidia_header_directory as find_nvidia_header_directory
 from cuda.pathfinder._headers.supported_nvidia_headers import SUPPORTED_HEADERS_CTK as _SUPPORTED_HEADERS_CTK
+from cuda.pathfinder._static_libs.find_nvidia_static_libs import find_nvidia_static_lib as find_nvidia_static_lib
+from cuda.pathfinder._static_libs.supported_nvidia_static_libs import (
+    SUPPORTED_STATIC_LIBS as SUPPORTED_NVIDIA_STATIC_LIBS,
+)
+from cuda.pathfinder._utils.toolchain_tracker import (
+    SearchContext as SearchContext,
+    ToolchainMismatchError as ToolchainMismatchError,
+    reset_default_context as reset_default_context,
+)
 
 # Indirections to help Sphinx find the docstrings.
 #: Mapping from short CUDA Toolkit (CTK) library names to their canonical

cuda_pathfinder/cuda/pathfinder/_binaries/find_nvidia_binaries.py

@@ -0,0 +1,110 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import functools
+import os
+from typing import Sequence
+
+from cuda.pathfinder._binaries.supported_nvidia_binaries import SITE_PACKAGES_BINDIRS, SUPPORTED_BINARIES
+from cuda.pathfinder._utils.env_vars import get_cuda_home_or_path
+from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages
+from cuda.pathfinder._utils.path_utils import _abs_norm
+from cuda.pathfinder._utils.toolchain_tracker import (
+    SearchContext,
+    SearchLocation,
+    ToolchainSource,
+    get_default_context,
+)
+
+
+def _binary_filename_variants(name: str) -> Sequence[str]:
+    """Generate filename variants for a binary (cross-platform).
+
+    Args:
+        name: Base binary name.
+
+    Returns:
+        Tuple of possible filenames (e.g., "nvcc", "nvcc.exe").
+    """
+    return (name, f"{name}.exe")
+
+
+def _get_site_packages_subdirs(binary_name: str) -> Sequence[str]:
+    """Get site-packages sub-directories for a binary.
+
+    Args:
+        binary_name: Name of the binary.
+
+    Returns:
+        List of sub-directories to search, or empty list if binary not in site-packages.
+    """
+    relative_directories = SITE_PACKAGES_BINDIRS.get(binary_name)
+    if not relative_directories:
+        return []
+
+    # Expand site-packages paths
+    sub_directories = []
+    for relative_directory in relative_directories:
+        for found_dir in find_sub_dirs_all_sitepackages(tuple(relative_directory.split("/"))):
+            sub_directories.append(found_dir)
+    return sub_directories
+
+
+# Define search locations for binaries
+def _create_search_locations(binary_name: str) -> list[SearchLocation]:
+    """Create search location configurations for a specific binary.
+
+    Args:
+        binary_name: Name of the binary to search for.
+
+    Returns:
+        List of SearchLocation objects to try.
+    """
+    return [
+        SearchLocation(
+            source=ToolchainSource.SITE_PACKAGES,
+            base_dir_func=lambda: None,  # Use subdirs for full paths
+            subdirs=_get_site_packages_subdirs(binary_name),
+            filename_variants=_binary_filename_variants,
+        ),
+        SearchLocation(
+            source=ToolchainSource.CONDA,
+            base_dir_func=lambda: os.environ.get("CONDA_PREFIX"),
+            subdirs=["Library/bin", "bin"],  # Windows and Unix layouts
+            filename_variants=_binary_filename_variants,
+        ),
+        SearchLocation(
+            source=ToolchainSource.CUDA_HOME,
+            base_dir_func=get_cuda_home_or_path,
+            subdirs=["bin"],
+            filename_variants=_binary_filename_variants,
+        ),
+    ]
+
+
+@functools.cache
+def find_nvidia_binary(binary_name: str, *, context: SearchContext | None = None) -> str | None:
+    """Locate a CUDA binary executable.
+
+    Args:
+        binary_name: Name of the binary (e.g., "nvdisasm", "cuobjdump").
+        context: Optional SearchContext for toolchain consistency tracking.
+            If None, uses the default module-level context.
+
+    Returns:
+        Absolute path to the binary, or None if not found.
+
+    Raises:
+        RuntimeError: If binary_name is not supported.
+        ToolchainMismatchError: If binary found in different source than
+            the context's preferred source.
+    """
+    if binary_name not in SUPPORTED_BINARIES:
+        raise RuntimeError(f"UNKNOWN {binary_name=}")
+
+    if context is None:
+        context = get_default_context()
+
+    locations = _create_search_locations(binary_name)
+    path = context.find(binary_name, locations)
+    return _abs_norm(path) if path else None

cuda_pathfinder/cuda/pathfinder/_binaries/supported_nvidia_binaries.py

@@ -0,0 +1,24 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# THIS FILE NEEDS TO BE REVIEWED/UPDATED FOR EACH CTK RELEASE
+# Likely candidates for updates are:
+#     SUPPORTED_BINARIES
+#     SITE_PACKAGES_BINDIRS
+
+from cuda.pathfinder._utils.platform_aware import IS_WINDOWS
+
+# Supported CUDA binaries that can be found
+SUPPORTED_BINARIES_COMMON = (
+    "nvdisasm",
+    "cuobjdump",
+)
+
+SUPPORTED_BINARIES = SUPPORTED_BINARIES_COMMON
+
+# Map from binary name to relative paths under site-packages
+# These are typically from cuda-toolkit[nvcc] wheels
+SITE_PACKAGES_BINDIRS = {
+    "nvdisasm": ["nvidia/cuda_nvcc/bin"],
+    "cuobjdump": ["nvidia/cuda_nvcc/bin"],
+}

cuda_pathfinder/cuda/pathfinder/_headers/find_nvidia_headers.py

@@ -8,15 +8,10 @@
 from cuda.pathfinder._headers import supported_nvidia_headers
 from cuda.pathfinder._utils.env_vars import get_cuda_home_or_path
 from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages
+from cuda.pathfinder._utils.path_utils import _abs_norm
 from cuda.pathfinder._utils.platform_aware import IS_WINDOWS
 
 
-def _abs_norm(path: str | None) -> str | None:
-    if path:
-        return os.path.normpath(os.path.abspath(path))
-    return None
-
-
 def _joined_isfile(dirpath: str, basename: str) -> bool:
     return os.path.isfile(os.path.join(dirpath, basename))
 

cuda_pathfinder/cuda/pathfinder/_static_libs/find_nvidia_static_libs.py

@@ -0,0 +1,164 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import functools
+import os
+from typing import Sequence
+
+from cuda.pathfinder._static_libs.supported_nvidia_static_libs import (
+    SITE_PACKAGES_STATIC_LIBDIRS,
+    SUPPORTED_STATIC_LIBS,
+)
+from cuda.pathfinder._utils.env_vars import get_cuda_home_or_path
+from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages
+from cuda.pathfinder._utils.path_utils import _abs_norm
+from cuda.pathfinder._utils.platform_aware import IS_WINDOWS
+from cuda.pathfinder._utils.toolchain_tracker import (
+    SearchContext,
+    SearchLocation,
+    ToolchainSource,
+    get_default_context,
+)
+
+
+# Generic search locations by toolchain source (platform-specific at import time)
+if IS_WINDOWS:
+    CONDA_LIB_SUBDIRS = ("Library/lib", "Library/lib/x64")
+    CONDA_NVVM_SUBDIRS = ("Library/nvvm/libdevice",)
+    CUDA_HOME_LIB_SUBDIRS = ("lib/x64", "lib")
+    CUDA_HOME_NVVM_SUBDIRS = ("nvvm/libdevice",)
+else:
+    CONDA_LIB_SUBDIRS = ("lib",)
+    CONDA_NVVM_SUBDIRS = ("nvvm/libdevice",)
+    CUDA_HOME_LIB_SUBDIRS = ("lib64", "lib")
+    CUDA_HOME_NVVM_SUBDIRS = ("nvvm/libdevice",)
+
+
+def _static_lib_filename_variants(artifact_name: str) -> Sequence[str]:
+    """Generate platform-appropriate filename variants for an artifact.
+
+    Args:
+        artifact_name: Canonical artifact name (e.g., "cudadevrt", "libdevice.10.bc").
+
+    Returns:
+        Sequence of filenames to search for on this platform.
+
+    Examples:
+        On Windows:
+            "cudadevrt" -> ("cudadevrt.lib",)
+            "libdevice.10.bc" -> ("libdevice.10.bc",)
+        On Linux:
+            "cudadevrt" -> ("libcudadevrt.a",)
+            "libdevice.10.bc" -> ("libdevice.10.bc",)
+    """
+    # Files that are the same on all platforms (e.g., .bc bitcode files)
+    if "." in artifact_name:
+        return (artifact_name,)
+
+    # Platform-specific library naming conventions
+    if IS_WINDOWS:
+        return (f"{artifact_name}.lib",)
+    else:
+        return (f"lib{artifact_name}.a",)
+
+
+def _get_cuda_home_subdirs_with_targets() -> tuple[str, ...]:
+    """Get CUDA_HOME subdirectories including expanded targets/* paths.
+
+    Returns:
+        Tuple of subdirectories to search under CUDA_HOME.
+    """
+    import glob
+
+    subdirs = list(CUDA_HOME_LIB_SUBDIRS + CUDA_HOME_NVVM_SUBDIRS)
+
+    # On Linux, also search targets/*/lib64 and targets/*/lib for cross-compilation
+    if not IS_WINDOWS:
+        cuda_home = get_cuda_home_or_path()
+        if cuda_home:
+            for lib_subdir in ("lib64", "lib"):
+                pattern = os.path.join(cuda_home, "targets", "*", lib_subdir)
+                for target_dir in sorted(glob.glob(pattern), reverse=True):
+                    # Make relative to cuda_home
+                    rel_path = os.path.relpath(target_dir, cuda_home)
+                    subdirs.append(rel_path)
+
+    return tuple(subdirs)
+
+
+def _create_search_locations(artifact_name: str) -> list[SearchLocation]:
+    """Create generic search location configurations.
+
+    Args:
+        artifact_name: Name of the artifact to search for.
+
+    Returns:
+        List of SearchLocation objects to try.
+    """
+    locations = []
+
+    # Site-packages: Create separate SearchLocation for each found directory
+    relative_directories = SITE_PACKAGES_STATIC_LIBDIRS.get(artifact_name)
+    if relative_directories:
+        for relative_directory in relative_directories:
+            for found_dir in find_sub_dirs_all_sitepackages(tuple(relative_directory.split("/"))):
+                locations.append(
+                    SearchLocation(
+                        source=ToolchainSource.SITE_PACKAGES,
+                        base_dir_func=lambda d=found_dir: d,
+                        subdirs=[""],
+                        filename_variants=_static_lib_filename_variants,
+                    )
+                )
+
+    # Conda: Generic lib and nvvm locations
+    locations.append(
+        SearchLocation(
+            source=ToolchainSource.CONDA,
+            base_dir_func=lambda: os.environ.get("CONDA_PREFIX"),
+            subdirs=CONDA_LIB_SUBDIRS + CONDA_NVVM_SUBDIRS,
+            filename_variants=_static_lib_filename_variants,
+        )
+    )
+
+    # CUDA_HOME: Generic lib and nvvm locations (including targets/* on Linux)
+    locations.append(
+        SearchLocation(
+            source=ToolchainSource.CUDA_HOME,
+            base_dir_func=get_cuda_home_or_path,
+            subdirs=_get_cuda_home_subdirs_with_targets(),
+            filename_variants=_static_lib_filename_variants,
+        )
+    )
+
+    return locations
+
+
+@functools.cache
+def find_nvidia_static_lib(artifact_name: str, *, context: SearchContext | None = None) -> str | None:
+    """Locate a CUDA static library or artifact file.
+
+    Args:
+        artifact_name: Canonical artifact name (e.g., "libdevice.10.bc", "cudadevrt").
+            Platform-specific filenames are resolved automatically:
+            - "cudadevrt" -> "libcudadevrt.a" on Linux, "cudadevrt.lib" on Windows
+        context: Optional SearchContext for toolchain consistency tracking.
+            If None, uses the default module-level context.
+
+    Returns:
+        Absolute path to the artifact, or None if not found.
+
+    Raises:
+        RuntimeError: If artifact_name is not supported.
+        ToolchainMismatchError: If artifact found in different source than
+            the context's preferred source.
+    """
+    if artifact_name not in SUPPORTED_STATIC_LIBS:
+        raise RuntimeError(f"UNKNOWN {artifact_name=}")
+
+    if context is None:
+        context = get_default_context()
+
+    locations = _create_search_locations(artifact_name)
+    path = context.find(artifact_name, locations)
+    return _abs_norm(path) if path else None

cuda_pathfinder/cuda/pathfinder/_static_libs/supported_nvidia_static_libs.py

@@ -0,0 +1,24 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# THIS FILE NEEDS TO BE REVIEWED/UPDATED FOR EACH CTK RELEASE
+# Likely candidates for updates are:
+#     SUPPORTED_STATIC_LIBS
+#     SITE_PACKAGES_STATIC_LIBDIRS
+
+# Supported CUDA static libraries and artifacts (canonical names)
+SUPPORTED_STATIC_LIBS_COMMON = (
+    "libdevice.10.bc",  # Bitcode file (same name on all platforms)
+    "cudadevrt",  # Static device runtime library (libcudadevrt.a on Linux, cudadevrt.lib on Windows)
+)
+
+SUPPORTED_STATIC_LIBS = SUPPORTED_STATIC_LIBS_COMMON
+
+# Map from canonical artifact name to relative paths under site-packages
+SITE_PACKAGES_STATIC_LIBDIRS = {
+    "libdevice.10.bc": ["nvidia/cuda_nvvm/nvvm/libdevice"],
+    "cudadevrt": [
+        "nvidia/cuda_cudart/lib",  # Linux
+        "nvidia/cuda_cudart/lib/x64",  # Windows (if present)
+    ],
+}

cuda_pathfinder/cuda/pathfinder/_utils/path_utils.py

@@ -0,0 +1,19 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+
+
+def _abs_norm(path: str | None) -> str | None:
+    """Normalize and convert a path to an absolute path.
+
+    Args:
+        path (str or None): The path to normalize and make absolute.
+
+    Returns:
+        str or None: The normalized absolute path, or None if the input is None
+        or empty.
+    """
+    if path:
+        return os.path.normpath(os.path.abspath(path))
+    return None