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

@@ -8,3 +8,5 @@ llvm-build
 .vscode
 .cache
 .idea/
+**/__pycache__
+*.whl

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()

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

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

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()
+)