From f7caf32465a2e9cf80fcf62a673910791d5bd92d Mon Sep 17 00:00:00 2001
From: Sunita Nadampalli <nadampal@amazon.com>
Date: Sun, 28 Apr 2024 11:11:55 -0500
Subject: [PATCH] aarch64: cd: switch from libomp to libgomp

In the current version of the CD scripts, torch libraries are
linked to llvm openmp because conda openblas-openmp is linked to it.
To switch to gnu libgomp, we are building the openblas from sources
instead of installing from conda.
Building openBLAS shared library instead of static library to
be able to discover LAPACK support in OpenBLAS.

cherrypicked from https://github.com/pytorch/builder/pull/1803
fixes: https://github.com/pytorch/builder/issues/1774
---
 aarch64_linux/aarch64_ci_setup.sh       |  2 +-
 aarch64_linux/aarch64_wheel_ci_build.py | 42 ++++++++++++++++++++++++-
 2 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/aarch64_linux/aarch64_ci_setup.sh b/aarch64_linux/aarch64_ci_setup.sh
index 87ed70a72..50bb2d1ea 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 c474184ed..7465eddeb 100755
--- a/aarch64_linux/aarch64_wheel_ci_build.py
+++ b/aarch64_linux/aarch64_wheel_ci_build.py
@@ -15,6 +15,43 @@ 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=0",
+        "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 +132,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 +140,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")