diff --git a/setup.py b/setup.py index 71491affb988..2aadd0ffe526 100644 --- a/setup.py +++ b/setup.py @@ -1,5 +1,4 @@ -# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. -# See https://llvm.org/LICENSE.txt for license information. +#!/usr/bin/env python # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception # Also available under a BSD-style license. See LICENSE. @@ -44,57 +43,49 @@ # directory that our CMake build produces. We go through quite a bit of effort # 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 multiprocessing - from distutils.command.build import build as _build from setuptools import setup, Extension from setuptools.command.build_ext import build_ext from setuptools.command.build_py import build_py -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: +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.getenv("TORCH_MLIR_PYTHON_PACKAGE_VERSION", "0.0.1") # If true, enable LTC build by default -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_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", True ) 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_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. +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) + + 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): @@ -108,14 +99,13 @@ def cmake_build(self, cmake_build_dir): llvm_dir = str(SRC_DIR / "externals" / "llvm-project" / "llvm") cmake_config_args = [ - f"cmake", + "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", @@ -126,24 +116,24 @@ def cmake_build(self, cmake_build_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}", + str(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}", + str(llvm_dir), ] cmake_build_args = [ - f"cmake", - f"--build", - f".", - f"--config", - f"{CMAKE_BUILD_TYPE}", - f"--target", - f"TorchMLIRPythonModules", - f"--", + "cmake", + "--build", + ".", + "--config", + CMAKE_BUILD_TYPE, + "--target", + "TorchMLIRPythonModules", + "--", f"-j{MAX_JOBS}", ] try: @@ -151,7 +141,7 @@ def cmake_build(self, 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:") + print("debug by following cmake command:") sys.exit(e.returncode) finally: print(f"cmake config: {' '.join(cmake_config_args)}") @@ -178,16 +168,9 @@ def run(self): cmake_cache_file = os.path.join(cmake_build_dir, "CMakeCache.txt") if os.path.exists(cmake_cache_file): os.remove(cmake_cache_file) - # NOTE: With repeated builds for different Python versions, the - # prior version binaries will continue to accumulate. IREE uses - # a separate install step and cleans the install directory to - # keep this from happening. That is the most robust. Here we just - # delete the directory where we build native extensions to keep - # this from happening but still take advantage of most of the - # build cache. mlir_libs_dir = os.path.join(python_package_dir, "torch_mlir", "_mlir_libs") if os.path.exists(mlir_libs_dir): - print(f"Removing _mlir_mlibs dir to force rebuild: {mlir_libs_dir}") + print(f"Removing _mlir_libs dir to force rebuild: {mlir_libs_dir}") shutil.rmtree(mlir_libs_dir) else: print(f"Not removing _mlir_libs dir (does not exist): {mlir_libs_dir}") @@ -213,9 +196,6 @@ def build_extension(self, ext): with open("README.md", "r", encoding="utf-8") as fh: long_description = fh.read() - -# Requires and extension modules depend on whether building PyTorch -# extensions. INSTALL_REQUIRES = [ "numpy", "packaging", @@ -225,7 +205,6 @@ def build_extension(self, ext): ] NAME = "torch-mlir-core" -# If building PyTorch extensions, customize. if not TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS: import torch @@ -241,10 +220,9 @@ def build_extension(self, ext): ] ) - setup( name=NAME, - version=f"{PACKAGE_VERSION}", + version=PACKAGE_VERSION, author="Sean Silva", author_email="silvasean@google.com", description="First-class interop between PyTorch and MLIR", @@ -253,21 +231,9 @@ def build_extension(self, ext): include_package_data=True, cmdclass={ "build": CustomBuild, - "built_ext": NoopBuildExtension, + "build_ext": NoopBuildExtension, "build_py": CMakeBuild, }, ext_modules=EXT_MODULES, - python_requires=">=3.8", install_requires=INSTALL_REQUIRES, - extras_require={ - "onnx": [ - "onnx>=1.15.0", - ], - }, - entry_points={ - "console_scripts": [ - "torch-mlir-import-onnx = torch_mlir.tools.import_onnx:_cli_main", - ], - }, - zip_safe=False, )