Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Update setup.py #3159

Open
wants to merge 6 commits into
base: main
Choose a base branch
from
Open

Update setup.py #3159

Changes from 3 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
e8df30d
Update setup.py
r0cketdyne Apr 13, 2024
a7a063a
Update setup.py
r0cketdyne Aug 16, 2024
082da93
Merge branch 'main' into patch-2
r0cketdyne Aug 22, 2024
997d39d
Update setup.py
r0cketdyne Sep 12, 2024
b9d3c91
Merge branch 'main' into patch-2
r0cketdyne Sep 17, 2024
1992b31
Update setup.py
r0cketdyne Sep 18, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
112 changes: 35 additions & 77 deletions setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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.

Expand Down Expand Up @@ -44,57 +43,45 @@
# 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_ONLY_MLIR_PYTHON_BINDINGS", False
)
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_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):
Expand All @@ -108,14 +95,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",
Expand All @@ -126,32 +112,32 @@ 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:
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:")
print("debug by following cmake command:")
sys.exit(e.returncode)
finally:
print(f"cmake config: {' '.join(cmake_config_args)}")
Expand All @@ -178,16 +164,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}")
Expand All @@ -213,9 +192,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",
Expand All @@ -225,26 +201,20 @@ 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

NAME = "torch-mlir"
INSTALL_REQUIRES.extend(
[
f"torch=={torch.__version__}".split("+", 1)[0],
]
)
EXT_MODULES.extend(
[
CMakeExtension("torch_mlir._mlir_libs._jit_ir_importer"),
]
)

INSTALL_REQUIRES.extend([
f"torch=={torch.__version__}".split("+", 1)[0],
])
EXT_MODULES.extend([
CMakeExtension("torch_mlir._mlir_libs._jit_ir_importer"),
])

setup(
name=NAME,
version=f"{PACKAGE_VERSION}",
version=PACKAGE_VERSION,
author="Sean Silva",
author_email="[email protected]",
description="First-class interop between PyTorch and MLIR",
Expand All @@ -253,21 +223,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,
)
Loading