[LICM] Make sinkRegion and hoistRegion non-recursive

[majnemer]

Summary

Large CFGs can cause us to blow up the stack because we would have a
recursive step for each basic block in a region.

Instead, create a worklist and iterate it. This limits the stack usage
to something more manageable.

[davide]

The idea is sound, I'll think a little bit more about the visitation order (tomorrow morning) but it seems you got it right.
Just a minor nit, I see this patch has some unrelated formatting, mind to revert that? (and/or commit separately)

[majnemer]

In D35609#814134, @davide wrote:

The idea is sound, I'll think a little bit more about the visitation order (tomorrow morning) but it seems you got it right.
Just a minor nit, I see this patch has some unrelated formatting, mind to revert that? (and/or commit separately)

AFAICT, iterating over the worklist caused the nesting to increase which caused the formatting to change.

[sanjoy]

lib/Transforms/Scalar/LICM.cpp

348 ↗

(On Diff #107249)

getDescendants seems to be doing a BFS. Given that (and that DT is a tree), can you iterate Worklist in reverse to get a post order? If you do this, we'd have to spec getDescendants as doing a BFS though (and not just have it be an implementation detail).

OTOH if you write the BFS by hand (instead of relying on getDescendents), you can add an early exit for CurLoop->contains(BB).

[majnemer]

lib/Transforms/Scalar/LICM.cpp

348 ↗

(On Diff #107249)

Do you mean BFS or DFS?

[sanjoy]

lib/Transforms/Scalar/LICM.cpp

348 ↗

(On Diff #107249)

I meant BFS; since DT is a tree you can do a (slightly) cheaper in-place BFS as:

Order.push_back(N);
for (int i = 0; i < Order.size(); i++)
  Order.push_back(N->children());

[majnemer]

lib/Transforms/Scalar/LICM.cpp

348 ↗

(On Diff #107249)

I don't believe reverse iterating a BFS gives you a post-order.

[davide]

If I understand the patch correctly, aren't we just iterating children->parent, therefore the order of siblings should be irrelevant in this case ?
(i.e. the worklist is sorted according to the dominance relation and shouldn't really matter if we process two siblings in the dominator in different order?)
I think Sanjoy's suggestion works if my reasoning is correct, although I'm also fine with the current version of the patch (I don't have a strong preference for one or the other).

[efriedma]

lib/Transforms/Scalar/LICM.cpp

354 ↗

(On Diff #107249)

properlyDominates is not a strict weak ordering, so this has undefined behavior.

[majnemer]

• Address review feedback

[davide]

LGTM, nice to have you back from the grave :)

lib/Transforms/Scalar/LICM.cpp

346 ↗

(On Diff #107431)

capital i maybe?

[sanjoy]

lib/Transforms/Scalar/LICM.cpp

349 ↗

(On Diff #107431)

Not sure that the lambda is actually useful here -- why not just:

if (CurLoop->contains(Child->getBlock()))
  Worklist.push_back(CurLoop);

?

If you wanted to go a bit further, I'd say even

for (size_t i = 0; i < Worklist.size(); i++)
  copy_if(Worklist[i]->getChildren(), <lambda to check if node is in loop>, std::back_inserter(Worklist));

may work.

[davide]

lib/Transforms/Scalar/LICM.cpp

349 ↗

(On Diff #107431)

I personally find the lambda version slightly more readable.

[sanjoy]

lib/Transforms/Scalar/LICM.cpp

349 ↗

(On Diff #107431)

In that case SGTM, but please use camel case. :)

llvm/trunk/lib/Transforms/Scalar/LICM.cpp

bool LoopInvariantCodeMotion::runOnLoop(Loop *L, AliasAnalysis *AA,
   else
     delete CurAST;
 
   if (Changed && SE)
     SE->forgetLoopDispositions(L);
   return Changed;
 }
 
+// Does a BFS from a given node to all of its children inside a given loop.
+// The returned vector of nodes includes the starting point.
+static SmallVector<DomTreeNode *, 16>
+collectChildrenInLoop(DomTreeNode *N, const Loop *CurLoop) {
+  SmallVector<DomTreeNode *, 16> Worklist;
+  auto add_region_to_worklist = [&](DomTreeNode *DTN) {
+    // Only include subregions in the top level loop.
+    BasicBlock *BB = DTN->getBlock();
+    if (CurLoop->contains(BB))
+      Worklist.push_back(DTN);
+  };
+
+  add_region_to_worklist(N);
+
+  for (size_t I = 0; I < Worklist.size(); I++) {
+    DomTreeNode *DTN = Worklist[I];
+    for (DomTreeNode *Child : DTN->getChildren())
+      add_region_to_worklist(Child);
+  }
+
+  return Worklist;
+}
+
 /// Walk the specified region of the CFG (defined by all blocks dominated by
 /// the specified block, and that are in the current loop) in reverse depth
 /// first order w.r.t the DominatorTree.  This allows us to visit uses before
 /// definitions, allowing us to sink a loop body in one pass without iteration.
 ///
 bool llvm::sinkRegion(DomTreeNode *N, AliasAnalysis *AA, LoopInfo *LI,
                       DominatorTree *DT, TargetLibraryInfo *TLI, Loop *CurLoop,
                       AliasSetTracker *CurAST, LoopSafetyInfo *SafetyInfo,
                       OptimizationRemarkEmitter *ORE) {
 
   // Verify inputs.
   assert(N != nullptr && AA != nullptr && LI != nullptr && DT != nullptr &&
          CurLoop != nullptr && CurAST != nullptr && SafetyInfo != nullptr &&
          "Unexpected input to sinkRegion");
 
-  BasicBlock *BB = N->getBlock();
-  // If this subregion is not in the top level loop at all, exit.
-  if (!CurLoop->contains(BB))
-    return false;
+  // We want to visit children before parents. We will enque all the parents
+  // before their children in the worklist and process the worklist in reverse
+  // order.
+  SmallVector<DomTreeNode *, 16> Worklist = collectChildrenInLoop(N, CurLoop);
 
-  // We are processing blocks in reverse dfo, so process children first.
   bool Changed = false;
-  const std::vector<DomTreeNode *> &Children = N->getChildren();
-  for (DomTreeNode *Child : Children)
-    Changed |=
-        sinkRegion(Child, AA, LI, DT, TLI, CurLoop, CurAST, SafetyInfo, ORE);
-
+  for (DomTreeNode *DTN : reverse(Worklist)) {
+    BasicBlock *BB = DTN->getBlock();
     // Only need to process the contents of this block if it is not part of a
     // subloop (which would already have been processed).
     if (inSubLoop(BB, CurLoop, LI))
-    return Changed;
+      continue;
 
     for (BasicBlock::iterator II = BB->end(); II != BB->begin();) {
       Instruction &I = *--II;
 
       // If the instruction is dead, we would try to sink it because it isn't used
       // in the loop, instead, just delete it.
       if (isInstructionTriviallyDead(&I, TLI)) {
         DEBUG(dbgs() << "LICM deleting dead inst: " << I << '\n');
         ++II;
         CurAST->deleteValue(&I);
         I.eraseFromParent();
         Changed = true;
         continue;
       }
 
       // Check to see if we can sink this instruction to the exit blocks
       // of the loop.  We can do this if the all users of the instruction are
       // outside of the loop.  In this case, it doesn't even matter if the
       // operands of the instruction are loop invariant.
       //
       if (isNotUsedInLoop(I, CurLoop, SafetyInfo) &&
           canSinkOrHoistInst(I, AA, DT, CurLoop, CurAST, SafetyInfo, ORE)) {
         ++II;
         Changed |= sink(I, LI, DT, CurLoop, CurAST, SafetyInfo, ORE);
       }
     }
+  }
   return Changed;
 }
 
 /// Walk the specified region of the CFG (defined by all blocks dominated by
 /// the specified block, and that are in the current loop) in depth first
 /// order w.r.t the DominatorTree.  This allows us to visit definitions before
 /// uses, allowing us to hoist a loop body in one pass without iteration.
 ///
 bool llvm::hoistRegion(DomTreeNode *N, AliasAnalysis *AA, LoopInfo *LI,
                        DominatorTree *DT, TargetLibraryInfo *TLI, Loop *CurLoop,
                        AliasSetTracker *CurAST, LoopSafetyInfo *SafetyInfo,
                        OptimizationRemarkEmitter *ORE) {
   // Verify inputs.
   assert(N != nullptr && AA != nullptr && LI != nullptr && DT != nullptr &&
          CurLoop != nullptr && CurAST != nullptr && SafetyInfo != nullptr &&
          "Unexpected input to hoistRegion");
 
-  BasicBlock *BB = N->getBlock();
-
-  // If this subregion is not in the top level loop at all, exit.
-  if (!CurLoop->contains(BB))
-    return false;
+  // We want to visit parents before children. We will enque all the parents
+  // before their children in the worklist and process the worklist in order.
+  SmallVector<DomTreeNode *, 16> Worklist = collectChildrenInLoop(N, CurLoop);
 
+  bool Changed = false;
+  for (DomTreeNode *DTN : Worklist) {
+    BasicBlock *BB = DTN->getBlock();
     // Only need to process the contents of this block if it is not part of a
     // subloop (which would already have been processed).
-  bool Changed = false;
     if (!inSubLoop(BB, CurLoop, LI))
       for (BasicBlock::iterator II = BB->begin(), E = BB->end(); II != E;) {
         Instruction &I = *II++;
         // Try constant folding this instruction.  If all the operands are
-      // constants, it is technically hoistable, but it would be better to just
-      // fold it.
+        // constants, it is technically hoistable, but it would be better to
+        // just fold it.
         if (Constant *C = ConstantFoldInstruction(
                 &I, I.getModule()->getDataLayout(), TLI)) {
           DEBUG(dbgs() << "LICM folding inst: " << I << "  --> " << *C << '\n');
           CurAST->copyValue(&I, C);
           I.replaceAllUsesWith(C);
           if (isInstructionTriviallyDead(&I, TLI)) {
             CurAST->deleteValue(&I);
             I.eraseFromParent();
           }
           Changed = true;
           continue;
         }
 
-      // Attempt to remove floating point division out of the loop by converting
-      // it to a reciprocal multiplication.
+        // Attempt to remove floating point division out of the loop by
+        // converting it to a reciprocal multiplication.
         if (I.getOpcode() == Instruction::FDiv &&
             CurLoop->isLoopInvariant(I.getOperand(1)) &&
             I.hasAllowReciprocal()) {
           auto Divisor = I.getOperand(1);
           auto One = llvm::ConstantFP::get(Divisor->getType(), 1.0);
           auto ReciprocalDivisor = BinaryOperator::CreateFDiv(One, Divisor);
           ReciprocalDivisor->setFastMathFlags(I.getFastMathFlags());
           ReciprocalDivisor->insertBefore(&I);
 
-        auto Product = BinaryOperator::CreateFMul(I.getOperand(0),
-                                                  ReciprocalDivisor);
+          auto Product =
+              BinaryOperator::CreateFMul(I.getOperand(0), ReciprocalDivisor);
           Product->setFastMathFlags(I.getFastMathFlags());
           Product->insertAfter(&I);
           I.replaceAllUsesWith(Product);
           I.eraseFromParent();
 
           hoist(*ReciprocalDivisor, DT, CurLoop, SafetyInfo, ORE);
           Changed = true;
           continue;
         }
 
-      // Try hoisting the instruction out to the preheader.  We can only do this
-      // if all of the operands of the instruction are loop invariant and if it
-      // is safe to hoist the instruction.
+        // Try hoisting the instruction out to the preheader.  We can only do
+        // this if all of the operands of the instruction are loop invariant and
+        // if it is safe to hoist the instruction.
         //
         if (CurLoop->hasLoopInvariantOperands(&I) &&
             canSinkOrHoistInst(I, AA, DT, CurLoop, CurAST, SafetyInfo, ORE) &&
             isSafeToExecuteUnconditionally(
                 I, DT, CurLoop, SafetyInfo, ORE,
                 CurLoop->getLoopPreheader()->getTerminator()))
           Changed |= hoist(I, DT, CurLoop, SafetyInfo, ORE);
       }
+  }
 
-  const std::vector<DomTreeNode *> &Children = N->getChildren();
-  for (DomTreeNode *Child : Children)
-    Changed |=
-        hoistRegion(Child, AA, LI, DT, TLI, CurLoop, CurAST, SafetyInfo, ORE);
   return Changed;
 }
 
 /// Computes loop safety information, checks loop body & header
 /// for the possibility of may throw exception.
 ///
 void llvm::computeLoopSafetyInfo(LoopSafetyInfo *SafetyInfo, Loop *CurLoop) {
   assert(CurLoop != nullptr && "CurLoop cant be null");