[Rematerialization] Reconsider to remat operations with control dependencies

We added a conservartive logic to not rematerialize operations with control dependencies since the rematerialized operations could result in undesired ordering. However, we now realize that when we remat an operation, we also copy the dependencies of them, which guarantees the rematerialized operation has the same constraint as the original operation.

PiperOrigin-RevId: 165654629

1 files changed, 31 insertions, 10 deletions

diff --git a/tensorflow/compiler/xla/service/hlo_rematerialization.cc b/tensorflow/compiler/xla/service/hlo_rematerialization.cc
index 9f65f1b851..a0e5bb7911 100644
--- a/tensorflow/compiler/xla/service/hlo_rematerialization.cc
+++ b/tensorflow/compiler/xla/service/hlo_rematerialization.cc
@@ -55,16 +55,6 @@ namespace {
 
 // Returns true if the given instruction is rematerializable.
 bool IsRematerializable(const HloInstruction* instruction) {
-  // Conservatively, don't rematerialize instruction with control
-  // dependencies. For one, control dependencies are added to prevent
-  // interference of aliased buffers (say, in while bodies) and
-  // rematerialization is ignorant of liveness and may break the intended
-  // ordering.
-  if (!instruction->control_predecessors().empty() ||
-      !instruction->control_successors().empty()) {
-    return false;
-  }
-
   // Don't rematerialize instructions with side effects or instructions which
   // cannot be cloned safely.
   switch (instruction->opcode()) {
@@ -906,6 +896,19 @@ Item* PickRematerializationCandidate(const MemoryUsageTracker& memory_tracker,
       continue;
     }
 
+    // If any of the candidate's control successor has been placed, we need to
+    // skip this candidate. Otherwise we will violate control dependency.
+    bool control_successor_placed =
+        std::any_of(candidate->control_successors().begin(),
+                    candidate->control_successors().end(),
+                    [&memory_tracker](const HloInstruction* inst) {
+                      return memory_tracker.IsPlaced(inst);
+                    });
+
+    if (control_successor_placed) {
+      continue;
+    }
+
     const int64 memory_reduced =
         memory_tracker.MemoryReducedIfRematerialized(item);
 
@@ -1047,6 +1050,15 @@ StatusOr<bool> HloRematerialization::RematerializeComputation(
 
       HloInstruction* remat =
           computation->AddInstruction(best->Clone(/*suffix=*/"remat"));
+
+      // Add control dependencies to the new operation.
+      for (auto successor : best->control_successors()) {
+        TF_RETURN_IF_ERROR(remat->AddControlDependencyTo(successor));
+      }
+      for (auto predecessor : best->control_predecessors()) {
+        TF_RETURN_IF_ERROR(predecessor->AddControlDependencyTo(remat));
+      }
+
       Item* remat_item = instruction_list.CreateItem(remat);
 
       // Replace each remaining use of 'best' with the rematerialization.
@@ -1082,6 +1094,15 @@ StatusOr<bool> HloRematerialization::RematerializeComputation(
           }
         }
       }
+      // Insert rematerialized instruction before any of its successors to
+      // preserve ordering regarding control dependency.
+      for (auto successor : remat->control_successors()) {
+        Item* successor_item = instruction_list.GetItem(successor);
+        // Assert to make sure we never remat an operation with control
+        // successor already placed.
+        CHECK(!successor_item->placed);
+        place_before.push_back(successor_item);
+      }
       instruction_list.InsertBeforeInstructions(remat_item, place_before);
 
       // If the rematerialized instruction is dead then rematerialization is