Introduced kDomain HLO instruction set isolation to bound connected sets of instructions with similar attributes (ie, sharding).

This CL simply adds the infrastructure, but leaves the wire-on to a separate CL.

PiperOrigin-RevId: 198503625

1 files changed, 23 insertions, 24 deletions

diff --git a/tensorflow/compiler/xla/service/hlo_computation.cc b/tensorflow/compiler/xla/service/hlo_computation.cc
index 63c3dc4a59..b61eabbbf5 100644
--- a/tensorflow/compiler/xla/service/hlo_computation.cc
+++ b/tensorflow/compiler/xla/service/hlo_computation.cc
@@ -752,22 +752,21 @@ Status HloComputation::Accept(
 }
 
 std::unique_ptr<HloComputation> HloComputation::Clone(
-    const string& suffix, HloModule* module,
-    HloInstruction::CloneMap* clone_map) {
+    const string& suffix, HloCloneContext* context) {
   return CloneWithReplacements(
       /*replacements=*/std::unordered_map<const HloInstruction*,
                                           std::unique_ptr<HloInstruction>>(),
-      module, clone_map, suffix);
+      context, suffix);
 }
 
 std::unique_ptr<HloComputation> HloComputation::CloneWithReplacements(
     std::unordered_map<const HloInstruction*, std::unique_ptr<HloInstruction>>
         replacements,
-    HloModule* module, HloInstruction::CloneMap* clone_map,
-    const string& suffix) {
-  HloInstruction::CloneMap local_clone_map;
-  if (clone_map == nullptr) {
-    clone_map = &local_clone_map;
+    HloCloneContext* context, const string& suffix) {
+  std::unique_ptr<HloCloneContext> context_ptr;
+  if (context == nullptr) {
+    context_ptr = MakeUnique<HloCloneContext>(parent(), suffix);
+    context = context_ptr.get();
   }
 
   // Look up instr in the replacements map, and return either the replacement,
@@ -792,18 +791,18 @@ std::unique_ptr<HloComputation> HloComputation::CloneWithReplacements(
   }
 
   std::vector<std::unique_ptr<HloInstruction>> instructions;
-  std::unique_ptr<HloInstruction> new_instr = nullptr;
+  std::unique_ptr<HloInstruction> new_instr;
   for (auto instr : postorder) {
     std::vector<HloInstruction*> new_operands;
     for (auto operand : instr->operands()) {
       auto replaced_operand = replace(operand);
       CHECK_NE(replaced_operand, nullptr)
-          << "Replacements map specifies to leave out " << operand->ToString()
-          << ", but it is used by " << instr->ToString() << ".";
-      new_operands.push_back(FindOrDie(*clone_map, replaced_operand));
+          << "replacements map tried to eliminate a used instruction "
+          << operand->ToString() << ", used by " << instr->ToString();
+      new_operands.push_back(context->GetInstruction(replaced_operand));
     }
-    new_instr = instr->CloneWithNewOperands(instr->shape(), new_operands,
-                                            module, clone_map);
+    new_instr =
+        instr->CloneWithNewOperands(instr->shape(), new_operands, context);
     instructions.push_back(std::move(new_instr));
   }
   Builder builder(name() + "." + suffix);
@@ -811,22 +810,23 @@ std::unique_ptr<HloComputation> HloComputation::CloneWithReplacements(
     builder.AddInstruction(std::move(instr));
   }
   auto result = builder.Build(
-      /*root_instruction=*/FindOrDie(*clone_map, replace(root_instruction())));
+      /*root_instruction=*/context->GetInstruction(
+          replace(root_instruction())));
 
   // Clone control dependencies.
   for (auto instr : postorder) {
-    HloInstruction* new_instr = FindOrDie(*clone_map, instr);
+    HloInstruction* new_instr = context->GetInstruction(instr);
     for (auto successor : instr->control_successors()) {
       auto replaced_successor = replace(successor);
-      CHECK_NE(replaced_successor, nullptr)
-          << "Replacements map specifies to leave out " << successor->ToString()
-          << ", but it is control-depended-on by " << instr->ToString() << ".";
-
-      TF_CHECK_OK(new_instr->AddControlDependencyTo(
-          FindOrDie(*clone_map, replaced_successor)));
+      // successor may not have been remapped, because it might have been
+      // removed by the replacements map.
+      if (replaced_successor != nullptr) {
+        TF_CHECK_OK(new_instr->AddControlDependencyTo(
+            context->GetInstruction(replaced_successor)));
+      }
     }
   }
-
+  context->MapComputation(this, result.get());
   // We cloned the elements of 'replacements', so they're all going to be
   // destroyed. HloInstructions need to be detached from their operands before
   // they're destroyed, otherwise they stick around in the operands' users lists
@@ -836,7 +836,6 @@ std::unique_ptr<HloComputation> HloComputation::CloneWithReplacements(
       new_instr->DetachFromOperands();
     }
   }
-
   return result;
 }