From b9b3af8003df35e7094525f9aea3ca605a0da9f9 Mon Sep 17 00:00:00 2001
From: powderluv <powderluv@users.noreply.github.com>
Date: Wed, 24 May 2023 13:52:13 -0700
Subject: [PATCH] [arm64] Fix release builds for ARM64 (#2157)

Tested on Ubuntu 23.04 on Ampere Altra instance.
---
 build-requirements.txt                            | 1 +
 build_tools/python_deploy/build_linux_packages.sh | 6 ++++--
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/build-requirements.txt b/build-requirements.txt
index db0c32051..1566aa676 100644
--- a/build-requirements.txt
+++ b/build-requirements.txt
@@ -4,6 +4,7 @@ wheel
 setuptools
 cmake
 ninja
+packaging
 
 # Workaround for what should be a torch dep
 # See discussion in #1174
diff --git a/build_tools/python_deploy/build_linux_packages.sh b/build_tools/python_deploy/build_linux_packages.sh
index cfb4dbfe5..580092789 100755
--- a/build_tools/python_deploy/build_linux_packages.sh
+++ b/build_tools/python_deploy/build_linux_packages.sh
@@ -38,8 +38,10 @@ set -eu -o errtrace
 
 this_dir="$(cd "$(dirname "$0")" && pwd)"
 repo_root="$(cd "$this_dir"/../../ && pwd)"
+arch="$(uname -m)"
+echo "Running on Arch: ${arch}"
 # This needs to be a manylinux image so we can ship pip packages
-TM_RELEASE_DOCKER_IMAGE="${TM_RELEASE_DOCKER_IMAGE:-gcr.io/iree-oss/manylinux2014_x86_64-release@sha256:d8994b87b45b7b2e6055fccc32db018ec73aeb05a4e43a9daa61b77cc34f846e}"
+TM_RELEASE_DOCKER_IMAGE="${TM_RELEASE_DOCKER_IMAGE:-quay.io/pypa/manylinux2014_${arch}}"
 # This assumes an Ubuntu LTS like image. You can build your own with
 # ./build_tools/docker/Dockerfile
 TM_CI_DOCKER_IMAGE="${TM_CI_DOCKER_IMAGE:-powderluv/torch-mlir-ci:latest}"
@@ -378,7 +380,7 @@ function build_torch_mlir() {
 function run_audit_wheel() {
   local wheel_basename="$1"
   local python_version="$2"
-  generic_wheel="/wheelhouse/${wheel_basename}-${TORCH_MLIR_PYTHON_PACKAGE_VERSION}-${python_version}-linux_x86_64.whl"
+  generic_wheel="/wheelhouse/${wheel_basename}-${TORCH_MLIR_PYTHON_PACKAGE_VERSION}-${python_version}-linux_${arch}.whl"
   echo ":::: Auditwheel $generic_wheel"
   auditwheel repair -w /wheelhouse "$generic_wheel"
   rm "$generic_wheel"