Add cuda build info again

Fixes #26395 According to https://github.com/tensorflow/tensorflow/blob/3385f66872ff56ac6468d6d2a9587ac9314e0180/tensorflow/contrib/cmake/CMakeLists.txt#L556 PiperOrigin-RevId: 238637381
tensorflow · Mar 15, 2019 · 9a43dfe · 9a43dfe
1 parent 6c4cccf
commit 9a43dfe
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 1 deletion.
diff --git a/tensorflow/tensorflow.bzl b/tensorflow/tensorflow.bzl
@@ -211,6 +211,12 @@ def if_windows(a, otherwise = []):
         "//conditions:default": otherwise,
     })
 
+def if_windows_cuda(a, otherwise = []):
+    return select({
+        clean_dep("//tensorflow:with_cuda_support_windows_override"): a,
+        "//conditions:default": otherwise,
+    })
+
 def if_not_windows_cuda(a):
     return select({
         clean_dep("//tensorflow:with_cuda_support_windows_override"): [],
@@ -2177,7 +2183,16 @@ def tf_py_build_info_genrule():
         name = "py_build_info_gen",
         outs = ["platform/build_info.py"],
         cmd =
-            "$(location //tensorflow/tools/build_info:gen_build_info) --raw_generate \"$@\" --build_config " + if_cuda("cuda", "cpu") + if_windows(" --key_value msvcp_dll_name=msvcp140.dll", ""),
+            "$(location //tensorflow/tools/build_info:gen_build_info) --raw_generate \"$@\" --build_config " +
+            if_cuda("cuda", "cpu") +
+            " --key_value " +
+            if_cuda(" cuda_version_number=$${TF_CUDA_VERSION} cudnn_version_number=$${TF_CUDNN_VERSION} ", "") +
+            if_windows(" msvcp_dll_name=msvcp140.dll ", "") +
+            if_windows_cuda(" ".join([
+                "nvcuda_dll_name=nvcuda.dll",
+                "cudart_dll_name=cudart64_$${TF_CUDA_VERSION/\\./}.dll",
+                "cudnn_dll_name=cudnn64_$${TF_CUDNN_VERSION}.dll",
+            ]), ""),
         local = 1,
         tools = [clean_dep("//tensorflow/tools/build_info:gen_build_info")],
     )

diff --git a/third_party/gpus/cuda_configure.bzl b/third_party/gpus/cuda_configure.bzl
@@ -1494,6 +1494,10 @@ def _cuda_autoconf_impl(repository_ctx):
     if not enable_cuda(repository_ctx):
         _create_dummy_repository(repository_ctx)
     elif _TF_CUDA_CONFIG_REPO in repository_ctx.os.environ:
+        if (_TF_CUDA_VERSION not in repository_ctx.os.environ or
+            _TF_CUDNN_VERSION not in repository_ctx.os.environ):
+            auto_configure_fail("%s and %s must also be set if %s is specified" %
+                                (_TF_CUDA_VERSION, _TF_CUDNN_VERSION, _TF_CUDA_CONFIG_REPO))
         _create_remote_cuda_repository(
             repository_ctx,
             repository_ctx.os.environ[_TF_CUDA_CONFIG_REPO],