[mlir] Fix stackoverflow in propagateLiveness

[tyb0807]

Summary

Recursion in propagateLiveness may lead to stackoverflow if the IR has many nested regions. Convert to an iterative algorithm to fix this.

mlir/lib/Transforms/Utils/RegionUtils.cpp

for (unsigned i = 0, e = op->getNumSuccessors(); i != e; ++i) {
     SuccessorOperands successorOperands =
         branchInterface.getSuccessorOperands(i);
     for (unsigned opI = 0, opE = successorOperands.getProducedOperandCount();
          opI != opE; ++opI)
       liveMap.setProvedLive(op->getSuccessor(i)->getArgument(opI));
   }
 }
 
-static void propagateLiveness(Operation *op, LiveMap &liveMap) {
-  // Recurse on any regions the op has.
-  for (Region &region : op->getRegions())
-    propagateLiveness(region, liveMap);
+static void propagateLiveness(Region &region, LiveMap &liveMap) {
+  if (region.empty())
+    return;
+
+  SmallVector<std::pair<Block *, SmallVector<std::pair<Operation *, bool>>>>
+      worklist;
+  auto addToWorklist = [&](Region &reg) {
+    for (Block &block : reg.getBlocks()) {
+      SmallVector<std::pair<Operation *, bool>> opsInBlock;
+      for (Operation &op : block.getOperations())
+        opsInBlock.push_back({&op, op.getNumRegions() != 0});
+      worklist.push_back({&block, opsInBlock});
+    }
+  };
+  addToWorklist(region);
+
+  while (!worklist.empty()) {
+    size_t lastIdx = worklist.size() - 1;
+    Block *block = worklist[lastIdx].first;
+    while (!worklist[lastIdx].second.empty()) {
+      Operation *op = worklist[lastIdx].second.back().first;
+      auto &processRegion = worklist[lastIdx].second.back().second;
+      if (processRegion) {
+        processRegion = false;
+        // Process any regions the op has if they have not been processed.
+        for (Region &reg : op->getRegions())
+          addToWorklist(reg);
+        break;
+      }
+
+      // Pop the op that is about to be processed.
+      worklist[lastIdx].second.pop_back();
 
       // Process terminator operations.
-  if (op->hasTrait<OpTrait::IsTerminator>())
-    return propagateTerminatorLiveness(op, liveMap);
+      if (op->hasTrait<OpTrait::IsTerminator>()) {
+        propagateTerminatorLiveness(op, liveMap);
+        continue;
+      }
 
       // Don't reprocess live operations.
       if (liveMap.wasProvenLive(op))
-    return;
+        continue;
 
       // Process the op itself.
-  if (!wouldOpBeTriviallyDead(op))
-    return liveMap.setProvedLive(op);
+      if (!wouldOpBeTriviallyDead(op)) {
+        liveMap.setProvedLive(op);
+        continue;
+      }
 
       // If the op isn't intrinsically alive, check it's results.
       for (Value value : op->getResults())
         processValue(value, liveMap);
     }
 
-static void propagateLiveness(Region &region, LiveMap &liveMap) {
-  if (region.empty())
-    return;
+    // Process the regions from ops in this block (newly added when processing
+    // these ops) first.
+    if (!worklist[lastIdx].second.empty())
+      continue;
 
-  for (Block *block : llvm::post_order(&region.front())) {
-    // We process block arguments after the ops in the block, to promote
-    // faster convergence to a fixed point (we try to visit uses before defs).
-    for (Operation &op : llvm::reverse(block->getOperations()))
-      propagateLiveness(&op, liveMap);
+    // Pop the block that is about to be processed.
+    worklist.pop_back();
 
     // We currently do not remove entry block arguments, so there is no need to
     // track their liveness.
     // TODO: We could track these and enable removing dead operands/arguments
     // from region control flow operations.
     if (block->isEntryBlock())
       continue;