aarch64: cd: switch from libomp to libgomp

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

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")