From 2bb46dff7b70d3bb63044539a39a2502d18cb23c Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lukas.geiger94@gmail.com>
Date: Fri, 10 Sep 2021 11:39:09 +0100
Subject: [PATCH] Add support for publishing macOS M1 ARM64 wheels (#2559)

* Add support for publishing macOS M1 ARM64 wheels

* Link against the tensorflow-macos when building for ARM Macs
---
 ...wheel_Linux.sh => make_wheel_Linux_x86.sh} |  0
 ...l_Windows.sh => make_wheel_Windows_x86.sh} |  0
 .github/workflows/make_wheel_macOS_arm64.sh   | 31 +++++++++++++++++++
 ...wheel_macOS.sh => make_wheel_macOS_x86.sh} |  0
 .github/workflows/release.yml                 | 27 ++++++++++++++--
 5 files changed, 55 insertions(+), 3 deletions(-)
 rename .github/workflows/{make_wheel_Linux.sh => make_wheel_Linux_x86.sh} (100%)
 rename .github/workflows/{make_wheel_Windows.sh => make_wheel_Windows_x86.sh} (100%)
 create mode 100644 .github/workflows/make_wheel_macOS_arm64.sh
 rename .github/workflows/{make_wheel_macOS.sh => make_wheel_macOS_x86.sh} (100%)

diff --git a/.github/workflows/make_wheel_Linux.sh b/.github/workflows/make_wheel_Linux_x86.sh
similarity index 100%
rename from .github/workflows/make_wheel_Linux.sh
rename to .github/workflows/make_wheel_Linux_x86.sh
diff --git a/.github/workflows/make_wheel_Windows.sh b/.github/workflows/make_wheel_Windows_x86.sh
similarity index 100%
rename from .github/workflows/make_wheel_Windows.sh
rename to .github/workflows/make_wheel_Windows_x86.sh
diff --git a/.github/workflows/make_wheel_macOS_arm64.sh b/.github/workflows/make_wheel_macOS_arm64.sh
new file mode 100644
index 0000000000..518c0cf663
--- /dev/null
+++ b/.github/workflows/make_wheel_macOS_arm64.sh
@@ -0,0 +1,31 @@
+set -e -x
+
+export TF_NEED_CUDA=0
+
+python --version
+python -m pip install --default-timeout=1000 delocate wheel setuptools tensorflow==$TF_VERSION
+
+python configure.py
+
+# For dynamic linking, we want the ARM version of TensorFlow.
+# Since we cannot run it on x86 so we need to force pip to install it regardless
+python -m pip install \
+  --platform=macosx_11_0_arm64 \
+  --no-deps \
+  --target=$(python -c 'import site; print(site.getsitepackages()[0])') \
+  --upgrade \
+  tensorflow-macos==$TF_VERSION
+
+bazel build \
+  --cpu=darwin_arm64 \
+  --copt -mmacosx-version-min=11.0 \
+  --linkopt -mmacosx-version-min=11.0 \
+  --noshow_progress \
+  --noshow_loading_progress \
+  --verbose_failures \
+  --test_output=errors \
+  build_pip_pkg
+
+bazel-bin/build_pip_pkg artifacts "--plat-name macosx_11_0_arm64 $NIGHTLY_FLAG"
+delocate-wheel -w wheelhouse artifacts/*.whl
+
diff --git a/.github/workflows/make_wheel_macOS.sh b/.github/workflows/make_wheel_macOS_x86.sh
similarity index 100%
rename from .github/workflows/make_wheel_macOS.sh
rename to .github/workflows/make_wheel_macOS_x86.sh
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
index c3d72d2537..52742c451d 100644
--- a/.github/workflows/release.yml
+++ b/.github/workflows/release.yml
@@ -41,6 +41,16 @@ jobs:
         os: ['macos-latest', 'windows-latest', 'ubuntu-18.04']
         py-version: ['3.6', '3.7', '3.8', '3.9']
         tf-version: ['2.4.2', '2.6.0']
+        cpu: ['x86']
+        include:
+          - os: 'macos-11'
+            cpu: 'arm64'
+            tf-version: '2.5.0'
+            py-version: '3.8'
+          - os: 'macos-11'
+            cpu: 'arm64'
+            tf-version: '2.5.0'
+            py-version: '3.9'
       fail-fast: false
     steps:
       - uses: actions/github-script@0.3.0
@@ -68,12 +78,13 @@ jobs:
           PY_VERSION: ${{ matrix.py-version }}
           TF_VERSION: ${{ matrix.tf-version }}
           NIGHTLY_TIME: ${{ steps.author-date.outputs.result }}
+          CPU: ${{ matrix.cpu }}
         shell: bash
-        run: bash .github/workflows/make_wheel_${OS}.sh
+        run: bash .github/workflows/make_wheel_${OS}_${CPU}.sh
       - if: matrix.py-version != '3.9' || matrix.tf-version != '2.4.2'
         uses: actions/upload-artifact@v1
         with:
-          name: ${{ runner.os }}-${{ matrix.py-version }}-tf${{ matrix.tf-version }}-wheel
+          name: ${{ runner.os }}-${{ matrix.py-version }}-tf${{ matrix.tf-version }}-${{ matrix.cpu }}-wheel
           path: wheelhouse
   upload-wheels:
     name: Publish wheels to PyPi
@@ -84,12 +95,22 @@ jobs:
         os: ['macOS', 'Windows', 'Linux']
         py-version: ['3.6', '3.7', '3.8', '3.9']
         tf-version: ['2.6.0']
+        cpu: ['x86']
+        include:
+          - os: 'macos-11'
+            cpu: 'arm64'
+            tf-version: '2.5.0'
+            py-version: '3.8'
+          - os: 'macos-11'
+            cpu: 'arm64'
+            tf-version: '2.5.0'
+            py-version: '3.9'
       fail-fast: false
     if: (github.event_name == 'push' && github.ref == 'refs/heads/master') || github.event_name == 'release'
     steps:
       - uses: actions/download-artifact@v1
         with:
-          name: ${{ matrix.os }}-${{ matrix.py-version }}-tf${{ matrix.tf-version }}-wheel
+          name: ${{ matrix.os }}-${{ matrix.py-version }}-tf${{ matrix.tf-version }}-${{ matrix.cpu }}-wheel
           path: ./dist
       - run: |
           set -e -x