Setuptools uses a separate build directory (llvm#3023)

* setuptools not steal the build directory name
llvm#3021 (comment)
* support pre-built LLVM
* support CMAKE_BUILD_TYPE env

.gitignore

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

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 @@ def run(self):
                 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)