[clang][CodeGen] Promote point of GRO(CWG2563) (#151067)

This patch implement piece of the proposed solution to
[CWG2563](https://cplusplus.github.io/CWG/issues/2563.html):

> [9.6.4 dcl.fct.def.coroutine.p8] This return exits the scope of gro.
It exits the scope of promise only if the coroutine completed without
suspending.

If a coroutine completes without suspending, it does not exit the scope
of the promise until GRO conversion is done, because GRO conversion is
considered part of the coroutine execution. The current behavior
performs conversion after coroutine state cleanup, which does not
conform to the standard:

``` LLVM
before.cleanup:
  ; ...
  br label %coro.cleanup

coro.cleanup:
  ; cleanup logics
  br %coro.end

coro.end:
  call void @llvm.coro.end(ptr null, i1 false, token none)
  ; GRO conversion
  ; ret GRO or void
```
This patch proposes the following codegen:
``` LLVM
any.suspend:
  %suspend = call i8 @llvm.coro.suspend(token %8, i1 true)
  switch i8 %suspend, label %pre.gro.conv [
    i8 0, label %ready
    i8 1, label %coro.cleanup
  ]

ready:
  ; ...

pre.gro.conv:
  %body.done = phi i1 [ false, %any.suspend ], [ true, %any.ready ]
  %InRamp = call i1 @llvm.coro.is_in_ramp()
  br i1 %InRamp, label %gro.conv, label %after.gro.conv

gro.conv:
  ; GRO conversion
  br label %after.gro.conv

after.gro.conv:
  br i1 %body.done, label %coro.cleanup, label %coro.ret

coro.cleanup:
  ; cleanup logics
  br %coro.ret

coro.ret:
  call void @llvm.coro.end(ptr null, i1 false, token none)
  ; ret GRO or void
```

Close #120200

Author: NewSigma

clang/lib/CodeGen/CGCoroutine.cpp

@@ -43,6 +43,8 @@ struct clang::CodeGen::CGCoroData {
 
   // A branch to this block is emitted when coroutine needs to suspend.
   llvm::BasicBlock *SuspendBB = nullptr;
+  // A branch to this block after final.cleanup or final.ready
+  llvm::BasicBlock *FinalExit = nullptr;
 
   // The promise type's 'unhandled_exception' handler, if it defines one.
   Stmt *ExceptionHandler = nullptr;
@@ -332,6 +334,8 @@ static LValueOrRValue emitSuspendExpression(CodeGenFunction &CGF, CGCoroData &Co
   // Emit cleanup for this suspend point.
   CGF.EmitBlock(CleanupBlock);
   CGF.EmitBranchThroughCleanup(Coro.CleanupJD);
+  if (IsFinalSuspend)
+    Coro.FinalExit = CleanupBlock->getSingleSuccessor();
 
   // Emit await_resume expression.
   CGF.EmitBlock(ReadyBlock);
@@ -687,11 +691,11 @@ struct GetReturnObjectManager {
   }
 
   // The gro variable has to outlive coroutine frame and coroutine promise, but,
-  // it can only be initialized after coroutine promise was created, thus, we
-  // split its emission in two parts. EmitGroAlloca emits an alloca and sets up
-  // cleanups. Later when coroutine promise is available we initialize the gro
-  // and sets the flag that the cleanup is now active.
-  void EmitGroAlloca() {
+  // it can only be initialized after coroutine promise was created. Thus,
+  // EmitGroActive emits a flag and sets it to false. Later when coroutine
+  // promise is available we initialize the gro and set the flag indicating that
+  // the cleanup is now active.
+  void EmitGroActive() {
     if (DirectEmit)
       return;
 
@@ -701,12 +705,23 @@ struct GetReturnObjectManager {
       return;
     }
 
-    auto *GroVarDecl = cast<VarDecl>(GroDeclStmt->getSingleDecl());
-
     // Set GRO flag that it is not initialized yet
     GroActiveFlag = CGF.CreateTempAlloca(Builder.getInt1Ty(), CharUnits::One(),
                                          "gro.active");
     Builder.CreateStore(Builder.getFalse(), GroActiveFlag);
+  }
+
+  void EmitGroAlloca() {
+    if (DirectEmit)
+      return;
+
+    auto *GroDeclStmt = dyn_cast_or_null<DeclStmt>(S.getResultDecl());
+    if (!GroDeclStmt) {
+      // If get_return_object returns void, no need to do an alloca.
+      return;
+    }
+
+    auto *GroVarDecl = cast<VarDecl>(GroDeclStmt->getSingleDecl());
 
     GroEmission = CGF.EmitAutoVarAlloca(*GroVarDecl);
 
@@ -768,6 +783,78 @@ struct GetReturnObjectManager {
     CGF.EmitAutoVarInit(GroEmission);
     Builder.CreateStore(Builder.getTrue(), GroActiveFlag);
   }
+  // The GRO returns either when it is first suspended or when it completes
+  // without ever being suspended. The EmitGroConv function evaluates these
+  // conditions and perform the conversion if needed.
+  //
+  // Before EmitGroConv():
+  //   final.exit:
+  //     switch i32 %cleanup.dest, label %destroy [
+  //        i32 0, label %after.ready
+  //     ]
+  //
+  //   after.ready:
+  //     ; (empty)
+  //
+  // After EmitGroConv():
+  //   final.exit:
+  //     switch i32 %cleanup.dest, label %destroy [
+  //        i32 0, label %pre.gro.conv
+  //     ]
+  //
+  //   pre.gro.conv:
+  //     %IsFinalExit = phi i1 [ false, %any.suspend ], [ true, %final.exit ]
+  //     %InRamp = call i1 @llvm.coro.is_in_ramp()
+  //     br i1 %InRamp, label %gro.conv, label %after.gro.conv
+  //
+  //   gro.conv:
+  //     ; GRO conversion
+  //     br label %after.gro.conv
+  //
+  //   after.gro.conv:
+  //     br i1 %IsFinalExit, label %after.ready, label %coro.ret
+  void EmitGroConv(BasicBlock *RetBB) {
+    auto *AfterReadyBB = Builder.GetInsertBlock();
+    Builder.ClearInsertionPoint();
+
+    auto *PreConvBB = CGF.CurCoro.Data->SuspendBB;
+    CGF.EmitBlock(PreConvBB);
+    // If final.exit exists, redirect it to PreConvBB
+    llvm::PHINode *IsFinalExit = nullptr;
+    if (BasicBlock *FinalExit = CGF.CurCoro.Data->FinalExit) {
+      assert(AfterReadyBB &&
+             AfterReadyBB->getSinglePredecessor() == FinalExit &&
+             "Expect fallthrough from final.exit block");
+      AfterReadyBB->replaceAllUsesWith(PreConvBB);
+      PreConvBB->moveBefore(AfterReadyBB);
+
+      // If true, coroutine completes and should be destroyed after conversion
+      IsFinalExit =
+          Builder.CreatePHI(Builder.getInt1Ty(), llvm::pred_size(PreConvBB));
+      for (auto *Pred : llvm::predecessors(PreConvBB)) {
+        auto *V = (Pred == FinalExit) ? Builder.getTrue() : Builder.getFalse();
+        IsFinalExit->addIncoming(V, Pred);
+      }
+    }
+    auto *InRampFn = CGF.CGM.getIntrinsic(llvm::Intrinsic::coro_is_in_ramp);
+    auto *InRamp = Builder.CreateCall(InRampFn, {}, "InRamp");
+    auto *ConvBB = CGF.createBasicBlock("gro.conv");
+    auto *AfterConvBB = CGF.createBasicBlock("after.gro.conv");
+    Builder.CreateCondBr(InRamp, ConvBB, AfterConvBB);
+
+    CGF.EmitBlock(ConvBB);
+    CGF.EmitAnyExprToMem(S.getReturnValue(), CGF.ReturnValue,
+                         S.getReturnValue()->getType().getQualifiers(),
+                         /*IsInit*/ true);
+    Builder.CreateBr(AfterConvBB);
+
+    CGF.EmitBlock(AfterConvBB);
+    if (IsFinalExit)
+      Builder.CreateCondBr(IsFinalExit, AfterReadyBB, RetBB);
+    else
+      Builder.CreateBr(RetBB);
+    Builder.SetInsertPoint(AfterReadyBB);
+  }
 };
 } // namespace
 
@@ -795,7 +882,10 @@ void CodeGenFunction::EmitCoroutineBody(const CoroutineBodyStmt &S) {
       CGM.getIntrinsic(llvm::Intrinsic::coro_id),
       {Builder.getInt32(NewAlign), NullPtr, NullPtr, NullPtr});
   createCoroData(*this, CurCoro, CoroId);
-  CurCoro.Data->SuspendBB = RetBB;
+
+  GetReturnObjectManager GroManager(*this, S);
+  CurCoro.Data->SuspendBB =
+      GroManager.DirectEmit ? RetBB : createBasicBlock("pre.gvo.conv");
   assert(ShouldEmitLifetimeMarkers &&
          "Must emit lifetime intrinsics for coroutines");
 
@@ -839,9 +929,6 @@ void CodeGenFunction::EmitCoroutineBody(const CoroutineBodyStmt &S) {
       CGM.getIntrinsic(llvm::Intrinsic::coro_begin), {CoroId, Phi});
   CurCoro.Data->CoroBegin = CoroBegin;
 
-  GetReturnObjectManager GroManager(*this, S);
-  GroManager.EmitGroAlloca();
-
   CurCoro.Data->CleanupJD = getJumpDestInCurrentScope(RetBB);
   {
     CGDebugInfo *DI = getDebugInfo();
@@ -884,6 +971,7 @@ void CodeGenFunction::EmitCoroutineBody(const CoroutineBodyStmt &S) {
       // not needed.
     }
 
+    GroManager.EmitGroActive();
     EmitStmt(S.getPromiseDeclStmt());
 
     Address PromiseAddr = GetAddrOfLocalVar(S.getPromiseDecl());
@@ -895,6 +983,7 @@ void CodeGenFunction::EmitCoroutineBody(const CoroutineBodyStmt &S) {
     CoroId->setArgOperand(1, PromiseAddrVoidPtr);
 
     // Now we have the promise, initialize the GRO
+    GroManager.EmitGroAlloca();
     GroManager.EmitGroInit();
 
     EHStack.pushCleanup<CallCoroEnd>(EHCleanup);
@@ -950,31 +1039,31 @@ void CodeGenFunction::EmitCoroutineBody(const CoroutineBodyStmt &S) {
       // We don't need FinalBB. Emit it to make sure the block is deleted.
       EmitBlock(FinalBB, /*IsFinished=*/true);
     }
+
+    // We need conversion if get_return_object's type doesn't matches the
+    // coroutine return type.
+    if (!GroManager.DirectEmit)
+      GroManager.EmitGroConv(RetBB);
   }
 
   EmitBlock(RetBB);
-  // Emit coro.end before getReturnStmt (and parameter destructors), since
-  // resume and destroy parts of the coroutine should not include them.
+  // Emit coro.end before ret instruction, since resume and destroy parts of the
+  // coroutine should return void.
   llvm::Function *CoroEnd = CGM.getIntrinsic(llvm::Intrinsic::coro_end);
   Builder.CreateCall(CoroEnd,
                      {NullPtr, Builder.getFalse(),
                       llvm::ConstantTokenNone::get(CoroEnd->getContext())});
 
-  if (Stmt *Ret = S.getReturnStmt()) {
+  if (auto *Ret = cast_or_null<ReturnStmt>(S.getReturnStmt())) {
     // Since we already emitted the return value above, so we shouldn't
     // emit it again here.
-    Expr *PreviousRetValue = nullptr;
-    if (GroManager.DirectEmit) {
-      PreviousRetValue = cast<ReturnStmt>(Ret)->getRetValue();
-      cast<ReturnStmt>(Ret)->setRetValue(nullptr);
-    }
+    Expr *PreviousRetValue = Ret->getRetValue();
+    Ret->setRetValue(nullptr);
     EmitStmt(Ret);
     // Set the return value back. The code generator, as the AST **Consumer**,
     // shouldn't change the AST.
-    if (PreviousRetValue)
-      cast<ReturnStmt>(Ret)->setRetValue(PreviousRetValue);
+    Ret->setRetValue(PreviousRetValue);
   }
-
   // LLVM require the frontend to mark the coroutine.
   CurFn->setPresplitCoroutine();
 

clang/test/CodeGenCoroutines/coro-gro.cpp

@@ -29,47 +29,72 @@ void doSomething() noexcept;
 // CHECK: define{{.*}} i32 @_Z1fv(
 int f() {
   // CHECK: %[[RetVal:.+]] = alloca i32
-  // CHECK: %[[GroActive:.+]] = alloca i1
-  // CHECK: %[[CoroGro:.+]] = alloca %struct.GroType, {{.*}} !coro.outside.frame ![[OutFrameMetadata:.+]]
+  // CHECK-NEXT: %[[GroActive:.+]] = alloca i1
+  // CHECK-NEXT: %[[Promise:.+]] = alloca %"struct.std::coroutine_traits<int>::promise_type", align 1
+  // CHECK-NEXT: %[[CoroGro:.+]] = alloca %struct.GroType, {{.*}} !coro.outside.frame ![[OutFrameMetadata:.+]]
 
   // CHECK: %[[Size:.+]] = call i64 @llvm.coro.size.i64()
-  // CHECK: call noalias noundef nonnull ptr @_Znwm(i64 noundef %[[Size]])
+  // CHECK-NEXT: call noalias noundef nonnull ptr @_Znwm(i64 noundef %[[Size]])
+
   // CHECK: store i1 false, ptr %[[GroActive]]
-  // CHECK: call void @_ZNSt16coroutine_traitsIiJEE12promise_typeC1Ev(
-  // CHECK: call void @_ZNSt16coroutine_traitsIiJEE12promise_type17get_return_objectEv({{.*}} %[[CoroGro]]
-  // CHECK: store i1 true, ptr %[[GroActive]]
+  // CHECK-NEXT: call void @llvm.lifetime.start.p0(ptr %[[Promise]])
+  // CHECK-NEXT: call void @_ZNSt16coroutine_traitsIiJEE12promise_typeC1Ev(
+  // CHECK-NEXT: call void @llvm.lifetime.start.p0(ptr %[[CoroGro]])
+  // CHECK-NEXT: call void @_ZNSt16coroutine_traitsIiJEE12promise_type17get_return_objectEv({{.*}} %[[CoroGro]]
+  // CHECK-NEXT: store i1 true, ptr %[[GroActive]]
 
   Cleanup cleanup;
   doSomething();
   co_return;
 
   // CHECK: call void @_Z11doSomethingv(
-  // CHECK: call void @_ZNSt16coroutine_traitsIiJEE12promise_type11return_voidEv(
-  // CHECK: call void @_ZN7CleanupD1Ev(
-
-  // Destroy promise and free the memory.
-
-  // CHECK: call void @_ZNSt16coroutine_traitsIiJEE12promise_typeD1Ev(
-  // CHECK: %[[Mem:.+]] = call ptr @llvm.coro.free(
-  // CHECK: %[[SIZE:.+]] = call i64 @llvm.coro.size.i64()
-  // CHECK: call void @_ZdlPvm(ptr noundef %[[Mem]], i64 noundef %[[SIZE]])
+  // CHECK-NEXT: call void @_ZNSt16coroutine_traitsIiJEE12promise_type11return_voidEv(
+  // CHECK-NEXT: call void @_ZN7CleanupD1Ev(
 
   // Initialize retval from Gro and destroy Gro
   // Note this also tests delaying initialization when Gro and function return
   // types mismatch (see cwg2563).
 
-  // CHECK: %[[Conv:.+]] = call noundef i32 @_ZN7GroTypecviEv(
-  // CHECK: store i32 %[[Conv]], ptr %[[RetVal]]
-  // CHECK: %[[IsActive:.+]] = load i1, ptr %[[GroActive]]
-  // CHECK: br i1 %[[IsActive]], label %[[CleanupGro:.+]], label %[[Done:.+]]
+  // CHECK: pre.gvo.conv:
+  // CHECK-NEXT: %10 = phi i1 [ true, %cleanup8 ], [ false, %final.suspend ], [ false, %init.suspend ]
+  // CHECK-NEXT: %InRamp = call i1 @llvm.coro.is_in_ramp()
+  // CHECK-NEXT: br i1 %InRamp, label %[[GroConv:.+]], label %[[AfterGroConv:.+]]
+
+  // CHECK: [[GroConv]]:
+  // CHECK-NEXT: %[[Conv:.+]] = call noundef i32 @_ZN7GroTypecviEv(
+  // CHECK-NEXT: store i32 %[[Conv]], ptr %[[RetVal]]
+  // CHECK-NEXT: br label %[[AfterGroConv]]
+
+  // CHECK: [[AfterGroConv]]:
+  // CHECK-NEXT: br i1 %10, label %cleanup.cont10, label %[[CoroRet:.+]]
+
+  // CHECK: cleanup.cont10:
+  // CHECK-NEXT: br label %[[Cleanup:.+]]
+
+  // CHECK: [[Cleanup]]:
+  // CHECK-NEXT: %{{.*}} = phi i32
+  // CHECK-NEXT: %[[IsActive:.+]] = load i1, ptr %[[GroActive]]
+  // CHECK-NEXT: br i1 %[[IsActive]], label %[[CleanupGro:.+]], label %[[Done:.+]]
 
   // CHECK: [[CleanupGro]]:
-  // CHECK:   call void @_ZN7GroTypeD1Ev(
-  // CHECK:   br label %[[Done]]
+  // CHECK-NEXT: call void @_ZN7GroTypeD1Ev(
+  // CHECK-NEXT: br label %[[Done]]
+
+  // Destroy promise and free the memory.
 
   // CHECK: [[Done]]:
-  // CHECK:   %[[LoadRet:.+]] = load i32, ptr %[[RetVal]]
-  // CHECK:   ret i32 %[[LoadRet]]
+  // CHECK-NEXT: call void @llvm.lifetime.end.p0(ptr %[[CoroGro]])
+  // CHECK-NEXT: call void @_ZNSt16coroutine_traitsIiJEE12promise_typeD1Ev(
+  // CHECK-NEXT: call void @llvm.lifetime.end.p0(ptr %[[Promise]])
+  // CHECK-NEXT: %[[Mem:.+]] = call ptr @llvm.coro.free(
+
+  // CHECK: %[[SIZE:.+]] = call i64 @llvm.coro.size.i64()
+  // CHECK-NEXT: call void @_ZdlPvm(ptr noundef %[[Mem]], i64 noundef %[[SIZE]])
+
+  // CHECK: [[CoroRet]]:
+  // CHECK-NEXT: call void @llvm.coro.end(
+  // CHECK-NEXT: %[[LoadRet:.+]] = load i32, ptr %[[RetVal]]
+  // CHECK-NEXT: ret i32 %[[LoadRet]]
 }
 
 class invoker {