experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Skip including cholesky_thunk in thunk_emitter on ROCm build.

Browse Source 
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.
This commit is contained in:
Wen-Heng (Jack) Chung 2019-08-22 16:41:57 +00:00
parent b7ab9b9e5f
commit 4fd5b07bd7
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
tensorflow/compiler/xla/service/gpu/thunk_emitter.cc
View File

@ -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));