[C++20] [Coroutines] Mark coroutine done if unhandled_exception throws

[ChuanqiXu]

Summary

[[dcl.fct.def.coroutine]/p14](https://eel.is/c++draft/dcl.fct.def.coroutine#14) says:

If the evaluation of the expression promise.unhandled_­exception() exits via an exception, the coroutine is considered suspended at the final suspend point.

However, this is not implemented in clang. We could observe this from: https://godbolt.org/z/Edr59d5Y6.

This patch would implement this feature by marking the coroutine as done at the place of coro.end(frame, /*InUnwindPath=*/true ).

After this patch, the behavior of this example would be the same with GCC: https://godbolt.org/z/rh86xKf85.

Test Plan: check-all, https://godbolt.org/z/rh86xKf85 and an internal coroutine library

[rjmccall]

I agree that you shouldn't call suspend, but doesn't coro.end have the behavior of marking the coroutine done? Should we just be calling coro.end on this path?

llvm/docs/Coroutines.rst

1667

"The '`llvm.coro.mark.done' intrinsic marks that the current coroutine is complete. 'llvm.coro.mark.done`' is lowered during coroutine splitting and must not appear outside of an unsplit coroutine. This intrinsic is only supported for switched-resume coroutines."

1674

"'`llvm.coro.mark.done`' doesn't require any arguments. There is no ambiguity about which coroutine is meant because unsplit coroutines in LLVM cannot be broken up or inlined."

[ChuanqiXu]

Use llvm.coro.end(frame, /*InUnwindPath=*/true) instead of new coroutine intrinsics.

[ChuanqiXu]

In D115219#3175582, @rjmccall wrote:

I agree that you shouldn't call suspend, but doesn't coro.end have the behavior of marking the coroutine done? Should we just be calling coro.end on this path?

@rjmccall great insight! coro.end wouldn't marking the coroutine done previously. But its place is perfect to do so. I have added the codes to mark the coroutine as done for coro.end in the unwind path. And I have checked the behavior. The exception happened earlier wouldn't run into the path of coro.end exists. So the behavior is satisfying. The only defect might be that the behavior is C++'s semantic. Although, there might be a problem if there is another language also uses switch-resumed ABI one day. I think the current approach is much better at least it has and would generate less code.

[rjmccall]

In D115219#3175660, @ChuanqiXu wrote:

In D115219#3175582, @rjmccall wrote:

I agree that you shouldn't call suspend, but doesn't coro.end have the behavior of marking the coroutine done? Should we just be calling coro.end on this path?

@rjmccall great insight! coro.end wouldn't marking the coroutine done previously. But its place is perfect to do so. I have added the codes to mark the coroutine as done for coro.end in the unwind path. And I have checked the behavior. The exception happened earlier wouldn't run into the path of coro.end exists. So the behavior is satisfying. The only defect might be that the behavior is C++'s semantic. Although, there might be a problem if there is another language also uses switch-resumed ABI one day. I think the current approach is much better at least it has and would generate less code.

Okay. Well, I'm glad it works. I guess I find it a little strange that coro.end doesn't already mark the coroutine done. I guess the normal C++ lowering always generates a final suspension before reaching a non-unwind coro.end? There are a lot of semantic differences between throwing out of a C++ coroutine and returning out of it that I always find surprising.

Something ultimately invalidates the coroutine handle by deallocating the frame, right? When does that happen when throwing out of a coroutine?

[ChuanqiXu]

In D115219#3177213, @rjmccall wrote:

Okay. Well, I'm glad it works. I guess I find it a little strange that coro.end doesn't already mark the coroutine done. I guess the normal C++ lowering always generates a final suspension before reaching a non-unwind coro.end?

Yes, there is always a final suspend before coro.end in normal path. So that the coroutine would be marked done all the time except unhandled_exception throws.

There are a lot of semantic differences between throwing out of a C++ coroutine and returning out of it that I always find surprising.

Yes, I guess this is the reason why the coroutine should be marked done if unhandled_exception throws.

Something ultimately invalidates the coroutine handle by deallocating the frame, right?

Yes, the coroutine frame is invalided by deallocating. Users (coroutine library writers) need to deallocate it by hand usually (by call coroutine_handle<>::destroy()). The coroutine would get destroyed automatically if the coroutine completes and not suspended at the final suspend point.

When does that happen when throwing out of a coroutine?

After throwing out of a coroutine, the coroutine frame is still valid. And the owner is responsible to destroy it.

[rjmccall]

In D115219#3178370, @ChuanqiXu wrote:

In D115219#3177213, @rjmccall wrote:

Okay. Well, I'm glad it works. I guess I find it a little strange that coro.end doesn't already mark the coroutine done. I guess the normal C++ lowering always generates a final suspension before reaching a non-unwind coro.end?

Yes, there is always a final suspend before coro.end in normal path. So that the coroutine would be marked done all the time except unhandled_exception throws.

Ah, I see. This is the one place where an exception can be thrown before we reach that point.

You should update the coroutines documentation to mention this effect of llvm.coro.end(true).

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

313

314

325

We need to do this store elsewhere, right, like during final suspends? Can we make a common function for it?

[ChuanqiXu]

Address comments

[ChuanqiXu]

llvm/docs/Coroutines.rst

1396

Maybe I need to explain why it would work in start (ramp) function. Generally, the ramp function is very small. But in case the initial_suspend of the coroutine wouldn't always suspend, the start function could be as large as the original function. In this case, it should remain the original semantics.

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

325

I am not sure if I understand your point. I think the place where the coro.end lives should be the right place. The common function is made.

[rjmccall]

Minor suggestion in the docs, but otherwise LGTM.

llvm/docs/Coroutines.rst

1382

How about:

In the unwind path (when the argument is `true`), `coro.end` will mark the coroutine as done, making it undefined behavior to resume the coroutine again and causing `llvm.coro.done` to return `true`.  This is not necessary in the normal path because the coroutine will already be marked as done by the final suspend.

1396

That makes sense. Honestly, it feels very strange to me that there are any differences between the ramp function and the resume/destroy functions in the handling of coro.end.

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

325

Yes, thank you, this is all I wanted.

[ChuanqiXu]

Address comments.

[ChuanqiXu]

Thanks for reviewing this!

llvm/docs/Coroutines.rst

ehcleanup:
       %tok = cleanuppad within none []
       %unused = call i1 @llvm.coro.end(i8* null, i1 true) [ "funclet"(token %tok) ]
       cleanupret from %tok unwind label %RestOfTheCleanup
 
 The `CoroSplit` pass, if the funclet bundle is present, will insert
 ``cleanupret from %tok unwind to caller`` before
 the `coro.end`_ intrinsic and will remove the rest of the block.
 
+In the unwind path (when the argument is `true`), `coro.end` will mark the coroutine
+as done, making it undefined behavior to resume the coroutine again and causing 
+`llvm.coro.done` to return `true`.  This is not necessary in the normal path because
+the coroutine will already be marked as done by the final suspend.
+
 The following table summarizes the handling of `coro.end`_ intrinsic.
 
-+--------------------------+-------------------+-------------------------------+
-|                          | In Start Function | In Resume/Destroy Functions   |
-+--------------------------+-------------------+-------------------------------+
-|unwind=false              | nothing           |``ret void``                   |
-+------------+-------------+-------------------+-------------------------------+
-|            | WinEH       | nothing           |``cleanupret unwind to caller``|
-|unwind=true +-------------+-------------------+-------------------------------+
-|            | Landingpad  | nothing           | nothing                       |
-+------------+-------------+-------------------+-------------------------------+
++--------------------------+------------------------+-------------------------------+
+|                          | In Start Function      | In Resume/Destroy Functions   |
++--------------------------+------------------------+-------------------------------+
+|unwind=false              | nothing                |``ret void``                   |
++------------+-------------+------------------------+-------------------------------+
+|            | WinEH       | mark coroutine as done |``cleanupret unwind to caller``|
+|            |             |                        |mark coroutine done            |
+|unwind=true +-------------+------------------------+-------------------------------+
+|            | Landingpad  | mark coroutine as done | mark coroutine done           |
++------------+-------------+------------------------+-------------------------------+
 
 
 'llvm.coro.end.async' Intrinsic
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 ::
 
   declare i1 @llvm.coro.end.async(i8* <handle>, i1 <unwind>, ...)
 

 help later coroutine passes. This pass lowers `coro.frame`_, `coro.done`_,
 and `coro.promise`_ intrinsics.
 
 .. _CoroSplit:
 
 CoroSplit
 ---------
 The pass CoroSplit buides coroutine frame and outlines resume and destroy parts
 into separate functions.
 
 CoroElide
 ---------
 The pass CoroElide examines if the inlined coroutine is eligible for heap
 allocation elision optimization. If so, it replaces
 `coro.begin` intrinsic with an address of a coroutine frame placed on its caller
 and replaces `coro.alloc` and `coro.free` intrinsics with `false` and `null`
 respectively to remove the deallocation code.
 This pass also replaces `coro.resume` and `coro.destroy` intrinsics with direct
 calls to resume and destroy functions for a particular coroutine where possible.
 
 CoroCleanup
 -----------
 This pass runs late to lower all coroutine related intrinsics not replaced by
 earlier passes.

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

static void replaceFallthroughCoroEnd(AnyCoroEndInst *End,
   }
 
   // Remove the rest of the block, by splitting it into an unreachable block.
   auto *BB = End->getParent();
   BB->splitBasicBlock(End);
   BB->getTerminator()->eraseFromParent();
 }
 
+// Mark a coroutine as done, which implies that the coroutine is finished and
+// never get resumed.
+//
+// In resume-switched ABI, the done state is represented by storing zero in
+// ResumeFnAddr.
+//
+// NOTE: We couldn't omit the argument `FramePtr`. It is necessary because the
+// pointer to the frame in splitted function is not stored in `Shape`.
+static void markCoroutineAsDone(IRBuilder<> &Builder, const coro::Shape &Shape,
+                                Value *FramePtr) {
+  assert(
+      Shape.ABI == coro::ABI::Switch &&
+      "markCoroutineAsDone is only supported for Switch-Resumed ABI for now.");
+  auto *GepIndex = Builder.CreateStructGEP(
+      Shape.FrameTy, FramePtr, coro::Shape::SwitchFieldIndex::Resume,
+      "ResumeFn.addr");
+  auto *NullPtr = ConstantPointerNull::get(cast<PointerType>(
+      Shape.FrameTy->getTypeAtIndex(coro::Shape::SwitchFieldIndex::Resume)));
+  Builder.CreateStore(NullPtr, GepIndex);
+}
+
 /// Replace an unwind call to llvm.coro.end.
 static void replaceUnwindCoroEnd(AnyCoroEndInst *End, const coro::Shape &Shape,
                                  Value *FramePtr, bool InResume,
                                  CallGraph *CG) {
   IRBuilder<> Builder(End);
 
   switch (Shape.ABI) {
   // In switch-lowering, this does nothing in the main function.
-  case coro::ABI::Switch:
+  case coro::ABI::Switch: {
+    // In C++'s specification, the coroutine should be marked as done
+    // if promise.unhandled_exception() throws.  The frontend will
+    // call coro.end(true) along this path.
+    //
+    // FIXME: We should refactor this once there is other language
+    // which uses Switch-Resumed style other than C++.
+    markCoroutineAsDone(Builder, Shape, FramePtr);
     if (!InResume)
       return;
     break;
+  }
   // In async lowering this does nothing.
   case coro::ABI::Async:
     break;
   // In continuation-lowering, this frees the continuation storage.
   case coro::ABI::Retcon:
   case coro::ABI::RetconOnce:
     maybeFreeRetconStorage(Builder, Shape, FramePtr, CG);
     break;
   }
 
for (auto *AnyS : Shape.CoroSuspends) {
     ConstantInt *IndexVal = Shape.getIndex(SuspendIndex);
 
     // Replace CoroSave with a store to Index:
     //    %index.addr = getelementptr %f.frame... (index field number)
     //    store i32 0, i32* %index.addr1
     auto *Save = S->getCoroSave();
     Builder.SetInsertPoint(Save);
     if (S->isFinal()) {
-      // Final suspend point is represented by storing zero in ResumeFnAddr.
-      auto *GepIndex = Builder.CreateStructGEP(FrameTy, FramePtr,
-                                 coro::Shape::SwitchFieldIndex::Resume,
-                                  "ResumeFn.addr");
-      auto *NullPtr = ConstantPointerNull::get(cast<PointerType>(
-          FrameTy->getTypeAtIndex(coro::Shape::SwitchFieldIndex::Resume)));
-      Builder.CreateStore(NullPtr, GepIndex);
+      // The coroutine should be marked done if it reaches the final suspend
+      // point.
+      markCoroutineAsDone(Builder, Shape, FramePtr);
     } else {
       auto *GepIndex = Builder.CreateStructGEP(
           FrameTy, FramePtr, Shape.getSwitchIndexField(), "index.addr");
       Builder.CreateStore(IndexVal, GepIndex);
     }
     Save->replaceAllUsesWith(ConstantTokenNone::get(C));
     Save->eraseFromParent();
 

llvm/test/Transforms/Coroutines/coro-split-eh-00.ll

 
 ; CHECK:      CoroEnd:
 ; CHECK-NEXT:   ret void
 
 ; CHECK:         lpad:
 ; CHECK-NEXT:      %lpval = landingpad { i8*, i32 }
 ; CHECK-NEXT:         cleanup
 ; CHECK-NEXT:      call void @print(i32 2)
+; Checks that the coroutine would be marked as done if it exits in unwinding path.
+; CHECK-NEXT:      %[[RESUME_ADDR:.+]] = getelementptr inbounds %[[FRAME_TY:.+]], %[[FRAME_TY]]* %FramePtr, i32 0, i32 0
+; CHECK-NEXT:      store void (%[[FRAME_TY]]*)* null, void (%[[FRAME_TY]]*)** %[[RESUME_ADDR]], align 8
 ; CHECK-NEXT:      resume { i8*, i32 } %lpval
 
 declare i8* @llvm.coro.free(token, i8*)
 declare i32 @llvm.coro.size.i32()
 declare i8  @llvm.coro.suspend(token, i1)
 declare void @llvm.coro.resume(i8*)
 declare void @llvm.coro.destroy(i8*)
 
 declare token @llvm.coro.id(i32, i8*, i8*, i8*)
 declare i8* @llvm.coro.alloc(token)
 declare i8* @llvm.coro.begin(token, i8*)
 declare i1 @llvm.coro.end(i8*, i1)
 
 declare noalias i8* @malloc(i32)
 declare void @print(i32)
 declare void @free(i8*)

llvm/test/Transforms/Coroutines/coro-split-eh-01.ll

 ; CHECK:   to label %CoroEnd unwind label %lpad
 
 ; CHECK:      CoroEnd:
 ; CHECK-NEXT:   ret void
 
 ; CHECK:      lpad:
 ; CHECK-NEXT:   %tok = cleanuppad within none []
 ; CHECK-NEXT:   call void @print(i32 2)
+; Checks that the coroutine would be marked as done if it exits in unwinding path.
+; CHECK-NEXT:   %[[RESUME_ADDR:.+]] = getelementptr inbounds %[[FRAME_TY:.+]], %[[FRAME_TY]]* %FramePtr, i32 0, i32 0
+; CHECK-NEXT:   store void (%[[FRAME_TY]]*)* null, void (%[[FRAME_TY]]*)** %[[RESUME_ADDR]], align 8
 ; CHECK-NEXT:   cleanupret from %tok unwind to caller
 
 declare i8* @llvm.coro.free(token, i8*)
 declare i32 @llvm.coro.size.i32()
 declare i8  @llvm.coro.suspend(token, i1)
 declare void @llvm.coro.resume(i8*)
 declare void @llvm.coro.destroy(i8*)
 
 declare token @llvm.coro.id(i32, i8*, i8*, i8*)
 declare i8* @llvm.coro.alloc(token)
 declare i8* @llvm.coro.begin(token, i8*)
 declare i1 @llvm.coro.end(i8*, i1)
 
 declare noalias i8* @malloc(i32)
 declare void @print(i32)
 declare void @free(i8*)