setup.py

from copy import deepcopy
import typing as tp
import textwrap
import os
import subprocess
import sys

from setuptools import setup


def maybe_download_cub(torch_include_dirs: tp.Iterable[str]) -> str:
    import torch

    cuda_ver = float(torch.version.cuda) if torch.version.cuda is not None else 0
    # Cub is not required for cuda 12.8 or higher
    if cuda_ver >= 12.8:
        return ""
    print("-" * 75)
    print("The CUB library is needed to build the cuAEV extension")
    for d in torch_include_dirs:
        cubdir = os.path.join(d, "cub")
        print(f"Searching for CUB at {cubdir}...")
        if os.path.isdir(cubdir):
            print(f"Found CUB in {cubdir}")
            return ""
    # if no cub, download it to include dir from github
    if not os.path.isdir("./include/cub"):
        if not os.path.exists("./include"):
            os.makedirs("include")
        commands = """
        echo "Downloading CUB library";
        wget -q https://github.com/NVIDIA/cub/archive/main.zip;
        unzip -q main.zip -d include;
        mv include/cub-main/cub include;
        echo "Removing unnecessary files";
        rm main.zip;
        rm -rf include/cub-main;
        """
        subprocess.run(commands, shell=True, check=True, universal_newlines=True)
    return os.path.abspath("./include")


def cuaev_extension_kwargs(
    sms: set[str],
    debug: bool,
    opt: bool,
) -> dict[str, tp.Any]:
    print("-" * 75)
    print(f"Will build cuAEV with support for SMs: {', '.join(sms)}")

    nvcc_args = ["--expt-extended-lambda"]
    # The following defs are not necessarily required to use CUB safely, but are
    # harmless
    nvcc_args.extend(
        [
            "-DCUB_NS_QUALIFIER=::cuaev::cub",
            "-DCUB_NS_PREFIX=namespace cuaev {",
            "-DCUB_NS_POSTFIX=}",
        ]
    )
    if debug:
        nvcc_args.append("-DTORCHANI_DEBUG")
    if opt:
        nvcc_args.extend(["-DTORCHANI_OPT", "-use_fast_math"])
    nvcc_args.extend([f"-gencode=arch=compute_{sm},code=sm_{sm}" for sm in sms])
    print("NVCC compiler args:")
    for arg in nvcc_args:
        print(f"    {arg}")

    if os.name == "nt":
        cxx_args = ["/std:c++20"]
    else:
        cxx_args = ["-std=c++17"]
    print("C++ compiler args:")
    for arg in cxx_args:
        print(f"    {arg}")
    return dict(
        name="torchani.cuaev",
        sources=["torchani/csrc/cuaev.cpp", "torchani/csrc/aev.cu"],
        include_dirs=[os.path.abspath("torchani/csrc/")],
        extra_compile_args={"cxx": cxx_args, "nvcc": nvcc_args},
    )


def mnp_extension_kwargs(debug: bool) -> dict[str, tp.Any]:
    print("-" * 75)
    print("Will build MNP")

    if os.name == "nt":
        cxx_args = ["/std:c++20", "/openmp"]
    else:
        cxx_args = ["-std=c++17", "-fopenmp"]
    if debug:
        cxx_args.append("/DTORCHANI_DEBUG" if os.name == "nt" else "-DTORCHANI_DEBUG")

    print("C++ compiler args:")
    for arg in cxx_args:
        print(f"    {arg}")
    return dict(
        name="torchani.mnp",
        sources=["torchani/csrc/mnp.cpp"],
        extra_compile_args={"cxx": cxx_args},
    )


def clist_extension_kwargs(debug: bool) -> dict[str, tp.Any]:
    print("-" * 75)
    print("Will build Cell List")

    if os.name == "nt":
        cxx_args = ["/std:c++20", "/openmp"]
    else:
        cxx_args = ["-std=c++17", "-fopenmp"]
    if debug:
        cxx_args.append("/DTORCHANI_DEBUG" if os.name == "nt" else "-DTORCHANI_DEBUG")

    print("C++ compiler args:")
    for arg in cxx_args:
        print(f"    {arg}")
    return dict(
        name="torchani.cell_list",
        sources=["torchani/csrc/cell_list.cpp"],
        extra_compile_args={"cxx": cxx_args},
    )


def will_not_build_extensions_warning(torch_import_error: bool = False) -> None:
    print("-" * 75)
    if torch_import_error:
        print("Torch could not be imported")
    print(
        textwrap.dedent(
            """
            Will not install TorchANI extensions (cuAEV and MNP)
            To build the extensions with the pip frontend:
            - Make sure Torch binaries compiled with CUDA support are installed
            - Make sure a compatible CUDA Toolkit version is available
            - Add the --no-build-isolation flag to pip
            - Add --config-settings=--global-option=ext (verbatim) flag to pip
            """
        ).strip()
    )
    print("-" * 75)


TORCHANI_FLAGS = {"ext", "ext-all-sms", "ext-debug", "ext-no-opt"}
# Pascal: 60, 61
# Volta: 70 (1st gen TensorCore)
# Turing: 75 (2nd gen TensorCore)
# Ampere: 80, 86 (3d gen TensorCore)
# Ada Lovelace, Hopper: 89 (4th gen TensorCore)
# Blackwell: 100 (5th gen TensorCore)
SUPPORTED_SMS = {"60", "61", "70", "75", "80", "86", "89", "100"}
for sm in SUPPORTED_SMS:
    TORCHANI_FLAGS.add(f"ext-sm{sm}")


def strip_argv():
    argv = deepcopy(sys.argv)
    for arg in argv:
        if arg in TORCHANI_FLAGS:
            sys.argv.remove(arg)


def setup_kwargs() -> dict[str, tp.Any]:
    # setuptools executes this file:
    # - 3 times in case of build-isolation mode
    #   (egg_info, dist_info, (editable|bdist)_wheel)
    # - 2 times in case of no-build-isolation mode
    #   (dist_info, (editable|bdist)_wheel)
    #
    # Extensions may only be built when building the actual wheel,
    # In other cases executing this file is a no-op.
    if "dist_info" in sys.argv or "egg_info" in sys.argv:
        # In this case setuptools is just building the metadata
        # --global-option is passed to all stages of the build, but
        # we need the options only when building the wheel, so we strip sys.argv
        # of the options in other cases
        strip_argv()
        return dict()

    # Building the actual wheel, so attempt to import torch:
    try:
        import torch
        from torch.utils.cpp_extension import CUDAExtension
        from torch.utils.cpp_extension import BuildExtension

        TORCH_AVAILABLE = True
    except ImportError:
        TORCH_AVAILABLE = False

    if not TORCH_AVAILABLE:
        will_not_build_extensions_warning(torch_import_error=True)
        strip_argv()
        return dict()

    def collect_all_sms() -> set[str]:
        print("-" * 75)
        print("Will add all SMs torch supports")
        sms = {"60", "61", "70"}
        cuda_ver = float(torch.version.cuda) if torch.version.cuda is not None else 0
        if cuda_ver >= 10:
            sms.add("75")
        if cuda_ver >= 11:
            sms.add("80")
        if cuda_ver >= 11.1:
            sms.add("86")
            sms.add("89")
        if cuda_ver >= 12.8:
            sms.add("100")
        return sms

    def collect_compatible_sms() -> set[str]:
        print("-" * 75)
        print("Will try to find compatible CUDA devices visible to torch")
        devices = torch.cuda.device_count()
        sms: set[str] = set()
        for i in range(devices):
            sm_tuple = torch.cuda.get_device_capability(i)
            if sm_tuple >= (5, 0):
                print("Found compatible device:")
                print(f'{i}: {torch.cuda.get_device_name(f"cuda:{i}")}')
                print(f"   {torch.cuda.get_device_properties(i)}")
                sms.add(f"{sm_tuple[0]}{sm_tuple[1]}")
        if sms:
            return sms
        print("No compatible devices found")
        return collect_all_sms()

    # Flags for requesting specific SMs
    sms: set[str] = set()
    TORCHANI_BUILD_SMS = os.getenv("TORCHANI_SMS", "").split(",")
    for sm in SUPPORTED_SMS:
        if f"ext-sm{sm}" in sys.argv or sm in TORCHANI_BUILD_SMS:
            try:
                sys.argv.remove(f"ext-sm{sm}")
            except ValueError:
                pass
            sms.add(sm)
    # Flag for requesting compatible SMs detection
    if "ext" in sys.argv or os.getenv("TORCHANI_BUILD_EXT"):
        try:
            sys.argv.remove("ext")
        except ValueError:
            pass
        sms.update(collect_compatible_sms())
    # Flag for requesting all sms
    if "ext-all-sms" in sys.argv or os.getenv("TORCHANI_BUILD_ALL_SMS"):
        try:
            sys.argv.remove("ext-all-sms")
        except ValueError:
            pass
        sms.update(collect_all_sms())

    # Compile extensions with DEBUG infomation
    debug = False
    if "ext-debug" in sys.argv:
        sys.argv.remove("ext-debug")
        debug = True

    # Compile optimized extensions
    # (intrinsic math fns and -use_fast_math nvcc flag)
    opt = True
    if "ext-no-opt" in sys.argv:
        sys.argv.remove("ext-no-opt")
        opt = False

    # At least 1 SM is always added to the "sms" set if extensions need to be built
    # If nothing is added, then don't build the extensions
    if not sms:
        will_not_build_extensions_warning()
        return dict()

    cuaev_kwargs = cuaev_extension_kwargs(sms, debug, opt)
    mnp_kwargs = mnp_extension_kwargs(debug=debug)
    clist_kwargs = clist_extension_kwargs(debug=debug)

    # CUB needed to build the cuAEV, download it if not found bundled with Torch
    include_paths_kwargs: dict[str, tp.Any]
    if tuple(map(int, torch.__version__.split(".")[:2])) >= (2, 7):
        include_paths_kwargs = {"device_type": "cuda"}
    else:
        include_paths_kwargs = {"cuda": True}
    torch_include_dirs = torch.utils.cpp_extension.include_paths(**include_paths_kwargs)
    cub_include_dir = maybe_download_cub(torch_include_dirs)
    if cub_include_dir:
        cuaev_kwargs["include_dirs"].append(cub_include_dir)

    # MNP extension doesn't need nvcc to be compiled, but it still uses torch
    # CUDA libraries, so CUDAExtension is needed
    print("-" * 75)
    return {
        "ext_modules": [
            CUDAExtension(**cuaev_kwargs),
            CUDAExtension(**mnp_kwargs),
            CUDAExtension(**clist_kwargs),
        ],
        "cmdclass": {
            "build_ext": BuildExtension.with_options(
                no_python_abi_suffix=True,
            )
        },
    }


setup(**setup_kwargs())