Run AddSpecialCaseCopies in HloRematerialization.

RemoveUnnecessaryCopies which runs in rematerialization to take advantage of scheduling can sometimes remove copies which are needed to non-interference reasons. This requires running AddSpecialCaseCopies to add them back in. Furthermore, the schedule needs to be updated to account for the changes to the module, so add an UpdateSchedule function which can patch up a schedule in light a limited set of transformations to the module (addition and deletion of instructions).

PiperOrigin-RevId: 210186375

1 files changed, 19 insertions, 5 deletions

diff --git a/tensorflow/compiler/xla/service/copy_insertion.h b/tensorflow/compiler/xla/service/copy_insertion.h
index f797ee7e4d..d308f6bc84 100644
--- a/tensorflow/compiler/xla/service/copy_insertion.h
+++ b/tensorflow/compiler/xla/service/copy_insertion.h
@@ -77,15 +77,29 @@ class CopyInsertion : public HloPassInterface {
   Status RemoveUnnecessaryCopies(const HloOrdering& ordering,
                                  HloModule* module);
 
- private:
-  // Verifies that no HLO values have interfering live ranged assuming the
-  // ordering used by copy insertion.
-  Status VerifyNoLiveRangeInterference(HloModule* module);
+  // Add copies to address special constraints on the roots of computations not
+  // related to live range interference:
+  //
+  //    (1) Entry computation root must be unambiguous and distinct.
+  //
+  //    (2) Any computation called by a kCall instruction must have an
+  //        unambiguous root.
+  //
+  //    (3) Constants and parameters cannot be live out of the entry computation
+  //
+  Status AddSpecialCaseCopies(HloModule* module);
 
-  Status AddCopiesToResolveInterference(HloModule* module);
+  // Verifies that no HLO values have interfering live ranges using the given
+  // ordering.
+  Status VerifyNoLiveRangeInterference(const HloOrdering& ordering,
+                                       HloModule* module);
 
+ private:
+  // Override which requires the caller to pass in a call graph.
   Status AddSpecialCaseCopies(const CallGraph& call_graph, HloModule* module);
 
+  Status AddCopiesToResolveInterference(HloModule* module);
+
   // Backend specific function that decides whether a fusion can share buffer
   // with its operand.
   HloDataflowAnalysis::FusionCanShareBufferFunction fusion_can_share_buffer_;