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 1 commit
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
96 changes: 10 additions & 86 deletions setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,84 +1,41 @@
# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# Also available under a BSD-style license. See LICENSE.

# Script for generating the torch-mlir wheel.
# ```
# $ python setup.py bdist_wheel
# ```
#
# It is recommended to build with Ninja and ccache. To do so, set environment
# variables by prefixing to above invocations:
# ```
# 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
# on the CMake side to organize that directory already, so we avoid duplicating
# that here, and just package up its contents.
#!/usr/bin/env python

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


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
PACKAGE_VERSION = os.environ.get("TORCH_MLIR_PYTHON_PACKAGE_VERSION", "0.0.1")
TORCH_MLIR_ENABLE_LTC_DEFAULT = True
TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS = check_env_flag(
'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())
max_jobs = os.getenv("MAX_JOBS", str(multiprocessing.cpu_count()))

cmake_config_args = [
f"cmake",
Expand All @@ -88,7 +45,6 @@ def cmake_build(self, cmake_build_dir):
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 Down Expand Up @@ -131,7 +87,6 @@ def cmake_build(self, cmake_build_dir):
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")
Expand All @@ -151,13 +106,6 @@ 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}")
Expand All @@ -173,26 +121,18 @@ def run(self):
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):

def build_extension(self, ext):
pass


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 @@ -202,7 +142,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
NAME = "torch-mlir"
Expand All @@ -213,10 +152,9 @@ def build_extension(self, ext):
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 @@ -228,18 +166,4 @@ def build_extension(self, ext):
"built_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,
)
ext
Loading