diff options
author | A. Unique TensorFlower <gardener@tensorflow.org> | 2018-10-10 03:29:54 -0700 |
---|---|---|
committer | TensorFlower Gardener <gardener@tensorflow.org> | 2018-10-10 03:34:17 -0700 |
commit | 028ca321cb7b476868dcb39585d5cd361d81f05f (patch) | |
tree | ff00131879d1f4eaff93bcceb39a8353eb6786b6 | |
parent | d6a3d6a8295359364c86aecc479e6392bcde0ce4 (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.cc | 2 | ||||
-rw-r--r-- | tensorflow/compiler/xla/service/hlo_computation.h | 7 |
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 |