Function find_nvshmem_paths

python/tvm/support/nvcc.py:765–848

Searches for the NVSHMEM include and library directories. Returns ------- A tuple containing the path to the include directory and the library directory.

()

def find_nvshmem_paths() -> tuple[str, str]:
    """
    Searches for the NVSHMEM include and library directories.

    Returns
    -------
    A tuple containing the path to the include directory and the library directory.
    """
    candidate_roots = []

    # 1. NVSHMEM_HOME env variable
    if "NVSHMEM_HOME" in os.environ:
        candidate_roots.append(os.environ["NVSHMEM_HOME"])

    # 2. CUDA Toolkit
    try:
        cuda_home = find_cuda_path()
        candidate_roots.append(cuda_home)
    except RuntimeError:
        pass

    # 3. Other common system installation paths
    candidate_roots.extend(["/usr/local", "/usr"])

    seen = set()
    unique_candidates = []
    for path in candidate_roots:
        if path and path not in seen:
            seen.add(path)
            unique_candidates.append(path)

    for root in unique_candidates:
        # Check both standard include path and versioned subdirectories (e.g., nvshmem_12)
        include_paths_to_check = [os.path.join(root, "include")]

        # Add versioned subdirectories like include/nvshmem_*
        versioned_includes = glob.glob(os.path.join(root, "include", "nvshmem_*"))
        include_paths_to_check.extend(versioned_includes)

        # Check standard and architecture-specific lib directories
        lib_paths_to_check = [
            os.path.join(root, "lib64"),
            os.path.join(root, "lib"),
        ]

        # Add architecture-specific lib paths (e.g., lib/x86_64-linux-gnu)
        machine = platform.machine()
        system = platform.system().lower()
        lib_paths_to_check.extend(
            [
                os.path.join(root, "lib", f"{machine}-{system}-gnu"),
                os.path.join(root, "lib64", f"{machine}-{system}-gnu"),
            ]
        )

        for include_path in include_paths_to_check:
            if os.path.isfile(os.path.join(include_path, "nvshmem.h")):
                for lib_path in lib_paths_to_check: