From d9255973eb279641ad2da36d540e8f1c994b2a9f Mon Sep 17 00:00:00 2001 From: Yunxing Dai Date: Fri, 20 Nov 2020 18:32:25 -0800 Subject: [PATCH] Condintional Simplifier: Return early before creating any instruction. The newly created instruction can activate DCE and then trigger HloFixPass, which creates long compilation time. PiperOrigin-RevId: 343602826 Change-Id: I46eec0bccad24205f8fe8ee59fd73f6d2875c9e4 --- .../xla/service/conditional_simplifier.cc | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/tensorflow/compiler/xla/service/conditional_simplifier.cc b/tensorflow/compiler/xla/service/conditional_simplifier.cc index 199bc787b83..bff68574e1a 100644 --- a/tensorflow/compiler/xla/service/conditional_simplifier.cc +++ b/tensorflow/compiler/xla/service/conditional_simplifier.cc @@ -144,6 +144,15 @@ StatusOr TryRemoveConditional(HloInstruction* conditional) { << conditional->ToShortString(); return false; } + + bool branch_empty = + ComputationIsEmptyWithArrayRoot(conditional->branch_computation(0)) || + ComputationIsEmptyWithArrayRoot(conditional->branch_computation(1)); + // Empty branch is faster to execute than select. + if (branch_empty) { + return false; + } + HloInstruction* true_call_op = create_call(0); HloInstruction* false_call_op = create_call(1); auto condition_broadcast = [&](const Shape& shape) { @@ -160,13 +169,6 @@ StatusOr TryRemoveConditional(HloInstruction* conditional) { hlo->shape().tuple_shapes(i), hlo, i)); }; - bool branch_empty = - ComputationIsEmptyWithArrayRoot(conditional->branch_computation(0)) || - ComputationIsEmptyWithArrayRoot(conditional->branch_computation(1)); - // Empty branch is faster to execute than select. - if (branch_empty) { - return false; - } std::function select = [&](HloInstruction* t, HloInstruction* f) { if (f->shape().IsToken()) {