Merge pull request #2 from ika-rwth-aachen/feature/arm64

lreiher · web-flow · commit 54ca7d607075 · 2022-11-11T17:17:05.000+01:00
Support arm64
diff --git a/Makefile b/Makefile
@@ -28,11 +28,13 @@ DEFAULT_TF_VERSION := 2.9.2
 TF_VERSIONS := 2.9.2 2.9.1 2.9.0 2.8.3 2.8.2 2.8.1 2.8.0 2.7.4 2.7.3 2.7.2 2.7.1 2.7.0 2.6.5 2.6.4 2.6.3 2.6.2 2.6.1 2.6.0 2.5.3 2.5.2 2.5.1 2.5.0 2.4.4 2.4.3 2.4.2 2.4.1 2.4.0 2.3.4 2.3.3 2.3.2 2.3.1 2.3.0 2.2.3 2.2.2 2.2.1 2.2.0 2.1.4 2.1.3 2.1.2 2.1.1 2.1.0 2.0.4 2.0.3 2.0.2 2.0.1 2.0.0
 DEFAULT_JOBS := $(shell nproc)
 DEFAULT_GPU := 1
+DEFAULT_ARCH := $(shell dpkg --print-architecture)
 
 # arguments
 TF_VERSION := $(if $(TF_VERSION),$(TF_VERSION),$(DEFAULT_TF_VERSION))
 JOBS := $(if $(JOBS),$(JOBS),$(DEFAULT_JOBS))
 GPU := $(if $(GPU),$(GPU),$(DEFAULT_GPU))
+ARCH := $(if $(ARCH),$(ARCH),$(DEFAULT_ARCH))
 
 # variables
 ifeq ($(GPU), 1)
@@ -42,7 +44,7 @@ else
 endif
 OFFICIAL_DEVEL_IMAGE := tensorflow/tensorflow:$(TF_VERSION)-devel$(GPU_POSTFIX)
 OFFICIAL_DEVEL_IMAGES := tensorflow/tensorflow:*-devel$(GPU_POSTFIX)
-CPP_IMAGE := rwthika/tensorflow-cc:$(TF_VERSION)$(GPU_POSTFIX)
+CPP_IMAGE := rwthika/tensorflow-cc:$(TF_VERSION)$(GPU_POSTFIX)-$(ARCH)
 CPP_IMAGES := rwthika/tensorflow-cc:*
 LIBTENSORFLOW_CC_IMAGE := rwthika/tensorflow-cc:$(TF_VERSION)-libtensorflow_cc*
 LIBTENSORFLOW_CC_IMAGES := rwthika/tensorflow-cc:*-libtensorflow_cc*
diff --git a/README.md b/README.md
diff --git a/docker/DEBIAN/control b/docker/DEBIAN/control
@@ -1,5 +1,5 @@
 Package: libtensorflow-cc
 Version: TF_VERSION
-Architecture: amd64
+Architecture: TARGETARCH
 Maintainer: Lennart Reiher <lennart.reiher@rwth-aachen.de>
 Description: TensorFlow C++ Library
diff --git a/docker/Dockerfile b/docker/Dockerfile
@@ -21,8 +21,10 @@
 # --- build stage for building libtensorflow_cc --------------------------------
 ARG TF_VERSION=master
 ARG GPU_POSTFIX=-gpu
-FROM tensorflow/tensorflow:${TF_VERSION}-devel${GPU_POSTFIX} AS build
+ARG TARGETARCH
+FROM tensorflow/tensorflow:${TF_VERSION}-devel${GPU_POSTFIX}-${TARGETARCH} as build
 
+ARG TARGETARCH
 ARG TF_VERSION
 ARG GPU_POSTFIX
 ARG JOBS="auto"
@@ -35,7 +37,7 @@ WORKDIR /tensorflow
 ENV PYTHON_BIN_PATH=/usr/bin/python3
 ENV PYTHON_LIB_PATH=/usr/lib/python3/dist-packages
 ENV TF_NEED_ROCM=0
-ENV TF_CUDA_COMPUTE_CAPABILITIES=5.3,6.0,6.1,7.0,7.2,7.5,8.0,8.6
+ENV TF_CUDA_COMPUTE_CAPABILITIES=5.3,6.0,6.1,7.0,7.2,7.5,8.0,8.6,8.7
 ENV TF_CUDA_CLANG=0
 ENV GCC_HOST_COMPILER_PATH=/usr/bin/gcc
 ENV CC_OPT_FLAGS="-march=native -Wno-sign-compare"
@@ -72,6 +74,7 @@ RUN apt-get install -y autoconf automake libtool curl make g++ unzip && \
 # --- deb-package stage providing libtensorflow-cc.deb -------------------------
 FROM ubuntu:focal as deb-package
 
+ARG TARGETARCH
 ARG TF_VERSION
 ARG GPU_POSTFIX
 
@@ -88,19 +91,27 @@ COPY --from=build /usr/local/lib                            usr/local/lib/protob
 RUN cp -d usr/local/lib/tensorflow/lib* usr/local/lib/protobuf/lib* usr/local/lib/ && \
     rm -rf usr/local/lib/tensorflow usr/local/lib/protobuf
 RUN sed -i "s/TF_VERSION/$TF_VERSION/" DEBIAN/control
+RUN sed -i "s/TARGETARCH/$TARGETARCH/" DEBIAN/control
 
 # build .deb
 WORKDIR /
 RUN dpkg-deb --build --root-owner-group libtensorflow-cc_${TF_VERSION}${GPU_POSTFIX} && \
     rm -rf libtensorflow-cc_${TF_VERSION}${GPU_POSTFIX}/
 
 # --- final stage with TensorFlow Python and C++ Library -----------------------
-FROM tensorflow/tensorflow:${TF_VERSION}${GPU_POSTFIX} as final
+FROM --platform=amd64 tensorflow/tensorflow:${TF_VERSION}${GPU_POSTFIX} as final-amd64
+ARG TARGETARCH
+
+FROM --platform=arm64 nvcr.io/nvidia/l4t-tensorflow:r35.1.0-tf2.9-py3 as final-arm64
+ARG TARGETARCH
+
+FROM "final-${TARGETARCH}" as final
 
 ARG TF_VERSION
 ARG GPU_POSTFIX
 
 # install TensorFlow C++ API incl. protobuf
 COPY --from=deb-package /libtensorflow-cc_${TF_VERSION}${GPU_POSTFIX}.deb /tmp
 RUN dpkg -i /tmp/libtensorflow-cc_${TF_VERSION}${GPU_POSTFIX}.deb && \
-    ldconfig
+    ldconfig && \
+    rm /tmp/libtensorflow-cc_${TF_VERSION}${GPU_POSTFIX}.deb
diff --git a/scripts/.common.sh b/scripts/.common.sh
@@ -24,11 +24,13 @@ set -o pipefail
 DEFAULT_TF_VERSION="2.9.2"
 DEFAULT_JOBS=$(nproc)
 DEFAULT_GPU=1
+DEFAULT_ARCH=$(dpkg --print-architecture)
 
 TF_VERSION=${TF_VERSION:-${DEFAULT_TF_VERSION}}
 JOBS=${JOBS:-${DEFAULT_JOBS}}
 GPU=${GPU:-${DEFAULT_GPU}}
 [[ $GPU == "1" ]] && GPU_POSTFIX="-gpu" || GPU_POSTFIX=""
+ARCH=${ARCH:-${DEFAULT_ARCH}}
 
 SCRIPT_NAME=$(basename "$0")
 SCRIPT_DIR=$(realpath $(dirname "$0"))
@@ -41,6 +43,7 @@ mkdir -p ${LOG_DIR}
 DOWNLOAD_DOCKERFILES_DIR=${DOCKER_DIR}/.Dockerfiles
 DOWNLOAD_DOCKERFILE_DIR=${DOWNLOAD_DOCKERFILES_DIR}/${TF_VERSION}
 
-IMAGE_DEVEL="tensorflow/tensorflow:${TF_VERSION}-devel${GPU_POSTFIX}"
+IMAGE_DEVEL_ARCH="tensorflow/tensorflow:${TF_VERSION}-devel${GPU_POSTFIX}-${ARCH}"
 IMAGE_CPP="rwthika/tensorflow-cc:${TF_VERSION}${GPU_POSTFIX}"
-IMAGE_LIBTENSORFLOW_CC="rwthika/tensorflow-cc:${TF_VERSION}-libtensorflow_cc${GPU_POSTFIX}"
+IMAGE_CPP_ARCH="${IMAGE_CPP}-${ARCH}"
+IMAGE_LIBTENSORFLOW_CC_ARCH="rwthika/tensorflow-cc:${TF_VERSION}-libtensorflow_cc${GPU_POSTFIX}-${ARCH}"
diff --git a/scripts/.versions.run.sh b/scripts/.versions.run.sh
@@ -19,10 +19,10 @@
 # SOFTWARE.
 # ==============================================================================
 
-PYTHON_VERSION=$(python --version 2>&1 | awk '{print $2}')
+PYTHON_VERSION=$(python3 --version 2>&1 | awk '{print $2}')
 
 if [[ $(command -v python) ]]; then
-    TENSORFLOW_PYTHON_VERSION=$(python -c "exec(\"try:\n  import os; os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'; import tensorflow as tf; print(tf.__version__);\n\rexcept ImportError:\n  pass\")" 2> /dev/null)
+    TENSORFLOW_PYTHON_VERSION=$(python3 -c "exec(\"try:\n  import os; os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'; import tensorflow as tf; print(tf.__version__);\n\rexcept ImportError:\n  pass\")" 2> /dev/null)
 fi
 
 if [[ -f /usr/local/include/tensorflow/tensorflow/core/public/version.h ]]; then
@@ -47,10 +47,10 @@ if [[ -f /usr/include/cudnn_version.h ]]; then
     CUDNN_VERSION=$CUDNN_MAJOR.$CUDNN_MINOR.$CUDNN_PATCH
 fi
 
-if [[ -f /usr/include/x86_64-linux-gnu/NvInferVersion.h ]]; then
-    TENSORRT_MAJOR=$(cat /usr/include/x86_64-linux-gnu/NvInferVersion.h | grep "#define NV_TENSORRT_MAJOR" | sed "s/#define NV_TENSORRT_MAJOR //" | sed "s#//.*##" | sed "s/ //")
-    TENSORRT_MINOR=$(cat /usr/include/x86_64-linux-gnu/NvInferVersion.h | grep "#define NV_TENSORRT_MINOR" | sed "s/#define NV_TENSORRT_MINOR //" | sed "s#//.*##" | sed "s/ //")
-    TENSORRT_PATCH=$(cat /usr/include/x86_64-linux-gnu/NvInferVersion.h | grep "#define NV_TENSORRT_PATCH" | sed "s/#define NV_TENSORRT_PATCH //" | sed "s#//.*##" | sed "s/ //")
+if [[ -f /usr/include/$(uname -m)-linux-gnu/NvInferVersion.h ]]; then
+    TENSORRT_MAJOR=$(cat /usr/include/$(uname -m)-linux-gnu/NvInferVersion.h | grep "#define NV_TENSORRT_MAJOR" | sed "s/#define NV_TENSORRT_MAJOR //" | sed "s#//.*##" | sed "s/ //")
+    TENSORRT_MINOR=$(cat /usr/include/$(uname -m)-linux-gnu/NvInferVersion.h | grep "#define NV_TENSORRT_MINOR" | sed "s/#define NV_TENSORRT_MINOR //" | sed "s#//.*##" | sed "s/ //")
+    TENSORRT_PATCH=$(cat /usr/include/$(uname -m)-linux-gnu/NvInferVersion.h | grep "#define NV_TENSORRT_PATCH" | sed "s/#define NV_TENSORRT_PATCH //" | sed "s#//.*##" | sed "s/ //")
     TENSORRT_VERSION=$TENSORRT_MAJOR.$TENSORRT_MINOR.$TENSORRT_PATCH
 fi
 
diff --git a/scripts/1-build-official-devel-image.sh b/scripts/1-build-official-devel-image.sh
@@ -22,8 +22,13 @@
 source  $(dirname "$0")/.common.sh
 
 CPU_GPU_POSTFIX=${GPU_POSTFIX:--cpu}
-DOCKERFILE=${DOWNLOAD_DOCKERFILE_DIR}/dockerfiles/devel${CPU_GPU_POSTFIX}.Dockerfile
+if [ "$ARCH" = "amd64" ]; then
+    DOCKERFILE=${DOWNLOAD_DOCKERFILE_DIR}/dockerfiles/devel${CPU_GPU_POSTFIX}.Dockerfile
+elif [ "$ARCH" = "arm64" ]; then
+    DOCKERFILE=${DOWNLOAD_DOCKERFILE_DIR}/dockerfiles/arm64v8/devel-cpu-arm64v8.Dockerfile
+    sed -i "s/ubuntu:\${UBUNTU_VERSION}/nvcr.io\/nvidia\/l4t-tensorflow:r35.1.0-tf2.9-py3/" $DOCKERFILE
+fi
 BUILD_DIR=${DOWNLOAD_DOCKERFILE_DIR}
 
-echo "Building ${IMAGE_DEVEL} ... "
-docker build -t ${IMAGE_DEVEL} -f ${DOCKERFILE} ${BUILD_DIR} | tee ${LOG_FILE}
+echo "Building ${IMAGE_DEVEL_ARCH} ... "
+docker build -t ${IMAGE_DEVEL_ARCH} -f ${DOCKERFILE} ${BUILD_DIR} | tee ${LOG_FILE}
diff --git a/scripts/2-build-cpp-image.sh b/scripts/2-build-cpp-image.sh
@@ -22,4 +22,5 @@
 source  $(dirname "$0")/.common.sh
 
 echo "Building ${IMAGE_CPP} ... "
-docker build --build-arg TF_VERSION=${TF_VERSION} --build-arg JOBS=${JOBS} --build-arg GPU_POSTFIX=${GPU_POSTFIX} -t ${IMAGE_CPP} ${DOCKER_DIR} | tee ${LOG_FILE}
+docker build --build-arg TARGETARCH=$ARCH --build-arg TF_VERSION=${TF_VERSION} --build-arg JOBS=${JOBS} --build-arg GPU_POSTFIX=${GPU_POSTFIX} -t ${IMAGE_CPP_ARCH} ${DOCKER_DIR} | tee ${LOG_FILE}
+docker tag ${IMAGE_CPP_ARCH} ${IMAGE_CPP}
diff --git a/scripts/3-export-libtensorflow-cc.sh b/scripts/3-export-libtensorflow-cc.sh
@@ -22,14 +22,14 @@
 source  $(dirname "$0")/.common.sh
 
 EXPORT_DIR=${REPOSITORY_DIR}/libtensorflow-cc
-CONTAINER_FILE="libtensorflow-cc_${TF_VERSION}.deb"
-EXPORT_FILE="libtensorflow-cc_${TF_VERSION}${GPU_POSTFIX}.deb"
+CONTAINER_FILE="libtensorflow-cc_${TF_VERSION}${GPU_POSTFIX}.deb"
+EXPORT_FILE="libtensorflow-cc_${TF_VERSION}${GPU_POSTFIX}_${ARCH}.deb"
 STAGE="deb-package"
 
-echo "Building ${IMAGE_LIBTENSORFLOW_CC} ... "
-docker build --build-arg TF_VERSION=${TF_VERSION} --build-arg JOBS=${JOBS} --build-arg GPU_POSTFIX=${GPU_POSTFIX} --target ${STAGE} -t ${IMAGE_LIBTENSORFLOW_CC} ${DOCKER_DIR} | tee ${LOG_FILE}
+echo "Building ${IMAGE_LIBTENSORFLOW_CC_ARCH} ... "
+docker build --build-arg TARGETARCH=$ARCH --build-arg TF_VERSION=${TF_VERSION} --build-arg JOBS=${JOBS} --build-arg GPU_POSTFIX=${GPU_POSTFIX} --target ${STAGE} -t ${IMAGE_LIBTENSORFLOW_CC_ARCH} ${DOCKER_DIR} | tee ${LOG_FILE}
 
 echo "Exporting to $(realpath ${EXPORT_DIR})/${EXPORT_FILE} ... "
-TMP_CONTAINER=$(docker create ${IMAGE_LIBTENSORFLOW_CC})
-docker cp ${TMP_CONTAINER}:/${EXPORT_FILE} ${EXPORT_DIR}/
+TMP_CONTAINER=$(docker create ${IMAGE_LIBTENSORFLOW_CC_ARCH})
+docker cp ${TMP_CONTAINER}:/${CONTAINER_FILE} ${EXPORT_DIR}/${EXPORT_FILE}
 docker rm -v ${TMP_CONTAINER}
diff --git a/scripts/4-test-libtensorflow-cc.sh b/scripts/4-test-libtensorflow-cc.sh
@@ -26,4 +26,11 @@ EXAMPLE_MOUNT="/example"
 CMD="./build-and-run.sh"
 
 echo "Testing libtensorflow_cc in ${IMAGE_CPP} ... "
-docker run --rm --gpus all -v ${EXAMPLE_DIR}:${EXAMPLE_MOUNT} -w ${EXAMPLE_MOUNT} ${IMAGE_CPP} ${CMD} | tee ${LOG_FILE}
+if [[ "$GPU" == "1" && "$ARCH" = "amd64" ]]; then
+    GPU_ARG = "--gpus all"
+elif [[ "$GPU" == "1" && "$ARCH" = "arm64" ]]; then
+    GPU_ARG="--runtime nvidia"
+else
+    GPU_ARG=""
+fi
+docker run --rm ${GPU_ARG} -v ${EXAMPLE_DIR}:${EXAMPLE_MOUNT} -w ${EXAMPLE_MOUNT} ${IMAGE_CPP} ${CMD} | tee ${LOG_FILE}
diff --git a/scripts/5-print-versions.sh b/scripts/5-print-versions.sh
@@ -26,6 +26,13 @@ SCRIPT_RUN=${SCRIPT_DIR}/.versions.run.sh
 SCRIPT_MOUNT=/.versions.sh
 CMD="bash ${SCRIPT_MOUNT}"
 
-echo "Getting version information from ${IMAGE_DEVEL} and ${IMAGE_CPP} ... "
-docker run --rm --gpus all -v ${SCRIPT_DEVEL}:${SCRIPT_MOUNT} ${IMAGE_DEVEL} ${CMD} | tee ${LOG_FILE}
-docker run --rm --gpus all -v ${SCRIPT_RUN}:${SCRIPT_MOUNT} ${IMAGE_CPP} ${CMD}| tee -a ${LOG_FILE}
+echo "Getting version information from ${IMAGE_DEVEL_ARCH} and ${IMAGE_CPP} ... "
+if [[ "$GPU" == "1" && "$ARCH" = "amd64" ]]; then
+    GPU_ARG = "--gpus all"
+elif [[ "$GPU" == "1" && "$ARCH" = "arm64" ]]; then
+    GPU_ARG="--runtime nvidia"
+else
+    GPU_ARG=""
+fi
+docker run --rm ${GPU_ARG} -v ${SCRIPT_DEVEL}:${SCRIPT_MOUNT} ${IMAGE_DEVEL_ARCH} ${CMD} | tee ${LOG_FILE}
+docker run --rm ${GPU_ARG} -v ${SCRIPT_RUN}:${SCRIPT_MOUNT} ${IMAGE_CPP} ${CMD}| tee -a ${LOG_FILE}
