Install CUDA 11.2 in conda-cuda and libtorch containers (#616)

Author: janeyx99

common/install_cuda.sh

@@ -146,6 +146,28 @@ function install_111 {
     ldconfig
 }
 
+function install_112 {
+    echo "Installing CUDA 11.2 and CuDNN"
+    rm -rf /usr/local/cuda-11.2 /usr/local/cuda
+    # install CUDA 11.2 in the same container
+    wget -q https://developer.download.nvidia.com/compute/cuda/11.2.0/local_installers/cuda_11.2.0_460.27.04_linux.run
+    chmod +x cuda_11.2.0_460.27.04_linux.run
+    ./cuda_11.2.0_460.27.04_linux.run --toolkit --silent
+    rm -f cuda_11.2.0_460.27.04_linux.run
+    rm -f /usr/local/cuda && ln -s /usr/local/cuda-11.2 /usr/local/cuda
+
+    # TODO: install CUDA 11.2 CuDNN 8.0.5, currently it's the 11.1 version
+    # cuDNN license: https://developer.nvidia.com/cudnn/license_agreement
+    mkdir tmp_cudnn && cd tmp_cudnn
+    wget -q https://developer.download.nvidia.com/compute/redist/cudnn/v8.0.5/cudnn-11.1-linux-x64-v8.0.5.39.tgz -O cudnn-8.0.tgz
+    tar xf cudnn-8.0.tgz
+    cp -a cuda/include/* /usr/local/cuda/include/
+    cp -a cuda/lib64/* /usr/local/cuda/lib64/
+    cd ..
+    rm -rf tmp_cudnn
+    ldconfig
+}
+
 function prune_92 {
     echo "Pruning CUDA 9.2 and CuDNN"
     #####################################################################################
@@ -307,6 +329,37 @@ function prune_111 {
     rm -rf $CUDA_BASE/libnvvp $CUDA_BASE/nsightee_plugins $CUDA_BASE/nsight-compute-2020.2.1 $CUDA_BASE/nsight-systems-2020.3.4
 }
 
+function prune_112 {
+    echo "Pruning CUDA 11.2 and CuDNN"
+    #####################################################################################
+    # CUDA 11.2 prune static libs
+    #####################################################################################
+    export NVPRUNE="/usr/local/cuda-11.2/bin/nvprune"
+    export CUDA_LIB_DIR="/usr/local/cuda-11.2/lib64"
+
+    export GENCODE="-gencode arch=compute_35,code=sm_35 -gencode arch=compute_50,code=sm_50 -gencode arch=compute_60,code=sm_60 -gencode arch=compute_70,code=sm_70 -gencode arch=compute_75,code=sm_75 -gencode arch=compute_80,code=sm_80 -gencode arch=compute_86,code=sm_86"
+    export GENCODE_CUDNN="-gencode arch=compute_35,code=sm_35 -gencode arch=compute_37,code=sm_37 -gencode arch=compute_50,code=sm_50 -gencode arch=compute_60,code=sm_60 -gencode arch=compute_61,code=sm_61 -gencode arch=compute_70,code=sm_70 -gencode arch=compute_75,code=sm_75 -gencode arch=compute_80,code=sm_80 -gencode arch=compute_86,code=sm_86"
+
+    if [[ -n "$OVERRIDE_GENCODE" ]]; then
+        export GENCODE=$OVERRIDE_GENCODE
+    fi
+
+    # all CUDA libs except CuDNN and CuBLAS (cudnn and cublas need arch 3.7 included)
+    ls $CUDA_LIB_DIR/ | grep "\.a" | grep -v "culibos" | grep -v "cudart" | grep -v "cudnn" | grep -v "cublas" | grep -v "metis"  \
+      | xargs -I {} bash -c \
+		"echo {} && $NVPRUNE $GENCODE $CUDA_LIB_DIR/{} -o $CUDA_LIB_DIR/{}"
+
+    # prune CuDNN and CuBLAS
+    $NVPRUNE $GENCODE_CUDNN $CUDA_LIB_DIR/libcudnn_static.a -o $CUDA_LIB_DIR/libcudnn_static.a
+    $NVPRUNE $GENCODE_CUDNN $CUDA_LIB_DIR/libcublas_static.a -o $CUDA_LIB_DIR/libcublas_static.a
+    $NVPRUNE $GENCODE_CUDNN $CUDA_LIB_DIR/libcublasLt_static.a -o $CUDA_LIB_DIR/libcublasLt_static.a
+
+    #####################################################################################
+    # CUDA 11.2 prune visual tools
+    #####################################################################################
+    export CUDA_BASE="/usr/local/cuda-11.2/"
+    rm -rf $CUDA_BASE/libnvvp $CUDA_BASE/nsightee_plugins $CUDA_BASE/nsight-compute-2020.3.0 $CUDA_BASE/nsight-systems-2020.4.3
+}
 
 # idiomatic parameter and option handling in sh
 while test $# -gt 0
@@ -322,6 +375,8 @@ do
 		;;
     11.1) install_111; prune_111
 		;;
+    11.2) install_112; prune_112
+		;;
 	*) echo "bad argument $1"; exit 1
 	   ;;
     esac

conda/Dockerfile

@@ -49,6 +49,9 @@ RUN bash ./install_cuda.sh 11.0
 FROM cuda as cuda111
 RUN bash ./install_cuda.sh 11.1
 
+FROM cuda as cuda112
+RUN bash ./install_cuda.sh 11.2
+
 # Install MNIST test data
 FROM base as mnist
 ADD ./common/install_mnist.sh install_mnist.sh
@@ -63,6 +66,7 @@ COPY --from=cuda101  /usr/local/cuda-10.1 /usr/local/cuda-10.1
 COPY --from=cuda102  /usr/local/cuda-10.2 /usr/local/cuda-10.2
 COPY --from=cuda110  /usr/local/cuda-11.0 /usr/local/cuda-11.0
 COPY --from=cuda111  /usr/local/cuda-11.1 /usr/local/cuda-11.1
+COPY --from=cuda112  /usr/local/cuda-11.2 /usr/local/cuda-11.2
 ADD ./java/jni.h     /usr/local/include/jni.h
 ENV PATH /opt/conda/bin:$PATH
 COPY --from=mnist  /usr/local/mnist /usr/local/mnist

libtorch/ubuntu16.04/Dockerfile

@@ -38,6 +38,10 @@ FROM cuda as cuda111
 RUN bash ./install_cuda.sh 11.1
 RUN bash ./install_magma.sh 11.1
 
+FROM cuda as cuda112
+RUN bash ./install_cuda.sh 11.2
+RUN bash ./install_magma.sh 11.2
+
 FROM base as final
 # Install patchelf
 ADD ./common/install_patchelf.sh install_patchelf.sh
@@ -54,5 +58,6 @@ COPY --from=cuda101  /usr/local/cuda-10.1 /usr/local/cuda-10.1
 COPY --from=cuda102  /usr/local/cuda-10.2 /usr/local/cuda-10.2
 COPY --from=cuda110  /usr/local/cuda-11.0 /usr/local/cuda-11.0
 COPY --from=cuda111  /usr/local/cuda-11.1 /usr/local/cuda-11.1
+COPY --from=cuda112  /usr/local/cuda-11.2 /usr/local/cuda-11.2
 COPY --from=intel /opt/intel /opt/intel
 ENV CUDA_HOME /usr/local/cuda