CodeGenPrep: insert return after tail calls that return their arg.

[t.p.northover]

Summary

When a call returns one of its arguments this is considered a cast by stripPointerCasts. So IR returning that directly from a tail call (which should be fine) can become muddied and seem to return an argument from earlier in the function.

This adds some simple checks to spot this situation and reclassify it as the original "tail call -> phi -> return" dataflow.

[nikic]

Possibly we should just drop the stripPointerCasts()? This is probably just not relevant anymore with opaque pointers.

llvm/lib/CodeGen/CodeGenPrepare.cpp

bool CodeGenPrepare::dupRetToEnableTailCallOpts(BasicBlock *BB,
   const Function *F = BB->getParent();
   SmallVector<BasicBlock *, 4> TailCallBBs;
   if (PN) {
     for (unsigned I = 0, E = PN->getNumIncomingValues(); I != E; ++I) {
       // Look through bitcasts.
       Value *IncomingVal = PN->getIncomingValue(I)->stripPointerCasts();
       CallInst *CI = dyn_cast<CallInst>(IncomingVal);
       BasicBlock *PredBB = PN->getIncomingBlock(I);
+
+      // LLVM treats calls that return one of their arguments as casts, and
+      // forwardoring/looking through that may split a tail call from the value
+      // it returns even more than a phi. Attempt to undo this by looking for a
+      // suitable call in the phi's predecessor block.
+      if (!CI) {
+        BranchInst *BI = dyn_cast<BranchInst>(PredBB->getTerminator());
+        if (!BI || !BI->isUnconditional())
+          continue;
+        CallInst *CastCI =
+            dyn_cast_or_null<CallInst>(BI->getPrevNonDebugInstruction());
+        if (!CastCI || CastCI->getReturnedArgOperand() != IncomingVal)
+          continue;
+        CI = CastCI;
+      }
+
       // Make sure the phi value is indeed produced by the tail call.
-      if (CI && CI->hasOneUse() && CI->getParent() == PredBB &&
+      if (CI && CI->getParent() == PredBB &&
           TLI->mayBeEmittedAsTailCall(CI) &&
           attributesPermitTailCall(F, CI, RetI, *TLI))
         TailCallBBs.push_back(PredBB);
     }
   } else {
     SmallPtrSet<BasicBlock *, 4> VisitedBBs;
     for (BasicBlock *Pred : predecessors(BB)) {
       if (!VisitedBBs.insert(Pred).second)

llvm/test/CodeGen/AArch64/tailcall-returned.ll

+; RUN: opt --codegenprepare -mtriple=arm64-apple-macosx %s -S -o - | FileCheck %s
+
+define ptr @test_already_opt(ptr %in, i1 %tst) {
+; CHECK-LABEL: @test_already_opt
+  br i1 %tst, label %true, label %false
+
+true:
+; CHECK: tail call ptr @return_input(ptr %in)
+; CHECK-NEXT: ret ptr %in
+
+  %call = tail call ptr @return_input(ptr %in)
+  br label %done
+
+false:
+  br label %done
+
+done:
+  %res = phi ptr [ %in, %true ],  [ null, %false ]
+  ret ptr %res
+}
+
+define ptr @test_simple(ptr %in, i1 %tst) {
+; CHECK-LABEL: @test_simple
+  br i1 %tst, label %true, label %false
+
+true:
+; CHECK: tail call ptr @return_input(ptr %in)
+; CHECK-NEXT: ret ptr %call
+
+  %call = tail call ptr @return_input(ptr %in)
+  br label %done
+
+false:
+  br label %done
+
+done:
+  %res = phi ptr [ %call, %true ],  [ null, %false ]
+  ret ptr %res
+}
+
+define ptr @test_unsafe(ptr %in, i1 %tst) {
+; CHECK-LABEL: @test_unsafe
+  br i1 %tst, label %true, label %false
+
+true:
+; CHECK: tail call ptr @return_rando(ptr %in)
+; CHECK-NEXT: br label %done
+
+  %call = tail call ptr @return_rando(ptr %in)
+  br label %done
+
+false:
+  br label %done
+
+done:
+  %res = phi ptr [ %in, %true ],  [ null, %false ]
+  ret ptr %res
+}
+
+declare ptr @return_input(ptr returned)
+declare ptr @return_rando(ptr)