From 4abc6e6273d240bf8305c07e767d21305d161338 Mon Sep 17 00:00:00 2001
From: Sunita Nadampalli <nadampal@amazon.com>
Date: Thu, 25 Apr 2024 23:26:08 +0000
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.
---
 aarch64_linux/aarch64_ci_setup.sh       |  2 +-
 aarch64_linux/aarch64_wheel_ci_build.py | 40 +++++++++++++++++++++++++
 2 files changed, 41 insertions(+), 1 deletion(-)

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 a7f513b26..e0b34c24b 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
@@ -186,6 +223,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")
@@ -195,6 +233,8 @@ def parse_arguments():
             "LD_LIBRARY_PATH=/pytorch/build/lib:/acl/build:$LD_LIBRARY_PATH "
             "ACL_INCLUDE_DIR=/acl/build "
             "ACL_LIBRARY=/acl/build "
+            "BLAS=OpenBLAS "
+            "OpenBLAS_HOME=/OpenBLAS "
         )
     else:
         print("build pytorch without mkldnn backend")