Avoid recursion in ExpandDomain() as stack is not happy.

PiperOrigin-RevId: 198629366

diff --git a/tensorflow/compiler/xla/service/hlo_domain_map.cc b/tensorflow/compiler/xla/service/hlo_domain_map.cc
index acb54c2..ebd5adb 100644 (file)
--- a/tensorflow/compiler/xla/service/hlo_domain_map.cc
+++ b/tensorflow/compiler/xla/service/hlo_domain_map.cc
@@ -93,31 +93,39 @@ Status HloDomainMap::InsertDomain(
 
 Status HloDomainMap::ExpandDomain(HloInstruction* instruction,
                                   DomainMetadata::Domain* domain) const {
-  if (domain->reach_set.insert(instruction).second) {
-    // We should not be finding instructions with assigned domain here.
-    // If we assigned a domain to the instruction, it means that all the
-    // instructions reached by it, should have a domain as well.
-    int64 domain_id = FindOrDefault(instruction_to_domain_, instruction, -1);
-    TF_RET_CHECK(domain_id < 0) << "Instruction " << instruction->ToString()
-                                << " already has domain " << domain_id;
-    for (HloInstruction* operand : instruction->operands()) {
-      if (IsDomainInstruction(operand)) {
-        // The reach set instruction is a user of the domain instruction
-        // (the instruction sees the kDomain as operand).
-        // IOW the dataflow enters the domain through the kDomain instruction.
-        domain->enter_domains.insert(operand);
-      } else {
-        TF_RETURN_IF_ERROR(ExpandDomain(operand, domain));
+  std::vector<HloInstruction*> in_queue;
+  in_queue.push_back(instruction);
+  while (!in_queue.empty()) {
+    HloInstruction* current_instruction = in_queue.back();
+    in_queue.pop_back();
+    if (domain->reach_set.insert(current_instruction).second) {
+      // We should not be finding instructions with assigned domain here.
+      // If we assigned a domain to the instruction, it means that all the
+      // instructions reached by it, should have a domain as well.
+      int64 domain_id =
+          FindOrDefault(instruction_to_domain_, current_instruction, -1);
+      TF_RET_CHECK(domain_id < 0)
+          << "Instruction " << current_instruction->ToString()
+          << " already has domain " << domain_id;
+      for (HloInstruction* operand : current_instruction->operands()) {
+        if (IsDomainInstruction(operand)) {
+          // The reach set instruction is a user of the domain instruction
+          // (the instruction sees the kDomain as operand).
+          // IOW the dataflow enters the domain through the kDomain instruction.
+          domain->enter_domains.insert(operand);
+        } else {
+          in_queue.push_back(operand);
+        }
       }
-    }
-    for (HloInstruction* user : instruction->users()) {
-      if (IsDomainInstruction(user)) {
-        // The reach set instruction is an operand of the domain instruction
-        // (the instruction sees the kDomain as user).
-        // IOW the dataflow exits the domain through the kDomain instruction.
-        domain->exit_domains.insert(user);
-      } else {
-        TF_RETURN_IF_ERROR(ExpandDomain(user, domain));
+      for (HloInstruction* user : current_instruction->users()) {
+        if (IsDomainInstruction(user)) {
+          // The reach set instruction is an operand of the domain instruction
+          // (the instruction sees the kDomain as user).
+          // IOW the dataflow exits the domain through the kDomain instruction.
+          domain->exit_domains.insert(user);
+        } else {
+          in_queue.push_back(user);
+        }
       }
     }
   }