From 9ac90ec7b2710dda9a30450d2c0ee308b31ee4e0 Mon Sep 17 00:00:00 2001 From: penguin_wwy <940375606@qq.com> Date: Mon, 15 Apr 2024 01:40:25 +0800 Subject: [PATCH] Refactor the parameters and usage instructions of the setup script. (#3162) As this issuecomment(https://github.com/llvm/torch-mlir/pull/3021#issuecomment-2031248199) suggests, `setup.py` should only be used for building Python packages, so: * disabled the develop command * refactor the environment variable parameters * add more doc for the usage and env var of setup.py --- docs/development.md | 9 +++++ setup.py | 96 ++++++++++++++++++++++++++++----------------- 2 files changed, 69 insertions(+), 36 deletions(-) diff --git a/docs/development.md b/docs/development.md index 27ff8b7bf..5e9a3f6b8 100644 --- a/docs/development.md +++ b/docs/development.md @@ -349,6 +349,15 @@ python -m pip install -r requirements.txt CMAKE_GENERATOR=Ninja python setup.py bdist_wheel ``` +To package a completed CMake build directory, +you can use the `TORCH_MLIR_CMAKE_BUILD_DIR` and `TORCH_MLIR_CMAKE_ALREADY_BUILT` environment variables: +```shell +TORCH_MLIR_CMAKE_BUILD_DIR=build/ TORCH_MLIR_CMAKE_ALREADY_BUILT=1 python setup.py bdist_wheel +``` + +Note: The setup.py script is only used for building the Python packages, +not support commands like `setup.py develop` to build the development environment. + # Testing Torch-MLIR has two types of tests: diff --git a/setup.py b/setup.py index 3cd5f2eca..d5e3d055c 100644 --- a/setup.py +++ b/setup.py @@ -7,6 +7,31 @@ # ``` # $ python setup.py bdist_wheel # ``` +# Environment variables you are probably interested in: +# +# TORCH_MLIR_PYTHON_PACKAGE_VERSION: +# specify the version of torch-mlir, for example, this can be "20220330.357" +# for a snapshot release on 2022-03-30 with build number 357. +# +# TORCH_MLIR_ENABLE_LTC: +# enables the Lazy Tensor Core Backend +# +# LLVM_INSTALL_DIR: +# build the project *out-of-tree* using the built llvm-project +# +# CMAKE_BUILD_TYPE: +# specify the build type: DEBUG/RelWithDebInfo/Release +# +# TORCH_MLIR_CMAKE_BUILD_DIR: +# specify the cmake build directory +# +# TORCH_MLIR_CMAKE_ALREADY_BUILT: +# the `TORCH_MLIR_CMAKE_BUILD_DIR` directory has already been compiled, +# and the CMake compilation process will not be executed again. +# On CIs, it is often advantageous to re-use/control the CMake build directory. +# +# MAX_JOBS: +# maximum number of compile jobs we should use to compile your code # # It is recommended to build with Ninja and ccache. To do so, set environment # variables by prefixing to above invocations: @@ -14,16 +39,6 @@ # CMAKE_GENERATOR=Ninja CMAKE_C_COMPILER_LAUNCHER=ccache CMAKE_CXX_COMPILER_LAUNCHER=ccache # ``` # -# On CIs, it is often advantageous to re-use/control the CMake build directory. -# This can be set with the TORCH_MLIR_CMAKE_BUILD_DIR env var. -# Additionally, the TORCH_MLIR_CMAKE_BUILD_DIR_ALREADY_BUILT env var will -# prevent this script from attempting to build the directory, and will simply -# use the (presumed already built) directory as-is. -# -# The package version can be set with the TORCH_MLIR_PYTHON_PACKAGE_VERSION -# environment variable. For example, this can be "20220330.357" for a snapshot -# release on 2022-03-30 with build number 357. -# # Implementation notes: # The contents of the wheel is just the contents of the `python_packages` # directory that our CMake build produces. We go through quite a bit of effort @@ -42,20 +57,33 @@ 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: +if "develop" in sys.argv: + print("Warning: The setup.py script is only used for building the wheel package.") + print("For initializing the development environment," + "please use the cmake commands introduced in the docs/development.md.") + sys.exit(1) + +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" +PACKAGE_VERSION = os.getenv("TORCH_MLIR_PYTHON_PACKAGE_VERSION", "0.0.1") # If true, enable LTC build by default -TORCH_MLIR_ENABLE_LTC_DEFAULT = True -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) +TORCH_MLIR_ENABLE_LTC = _check_env_flag("TORCH_MLIR_ENABLE_LTC", True) +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") +TORCH_MLIR_CMAKE_ALREADY_BUILT = _check_env_flag( + "TORCH_MLIR_CMAKE_ALREADY_BUILT", False +) +TORCH_MLIR_CMAKE_BUILD_DIR = os.getenv("TORCH_MLIR_CMAKE_BUILD_DIR") +MAX_JOBS = os.getenv("MAX_JOBS", str(multiprocessing.cpu_count())) + # Build phase discovery is unreliable. Just tell it what phases to run. class CustomBuild(_build): @@ -77,8 +105,6 @@ 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", @@ -92,7 +118,7 @@ class CMakeBuild(build_py): 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_LTC={'ON' if TORCH_MLIR_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: @@ -117,7 +143,7 @@ class CMakeBuild(build_py): f"--target", f"TorchMLIRPythonModules", f"--", - f"-j{max_jobs}" + f"-j{MAX_JOBS}", ] try: subprocess.check_call(cmake_config_args, cwd=cmake_build_dir) @@ -131,22 +157,22 @@ class CMakeBuild(build_py): 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") + cmake_build_dir = TORCH_MLIR_CMAKE_BUILD_DIR if not cmake_build_dir: cmake_build_dir = os.path.abspath( - os.path.join(target_dir, "..", "cmake_build")) + os.path.join(target_dir, "..", "cmake_build") + ) if LLVM_INSTALL_DIR: - python_package_dir = os.path.join(cmake_build_dir, - "python_packages", - "torch_mlir") + 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"): + python_package_dir = os.path.join( + cmake_build_dir, "tools", "torch-mlir", "python_packages", "torch_mlir" + ) + if not TORCH_MLIR_CMAKE_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): @@ -169,16 +195,14 @@ class CMakeBuild(build_py): if os.path.exists(target_dir): shutil.rmtree(target_dir, ignore_errors=False, onerror=None) - shutil.copytree(python_package_dir, - target_dir, - symlinks=False) + shutil.copytree(python_package_dir, target_dir, symlinks=False) class CMakeExtension(Extension): - def __init__(self, name, sourcedir=""): - Extension.__init__(self, name, sources=[]) - self.sourcedir = os.path.abspath(sourcedir) + def __init__(self, name, sourcedir=""): + Extension.__init__(self, name, sources=[]) + self.sourcedir = os.path.abspath(sourcedir) class NoopBuildExtension(build_ext):