Add ARM64 release builds (#2159)

Creates a build_linux_arm64 job that builds the release on an arm64 self-hosted runner. Drop Python 3.10 support Pass TM_TORCH_VERSION to choose the Stable PyTorch version (since arm64 doesn't have nightly builds) Borrows nightly / stable Pytorch switch from the WIP https://github.com/llvm/torch-mlir/pull/2038
2023-05-25 20:39:19 -07:00 · 2023-05-25 20:39:19 -07:00 · 9b7909b599
parent 69e993b03f
commit 9b7909b599
3 changed files with 125 additions and 17 deletions
--- a/.github/workflows/buildRelease.yml
+++ b/.github/workflows/buildRelease.yml
@ -12,12 +12,12 @@ on:

 jobs:
  build_linux:
-    name: Manylinux Build
+    name: Manylinux x86_64 Build
    runs-on: a100
    strategy:
      matrix:
        package: [ torch-mlir, torch-mlir-core ]
-        py_version: [ cp38-cp38, cp310-cp310, cp311-cp311 ]
+        py_version: [ cp38-cp38, cp311-cp311 ]
        exclude:
          - package: torch-mlir-core
            py_version: cp38-cp38
@ -85,6 +85,76 @@ jobs:
        name: wheels
        path: dist

+  build_linux_arm64:
+    name: Manylinux arm64 Build
+    runs-on: linux-arm64
+    strategy:
+      matrix:
+        package: [ torch-mlir, torch-mlir-core ]
+        py_version: [ cp311-cp311 ]
+
+    steps:
+
+    - name: Prepare workspace
+      run: |
+        # Clear the workspace directory so that we don't run into errors about
+        # existing lock files.
+        sudo rm -rf $GITHUB_WORKSPACE/*
+
+    - name: Get torch-mlir
+      uses: actions/checkout@v3
+      with:
+        submodules: 'true'
+        fetch-depth: 0
+
+    - uses: ./.github/actions/setup-build
+      with:
+        cache-enabled: 'false'
+    - name: Build Python wheels and smoke test.
+      run: |
+        cd $GITHUB_WORKSPACE
+        python -m pip install wheel
+        TM_PACKAGE_VERSION=${{ github.event.inputs.python_package_version }}
+        printf "TORCH_MLIR_PYTHON_PACKAGE_VERSION=%s\n" $TM_PACKAGE_VERSION > ./torch_mlir_package_version
+        TM_PYTHON_VERSIONS=${{ matrix.py_version }} TM_PACKAGES=${{ matrix.package }} TM_TORCH_VERSION="stable" ./build_tools/python_deploy/build_linux_packages.sh
+
+    # If we were given a release_id, then upload the package we just built
+    # to the github releases page.
+    - name: Upload Release Assets (if requested)
+      if: github.event.inputs.release_id != ''
+      id: upload-release-assets
+      uses: dwenegar/upload-release-assets@v1
+      env:
+        GITHUB_TOKEN: ${{ secrets.WORKFLOW_INVOCATION_TOKEN }}
+      with:
+        release_id: ${{ github.event.inputs.release_id }}
+        assets_path: ./build_tools/python_deploy/wheelhouse/torch*.whl
+    # Publishing is necessary to make the release visible to `pip`
+    # on the github releases page.
+    - name: Publish Release (if requested)
+      if: github.event.inputs.release_id != ''
+      id: publish_release
+      uses: eregon/publish-release@v1
+      env:
+        GITHUB_TOKEN: ${{ secrets.WORKFLOW_INVOCATION_TOKEN }}
+      with:
+        release_id: ${{ github.event.inputs.release_id }}
+    - name: Create dist directory
+      if: github.event.inputs.release_id != ''
+      run: mkdir dist
+    - name: Copy releases to publish to dist directory
+      if: github.event.inputs.release_id != ''
+      run: cp build_tools/python_deploy/wheelhouse/torch_mlir*.whl dist/
+
+    # Wheels must be published from a linux environment.
+    #
+    # See https://github.com/pypa/gh-action-pypi-publish/discussions/15
+    - name: Store the binary wheel
+      uses: actions/upload-artifact@v2
+      with:
+        name: wheels
+        path: dist
+
  build_macos:
    name: MacOS Build
    runs-on: macos-latest
@ -219,6 +289,7 @@ jobs:
    runs-on: ubuntu-latest
    needs:
    - build_linux
+    - build_linux_arm64
    - build_macos
    - build_windows

--- a/build_tools/python_deploy/build_linux_packages.sh
+++ b/build_tools/python_deploy/build_linux_packages.sh
@ -57,6 +57,8 @@ TM_USE_PYTORCH_BINARY="${TM_USE_PYTORCH_BINARY:-ON}"
 TM_SKIP_TESTS="${TM_SKIP_TESTS:-OFF}"
 # Update ODS and abstract interpretation library files
 TM_UPDATE_ODS_AND_ABSTRACT_INTERP_LIB="${TM_UPDATE_ODS_AND_ABSTRACT_INTERP_LIB:-OFF}"
+# Determine if we use a stable or a nightly torch build
+TM_TORCH_VERSION="${TM_TORCH_VERSION:-nightly}"

 PKG_VER_FILE="${repo_root}"/torch_mlir_package_version ; [ -f "$PKG_VER_FILE" ] && . "$PKG_VER_FILE"
 TORCH_MLIR_PYTHON_PACKAGE_VERSION="${TORCH_MLIR_PYTHON_PACKAGE_VERSION:-0.0.1}"
@ -131,6 +133,7 @@ function run_on_host() {
    -e "TORCH_MLIR_SRC_PYTORCH_REPO=${TORCH_MLIR_SRC_PYTORCH_REPO}" \
    -e "TORCH_MLIR_SRC_PYTORCH_BRANCH=${TORCH_MLIR_SRC_PYTORCH_BRANCH}" \
    -e "TM_PYTORCH_INSTALL_WITHOUT_REBUILD=${TM_PYTORCH_INSTALL_WITHOUT_REBUILD}" \
+    -e "TM_TORCH_VERSION=${TM_TORCH_VERSION}" \
    -e "CCACHE_DIR=/main_checkout/torch-mlir/.ccache" \
    "${TM_CURRENT_DOCKER_IMAGE}" \
    /bin/bash /main_checkout/torch-mlir/build_tools/python_deploy/build_linux_packages.sh
@ -157,7 +160,7 @@ function run_in_docker() {
      case "$package" in
        torch-mlir)
          clean_wheels torch_mlir "$python_version"
-          build_torch_mlir
+          build_torch_mlir "$TM_TORCH_VERSION"

          # Disable audit wheel until we can fix ODR torch issues.  See
          # https://github.com/llvm/torch-mlir/issues/1709
@ -173,14 +176,14 @@ function run_in_docker() {
          clean_build torch_mlir_core "$python_version"
          ;;
        out-of-tree)
-          setup_venv "$python_version"
+          setup_venv "$python_version" "$TM_TORCH_VERSION"
          build_out_of_tree "$TM_USE_PYTORCH_BINARY" "$python_version"
          if [ "${TM_SKIP_TESTS}" == "OFF" ]; then
            test_out_of_tree
          fi
          ;;
        in-tree)
-          setup_venv "$python_version"
+          setup_venv "$python_version" "$TM_TORCH_VERSION"
          build_in_tree "$TM_USE_PYTORCH_BINARY" "$python_version"
          if [ "${TM_UPDATE_ODS_AND_ABSTRACT_INTERP_LIB}" == "ON" ]; then
            pushd /main_checkout/torch-mlir
@ -295,7 +298,9 @@ function test_in_tree() {

 function setup_venv() {
  local python_version="$1"
-  echo ":::: Setting up VENV with Python: $python_version"
+  local torch_version="$2"
+
+  echo ":::: Setting up VENV with Python: $python_version PyTorch $torch_version"
  python3 -m venv /main_checkout/torch-mlir/docker_venv
  source /main_checkout/torch-mlir/docker_venv/bin/activate

@ -303,6 +308,23 @@ function setup_venv() {
  python3 -m pip install --no-cache-dir -r /main_checkout/torch-mlir/externals/llvm-project/mlir/python/requirements.txt
  python3 -m pip install --no-cache-dir -r /main_checkout/torch-mlir/requirements.txt

+  case $torch_version in
+    nightly)
+      echo ":::: Using nightly dependencies"
+      python3 -m pip install --no-cache-dir -r /main_checkout/torch-mlir/requirements.txt
+      ;;
+    stable)
+      echo ":::: Using stable dependencies"
+      python3 -m pip install --no-cache-dir torch torchvision
+      python3 -m pip install --no-cache-dir -r /main_checkout/torch-mlir/build-requirements.txt
+      python3 -m pip install --no-cache-dir -r /main_checkout/torch-mlir/test-requirements.txt
+      ;;
+    *)
+      echo "Unrecognized torch version '$torch_version'"
+      exit 1
+      ;;
+  esac
+
 }

 function build_out_of_tree() {
@ -368,13 +390,31 @@ function clean_build() {
 }

 function build_torch_mlir() {
-  python -m pip install --no-cache-dir -r /main_checkout/torch-mlir/requirements.txt \
-    --extra-index-url https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
-  CMAKE_GENERATOR=Ninja \
-  TORCH_MLIR_PYTHON_PACKAGE_VERSION=${TORCH_MLIR_PYTHON_PACKAGE_VERSION} \
-  python -m pip wheel -v -w /wheelhouse /main_checkout/torch-mlir \
-    -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html \
-    -r /main_checkout/torch-mlir/whl-requirements.txt
+  local torch_version="$1"
+  case $torch_version in
+    nightly)
+      echo ":::: Using nightly dependencies"
+      python -m pip install --no-cache-dir -r /main_checkout/torch-mlir/requirements.txt \
+        --extra-index-url https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
+      CMAKE_GENERATOR=Ninja \
+      TORCH_MLIR_PYTHON_PACKAGE_VERSION=${TORCH_MLIR_PYTHON_PACKAGE_VERSION} \
+      python -m pip wheel -v -w /wheelhouse /main_checkout/torch-mlir \
+        -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html \
+        -r /main_checkout/torch-mlir/whl-requirements.txt
+      ;;
+    stable)
+      echo ":::: Using stable dependencies"
+      python3 -m pip install --no-cache-dir torch torchvision
+      python3 -m pip install --no-cache-dir -r /main_checkout/torch-mlir/build-requirements.txt
+      CMAKE_GENERATOR=Ninja \
+      TORCH_MLIR_PYTHON_PACKAGE_VERSION=${TORCH_MLIR_PYTHON_PACKAGE_VERSION} \
+      python -m pip wheel -v -w /wheelhouse /main_checkout/torch-mlir
+      ;;
+    *)
+      echo "Unrecognized torch version '$torch_version'"
+      exit 1
+      ;;
+  esac
 }

 function run_audit_wheel() {
--- a/whl-requirements.txt
+++ b/whl-requirements.txt
@ -1,5 +1,2 @@
 -f build-requirements.txt
 -f pytorch-requirements.txt
-
-# Packaging requirements.
-packaging