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

.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
 

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 \
+  local torch_version="$1"
-    --extra-index-url https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
+  case $torch_version in
-  CMAKE_GENERATOR=Ninja \
+    nightly)
-  TORCH_MLIR_PYTHON_PACKAGE_VERSION=${TORCH_MLIR_PYTHON_PACKAGE_VERSION} \
+      echo ":::: Using nightly dependencies"
-  python -m pip wheel -v -w /wheelhouse /main_checkout/torch-mlir \
+      python -m pip install --no-cache-dir -r /main_checkout/torch-mlir/requirements.txt \
-    -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html \
+        --extra-index-url https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
-    -r /main_checkout/torch-mlir/whl-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 \
+        -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() {

whl-requirements.txt

@@ -1,5 +1,2 @@
 -f build-requirements.txt
 -f pytorch-requirements.txt
-
-# Packaging requirements.
-packaging