diff --git a/aarch64_linux/aarch64_wheel_ci_build.py b/aarch64_linux/aarch64_wheel_ci_build.py index 61efa10f2..df73970fe 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 @@ -222,7 +184,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..d10a500a1 --- /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} 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