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

Condintional Simplifier: Return early before creating any instruction.

Browse Source 
The newly created instruction can activate DCE and then trigger HloFixPass, which creates long compilation time.

PiperOrigin-RevId: 343602826
Change-Id: I46eec0bccad24205f8fe8ee59fd73f6d2875c9e4
This commit is contained in:
Yunxing Dai 2020-11-20 18:32:25 -08:00 committed by TensorFlower Gardener
parent 2f4d749c04
commit d9255973eb
1 changed files with 9 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
tensorflow/compiler/xla/service/conditional_simplifier.cc
View File

@ -144,6 +144,15 @@ StatusOr<bool> 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<bool> 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<HloInstruction*(HloInstruction*, HloInstruction*)> select =
[&](HloInstruction* t, HloInstruction* f) {
if (f->shape().IsToken()) {