From ad063dfeeed7a3bec5ec7e5b3c579ee431e95378 Mon Sep 17 00:00:00 2001
From: Ting Lu <tingl@nvidia.com>
Date: Thu, 30 May 2024 03:56:34 -0700
Subject: [PATCH 1/2] apply openblas cache for cpu-aarch64

---
 manywheel/Dockerfile_aarch64 | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/manywheel/Dockerfile_aarch64 b/manywheel/Dockerfile_aarch64
index abfc2fd84..b0716d158 100644
--- a/manywheel/Dockerfile_aarch64
+++ b/manywheel/Dockerfile_aarch64
@@ -78,9 +78,16 @@ ADD ./common/install_openssl.sh install_openssl.sh
 RUN bash ./install_openssl.sh && rm install_openssl.sh
 ENV SSL_CERT_FILE=/opt/_internal/certs.pem
 
+FROM base as openblas
+# Install openblas
+ADD ./common/install_openblas.sh install_openblas.sh
+RUN bash ./install_openblas.sh && rm install_openblas.sh
+
 FROM openssl as final
 # remove unncessary python versions
 RUN rm -rf /opt/python/cp26-cp26m /opt/_internal/cpython-2.6.9-ucs2
 RUN rm -rf /opt/python/cp26-cp26mu /opt/_internal/cpython-2.6.9-ucs4
 RUN rm -rf /opt/python/cp33-cp33m /opt/_internal/cpython-3.3.6
 RUN rm -rf /opt/python/cp34-cp34m /opt/_internal/cpython-3.4.6
+COPY --from=openblas     /opt/OpenBLAS/  /opt/OpenBLAS/
+ENV LD_LIBRARY_PATH=/opt/OpenBLAS/lib:$LD_LIBRARY_PATH
\ No newline at end of file

From dedb51dea4796986ce03f3f4ceb828e14686dcdf Mon Sep 17 00:00:00 2001
From: Ting Lu <tingl@nvidia.com>
Date: Thu, 30 May 2024 08:53:39 -0700
Subject: [PATCH 2/2] reapply for cuda-aarch64

---
 aarch64_linux/aarch64_wheel_ci_build.py | 39 -------------------------
 common/install_openblas.sh              | 21 +++++++++++++
 manywheel/Dockerfile_cuda_aarch64       |  9 +++++-
 3 files changed, 29 insertions(+), 40 deletions(-)
 create mode 100644 common/install_openblas.sh

diff --git a/aarch64_linux/aarch64_wheel_ci_build.py b/aarch64_linux/aarch64_wheel_ci_build.py
index 458de5464..4a3cfd38f 100755
--- a/aarch64_linux/aarch64_wheel_ci_build.py
+++ b/aarch64_linux/aarch64_wheel_ci_build.py
@@ -14,44 +14,6 @@ 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
@@ -227,7 +189,6 @@ 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")
diff --git a/common/install_openblas.sh b/common/install_openblas.sh
new file mode 100644
index 000000000..e2deec811
--- /dev/null
+++ b/common/install_openblas.sh
@@ -0,0 +1,21 @@
+#!/bin/bash
+
+set -ex
+
+cd /
+git clone https://github.com/OpenMathLib/OpenBLAS.git -b v0.3.25 --depth 1 --shallow-submodules
+
+
+OPENBLAS_BUILD_FLAGS="
+NUM_THREADS=128
+USE_OPENMP=1
+NO_SHARED=0
+DYNAMIC_ARCH=1
+TARGET=ARMV8
+CFLAGS=-O3
+"
+
+OPENBLAS_CHECKOUT_DIR="OpenBLAS"
+
+make -j8 ${OPENBLAS_BUILD_FLAGS} -C ${OPENBLAS_CHECKOUT_DIR}
+make -j8 ${OPENBLAS_BUILD_FLAGS} install -C ${OPENBLAS_CHECKOUT_DIR}
\ No newline at end of file
diff --git a/manywheel/Dockerfile_cuda_aarch64 b/manywheel/Dockerfile_cuda_aarch64
index 74c60b299..d243b06b8 100644
--- a/manywheel/Dockerfile_cuda_aarch64
+++ b/manywheel/Dockerfile_cuda_aarch64
@@ -74,10 +74,17 @@ ARG BASE_CUDA_VERSION
 ADD ./common/install_magma.sh install_magma.sh
 RUN bash ./install_magma.sh ${BASE_CUDA_VERSION} && rm install_magma.sh
 
+FROM base as openblas
+# Install openblas
+ADD ./common/install_openblas.sh install_openblas.sh
+RUN bash ./install_openblas.sh && rm install_openblas.sh
+
 FROM final as cuda_final
 ARG BASE_CUDA_VERSION
 RUN rm -rf /usr/local/cuda-${BASE_CUDA_VERSION}
 COPY --from=cuda     /usr/local/cuda-${BASE_CUDA_VERSION}  /usr/local/cuda-${BASE_CUDA_VERSION}
 COPY --from=magma    /usr/local/cuda-${BASE_CUDA_VERSION}  /usr/local/cuda-${BASE_CUDA_VERSION}
+COPY --from=openblas     /opt/OpenBLAS/  /opt/OpenBLAS/
 RUN ln -sf /usr/local/cuda-${BASE_CUDA_VERSION} /usr/local/cuda
-ENV PATH=/usr/local/cuda/bin:$PATH
\ No newline at end of file
+ENV PATH=/usr/local/cuda/bin:$PATH
+ENV LD_LIBRARY_PATH=/opt/OpenBLAS/lib:$LD_LIBRARY_PATH
\ No newline at end of file