[TRE] Allow elimination when the returned value is non-constant

[laytonio]

Summary

Currently we can only eliminate call return pairs that either return the result of the call or a dynamic constant. This patch removes that limitation.

[efriedma]

Can you add a few testcases for eliminating multiple tail calls?

llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp

466

Maybe describe what specific instructions RetPN and RetKnownPN actually refer to?

600

Might as well use UndefValue::get here?

829

You can use dropAllReferences() to clear out the operands of an instruction before you erase it.

[laytonio]

Address comments. Added a test for multiple eliminations.

[efriedma]

LGTM. (I'll merge this for you.)

[laytonio]

Thank you. I appreciate the help.

[laytonio]

Was there anything else I needed to do for this?

[efriedma]

Sorry about the delay; it's merged now.

llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp

class TailRecursionEliminator {
 
   // The below are shared state we want to have available when eliminating any
   // calls in the function. There values should be populated by
   // createTailRecurseLoopHeader the first time we find a call we can eliminate.
   BasicBlock *HeaderBB = nullptr;
   SmallVector<PHINode *, 8> ArgumentPHIs;
   bool RemovableCallsMustBeMarkedTail = false;
 
+  // PHI node to store our return value.
+  PHINode *RetPN = nullptr;
+
+  // i1 PHI node to track if we have a valid return value stored in RetPN.
+  PHINode *RetKnownPN = nullptr;
+
+  // Vector of select instructions we insereted. These selects use RetKnownPN
+  // to either propagate RetPN or select a new return value.
+  SmallVector<SelectInst *, 8> RetSelects;
+
   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,
                              bool CannotTailCallElimCallsMarkedTail);
 
void TailRecursionEliminator::createTailRecurseLoopHeader(CallInst *CI) {
   Instruction *InsertPos = &HeaderBB->front();
   for (Function::arg_iterator I = F.arg_begin(), E = F.arg_end(); I != E; ++I) {
     PHINode *PN =
         PHINode::Create(I->getType(), 2, I->getName() + ".tr", InsertPos);
     I->replaceAllUsesWith(PN); // Everyone use the PHI node now!
     PN->addIncoming(&*I, NewEntry);
     ArgumentPHIs.push_back(PN);
   }
+
+  // If the function doen't return void, create the RetPN and RetKnownPN PHI
+  // nodes to track our return value. We initialize RetPN with undef and
+  // RetKnownPN with false since we can't know our return value at function
+  // entry.
+  Type *RetType = F.getReturnType();
+  if (!RetType->isVoidTy()) {
+    Type *BoolType = Type::getInt1Ty(F.getContext());
+    RetPN = PHINode::Create(RetType, 2, "ret.tr", InsertPos);
+    RetKnownPN = PHINode::Create(BoolType, 2, "ret.known.tr", InsertPos);
+
+    RetPN->addIncoming(UndefValue::get(RetType), NewEntry);
+    RetKnownPN->addIncoming(ConstantInt::getFalse(BoolType), NewEntry);
+  }
+
   // The entry block was changed from HeaderBB to NewEntry.
   // The forward DominatorTree needs to be recalculated when the EntryBB is
   // changed. In this corner-case we recalculate the entire tree.
   DTU.recalculate(*NewEntry->getParent());
 }
 
 PHINode *TailRecursionEliminator::insertAccumulator(
     Value *AccumulatorRecursionEliminationInitVal) {

 bool TailRecursionEliminator::eliminateCall(CallInst *CI) {
   ReturnInst *Ret = cast<ReturnInst>(CI->getParent()->getTerminator());
 
   // If we are introducing accumulator recursion to eliminate operations after
   // the call instruction that are both associative and commutative, the initial
   // value for the accumulator is placed in this variable.  If this value is set
   // then we actually perform accumulator recursion elimination instead of
   // simple tail recursion elimination.  If the operation is an LLVM instruction
-  // (eg: "add") then it is recorded in AccumulatorRecursionInstr.  If not, then
-  // we are handling the case when the return instruction returns a constant C
-  // which is different to the constant returned by other return instructions
-  // (which is recorded in AccumulatorRecursionEliminationInitVal).  This is a
-  // special case of accumulator recursion, the operation being "return C".
+  // (eg: "add") then it is recorded in AccumulatorRecursionInstr.
   Value *AccumulatorRecursionEliminationInitVal = nullptr;
   Instruction *AccumulatorRecursionInstr = nullptr;
 
   // Ok, we found a potential tail call.  We can currently only transform the
   // tail call if all of the instructions between the call and the return are
   // movable to above the call itself, leaving the call next to the return.
   // Check that this is the case now.
   BasicBlock::iterator BBI(CI);
if ((AccumulatorRecursionEliminationInitVal =
       // Yes, this is accumulator recursion.  Remember which instruction
       // accumulates.
       AccumulatorRecursionInstr = &*BBI;
     } else {
       return false;   // Otherwise, we cannot eliminate the tail recursion!
     }
   }
 
-  // We can only transform call/return pairs that either ignore the return value
-  // of the call and return void, ignore the value of the call and return a
-  // constant, return the value returned by the tail call, or that are being
-  // accumulator recursion variable eliminated.
-  if (Ret->getNumOperands() == 1 && Ret->getReturnValue() != CI &&
-      !isa<UndefValue>(Ret->getReturnValue()) &&
-      AccumulatorRecursionEliminationInitVal == nullptr &&
-      !getCommonReturnValue(nullptr, CI)) {
-    // One case remains that we are able to handle: the current return
-    // instruction returns a constant, and all other return instructions
-    // return a different constant.
-    if (!isDynamicConstant(Ret->getReturnValue(), CI, Ret))
-      return false; // Current return instruction does not return a constant.
-    // Check that all other return instructions return a common constant.  If
-    // so, record it in AccumulatorRecursionEliminationInitVal.
-    AccumulatorRecursionEliminationInitVal = getCommonReturnValue(Ret, CI);
-    if (!AccumulatorRecursionEliminationInitVal)
-      return false;
-  }
-
   BasicBlock *BB = Ret->getParent();
 
   using namespace ore;
   ORE->emit([&]() {
     return OptimizationRemark(DEBUG_TYPE, "tailcall-recursion", CI)
            << "transforming tail recursion into loop";
   });
 
bool TailRecursionEliminator::eliminateCall(CallInst *CI) {
   // do so now.  Note that we _know_ that no subsequent tail recursion
   // eliminations will happen on this function because of the way the
   // accumulator recursion predicate is set up.
   //
   if (AccumulatorRecursionEliminationInitVal) {
     PHINode *AccPN = insertAccumulator(AccumulatorRecursionEliminationInitVal);
 
     Instruction *AccRecInstr = AccumulatorRecursionInstr;
-    if (AccRecInstr) {
-      // Add an incoming argument for the current block, which is computed by
-      // our associative and commutative accumulator instruction.
-      AccPN->addIncoming(AccRecInstr, BB);
+
 
-      // Next, rewrite the accumulator recursion instruction so that it does not
-      // use the result of the call anymore, instead, use the PHI node we just
-      // inserted.
-      AccRecInstr->setOperand(AccRecInstr->getOperand(0) != CI, AccPN);
-    } else {
-      // Add an incoming argument for the current block, which is just the
-      // constant returned by the current return instruction.
-      AccPN->addIncoming(Ret->getReturnValue(), BB);
-    }
 
     // Finally, rewrite any return instructions in the program to return the PHI
     // node instead of the "initval" that they do currently.  This loop will
     // actually rewrite the return value we are destroying, but that's ok.
     for (BasicBlock &BBI : F)
       if (ReturnInst *RI = dyn_cast<ReturnInst>(BBI.getTerminator()))
         RI->setOperand(0, AccPN);
     ++NumAccumAdded;
   }
 
+  // Update our return value tracking
+  if (RetPN) {
+    if (Ret->getReturnValue() == CI || AccumulatorRecursionEliminationInitVal) {
+      // Defer selecting a return value
+      RetPN->addIncoming(RetPN, BB);
+      RetKnownPN->addIncoming(RetKnownPN, BB);
+    } else {
+      // We found a return value we want to use, insert a select instruction to
+      // select it if we don't already know what our return value will be and
+      // store the result in our return value PHI node.
+      SelectInst *SI = SelectInst::Create(
+          RetKnownPN, RetPN, Ret->getReturnValue(), "current.ret.tr", Ret);
+      RetSelects.push_back(SI);
+
+      RetPN->addIncoming(SI, BB);
+      RetKnownPN->addIncoming(ConstantInt::getTrue(RetKnownPN->getType()), BB);
+    }
+  }
+
   // Now that all of the PHI nodes are in place, remove the call and
   // ret instructions, replacing them with an unconditional branch.
   BranchInst *NewBI = BranchInst::Create(HeaderBB, Ret);
   NewBI->setDebugLoc(CI->getDebugLoc());
 
   BB->getInstList().erase(Ret);  // Remove return.
   BB->getInstList().erase(CI);   // Remove call.
   DTU.applyUpdates({{DominatorTree::Insert, BB, HeaderBB}});
void TailRecursionEliminator::cleanupAndFinalize() {
   // call.
   for (PHINode *PN : ArgumentPHIs) {
     // If the PHI Node is a dynamic constant, replace it with the value it is.
     if (Value *PNV = SimplifyInstruction(PN, F.getParent()->getDataLayout())) {
       PN->replaceAllUsesWith(PNV);
       PN->eraseFromParent();
     }
   }
+
+  if (RetPN) {
+    if (RetSelects.empty()) {
+      // If we didn't insert any select instructions, then we know we didn't
+      // store a return value and we can remove the PHI nodes we inserted.
+      RetPN->dropAllReferences();
+      RetPN->eraseFromParent();
+
+      RetKnownPN->dropAllReferences();
+      RetKnownPN->eraseFromParent();
+    } else {
+      // We need to insert a select instruction before any return left in the
+      // function to select our stored return value if we have one.
+      for (BasicBlock &BB : F) {
+        ReturnInst *RI = dyn_cast<ReturnInst>(BB.getTerminator());
+        if (!RI)
+          continue;
+
+        SelectInst *SI = SelectInst::Create(
+            RetKnownPN, RetPN, RI->getOperand(0), "current.ret.tr", RI);
+        RI->setOperand(0, SI);
+      }
+    }
+  }
 }
 
 bool TailRecursionEliminator::eliminate(Function &F,
                                         const TargetTransformInfo *TTI,
                                         AliasAnalysis *AA,
                                         OptimizationRemarkEmitter *ORE,
                                         DomTreeUpdater &DTU) {
   if (F.getFnAttribute("disable-tail-calls").getValueAsString() == "true")

llvm/test/Transforms/TailCallElim/2010-06-26-MultipleReturnValues.ll

 ; RUN: opt < %s -tailcallelim -verify-dom-info -S | FileCheck %s
 ; PR7328
 ; PR7506
-define i32 @foo(i32 %x) {
-; CHECK-LABEL: define i32 @foo(
-; CHECK: %accumulator.tr = phi i32 [ 1, %entry ], [ 0, %body ]
+define i32 @test1_constants(i32 %x) {
 entry:
   %cond = icmp ugt i32 %x, 0                      ; <i1> [#uses=1]
   br i1 %cond, label %return, label %body
 
 body:                                             ; preds = %entry
   %y = add i32 %x, 1                              ; <i32> [#uses=1]
-  %tmp = call i32 @foo(i32 %y)                    ; <i32> [#uses=0]
-; CHECK-NOT: call
+  %recurse = call i32 @test1_constants(i32 %y)        ; <i32> [#uses=0]
   ret i32 0
-; CHECK: ret i32 %accumulator.tr
 
 return:                                           ; preds = %entry
   ret i32 1
 }
+
+; CHECK-LABEL: define i32 @test1_constants(
+; CHECK: tailrecurse:
+; CHECK: %ret.tr = phi i32 [ undef, %entry ], [ %current.ret.tr, %body ]
+; CHECK: %ret.known.tr = phi i1 [ false, %entry ], [ true, %body ]
+; CHECK: body:
+; CHECK-NOT: %recurse
+; CHECK: %current.ret.tr = select i1 %ret.known.tr, i32 %ret.tr, i32 0
+; CHECK-NOT: ret
+; CHECK: return:
+; CHECK: %current.ret.tr1 = select i1 %ret.known.tr, i32 %ret.tr, i32 1
+; CHECK: ret i32 %current.ret.tr1
+
+define i32 @test2_non_constants(i32 %x) {
+entry:
+  %cond = icmp ugt i32 %x, 0
+  br i1 %cond, label %return, label %body
+
+body:
+  %y = add i32 %x, 1
+  %helper1 = call i32 @test2_helper()
+  %recurse = call i32 @test2_non_constants(i32 %y)
+  ret i32 %helper1
+
+return:
+  %helper2 = call i32 @test2_helper()
+  ret i32 %helper2
+}
+
+declare i32 @test2_helper()
+
+; CHECK-LABEL: define i32 @test2_non_constants(
+; CHECK: tailrecurse:
+; CHECK: %ret.tr = phi i32 [ undef, %entry ], [ %current.ret.tr, %body ]
+; CHECK: %ret.known.tr = phi i1 [ false, %entry ], [ true, %body ]
+; CHECK: body:
+; CHECK-NOT: %recurse
+; CHECK: %current.ret.tr = select i1 %ret.known.tr, i32 %ret.tr, i32 %helper1
+; CHECK-NOT: ret
+; CHECK: return:
+; CHECK: %current.ret.tr1 = select i1 %ret.known.tr, i32 %ret.tr, i32 %helper2
+; CHECK: ret i32 %current.ret.tr1
+
+define i32 @test3_mixed(i32 %x) {
+entry:
+  switch i32 %x, label %default [
+    i32 0, label %case0
+    i32 1, label %case1
+    i32 2, label %case2
+  ]
+
+case0:
+  %helper1 = call i32 @test3_helper()
+  br label %return
+
+case1:
+  %y1 = add i32 %x, -1
+  %recurse1 = call i32 @test3_mixed(i32 %y1)
+  br label %return
+
+case2:
+  %y2 = add i32 %x, -1
+  %helper2 = call i32 @test3_helper()
+  %recurse2 = call i32 @test3_mixed(i32 %y2)
+  br label %return
+
+default:
+  %y3 = urem i32 %x, 3
+  %recurse3 = call i32 @test3_mixed(i32 %y3)
+  br label %return
+
+return:
+  %retval = phi i32 [ %recurse3, %default ], [ %helper2, %case2 ], [ 9, %case1 ], [ %helper1, %case0 ]
+  ret i32 %retval
+}
+
+declare i32 @test3_helper()
+
+; CHECK-LABEL: define i32 @test3_mixed(
+; CHECK: tailrecurse:
+; CHECK: %ret.tr = phi i32 [ undef, %entry ], [ %current.ret.tr, %case1 ], [ %current.ret.tr1, %case2 ], [ %ret.tr, %default ]
+; CHECK: %ret.known.tr = phi i1 [ false, %entry ], [ true, %case1 ], [ true, %case2 ], [ %ret.known.tr, %default ]
+; CHECK: case1:
+; CHECK-NOT: %recurse
+; CHECK: %current.ret.tr = select i1 %ret.known.tr, i32 %ret.tr, i32 9
+; CHECK: br label %tailrecurse
+; CHECK: case2:
+; CHECK-NOT: %recurse
+; CHECK: %current.ret.tr1 = select i1 %ret.known.tr, i32 %ret.tr, i32 %helper2
+; CHECK: br label %tailrecurse
+; CHECK: default:
+; CHECK-NOT: %recurse
+; CHECK: br label %tailrecurse
+; CHECK: return:
+; CHECK: %current.ret.tr2 = select i1 %ret.known.tr, i32 %ret.tr, i32 %helper1
+; CHECK: ret i32 %current.ret.tr2

llvm/test/Transforms/TailCallElim/basic.ll

endif.0:    ; preds = %entry
   %DUMMY = add i32 %X, 1    ; <i32> [#uses=0]
   ret i32 %tmp.8
 }
 
 ; Though this case seems to be fairly unlikely to occur in the wild, someone
 ; plunked it into the demo script, so maybe they care about it.
 define i32 @test3(i32 %c) {
 ; CHECK: i32 @test3
+; CHECK: tailrecurse:
+; CHECK: %ret.tr = phi i32 [ undef, %entry ], [ %current.ret.tr, %else ]
+; CHECK: %ret.known.tr = phi i1 [ false, %entry ], [ true, %else ]
+; CHECK: else:
 ; CHECK-NOT: call
-; CHECK: ret i32 0
+; CHECK: %current.ret.tr = select i1 %ret.known.tr, i32 %ret.tr, i32 0
+; CHECK-NOT: ret
+; CHECK: return:
+; CHECK: %current.ret.tr1 = select i1 %ret.known.tr, i32 %ret.tr, i32 0
+; CHECK: ret i32 %current.ret.tr1
 entry:
   %tmp.1 = icmp eq i32 %c, 0    ; <i1> [#uses=1]
   br i1 %tmp.1, label %return, label %else
 else:   ; preds = %entry
   %tmp.5 = add i32 %c, -1   ; <i32> [#uses=1]
   %tmp.3 = call i32 @test3(i32 %tmp.5)    ; <i32> [#uses=0]
   ret i32 0
 return:   ; preds = %entry