From a709a894c2cdc6fe069ba78ac20df7b916e9579a Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 29 Apr 2019 05:32:51 -0700
Subject: [PATCH] Improve performance of domain isolator
Previously we used HloInstruction::ReplaceUseWith what involves a sahpe
check whayt can be expensive for large tuples. After this change we swith
to HloInstruction::ReplaceUseWithDifferentShape what is an equivalent method
just without the shape check to save compile time.
PiperOrigin-RevId: 245729563
---
tensorflow/compiler/xla/service/hlo_domain_isolator.cc | 6 +++++-
1 file changed, 5 insertions(+), 1 deletion(-)
diff --git a/tensorflow/compiler/xla/service/hlo_domain_isolator.cc b/tensorflow/compiler/xla/service/hlo_domain_isolator.cc
index 3746fbbda02..5b388bc0bd8 100644
--- a/tensorflow/compiler/xla/service/hlo_domain_isolator.cc
+++ b/tensorflow/compiler/xla/service/hlo_domain_isolator.cc
@@ -47,7 +47,11 @@ StatusOr<bool> RunInternal(HloModule* module,
HloInstruction* domain = (*creator)(instruction, root, operand);
if (domain != nullptr) {
VLOG(4) << "New domain: " << domain->ToString();
- TF_RETURN_IF_ERROR(operand->ReplaceUseWith(instruction, domain));
+ // Call ReplaceUseWithDifferentShape even though the shapes are
+ // expected to match to avoid an expensive shape check between the
+ // original and the new instruction.
+ TF_RETURN_IF_ERROR(
+ operand->ReplaceUseWithDifferentShape(instruction, domain));
++added_domains;
}
}