diff options
Diffstat (limited to 'third_party/gpus/cuda/cuda_config.sh')
| -rwxr-xr-x | third_party/gpus/cuda/cuda_config.sh | 169 |
1 files changed, 169 insertions, 0 deletions
diff --git a/third_party/gpus/cuda/cuda_config.sh b/third_party/gpus/cuda/cuda_config.sh new file mode 100755 index 0000000000..034298590b --- /dev/null +++ b/third_party/gpus/cuda/cuda_config.sh @@ -0,0 +1,169 @@ +#!/bin/bash + +# A simple script to configure the Cuda tree needed for the TensorFlow GPU +# build. We need both Cuda toolkit 7.0 and Cudnn 6.5. +# Useage: +# * User edit cuda.config to point both Cuda toolkit and Cudnn libraries to their local path +# * run cuda_config.sh to generate symbolic links in the source tree to reflect +# * the file organizations needed by TensorFlow. + +print_usage() { +cat << EOF +Usage: $0 [--check] + Configure TensorFlow's canonical view of Cuda libraries using cuda.config. +Arguments: + --check: Only check that the proper Cuda dependencies has already been + properly configured in the source tree. It also creates symbolic links to + the files in the gen-tree to make bazel happy. +EOF +} + +CHECK_ONLY=0 +# Parse the arguments. Add more arguments as the "case" line when needed. +while [[ $# > 0 ]]; do + argument="$1" + shift + case $argument in + --check) + CHECK_ONLY=1 + ;; + *) + echo "Error: unknown arguments" + print_usage + exit -1 + ;; + esac +done + +source cuda.config || exit -1 + +OUTPUTDIR=${OUTPUTDIR:-../../..} +CUDA_TOOLKIT_PATH=${CUDA_TOOLKIT_PATH:-/usr/local/cuda} +CUDNN_INSTALL_PATH=${CUDNN_INSTALL_PATH:-/usr/local/cuda} + +# An error message when the Cuda toolkit is not found +function CudaError { + echo ERROR: $1 +cat << EOF +############################################################################## +############################################################################## +Cuda 7.0 toolkit is missing. +1. Download and install the CUDA 7.0 toolkit and CUDNN 6.5 library; +2. Run configure from the root of the source tree, before rerunning bazel; +Please refer to README.md for more details. +############################################################################## +############################################################################## +EOF + exit -1 +} + +# An error message when CUDNN is not found +function CudnnError { + echo ERROR: $1 +cat << EOF +############################################################################## +############################################################################## +Cudnn 6.5 is missing. +1. Download and install the CUDA 7.0 toolkit and CUDNN 6.5 library; +2. Run configure from the root of the source tree, before rerunning bazel; +Please refer to README.md for more details. +############################################################################## +############################################################################## +EOF + exit -1 +} + +# Check that Cuda libraries has already been properly configured in the source tree. +# We still need to create links to the gen-tree to make bazel happy. +function CheckAndLinkToSrcTree { + ERROR_FUNC=$1 + FILE=$2 + if test ! -e $FILE; then + $ERROR_FUNC "$PWD/$FILE cannot be found" + fi + + # Link the output file to the source tree, avoiding self links if they are + # the same. This could happen if invoked from the source tree by accident. + if [ ! `readlink -f $PWD` == `readlink -f $OUTPUTDIR/third_party/gpus/cuda` ]; then + mkdir -p `dirname $OUTPUTDIR/third_party/gpus/cuda/$FILE` + ln -sf $PWD/$FILE $OUTPUTDIR/third_party/gpus/cuda/$FILE + fi +} + +if [ "$CHECK_ONLY" == "1" ]; then + CheckAndLinkToSrcTree CudaError include/cuda.h + CheckAndLinkToSrcTree CudaError include/cublas.h + CheckAndLinkToSrcTree CudnnError include/cudnn.h + CheckAndLinkToSrcTree CudaError lib64/libcudart_static.a + CheckAndLinkToSrcTree CudaError lib64/libcublas.so.7.0 + CheckAndLinkToSrcTree CudnnError lib64/libcudnn.so.6.5 + CheckAndLinkToSrcTree CudaError lib64/libcudart.so.7.0 + exit 0 +fi + +# Actually configure the source tree for TensorFlow's canonical view of Cuda +# libraries. + +if test ! -e ${CUDA_TOOLKIT_PATH}/lib64/libcudart.so.7.0; then + CudaError "cannot find ${CUDA_TOOLKIT_PATH}/lib64/libcudart.so.7.0" +fi + +if test ! -d ${CUDNN_INSTALL_PATH}; then + CudnnError "cannot find dir: ${CUDNN_INSTALL_PATH}" +fi + +# Locate cudnn.h +if test -e ${CUDNN_INSTALL_PATH}/cudnn.h; then + CUDNN_HEADER_PATH=${CUDNN_INSTALL_PATH} +elif test -e ${CUDNN_INSTALL_PATH}/include/cudnn.h; then + CUDNN_HEADER_PATH=${CUDNN_INSTALL_PATH}/include +else + CudnnError "cannot find cudnn.h under: ${CUDNN_INSTALL_PATH}" +fi + +# Locate libcudnn.so.6.5 +if test -e ${CUDNN_INSTALL_PATH}/libcudnn.so.6.5; then + CUDNN_LIB_PATH=${CUDNN_INSTALL_PATH} +elif test -e ${CUDNN_INSTALL_PATH}/lib64/libcudnn.so.6.5; then + CUDNN_LIB_PATH=${CUDNN_INSTALL_PATH}/lib64 +else + CudnnError "cannot find libcudnn.so.6.5 under: ${CUDNN_INSTALL_PATH}" +fi + +# Helper function to build symbolic links for all files under a directory. +function LinkOneDir { + SRC_PREFIX=$1 + DST_PREFIX=$2 + SRC_DIR=$3 + DST_DIR=`echo $SRC_DIR | sed "s,^$SRC_PREFIX,$DST_PREFIX,"` + mkdir -p $DST_DIR + FILE_LIST=`find -L $SRC_DIR -maxdepth 1 -type f` + if test "$FILE_LIST" != ""; then + ln -sf $FILE_LIST $DST_DIR/ || exit -1 + fi +} +export -f LinkOneDir + +# Build links for all files under the directory, including subdirectoreis. +function LinkAllFiles { + SRC_DIR=$1 + DST_DIR=$2 + find -L $SRC_DIR -type d | xargs -I {} bash -c "LinkOneDir $SRC_DIR $DST_DIR {}" || exit -1 +} + +# Set up the symbolic links for cuda toolkit. We link at individual file level, +# not at the directory level. +# This is because the external library may have different file layout from our desired structure. +mkdir -p $OUTPUTDIR/third_party/gpus/cuda +echo "Setting up Cuda include" +LinkAllFiles ${CUDA_TOOLKIT_PATH}/include $OUTPUTDIR/third_party/gpus/cuda/include || exit -1 +echo "Setting up Cuda lib64" +LinkAllFiles ${CUDA_TOOLKIT_PATH}/lib64 $OUTPUTDIR/third_party/gpus/cuda/lib64 || exit -1 +echo "Setting up Cuda bin" +LinkAllFiles ${CUDA_TOOLKIT_PATH}/bin $OUTPUTDIR/third_party/gpus/cuda/bin || exit -1 +echo "Setting up Cuda nvvm" +LinkAllFiles ${CUDA_TOOLKIT_PATH}/nvvm $OUTPUTDIR/third_party/gpus/cuda/nvvm || exit -1 + +# Set up symbolic link for cudnn +ln -sf $CUDNN_HEADER_PATH/cudnn.h $OUTPUTDIR/third_party/gpus/cuda/include/cudnn.h || exit -1 +ln -sf $CUDNN_LIB_PATH/libcudnn.so.6.5 $OUTPUTDIR/third_party/gpus/cuda/lib64/libcudnn.so.6.5 || exit -1 |