forked from llvm/torch-mlir
-
Notifications
You must be signed in to change notification settings - Fork 1
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Will add a CI job that builds and uploads snapshot packages next.
- Loading branch information
Showing
3 changed files
with
139 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -13,3 +13,4 @@ __pycache__ | |
# Pip artifacts. | ||
*.egg-info | ||
*.whl | ||
/wheelhouse |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,36 @@ | ||
#!/bin/bash | ||
set -e | ||
|
||
if [ -z "$PYTHON" ]; then | ||
PYTHON="$(which python)" | ||
fi | ||
version="$("$PYTHON" --version)" | ||
echo "Using python: $PYTHON (version $version)" | ||
|
||
repo_root="$(cd $(dirname $0)/.. && pwd)" | ||
wheelhouse="$repo_root/wheelhouse" | ||
package_test_venv="$wheelhouse/package-test.venv" | ||
mkdir -p $wheelhouse | ||
cd $wheelhouse | ||
|
||
echo "---- BUILDING torch-mlir ----" | ||
CMAKE_GENERATOR=Ninja CMAKE_C_COMPILER_LAUNCHER=ccache CMAKE_CXX_COMPILER_LAUNCHER=ccache \ | ||
$PYTHON "${repo_root}/setup.py" bdist_wheel \ | ||
--dist-dir "$wheelhouse" -v | ||
|
||
|
||
# Smoke test: create a venv, install the package, and run an example. | ||
|
||
echo "---- CREATING VENV ----" | ||
python -m venv "$package_test_venv" | ||
VENV_PYTHON="$package_test_venv/bin/python" | ||
|
||
echo "---- INSTALLING torch ----" | ||
$VENV_PYTHON -m pip install --pre torch torchvision pybind11 -f "https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html" | ||
echo "---- INSTALLING other deps for smoke test script ----" | ||
$VENV_PYTHON -m pip install requests pillow | ||
echo "---- INSTALLING torch-mlir ----" | ||
$VENV_PYTHON -m pip install -f "$wheelhouse" --force-reinstall torch-mlir | ||
|
||
echo "---- RUNNING SMOKE TEST ----" | ||
$VENV_PYTHON "$repo_root/examples/torchscript_resnet18_e2e.py" |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,102 @@ | ||
# 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. | ||
# | ||
# 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. | ||
import os | ||
import shutil | ||
import subprocess | ||
import sys | ||
|
||
from distutils.command.build import build as _build | ||
from setuptools import setup | ||
from setuptools.command.build_ext import build_ext | ||
from setuptools.command.build_py import build_py | ||
|
||
|
||
# Build phase discovery is unreliable. Just tell it what phases to run. | ||
class CustomBuild(_build): | ||
|
||
def run(self): | ||
self.run_command("build_py") | ||
self.run_command("build_ext") | ||
self.run_command("build_scripts") | ||
|
||
|
||
class CMakeBuild(build_py): | ||
|
||
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.join(target_dir, "..", "cmake_build") | ||
src_dir = os.path.abspath(os.path.dirname(__file__)) | ||
llvm_dir = os.path.join(src_dir, "external", "llvm-project", "llvm") | ||
cmake_args = [ | ||
f"-DCMAKE_BUILD_TYPE=Release", | ||
f"-DPython3_EXECUTABLE={sys.executable}", | ||
f"-DLLVM_TARGETS_TO_BUILD=host", | ||
f"-DMLIR_ENABLE_BINDINGS_PYTHON=ON", | ||
f"-DLLVM_ENABLE_PROJECTS=mlir", | ||
f"-DLLVM_EXTERNAL_PROJECTS=torch-mlir", | ||
f"-DLLVM_EXTERNAL_TORCH_MLIR_SOURCE_DIR={src_dir}", | ||
] | ||
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): | ||
os.remove(cmake_cache_file) | ||
subprocess.check_call(["cmake", llvm_dir] + | ||
cmake_args, cwd=cmake_build_dir) | ||
subprocess.check_call(["cmake", | ||
"--build", ".", | ||
"--target", "TorchMLIRPythonModules"], | ||
cwd=cmake_build_dir) | ||
python_package_dir = os.path.join(cmake_build_dir, | ||
"tools", "torch-mlir", "python_packages", | ||
"torch_mlir") | ||
shutil.copytree(python_package_dir, | ||
target_dir, | ||
symlinks=False, | ||
dirs_exist_ok=True) | ||
|
||
|
||
class NoopBuildExtension(build_ext): | ||
|
||
def build_extension(self, ext): | ||
pass | ||
|
||
|
||
setup( | ||
name="torch-mlir", | ||
version="0.0.1", | ||
author="Sean Silva", | ||
author_email="silvasean@google.com", | ||
description="First-class interop between PyTorch and MLIR", | ||
long_description="", | ||
include_package_data=True, | ||
cmdclass={ | ||
"build": CustomBuild, | ||
"built_ext": NoopBuildExtension, | ||
"build_py": CMakeBuild, | ||
}, | ||
zip_safe=False, | ||
) |