Introduce capability to build and archive commit level Python wheels (o…

…penxla#1924)

This PR introduces a new GitHub action which triggers on commits to main
(or manual execution) that build Pythonn (3.10 & 3.11) wheels
(distribution format) and stores the wheels under a GitHub releases.

```shell'
pip install stablehlo -f https://github.com/openxla/stablehlo/releases/expanded_assets/dev-wheels
``` 

The 'trick' of using GitHub releases as a PyPI storage mechanism is
outlined in [this blog
post](https://fzakaria.com/2024/01/15/abusing-github-as-a-pypi-server.html).

At a high level the changes are:
* Introduce a `setup.py` that covers the necessary files to include in
the wheel file. The setup.py **presumes** that a StableHLO was built
already **and** with Python bindings
* A new GitHub action is triggered on commits to main.
* Made minor changes to all the scripts to make a few CMake variables
configurable such as CMAKE_BUILD_TYPE and whether to build the Python
bindings

An example of this working can be seen on my fork
https://github.com/fzakaria/stablehlo release page
https://github.com/fzakaria/stablehlo/releases/tag/dev-wheels

Many thanks to @makslevental who helped me a lot on the Discord channel.
He pointed me to a few other solutions in MLIR that do similar approach
which was the inspiration such as:
- https://github.com/makslevental/pristine-llvm-release
- https://github.com/Xilinx/mlir-aie
- https://github.com/llvm/torch-mlir

.github/actions/setup-build/action.yml

@@ -14,13 +14,18 @@
 
 # The setup-build action gets everything needed by buildAndTest into the workspace.
 name: "Setup build environment (ninja, ccache, llvm, lld)"
+description: "Setup build environment (ninja, ccache, llvm, lld)"
 
 inputs:
   llvm-version:
     description: |
       LLVM version for checkout and build. Used for ccache value and checkout.
     required: true
-
+  python-version:
+    description: |
+      Python version to install
+    required: false
+    default: '3.10.6'
 runs:
   # This is a composite action - has a list of steps to execute.
   using: "composite"
@@ -54,9 +59,9 @@ runs:
 
   # Install Python/Numpy for API tests
   - name: Install Python and Pip
-    uses: actions/setup-python@v4
+    uses: actions/setup-python@v5
     with:
-      python-version: '3.10.6'
+      python-version: ${{ inputs.python-version }}
       cache: 'pip' # caching pip dependencies
   - name: Install MLIR python requirements
     shell: bash

.github/workflows/publishWheelRelease.yml

@@ -0,0 +1,132 @@
+# Copyright 2024 The StableHLO Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+name: Publish Python Wheel Release
+
+on:
+  workflow_dispatch:
+  schedule:
+    # Runs at 6:00 AM UTC, which is 10:00 PM previous day PST (UTC-8)
+    - cron: '0 6 * * *' 
+
+# Ensure that only a single job or workflow using the same
+# concurrency group will run at a time. This would cancel
+# any in-progress jobs in the same github workflow and github
+# ref (e.g. refs/heads/main or refs/pull/<pr_number>/merge).
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  # TODO(fzakaria): Consider running this job only if a new commit for the day has been done
+  cmake-build-wheel:
+    # Only run scheduled CI on main repo
+    if: (github.repository == 'openxla/stablehlo' || github.event_name != 'schedule')
+    name: "cmake-build ${{ github.event_name == 'schedule' && '(llvm-project@HEAD)' || ''}}"
+    env:
+      LLVM_PROJECT_DIR: "llvm-project"
+      LLVM_BUILD_DIR: "llvm-build"
+      # The setup.py script expects the StableHLO build to be in 'build'
+      STABLEHLO_BUILD_DIR: "build"
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: [ "3.10", "3.11" ]
+    runs-on: ${{ github.repository == 'openxla/stablehlo' && 'ubuntu-22.04-64core' ||  'ubuntu-22.04'  }}
+
+    steps:
+    - name: Checkout StableHLO
+      uses: actions/checkout@v4
+      # TODO(fzakaria): We need all the commits and the tags for the
+      # git describe` command to work to find the correct version.
+      # This can be removed if the version can be determmined in another way.
+      with:
+        fetch-depth: 0
+    - name: Get LLVM Version
+      id: llvm-version
+      shell: bash
+      run: |
+        echo "version=$(cat ${{ github.workspace }}/build_tools/llvm_version.txt)" >> $GITHUB_OUTPUT
+
+    - name: Setup workspace
+      uses: ./.github/actions/setup-build
+      with:
+        llvm-version: ${{ steps.llvm-version.outputs.version }}
+        python-version: ${{ matrix.python-version }}
+
+    # The actual tool need to execute the python build
+    # TODO(fzakaria): Consider using cibuildwheel to build multiple wheels at
+    # once and mark them as manylinux compatible.
+    - name: Install build dependencies
+      shell: bash
+      run: pip install build
+
+    - name: Configure and Build LLVM
+      shell: bash
+      run: |
+        ./build_tools/github_actions/ci_build_cmake_llvm.sh "$LLVM_PROJECT_DIR" "$LLVM_BUILD_DIR"
+      env:
+          CMAKE_BUILD_TYPE: Release
+          MLIR_ENABLE_BINDINGS_PYTHON: ON
+
+    - name: Build and Test StableHLO with Python
+      shell: bash
+      run: |
+        ./build_tools/github_actions/ci_build_cmake.sh "$LLVM_BUILD_DIR" "$STABLEHLO_BUILD_DIR"
+      env:
+          CMAKE_BUILD_TYPE: Release
+          STABLEHLO_ENABLE_BINDINGS_PYTHON: ON
+
+    - name: Build and Test Python Wheel
+      shell: bash
+      run: |
+        ./build_tools/github_actions/ci_build_python_wheel.sh "$GITHUB_WORKSPACE/${{ github.sha }}"
+    
+    - name: Upload an artifact
+      uses: actions/upload-artifact@v4
+      with:
+        if-no-files-found: error
+        name: build_artifact-${{ matrix.python-version }}
+        path: ${{ github.sha }}
+
+  upload-tarballs:
+    runs-on: ubuntu-20.04
+    needs: cmake-build-wheel
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: [ "3.10", "3.11" ]
+
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+
+      - name: Download artifacts
+        uses: actions/download-artifact@v4
+        with:
+          name: build_artifact-${{ matrix.python-version }}
+          path: ${{ github.sha }}
+
+      - name: Release current commit
+        uses: ncipollo/[email protected]
+        with:
+          artifacts: "${{ github.sha }}/*.whl"
+          token: "${{ secrets.GITHUB_TOKEN }}"
+          # git describe in setup.py explicitly filters out this tag
+          tag: "dev-wheels"
+          name: "dev-wheels"
+          removeArtifacts: false
+          allowUpdates: true
+          replacesArtifacts: true
+          makeLatest: true

README.md

@@ -68,7 +68,7 @@ Here's how to build the StableHLO repo on Linux or macOS:
 5. Configure and build MLIR:
 
    ```sh
-   build_tools/build_mlir.sh ${PWD}/llvm-project/ ${PWD}/llvm-build
+   MLIR_ENABLE_BINDINGS_PYTHON=ON build_tools/build_mlir.sh ${PWD}/llvm-project/ ${PWD}/llvm-build
    ```
 
    This will take a considerable amount of time. For example, on a MacBook Pro
@@ -85,7 +85,8 @@ Here's how to build the StableHLO repo on Linux or macOS:
    cmake .. -GNinja \
      -DLLVM_ENABLE_LLD="$LLVM_ENABLE_LLD" \
      -DCMAKE_BUILD_TYPE=Release \
-     -DLLVM_ENABLE_ASSERTIONS=On \
+     -DLLVM_ENABLE_ASSERTIONS=ON \
+     -DSTABLEHLO_ENABLE_BINDINGS_PYTHON=ON \
      -DMLIR_DIR=${PWD}/../llvm-build/lib/cmake/mlir
    ```
 
@@ -104,6 +105,37 @@ 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
+pip install  install -r ./llvm-project/mlir/python/requirements.txt
+```
+
+If you've built MLIR & StableHLO using the script above, the Python bindings
+for MLIR may 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
+>>> from mlir.ir import Context, Location
+>>> import mlir.dialects.arith
+```
+
+You can also build a wheel yourself using the `setup.py` file.
+We also make nightly wheels available on our GitHub Releases page.
+
+```shell
+pip install stablehlo -f https://github.com/openxla/stablehlo/releases/expanded_assets/dev-wheels
+```
+
 ## Community
 
 Building an amazing portability layer between ML frameworks and ML compilers

build_tools/build_mlir.sh

@@ -23,6 +23,9 @@ fi
 # LLVM source
 LLVM_SRC_DIR="$1"
 build_dir="$2"
+CMAKE_BUILD_TYPE="${CMAKE_BUILD_TYPE:-RelWithDebInfo}"
+# Turn on building Python bindings
+MLIR_ENABLE_BINDINGS_PYTHON="${MLIR_ENABLE_BINDINGS_PYTHON:-OFF}"
 
 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"
@@ -46,10 +49,14 @@ cmake -GNinja \
   -DLLVM_ENABLE_PROJECTS=mlir \
   -DLLVM_TARGETS_TO_BUILD=host \
   -DLLVM_INCLUDE_TOOLS=ON \
+  -DMLIR_ENABLE_BINDINGS_PYTHON="${MLIR_ENABLE_BINDINGS_PYTHON}" \
   -DLLVM_ENABLE_BINDINGS=OFF \
+  -DMLIR_ENABLE_BINDINGS_PYTHON=ON \
+  -DLLVM_VERSION_SUFFIX="" \
+  -DCMAKE_PLATFORM_NO_VERSIONED_SONAME:BOOL=ON \
   -DLLVM_BUILD_TOOLS=OFF \
   -DLLVM_INCLUDE_TESTS=OFF \
-  -DCMAKE_BUILD_TYPE=RelWithDebInfo \
-  -DLLVM_ENABLE_ASSERTIONS=On
+  -DCMAKE_BUILD_TYPE="$CMAKE_BUILD_TYPE" \
+  -DLLVM_ENABLE_ASSERTIONS=ON
 
 cmake --build "$build_dir" --target all

build_tools/github_actions/ci_build_cmake.sh

@@ -33,18 +33,22 @@ CMAKE_BUILD_TYPE="${CMAKE_BUILD_TYPE:-RelWithDebInfo}"
 STABLEHLO_ENABLE_BINDINGS_PYTHON="${STABLEHLO_ENABLE_BINDINGS_PYTHON:-OFF}"
 
 # Configure StableHLO
+# CMAKE_PLATFORM_NO_VERSIONED_SONAME Disables generation of "version soname"
+#                         (i.e. libFoo.so.<version>), which causes pure
+#                         duplication of various shlibs for Python wheels.
 cmake -GNinja \
   -B"$STABLEHLO_BUILD_DIR" \
   -DLLVM_ENABLE_LLD=ON \
-  -DCMAKE_BUILD_TYPE="${CMAKE_BUILD_TYPE}" \
+  -DCMAKE_BUILD_TYPE="$CMAKE_BUILD_TYPE" \
   -DLLVM_ENABLE_ASSERTIONS=ON \
   -DMLIR_DIR="$LLVM_BUILD_DIR/lib/cmake/mlir" \
   -DCMAKE_CXX_COMPILER=clang++ \
   -DCMAKE_CXX_COMPILER_LAUNCHER=ccache \
   -DCMAKE_C_COMPILER=clang \
   -DCMAKE_C_COMPILER_LAUNCHER=ccache \
   -DSTABLEHLO_ENABLE_STRICT_BUILD=ON \
-  -DSTABLEHLO_ENABLE_BINDINGS_PYTHON="${STABLEHLO_ENABLE_BINDINGS_PYTHON}"
+  -DCMAKE_PLATFORM_NO_VERSIONED_SONAME:BOOL=ON \
+  -DSTABLEHLO_ENABLE_BINDINGS_PYTHON="$STABLEHLO_ENABLE_BINDINGS_PYTHON"
 
 # Build and Test StableHLO
 cd "$STABLEHLO_BUILD_DIR" || exit

build_tools/github_actions/ci_build_cmake_llvm.sh

@@ -30,6 +30,10 @@ CMAKE_BUILD_TYPE="${CMAKE_BUILD_TYPE:-RelWithDebInfo}"
 MLIR_ENABLE_BINDINGS_PYTHON="${MLIR_ENABLE_BINDINGS_PYTHON:-OFF}"
 
 # Configure LLVM
+# LLVM_VERSION_SUFFIX to get rid of that annoying af git on the end of .17git
+# CMAKE_PLATFORM_NO_VERSIONED_SONAME Disables generation of "version soname"
+#                         (i.e. libFoo.so.<version>), which causes pure
+#                         duplication of various shlibs for Python wheels.
 cmake -GNinja \
   "-H$LLVM_SRC_DIR/llvm" \
   "-B$LLVM_BUILD_DIR" \
@@ -42,8 +46,10 @@ cmake -GNinja \
   -DLLVM_ENABLE_BINDINGS=OFF \
   -DLLVM_BUILD_TOOLS=OFF \
   -DLLVM_INCLUDE_TESTS=OFF \
-  -DCMAKE_BUILD_TYPE="${CMAKE_BUILD_TYPE}" \
+  -DCMAKE_BUILD_TYPE="$CMAKE_BUILD_TYPE" \
   -DLLVM_ENABLE_ASSERTIONS=On \
+  -DLLVM_VERSION_SUFFIX="" \
+  -DCMAKE_PLATFORM_NO_VERSIONED_SONAME:BOOL=ON \
   -DCMAKE_CXX_COMPILER=clang++ \
   -DCMAKE_CXX_COMPILER_LAUNCHER=ccache \
   -DCMAKE_C_COMPILER=clang \

build_tools/github_actions/ci_build_python_wheel.sh

@@ -0,0 +1,43 @@
+#!/usr/bin/env bash
+# Copyright 2024 The StableHLO Authors.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# This runs the necessary step to build a Python wheel for StableHLO
+# At the moment it only builds a Linux variant of Python wheel for the
+# default python3 version present on the system (set via GitHub action typically)
+# TODO: Incorporate cibuildwheel and build manylinux wheels
+set -o errexit
+set -o nounset
+set -o pipefail
+
+# Collect the root of the project directory relative to this file
+PROJECT_DIR="$(realpath "$(dirname "$0")"/../../)"
+
+# Set the source directory and destination directory for the distribution
+SRC_DIR="${PROJECT_DIR}/stablehlo/integrations/python"
+OUT_DIR="${1:-$SRC_DIR/dist}"
+
+echo "Building python wheel. You will find it at ${OUT_DIR}"
+python3 -m build -w --outdir "${OUT_DIR}" "${SRC_DIR}"
+
+echo "Testing that the python wheel works correctly"
+# Create a new virtual environment
+python3 -m venv venv
+# Activate the virtual environment
+# shellcheck disable=SC1091
+source venv/bin/activate
+# Install the wheel
+pip install --no-index --find-links="${OUT_DIR}" stablehlo
+# Run the smoke test script
+python "${SRC_DIR}/tests/smoketest.py"
+

stablehlo/integrations/python/setup.py

@@ -0,0 +1,58 @@
+# Copyright 2024 The StableHLO Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""This setup.py builds a wheel file assuming that StableHLO is already built
+"""
+from setuptools import find_namespace_packages, setup, Distribution
+import os
+import subprocess
+import time
+
+class BinaryDistribution(Distribution):
+  """Force distribution which always forces a binary package"""
+
+  def has_ext_modules(foo):
+    return True
+
+
+def get_version():
+  # get the latest tag without the leading v
+  latest_tag = subprocess.check_output(
+      ["git", "describe", "--tags", "--abbrev=0", "--exclude", "dev-wheels"], text=True).strip('v').strip()
+  latest_commit = subprocess.check_output(
+      ["git", "rev-parse", "--short", "HEAD"], text=True).strip()
+  # in order for the wheels to be ordered chronologically
+  # include the epoch seconds as a portion of the version
+  return f"{latest_tag}.{int(time.time())}+{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',
+    # TODO(fzakaria): The CMake build path is fixed here which kind of sucks.
+    # Ideally it should be passed in or setup.py should do the build itself.
+    packages=find_namespace_packages(where=os.path.normpath("../../../build/python_packages/stablehlo")),
+    package_dir={
+        "": os.path.normpath("../../../build/python_packages/stablehlo")},
+    package_data={'mlir': ['_mlir_libs/*.so']},
+    include_package_data=True,
+    distclass=BinaryDistribution,
+    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()
+)