[coroutines] Add handling for unwind coro.ends

[GorNishanov]

Summary

Hi David:

This is an implementation of a hybrid idea, where for WinEH I use bundles and replace the coro.end with cleanupret and for landing pad model I use explicit control flow that jumps to EHresume block that clang set up for me.

clang part: https://reviews.llvm.org/D25444

[majnemer]

Can we have IR tests for this?

lib/Transforms/Coroutines/CoroSplit.cpp

165

Does FromPad have any CleanupReturnInsts at this point?

[GorNishanov]

In D25445#566982, @majnemer wrote:

Can we have IR tests for this?

Absolutely. I am working on them and will send an official CR with llvm-commits added on creation. I just wanted to make sure that the approach is reasonable before investing more time in this direction.

[majnemer]

I think its fine but int_coro_end should probably be marked as may_throw so we don't sink after it. We could append nounwind after we do the slicing.

[GorNishanov]

lib/Transforms/Coroutines/CoroSplit.cpp

165

By construction, it should not have any between cleanuppad and coro.end. It might have afterwards, but, those will get cleaned up by simplify cfg.

I will update the documentation for coroutine intrinsics and will indicate that expectation is that for WinEH exception hanlding, frontend should not emit any cleanuprets between cleanuppad and coro.end.

[majnemer]

lib/Transforms/Coroutines/CoroSplit.cpp

165

I wasn't thinking of cleanuprets between the pad and the coro.end but rets on paths which never reach coro.end.

[GorNishanov]

Opened official CR here https://reviews.llvm.org/D25543

include/llvm/IR/Intrinsics.td

def int_coro_id : Intrinsic<[llvm_token_ty], [llvm_i32_ty, llvm_ptr_ty,
                              ReadNone<1>, ReadOnly<2>, NoCapture<2>]>;
 def int_coro_alloc : Intrinsic<[llvm_i1_ty], [llvm_token_ty], []>;
 def int_coro_begin : Intrinsic<[llvm_ptr_ty], [llvm_token_ty, llvm_ptr_ty],
                                [WriteOnly<1>]>;
 
 def int_coro_free : Intrinsic<[llvm_ptr_ty], [llvm_token_ty, llvm_ptr_ty], 
                               [IntrReadMem, IntrArgMemOnly, ReadOnly<1>, 
                                NoCapture<1>]>;
-def int_coro_end : Intrinsic<[], [llvm_ptr_ty, llvm_i1_ty], []>;
+def int_coro_end : Intrinsic<[llvm_i1_ty], [llvm_ptr_ty, llvm_i1_ty], []>;
 
 def int_coro_frame : Intrinsic<[llvm_ptr_ty], [], [IntrNoMem]>;
 def int_coro_size : Intrinsic<[llvm_anyint_ty], [], [IntrNoMem]>;
 
 def int_coro_save : Intrinsic<[llvm_token_ty], [llvm_ptr_ty], []>;
 def int_coro_suspend : Intrinsic<[llvm_i8_ty], [llvm_token_ty, llvm_i1_ty], []>;
 
 def int_coro_param : Intrinsic<[llvm_i1_ty], [llvm_ptr_ty, llvm_ptr_ty],

lib/Transforms/Coroutines/CoroSplit.cpp

static void replaceFallthroughCoroEnd(IntrinsicInst *End,
   ReturnInst::Create(NewE->getContext(), nullptr, NewE);
 
   // Remove the rest of the block, by splitting it into an unreachable block.
   auto *BB = NewE->getParent();
   BB->splitBasicBlock(NewE);
   BB->getTerminator()->eraseFromParent();
 }
 
+// In Resumers, we replace unwind coro.end with True to force the immediate
+// unwind to caller.
+static void replaceUnwindCoroEnds(coro::Shape &Shape, ValueToValueMapTy &VMap) {
+  if (Shape.CoroEnds.empty())
+    return;
+
+  LLVMContext &Context = Shape.CoroEnds.front()->getContext();
+  auto *True = ConstantInt::getTrue(Context);
+  for (CoroEndInst *CE : Shape.CoroEnds) {
+    if (!CE->isUnwind())
+      continue;
+
+    auto *NewCE = cast<IntrinsicInst>(VMap[CE]);
+
+    // If coro.end has an associated bundle, add cleanupret instruction.
+    if (auto Bundle = NewCE->getOperandBundle(LLVMContext::OB_funclet)) {
+      Value *FromPad = Bundle->Inputs[0];
+      auto *CleanupRet = CleanupReturnInst::Create(FromPad, nullptr, NewCE);
+      NewCE->getParent()->splitBasicBlock(NewCE);
+      CleanupRet->getParent()->getTerminator()->eraseFromParent();
+    }
+
+    NewCE->replaceAllUsesWith(True);
+    NewCE->eraseFromParent();
+  }
+}
+
 // Rewrite final suspend point handling. We do not use suspend index to
 // represent the final suspend point. Instead we zero-out ResumeFnAddr in the
 // coroutine frame, since it is undefined behavior to resume a coroutine
 // suspended at the final suspend point. Thus, in the resume function, we can
 // simply remove the last case (when coro::Shape is built, the final suspend
 // point (if present) is always the last element of CoroSuspends array).
 // In the destroy function, we add a code sequence to check if ResumeFnAddress
 // is Null, and if so, jump to the appropriate label to handle cleanup from the
static Function *createClone(Function &F, Twine Suffix, coro::Shape &Shape,
   for (CoroSuspendInst *CS : Shape.CoroSuspends) {
     auto *MappedCS = cast<CoroSuspendInst>(VMap[CS]);
     MappedCS->replaceAllUsesWith(NewValue);
     MappedCS->eraseFromParent();
   }
 
   // Remove coro.end intrinsics.
   replaceFallthroughCoroEnd(Shape.CoroEnds.front(), VMap);
-  // FIXME: coming in upcoming patches:
-  // replaceUnwindCoroEnds(Shape.CoroEnds, VMap);
+  replaceUnwindCoroEnds(Shape, VMap);
 
   // Eliminate coro.free from the clones, replacing it with 'null' in cleanup,
   // to suppress deallocation code.
   coro::replaceCoroFree(cast<CoroIdInst>(VMap[Shape.CoroBegin->getId()]),
                         /*Elide=*/FnIndex == 2);
 
   NewF->setCallingConv(CallingConv::Fast);
 
   return NewF;
 }
 
 static void removeCoroEnds(coro::Shape &Shape) {
-  for (CoroEndInst *CE : Shape.CoroEnds)
+  if (Shape.CoroEnds.empty())
+    return;
+
+  LLVMContext &Context = Shape.CoroEnds.front()->getContext();
+  auto *False = ConstantInt::getFalse(Context);
+
+  for (CoroEndInst *CE : Shape.CoroEnds) {
+    CE->replaceAllUsesWith(False);
     CE->eraseFromParent();
   }
+}
 
 static void replaceFrameSize(coro::Shape &Shape) {
   if (Shape.CoroSizes.empty())
     return;
 
   // In the same function all coro.sizes should have the same result type.
   auto *SizeIntrin = Shape.CoroSizes.back();
   Module *M = SizeIntrin->getModule();
static void splitCoroutine(Function &F, CallGraph &CG, CallGraphSCC &SCC) {
   if (Shape.CoroSuspends.empty()) {
     handleNoSuspendCoroutine(Shape.CoroBegin, Shape.FrameTy);
     removeCoroEnds(Shape);
     postSplitCleanup(F);
     coro::updateCallGraph(F, {}, CG, SCC);
     return;
   }
 
-  // If there are no suspend points, no split required, just remove
-  // the allocation and deallocation blocks, they are not needed.
-  if (Shape.CoroSuspends.empty()) {
-    handleNoSuspendCoroutine(Shape.CoroBegin, Shape.FrameTy);
-    removeCoroEnds(Shape);
-    postSplitCleanup(F);
-    coro::updateCallGraph(F, {}, CG, SCC);
-    return;
-  }
-
   auto *ResumeEntry = createResumeEntryBlock(F, Shape);
   auto ResumeClone = createClone(F, ".resume", Shape, ResumeEntry, 0);
   auto DestroyClone = createClone(F, ".destroy", Shape, ResumeEntry, 1);
   auto CleanupClone = createClone(F, ".cleanup", Shape, ResumeEntry, 2);
 
   // We no longer need coro.end in F.
   removeCoroEnds(Shape);