From 2642e93e6cbb7a3a1e916abf1ab8e18fa2735237 Mon Sep 17 00:00:00 2001
From: Nathan Luehr <nluehr@nvidia.com>
Date: Fri, 14 Aug 2020 13:21:58 -0700
Subject: [PATCH] Use correct cudart soname in GetDsoHandle

---
 tensorflow/stream_executor/platform/default/dso_loader.cc | 3 ++-
 third_party/gpus/cuda/cuda_config.h.tpl                   | 1 +
 third_party/gpus/cuda_configure.bzl                       | 2 ++
 3 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/tensorflow/stream_executor/platform/default/dso_loader.cc b/tensorflow/stream_executor/platform/default/dso_loader.cc
index 84293b7767a..a78c738f32c 100644
--- a/tensorflow/stream_executor/platform/default/dso_loader.cc
+++ b/tensorflow/stream_executor/platform/default/dso_loader.cc
@@ -31,6 +31,7 @@ namespace internal {
 
 namespace {
 string GetCudaVersion() { return TF_CUDA_VERSION; }
+string GetCudaRtVersion() { return TF_CUDART_VERSION; }
 string GetCudnnVersion() { return TF_CUDNN_VERSION; }
 string GetCublasVersion() { return TF_CUBLAS_VERSION; }
 string GetCusolverVersion() { return TF_CUSOLVER_VERSION; }
@@ -77,7 +78,7 @@ port::StatusOr<void*> GetCudaDriverDsoHandle() {
 }
 
 port::StatusOr<void*> GetCudaRuntimeDsoHandle() {
-  return GetDsoHandle("cudart", GetCudaVersion());
+  return GetDsoHandle("cudart", GetCudaRtVersion());
 }
 
 port::StatusOr<void*> GetCublasDsoHandle() {
diff --git a/third_party/gpus/cuda/cuda_config.h.tpl b/third_party/gpus/cuda/cuda_config.h.tpl
index b59889938b1..ab26686ccb8 100644
--- a/third_party/gpus/cuda/cuda_config.h.tpl
+++ b/third_party/gpus/cuda/cuda_config.h.tpl
@@ -17,6 +17,7 @@ limitations under the License.
 #define CUDA_CUDA_CONFIG_H_
 
 #define TF_CUDA_VERSION "%{cuda_version}"
+#define TF_CUDART_VERSION "%{cudart_version}"
 #define TF_CUBLAS_VERSION "%{cublas_version}"
 #define TF_CUSOLVER_VERSION "%{cusolver_version}"
 #define TF_CURAND_VERSION "%{curand_version}"
diff --git a/third_party/gpus/cuda_configure.bzl b/third_party/gpus/cuda_configure.bzl
index 3e6bdc9d8eb..f85a53b1593 100644
--- a/third_party/gpus/cuda_configure.bzl
+++ b/third_party/gpus/cuda_configure.bzl
@@ -824,6 +824,7 @@ filegroup(name="cudnn-include")
         "cuda:cuda_config.h",
         {
             "%{cuda_version}": "",
+            "%{cudart_version}": "",
             "%{cublas_version}": "",
             "%{cusolver_version}": "",
             "%{curand_version}": "",
@@ -1289,6 +1290,7 @@ def _create_local_cuda_repository(repository_ctx):
         tpl_paths["cuda:cuda_config.h"],
         {
             "%{cuda_version}": cuda_config.cuda_version,
+            "%{cudart_version}": cuda_config.cudart_version,
             "%{cublas_version}": cuda_config.cublas_version,
             "%{cusolver_version}": cuda_config.cusolver_version,
             "%{curand_version}": cuda_config.curand_version,