Bring back Python packaging.

Will add a CI job that builds and uploads snapshot packages next.

.gitignore

@@ -13,3 +13,4 @@ __pycache__
 # Pip artifacts.
 *.egg-info
 *.whl
+/wheelhouse

build_tools/build_python_wheels.sh

@@ -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"

setup.py

@@ -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,
+)