From 29ac23a7903cd6e74f113020f1c037c2c27f8fac Mon Sep 17 00:00:00 2001 From: penguin_wwy <940375606@qq.com> Date: Thu, 14 Mar 2024 11:41:48 +0800 Subject: [PATCH] Setuptools uses a separate build directory (#3023) * setuptools not steal the build directory name https://github.com/llvm/torch-mlir/pull/3021#issuecomment-1994447855 * support pre-built LLVM * support CMAKE_BUILD_TYPE env --- .gitignore | 1 + setup.py | 124 +++++++++++++++++++++++++++++++++++++---------------- 2 files changed, 87 insertions(+), 38 deletions(-) diff --git a/.gitignore b/.gitignore index 5c4074289..00a5bc96f 100644 --- a/.gitignore +++ b/.gitignore @@ -11,6 +11,7 @@ externals/pytorch/ libtorch* /build/ +/setup_build/ __pycache__ *.pyc diff --git a/setup.py b/setup.py index 77c8b2ad0..4863a9807 100644 --- a/setup.py +++ b/setup.py @@ -30,68 +30,123 @@ # on the CMake side to organize that directory already, so we avoid duplicating # that here, and just package up its contents. import os +import pathlib import shutil import subprocess import sys -import sysconfig +import multiprocessing from distutils.command.build import build as _build -from distutils.sysconfig import get_python_inc from setuptools import setup, Extension from setuptools.command.build_ext import build_ext from setuptools.command.build_py import build_py +def check_env_flag(name: str, default=None) -> bool: + return str(os.getenv(name, default)).upper() in ["ON", "1", "YES", "TRUE", "Y"] + + PACKAGE_VERSION = os.environ.get("TORCH_MLIR_PYTHON_PACKAGE_VERSION") or "0.0.1" # If true, enable LTC build by default TORCH_MLIR_ENABLE_LTC_DEFAULT = True -TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS = int(os.environ.get('TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS', False)) +TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS = check_env_flag( + 'TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS', False) +LLVM_INSTALL_DIR = os.getenv('LLVM_INSTALL_DIR', None) +SRC_DIR = pathlib.Path(__file__).parent.absolute() +CMAKE_BUILD_TYPE = os.getenv("CMAKE_BUILD_TYPE", "Release") + # Build phase discovery is unreliable. Just tell it what phases to run. class CustomBuild(_build): + def initialize_options(self): + _build.initialize_options(self) + # Make setuptools not steal the build directory name, + # because the mlir c++ developers are quite + # used to having build/ be for cmake + self.build_base = "setup_build" + def run(self): self.run_command("build_py") self.run_command("build_ext") self.run_command("build_scripts") + class CMakeBuild(build_py): + def cmake_build(self, cmake_build_dir): + llvm_dir = str(SRC_DIR / "externals" / "llvm-project" / "llvm") + enable_ltc = check_env_flag('TORCH_MLIR_ENABLE_LTC', TORCH_MLIR_ENABLE_LTC_DEFAULT) + max_jobs = os.getenv("MAX_JOBS") or str(multiprocessing.cpu_count()) + + cmake_config_args = [ + f"cmake", + f"-DCMAKE_BUILD_TYPE={CMAKE_BUILD_TYPE}", + f"-DPython3_EXECUTABLE={sys.executable}", + f"-DPython3_FIND_VIRTUALENV=ONLY", + f"-DMLIR_ENABLE_BINDINGS_PYTHON=ON", + f"-DLLVM_TARGETS_TO_BUILD=host", + f"-DLLVM_ENABLE_ZSTD=OFF", + # Optimization options for building wheels. + f"-DCMAKE_VISIBILITY_INLINES_HIDDEN=ON", + f"-DCMAKE_C_VISIBILITY_PRESET=hidden", + f"-DCMAKE_CXX_VISIBILITY_PRESET=hidden", + f"-DTORCH_MLIR_ENABLE_LTC={'ON' if enable_ltc else 'OFF'}", + f"-DTORCH_MLIR_ENABLE_PYTORCH_EXTENSIONS={'OFF' if TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS else 'ON'}", + ] + if LLVM_INSTALL_DIR: + cmake_config_args += [ + f"-DMLIR_DIR='{LLVM_INSTALL_DIR}/lib/cmake/mlir/'", + f"-DLLVM_DIR='{LLVM_INSTALL_DIR}/lib/cmake/llvm/'", + f"{SRC_DIR}", + ] + else: + cmake_config_args += [ + f"-DLLVM_ENABLE_PROJECTS=mlir", + f"-DLLVM_EXTERNAL_PROJECTS='torch-mlir'", + f"-DLLVM_EXTERNAL_TORCH_MLIR_SOURCE_DIR={SRC_DIR}", + f"{llvm_dir}", + ] + cmake_build_args = [ + f"cmake", + f"--build", + f".", + f"--config", + f"{CMAKE_BUILD_TYPE}", + f"--target", + f"TorchMLIRPythonModules", + f"--", + f"-j{max_jobs}" + ] + try: + subprocess.check_call(cmake_config_args, cwd=cmake_build_dir) + subprocess.check_call(cmake_build_args, cwd=cmake_build_dir) + except subprocess.CalledProcessError as e: + print("cmake build failed with\n", e) + print("debug by follow cmake command:") + sys.exit(e.returncode) + finally: + print(f"cmake config: {' '.join(cmake_config_args)}") + print(f"cmake build: {' '.join(cmake_build_args)}") + print(f"cmake workspace: {cmake_build_dir}") + + def run(self): target_dir = self.build_lib cmake_build_dir = os.getenv("TORCH_MLIR_CMAKE_BUILD_DIR") if not cmake_build_dir: cmake_build_dir = os.path.abspath( os.path.join(target_dir, "..", "cmake_build")) - python_package_dir = os.path.join(cmake_build_dir, - "tools", "torch-mlir", "python_packages", - "torch_mlir") + if LLVM_INSTALL_DIR: + python_package_dir = os.path.join(cmake_build_dir, + "python_packages", + "torch_mlir") + else: + python_package_dir = os.path.join(cmake_build_dir, + "tools", "torch-mlir", "python_packages", + "torch_mlir") if not os.getenv("TORCH_MLIR_CMAKE_BUILD_DIR_ALREADY_BUILT"): - src_dir = os.path.abspath(os.path.dirname(__file__)) - llvm_dir = os.path.join( - src_dir, "externals", "llvm-project", "llvm") - - enable_ltc = int(os.environ.get('TORCH_MLIR_ENABLE_LTC', TORCH_MLIR_ENABLE_LTC_DEFAULT)) - - cmake_args = [ - f"-DCMAKE_BUILD_TYPE=Release", - f"-DPython3_EXECUTABLE={sys.executable}", - f"-DPython3_FIND_VIRTUALENV=ONLY", - f"-DLLVM_TARGETS_TO_BUILD=host", - f"-DMLIR_ENABLE_BINDINGS_PYTHON=ON", - f"-DLLVM_ENABLE_PROJECTS=mlir", - f"-DLLVM_ENABLE_ZSTD=OFF", - f"-DLLVM_EXTERNAL_PROJECTS=torch-mlir", - f"-DLLVM_EXTERNAL_TORCH_MLIR_SOURCE_DIR={src_dir}", - # Optimization options for building wheels. - f"-DCMAKE_VISIBILITY_INLINES_HIDDEN=ON", - f"-DCMAKE_C_VISIBILITY_PRESET=hidden", - f"-DCMAKE_CXX_VISIBILITY_PRESET=hidden", - f"-DTORCH_MLIR_ENABLE_LTC={'ON' if enable_ltc else 'OFF'}", - f"-DTORCH_MLIR_ENABLE_PYTORCH_EXTENSIONS={'OFF' if TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS else 'ON'}", - ] - os.makedirs(cmake_build_dir, exist_ok=True) cmake_cache_file = os.path.join(cmake_build_dir, "CMakeCache.txt") if os.path.exists(cmake_cache_file): @@ -109,14 +164,7 @@ class CMakeBuild(build_py): shutil.rmtree(mlir_libs_dir) else: print(f"Not removing _mlir_libs dir (does not exist): {mlir_libs_dir}") - - subprocess.check_call(["cmake", llvm_dir] + - cmake_args, cwd=cmake_build_dir) - subprocess.check_call(["cmake", - "--build", ".", - "--config", "Release", - "--target", "TorchMLIRPythonModules"], - cwd=cmake_build_dir) + self.cmake_build(cmake_build_dir) if os.path.exists(target_dir): shutil.rmtree(target_dir, ignore_errors=False, onerror=None)