Change the logic of CUDA detection

Author: FloopCZ

tensorflow_cc/cmake/build_tensorflow.sh.in

@@ -32,38 +32,16 @@ export NCCL_INSTALL_PATH=${NCCL_INSTALL_PATH:-/usr}
 export PYTHON_BIN_PATH=${PYTHON_BIN_PATH:-"$(which python3)"}
 export PYTHON_LIB_PATH="$($PYTHON_BIN_PATH -c 'import site; print(site.getsitepackages()[0])')"
 
-# configure cuda environmental variables
-
-if [ -e /opt/cuda ]; then
-    echo "Using CUDA from /opt/cuda"
-    export CUDA_TOOLKIT_PATH=/opt/cuda
-elif [ -e /usr/local/cuda ]; then
-    echo "Using CUDA from /usr/local/cuda"
-    export CUDA_TOOLKIT_PATH=/usr/local/cuda
-fi
-
-if [ -e /opt/cuda/include/cudnn.h ]; then
-    echo "Using CUDNN from /opt/cuda"
-    export CUDNN_INSTALL_PATH=/opt/cuda
-elif [ -e /usr/local/cuda/include/cudnn.h ]; then
-    echo "Using CUDNN from /usr/local/cuda"
-    export CUDNN_INSTALL_PATH=/usr/local/cuda
-elif [ -e /usr/include/cudnn.h ]; then
-    echo "Using CUDNN from /usr"
-    export CUDNN_INSTALL_PATH=/usr
-fi
-
-if [ "@ALLOW_CUDA@" = "ON" ] && [ -n "${CUDA_TOOLKIT_PATH}" ]; then
-    if [[ -z "${CUDNN_INSTALL_PATH}" ]]; then
-        echo "CUDA found but no cudnn.h found. Please install cuDNN."
-        exit 1
-    fi
+# check if cuda support requested and supported
+if [ "@ALLOW_CUDA@" = "ON" ] && hash nvcc 2>/dev/null; then
     echo "CUDA support enabled"
     cuda_config_opts="--config=cuda"
     export TF_NEED_CUDA=1
     export TF_CUDA_COMPUTE_CAPABILITIES=${TF_CUDA_COMPUTE_CAPABILITIES:-"3.5,7.0"}  # default from configure.py
-    export TF_CUDA_VERSION="$($CUDA_TOOLKIT_PATH/bin/nvcc --version | sed -n 's/^.*release \(.*\),.*/\1/p')"
-    export TF_CUDNN_VERSION="$(sed -n 's/^#define CUDNN_MAJOR\s*\(.*\).*/\1/p' $CUDNN_INSTALL_PATH/include/cudnn.h)"
+    export TF_CUDA_PATHS=${TF_CUDA_PATHS:-"/opt/cuda,/usr/local/cuda,/usr/local,/usr/cuda,/usr"}
+    export TF_CUDA_VERSION="$(nvcc --version | sed -n 's/^.*release \(.*\),.*/\1/p')"
+    export TF_NCCL_VERSION="$(find / -name 'libnccl.so.*' | tail -n1 | sed -r 's/^.*\.so\.//')"
+    export TF_CUDNN_VERSION="$(find / -name 'libcudnn.so.*' | tail -n1 | sed -r 's/^.*\.so\.//')"
 
     # choose the right version of CUDA compiler
     if [ -z "$GCC_HOST_COMPILER_PATH" ]; then