diff options
Diffstat (limited to 'third_party/gpus')
| -rw-r--r-- | third_party/gpus/cuda/BUILD | 36 | ||||
| -rwxr-xr-x | third_party/gpus/cuda/cuda_config.sh | 32 |
2 files changed, 36 insertions, 32 deletions
diff --git a/third_party/gpus/cuda/BUILD b/third_party/gpus/cuda/BUILD index 690728c5d2..47adb1d9d0 100644 --- a/third_party/gpus/cuda/BUILD +++ b/third_party/gpus/cuda/BUILD @@ -1,6 +1,10 @@ licenses(["restricted"]) # MPL2, portions GPL v3, LGPL v3, BSD-like load("/tensorflow/tensorflow", "if_cuda") +load("//tensorflow/core:platform/default/build_config.bzl", + "tf_get_cuda_version", + "tf_get_cudnn_version", + ) package(default_visibility = ["//visibility:public"]) @@ -53,10 +57,10 @@ cc_library( cc_library( name = "cudart", srcs = [ - "lib64/libcudart.so.7.0", + "lib64/libcudart.so." + tf_get_cuda_version(), ], data = [ - "lib64/libcudart.so.7.0", + "lib64/libcudart.so." + tf_get_cuda_version(), ], includes = ["include/"], visibility = ["//visibility:public"], @@ -66,10 +70,10 @@ cc_library( cc_library( name = "cublas", srcs = [ - "lib64/libcublas.so.7.0", + "lib64/libcublas.so." + tf_get_cuda_version(), ], data = [ - "lib64/libcublas.so.7.0", + "lib64/libcublas.so." + tf_get_cuda_version(), ], includes = ["include/"], visibility = ["//visibility:public"], @@ -79,10 +83,10 @@ cc_library( cc_library( name = "cudnn", srcs = [ - "lib64/libcudnn.so.6.5", + "lib64/libcudnn.so." + tf_get_cudnn_version(), ], data = [ - "lib64/libcudnn.so.6.5", + "lib64/libcudnn.so." + tf_get_cudnn_version(), ], includes = ["include/"], visibility = ["//visibility:public"], @@ -92,10 +96,10 @@ cc_library( cc_library( name = "cufft", srcs = [ - "lib64/libcufft.so.7.0", + "lib64/libcufft.so." + tf_get_cuda_version(), ], data = [ - "lib64/libcufft.so.7.0", + "lib64/libcufft.so." + tf_get_cuda_version(), ], includes = ["include/"], visibility = ["//visibility:public"], @@ -130,10 +134,10 @@ genrule( "include/cublas.h", "include/cudnn.h", "lib64/libcudart_static.a", - "lib64/libcublas.so.7.0", - "lib64/libcudnn.so.6.5", - "lib64/libcudart.so.7.0", - "lib64/libcufft.so.7.0", + "lib64/libcublas.so." + tf_get_cuda_version(), + "lib64/libcudnn.so." + tf_get_cudnn_version(), + "lib64/libcudart.so." + tf_get_cuda_version(), + "lib64/libcufft.so." + tf_get_cuda_version(), ], cmd = if_cuda( # Under cuda config, create all the symbolic links to the actual cuda files @@ -147,10 +151,10 @@ genrule( "touch $(@D)/include/cublas.h", "touch $(@D)/include/cudnn.h", "touch $(@D)/lib64/libcudart_static.a", - "touch $(@D)/lib64/libcublas.so.7.0", - "touch $(@D)/lib64/libcudnn.so.6.5", - "touch $(@D)/lib64/libcudart.so.7.0", - "touch $(@D)/lib64/libcufft.so.7.0" + "touch $(@D)/lib64/libcublas.so." + tf_get_cuda_version(), + "touch $(@D)/lib64/libcudnn.so." + tf_get_cudnn_version(), + "touch $(@D)/lib64/libcudart.so." + tf_get_cuda_version(), + "touch $(@D)/lib64/libcufft.so." + tf_get_cuda_version(), ]), ), local = 1, diff --git a/third_party/gpus/cuda/cuda_config.sh b/third_party/gpus/cuda/cuda_config.sh index 44e07433a0..87c35349c0 100755 --- a/third_party/gpus/cuda/cuda_config.sh +++ b/third_party/gpus/cuda/cuda_config.sh @@ -16,7 +16,7 @@ # 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. +# build. We need both Cuda toolkit $TF_CUDA_VERSION and Cudnn $TF_CUDNN_VERSION. # 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 @@ -62,8 +62,8 @@ function CudaError { cat << EOF ############################################################################## ############################################################################## -Cuda 7.0 toolkit is missing. -1. Download and install the CUDA 7.0 toolkit and CUDNN 6.5 library; +Cuda $TF_CUDA_VERSION toolkit is missing. +1. Download and install the CUDA $TF_CUDA_VERSION toolkit and CUDNN $TF_CUDNN_VERSION library; 2. Run configure from the root of the source tree, before rerunning bazel; Please refer to README.md for more details. ############################################################################## @@ -78,8 +78,8 @@ function CudnnError { cat << EOF ############################################################################## ############################################################################## -Cudnn 6.5 is missing. -1. Download and install the CUDA 7.0 toolkit and CUDNN 6.5 library; +Cudnn $TF_CUDNN_VERSION is missing. +1. Download and install the CUDA $TF_CUDA_VERSION toolkit and CUDNN $TF_CUDNN_VERSION library; 2. Run configure from the root of the source tree, before rerunning bazel; Please refer to README.md for more details. ############################################################################## @@ -110,18 +110,18 @@ if [ "$CHECK_ONLY" == "1" ]; then 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 - CheckAndLinkToSrcTree CudaError lib64/libcufft.so.7.0 + CheckAndLinkToSrcTree CudaError lib64/libcublas.so.$TF_CUDA_VERSION + CheckAndLinkToSrcTree CudnnError lib64/libcudnn.so.$TF_CUDNN_VERSION + CheckAndLinkToSrcTree CudaError lib64/libcudart.so.$TF_CUDA_VERSION + CheckAndLinkToSrcTree CudaError lib64/libcufft.so.$TF_CUDA_VERSION 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" +if test ! -e ${CUDA_TOOLKIT_PATH}/lib64/libcudart.so.$TF_CUDA_VERSION; then + CudaError "cannot find ${CUDA_TOOLKIT_PATH}/lib64/libcudart.so.$TF_CUDA_VERSION" fi if test ! -d ${CUDNN_INSTALL_PATH}; then @@ -137,13 +137,13 @@ 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 +# Locate libcudnn.so.${$TF_CUDNN_VERSION} +if test -e ${CUDNN_INSTALL_PATH}/libcudnn.so.$TF_CUDNN_VERSION; then CUDNN_LIB_PATH=${CUDNN_INSTALL_PATH} -elif test -e ${CUDNN_INSTALL_PATH}/lib64/libcudnn.so.6.5; then +elif test -e ${CUDNN_INSTALL_PATH}/lib64/libcudnn.so.$TF_CUDNN_VERSION; then CUDNN_LIB_PATH=${CUDNN_INSTALL_PATH}/lib64 else - CudnnError "cannot find libcudnn.so.6.5 under: ${CUDNN_INSTALL_PATH}" + CudnnError "cannot find libcudnn.so.$TF_CUDNN_VERSION under: ${CUDNN_INSTALL_PATH}" fi # Helper function to build symbolic links for all files under a directory. @@ -182,4 +182,4 @@ LinkAllFiles ${CUDA_TOOLKIT_PATH}/nvvm $OUTPUTDIR/third_party/gpus/cuda/nvvm || # 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 +ln -sf $CUDNN_LIB_PATH/libcudnn.so.$TF_CUDNN_VERSION $OUTPUTDIR/third_party/gpus/cuda/lib64/libcudnn.so.$TF_CUDNN_VERSION || exit -1 |