From e7f306ec2fc532a43958e4381b3f6f8212a908ea Mon Sep 17 00:00:00 2001
From: powderluv <powderluv@users.noreply.github.com>
Date: Thu, 12 May 2022 14:16:32 -0700
Subject: [PATCH] Use delocate to make portable wheels on OSX (#850)

Fix up wheels per python version on OSX
---
 .../python_deploy/build_macos_packages.sh     | 26 +++++++++++++++++--
 .../python_deploy/build_upload_m1_snapshot.sh |  8 +++++-
 2 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/build_tools/python_deploy/build_macos_packages.sh b/build_tools/python_deploy/build_macos_packages.sh
index 389c3e0d3..7b2e647fb 100755
--- a/build_tools/python_deploy/build_macos_packages.sh
+++ b/build_tools/python_deploy/build_macos_packages.sh
@@ -20,7 +20,7 @@ set -eu -o errtrace
 
 this_dir="$(cd $(dirname $0) && pwd)"
 repo_root="$(cd $this_dir/../../ && pwd)"
-python_versions="${python_versions:-3.9 3.10}"
+python_versions="${TORCH_MLIR_PYTHON_VERSIONS:-3.9 3.10}"
 output_dir="${output_dir:-${this_dir}/wheelhouse}"
 packages="${packages:-torch-mlir}"
 
@@ -55,6 +55,7 @@ function run() {
         torch-mlir)
           clean_wheels torch-mlir $python_version
           build_torch_mlir
+          run_audit_wheel torch_mlir $python_version
           ;;
         *)
           echo "Unrecognized package '$package'"
@@ -78,7 +79,28 @@ function clean_wheels() {
   local wheel_basename="$1"
   local python_version="$2"
   echo ":::: Clean wheels $wheel_basename $python_version"
-  rm -f $output_dir/${wheel_basename}-*-${python_version}-*.whl
+  rm -rf $repo_root/build/
+  rm -f $output_dir/${wheel_basename}-*-${python_version//./}-*.whl
+}
+
+function run_audit_wheel() {
+  set +x
+  local wheel_basename="$1"
+  local python_version="$2"
+  generic_wheel=`ls  $output_dir/${wheel_basename}-* | grep ${python_version//./}`
+  echo "Looking for $generic_wheel"
+  if [ -f "$generic_wheel" ]; then
+    echo "$generic_wheel found. Delocating it.."
+    rm -rf $output_dir/test_venv
+    python${python_version} -m venv $output_dir/test_venv
+    source  $output_dir/test_venv/bin/activate
+    python${python_version} -m pip install -U pip
+    python${python_version} -m pip install -r $repo_root/requirements.txt --extra-index-url https://download.pytorch.org/whl/nightly/cpu
+    python${python_version} -m pip install $generic_wheel --extra-index-url https://download.pytorch.org/whl/nightly/cpu
+    DYLD_LIBRARY_PATH=$output_dir/test_venv/lib/python${python_version}/site-packages/torch/lib delocate-wheel -v $generic_wheel
+    deactivate
+    rm -rf $output_dir/test_venv
+  fi
 }
 
 run
diff --git a/build_tools/python_deploy/build_upload_m1_snapshot.sh b/build_tools/python_deploy/build_upload_m1_snapshot.sh
index fc8cc821b..9a6a8bf38 100755
--- a/build_tools/python_deploy/build_upload_m1_snapshot.sh
+++ b/build_tools/python_deploy/build_upload_m1_snapshot.sh
@@ -26,9 +26,15 @@ package_version=${latest_snapshot#"snapshot-"}
 echo "Latest snapshot tag is: ${latest_snapshot}"
 echo "Latest version is: ${package_version}"
 
+export TORCH_MLIR_PYTHON_VERSIONS="3.9 3.10"
+echo "Using Python Versions: ${TORCH_MLIR_PYTHON_VERSIONS}"
 export TORCH_MLIR_PYTHON_PACKAGE_VERSION="${package_version}"
 echo "Setting torch-mlir Python Package version to: ${TORCH_MLIR_PYTHON_PACKAGE_VERSION}"
 
-TORCH_MLIR_PYTHON_PACKAGE_VERSION="${package_version}" ${repo_root}/build_tools/python_deploy/build_macos_packages.sh
+TORCH_MLIR_OSX_ARCH=arm64 \
+TORCH_MLIR_OSX_TARGET=11.0 \
+TORCH_MLIR_PYTHON_PACKAGE_VERSION="${package_version}" \
+TORCH_MLIR_PYTHON_VERSIONS="${TORCH_MLIR_PYTHON_VERSIONS}" \
+${repo_root}/build_tools/python_deploy/build_macos_packages.sh
 
 gh release upload ${latest_snapshot} ${repo_root}/build_tools/python_deploy/wheelhouse/torch*.whl