From 4fd5b07bd7c0db869509803c6c2b0d9d13a82bb1 Mon Sep 17 00:00:00 2001
From: "Wen-Heng (Jack) Chung" <whchung@gmail.com>
Date: Thu, 22 Aug 2019 16:41:57 +0000
Subject: [PATCH] Skip including cholesky_thunk in thunk_emitter on ROCm build.

cholesky_thunk has source code dependency to cusolver_context which has source
code dependency to CUDA-specific headers. Disable them on ROCm build.

Since GOOGLE_CUDA macros is discouraged in XLA, use !TENSORFLOW_USE_ROCM
instead.
---
 tensorflow/compiler/xla/service/gpu/thunk_emitter.cc | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/tensorflow/compiler/xla/service/gpu/thunk_emitter.cc b/tensorflow/compiler/xla/service/gpu/thunk_emitter.cc
index 13d32672a95..49157bbc9d4 100644
--- a/tensorflow/compiler/xla/service/gpu/thunk_emitter.cc
+++ b/tensorflow/compiler/xla/service/gpu/thunk_emitter.cc
@@ -17,7 +17,9 @@ limitations under the License.
 
 #include "tensorflow/compiler/xla/service/custom_call_target_registry.h"
 #include "tensorflow/compiler/xla/service/gpu/backend_configs.pb.h"
-#include "tensorflow/compiler/xla/service/gpu/cholesky_thunk.h"
+#if !TENSORFLOW_USE_ROCM
+  #include "tensorflow/compiler/xla/service/gpu/cholesky_thunk.h"
+#endif
 #include "tensorflow/compiler/xla/service/gpu/convolution_thunk.h"
 #include "tensorflow/compiler/xla/service/gpu/copy_thunk.h"
 #include "tensorflow/compiler/xla/service/gpu/cudnn_batchnorm_thunk.h"
@@ -237,6 +239,7 @@ Status ThunkEmitter::HandleCustomCall(HloInstruction* custom_call) {
     return Status::OK();
   }
 
+#if !TENSORFLOW_USE_ROCM
   if (custom_call->custom_call_target() == kCusolverCholeskyCallTarget) {
     TF_ASSIGN_OR_RETURN(CholeskyOptions options,
                         custom_call->backend_config<CholeskyOptions>());
@@ -280,6 +283,7 @@ Status ThunkEmitter::HandleCustomCall(HloInstruction* custom_call) {
 
     return Status::OK();
   }
+#endif
 
   if (IsCublasGemm(*custom_call)) {
     AddThunkToThunkSequence(BuildGemmThunk(custom_call));