[CodeGenPrepare] Also skip lifetime.end intrinsic when check return block in dupRetToEnableTailCallOpts.

Differential Revision: https://reviews.llvm.org/D95424

lib/CodeGen/CodeGenPrepare.cpp

@@ -2242,21 +2242,25 @@ bool CodeGenPrepare::dupRetToEnableTailCallOpts(BasicBlock *BB, bool &ModifiedDT
   if (PN && PN->getParent() != BB)
     return false;
 
-  // Make sure there are no instructions between the PHI and return, or that the
-  // return is the first instruction in the block.
-  if (PN) {
-    BasicBlock::iterator BI = BB->begin();
-    // Skip over debug and the bitcast.
-    do {
-      ++BI;
-    } while (isa<DbgInfoIntrinsic>(BI) || &*BI == BCI || &*BI == EVI ||
-             isa<PseudoProbeInst>(BI));
-    if (&*BI != RetI)
-      return false;
-  } else {
-    if (BB->getFirstNonPHIOrDbg(true) != RetI)
-      return false;
-  }
+  auto isLifetimeEndOrBitCastFor = [](const Instruction *Inst) {
+    const BitCastInst *BC = dyn_cast<BitCastInst>(Inst);
+    if (BC && BC->hasOneUse())
+      Inst = BC->user_back();
+
+    if (const IntrinsicInst *II = dyn_cast<IntrinsicInst>(Inst))
+      return II->getIntrinsicID() == Intrinsic::lifetime_end;
+    return false;
+  };
+
+  // Make sure there are no instructions between the first instruction
+  // and return.
+  const Instruction *BI = BB->getFirstNonPHI();
+  // Skip over debug and the bitcast.
+  while (isa<DbgInfoIntrinsic>(BI) || BI == BCI || BI == EVI ||
+         isa<PseudoProbeInst>(BI) || isLifetimeEndOrBitCastFor(BI))
+    BI = BI->getNextNode();
+  if (BI != RetI)
+    return false;
 
   /// Only dup the ReturnInst if the CallInst is likely to be emitted as a tail
   /// call.

test/Transforms/CodeGenPrepare/ARM/tailcall-dup.ll

@@ -4,6 +4,8 @@ target triple = "armv8m.main-none-eabi"
 
 declare i8* @f0()
 declare i8* @f1()
+declare void @llvm.lifetime.start.p0i8(i64, i8* nocapture) nounwind
+declare void @llvm.lifetime.end.p0i8(i64, i8* nocapture) nounwind
 
 define i8* @tail_dup() {
 ; CHECK-LABEL: tail_dup
@@ -12,6 +14,9 @@ define i8* @tail_dup() {
 ; CHECK: tail call i8* @f1()
 ; CHECK-NEXT: ret i8*
 bb0:
+  %a = alloca i32
+  %a1 = bitcast i32* %a to i8*
+  call void @llvm.lifetime.start.p0i8(i64 -1, i8* %a1) nounwind
   %tmp0 = tail call i8* @f0()
   br label %return
 bb1:
@@ -19,6 +24,8 @@ bb1:
   br label %return
 return:
   %retval = phi i8* [ %tmp0, %bb0 ], [ %tmp1, %bb1 ]
+  %a2 = bitcast i32* %a to i8*
+  call void @llvm.lifetime.end.p0i8(i64 -1, i8* %a2) nounwind
   ret i8* %retval
 }
 

test/Transforms/CodeGenPrepare/X86/tailcall-assume-xbb.ll

@@ -13,6 +13,8 @@
 
 define i8* @foo(i64 %size, i64 %v1, i64 %v2) {
 entry:
+  %a = alloca i8
+  call void @llvm.lifetime.start.p0i8(i64 -1, i8* %a) nounwind
   %cmp1 = icmp ult i64 %size, 1025
   br i1 %cmp1, label %if.end, label %case1
 
@@ -40,9 +42,12 @@ exit1:
 
 exit2:
   %retval2 = phi i8* [ %ret1, %case1 ], [ %retval1, %exit1 ]
+  call void @llvm.lifetime.end.p0i8(i64 -1, i8* %a) nounwind
   ret i8* %retval2
 }
 
 declare void @llvm.assume(i1)
 declare i8* @qux()
 declare i8* @bar()
+declare void @llvm.lifetime.start.p0i8(i64, i8* nocapture) nounwind
+declare void @llvm.lifetime.end.p0i8(i64, i8* nocapture) nounwind