pytorch.py

# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.

"""PyTorch related extensions."""
import os
from pathlib import Path

import setuptools

from .utils import all_files_in_dir, cuda_version, get_cuda_include_dirs, debug_build_enabled, rocm_build, hipify


def setup_pytorch_extension(
    csrc_source_files,
    csrc_header_files,
    common_header_files,
) -> setuptools.Extension:
    """Setup CUDA extension for PyTorch support"""

    # Source files
    sources = all_files_in_dir(Path(csrc_source_files), name_extension="cpp")

    # Header files
    if rocm_build():
        include_dirs = [
            common_header_files,
            common_header_files / "common",
            common_header_files / "common" / "include",
            csrc_header_files,
        ]
    else:
        include_dirs = get_cuda_include_dirs()
        include_dirs.extend(
            [
                common_header_files,
                common_header_files / "common",
                common_header_files / "common" / "include",
                csrc_header_files,
            ]
        )

    if rocm_build():
        current_file_path = Path(__file__).parent.resolve()
        base_dir = current_file_path.parent
        sources = hipify(base_dir, csrc_source_files, sources, include_dirs)

    # Compiler flags
    cxx_flags = [
        "-O3",
        "-fvisibility=hidden",
    ]
    if rocm_build():
        nvcc_flags = [
            "-O3",
            "-U__HIP_NO_HALF_OPERATORS__",
            "-U__HIP_NO_HALF_CONVERSIONS__",
            "-U__HIP_NO_BFLOAT16_OPERATORS__",
            "-U__HIP_NO_BFLOAT16_CONVERSIONS__",
            "-U__HIP_NO_BFLOAT162_OPERATORS__",
            "-U__HIP_NO_BFLOAT162_CONVERSIONS__",
            "-DUSE_ROCM",
        ]
        if bool(int(os.getenv("NVTE_BUILD_SUPPRESS_UNUSED_WARNING", "1"))):
            nvcc_flags.extend(
                [
                    "-Wno-unused-result",
                    "-Wno-unused-function",
                    "-Wno-unused-private-field",
                    "-Wno-unused-variable",
                ]
            )
            cxx_flags.extend(
                [
                    "-Wno-unused-result",
                    "-Wno-unused-function",
                    "-Wno-unused-private-field",
                    "-Wno-unused-variable",
                ]
            )

        if bool(int(os.getenv("NVTE_BUILD_SUPPRESS_RETURN_TYPE_WARNING", "1"))):
            nvcc_flags.append("-Wno-return-type")
            cxx_flags.append("-Wno-return-type")

        if bool(int(os.getenv("NVTE_BUILD_SUPPRESS_SIGN_COMPARE", "1"))):
            nvcc_flags.append("-Wno-sign-compare")
            cxx_flags.append("-Wno-sign-compare")
            
    else:
        pass

    if debug_build_enabled():
        cxx_flags.append("-g")
        cxx_flags.append("-UNDEBUG")
        if rocm_build():
            nvcc_flags.append("-g")
            nvcc_flags.append("-UNDEBUG")

    # Version-dependent CUDA options
    if rocm_build():
        ##TODO: Figure out which hipcc version starts to support this parallel compilation
        nvcc_flags.extend(["-parallel-jobs=4"])
    else:
        try:
            version = cuda_version()
        except FileNotFoundError:
            print("Could not determine CUDA version")
        else:
            if version < (12, 0):
                raise RuntimeError("Transformer Engine requires CUDA 12.0 or newer")

    # Libraries
    library_dirs = []
    libraries = []
    if bool(int(os.getenv("NVTE_UB_WITH_MPI", "0"))):
        assert (
            os.getenv("MPI_HOME") is not None
        ), "MPI_HOME=/path/to/mpi must be set when compiling with NVTE_UB_WITH_MPI=1!"
        mpi_path = Path(os.getenv("MPI_HOME"))
        include_dirs.append(mpi_path / "include")
        cxx_flags.append("-DNVTE_UB_WITH_MPI")
        if rocm_build():
            nvcc_flags.append("-DNVTE_UB_WITH_MPI")
            library_dirs.append(mpi_path / "lib")
            libraries.append("mpi")

    if bool(int(os.getenv("NVTE_ENABLE_NVSHMEM", 0))):
        assert (
            os.getenv("NVSHMEM_HOME") is not None
        ), "NVSHMEM_HOME must be set when compiling with NVTE_ENABLE_NVSHMEM=1"
        nvshmem_home = Path(os.getenv("NVSHMEM_HOME"))
        include_dirs.append(nvshmem_home / "include")
        library_dirs.append(nvshmem_home / "lib")
        libraries.append("nvshmem_host")
        cxx_flags.append("-DNVTE_ENABLE_NVSHMEM")
        if rocm_build():
            nvcc_flags.append("-DNVTE_ENABLE_NVSHMEM")

    # Construct PyTorch CUDA extension
    sources = [str(path) for path in sources]
    include_dirs = [str(path) for path in include_dirs]
    from torch.utils.cpp_extension import CppExtension, CUDAExtension

    if rocm_build():
        return CUDAExtension(
            name="transformer_engine_torch",
            sources=[str(src) for src in sources],
            include_dirs=[str(inc) for inc in include_dirs],
            extra_compile_args={
                "cxx": cxx_flags,
                "nvcc": nvcc_flags,
            },
            libraries=[str(lib) for lib in libraries],
            library_dirs=[str(lib_dir) for lib_dir in library_dirs],
        )
    else:
        return CppExtension(
            name="transformer_engine_torch",
            sources=[str(src) for src in sources],
            include_dirs=[str(inc) for inc in include_dirs],
            extra_compile_args={"cxx": cxx_flags},
            libraries=[str(lib) for lib in libraries],
            library_dirs=[str(lib_dir) for lib_dir in library_dirs],
        )