[Windows] Improve import self-check with tests for GPU-related DLLs.

This change incorporates the full logic of the [Windows self-check
script](https://gist.github.com/mrry/ee5dbcfdd045fa48a27d56664411d41c)
into core TensorFlow.

Fixes #9170.

PiperOrigin-RevId: 170746452

tensorflow/contrib/cmake/CMakeLists.txt

@@ -245,7 +245,7 @@ if (tensorflow_ENABLE_GPU)
       "#define CUDA_CUDA_CONFIG_H_\n"
       "#define TF_CUDA_CAPABILITIES CudaVersion(\"3.0\"),CudaVersion(\"3.5\"),CudaVersion(\"5.2\")\n"
       "#define TF_CUDA_VERSION \"64_80\"\n"
-      "#define TF_CUDNN_VERSION \"64_5\"\n"
+      "#define TF_CUDNN_VERSION \"64_6\"\n"
       "#define TF_CUDA_TOOLKIT_PATH \"${CUDA_TOOLKIT_ROOT_DIR}\"\n"
       "#endif  // CUDA_CUDA_CONFIG_H_\n"
     )
@@ -264,8 +264,23 @@ if (tensorflow_ENABLE_GPU)
     include_directories(${tensorflow_source_dir}/third_party/gpus)
     # add cuda libraries to tensorflow_EXTERNAL_LIBRARIES
     list(APPEND tensorflow_EXTERNAL_LIBRARIES ${CUDA_LIBRARIES})
-  endif()
-endif()
+
+    # NOTE(mrry): Update these flags when the version of CUDA or cuDNN used
+    # in the default build is upgraded.
+    set(tensorflow_BUILD_INFO_FLAGS --build_config cuda --key_value
+      msvcp_dll_name=msvcp140.dll
+      cudart_dll_name=cudart64_80.dll
+      cuda_version_number=8.0
+      nvcuda_dll_name=nvcuda.dll
+      cudnn_dll_name=cudnn64_6.dll
+      cudnn_version_number=6)
+  else(WIN32)
+    message(FATAL_ERROR "CMake GPU build is currently only supported on Windows.")
+  endif(WIN32)
+else(tensorflow_ENABLE_GPU)
+  set(tensorflow_BUILD_INFO_FLAGS --build_config cpu --key_value
+    msvcp_dll_name=msvcp140.dll)
+endif(tensorflow_ENABLE_GPU)
 
 # Find python executable
 include(FindPythonInterp)

tensorflow/contrib/cmake/tf_python.cmake

@@ -638,13 +638,8 @@ add_python_module("tensorflow/contrib/reduce_slice_ops/python/ops")
 
 # Generate the tensorflow.python.platform.build_info module.
 set(BUILD_INFO_PY "${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/python/platform/build_info.py")
-if(tensorflow_ENABLE_GPU)
-  set(BUILD_CONFIG_STRING "cuda")
-else(tensorflow_ENABLE_GPU)
-  set(BUILD_CONFIG_STRING "cpu")
-endif(tensorflow_ENABLE_GPU)
 add_custom_command(TARGET tf_python_copy_scripts_to_destination PRE_BUILD
-  COMMAND ${PYTHON_EXECUTABLE} ${tensorflow_source_dir}/tensorflow/tools/build_info/gen_build_info.py --build_config ${BUILD_CONFIG_STRING} --raw_generate ${BUILD_INFO_PY})
+  COMMAND ${PYTHON_EXECUTABLE} ${tensorflow_source_dir}/tensorflow/tools/build_info/gen_build_info.py --raw_generate ${BUILD_INFO_PY} ${tensorflow_BUILD_INFO_FLAGS})
 
 
 ########################################################

tensorflow/python/platform/self_check.py

@@ -21,6 +21,9 @@ from __future__ import print_function
 import os
 
 
+from tensorflow.python.platform import build_info
+
+
 def preload_check():
   """Raises an exception if the environment is not correctly configured.
 
@@ -33,17 +36,60 @@ def preload_check():
     # we load the Python extension, so that we can raise an actionable error
     # message if they are not found.
     import ctypes  # pylint: disable=g-import-not-at-top
-    try:
-      ctypes.WinDLL("msvcp140.dll")
-    except OSError:
-      raise ImportError(
-          "Could not find 'msvcp140.dll'. TensorFlow requires that this DLL be "
-          "installed in a directory that is named in your %PATH% environment "
-          "variable. You may install this DLL by downloading Visual C++ 2015 "
-          "Redistributable Update 3 from this URL: "
-          "https://www.microsoft.com/en-us/download/details.aspx?id=53587")
-    # TODO(mrry): Add specific checks for GPU DLLs if build_info indicates
-    # that this is a GPU build.
+    if hasattr(build_info, "msvcp_dll_name"):
+      try:
+        ctypes.WinDLL(build_info.msvcp_dll_name)
+      except OSError:
+        raise ImportError(
+            "Could not find %r. TensorFlow requires that this DLL be "
+            "installed in a directory that is named in your %%PATH%% "
+            "environment variable. You may install this DLL by downloading "
+            "Visual C++ 2015 Redistributable Update 3 from this URL: "
+            "https://www.microsoft.com/en-us/download/details.aspx?id=53587"
+            % build_info.msvcp_dll_name)
+
+    if build_info.is_cuda_build:
+      # Attempt to check that the necessary CUDA DLLs are loadable.
+
+      if hasattr(build_info, "nvcuda_dll_name"):
+        try:
+          ctypes.WinDLL(build_info.nvcuda_dll_name)
+        except OSError:
+          raise ImportError(
+              "Could not find %r. TensorFlow requires that this DLL "
+              "be installed in a directory that is named in your %%PATH%% "
+              "environment variable. Typically it is installed in "
+              "'C:\\Windows\\System32'. If it is not present, ensure that you "
+              "have a CUDA-capable GPU with the correct driver installed."
+              % build_info.nvcuda_dll_name)
+
+      if hasattr(build_info, "cudart_dll_name") and hasattr(
+          build_info, "cuda_version_number"):
+        try:
+          ctypes.WinDLL(build_info.cudart_dll_name)
+        except OSError:
+          raise ImportError(
+              "Could not find %r. TensorFlow requires that this DLL be "
+              "installed in a directory that is named in your %%PATH%% "
+              "environment variable. Download and install CUDA %s from "
+              "this URL: https://developer.nvidia.com/cuda-toolkit"
+              % (build_info.cudart_dll_name, build_info.cuda_version_number))
+
+      if hasattr(build_info, "cudnn_dll_name") and hasattr(
+          build_info, "cudnn_version_number"):
+        try:
+          ctypes.WinDLL(build_info.cudnn_dll_name)
+        except OSError:
+          raise ImportError(
+              "Could not find %r. TensorFlow requires that this DLL be "
+              "installed in a directory that is named in your %%PATH%% "
+              "environment variable. Note that installing cuDNN is a separate "
+              "step from installing CUDA, and this DLL is often found in a "
+              "different directory from the CUDA DLLs. You may install the "
+              "necessary DLL by downloading cuDNN %s from this URL: "
+              "https://developer.nvidia.com/cudnn"
+              % (build_info.cudnn_dll_name, build_info.cudnn_version_number))
+
   else:
     # TODO(mrry): Consider adding checks for the Linux and Mac OS X builds.
     pass

tensorflow/tools/build_info/gen_build_info.py

@@ -20,12 +20,19 @@ from __future__ import print_function
 import argparse
 
 
-def write_build_info(filename, build_config):
+def write_build_info(filename, build_config, key_value_list):
   """Writes a Python that describes the build.
 
   Args:
     filename: filename to write to.
-    build_config: A string containinggit_version: the result of a git describe.
+    build_config: A string that represents the config used in this build (e.g.
+      "cuda").
+    key_value_list: A list of "key=value" strings that will be added to the
+      module as additional fields.
+
+  Raises:
+    ValueError: If `key_value_list` includes the key "is_cuda_build", which
+      would clash with one of the default fields.
   """
   module_docstring = "\"\"\"Generates a Python module containing information "
   module_docstring += "about the build.\"\"\""
@@ -34,6 +41,16 @@ def write_build_info(filename, build_config):
   else:
     build_config_bool = "False"
 
+  key_value_pair_stmts = []
+  if key_value_list:
+    for arg in key_value_list:
+      key, value = arg.split("=")
+      if key == "is_cuda_build":
+        raise ValueError("The key \"is_cuda_build\" cannot be passed as one of "
+                         "the --key_value arguments.")
+      key_value_pair_stmts.append("%s = %r" % (key, value))
+  key_value_pair_content = "\n".join(key_value_pair_stmts)
+
   contents = """
 # Copyright 2017 The TensorFlow Authors. All Rights Reserved.
 #
@@ -55,7 +72,9 @@ from __future__ import division
 from __future__ import print_function
 
 is_cuda_build = %s
-""" % (module_docstring, build_config_bool)
+
+%s
+""" % (module_docstring, build_config_bool, key_value_pair_content)
   open(filename, "w").write(contents)
 
 
@@ -69,9 +88,12 @@ parser.add_argument(
 
 parser.add_argument("--raw_generate", type=str, help="Generate build_info.py")
 
+parser.add_argument("--key_value", type=str, nargs="*",
+                    help="List of key=value pairs.")
+
 args = parser.parse_args()
 
 if args.raw_generate is not None and args.build_config is not None:
-  write_build_info(args.raw_generate, args.build_config)
+  write_build_info(args.raw_generate, args.build_config, args.key_value)
 else:
   raise RuntimeError("--raw_generate and --build_config must be used")