Merge changes from github.

Change: 121586635

3 files changed, 175 insertions, 70 deletions

diff --git a/third_party/gpus/cuda/BUILD b/third_party/gpus/cuda/BUILD
index 28b49b8cc7..792dbb4268 100644
--- a/third_party/gpus/cuda/BUILD
+++ b/third_party/gpus/cuda/BUILD
@@ -1,10 +1,11 @@
 licenses(["restricted"])  # MPL2, portions GPL v3, LGPL v3, BSD-like
 
 load("//tensorflow:tensorflow.bzl", "if_cuda")
-load("//tensorflow/core:platform/default/build_config.bzl",
-     "tf_get_cuda_version",
-     "tf_get_cudnn_version",
-    )
+load("platform", "cuda_library_path")
+load("platform", "cuda_static_library_path")
+load("platform", "cudnn_library_path")
+load("platform", "cupti_library_path")
+load("platform", "readlink_command")
 
 package(default_visibility = ["//visibility:public"])
 
@@ -31,6 +32,12 @@ config_setting(
     },
 )
 
+config_setting(
+    name = "darwin",
+    values = {"cpu": "darwin"},
+    visibility = ["//visibility:public"],
+)
+
 cc_library(
     name = "cuda_headers",
     hdrs = glob([
@@ -43,24 +50,26 @@ cc_library(
 cc_library(
     name = "cudart_static",
     srcs = [
-        "lib64/libcudart_static.a",
+        cuda_static_library_path("cudart"),
     ],
     includes = ["include/"],
     linkopts = [
         "-ldl",
-        "-lrt",
         "-lpthread",
-    ],
+    ] + select({
+        "//tensorflow:darwin": [],
+        "//conditions:default": ["-lrt"]
+    }),
     visibility = ["//visibility:public"],
 )
 
 cc_library(
     name = "cudart",
     srcs = [
-        "lib64/libcudart.so" + tf_get_cuda_version(),
+        cuda_library_path("cudart")
     ],
     data = [
-        "lib64/libcudart.so" + tf_get_cuda_version(),
+        cuda_library_path("cudart")
     ],
     includes = ["include/"],
     visibility = ["//visibility:public"],
@@ -70,10 +79,10 @@ cc_library(
 cc_library(
     name = "cublas",
     srcs = [
-        "lib64/libcublas.so" + tf_get_cuda_version(),
+        cuda_library_path("cublas")
     ],
     data = [
-        "lib64/libcublas.so" + tf_get_cuda_version(),
+        cuda_library_path("cublas")
     ],
     includes = ["include/"],
     visibility = ["//visibility:public"],
@@ -83,10 +92,10 @@ cc_library(
 cc_library(
     name = "cudnn",
     srcs = [
-        "lib64/libcudnn.so" + tf_get_cudnn_version(),
+        cudnn_library_path()
     ],
     data = [
-        "lib64/libcudnn.so" + tf_get_cudnn_version(),
+        cudnn_library_path()
     ],
     includes = ["include/"],
     visibility = ["//visibility:public"],
@@ -96,10 +105,10 @@ cc_library(
 cc_library(
     name = "cufft",
     srcs = [
-        "lib64/libcufft.so" + tf_get_cuda_version(),
+        cuda_library_path("cufft")
     ],
     data = [
-        "lib64/libcufft.so" + tf_get_cuda_version(),
+        cuda_library_path("cufft")
     ],
     includes = ["include/"],
     visibility = ["//visibility:public"],
@@ -130,7 +139,7 @@ cc_library(
 cc_library(
     name = "cupti_dsos",
     data = [
-        "extras/CUPTI/lib64/libcupti.so" + tf_get_cuda_version(),
+        cupti_library_path(),
     ],
     visibility = ["//visibility:public"],
 )
@@ -152,34 +161,34 @@ genrule(
         "include/cublas.h",
         "include/cudnn.h",
         "extras/CUPTI/include/cupti.h",
-        "lib64/libcudart_static.a",
-        "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(),
-        "extras/CUPTI/lib64/libcupti.so" + tf_get_cuda_version(),
+        cuda_static_library_path("cudart"),
+        cuda_library_path("cublas"),
+        cudnn_library_path(),
+        cuda_library_path("cudart"),
+        cuda_library_path("cufft"),
+        cupti_library_path(),
     ],
     cmd = if_cuda(
         # Under cuda config, create all the symbolic links to the actual cuda files
-        "OUTPUTDIR=`readlink -f $(@D)/../../..`; cd `dirname $(location :cuda_config.sh)`; OUTPUTDIR=$$OUTPUTDIR ./cuda_config.sh --check;",
+        "OUTPUTDIR=`{} -f $(@D)/../../..`; cd third_party/gpus/cuda; OUTPUTDIR=$$OUTPUTDIR ./cuda_config.sh --check;".format(readlink_command()),
 
         # Under non-cuda config, create all dummy files to make the build go through
         ";".join([
-          "mkdir -p $(@D)/include",
-          "mkdir -p $(@D)/lib64",
-          "mkdir -p $(@D)/extras/CUPTI/include",
-          "mkdir -p $(@D)/extras/CUPTI/lib64",
-          "touch $(@D)/include/cuda.h",
-          "touch $(@D)/include/cublas.h",
-          "touch $(@D)/include/cudnn.h",
-          "touch $(@D)/extras/CUPTI/include/cupti.h",
-          "touch $(@D)/lib64/libcudart_static.a",
-          "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(),
-          "touch $(@D)/extras/CUPTI/lib64/libcupti.so" + tf_get_cuda_version(),
-            ]),
+           "mkdir -p $(@D)/include",
+           "mkdir -p $(@D)/lib64",
+           "mkdir -p $(@D)/extras/CUPTI/include",
+           "mkdir -p $(@D)/extras/CUPTI/lib64",
+           "touch $(@D)/include/cuda.h",
+           "touch $(@D)/include/cublas.h",
+           "touch $(@D)/include/cudnn.h",
+           "touch $(@D)/extras/CUPTI/include/cupti.h",
+           "touch $(@D)/{}".format(cuda_static_library_path("cudart")),
+           "touch $(@D)/{}".format(cuda_library_path("cublas")),
+           "touch $(@D)/{}".format(cudnn_library_path()),
+           "touch $(@D)/{}".format(cuda_library_path("cudart")),
+           "touch $(@D)/{}".format(cuda_library_path("cufft")),
+           "touch $(@D)/{}".format(cupti_library_path()),
+         ]),
     ),
     local = 1,
 )
@@ -191,7 +200,7 @@ genrule(
     ],
     cmd = if_cuda(
         # Under cuda config, create the symbolic link to the actual cuda.config
-        "configfile=$(location :cuda.config); ln -sf `readlink -f $${configfile#*/*/*/}` $(@D)/;",
+        "configfile=$(location :cuda.config); ln -sf `{} -f $${{configfile#*/*/*/}}` $(@D)/;".format(readlink_command()),
 
         # Under non-cuda config, create the dummy file
         ";".join([
diff --git a/third_party/gpus/cuda/cuda_config.sh b/third_party/gpus/cuda/cuda_config.sh
index e93a7ed741..0e1106bb70 100755
--- a/third_party/gpus/cuda/cuda_config.sh
+++ b/third_party/gpus/cuda/cuda_config.sh
@@ -54,7 +54,19 @@ 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}
+CUDNN_INSTALL_BASEDIR=${CUDNN_INSTALL_PATH:-/usr/local/cuda}
+
+if [[ -z "$TF_CUDA_VERSION" ]]; then
+  TF_CUDA_EXT=""
+else
+  TF_CUDA_EXT=".$TF_CUDA_VERSION"
+fi
+
+if [[ -z "$TF_CUDNN_VERSION" ]]; then
+  TF_CUDNN_EXT=""
+else
+  TF_CUDNN_EXT=".$TF_CUDNN_VERSION"
+fi
 
 # An error message when the Cuda toolkit is not found
 function CudaError {
@@ -99,59 +111,84 @@ function CheckAndLinkToSrcTree {
 
   # 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
+  if [ ! $($READLINK_CMD -f $PWD) == $($READLINK_CMD -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
 }
 
+OSNAME=`uname -s`
+if [ "$OSNAME" == "Linux" ]; then
+  CUDA_LIB_PATH="lib64"
+  CUDA_CUPTI_LIB_DIR="extras/CUPTI/lib64"
+  CUDA_RT_LIB_PATH="lib64/libcudart.so${TF_CUDA_EXT}"
+  CUDA_RT_LIB_STATIC_PATH="lib64/libcudart_static.a"
+  CUDA_BLAS_LIB_PATH="lib64/libcublas.so${TF_CUDA_EXT}"
+  CUDA_DNN_LIB_PATH="lib64/libcudnn.so${TF_CUDNN_EXT}"
+  CUDA_DNN_LIB_ALT_PATH="libcudnn.so${TF_CUDNN_EXT}"
+  CUDA_FFT_LIB_PATH="lib64/libcufft.so${TF_CUDA_EXT}"
+  CUDA_CUPTI_LIB_PATH="extras/CUPTI/lib64/libcupti.so${TF_CUDA_EXT}"
+  READLINK_CMD="readlink"
+elif [ "$OSNAME" == "Darwin" ]; then
+  CUDA_LIB_PATH="lib"
+  CUDA_CUPTI_LIB_DIR="extras/CUPTI/lib"
+  CUDA_RT_LIB_PATH="lib/libcudart${TF_CUDA_EXT}.dylib"
+  CUDA_RT_LIB_STATIC_PATH="lib/libcudart_static.a"
+  CUDA_BLAS_LIB_PATH="lib/libcublas${TF_CUDA_EXT}.dylib"
+  CUDA_DNN_LIB_PATH="lib/libcudnn${TF_CUDNN_EXT}.dylib"
+  CUDA_DNN_LIB_ALT_PATH="libcudnn${TF_CUDNN_EXT}.dylib"
+  CUDA_FFT_LIB_PATH="lib/libcufft${TF_CUDA_EXT}.dylib"
+  CUDA_CUPTI_LIB_PATH="extras/CUPTI/lib/libcupti${TF_CUDA_EXT}.dylib"
+  READLINK_CMD="greadlink"
+fi
+
 if [ "$CHECK_ONLY" == "1" ]; then
   CheckAndLinkToSrcTree CudaError include/cuda.h
   CheckAndLinkToSrcTree CudaError include/cublas.h
   CheckAndLinkToSrcTree CudnnError include/cudnn.h
   CheckAndLinkToSrcTree CudaError extras/CUPTI/include/cupti.h
-  CheckAndLinkToSrcTree CudaError lib64/libcudart_static.a
-  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
-  CheckAndLinkToSrcTree CudaError extras/CUPTI/lib64/libcupti.so$TF_CUDA_VERSION
+  CheckAndLinkToSrcTree CudaError $CUDA_RT_LIB_STATIC_PATH
+  CheckAndLinkToSrcTree CudaError $CUDA_BLAS_LIB_PATH
+  CheckAndLinkToSrcTree CudnnError $CUDA_DNN_LIB_PATH
+  CheckAndLinkToSrcTree CudaError $CUDA_RT_LIB_PATH
+  CheckAndLinkToSrcTree CudaError $CUDA_FFT_LIB_PATH
+  CheckAndLinkToSrcTree CudaError $CUDA_CUPTI_LIB_PATH
   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$TF_CUDA_VERSION; then
-  CudaError "cannot find ${CUDA_TOOLKIT_PATH}/lib64/libcudart.so$TF_CUDA_VERSION"
+if test ! -e ${CUDA_TOOLKIT_PATH}/${CUDA_RT_LIB_PATH}; then
+  CudaError "cannot find ${CUDA_TOOLKIT_PATH}/${CUDA_RT_LIB_PATH}"
 fi
 
-if test ! -e ${CUDA_TOOLKIT_PATH}/extras/CUPTI/lib64/libcupti.so$TF_CUDA_VERSION; then
-  CudaError "cannot find ${CUDA_TOOLKIT_PATH}/extras/CUPTI/lib64/libcupti.so$TF_CUDA_VERSION"
+if test ! -e ${CUDA_TOOLKIT_PATH}/${CUDA_CUPTI_LIB_PATH}; then
+  CudaError "cannot find ${CUDA_TOOLKIT_PATH}/${CUDA_CUPTI_LIB_PATH}"
 fi
 
-if test ! -d ${CUDNN_INSTALL_PATH}; then
-  CudnnError "cannot find dir: ${CUDNN_INSTALL_PATH}"
+if test ! -d ${CUDNN_INSTALL_BASEDIR}; then
+  CudnnError "cannot find dir: ${CUDNN_INSTALL_BASEDIR}"
 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
+if test -e ${CUDNN_INSTALL_BASEDIR}/cudnn.h; then
+  CUDNN_HEADER_DIR=${CUDNN_INSTALL_BASEDIR}
+elif test -e ${CUDNN_INSTALL_BASEDIR}/include/cudnn.h; then
+  CUDNN_HEADER_DIR=${CUDNN_INSTALL_BASEDIR}/include
 elif test -e /usr/include/cudnn.h; then
-  CUDNN_HEADER_PATH=/usr/include
+  CUDNN_HEADER_DIR=/usr/include
 else
-  CudnnError "cannot find cudnn.h under: ${CUDNN_INSTALL_PATH} or /usr/include"
+  CudnnError "cannot find cudnn.h under: ${CUDNN_INSTALL_BASEDIR}"
 fi
 
-# 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$TF_CUDNN_VERSION; then
-  CUDNN_LIB_PATH=${CUDNN_INSTALL_PATH}/lib64
+# Locate libcudnn
+if test -e ${CUDNN_INSTALL_BASEDIR}/${CUDA_DNN_LIB_PATH}; then
+  CUDNN_LIB_INSTALL_PATH=${CUDNN_INSTALL_BASEDIR}/${CUDA_DNN_LIB_PATH}
+elif test -e ${CUDNN_INSTALL_BASEDIR}/${CUDA_DNN_LIB_ALT_PATH}; then
+  CUDNN_LIB_INSTALL_PATH=${CUDNN_INSTALL_BASEDIR}/${CUDA_DNN_LIB_ALT_PATH}
 else
-  CudnnError "cannot find libcudnn.so.$TF_CUDNN_VERSION under: ${CUDNN_INSTALL_PATH}"
+  CudnnError "cannot find ${CUDA_DNN_LIB_PATH} or ${CUDA_DNN_LIB_ALT_PATH} under: ${CUDNN_INSTALL_BASEDIR}"
 fi
 
 # Helper function to build symbolic links for all files under a directory.
@@ -181,8 +218,8 @@ function LinkAllFiles {
 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 ${CUDA_LIB_PATH}"
+LinkAllFiles ${CUDA_TOOLKIT_PATH}/${CUDA_LIB_PATH} $OUTPUTDIR/third_party/gpus/cuda/${CUDA_LIB_PATH} || 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"
@@ -190,8 +227,8 @@ LinkAllFiles ${CUDA_TOOLKIT_PATH}/nvvm $OUTPUTDIR/third_party/gpus/cuda/nvvm ||
 echo "Setting up CUPTI include"
 LinkAllFiles ${CUDA_TOOLKIT_PATH}/extras/CUPTI/include $OUTPUTDIR/third_party/gpus/cuda/extras/CUPTI/include || exit -1
 echo "Setting up CUPTI lib64"
-LinkAllFiles ${CUDA_TOOLKIT_PATH}/extras/CUPTI/lib64 $OUTPUTDIR/third_party/gpus/cuda/extras/CUPTI/lib64 || exit -1
+LinkAllFiles ${CUDA_TOOLKIT_PATH}/${CUDA_CUPTI_LIB_DIR} $OUTPUTDIR/third_party/gpus/cuda/${CUDA_CUPTI_LIB_DIR} || 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$TF_CUDNN_VERSION $OUTPUTDIR/third_party/gpus/cuda/lib64/libcudnn.so$TF_CUDNN_VERSION || exit -1
+ln -sf $CUDNN_HEADER_DIR/cudnn.h $OUTPUTDIR/third_party/gpus/cuda/include/cudnn.h || exit -1
+ln -sf $CUDNN_LIB_INSTALL_PATH $OUTPUTDIR/third_party/gpus/cuda/$CUDA_DNN_LIB_PATH || exit -1
diff --git a/third_party/gpus/cuda/platform.bzl b/third_party/gpus/cuda/platform.bzl
new file mode 100644
index 0000000000..20ab441bf4
--- /dev/null
+++ b/third_party/gpus/cuda/platform.bzl
@@ -0,0 +1,59 @@
+CUDA_VERSION = ""
+
+CUDNN_VERSION = ""
+
+PLATFORM = ""
+
+def cuda_sdk_version():
+  return CUDA_VERSION
+
+def cudnn_sdk_version():
+  return CUDNN_VERSION
+
+def cuda_library_path(name, version = cuda_sdk_version()):
+  if PLATFORM == "Darwin":
+    if not version:
+      return "lib/lib{}.dylib".format(name)
+    else:
+      return "lib/lib{}.{}.dylib".format(name, version)
+  else:
+    if not version:
+      return "lib64/lib{}.so".format(name)
+    else:
+      return "lib64/lib{}.so.{}".format(name, version)
+
+def cuda_static_library_path(name):
+  if PLATFORM == "Darwin":
+    return "lib/lib{}_static.a".format(name)
+  else:
+    return "lib64/lib{}_static.a".format(name)
+
+def cudnn_library_path(version = cudnn_sdk_version()):
+  if PLATFORM == "Darwin":
+    if not version:
+      return "lib/libcudnn.dylib"
+    else:
+      return "lib/libcudnn.{}.dylib".format(version)
+  else:
+    if not version:
+      return "lib64/libcudnn.so"
+    else:
+      return "lib64/libcudnn.so.{}".format(version)
+
+def cupti_library_path(version = cuda_sdk_version()):
+  if PLATFORM == "Darwin":
+    if not version:
+      return "extras/CUPTI/lib/libcupti.dylib"
+    else:
+      return "extras/CUPTI/lib/libcupti.{}.dylib".format(version)
+  else:
+    if not version:
+      return "extras/CUPTI/lib64/libcupti.so"
+    else:
+      return "extras/CUPTI/lib64/libcupti.so.{}".format(version)
+
+def readlink_command():
+  if PLATFORM == "Darwin":
+    return "greadlink"
+  else:
+    return "readlink"