[TRE][NFC] Refactor Basic Block Processing

[laytonio]

Summary

Simplify and improve readability.

[efriedma]

Do we want to land this before D82085?

Otherwise, LGTM

llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp

774

I'm pretty sure the address-taken check here is wrong, and this doesn't interact correctly with unreachable code. Probably we should teach eliminateCall to RAUW away dead uses instead, and then let some other code call removeUnreachableBlocks to do the cleanup properly.

I guess this isn't new code, though.

[laytonio]

I'll wait and rebase after D82085.

llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp

774

I don't think there is anything we could use to RAUW that would be semantic preserving, not that it matters much I guess since this block is now dead. I looked at the other uses of FoldReturnIntoUncondBranch (there are only 2) to see how other passes handle the possibility of creating a dead block. SimplifyCFG just checks the predecessors and deletes the block. CodeGenPrepare checks the predecessors and the address-taken and deletes the block. The BasicBlock destructor seems to correctly handle the case where a dead block has its address taken, and even a nice comment explaining how you could end up in that situation. Therefore, I think we are safe to remove the address-taken check and just delete the block if it has no predecessors.

[efriedma]

llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp

774

I don't think there is anything we could use to RAUW that would be semantic preserving, not that it matters much I guess since this block is now dead.

Exactly, the value doesn't matter because it's dead.

I looked at the other uses of FoldReturnIntoUncondBranch (there are only 2) to see how other passes handle the possibility of creating a dead block. SimplifyCFG just checks the predecessors and deletes the block. CodeGenPrepare checks the predecessors and the address-taken and deletes the block

Consider the following example, which currently crashes tailcallelim:

define i32 @foo() {
  %z = call i32 @foo()
  ret i32 %z

  ret i32 %z
}

removeUnreachableBlocks avoids the problem by calling dropAllReferences() (which only works because it's erasing all the unreachable blocks at once). CodeGenPrepare avoids the problem by accident, I think; some of the checks implicitly exclude or manipulate dead code.

[laytonio]

llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp

774

I feel like I may be misunderstanding what you're saying. In your example you are passing in a function with a dead block. Where as this only applies to something like the following:

define i32 @foo() {
  %z = call i32 @foo()
  br label %retblock

retblock:
  ret i32 %z
}

Where we transform to this:

define i32 @foo() {
  %z = call i32 @foo()
  ret i32 %z

retblock:
  ret i32 %z
}

In order to eliminate the call. Is it not valid to just delete the now dead block? It seems a little bug prone to just RAUW when the replacement doesn't maintain the semantics, but that is just my opinion.

[efriedma]

llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp

774

The general point is just that if you haven't run removeUnreachableBlocks, there can unreachable uses floating around of any instruction that defines a value. Before you erase an instruction, you have to erase/fix any uses of that instruction. To erase a basic block, you have to erase every instruction in that basic block.

Actually, taking another look at this, I guess it isn't possible to have an instruction with uses at this point. We checked that the only instructions in the block are phi, debug intrinsics, and return. FoldReturnIntoUncondBranch eliminates any phi nodes, and debug intrinsics and return don't define a value. So I guess you're right, we're okay here.

[avl]

I'll wait and rebase after D82085.

It looks like D82085 needs more work(it did not pass two staged build). I suggest not to wait D82085 and go head with D82269.

[laytonio]

Remove hasAddressTaken check
Update comment
Use pred_empty

[laytonio]

ping

[laytonio]

ping

[laytonio]

ping

[efriedma]

LGTM. Sorry about the delay.

[laytonio]

Thanks for the review! Could you commit this for me?

llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp

 
 STATISTIC(NumEliminated, "Number of tail calls removed");
 STATISTIC(NumRetDuped,   "Number of return duplicated");
 STATISTIC(NumAccumAdded, "Number of accumulators introduced");
 
 /// Scan the specified function for alloca instructions.
 /// If it contains any dynamic allocas, returns false.
 static bool canTRE(Function &F) {
-  // Because of PR962, we don't TRE dynamic allocas.
+  // FIXME: The code generator produces really bad code when an 'escaping
+  // alloca' is changed from being a static alloca to being a dynamic alloca.
+  // Until this is resolved, disable this transformation if that would ever
+  // happen.  This bug is PR962.
   return llvm::all_of(instructions(F), [](Instruction &I) {
     auto *AI = dyn_cast<AllocaInst>(&I);
     return !AI || AI->isStaticAlloca();
   });
 }
 
 namespace {
 struct AllocaDerivedValueTracker {
class TailRecursionEliminator {
   // The instruction doing the accumulating.
   Instruction *AccumulatorRecursionInstr = nullptr;
 
   TailRecursionEliminator(Function &F, const TargetTransformInfo *TTI,
                           AliasAnalysis *AA, OptimizationRemarkEmitter *ORE,
                           DomTreeUpdater &DTU)
       : F(F), TTI(TTI), AA(AA), ORE(ORE), DTU(DTU) {}
 
-  CallInst *findTRECandidate(Instruction *TI,
+  CallInst *findTRECandidate(BasicBlock *BB,
                              bool CannotTailCallElimCallsMarkedTail);
 
   void createTailRecurseLoopHeader(CallInst *CI);
 
   void insertAccumulator(Instruction *AccRecInstr);
 
   bool eliminateCall(CallInst *CI);
 
-  bool foldReturnAndProcessPred(ReturnInst *Ret,
-                                bool CannotTailCallElimCallsMarkedTail);
-
-  bool processReturningBlock(ReturnInst *Ret,
-                             bool CannotTailCallElimCallsMarkedTail);
-
   void cleanupAndFinalize();
 
+  bool processBlock(BasicBlock &BB, bool CannotTailCallElimCallsMarkedTail);
+
 public:
   static bool eliminate(Function &F, const TargetTransformInfo *TTI,
                         AliasAnalysis *AA, OptimizationRemarkEmitter *ORE,
                         DomTreeUpdater &DTU);
 };
 } // namespace
 
 CallInst *TailRecursionEliminator::findTRECandidate(
-    Instruction *TI, bool CannotTailCallElimCallsMarkedTail) {
-  BasicBlock *BB = TI->getParent();
+    BasicBlock *BB, bool CannotTailCallElimCallsMarkedTail) {
+  Instruction *TI = BB->getTerminator();
 
   if (&BB->front() == TI) // Make sure there is something before the terminator.
     return nullptr;
 
   // Scan backwards from the return, checking to see if there is a tail call in
   // this block.  If so, set CI to it.
   CallInst *CI = nullptr;
   BasicBlock::iterator BBI(TI);
bool TailRecursionEliminator::eliminateCall(CallInst *CI) {
 
   BB->getInstList().erase(Ret);  // Remove return.
   BB->getInstList().erase(CI);   // Remove call.
   DTU.applyUpdates({{DominatorTree::Insert, BB, HeaderBB}});
   ++NumEliminated;
   return true;
 }
 
-bool TailRecursionEliminator::foldReturnAndProcessPred(
-    ReturnInst *Ret, bool CannotTailCallElimCallsMarkedTail) {
-  BasicBlock *BB = Ret->getParent();
-
-  bool Change = false;
-
-  // Make sure this block is a trivial return block.
-  assert(BB->getFirstNonPHIOrDbg() == Ret &&
-         "Trying to fold non-trivial return block");
-
-  // If the return block contains nothing but the return and PHI's,
-  // there might be an opportunity to duplicate the return in its
-  // predecessors and perform TRE there. Look for predecessors that end
-  // in unconditional branch and recursive call(s).
-  SmallVector<BranchInst*, 8> UncondBranchPreds;
-  for (pred_iterator PI = pred_begin(BB), E = pred_end(BB); PI != E; ++PI) {
-    BasicBlock *Pred = *PI;
-    Instruction *PTI = Pred->getTerminator();
-    if (BranchInst *BI = dyn_cast<BranchInst>(PTI))
-      if (BI->isUnconditional())
-        UncondBranchPreds.push_back(BI);
-  }
-
-  while (!UncondBranchPreds.empty()) {
-    BranchInst *BI = UncondBranchPreds.pop_back_val();
-    BasicBlock *Pred = BI->getParent();
-    if (CallInst *CI =
-            findTRECandidate(BI, CannotTailCallElimCallsMarkedTail)) {
-      LLVM_DEBUG(dbgs() << "FOLDING: " << *BB
-                        << "INTO UNCOND BRANCH PRED: " << *Pred);
-      FoldReturnIntoUncondBranch(Ret, BB, Pred, &DTU);
-
-      // Cleanup: if all predecessors of BB have been eliminated by
-      // FoldReturnIntoUncondBranch, delete it.  It is important to empty it,
-      // because the ret instruction in there is still using a value which
-      // eliminateRecursiveTailCall will attempt to remove.
-      if (!BB->hasAddressTaken() && pred_begin(BB) == pred_end(BB))
-        DTU.deleteBB(BB);
-
-      eliminateCall(CI);
-      ++NumRetDuped;
-      Change = true;
-    }
-  }
-
-  return Change;
-}
-
-bool TailRecursionEliminator::processReturningBlock(
-    ReturnInst *Ret, bool CannotTailCallElimCallsMarkedTail) {
-  CallInst *CI = findTRECandidate(Ret, CannotTailCallElimCallsMarkedTail);
-  if (!CI)
-    return false;
-
-  return eliminateCall(CI);
-}
-
 void TailRecursionEliminator::cleanupAndFinalize() {
   // If we eliminated any tail recursions, it's possible that we inserted some
   // silly PHI nodes which just merge an initial value (the incoming operand)
   // with themselves.  Check to see if we did and clean up our mess if so.  This
   // occurs when a function passes an argument straight through to its tail
   // call.
   for (PHINode *PN : ArgumentPHIs) {
     // If the PHI Node is a dynamic constant, replace it with the value it is.
if (RetSelects.empty()) {
           AccRecInstrNew->insertBefore(SI);
           SI->setFalseValue(AccRecInstrNew);
         }
       }
     }
   }
 }
 
+bool TailRecursionEliminator::processBlock(
+    BasicBlock &BB, bool CannotTailCallElimCallsMarkedTail) {
+  Instruction *TI = BB.getTerminator();
+
+  if (BranchInst *BI = dyn_cast<BranchInst>(TI)) {
+    if (BI->isConditional())
+      return false;
+
+    BasicBlock *Succ = BI->getSuccessor(0);
+    ReturnInst *Ret = dyn_cast<ReturnInst>(Succ->getFirstNonPHIOrDbg());
+
+    if (!Ret)
+      return false;
+
+    CallInst *CI = findTRECandidate(&BB, CannotTailCallElimCallsMarkedTail);
+
+    if (!CI)
+      return false;
+
+    LLVM_DEBUG(dbgs() << "FOLDING: " << *Succ
+                      << "INTO UNCOND BRANCH PRED: " << BB);
+    FoldReturnIntoUncondBranch(Ret, Succ, &BB, &DTU);
+    ++NumRetDuped;
+
+    // Cleanup: if all predecessors of Succ have been eliminated by
+    // FoldReturnIntoUncondBranch, delete it.  It is important to empty it,
+    // because the ret instruction in there is still using a value which
+    // eliminateCall will attempt to remove.
+    if (!Succ->hasAddressTaken() && pred_begin(Succ) == pred_end(Succ))
+      DTU.deleteBB(Succ);
+
+    eliminateCall(CI);
+    return true;
+  } else if (isa<ReturnInst>(TI)) {
+    CallInst *CI = findTRECandidate(&BB, CannotTailCallElimCallsMarkedTail);
+
+    if (CI)
+      return eliminateCall(CI);
+  }
+
+  return false;
+}
+
 bool TailRecursionEliminator::eliminate(Function &F,
                                         const TargetTransformInfo *TTI,
                                         AliasAnalysis *AA,
                                         OptimizationRemarkEmitter *ORE,
                                         DomTreeUpdater &DTU) {
   if (F.getFnAttribute("disable-tail-calls").getValueAsString() == "true")
     return false;
 
   bool MadeChange = false;
   bool AllCallsAreTailCalls = false;
   MadeChange |= markTails(F, AllCallsAreTailCalls, ORE);
   if (!AllCallsAreTailCalls)
     return MadeChange;
 
   // If this function is a varargs function, we won't be able to PHI the args
   // right, so don't even try to convert it...
   if (F.getFunctionType()->isVarArg())
     return MadeChange;
 
   // If false, we cannot perform TRE on tail calls marked with the 'tail'
   // attribute, because doing so would cause the stack size to increase (real
   // TRE would deallocate variable sized allocas, TRE doesn't).
   bool CanTRETailMarkedCall = canTRE(F);
 
+  // Change any tail recursive calls to loops.
   TailRecursionEliminator TRE(F, TTI, AA, ORE, DTU);
 
-  // Change any tail recursive calls to loops.
-  //
-  // FIXME: The code generator produces really bad code when an 'escaping
-  // alloca' is changed from being a static alloca to being a dynamic alloca.
-  // Until this is resolved, disable this transformation if that would ever
-  // happen.  This bug is PR962.
-  for (Function::iterator BBI = F.begin(), E = F.end(); BBI != E; /*in loop*/) {
-    BasicBlock *BB = &*BBI++; // foldReturnAndProcessPred may delete BB.
-    if (ReturnInst *Ret = dyn_cast<ReturnInst>(BB->getTerminator())) {
-      bool Change = TRE.processReturningBlock(Ret, !CanTRETailMarkedCall);
-      if (!Change && BB->getFirstNonPHIOrDbg() == Ret)
-        Change = TRE.foldReturnAndProcessPred(Ret, !CanTRETailMarkedCall);
-      MadeChange |= Change;
-    }
-  }
+  for (BasicBlock &BB : F)
+    MadeChange |= TRE.processBlock(BB, !CanTRETailMarkedCall);
 
   TRE.cleanupAndFinalize();
 
   return MadeChange;
 }
 
 namespace {
 struct TailCallElim : public FunctionPass {