diff --git a/tensorflow/compiler/xla/service/BUILD b/tensorflow/compiler/xla/service/BUILD index acd35cbc153..2fd457e8e47 100644 --- a/tensorflow/compiler/xla/service/BUILD +++ b/tensorflow/compiler/xla/service/BUILD @@ -1679,6 +1679,7 @@ cc_library( hdrs = ["multi_output_fusion.h"], deps = [ ":hlo", + ":hlo_dce", ":hlo_pass", ":hlo_reachability", "//tensorflow/compiler/xla:debug_options_flags", diff --git a/tensorflow/compiler/xla/service/multi_output_fusion.cc b/tensorflow/compiler/xla/service/multi_output_fusion.cc index b95b27d6291..a21cec538d1 100644 --- a/tensorflow/compiler/xla/service/multi_output_fusion.cc +++ b/tensorflow/compiler/xla/service/multi_output_fusion.cc @@ -17,6 +17,7 @@ limitations under the License. #include "absl/container/flat_hash_set.h" #include "tensorflow/compiler/xla/debug_options_flags.h" +#include "tensorflow/compiler/xla/service/hlo_dce.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" #include "tensorflow/compiler/xla/service/hlo_reachability.h" @@ -126,6 +127,10 @@ StatusOr MultiOutputFusion::Run(HloModule* module) { candidates_index_.clear(); all_fusion_candidates_.clear(); reachability_.reset(); + if (changed) { + HloDCE dce; + TF_RETURN_IF_ERROR(dce.Run(module).status()); + } return changed; }