[aarch64] Adding CI Scripts to build aarch64 wheels

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=<PythonVersion> 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 <YourPemKey> --use-docker --python 3.7 --branch <RCtag>```

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

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