find_bitcode_lib.py

# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

import functools
import ntpath
import os
from dataclasses import dataclass
from typing import NoReturn, TypedDict

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


class BitcodeLibNotFoundError(RuntimeError):
    """Raised when a bitcode library cannot be found."""


@dataclass(frozen=True)
class LocatedBitcodeLib:
    """Information about a located bitcode library."""

    name: str
    abs_path: str
    filename: str
    found_via: str


class _BitcodeLibInfo(TypedDict):
    filename: str
    rel_path: str
    site_packages_dirs: tuple[str, ...]
    available_on_windows: bool


_SUPPORTED_BITCODE_LIBS_INFO: dict[str, _BitcodeLibInfo] = {
    "device": {
        "filename": "libdevice.10.bc",
        "rel_path": os.path.join("nvvm", "libdevice"),
        "site_packages_dirs": (
            "nvidia/cu13/nvvm/libdevice",
            "nvidia/cuda_nvcc/nvvm/libdevice",
        ),
        "available_on_windows": True,
    },
    "nccl_device": {
        "filename": "libnccl_device.bc",
        "rel_path": "lib",
        "site_packages_dirs": ("nvidia/nccl/lib",),
        "available_on_windows": False,
    },
    "nvshmem_device": {
        "filename": "libnvshmem_device.bc",
        "rel_path": "lib",
        "site_packages_dirs": ("nvidia/nvshmem/lib",),
        "available_on_windows": False,
    },
}

# Public API: just the supported library names
SUPPORTED_BITCODE_LIBS: tuple[str, ...] = tuple(
    sorted(
        name for name, info in _SUPPORTED_BITCODE_LIBS_INFO.items() if not IS_WINDOWS or info["available_on_windows"]
    )
)


def _no_such_file_in_dir(dir_path: str, filename: str, error_messages: list[str], attachments: list[str]) -> None:
    error_messages.append(f"No such file: {os.path.join(dir_path, filename)}")
    if os.path.isdir(dir_path):
        attachments.append(f'  listdir("{dir_path}"):')
        for node in sorted(os.listdir(dir_path)):
            attachments.append(f"    {node}")
    else:
        attachments.append(f'  Directory does not exist: "{dir_path}"')


def _validate_filename(filename: str) -> str:
    if not isinstance(filename, str):
        raise TypeError(f"filename must be a string, got {type(filename).__name__}")
    if filename in ("", ".", "..") or "\0" in filename or ntpath.basename(filename) != filename:
        raise ValueError(f"filename must be a file name without a directory: {filename!r}")
    return filename


class _FindBitcodeLib:
    def __init__(self, name: str, *, filename: str | None = None) -> None:
        if name not in _SUPPORTED_BITCODE_LIBS_INFO:  # Updated reference
            raise ValueError(f"Unknown bitcode library: '{name}'. Supported: {', '.join(SUPPORTED_BITCODE_LIBS)}")
        self.name: str = name
        self.config: _BitcodeLibInfo = _SUPPORTED_BITCODE_LIBS_INFO[name]  # Updated reference
        self.filename: str = self.config["filename"] if filename is None else _validate_filename(filename)
        self.rel_path: str = self.config["rel_path"]
        self.site_packages_dirs: tuple[str, ...] = self.config["site_packages_dirs"]
        self.error_messages: list[str] = []
        self.attachments: list[str] = []

    def try_site_packages(self) -> str | None:
        for rel_dir in self.site_packages_dirs:
            sub_dir = tuple(rel_dir.split("/"))
            for abs_dir in find_sub_dirs_all_sitepackages(sub_dir):
                file_path = os.path.join(abs_dir, self.filename)
                if os.path.isfile(file_path):
                    return file_path
        return None

    def try_with_conda_prefix(self) -> str | None:
        conda_prefix = os.environ.get("CONDA_PREFIX")
        if not conda_prefix:
            return None

        anchor = os.path.join(conda_prefix, "Library") if IS_WINDOWS else conda_prefix
        file_path = os.path.join(anchor, self.rel_path, self.filename)
        if os.path.isfile(file_path):
            return file_path
        return None

    def try_with_cuda_home(self) -> str | None:
        cuda_home = get_cuda_path_or_home()
        if cuda_home is None:
            self.error_messages.append("CUDA_HOME/CUDA_PATH not set")
            return None

        file_path = os.path.join(cuda_home, self.rel_path, self.filename)
        if os.path.isfile(file_path):
            return file_path

        _no_such_file_in_dir(
            os.path.join(cuda_home, self.rel_path),
            self.filename,
            self.error_messages,
            self.attachments,
        )
        return None

    def raise_not_found_error(self) -> NoReturn:
        err = ", ".join(self.error_messages) if self.error_messages else "No search paths available"
        att = "\n".join(self.attachments) if self.attachments else ""
        raise BitcodeLibNotFoundError(f'Failure finding "{self.filename}": {err}\n{att}')


def _locate_bitcode_lib(name: str, *, filename: str | None = None) -> LocatedBitcodeLib:
    finder = _FindBitcodeLib(name, filename=filename)

    abs_path = finder.try_site_packages()
    if abs_path is not None:
        return LocatedBitcodeLib(
            name=name,
            abs_path=abs_path,
            filename=finder.filename,
            found_via="site-packages",
        )

    abs_path = finder.try_with_conda_prefix()
    if abs_path is not None:
        return LocatedBitcodeLib(
            name=name,
            abs_path=abs_path,
            filename=finder.filename,
            found_via="conda",
        )

    abs_path = finder.try_with_cuda_home()
    if abs_path is not None:
        return LocatedBitcodeLib(
            name=name,
            abs_path=abs_path,
            filename=finder.filename,
            found_via="CUDA_PATH",
        )

    finder.raise_not_found_error()


def locate_bitcode_lib(name: str) -> LocatedBitcodeLib:
    """Locate a bitcode library by name.

    Raises:
        ValueError: If ``name`` is not a supported bitcode library.
        BitcodeLibNotFoundError: If the bitcode library cannot be found.
    """
    return _locate_bitcode_lib(name)


def locate_bitcode_lib_by_name(libname: str, filename: str) -> LocatedBitcodeLib:
    """Locate an exact bitcode filename using a supported library's search paths.

    This lets a library define its own filename convention, including any
    architecture or other attributes encoded in the filename.

    Args:
        libname: Supported bitcode library whose configured directories to search.
        filename: Exact file name to locate. Directory components are not allowed.

    Raises:
        TypeError: If ``filename`` is not a string.
        ValueError: If ``libname`` is unsupported or ``filename`` is not a file name.
        BitcodeLibNotFoundError: If ``filename`` cannot be found.
    """
    _validate_filename(filename)
    return _locate_bitcode_lib(libname, filename=filename)


@functools.cache
def find_bitcode_lib(name: str) -> str:
    """Find the absolute path to a bitcode library.

    Raises:
        ValueError: If ``name`` is not a supported bitcode library.
        BitcodeLibNotFoundError: If the bitcode library cannot be found.
    """
    return locate_bitcode_lib(name).abs_path


@functools.cache
def find_bitcode_lib_by_name(libname: str, filename: str) -> str:
    """Find an exact bitcode filename using a supported library's search paths.

    This lets a library define its own filename convention, including any
    architecture or other attributes encoded in the filename.

    Args:
        libname: Supported bitcode library whose configured directories to search.
        filename: Exact file name to find. Directory components are not allowed.

    Raises:
        TypeError: If ``filename`` is not a string.
        ValueError: If ``libname`` is unsupported or ``filename`` is not a file name.
        BitcodeLibNotFoundError: If ``filename`` cannot be found.
    """
    return locate_bitcode_lib_by_name(libname, filename).abs_path