diff --git a/aarch64_linux/README.md b/aarch64_linux/README.md new file mode 100644 index 000000000..0f899b444 --- /dev/null +++ b/aarch64_linux/README.md @@ -0,0 +1,19 @@ +# Aarch64 (ARM/Graviton) Support Scripts +Scripts for building aarch64 PyTorch PIP Wheels. These scripts build the following wheels: +* torch +* torchvision +* torchaudio +* torchtext +* torchdata +## Aarch64_ci_build.sh +This script is design to support CD operations within PyPi manylinux aarch64 container, and be executed in the container. It prepares the container and then executes __aarch64_wheel_ci_build.py__ to build the wheels. The script "assumes" the PyTorch repo is located at: ```/pytorch``` and will put the wheels into ```/artifacts```. +### Usage +```DESIRED_PYTHON= aarch64_ci_build.sh``` + +__NOTE:__ CI build is currently __EXPERMINTAL__ + +## Build_aarch64_wheel.py +This app allows a person to build using AWS EC3 resources and requires AWS-CLI and Boto3 with AWS credentials to support building EC2 instances for the wheel builds. Can be used in a codebuild CD or from a local system. + +### Usage +```build_aarch64_wheel.py --key-name --use-docker --python 3.7 --branch ``` \ No newline at end of file diff --git a/aarch64_linux/aarch64_ci_build.sh b/aarch64_linux/aarch64_ci_build.sh new file mode 100644 index 000000000..c72698389 --- /dev/null +++ b/aarch64_linux/aarch64_ci_build.sh @@ -0,0 +1,52 @@ +#!/bin/bash +set -eux -o pipefail + +# This script is used to prepare the Docker container for aarch64_ci_wheel_build.py python script +# as we need to install conda and setup the python version for the build. + +CONDA_PYTHON_EXE=/opt/conda/bin/python +CONDA_EXE=/opt/conda/bin/conda +PATH=/opt/conda/bin:$PATH + +############################################################################### +# Install OS dependent packages +############################################################################### +yum -y install epel-release +yum -y install less zstd + +############################################################################### +# Install conda +# disable SSL_verify due to getting "Could not find a suitable TLS CA certificate bundle, invalid path" +# when using Python version, less than the conda latest +############################################################################### +echo 'Installing conda-forge' +curl -L -o /mambaforge.sh https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-aarch64.sh +chmod +x /mambaforge.sh +/mambaforge.sh -b -p /opt/conda +rm /mambaforge.sh +/opt/conda/bin/conda config --set ssl_verify False +/opt/conda/bin/conda install -y -c conda-forge python=${DESIRED_PYTHON} numpy pyyaml setuptools patchelf +python --version +conda --version + +############################################################################### +# Exec libglfortran.a hack +# +# libgfortran.a from quay.io/pypa/manylinux2014_aarch64 is not compiled with -fPIC. +# This causes __stack_chk_guard@@GLIBC_2.17 on pytorch build. To solve, get +# ubuntu's libgfortran.a which is compiled with -fPIC +############################################################################### +cd ~/ +curl -L -o ~/libgfortran-10-dev.deb http://ports.ubuntu.com/ubuntu-ports/pool/universe/g/gcc-10/libgfortran-10-dev_10.4.0-6ubuntu1_arm64.deb +ar x ~/libgfortran-10-dev.deb +tar --use-compress-program=unzstd -xvf data.tar.zst -C ~/ +cp -f ~/usr/lib/gcc/aarch64-linux-gnu/10/libgfortran.a /opt/rh/devtoolset-10/root/usr/lib/gcc/aarch64-redhat-linux/10/ + +############################################################################### +# Run aarch64 builder python +############################################################################### +cd / +# adding safe directory for git as the permissions will be +# on the mounted pytorch repo +git config --global --add safe.directory /pytorch +python /builder/aarch64_linux/aarch64_wheel_ci_build.py --enable-mkldnn diff --git a/aarch64_linux/aarch64_wheel_ci_build.py b/aarch64_linux/aarch64_wheel_ci_build.py new file mode 100755 index 000000000..c76f6d647 --- /dev/null +++ b/aarch64_linux/aarch64_wheel_ci_build.py @@ -0,0 +1,309 @@ +#!/usr/bin/env python3 + +import os +import subprocess +from typing import Dict, List, Optional, Tuple + + +'''' +Helper for getting paths for Python +''' +def list_dir(path: str) -> List[str]: + return subprocess.check_output(["ls", "-1", path]).decode().split("\n") + + +''' +Helper to get repo branches for specific versions +''' +def checkout_repo(branch: str = "main", + url: str = "", + git_clone_flags: str = "", + mapping: Dict[str, Tuple[str, str]] = []) -> Optional[str]: + for prefix in mapping: + if not branch.startswith(prefix): + continue + tag = f"v{mapping[prefix][0]}-{mapping[prefix][1]}" + os.system(f"git clone {url} -b {tag} {git_clone_flags}") + return mapping[prefix][0] + + os.system(f"git clone {url} {git_clone_flags}") + return None + + +''' +Using OpenBLAS with PyTorch +''' +def build_OpenBLAS(git_clone_flags: str = "") -> None: + print('Building OpenBLAS') + os.system(f"cd /; git clone https://github.com/xianyi/OpenBLAS -b v0.3.21 {git_clone_flags}") + make_flags = "NUM_THREADS=64 USE_OPENMP=1 NO_SHARED=1 DYNAMIC_ARCH=1 TARGET=ARMV8 " + os.system(f"cd OpenBLAS; make {make_flags} -j8; make {make_flags} install; cd /; rm -rf OpenBLAS") + + +''' +Using ArmComputeLibrary for aarch64 PyTorch +''' +def build_ArmComputeLibrary(git_clone_flags: str = "") -> None: + print('Building Arm Compute Library') + os.system("cd / && mkdir /acl") + os.system(f"git clone https://github.com/ARM-software/ComputeLibrary.git -b v22.11 {git_clone_flags}") + os.system(f"cd ComputeLibrary; export acl_install_dir=/acl; " \ + f"scons Werror=1 -j8 debug=0 neon=1 opencl=0 os=linux openmp=1 cppthreads=0 arch=armv8.2-a multi_isa=1 build=native build_dir=$acl_install_dir/build; " \ + f"cp -r arm_compute $acl_install_dir; " \ + f"cp -r include $acl_install_dir; " \ + f"cp -r utils $acl_install_dir; " \ + f"cp -r support $acl_install_dir; " \ + f"cp -r src $acl_install_dir; cd /") + + +''' +Script to embed libgomp to the wheels +''' +def embed_libgomp(wheel_name) -> None: + print('Embedding libgomp into wheel') + os.system(f"python3 /builder/aarch64_linux/embed_library.py {wheel_name} --update-tag") + + +''' +Build TorchVision wheel +''' +def build_torchvision(branch: str = "main", + git_clone_flags: str = "") -> str: + print('Checking out TorchVision repo') + build_version = checkout_repo(branch=branch, + url="https://github.com/pytorch/vision", + git_clone_flags=git_clone_flags, + mapping={ + "v1.7.1": ("0.8.2", "rc2"), + "v1.8.0": ("0.9.0", "rc3"), + "v1.8.1": ("0.9.1", "rc1"), + "v1.9.0": ("0.10.0", "rc1"), + "v1.10.0": ("0.11.1", "rc1"), + "v1.10.1": ("0.11.2", "rc1"), + "v1.10.2": ("0.11.3", "rc1"), + "v1.11.0": ("0.12.0", "rc1"), + "v1.12.0": ("0.13.0", "rc4"), + "v1.12.1": ("0.13.1", "rc6"), + "v1.13.0": ("0.14.0", "rc4"), + "v1.13.1": ("0.14.1", "rc2"), + "v2.0.0": ("0.15.0", "rc2"), + }) + print('Building TorchVision wheel') + build_vars = "CMAKE_SHARED_LINKER_FLAGS=-Wl,-z,max-page-size=0x10000 " + if branch == 'nightly': + version = '' + if os.path.exists('/vision/version.txt'): + version = subprocess.check_output(['cat', '/vision/version.txt']).decode().strip() + if len(version) == 0: + # In older revisions, version was embedded in setup.py + version = subprocess.check_output(['grep', 'version', 'setup.py']).decode().strip().split('\'')[1][:-2] + build_date = subprocess.check_output(['git','log','--pretty=format:%cs','-1'], cwd='/vision').decode().replace('-','') + build_vars += f"BUILD_VERSION={version}.dev{build_date}" + elif build_version is not None: + build_vars += f"BUILD_VERSION={build_version}" + + os.system(f"cd /vision; {build_vars} python3 setup.py bdist_wheel") + wheel_name = list_dir("/vision/dist")[0] + embed_libgomp(f"/vision/dist/{wheel_name}") + + print('Move TorchVision wheel to artfacts') + os.system(f"mv /vision/dist/{wheel_name} /artifacts/") + return wheel_name + + +''' +Build TorchAudio wheel +''' +def build_torchaudio(branch: str = "main", + git_clone_flags: str = "") -> str: + print('Checking out TorchAudio repo') + git_clone_flags += " --recurse-submodules" + build_version = checkout_repo(branch=branch, + url="https://github.com/pytorch/audio", + git_clone_flags=git_clone_flags, + mapping={ + "v1.9.0": ("0.9.0", "rc2"), + "v1.10.0": ("0.10.0", "rc5"), + "v1.10.1": ("0.10.1", "rc1"), + "v1.10.2": ("0.10.2", "rc1"), + "v1.11.0": ("0.11.0", "rc1"), + "v1.12.0": ("0.12.0", "rc3"), + "v1.12.1": ("0.12.1", "rc5"), + "v1.13.0": ("0.13.0", "rc4"), + "v1.13.1": ("0.13.1", "rc2"), + "v2.0.0": ("2.0.0", "rc2"), + }) + print('Building TorchAudio wheel') + build_vars = "CMAKE_SHARED_LINKER_FLAGS=-Wl,-z,max-page-size=0x10000 " + if branch == 'nightly': + version = '' + if os.path.exists('/audio/version.txt'): + version = subprocess.check_output(['cat', '/audio/version.txt']).decode().strip() + build_date = subprocess.check_output(['git','log','--pretty=format:%cs','-1'], cwd='/audio').decode().replace('-','') + build_vars += f"BUILD_VERSION={version}.dev{build_date}" + elif build_version is not None: + build_vars += f"BUILD_VERSION={build_version}" + + os.system(f"cd /audio; {build_vars} python3 setup.py bdist_wheel") + wheel_name = list_dir("/audio/dist")[0] + embed_libgomp(f"/audio/dist/{wheel_name}") + + print('Move TorchAudio wheel to artfacts') + os.system(f"mv /audio/dist/{wheel_name} /artifacts/") + return wheel_name + + +''' +Build TorchText wheel +''' +def build_torchtext(branch: str = "main", + git_clone_flags: str = "") -> str: + print('Checking out TorchText repo') + os.system(f"cd /") + git_clone_flags += " --recurse-submodules" + build_version = checkout_repo(branch=branch, + url="https://github.com/pytorch/text", + git_clone_flags=git_clone_flags, + mapping={ + "v1.9.0": ("0.10.0", "rc1"), + "v1.10.0": ("0.11.0", "rc2"), + "v1.10.1": ("0.11.1", "rc1"), + "v1.10.2": ("0.11.2", "rc1"), + "v1.11.0": ("0.12.0", "rc1"), + "v1.12.0": ("0.13.0", "rc2"), + "v1.12.1": ("0.13.1", "rc5"), + "v1.13.0": ("0.14.0", "rc3"), + "v1.13.1": ("0.14.1", "rc1"), + "v2.0.0": ("0.15.0", "rc2"), + }) + print('Building TorchText wheel') + build_vars = "CMAKE_SHARED_LINKER_FLAGS=-Wl,-z,max-page-size=0x10000 " + if branch == 'nightly': + version = '' + if os.path.exists('/text/version.txt'): + version = subprocess.check_output(['cat', '/text/version.txt']).decode().strip() + build_date = subprocess.check_output(['git','log','--pretty=format:%cs','-1'], cwd='/text').decode().replace('-','') + build_vars += f"BUILD_VERSION={version}.dev{build_date}" + elif build_version is not None: + build_vars += f"BUILD_VERSION={build_version}" + + os.system(f"cd text; {build_vars} python3 setup.py bdist_wheel") + wheel_name = list_dir("/text/dist")[0] + embed_libgomp(f"/text/dist/{wheel_name}") + + print('Move TorchText wheel to artfacts') + os.system(f"mv /text/dist/{wheel_name} /artifacts/") + return wheel_name + + +''' +Build TorchData wheel +''' +def build_torchdata(branch: str = "main", + git_clone_flags: str = "") -> str: + print('Checking out TorchData repo') + git_clone_flags += " --recurse-submodules" + build_version = checkout_repo(branch=branch, + url="https://github.com/pytorch/data", + git_clone_flags=git_clone_flags, + mapping={ + "v1.11.0": ("0.3.0", "rc1"), + "v1.12.0": ("0.4.0", "rc3"), + "v1.12.1": ("0.4.1", "rc5"), + "v1.13.1": ("0.5.1", "rc2"), + "v2.0.0": ("0.6.0", "rc2"), + }) + print('Building TorchData wheel') + build_vars = "CMAKE_SHARED_LINKER_FLAGS=-Wl,-z,max-page-size=0x10000 " + if branch == 'nightly': + version = '' + if os.path.exists('/data/version.txt'): + version = subprocess.check_output(['cat', '/data/version.txt']).decode().strip() + build_date = subprocess.check_output(['git','log','--pretty=format:%cs','-1'], cwd='/data').decode().replace('-','') + build_vars += f"BUILD_VERSION={version}.dev{build_date}" + elif build_version is not None: + build_vars += f"BUILD_VERSION={build_version}" + + os.system(f"cd /data; {build_vars} python3 setup.py bdist_wheel") + wheel_name = list_dir("/data/dist")[0] + embed_libgomp(f"/data/dist/{wheel_name}") + + print('Move TorchAudio wheel to artfacts') + os.system(f"mv /data/dist/{wheel_name} /artifacts/") + return wheel_name + + +def parse_arguments(): + from argparse import ArgumentParser + parser = ArgumentParser("AARCH64 wheels python CD") + parser.add_argument("--debug", action="store_true") + parser.add_argument("--build-only", action="store_true") + parser.add_argument("--test-only", type=str) + parser.add_argument("--enable-mkldnn", action="store_true") + return parser.parse_args() + + +''' +Entry Point +''' +if __name__ == '__main__': + + args = parse_arguments() + enable_mkldnn = args.enable_mkldnn + os.system("cd /pytorch") + branch = subprocess.check_output("git rev-parse --abbrev-ref HEAD") + + git_clone_flags = " --depth 1 --shallow-submodules" + os.system(f"conda install -y ninja scons") + + print("Build and Install OpenBLAS") + build_OpenBLAS(git_clone_flags) + + print('Building PyTorch wheel') + build_vars = "CMAKE_SHARED_LINKER_FLAGS=-Wl,-z,max-page-size=0x10000 " + os.system(f"cd /pytorch; pip install -r requirements.txt") + os.system(f"pip install auditwheel") + os.system(f"python setup.py clean") + + if branch == 'nightly' or branch == 'master': + build_date = subprocess.check_output(['git','log','--pretty=format:%cs','-1'], cwd='/pytorch').decode().replace('-','') + version = subprocess.check_output(['cat','version.txt'], cwd='/pytorch').decode().strip()[:-2] + build_vars += f"BUILD_TEST=0 PYTORCH_BUILD_VERSION={version}.dev{build_date} PYTORCH_BUILD_NUMBER=1" + if branch.startswith("v1.") or branch.startswith("v2."): + build_vars += f"BUILD_TEST=0 PYTORCH_BUILD_VERSION={branch[1:branch.find('-')]} PYTORCH_BUILD_NUMBER=1" + if enable_mkldnn: + build_ArmComputeLibrary(git_clone_flags) + print("build pytorch with mkldnn+acl backend") + os.system(f"export ACL_ROOT_DIR=/acl; export LD_LIBRARY_PATH=/acl/build; export ACL_LIBRARY=/acl/build") + build_vars += " USE_MKLDNN=ON USE_MKLDNN_ACL=ON" + os.system(f"cd /pytorch; {build_vars} python3 setup.py bdist_wheel") + print('Repair the wheel') + pytorch_wheel_name = list_dir("pytorch/dist")[0] + os.system(f"export LD_LIBRARY_PATH=/pytorch/build/lib:$LD_LIBRARY_PATH; auditwheel repair /pytorch/dist/{pytorch_wheel_name}") + print('replace the original wheel with the repaired one') + pytorch_repaired_wheel_name = list_dir("wheelhouse")[0] + os.system(f"cp /wheelhouse/{pytorch_repaired_wheel_name} /pytorch/dist/{pytorch_wheel_name}") + else: + print("build pytorch without mkldnn backend") + os.system(f"cd pytorch ; {build_vars} python3 setup.py bdist_wheel") + + print("Deleting build folder") + os.system("cd /pytorch; rm -rf build") + pytorch_wheel_name = list_dir("/pytorch/dist")[0] + embed_libgomp(f"/pytorch/dist/{pytorch_wheel_name}") + print('Move PyTorch wheel to artfacts') + os.system(f"mv /pytorch/dist/{pytorch_wheel_name} /artifacts/") + print("Installing Pytorch wheel") + os.system(f"pip install /artifacts/{pytorch_wheel_name}") + + vision_wheel_name = build_torchvision(branch=branch, git_clone_flags=git_clone_flags) + audio_wheel_name = build_torchaudio(branch=branch, git_clone_flags=git_clone_flags) + text_wheel_name = build_torchtext(branch=branch, git_clone_flags=git_clone_flags) + data_wheel_name = build_torchdata(branch=branch, git_clone_flags=git_clone_flags) + + print(f"Wheels Created:\n" \ + f"{pytorch_wheel_name}\n" \ + f"{vision_wheel_name}\n" \ + f"{audio_wheel_name}\n" \ + f"{text_wheel_name}\n" \ + f"{data_wheel_name}\n") diff --git a/build_aarch64_wheel.py b/aarch64_linux/build_aarch64_wheel.py similarity index 100% rename from build_aarch64_wheel.py rename to aarch64_linux/build_aarch64_wheel.py diff --git a/aarch64_linux/embed_library.py b/aarch64_linux/embed_library.py new file mode 100644 index 000000000..978970d45 --- /dev/null +++ b/aarch64_linux/embed_library.py @@ -0,0 +1,72 @@ +#!/usr/bin/env python3 + +from auditwheel.patcher import Patchelf +from auditwheel.wheeltools import InWheelCtx +from auditwheel.elfutils import elf_file_filter +from auditwheel.repair import copylib +from auditwheel.lddtree import lddtree +from subprocess import check_call +import os +import shutil +import sys +from tempfile import TemporaryDirectory + + +def replace_tag(filename): + with open(filename, 'r') as f: + lines = f.read().split("\\n") + for i,line in enumerate(lines): + if not line.startswith("Tag: "): + continue + lines[i] = line.replace("-linux_", "-manylinux2014_") + print(f'Updated tag from {line} to {lines[i]}') + + with open(filename, 'w') as f: + f.write("\\n".join(lines)) + + +class AlignedPatchelf(Patchelf): + def set_soname(self, file_name: str, new_soname: str) -> None: + check_call(['patchelf', '--page-size', '65536', '--set-soname', new_soname, file_name]) + + def replace_needed(self, file_name: str, soname: str, new_soname: str) -> None: + check_call(['patchelf', '--page-size', '65536', '--replace-needed', soname, new_soname, file_name]) + + +def embed_library(whl_path, lib_soname, update_tag=False): + patcher = AlignedPatchelf() + out_dir = TemporaryDirectory() + whl_name = os.path.basename(whl_path) + tmp_whl_name = os.path.join(out_dir.name, whl_name) + with InWheelCtx(whl_path) as ctx: + torchlib_path = os.path.join(ctx._tmpdir.name, 'torch', 'lib') + ctx.out_wheel=tmp_whl_name + new_lib_path, new_lib_soname = None, None + for filename, elf in elf_file_filter(ctx.iter_files()): + if not filename.startswith('torch/lib'): + continue + libtree = lddtree(filename) + if lib_soname not in libtree['needed']: + continue + lib_path = libtree['libs'][lib_soname]['path'] + if lib_path is None: + print(f"Can't embed {lib_soname} as it could not be found") + break + if lib_path.startswith(torchlib_path): + continue + + if new_lib_path is None: + new_lib_soname, new_lib_path = copylib(lib_path, torchlib_path, patcher) + patcher.replace_needed(filename, lib_soname, new_lib_soname) + print(f'Replacing {lib_soname} with {new_lib_soname} for {filename}') + if update_tag: + # Add manylinux2014 tag + for filename in ctx.iter_files(): + if os.path.basename(filename) != 'WHEEL': + continue + replace_tag(filename) + shutil.move(tmp_whl_name, whl_path) + + +if __name__ == '__main__': + embed_library(sys.argv[1], 'libgomp.so.1', len(sys.argv) > 2 and sys.argv[2] == '--update-tag')