From 5e23e0e67ac565d56de7680ccb8d8ccc6a0d2179 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 23 Oct 2017 11:12:26 -0700
Subject: [PATCH] [XLA] Erase cloned instructions on the fly when merging
fusion nodes.
This avoids the awkward situation where an RNG which is clearly eligible for fusion becomes ineligible mid-fusion because it suddenly has an extra (dead) user.
PiperOrigin-RevId: 173141716
---
tensorflow/compiler/xla/service/hlo_instruction.cc | 2 ++
1 file changed, 2 insertions(+)
diff --git a/tensorflow/compiler/xla/service/hlo_instruction.cc b/tensorflow/compiler/xla/service/hlo_instruction.cc
index f24953051ac..0669a868630 100644
--- a/tensorflow/compiler/xla/service/hlo_instruction.cc
+++ b/tensorflow/compiler/xla/service/hlo_instruction.cc
@@ -716,10 +716,12 @@ void HloInstruction::MergeFusionInstructionIntoMultiOutput(
// Fuse the root instruction and generate multiple outputs.
FuseInstructionIntoMultiOutput(unfused_root);
+ TF_CHECK_OK(unfused_root->parent()->RemoveInstruction(unfused_root));
// The rest instructions are of normal fusing.
for (int64 i = 1; i < unfused_instructions.size(); i++) {
auto instruction = unfused_instructions[i];
FuseInstruction(instruction);
+ TF_CHECK_OK(instruction->parent()->RemoveInstruction(instruction));
}
}