Assemble CUDA and cuDNN DLL name in build_info from version.

The version strings from cuda_configure.bzl are
cuda_version: x.y on Linux, 64_xy on Windows
cudnn_version: x on Linux, 64_x on Windows

The DLL names can be assembled from those, instead of passing them in from environment variables.

This hopefully fixes the Windows build.

PiperOrigin-RevId: 315212567
Change-Id: Ia9709297f8eff1baece1722e58e58f5bb469bbbf
This commit is contained in:
Christian Sigg 2020-06-07 22:03:55 -07:00 committed by TensorFlower Gardener
parent 5875e852b2
commit 7a239d4952
2 changed files with 7 additions and 6 deletions

View File

@ -2633,8 +2633,8 @@ def tf_py_build_info_genrule(name, out):
"msvcp_dll_names": "msvcp140.dll,msvcp140_1.dll", "msvcp_dll_names": "msvcp140.dll,msvcp140_1.dll",
}), "") + if_windows_cuda(_dict_to_kv({ }), "") + if_windows_cuda(_dict_to_kv({
"nvcuda_dll_name": "nvcuda.dll", "nvcuda_dll_name": "nvcuda.dll",
"cudart_dll_name": "cudart64_$$(echo $${TF_CUDA_VERSION:-} | sed \"s/\\.//\").dll", "cudart_dll_name": "cudart{cuda_version}.dll",
"cudnn_dll_name": "cudnn64_$${TF_CUDNN_VERSION:-}.dll", "cudnn_dll_name": "cudnn{cudnn_version}.dll",
}), ""), }), ""),
) )

View File

@ -39,6 +39,10 @@ def write_build_info(filename, key_value_list):
""" """
build_info = {} build_info = {}
if cuda_config:
build_info.update(cuda_config.config)
for arg in key_value_list: for arg in key_value_list:
key, value = six.ensure_str(arg).split("=") key, value = six.ensure_str(arg).split("=")
if value.lower() == "true": if value.lower() == "true":
@ -46,10 +50,7 @@ def write_build_info(filename, key_value_list):
elif value.lower() == "false": elif value.lower() == "false":
build_info[key] = False build_info[key] = False
else: else:
build_info[key] = value build_info[key] = value.format(**build_info)
if cuda_config:
build_info.update(cuda_config.config)
contents = """ contents = """
# Copyright 2020 The TensorFlow Authors. All Rights Reserved. # Copyright 2020 The TensorFlow Authors. All Rights Reserved.