feat(aarch64/whl): support aarch64 whl

GitOrigin-RevId: 656a27d62b

CMakeLists.txt

@@ -808,7 +808,7 @@ if(MGE_ARCH STREQUAL "x86_64" OR MGE_ARCH STREQUAL "i386")
     endif()
 endif()
 # dotprod is not enable by default on APPLE, cpuinfo has some problem on APPLE
-if(NOT APPLE)
+if(NOT APPLE AND ${CMAKE_C_COMPILER_ID} STREQUAL "Clang")
     CHECK_CXX_COMPILER_FLAG("-march=armv8.2-a+dotprod" CXX_COMPILER_SUPPORT_DOT)
     if(CXX_COMPILER_SUPPORT_DOT)
         message(STATUS "Enable dotprod feature in armv8.2-a using MGB_ENABLE_DOT")

scripts/whl/BUILD_PYTHON_WHL_README.md

@@ -10,6 +10,15 @@
 1: please refer to: https://docs.docker.com/engine/security/rootless/ to enable rootless docker env
 2: cd ./scripts/whl/manylinux2014
 3: ./build_image.sh
+4: as aarch64-linux python3.5 pip do not provide megengine depends prebuild binary package, which definition
+    in imperative/python/requires.txt, before install megengine wheel package(just python3.5), you need
+    do follow step to init python3.5 arrow build env:
+        4a: wget https://apache.bintray.com/arrow/debian/apache-arrow-archive-keyring-latest-buster.deb
+        4b: sudo apt install ./apache-arrow-archive-keyring-latest-buster.deb
+        4c: sudo apt update
+        4d: sudo apt install libarrow-dev
+        4c: sudo apt install libarrow-python-dev
+        4e: sudo ln -s /usr/include/locale.h /usr/include/xlocale.h
 
 ```
 

scripts/whl/macos/macos_build_whl.sh

@@ -41,6 +41,8 @@ FULL_PYTHON_VER="3.5.9 3.6.10 3.7.7 3.8.3"
 if [[ -z ${ALL_PYTHON} ]]
 then
     ALL_PYTHON=${FULL_PYTHON_VER}
+else
+    check_python_version_is_valid "${ALL_PYTHON}" "${FULL_PYTHON_VER}"
 fi
 
 PYTHON_DIR=

scripts/whl/manylinux2014/Dockerfile

@@ -11,5 +11,5 @@ COPY setup_mirror.sh .
 RUN ./setup_mirror.sh "$platform"
 
 ADD init_image.sh /tmp
-RUN /tmp/init_image.sh && rm -f /tmp/init_image.sh
+RUN /tmp/init_image.sh x86_64 && rm -f /tmp/init_image.sh
 

scripts/whl/manylinux2014/Dockerfile_aarch64

@@ -0,0 +1,11 @@
+FROM quay.io/pypa/manylinux2014_aarch64:2020-12-31-56195b3
+
+ENV UID=1024 PATH=${PATH}:/usr/local/cuda/bin
+
+ARG platform
+COPY setup_mirror.sh .
+RUN ./setup_mirror.sh "$platform"
+
+ADD init_image.sh /tmp
+RUN /tmp/init_image.sh aarch64 && rm -f /tmp/init_image.sh
+

scripts/whl/manylinux2014/build_image.sh

@@ -1,5 +1,9 @@
 #!/bin/bash -e
 
-cd $(dirname $0)
+SRC_DIR=$(readlink -f "`dirname $0`/../../../")
+source ${SRC_DIR}/scripts/whl/utils/utils.sh
+config_docker_file
 
-docker build -t env_manylinux2014:latest .
+cd $(dirname $0)
+echo "docker_file is ${docker_file}"
+docker build -t env_manylinux2014:latest -f ${docker_file} .

scripts/whl/manylinux2014/build_wheel_common.sh

@@ -10,23 +10,48 @@ CUDNN_LIB_DIR="/opt/cudnn/lib64/"
 CUDA_LIB_DIR="/usr/local/cuda/lib64/"
 
 SDK_NAME="unknown"
+x86_64_support_version="cu101 cu111 cu112 cpu"
+aarch64_support_version="cu111 cpu"
 function usage() {
-    echo "use '-sdk cu111' to specify cuda toolkit config, also support cu101, cu112, cpu"
+    echo "use -sdk sdk_version to specify sdk toolkit config!"
+    echo "now x86_64 sdk_version support ${x86_64_support_version}"
+    echo "now aarch64 sdk_version support ${aarch64_support_version}"
 }
 
 while [ "$1" != "" ]; do
     case $1 in
         -sdk)
-           shift
-           SDK_NAME=$1
-           shift
-           ;;
+            shift
+            SDK_NAME=$1
+            shift
+            ;;
         *)
             usage
-            exit 1
+            exit -1
     esac
 done
 
+is_valid_sdk="false"
+all_sdk=""
+machine=$(uname -m)
+case ${machine} in
+    x86_64) all_sdk=${x86_64_support_version} ;;
+    aarch64) all_sdk=${aarch64_support_version} ;;
+    *) echo "nonsupport env!!!";exit -1 ;;
+esac
+
+for i_sdk in ${all_sdk}
+do
+    if [ ${i_sdk} == ${SDK_NAME} ];then
+        is_valid_sdk="true"
+    fi
+done
+if [ ${is_valid_sdk} == "false" ];then
+    echo "invalid sdk: ${SDK_NAME}"
+    usage
+    exit -1
+fi
+
 echo "Build with ${SDK_NAME}"
 
 if [ $SDK_NAME == "cu101" ];then
@@ -38,6 +63,21 @@ if [ $SDK_NAME == "cu101" ];then
     REQUIR_TENSORRT_VERSION="6.0.1.5" 
     REQUIR_CUBLAS_VERSION="10.2.1.243"
 elif [ $SDK_NAME == "cu111" ];then
+    if [ ${machine} == "aarch64" ];then
+        REQUIR_CUDA_VERSION="11010"
+        REQUIR_CUDNN_VERSION="8.0.5"
+        REQUIR_TENSORRT_VERSION="7.2.1.6"
+        REQUIR_CUBLAS_VERSION="11.3.0.106"
+    elif [ ${machine} == "x86_64" ];then
+        REQUIR_CUDA_VERSION="11010"
+        REQUIR_CUDNN_VERSION="8.0.4"
+        REQUIR_TENSORRT_VERSION="7.2.2.3"
+        REQUIR_CUBLAS_VERSION="11.2.1.74"
+    else
+        echo "no support machine: ${machine}"
+        exit -1
+    fi
+
     CUDA_COPY_LIB_LIST="\
         ${CUDA_LIB_DIR}/libnvrtc.so.11.1:\
         ${CUDA_LIB_DIR}/libcublasLt.so.11:\
@@ -49,6 +89,7 @@ elif [ $SDK_NAME == "cu111" ];then
         ${CUDNN_LIB_DIR}/libcudnn_ops_infer.so.8:\
         ${CUDNN_LIB_DIR}/libcudnn_ops_train.so.8:\
         ${CUDNN_LIB_DIR}/libcudnn.so.8"
+
     EXTRA_CMAKE_FLAG=" -DMGE_WITH_CUDNN_SHARED=ON -DMGE_WITH_CUBLAS_SHARED=ON \
         -DMGE_CUDA_GENCODE=\"-gencode arch=compute_61,code=sm_61 \
         -gencode arch=compute_70,code=sm_70 \
@@ -57,10 +98,6 @@ elif [ $SDK_NAME == "cu111" ];then
         -gencode arch=compute_86,code=sm_86 \
         -gencode arch=compute_86,code=compute_86\" "
 
-    REQUIR_CUDA_VERSION="11010" 
-    REQUIR_CUDNN_VERSION="8.0.4" 
-    REQUIR_TENSORRT_VERSION="7.2.2.3" 
-    REQUIR_CUBLAS_VERSION="11.2.1.74"
 elif [ $SDK_NAME == "cu112" ];then
     CUDA_COPY_LIB_LIST="\
         ${CUDA_LIB_DIR}/libnvrtc.so.11.2:\
@@ -73,6 +110,7 @@ elif [ $SDK_NAME == "cu112" ];then
         ${CUDNN_LIB_DIR}/libcudnn_ops_infer.so.8:\
         ${CUDNN_LIB_DIR}/libcudnn_ops_train.so.8:\
         ${CUDNN_LIB_DIR}/libcudnn.so.8"
+
     EXTRA_CMAKE_FLAG=" -DMGE_WITH_CUDNN_SHARED=ON -DMGE_WITH_CUBLAS_SHARED=ON \
         -DMGE_CUDA_GENCODE=\"-gencode arch=compute_61,code=sm_61 \
         -gencode arch=compute_70,code=sm_70 \
@@ -85,11 +123,13 @@ elif [ $SDK_NAME == "cu112" ];then
     REQUIR_CUDNN_VERSION="8.0.4" 
     REQUIR_TENSORRT_VERSION="7.2.2.3" 
     REQUIR_CUBLAS_VERSION="11.3.1.68"
+
 elif [ $SDK_NAME == "cpu" ];then
     echo "use $SDK_NAME without cuda support"
     BUILD_WHL_CPU_ONLY="ON"
 else
-    echo "no support sdk ${SDK_NAME}, please set by '-sdk cu111'"
+    echo "no support sdk ${SDK_NAME}"
+    usage
     exit -1
 fi
 
@@ -100,7 +140,7 @@ fi
 
 echo ${BASEDIR}
 pushd ${BASEDIR}/third_party >/dev/null
-    ./prepare.sh
+./prepare.sh
 popd >/dev/null
 
 cd ${CWD}
@@ -121,8 +161,6 @@ if [ ${BUILD_WHL_CPU_ONLY} = "OFF" ]; then
     fi
 
     ## YOU SHOULD MODIFY CUDA VERSION AS BELOW WHEN UPGRADE
-    
-
     CUDA_ROOT_DIR_=${CUDA_ROOT_DIR%*/}
     CUDNN_ROOT_DIR_=${CUDNN_ROOT_DIR%*/}
     TENSORRT_ROOT_DIR_=${TENSORRT_ROOT_DIR%*/}
@@ -190,7 +228,7 @@ if [ ${BUILD_WHL_CPU_ONLY} = "OFF" ]; then
     echo CUBLAS_VERSION:$CUBLAS_VERSION
 
     if [ $CUDA_VERSION != $REQUIR_CUDA_VERSION ] ; then
-        echo please check the Environment must use CUDA-10.1 NO.$REQUIR_CUDA_VERSION
+        echo please check the Environment must use CUDA NO.$REQUIR_CUDA_VERSION
         exit -1
     fi
 
@@ -214,6 +252,14 @@ if [[ -z ${BUILD_GCC8} ]];then
     BUILD_GCC8=OFF
 fi
 
+if [ ${machine} == "aarch64" ];then
+    # manylinux on aarch64 gcc9 is: (GCC) 9.3.1 20200408 (Red Hat 9.3.1-2)
+    # which version has issue: 'as' take a long long long time for some dnn kernel!
+    # infact ubuntu gcc version: gcc (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0 is OK
+    echo "force use gcc8 on aarch64 linux"
+    BUILD_GCC8="ON"
+fi
+
 if [ "$BUILD_GCC8" == "ON" ];then
     run_cmd="scl enable devtoolset-8 /home/code/scripts/whl/manylinux2014/do_build_common.sh"
 else
@@ -230,12 +276,12 @@ docker run --rm -it $TMPFS_ARGS \
     -e EXTRA_CMAKE_FLAG="$EXTRA_CMAKE_FLAG" \
     -e CUDA_COPY_LIB_LIST="$CUDA_COPY_LIB_LIST"  \
     -e SDK_NAME="$SDK_NAME"  \
+    -e CUDA_ROOT_DIR="/usr/local/cuda" \
+    -e CUDNN_ROOT_DIR="/opt/cudnn" \
+    -e TRT_ROOT_DIR="/opt/tensorrt" \
     -v ${CUDA_ROOT_DIR}:/usr/local/cuda \
     -v ${CUDNN_ROOT_DIR}:/opt/cudnn \
     -v ${TENSORRT_ROOT_DIR}:/opt/tensorrt \
     -v ${BASEDIR}:/home/code \
     -v ${OUTPUTDIR}:/home/output:rw \
     env_manylinux2014:latest /bin/bash -c "$run_cmd"
-
-
-

scripts/whl/manylinux2014/do_build_common.sh

@@ -62,10 +62,16 @@ function patch_elf_depend_lib_mgb_mge() {
 }
 
 
+SRC_DIR=$(readlink -f "`dirname $0`/../../../")
+source ${SRC_DIR}/scripts/whl/utils/utils.sh
+
+SUPPORT_ALL_VERSION="35m 36m 37m 38"
 ALL_PYTHON=${ALL_PYTHON}
 if [[ -z ${ALL_PYTHON} ]]
 then
-    ALL_PYTHON="35m 36m 37m 38"
+    ALL_PYTHON=${SUPPORT_ALL_VERSION}
+else
+    check_python_version_is_valid "${ALL_PYTHON}" "${SUPPORT_ALL_VERSION}"
 fi
 
 BUILD_WHL_CPU_ONLY=${BUILD_WHL_CPU_ONLY}
@@ -74,9 +80,6 @@ then
     BUILD_WHL_CPU_ONLY="OFF"
 fi
 
-SRC_DIR=$(readlink -f "`dirname $0`/../../../")
-source ${SRC_DIR}/scripts/whl/utils/utils.sh
-
 BUILD_DIR=${SRC_DIR}/build_dir/host/MGE_WITH_CUDA_OFF/MGE_INFERENCE_ONLY_OFF/Release/build/
 if [ ${BUILD_WHL_CPU_ONLY} = "OFF" ]; then
     BUILD_DIR=${SRC_DIR}/build_dir/host/MGE_WITH_CUDA_ON/MGE_INFERENCE_ONLY_OFF/Release/build/

scripts/whl/manylinux2014/init_image.sh

@@ -2,18 +2,28 @@
 
 GET_PIP_URL='https://bootstrap.pypa.io/get-pip.py'
 GET_PIP_URL_35='https://bootstrap.pypa.io/pip/3.5/get-pip.py'
-SWIG_URL='https://downloads.sourceforge.net/project/swig/swig/swig-3.0.12/swig-3.0.12.tar.gz?use_mirror=autoselect'
+SWIG_URL='https://codeload.github.com/swig/swig/tar.gz/refs/tags/rel-3.0.12'
 LLVM_URL='https://github.com/llvm-mirror/llvm/archive/release_60.tar.gz' 
 CLANG_URL='https://github.com/llvm-mirror/clang/archive/release_60.tar.gz'
+NINJA_URL='https://codeload.github.com/ninja-build/ninja/tar.gz/refs/tags/v1.10.0'
 
 
-yum install -y pcre-devel devtoolset-9-libatomic-devel.x86_64 ninja-build
-yum install -y devtoolset-8 devtoolset-8-libatomic-devel.x86_64
+ARCH=$1
+echo "ARCH: ${ARCH}"
+yum install -y pcre-devel devtoolset-9-libatomic-devel.${ARCH}
+yum install -y devtoolset-8 devtoolset-8-libatomic-devel.${ARCH}
 # install a default python3 for cmake PYTHON3_EXECUTABLE_WITHOUT_VERSION
-yum install -y python3
+yum install -y python3 python3-devel
+python3 -m pip install cython
 python3 -m pip install numpy
 
-for ver in 35m 36m 37m 38 
+ALL_PYTHON="35m 36m 37m 38"
+numpy_version="1.18.1"
+if [ ${ARCH} = "aarch64" ];then
+    # numpy do not have 1.18.1 on aarch64 linux, so we use another fix version
+    numpy_version="1.19.5"
+fi
+for ver in ${ALL_PYTHON}
 do
     python_ver=${ver:0:2}
     PIP_URL=${GET_PIP_URL}
@@ -23,14 +33,20 @@ do
     echo "use pip url: ${PIP_URL}"
     curl ${PIP_URL} | /opt/python/cp${python_ver}-cp${ver}/bin/python - \
 	--no-cache-dir --only-binary :all:
-    /opt/python/cp${python_ver}-cp${ver}/bin/pip install \
-	--no-cache-dir --only-binary :all: numpy==1.18.1 setuptools==46.1.3
+    if [ ${ARCH} = "aarch64" ] && [ ${ver} = "35m" ];then
+        # aarch64 linux python3.5 pip do not provide binary package
+        /opt/python/cp${python_ver}-cp${ver}/bin/pip install --no-cache-dir numpy setuptools==46.1.3
+    else
+        /opt/python/cp${python_ver}-cp${ver}/bin/pip install \
+        --no-cache-dir --only-binary :all: numpy==${numpy_version} setuptools==46.1.3
+    fi
 done
 
 pushd /home >/dev/null
     echo "Install swig"
     curl -sSL ${SWIG_URL} | tar xz
-    pushd swig-3.0.12 >/dev/null
+    pushd swig-rel-3.0.12 >/dev/null
+        ./autogen.sh
         mkdir build
        	pushd build >/dev/null
 	    ../configure
@@ -65,6 +81,17 @@ pushd /home >/dev/null
 	popd >/dev/null
     popd >/dev/null
     rm -rf clang-release_60 
+    echo "Install ninja build"
+    curl -sSL ${NINJA_URL} | tar xz
+    pushd ninja-1.10.0 >/dev/null
+        mkdir build
+	pushd build >/dev/null
+            cmake .. -DCMAKE_BUILD_TYPE=Release
+	    make -j$(nproc)
+	    cp ninja /usr/bin/
+	popd >/dev/null
+    popd >/dev/null
+    rm -rf ninja-1.10.0
 popd >/dev/null
 
 pushd /tmp >/dev/null

scripts/whl/utils/utils.sh

@@ -3,6 +3,15 @@ set -e
 
 OS=$(uname -s)
 
+docker_file=""
+function config_docker_file() {
+    case $(uname -m) in
+        x86_64) docker_file=Dockerfile ;;
+        aarch64) docker_file=Dockerfile_aarch64 ;;
+        *) echo "nonsupport env!!!";exit -1 ;;
+    esac
+}
+
 function ninja_dry_run_and_check_increment() {
     if [ $# -eq 3 ]; then
         _BUILD_SHELL=$1
@@ -79,3 +88,29 @@ function check_build_ninja_python_api() {
         exit -1
     fi
 }
+
+function check_python_version_is_valid() {
+    want_build_version=$1
+    support_version=$2
+    if [ $# -eq 2 ]; then
+        ver=$1
+    else
+        echo "err call check_python_version_is_valid"
+        exit -1
+    fi
+    is_valid="false"
+    for i_want_build_version in ${want_build_version}
+    do
+        is_valid="false"
+        for i_support_version in ${support_version}
+        do
+            if [ ${i_want_build_version} == ${i_support_version} ];then
+                is_valid="true"
+            fi
+        done
+        if [ ${is_valid} == "false" ];then
+            echo "invalid build python version : \"${want_build_version}\", now support party of \"${support_version}\""
+            exit -1
+        fi
+    done
+}

scripts/whl/windows/windows_build_whl.sh

@@ -30,6 +30,8 @@ FULL_PYTHON_VER="3.5.4 3.6.8 3.7.7 3.8.3"
 if [[ -z ${ALL_PYTHON} ]]
 then
     ALL_PYTHON=${FULL_PYTHON_VER}
+else
+    check_python_version_is_valid "${ALL_PYTHON}" "${FULL_PYTHON_VER}"
 fi
 
 PYTHON_DIR=