aarch64: cd: switch from libomp to libgomp (#1805)

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 #1803 fixes: #1774
pytorch · May 8, 2024 · ba10bae · ba10bae
1 parent e5d769d
commit ba10bae
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 2 deletions.
diff --git 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
@@ -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,14 +132,17 @@ 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")
         build_vars += "USE_MKLDNN=ON USE_MKLDNN_ACL=ON " \
                       "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")