[XLA] Add dead tuple elem removal to WhileLoopSimplifier.

Specifically, if a while loop has tuple element that

 - is not used by the while condition, and
 - is not used by the while body, except to pass it along to the next
   iteration of the loop,

then we can reshape the while loop's computations to eliminate this
tuple element.

PiperOrigin-RevId: 174413683

1 files changed, 10 insertions, 0 deletions

diff --git a/tensorflow/compiler/xla/service/hlo_computation.h b/tensorflow/compiler/xla/service/hlo_computation.h
index 1ff7004c4c..0754a9024c 100644
--- a/tensorflow/compiler/xla/service/hlo_computation.h
+++ b/tensorflow/compiler/xla/service/hlo_computation.h
@@ -298,6 +298,16 @@ class HloComputation {
   std::unique_ptr<HloComputation> Clone(const string& suffix = "clone",
                                         HloModule* module = nullptr);
 
+  // Like Clone(), but if an instruction is present in replacement_map, we use
+  // the map's value to replace that instruction in the cloned computation.
+  //
+  // If replacements maps a key to nullptr, we remove that instruction from the
+  // new computation.
+  std::unique_ptr<HloComputation> CloneWithReplacements(
+      std::unordered_map<const HloInstruction*, std::unique_ptr<HloInstruction>>
+          replacements,
+      HloModule* module = nullptr, const string& suffix = "clone");
+
   // Returns true if the given instruction can be removed from the
   // computation. Instructions such as parameters and send/receive instructions
   // cannot be removed without violating invariants of the HLO computation or