llvm
/
torch-mlir
mirror of https://github.com/llvm/torch-mlir
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

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
pull/3026/head
penguin_wwy 2024-03-14 11:41:48 +08:00 committed by GitHub
parent 870e63bc3c
commit 29ac23a790
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 87 additions and 38 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
.gitignore vendored
View File

@ -11,6 +11,7 @@ externals/pytorch/
libtorch*
/build/
/setup_build/
__pycache__
*.pyc

112
setup.py
View File

@ -30,60 +30,64 @@
# 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 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 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")
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())
enable_ltc = int(os.environ.get('TORCH_MLIR_ENABLE_LTC', TORCH_MLIR_ENABLE_LTC_DEFAULT))
cmake_args = [
f"-DCMAKE_BUILD_TYPE=Release",
cmake_config_args = [
f"cmake",
f"-DCMAKE_BUILD_TYPE={CMAKE_BUILD_TYPE}",
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_TARGETS_TO_BUILD=host",
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",
@ -91,7 +95,58 @@ class CMakeBuild(build_py):
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"))
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"):
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)