From 3e8935149172bf6d60613ecdf7739d790fe9b9dc Mon Sep 17 00:00:00 2001 From: Farid Zakaria Date: Sun, 14 Jan 2024 19:44:32 +0000 Subject: [PATCH] Beginning framework to support building python wheel * Added more to gitignore for python wheel building * Removed check that you can't build wheel in isolated use of project * Updated README * Made CMAKE_BUILD_TYPE for mlir build configurable * Added setup.py --- .gitignore | 2 + CMakeLists.txt | 4 -- README.md | 25 ++++++++ build_tools/build_mlir.sh | 4 +- stablehlo/integrations/python/setup.py | 84 ++++++++++++++++++++++++++ 5 files changed, 114 insertions(+), 5 deletions(-) create mode 100644 stablehlo/integrations/python/setup.py diff --git a/.gitignore b/.gitignore index 17efa3b84ae..9cb2619d7b0 100644 --- a/.gitignore +++ b/.gitignore @@ -8,3 +8,5 @@ llvm-build .vscode .cache .idea/ +**/__pycache__ +*.whl diff --git a/CMakeLists.txt b/CMakeLists.txt index ed639ba25f6..ad9cbe92b5f 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -135,10 +135,6 @@ add_definitions(${LLVM_DEFINITIONS}) #------------------------------------------------------------------------------- if(STABLEHLO_ENABLE_BINDINGS_PYTHON) - if(NOT STABLEHLO_EXTERNAL_PROJECT_BUILD) - message(WARNING "StableHLO Python bindings are not supported in standalone mode") - endif() - include(MLIRDetectPythonEnv) mlir_configure_python_dev_packages() endif() diff --git a/README.md b/README.md index a94813dc283..6615daa7532 100644 --- a/README.md +++ b/README.md @@ -104,6 +104,31 @@ Here's how to build the StableHLO repo on Linux or macOS: This runs all the tests in `stablehlo/tests/`. +## Python + +If you'd like to build the Python bindings, you'll need to install a few +additional dependencies. + +```sh +sudo apt -y install python3-pybind11 +``` + +If you've built MLIR using the script above, the Python bindings for MLIR are +already built. + +After you have built the project you can import the Python bindings to begin +by modifying your Python path variable + +```shell +$ PYTHONPATH="./build/python_packages/stablehlo" python3 +Python 3.11.6 (main, Oct 8 2023, 05:06:43) [GCC 13.2.0] on linux +Type "help", "copyright", "credits" or "license" for more information. +>>> import mlir.dialects.stablehlo +>>> import mlir.dialects.stablehlo +>>> from mlir.ir import Context, Location +>>> import mlir.dialects.arith +``` + ## Community Building an amazing portability layer between ML frameworks and ML compilers diff --git a/build_tools/build_mlir.sh b/build_tools/build_mlir.sh index b5481939276..c6ddab9a139 100755 --- a/build_tools/build_mlir.sh +++ b/build_tools/build_mlir.sh @@ -23,6 +23,7 @@ fi # LLVM source LLVM_SRC_DIR="$1" build_dir="$2" +MLIR_CMAKE_BUILD_TYPE="${CMAKE_BUILD_TYPE:-RelWithDebInfo}" if ! [ -f "$LLVM_SRC_DIR/llvm/CMakeLists.txt" ]; then echo "Expected the path to LLVM to be set correctly (got '$LLVM_SRC_DIR'): can't find CMakeLists.txt" @@ -47,9 +48,10 @@ cmake -GNinja \ -DLLVM_TARGETS_TO_BUILD=host \ -DLLVM_INCLUDE_TOOLS=ON \ -DLLVM_ENABLE_BINDINGS=OFF \ + -DMLIR_ENABLE_BINDINGS_PYTHON=ON \ -DLLVM_BUILD_TOOLS=OFF \ -DLLVM_INCLUDE_TESTS=OFF \ - -DCMAKE_BUILD_TYPE=RelWithDebInfo \ + -DCMAKE_BUILD_TYPE="$MLIR_CMAKE_BUILD_TYPE" \ -DLLVM_ENABLE_ASSERTIONS=On cmake --build "$build_dir" --target all diff --git a/stablehlo/integrations/python/setup.py b/stablehlo/integrations/python/setup.py new file mode 100644 index 00000000000..a8b26926c10 --- /dev/null +++ b/stablehlo/integrations/python/setup.py @@ -0,0 +1,84 @@ +import os +import shutil +from typing import Any +from distutils.command.clean import clean +from setuptools import Extension +from setuptools import find_namespace_packages +from setuptools import setup +from setuptools.command.build_ext import build_ext +import pathlib +import subprocess +from glob import glob + +# Assuming your shared object files are in the '_mlir_libs' directory +# and that the CMake directory you've built into is called 'build' +# TODO(fzakaria): Consider using the CMake extension module +# similar to https://github.com/makslevental/mlir-wheels/blob/2a74f388d987bb62f832660a975b51168e30d04f/setup.py#L15 +lib_dir = os.path.normpath( + '../../../build/python_packages/stablehlo/mlir/_mlir_libs') + + +class CMakeExtension(Extension): + def __init__(self, name): + # don't invoke the original build_ext for this special extension + super().__init__(name, sources=[]) + + +class CMakeBuild(build_ext): + """This is a fake build_ext that doesn't do any building. It expects that the + shared object files have already been built externally and are available at + the lib_dir variable above. It then copies them into the extension directory.""" + def build_extension(self, ext: Any) -> None: + # copy _ml_ir_libs into extension directory + # self.get_ext_fullpath('xxx') gives you something like: + # build/lib.linux-x86_64-cpython-311/xxx.cpython-311-x86_64-linux-gnu.so + # we then take the parent directory and copy the contents of the + # real _mlir_libs directory into the parent. + # we make sure to also recreate the mlir/_mlir_libs directory + ext_dir = pathlib.Path(self.get_ext_fullpath('ignored')) + target_dir = ext_dir.parent / 'mlir' / '_mlir_libs' + shutil.copytree(lib_dir, target_dir, dirs_exist_ok=True) + + +class CleanCommand(clean): + """ + Custom implementation of ``clean`` setuptools command. + """ + + def run(self): + """After calling the super class implementation, this function removes + the dist directory if it exists and any egg files.""" + self.all = True # --all by default when cleaning + super().run() + shutil.rmtree("dist", ignore_errors=True) + for egg in glob('*.egg-info'): + shutil.rmtree(egg, ignore_errors=True) + +def get_version(): + # get the latest tag without the leading v + latest_tag = subprocess.check_output(["git", "describe", "--tags", "--abbrev=0"], text=True).strip('v').strip() + latest_commit = subprocess.check_output(["git", "rev-parse", "--short", "HEAD"], text=True).strip() + return f"{latest_tag}+{latest_commit}" + +# TODO(fzakaria): The distribution (wheel) of this package is not manylinux +# conformant. Consider also running auditwheel similar to +# https://github.com/makslevental/mlir-wheels to make it a smoother installation +# experience. +setup( + name='stablehlo', + packages=find_namespace_packages( + os.path.normpath("../../build/python_packages/stablehlo")), + package_dir={ + "": os.path.normpath("../../../build/python_packages/stablehlo")}, + + # Define extensions if your package needs to compile anything + ext_modules=[CMakeExtension(name="_mlir_libs")], + cmdclass={"build_ext": CMakeBuild, "clean": CleanCommand}, + + author='Your Name', + author_email='your.email@example.com', + description='Backward compatible ML compute opset inspired by HLO/MHLO', + url='https://github.com/openxla/stablehlo', + # TODO(fzakaria): Figure out how to get version same as code; os.environ ? + version = get_version() +)