[pull] master from pytorch:master

common/install_cuda.sh

@@ -131,6 +131,27 @@ function install_102 {
 
 }
 
+function install_110 {
+    echo "Installing CUDA 11.0 and CuDNN"
+    rm -rf /usr/local/cuda-11.0 /usr/local/cuda
+    # # install CUDA 11.0 in the same container
+    wget -q http://developer.download.nvidia.com/compute/cuda/11.0.1/local_installers/cuda_11.0.1_450.36.06_linux.run
+    chmod +x cuda_11.0.1_450.36.06_linux.run
+    ./cuda_11.0.1_450.36.06_linux.run --toolkit --silent
+    rm -f cuda_11.0.1_450.36.06_linux.run
+    rm -f /usr/local/cuda && ln -s /usr/local/cuda-11.0 /usr/local/cuda
+
+    # install CUDA 11.0 CuDNN
+    # cuDNN license: https://developer.nvidia.com/cudnn/license_agreement
+    mkdir tmp_cudnn && cd tmp_cudnn
+    wget -q http://developer.download.nvidia.com/compute/redist/cudnn/v8.0.0/cudnn-11.0-linux-x64-v8.0.0.180.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"
@@ -247,6 +268,34 @@ function prune_102 {
 
 }
 
+function prune_110 {
+    echo "Pruning CUDA 11.0 and CuDNN"
+    #####################################################################################
+    # CUDA 11.0 prune static libs
+    #####################################################################################
+    export NVPRUNE="/usr/local/cuda-11.0/bin/nvprune"
+    export CUDA_LIB_DIR="/usr/local/cuda-11.0/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"
+    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"
+
+
+    # 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.0 prune visual tools
+    #####################################################################################
+    export CUDA_BASE="/usr/local/cuda-11.0/"
+    rm -rf $CUDA_BASE/libnsight $CUDA_BASE/libnvvp $CUDA_BASE/nsightee_plugins $CUDA_BASE/nsight-compute-2020.1.0 $CUDA_BASE/nsight-systems-2020.2.5
+}
 
 # idiomatic parameter and option handling in sh
 while test $# -gt 0
@@ -260,6 +309,8 @@ do
 		;;
 	10.2) install_102; prune_102
 		;;
+	11.0) install_110; prune_110
+		;;
 	*) echo "bad argument $1"; exit 1
 	   ;;
     esac

conda/Dockerfile

@@ -46,6 +46,9 @@ RUN bash ./install_cuda.sh 10.1
 FROM cuda as cuda102
 RUN bash ./install_cuda.sh 10.2
 
+FROM cuda as cuda110
+RUN bash ./install_cuda.sh 11.0
+
 # Instal MNIST test data
 FROM base as mnist
 ADD ./common/install_mnist.sh install_mnist.sh
@@ -59,6 +62,7 @@ COPY --from=cuda92   /usr/local/cuda-9.2  /usr/local/cuda-9.2
 COPY --from=cuda100  /usr/local/cuda-10.0 /usr/local/cuda-10.0
 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
 ADD ./java/jni.h     /usr/local/include/jni.h
 ENV PATH /opt/conda/bin:$PATH
 COPY --from=mnist  /usr/local/mnist /usr/local/mnist