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

[junparser]

Summary

As we find in many real workloads that there are maybe some lifetime.end intrinsics existing in return block, so it is better to skip lifetime.end intrinsic as well. This patch just do such thing.

TestPlan: check-llvm

[junparser]

Update testcase.

[junparser]

hi @aqjune @bkramer kindly ping~

[aqjune]

Hi,
This transformation is basically removing lifetime.end, is my understanding right?

In terms of correctness, this is fine IMO because removing lifetime.end makes the program more defined.
It is also syntactically fine to have an unmatched lifetime.start because LLVM already removes unreachable code that may contain lifetime.end.

[junparser]

In D95424#2527105, @aqjune wrote:

Hi,
This transformation is basically removing lifetime.end, is my understanding right?

Yes, It remove lifetime.end in the new duplicated return block.

In terms of correctness, this is fine IMO because removing lifetime.end makes the program more defined.
It is also syntactically fine to have an unmatched lifetime.start because LLVM already removes unreachable code that may contain lifetime.end.

Yes, you are right, also IMO lifetime.end does not have any effect in return block.

[junparser]

@xbolva00 @Carrot Would you take a look at this patch? Since I have seen you also deal with tailcall in D76539.

llvm/lib/CodeGen/CodeGenPrepare.cpp

if (V) {
     PN = dyn_cast<PHINode>(V);
     if (!PN)
       return false;
   }
 
   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;
+  };
+
+  // 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)
 
   /// Only dup the ReturnInst if the CallInst is likely to be emitted as a tail
   /// call.
   const Function *F = BB->getParent();
   SmallVector<BasicBlock*, 4> TailCallBBs;
   if (PN) {
     for (unsigned I = 0, E = PN->getNumIncomingValues(); I != E; ++I) {
       // Look through bitcasts.

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

 ; RUN: opt -codegenprepare -S < %s | FileCheck %s
 
 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
 ; CHECK: tail call i8* @f0()
 ; CHECK-NEXT: ret i8*
 ; 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:
   %tmp1 = tail call i8* @f1()
   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
 }
 
 define nonnull i8* @nonnull_dup() {
 ; CHECK-LABEL: nonnull_dup
 ; CHECK: tail call i8* @f0()
 ; CHECK-NEXT: ret i8*
 ; CHECK: tail call i8* @f1()

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

 ; RUN: opt -codegenprepare -S -mtriple=x86_64-linux < %s | FileCheck %s
 
 ; The ret instruction can be duplicated into BB case2 even though there is an
 ; intermediate BB exit1 and call to llvm.assume.
 
 @ptr = external global i8*, align 8
 
 ; CHECK:       %ret1 = tail call i8* @qux()
 ; CHECK-NEXT:  ret i8* %ret1
 
 ; CHECK:       %ret2 = tail call i8* @bar()
 ; CHECK-NEXT:  ret i8* %ret2
 
 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
 
 case1:
   %ret1 = tail call i8* @qux()
   br label %exit2
 
 if.end:

 exit1:
   %retval1 = phi i8* [ %ret2, %case2 ], [ %ret3, %case3 ]
   %cmp3 = icmp ne i8* %retval1, null
   tail call void @llvm.assume(i1 %cmp3)
   br label %exit2
 
 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