aboutsummaryrefslogtreecommitdiffhomepage
diff options
context:
space:
mode:
authorGravatar A. Unique TensorFlower <gardener@tensorflow.org>2018-10-10 03:29:54 -0700
committerGravatar TensorFlower Gardener <gardener@tensorflow.org>2018-10-10 03:34:17 -0700
commit028ca321cb7b476868dcb39585d5cd361d81f05f (patch)
treeff00131879d1f4eaff93bcceb39a8353eb6786b6
parentd6a3d6a8295359364c86aecc479e6392bcde0ce4 (diff)
Support removing side effecting instructions with RemoveInstructionAndUnusedOperands
If the caller explicitly asks to remove a side effceting instruction (e.g. all-reduce) then we should respect it instead of silently ignoring the request. PiperOrigin-RevId: 216505133
-rw-r--r--tensorflow/compiler/xla/service/hlo_computation.cc2
-rw-r--r--tensorflow/compiler/xla/service/hlo_computation.h7
2 files changed, 5 insertions, 4 deletions
diff --git a/tensorflow/compiler/xla/service/hlo_computation.cc b/tensorflow/compiler/xla/service/hlo_computation.cc
index c2041c4667..b0f7cd91ad 100644
--- a/tensorflow/compiler/xla/service/hlo_computation.cc
+++ b/tensorflow/compiler/xla/service/hlo_computation.cc
@@ -215,7 +215,7 @@ Status HloComputation::RemoveInstructionAndUnusedOperands(
if (removed.count(item) != 0 || item->user_count() != 0 ||
item == root_instruction() || !IsRemovable(item) ||
- item->HasSideEffect()) {
+ (item->HasSideEffect() && item != instruction)) {
continue;
}
for (int i = 0; i < item->operand_count(); ++i) {
diff --git a/tensorflow/compiler/xla/service/hlo_computation.h b/tensorflow/compiler/xla/service/hlo_computation.h
index d87ab4bda1..dec96d11a9 100644
--- a/tensorflow/compiler/xla/service/hlo_computation.h
+++ b/tensorflow/compiler/xla/service/hlo_computation.h
@@ -128,9 +128,10 @@ class HloComputation {
// users. Instruction is deallocated with this call.
Status RemoveInstruction(HloInstruction* instruction);
- // Remove an instruction from the computation and also transitively any
- // operand that has no users post removing an instruction. The instruction
- // must have no users. Instruction is deallocated with this call.
+ // Remove an instruction (including side effecting ones) from the computation
+ // and also transitively any operand that has no side effect and no users post
+ // removing an instruction. The instruction must have no users. Instruction is
+ // deallocated with this call.
Status RemoveInstructionAndUnusedOperands(HloInstruction* instruction);
// Set the root of the computation to the given instruction. The instruction