From aba15015bf098b36ff0354ec4a7afb5e055a8b97 Mon Sep 17 00:00:00 2001
From: anakinxc <103552181+anakinxc@users.noreply.github.com>
Date: Mon, 11 Dec 2023 13:32:01 +0800
Subject: [PATCH] Add experimental Linux arm build (#420)

---
 .circleci/continue-config.yml | 33 ++++++++++-----------------------
 CHANGELOG.md                  |  5 +++--
 bazel/repositories.bzl        |  2 +-
 requirements-dev.txt          |  4 ++--
 4 files changed, 16 insertions(+), 28 deletions(-)

diff --git a/.circleci/continue-config.yml b/.circleci/continue-config.yml
index f3c40f2d..b5579e7f 100644
--- a/.circleci/continue-config.yml
+++ b/.circleci/continue-config.yml
@@ -29,7 +29,10 @@ jobs:
     # See: https://circleci.com/docs/2.0/configuration-reference/#docker-machine-macos-windows-executor
     docker:
       - image: secretflow/ubuntu-base-ci:latest
-    resource_class: 2xlarge+
+    parameters:
+      resource_class:
+        type: string
+    resource_class: << parameters.resource_class >>
     shell: /bin/bash --login -eo pipefail
     # Add steps to the job
     # See: https://circleci.com/docs/2.0/configuration-reference/#steps
@@ -60,7 +63,6 @@ jobs:
           command: |
             python3 -m pip install -U pip
             python3 -m pip install -r requirements.txt
-            sed -i "s/tensorflow /tensorflow-cpu /g" ./requirements-dev.txt
             python3 -m pip install -r requirements-dev.txt
       - run:
           name: "build"
@@ -104,15 +106,9 @@ jobs:
     resource_class: macos.m1.large.gen1
     steps:
       - checkout
-      - restore_cache:
-          name: "Restore build cache"
-          key: spu-build-comp-{{ arch }}-
       - run:
           name: Checkout devtools
           command: git clone https://github.com/secretflow/devtools.git ../devtools
-      - run:
-          name: Decompress cache
-          command: sh ../devtools/decompress-build-cache.sh spu_build_cache
       - run:
           name: "Install homebrew dependencies"
           command: brew install bazelisk cmake ninja libomp wget go md5sha1sum
@@ -127,42 +123,33 @@ jobs:
             pip install -r requirements-dev.txt
       - run:
           name: "build"
-          command: bazel build //... -c opt --ui_event_filters=-info,-debug,-warning --disk_cache=~/.cache/spu_build_cache
+          command: bazel build //... -c opt --ui_event_filters=-info,-debug,-warning
       - run:
           name: "test"
           command: |
             set +e
             declare -i test_status
-            bazel test //... -c opt --ui_event_filters=-info,-debug,-warning --test_output=errors --disk_cache=~/.cache/spu_build_cache | tee test_result.log; test_status=${PIPESTATUS[0]}
+            bazel test //... -c opt --ui_event_filters=-info,-debug,-warning --test_output=errors | tee test_result.log; test_status=${PIPESTATUS[0]}
 
             sh ../devtools/rename-junit-xml.sh
             find bazel-bin/ -perm +111 -type f -name "*_test"  -print0 | xargs -0 tar -cvzf test_binary.tar.gz
             find bazel-testlogs/ -type f -name "test.log"  -print0 | xargs -0 tar -cvzf test_logs.tar.gz
 
             exit ${test_status}
-      - run:
-          name: Cleanup and compress cache
-          command: |
-            sh ../devtools/clean_disk_cache.sh ~/.cache/spu_build_cache
-            sh ../devtools/compress-build-cache.sh spu_build_cache
-          when: always
-      - save_cache:
-          key: spu-build-comp-{{ arch }}-{{ .Environment.CIRCLE_BUILD_NUM }}
-          paths:
-            - ~/.cache/spu_build_cache.tar.gz
-          when: always
       - store_test_results:
           path: test-results
       - store_artifacts:
           path: test_binary.tar.gz
       - store_artifacts:
           path: test_logs.tar.gz
-
 # Invoke jobs via workflows
 # See: https://circleci.com/docs/2.0/configuration-reference/#workflows
 workflows:
   unittest:
     when: << pipeline.parameters.build-and-run >>
     jobs:
-      - linux_ut
+      - linux_ut:
+          matrix:
+            parameters:
+              resource_class: ["2xlarge", "arm-xlarge"]
       - macOS_ut
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 5320f730..efc46dea 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,10 +11,11 @@
 > please add your unreleased change here.
 
 - [Feature] Add Odd-Even Merge Sort to replace the bitonic sort.
+- [Feature] Add radix sort support for ABY3.
+- [Feature] Integrate with secretflow/psi.
+- [Feature] Add Linux aarch64 support.
 - [Improvement] Optimize sort memory usage.
-- [Feature] Add radix sort support for ABY3
 - [Deprecated] macOS 11.x is no longer supported
-- [Feature] Integrate with secretflow/psi
 
 ## 20231108
 
diff --git a/bazel/repositories.bzl b/bazel/repositories.bzl
index 6dac0115..11dcb110 100644
--- a/bazel/repositories.bzl
+++ b/bazel/repositories.bzl
@@ -20,7 +20,7 @@ SECRETFLOW_GIT = "https://github.com/secretflow"
 
 YACL_COMMIT_ID = "d9f18c74ba2a1bbbe52bd6beac644adf81fb9925"
 
-LIBSPI_COMMIT_ID = "182a2e12850187790309331997859b0b2de87f4b"
+LIBSPI_COMMIT_ID = "575cb5999b5bc52926d08de6944e8d54f0a38f9f"
 
 def spu_deps():
     _rules_cuda()
diff --git a/requirements-dev.txt b/requirements-dev.txt
index f3a5c515..515765f4 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -1,6 +1,6 @@
 matplotlib
-tensorflow >= 2.7.0, < 2.12.0; platform_machine != "arm64"
-tensorflow-macos >= 2.7.0, < 2.12.0; platform_machine == "arm64" and platform_system == "Darwin"
+tensorflow >= 2.7.0; platform_machine != "arm64"
+tensorflow-macos >= 2.7.0; platform_machine == "arm64" and platform_system == "Darwin"
 tensorflow_datasets
 pandas>=1.4.2
 flax