From 57ad1cdcea96e2aac606883dad42c8ecf0f9a039 Mon Sep 17 00:00:00 2001
From: Sunita Nadampalli <nadampal@amazon.com>
Date: Thu, 18 Apr 2024 00:11:26 +0000
Subject: [PATCH] aarch64: cd: switch from libomp to libgomp

In the current version of the scripts, torch libraries are
linked to llvm openmp becasue conda openblas-openmp is linked to it.
to switch to gnu libgomp, we are building the openblas from sources
instead of installing from conda.
---
 aarch64_linux/aarch64_ci_setup.sh       |  2 +-
 aarch64_linux/aarch64_wheel_ci_build.py | 21 ++++++++++++++++++++-
 2 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/aarch64_linux/aarch64_ci_setup.sh b/aarch64_linux/aarch64_ci_setup.sh
index 87ed70a72d..50bb2d1ea2 100755
--- a/aarch64_linux/aarch64_ci_setup.sh
+++ b/aarch64_linux/aarch64_ci_setup.sh
@@ -30,7 +30,7 @@ if [[ "$DESIRED_PYTHON"  == "3.8" ]]; then
 else
     pip install -q --pre numpy==2.0.0rc1
 fi
-conda install -y -c conda-forge pyyaml==6.0.1 patchelf==0.17.2 pygit2==1.13.2 openblas==0.3.25=*openmp* ninja==1.11.1 scons==4.5.2
+conda install -y -c conda-forge pyyaml==6.0.1 patchelf==0.17.2 pygit2==1.13.2 ninja==1.11.1 scons==4.5.2
 
 python --version
 conda --version
diff --git a/aarch64_linux/aarch64_wheel_ci_build.py b/aarch64_linux/aarch64_wheel_ci_build.py
index c474184edd..4fa13eadf8 100755
--- a/aarch64_linux/aarch64_wheel_ci_build.py
+++ b/aarch64_linux/aarch64_wheel_ci_build.py
@@ -15,6 +15,22 @@ def list_dir(path: str) -> List[str]:
     return check_output(["ls", "-1", path]).decode().split("\n")
 
 
+def build_OpenBLAS() -> None:
+    '''
+    Building OpenBLAS, because the package in many linux is old
+    '''
+    print('Building OpenBLAS')
+    openblas_build_flags=["NUM_THREADS=128", "USE_OPENMP=1", "NO_SHARED=1", "DYNAMIC_ARCH=1", "TARGET=ARMV8", "CFLAGS=-O3"]
+    openblas_checkout_dir="OpenBLAS"
+
+    check_call(["git", "clone", "https://github.com/OpenMathLib/OpenBLAS.git", "-b", "v0.3.25",
+                "--depth", "1", "--shallow-submodules"])
+
+    check_call(["make", "-j8"]  + openblas_build_flags,
+               cwd=openblas_checkout_dir)
+    check_call(["make", "-j8"] + openblas_build_flags + ["install"], cwd=openblas_checkout_dir)
+
+
 def build_ArmComputeLibrary() -> None:
     '''
     Using ArmComputeLibrary for aarch64 PyTorch
@@ -95,6 +111,7 @@ def parse_arguments():
     elif branch.startswith(("v1.", "v2.")):
         build_vars += f"BUILD_TEST=0 PYTORCH_BUILD_VERSION={branch[1:branch.find('-')]} PYTORCH_BUILD_NUMBER=1 "
 
+    build_OpenBLAS()
     if enable_mkldnn:
         build_ArmComputeLibrary()
         print("build pytorch with mkldnn+acl backend")
@@ -102,7 +119,9 @@ def parse_arguments():
                       "ACL_ROOT_DIR=/acl " \
                       "LD_LIBRARY_PATH=/pytorch/build/lib:/acl/build:$LD_LIBRARY_PATH " \
                       "ACL_INCLUDE_DIR=/acl/build " \
-                      "ACL_LIBRARY=/acl/build "
+                      "ACL_LIBRARY=/acl/build " \
+                      "BLAS=OpenBLAS " \
+                      "OpenBLAS_HOME=/OpenBLAS "
     else:
         print("build pytorch without mkldnn backend")