[TRE] Fix bug in handling of switch statements

[laytonio]

Summary

Currently, isDynamicConstant tries (incorrectly) to check whether we could only have reached the current block from a switch on the value in the previous block. If the block is determined to only be reachable from one case, the value is being used as if it were a dynamic constant (which it may not be), instead of being replaced by the constant. This patch fixes the check, and also does the replacement if the check succeeds. Also, disable a test that relied on the buggy behavior.

[efriedma]

Are you sure that actually fixes the issue? On trunk, the following also crashes:

define i32 @f() local_unnamed_addr {
entry:
  %call = call i32 @g()
  switch i32 %call, label %sw.default [
    i32 1, label %cleanup
  ]

sw.default:
  %call1 = call i32 @f()
  %add = add nsw i32 %call1, 1
  br label %cleanup

cleanup:
  %retval.0 = phi i32 [ %add, %sw.default ], [ %call, %entry ]
  ret i32 %retval.0
}

declare i32 @g()

Also, missing testcase.

[laytonio]

This now fixes both repos of the crash. I removed a test that relied on the bug and actually didn't produce valid output. Added the case with a single branch as a test case. However, all this switch handling really does is propagate a constant. Maybe we should remove it entirely and expect that the constant propagation passes do this for us.

[laytonio]

llvm/test/Transforms/TailCallElim/accum_recursion.ll

48

Its not valid for us to use %n here as our initial value, since %n would change with each iteration of the recursion.

52–53

If only one of these cases resulted in this branch we could replace the %n (see above comment) with the constant. However, we don't whether to choose the 1 or the 0.

[efriedma]

However, all this switch handling really does is propagate a constant. Maybe we should remove it entirely and expect that the constant propagation passes do this for us.

This sort of depends on how the interaction with other passes works out. A switch with one case will be turned into a branch by SimplifyCFG, so it's not worth optimizing. If it's realistic to have a switch where only one of the cases goes to a return instruction, maybe worth keeping the switch handling around; SimplifyCFG will often prefer a unified "ret" instruction. Probably want to check what the input to -tailcallelim looks like, realistically.

---

In terms of whether we could actually transform my original example... well, it sort of goes back to what I was saying on D78259. Currently, we set the initial value of the "accumulator" to the value returned in the base case. So if the base case isn't a value we can materialize in the entry block, we can't do the tail call transform. But I'm not sure that's a fundamental limitation. Suppose, instead, the accumulator always starts at zero. Then, just before we return, we add the value of the base case. It's still basically the same transform, but it's more flexible: you can drop the whole isDynamicConstant() check, I think. Does that make sense?

To write this out:

Input

define i32 @f() local_unnamed_addr {
entry:
  %call = call i32 @g()
  switch i32 %call, label %sw.default [
    i32 1, label %cleanup
    i32 2, label %cleanup
  ]

sw.default:
  %call1 = call i32 @f()
  %add = add nsw i32 %call1, 1
  ret i32 %add

cleanup:
  ret i32 %call
}

declare i32 @g()

Current, invalid -tailcallelim output on trunk:

define i32 @f() local_unnamed_addr {
entry:
  br label %tailrecurse

tailrecurse:                                      ; preds = %sw.default, %entry
  %accumulator.tr = phi i32 [ %call, %entry ], [ %add, %sw.default ]
  %call = tail call i32 @g()
  switch i32 %call, label %sw.default [
    i32 1, label %cleanup
    i32 2, label %cleanup
  ]

sw.default:                                       ; preds = %tailrecurse
  %add = add nsw i32 %accumulator.tr, 1
  br label %tailrecurse

cleanup:                                          ; preds = %tailrecurse, %tailrecurse
  ret i32 %accumulator.tr
}

declare i32 @g()

Alternative transform:

define i32 @f() local_unnamed_addr {
entry:
  br label %tailrecurse

tailrecurse:                                      ; preds = %sw.default, %entry
  %accumulator.tr = phi i32 [ 0, %entry ], [ %add, %sw.default ]
  %call = tail call i32 @g()
  switch i32 %call, label %sw.default [
    i32 1, label %cleanup
    i32 2, label %cleanup
  ]

sw.default:                                       ; preds = %tailrecurse
  %add = add nsw i32 %accumulator.tr, 1
  br label %tailrecurse

cleanup:                                          ; preds = %tailrecurse, %tailrecurse
  %accumulator.ret = add i32 %accumulator.tr, %call
  ret i32 %accumulator.ret
}

declare i32 @g()

llvm/test/Transforms/TailCallElim/accum_recursion.ll

44

Please keep this function as a testcase, even if we're just keeping it an example of something we currently can't transform.

[laytonio]

Suppose, instead, the accumulator always starts at zero.

The accumulator has to start with a value that won't effect the computation. Zero works in the case where we are doing addition, but would not work if we were doing multiplication. We could possibly implement a version that looks at the operation and then decides a starting value. Alternatively, and what I am working on is, we could basically inline one iteration though the loop and start the accumulator with the value you get from the first iteration. I have a version of this working, that does transform your original example. I still wanted to keep the current implementation where it applies though, because there is no need to do the inlining if we can already find the value.

[efriedma]

The accumulator has to start with a value that won't effect the computation. Zero works in the case where we are doing addition, but would not work if we were doing multiplication. We could possibly implement a version that looks at the operation and then decides a starting value.

tailcallelim currently only recognizes a very restrictive set of accumulator operations: ones where BinaryOperator::isAssociative returns true. ConstantExpr::getBinOpIdentity() can compute an appropriate identity value for all those operations.

Alternatively, and what I am working on is, we could basically inline one iteration though the loop and start the accumulator with the value you get from the first iteration.

Compiler jargon for making a copy of the loop for the first iteration is "loop peeling".

In general, it's more restrictive to require an identity value, yes. Peeling the first iteration of the loop provides an alternative transform in those cases. But I can't think of any realistic case where we could prove an operation is associative without being able to compute an identity value. And we don't really want to peel loops unless the peeled loop is significantly faster.

If we're not peeling the loop, I can't think of any other reason to keep the isDynamicConstant path around.

[laytonio]

I'm trying this approach now, but I'm having a hard time seeing a good way to transform this case:

define i32 @func(i32 %index) local_unnamed_addr {
entry:
  %0 = icmp eq i32 %index, 0
  br i1 %0, label %then, label %else

then:
  ret i32 12

else:
  %1 = call i32 @func(i32 0)
  ret i32 3
}

Current transform:

define i32 @func(i32 %index) local_unnamed_addr {
entry:
  br label %tailrecurse

tailrecurse:                                      ; preds = %else, %entry
  %accumulator.tr = phi i32 [ 12, %entry ], [ 3, %else ]
  %index.tr = phi i32 [ %index, %entry ], [ 0, %else ]
  %0 = icmp eq i32 %index.tr, 0
  br i1 %0, label %then, label %else

then:                                             ; preds = %tailrecurse
  ret i32 %accumulator.tr

else:                                             ; preds = %tailrecurse
  br label %tailrecurse
}

We could do something like this:

define i32 @func(i32 %index) local_unnamed_addr {
entry:
  br label %tailrecurse

tailrecurse:                                      ; preds = %else, %entry
  %selector.tr = phi i1 [ true, %entry ], [ false, %else ]
  %accumulator.tr = phi i32 [ 0, %entry ], [ %selection.tr1, %else ]
  %index.tr = phi i32 [ %index, %entry ], [ 0, %else ]
  %0 = icmp eq i32 %index.tr, 0
  br i1 %0, label %then, label %else

then:                                             ; preds = %tailrecurse
  %selection.tr = select i1 %selector.tr, i32 12, i32 %accumulator.tr
  ret i32 %selection.tr

else:                                             ; preds = %tailrecurse
  %selection.tr1 = select i1 %selector.tr, i32 3, i32 %accumulator.tr
  br label %tailrecurse
}

But that really doesn't seem all that clean to me. Is there a better way to do this that I am missing?

[efriedma]

In general, I think it has to look something like that, yes. You have to keep the value from the first iteration in a PHI node like %accumulator.tr. You need a select in the loop to consistently take the value from the first iteration. And you need a select before the return to check if there's a valid value in %accumulator.tr. Otherwise, the dominance doesn't work out in general, and you're back in isDynamicConstant() territory. Later loop optimizations should be able to simplify it to something sane in most cases.

[laytonio]

I have gotten a version of this that does the accumulation at the end mostly working. I will finish cleaning it up and try to split it into a few different patches as I reworked a fair bit, and found a few addition improvements I think we can make. I went ahead and updated this diff in case we want to consider it a first step of incremental improvement.

[laytonio]

Fixed by D80844

llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp

static bool canMoveAboveCall(Instruction *I, CallInst *CI, AliasAnalysis *AA) {
   return !is_contained(I->operands(), CI);
 }
 
 /// Return true if the specified value is the same when the return would exit
 /// as it was when the initial iteration of the recursive function was executed.
 ///
 /// We currently handle static constants and arguments that are not modified as
 /// part of the recursion.
-static bool isDynamicConstant(Value *V, CallInst *CI, ReturnInst *RI) {
-  if (isa<Constant>(V)) return true; // Static constants are always dyn consts
+static Value *getDynamicConstantOrReplacement(Value *V, CallInst *CI,
+                                              ReturnInst *RI) {
+  if (isa<Constant>(V))
+    return V; // Static constants are always dyn consts
 
   // Check to see if this is an immutable argument, if so, the value
   // will be available to initialize the accumulator.
   if (Argument *Arg = dyn_cast<Argument>(V)) {
     // Figure out which argument number this is...
     unsigned ArgNo = 0;
     Function *F = CI->getParent()->getParent();
     for (Function::arg_iterator AI = F->arg_begin(); &*AI != Arg; ++AI)
       ++ArgNo;
 
     // If we are passing this argument into call as the corresponding
     // argument operand, then the argument is dynamically constant.
     // Otherwise, we cannot transform this function safely.
     if (CI->getArgOperand(ArgNo) == Arg)
-      return true;
+      return Arg;
   }
 
   // Switch cases are always constant integers. If the value is being switched
   // on and the return is only reachable from one of its cases, it's
   // effectively constant.
   if (BasicBlock *UniquePred = RI->getParent()->getUniquePredecessor())
     if (SwitchInst *SI = dyn_cast<SwitchInst>(UniquePred->getTerminator()))
       if (SI->getCondition() == V)
-        return SI->getDefaultDest() != RI->getParent();
+        return SI->findCaseDest(RI->getParent());
 
   // Not a constant or immutable argument, we can't safely transform.
-  return false;
+  return nullptr;
 }
 
 /// Check to see if the function containing the specified tail call consistently
 /// returns the same runtime-constant value at all exit points except for
 /// IgnoreRI. If so, return the returned value.
 static Value *getCommonReturnValue(ReturnInst *IgnoreRI, CallInst *CI) {
   Function *F = CI->getParent()->getParent();
   Value *ReturnedValue = nullptr;
 
   for (BasicBlock &BBI : *F) {
     ReturnInst *RI = dyn_cast<ReturnInst>(BBI.getTerminator());
     if (RI == nullptr || RI == IgnoreRI) continue;
 
     // We can only perform this transformation if the value returned is
     // evaluatable at the start of the initial invocation of the function,
     // instead of at the end of the evaluation.
     //
-    Value *RetOp = RI->getOperand(0);
-    if (!isDynamicConstant(RetOp, CI, RI))
+    Value *V = getDynamicConstantOrReplacement(RI->getReturnValue(), CI, RI);
+    if (!V)
       return nullptr;
 
-    if (ReturnedValue && RetOp != ReturnedValue)
+    if (ReturnedValue && V != ReturnedValue)
       return nullptr;     // Cannot transform if differing values are returned.
-    ReturnedValue = RetOp;
+    ReturnedValue = V;
   }
   return ReturnedValue;
 }
 
 /// If the specified instruction can be transformed using accumulator recursion
 /// elimination, return the constant which is the start of the accumulator
 /// value.  Otherwise return null.
 static Value *canTransformAccumulatorRecursion(Instruction *I, CallInst *CI) {

 static bool eliminateRecursiveTailCall(
     CallInst *CI, ReturnInst *Ret, BasicBlock *&OldEntry,
     bool &TailCallsAreMarkedTail, SmallVectorImpl<PHINode *> &ArgumentPHIs,
     AliasAnalysis *AA, OptimizationRemarkEmitter *ORE, DomTreeUpdater &DTU) {
   // 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".
+  // simple tail recursion elimination.
   Value *AccumulatorRecursionEliminationInitVal = nullptr;
+
+  // If the operation is an LLVM instruction (eg: "add") then it is recorded in
+  // AccumulatorRecursionInstr.
   Instruction *AccumulatorRecursionInstr = nullptr;
 
+  // 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". In this case store the returned constant C in
+  // AccumulatorRecursionEliminationReturnedConstant
+  Value *AccumulatorRecursionEliminationReturnedConstant = 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);
   for (++BBI; &*BBI != Ret; ++BBI) {
     if (canMoveAboveCall(&*BBI, CI, AA))
       continue;
static bool eliminateRecursiveTailCall(
   // 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.
+    AccumulatorRecursionEliminationReturnedConstant =
+        getDynamicConstantOrReplacement(Ret->getReturnValue(), CI, Ret);
+
+    if (!AccumulatorRecursionEliminationReturnedConstant)
+      return false;
+
     // 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();
if (AccRecInstr) {
 
       // 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);
+      AccPN->addIncoming(AccumulatorRecursionEliminationReturnedConstant, 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);

llvm/test/Transforms/TailCallElim/accum_recursion.ll

return:   ; preds = %entry
   ret i32 %x
 }
 
 ; CHECK-LABEL: define i32 @test2_mul(
 ; CHECK: phi i32
 ; CHECK-NOT: call i32
 ; CHECK: return:
 
+; We don't currently transform the below function but would like to.
 
 define i64 @test3_fib(i64 %n) nounwind readnone {
-; CHECK-LABEL: @test3_fib(
 entry:
-; CHECK: tailrecurse:
-; CHECK: %accumulator.tr = phi i64 [ %n, %entry ], [ %3, %bb1 ]
-; CHECK: %n.tr = phi i64 [ %n, %entry ], [ %2, %bb1 ]
   switch i64 %n, label %bb1 [
-; CHECK: switch i64 %n.tr, label %bb1 [
     i64 0, label %bb2
     i64 1, label %bb2
   ]
 
 bb1:
-; CHECK: bb1:
   %0 = add i64 %n, -1
-; CHECK: %0 = add i64 %n.tr, -1
   %1 = tail call i64 @test3_fib(i64 %0) nounwind
-; CHECK: %1 = tail call i64 @test3_fib(i64 %0)
   %2 = add i64 %n, -2
-; CHECK: %2 = add i64 %n.tr, -2
   %3 = tail call i64 @test3_fib(i64 %2) nounwind
-; CHECK-NOT: tail call i64 @test3_fib
   %4 = add nsw i64 %3, %1
-; CHECK: add nsw i64 %accumulator.tr, %1
   ret i64 %4
-; CHECK: br label %tailrecurse
 
 bb2:
-; CHECK: bb2:
   ret i64 %n
-; CHECK: ret i64 %accumulator.tr
 }
+
+define i32 @test4_switch() local_unnamed_addr {
+entry:
+  %call = call i32 @test4_helper()
+  switch i32 %call, label %sw.default [
+    i32 1, label %cleanup
+  ]
+
+sw.default:
+  %call1 = call i32 @test4_switch()
+  %add = add nsw i32 %call1, 1
+  br label %cleanup
+
+cleanup:
+  %retval.0 = phi i32 [ %add, %sw.default ], [ %call, %entry ]
+  ret i32 %retval.0
+}
+
+declare i32 @test4_helper()
+
+; CHECK-LABEL: define i32 @test4_switch(
+; CHECK: tailrecurse:
+; CHECK: %accumulator.tr = phi i32 [ 1, %entry ],
+; CHECK: sw.default:
+; CHECK-NOT: call i32
+; CHECK: cleanup: