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, 0 insertions, 13 deletions

diff --git a/tensorflow/compiler/xla/service/hlo_instruction.h b/tensorflow/compiler/xla/service/hlo_instruction.h
index 566c1c449a..948e33a0a3 100644
--- a/tensorflow/compiler/xla/service/hlo_instruction.h
+++ b/tensorflow/compiler/xla/service/hlo_instruction.h
@@ -1092,19 +1092,6 @@ class HloInstruction {
   // instruction.
   void SetupDerivedInstruction(HloInstruction* derived_instruction) const;
 
-  // TODO(b/80249101): Remove these methods once HLO scheduling and copy
-  // insertion are integrated, and we don't need to run a separate pass
-  // of copy elision anymore.
-  bool CopyElisionAllowed() const {
-    CHECK_EQ(HloOpcode::kCopy, opcode_);
-    return copy_elision_allowed_;
-  }
-
-  void SetCopyElisionAllowed(bool value) {
-    CHECK_EQ(HloOpcode::kCopy, opcode_);
-    copy_elision_allowed_ = value;
-  }
-
   // Returns data on the dimension numbers used for a dot operation.
   const DotDimensionNumbers& dot_dimension_numbers() const {
     CHECK(dot_dimension_numbers_ != nullptr);