[Coroutines] Handle lifetime markers, bitcast and unused instruciton for symmetric transfer

[ChuanqiXu]

Summary

This fixes bug49888. The root cause for this is that simplifyTerminatorLeadingToRet didn't handle lifetime markers well. Another issue also noted in D116327 is that we deleted some inlined optimization pass in CoroSplit so that simplifyTerminatorLeadingToRet need to remove dead instructions by hand.

This patch fixes bug49888 by skipping lifetime markers and bitcast instruction and removing dead instructions by hand in simplifyTerminatorLeadingToRet.

This depends on https://reviews.llvm.org/D116327

[ChuanqiXu]

@rjmccall gentle ping~

[ChuanqiXu]

Rebased.

[junparser]

LGTM with some comments.

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

1251

GetFirstValidInstructionInBB seems can be removed. GetFirstValidInstruction(Succ->getFirstNonPHI()) is enough.

1303

ditto.

llvm/test/Transforms/Coroutines/coro-split-musttail6.ll

43

FIXME should be removed？

[ChuanqiXu]

Address comments.

[ChuanqiXu]

Clean codes

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

 // instruction. Suspend instruction represented by a switch, track the PHI
 // values and select the correct case successor when possible.
 static bool simplifyTerminatorLeadingToRet(Instruction *InitialInst) {
   DenseMap<Value *, Value *> ResolvedValues;
   BasicBlock *UnconditionalSucc = nullptr;
   assert(InitialInst->getModule());
   const DataLayout &DL = InitialInst->getModule()->getDataLayout();
 
+  auto GetFirstValidInstruction = [](Instruction *I) {
+    while (I) {
+      // BitCastInst wouldn't generate actual code so that we could skip it.
+      if (isa<BitCastInst>(I) || I->isDebugOrPseudoInst() ||
+          I->isLifetimeStartOrEnd())
+        I = I->getNextNode();
+      else if (isInstructionTriviallyDead(I))
+        // Duing we are in the middle of the transformation, we need to erase
+        // the dead instruction manually.
+        I = &*I->eraseFromParent();
+      else
+        break;
+    }
+    assert(I);
+    return I;
+  };
+
+  auto GetFirstValidInstructionInBB =
+      [&GetFirstValidInstruction](BasicBlock *BB) {
+        return GetFirstValidInstruction(BB->getFirstNonPHIOrDbgOrLifetime());
+      };
+
   auto TryResolveConstant = [&ResolvedValues](Value *V) {
     auto It = ResolvedValues.find(V);
     if (It != ResolvedValues.end())
       V = It->second;
     return dyn_cast<ConstantInt>(V);
   };
 
   Instruction *I = InitialInst;
-  while (I->isTerminator() ||
-         (isa<CmpInst>(I) && I->getNextNode()->isTerminator())) {
+  while (I->isTerminator() || isa<CmpInst>(I)) {
     if (isa<ReturnInst>(I)) {
       if (I != InitialInst) {
         // If InitialInst is an unconditional branch,
         // remove PHI values that come from basic block of InitialInst
         if (UnconditionalSucc)
           UnconditionalSucc->removePredecessor(InitialInst->getParent(), true);
         ReplaceInstWithInst(InitialInst, I->clone());
       }
       return true;
     }
     if (auto *BR = dyn_cast<BranchInst>(I)) {
       if (BR->isUnconditional()) {
         BasicBlock *Succ = BR->getSuccessor(0);
         if (I == InitialInst)
           UnconditionalSucc = Succ;
         scanPHIsAndUpdateValueMap(I, Succ, ResolvedValues);
-        I = Succ->getFirstNonPHIOrDbgOrLifetime();
+        I = GetFirstValidInstructionInBB(Succ);
         continue;
       }
 
       BasicBlock *BB = BR->getParent();
       // Handle the case the condition of the conditional branch is constant.
       // e.g.,
       //
       //     br i1 false, label %cleanup, label %CoroEnd
       //
       // It is possible during the transformation. We could continue the
       // simplifying in this case.
       if (ConstantFoldTerminator(BB, /*DeleteDeadConditions=*/true)) {
         // Handle this branch in next iteration.
         I = BB->getTerminator();
         continue;
       }
     } else if (auto *CondCmp = dyn_cast<CmpInst>(I)) {
-      // If the case number of suspended switch instruction is reduced to
-      // 1, then it is simplified to CmpInst in llvm::ConstantFoldTerminator.
-      auto *BR = dyn_cast<BranchInst>(I->getNextNode());
+      auto *BR = dyn_cast<BranchInst>(
+          GetFirstValidInstruction(CondCmp->getNextNode()));
       if (!BR || !BR->isConditional() || CondCmp != BR->getCondition())
         return false;
 
       // And the comparsion looks like : %cond = icmp eq i8 %V, constant.
       // So we try to resolve constant for the first operand only since the
       // second operand should be literal constant by design.
       ConstantInt *Cond0 = TryResolveConstant(CondCmp->getOperand(0));
       auto *Cond1 = dyn_cast<ConstantInt>(CondCmp->getOperand(1));
if (auto *BR = dyn_cast<BranchInst>(I)) {
       continue;
     } else if (auto *SI = dyn_cast<SwitchInst>(I)) {
       ConstantInt *Cond = TryResolveConstant(SI->getCondition());
       if (!Cond)
         return false;
 
       BasicBlock *BB = SI->findCaseValue(Cond)->getCaseSuccessor();
       scanPHIsAndUpdateValueMap(I, BB, ResolvedValues);
-      I = BB->getFirstNonPHIOrDbgOrLifetime();
+      I = GetFirstValidInstructionInBB(BB);
       continue;
     }
+
     return false;
   }
   return false;
 }
 
 // Check whether CI obeys the rules of musttail attribute.
 static bool shouldBeMustTail(const CallInst &CI, const Function &F) {
   if (CI.isInlineAsm())

llvm/test/Transforms/Coroutines/coro-split-musttail5.ll

await.ready:
   call void @consume(i8* %alloc.var)
   call void @llvm.lifetime.end.p0i8(i64 1, i8* %alloc.var)
   br label %exit
 exit:
   call i1 @llvm.coro.end(i8* null, i1 false)
   ret void
 }
 
-; FIXME: The fakeresume1 here should be marked as musttail.
 ; Verify that in the resume part resume call is marked with musttail.
 ; CHECK-LABEL: @g.resume(
-; CHECK-NOT: musttail call fastcc void @fakeresume1(i64* align 8 null)
+; CHECK:          musttail call fastcc void @fakeresume1(i64* align 8 null)
+; CHECK-NEXT:     ret void
 
 declare token @llvm.coro.id(i32, i8* readnone, i8* nocapture readonly, i8*) #1
 declare i1 @llvm.coro.alloc(token) #2
 declare i64 @llvm.coro.size.i64() #3
 declare i8* @llvm.coro.begin(token, i8* writeonly) #2
 declare token @llvm.coro.save(i8*) #2
 declare i8* @llvm.coro.frame() #3
 declare i8 @llvm.coro.suspend(token, i1) #2

llvm/test/Transforms/Coroutines/coro-split-musttail6.ll

await.ready:
   call void @consume(i64* %alloc.var)
   call void @llvm.lifetime.end.p0i8(i64 1, i8* %alloca.var.i8)
   br label %exit
 exit:
   call i1 @llvm.coro.end(i8* null, i1 false)
   ret void
 }
 
 ; FIXME: The fakeresume1 here should be marked as musttail.
 ; Verify that in the resume part resume call is marked with musttail.
 ; CHECK-LABEL: @g.resume(
-; CHECK-NOT: musttail call fastcc void @fakeresume1(i64* align 8 null)
+; CHECK:      musttail call fastcc void @fakeresume1(i64* align 8 null)
+; CHECK-NEXT: ret void
 
 ; It has a cleanup bb.
 define void @f() #0 {
 entry:
   %id = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* null)
   %alloc = call i8* @malloc(i64 16) #3
   %alloc.var = alloca i64
   %alloca.var.i8 = bitcast i64* %alloc.var to i8*

 exit:
   call i1 @llvm.coro.end(i8* null, i1 false)
   ret void
 }
 
 ; FIXME: The fakeresume1 here should be marked as musttail.
 ; Verify that in the resume part resume call is marked with musttail.
 ; CHECK-LABEL: @f.resume(
-; CHECK-NOT: musttail call fastcc void @fakeresume1(i64* align 8 null)
+; CHECK:      musttail call fastcc void @fakeresume1(i64* align 8 null)
+; CHECK-NEXT: ret void
 
 declare token @llvm.coro.id(i32, i8* readnone, i8* nocapture readonly, i8*) #1
 declare i1 @llvm.coro.alloc(token) #2
 declare i64 @llvm.coro.size.i64() #3
 declare i8* @llvm.coro.begin(token, i8* writeonly) #2
 declare token @llvm.coro.save(i8*) #2
 declare i8* @llvm.coro.frame() #3
 declare i8 @llvm.coro.suspend(token, i1) #2

llvm/test/Transforms/Coroutines/coro-split-musttail7.ll

exit:
   %.unused = getelementptr inbounds i8, i8* %vFrame, i32 0
   call i1 @llvm.coro.end(i8* null, i1 false)
   ret void
 }
 
 ; FIXME: The fakeresume1 here should be marked as musttail.
 ; Verify that in the resume part resume call is marked with musttail.
 ; CHECK-LABEL: @g.resume(
-; CHECK-NOT: musttail call fastcc void @fakeresume1(i64* align 8 null)
+; CHECK:         musttail call fastcc void @fakeresume1(i64* align 8 null)
+; CHECK-NEXT:    ret void
 
 ; It has a cleanup bb.
 define void @f() #0 {
 entry:
   %id = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* null)
   %alloc = call i8* @malloc(i64 16) #3
   %alloc.var = alloca i64
   %alloca.var.i8 = bitcast i64* %alloc.var to i8*
exit:
   %.unused = getelementptr inbounds i8, i8* %vFrame, i32 0
   call i1 @llvm.coro.end(i8* null, i1 false)
   ret void
 }
 
 ; FIXME: The fakeresume1 here should be marked as musttail.
 ; Verify that in the resume part resume call is marked with musttail.
 ; CHECK-LABEL: @f.resume(
-; CHECK-NOT: musttail call fastcc void @fakeresume1(i64* align 8 null)
+; CHECK:         musttail call fastcc void @fakeresume1(i64* align 8 null)
+; CHECK-NEXT:    ret void
 
 declare token @llvm.coro.id(i32, i8* readnone, i8* nocapture readonly, i8*) #1
 declare i1 @llvm.coro.alloc(token) #2
 declare i64 @llvm.coro.size.i64() #3
 declare i8* @llvm.coro.begin(token, i8* writeonly) #2
 declare token @llvm.coro.save(i8*) #2
 declare i8* @llvm.coro.frame() #3
 declare i8 @llvm.coro.suspend(token, i1) #2