[clang] [llvm] [coroutine] Implement llvm.coro.await.suspend intrinsic (PR #79712)

via cfe-commits cfe-commits at lists.llvm.org
Mon Mar 4 13:27:28 PST 2024


https://github.com/fpasserby updated https://github.com/llvm/llvm-project/pull/79712

>From e7f3b6d51e443167688ddc8114304184f24597a3 Mon Sep 17 00:00:00 2001
From: fpasserby <125797601+fpasserby at users.noreply.github.com>
Date: Sun, 28 Jan 2024 00:02:15 +0100
Subject: [PATCH 01/17] Implement llvm.coro.await.suspend intrinsic

---
 clang/include/clang/AST/ExprCXX.h             |  28 ++-
 clang/lib/CodeGen/CGCoroutine.cpp             | 144 +++++++++++++--
 clang/lib/CodeGen/CodeGenFunction.h           |   4 +
 clang/lib/Sema/SemaCoroutine.cpp              | 160 +++++++++--------
 clang/test/AST/coroutine-locals-cleanup.cpp   |  10 +-
 .../CodeGenCoroutines/coro-always-inline.cpp  |   2 +-
 clang/test/CodeGenCoroutines/coro-await.cpp   |  79 ++++++--
 .../coro-awaiter-noinline-suspend.cpp         | 168 ------------------
 clang/test/CodeGenCoroutines/coro-dwarf.cpp   |  12 ++
 .../coro-function-try-block.cpp               |   2 +-
 .../coro-simplify-suspend-point.cpp           |  66 +++++++
 .../coro-symmetric-transfer-01.cpp            |   8 +-
 .../coro-symmetric-transfer-02.cpp            |  12 +-
 clang/test/CodeGenCoroutines/pr59181.cpp      |   9 +-
 clang/test/SemaCXX/co_await-ast.cpp           |   7 +-
 llvm/include/llvm/IR/Intrinsics.td            |  12 ++
 llvm/lib/IR/Verifier.cpp                      |   3 +
 llvm/lib/Transforms/Coroutines/CoroInstr.h    |  30 ++++
 llvm/lib/Transforms/Coroutines/CoroSplit.cpp  | 132 +++++++++++++-
 llvm/lib/Transforms/Coroutines/Coroutines.cpp |   3 +
 20 files changed, 573 insertions(+), 318 deletions(-)
 delete mode 100644 clang/test/CodeGenCoroutines/coro-awaiter-noinline-suspend.cpp
 create mode 100644 clang/test/CodeGenCoroutines/coro-simplify-suspend-point.cpp

diff --git a/clang/include/clang/AST/ExprCXX.h b/clang/include/clang/AST/ExprCXX.h
index a0e467b35778c5c..57a505036def409 100644
--- a/clang/include/clang/AST/ExprCXX.h
+++ b/clang/include/clang/AST/ExprCXX.h
@@ -5035,15 +5035,19 @@ class CoroutineSuspendExpr : public Expr {
   enum SubExpr { Operand, Common, Ready, Suspend, Resume, Count };
 
   Stmt *SubExprs[SubExpr::Count];
+  bool IsSuspendNoThrow = false;
   OpaqueValueExpr *OpaqueValue = nullptr;
+  OpaqueValueExpr *OpaqueFramePtr = nullptr;
 
 public:
   CoroutineSuspendExpr(StmtClass SC, SourceLocation KeywordLoc, Expr *Operand,
                        Expr *Common, Expr *Ready, Expr *Suspend, Expr *Resume,
-                       OpaqueValueExpr *OpaqueValue)
+                       bool IsSuspendNoThrow, OpaqueValueExpr *OpaqueValue,
+                       OpaqueValueExpr *OpaqueFramePtr)
       : Expr(SC, Resume->getType(), Resume->getValueKind(),
              Resume->getObjectKind()),
-        KeywordLoc(KeywordLoc), OpaqueValue(OpaqueValue) {
+        KeywordLoc(KeywordLoc), IsSuspendNoThrow(IsSuspendNoThrow),
+        OpaqueValue(OpaqueValue), OpaqueFramePtr(OpaqueFramePtr) {
     SubExprs[SubExpr::Operand] = Operand;
     SubExprs[SubExpr::Common] = Common;
     SubExprs[SubExpr::Ready] = Ready;
@@ -5080,6 +5084,9 @@ class CoroutineSuspendExpr : public Expr {
   /// getOpaqueValue - Return the opaque value placeholder.
   OpaqueValueExpr *getOpaqueValue() const { return OpaqueValue; }
 
+  /// getOpaqueFramePtr - Return coroutine frame pointer placeholder.
+  OpaqueValueExpr *getOpaqueFramePtr() const { return OpaqueFramePtr; }
+
   Expr *getReadyExpr() const {
     return static_cast<Expr*>(SubExprs[SubExpr::Ready]);
   }
@@ -5097,6 +5104,8 @@ class CoroutineSuspendExpr : public Expr {
     return static_cast<Expr *>(SubExprs[SubExpr::Operand]);
   }
 
+  bool isSuspendNoThrow() const { return IsSuspendNoThrow; }
+
   SourceLocation getKeywordLoc() const { return KeywordLoc; }
 
   SourceLocation getBeginLoc() const LLVM_READONLY { return KeywordLoc; }
@@ -5125,10 +5134,12 @@ class CoawaitExpr : public CoroutineSuspendExpr {
 
 public:
   CoawaitExpr(SourceLocation CoawaitLoc, Expr *Operand, Expr *Common,
-              Expr *Ready, Expr *Suspend, Expr *Resume,
-              OpaqueValueExpr *OpaqueValue, bool IsImplicit = false)
+              Expr *Ready, Expr *Suspend, Expr *Resume, bool IsSuspendNoThrow,
+              OpaqueValueExpr *OpaqueValue, OpaqueValueExpr *OpaqueFramePtr,
+              bool IsImplicit = false)
       : CoroutineSuspendExpr(CoawaitExprClass, CoawaitLoc, Operand, Common,
-                             Ready, Suspend, Resume, OpaqueValue) {
+                             Ready, Suspend, Resume, IsSuspendNoThrow,
+                             OpaqueValue, OpaqueFramePtr) {
     CoawaitBits.IsImplicit = IsImplicit;
   }
 
@@ -5206,10 +5217,11 @@ class CoyieldExpr : public CoroutineSuspendExpr {
 
 public:
   CoyieldExpr(SourceLocation CoyieldLoc, Expr *Operand, Expr *Common,
-              Expr *Ready, Expr *Suspend, Expr *Resume,
-              OpaqueValueExpr *OpaqueValue)
+              Expr *Ready, Expr *Suspend, Expr *Resume, bool IsSuspendNoThrow,
+              OpaqueValueExpr *OpaqueValue, OpaqueValueExpr *OpaqueFramePtr)
       : CoroutineSuspendExpr(CoyieldExprClass, CoyieldLoc, Operand, Common,
-                             Ready, Suspend, Resume, OpaqueValue) {}
+                             Ready, Suspend, Resume, IsSuspendNoThrow,
+                             OpaqueValue, OpaqueFramePtr) {}
   CoyieldExpr(SourceLocation CoyieldLoc, QualType Ty, Expr *Operand,
               Expr *Common)
       : CoroutineSuspendExpr(CoyieldExprClass, CoyieldLoc, Ty, Operand,
diff --git a/clang/lib/CodeGen/CGCoroutine.cpp b/clang/lib/CodeGen/CGCoroutine.cpp
index 888d30bfb3e1d6a..233a0137c751bf4 100644
--- a/clang/lib/CodeGen/CGCoroutine.cpp
+++ b/clang/lib/CodeGen/CGCoroutine.cpp
@@ -212,9 +212,10 @@ static LValueOrRValue emitSuspendExpression(CodeGenFunction &CGF, CGCoroData &Co
                                     bool ignoreResult, bool forLValue) {
   auto *E = S.getCommonExpr();
 
-  auto Binder =
+  auto CommonBinder =
       CodeGenFunction::OpaqueValueMappingData::bind(CGF, S.getOpaqueValue(), E);
-  auto UnbindOnExit = llvm::make_scope_exit([&] { Binder.unbind(CGF); });
+  auto UnbindCommonOnExit =
+      llvm::make_scope_exit([&] { CommonBinder.unbind(CGF); });
 
   auto Prefix = buildSuspendPrefixStr(Coro, Kind);
   BasicBlock *ReadyBlock = CGF.createBasicBlock(Prefix + Twine(".ready"));
@@ -232,16 +233,57 @@ static LValueOrRValue emitSuspendExpression(CodeGenFunction &CGF, CGCoroData &Co
   auto *NullPtr = llvm::ConstantPointerNull::get(CGF.CGM.Int8PtrTy);
   auto *SaveCall = Builder.CreateCall(CoroSave, {NullPtr});
 
-  CGF.CurCoro.InSuspendBlock = true;
-  auto *SuspendRet = CGF.EmitScalarExpr(S.getSuspendExpr());
-  CGF.CurCoro.InSuspendBlock = false;
+  auto SuspendHelper = CodeGenFunction(CGF.CGM).generateAwaitSuspendHelper(
+      CGF.CurFn->getName(), Prefix, S);
 
-  if (SuspendRet != nullptr && SuspendRet->getType()->isIntegerTy(1)) {
-    // Veto suspension if requested by bool returning await_suspend.
-    BasicBlock *RealSuspendBlock =
-        CGF.createBasicBlock(Prefix + Twine(".suspend.bool"));
-    CGF.Builder.CreateCondBr(SuspendRet, RealSuspendBlock, ReadyBlock);
-    CGF.EmitBlock(RealSuspendBlock);
+  llvm::CallBase *SuspendRet = nullptr;
+
+  {
+    CGF.CurCoro.InSuspendBlock = true;
+
+    auto FramePtrBinder = CodeGenFunction::OpaqueValueMappingData::bind(
+        CGF, S.getOpaqueFramePtr(), S.getOpaqueFramePtr()->getSourceExpr());
+    auto UnbindFramePtrOnExit =
+        llvm::make_scope_exit([&] { FramePtrBinder.unbind(CGF); });
+
+    SmallVector<llvm::Value *, 3> SuspendHelperCallArgs;
+    SuspendHelperCallArgs.push_back(
+        CGF.getOrCreateOpaqueLValueMapping(S.getOpaqueValue()).getPointer(CGF));
+    SuspendHelperCallArgs.push_back(
+        CGF.getOrCreateOpaqueRValueMapping(S.getOpaqueFramePtr())
+            .getScalarVal());
+    SuspendHelperCallArgs.push_back(SuspendHelper);
+
+    auto IID = llvm::Intrinsic::coro_await_suspend;
+    if (S.getSuspendExpr()->getType()->isBooleanType())
+      IID = llvm::Intrinsic::coro_await_suspend_bool;
+    else if (S.getSuspendExpr()->getType()->isVoidPointerType())
+      IID = llvm::Intrinsic::coro_await_suspend_handle;
+
+    llvm::Function *AwaitSuspendIntrinsic = CGF.CGM.getIntrinsic(IID);
+    // FIXME: add call attributes?
+    if (S.isSuspendNoThrow()) {
+      SuspendRet = CGF.EmitNounwindRuntimeCall(AwaitSuspendIntrinsic,
+                                               SuspendHelperCallArgs);
+    } else {
+      SuspendRet =
+          CGF.EmitCallOrInvoke(AwaitSuspendIntrinsic, SuspendHelperCallArgs);
+    }
+
+    CGF.CurCoro.InSuspendBlock = false;
+  }
+
+  if (SuspendRet != nullptr) {
+    if (SuspendRet->getType()->isIntegerTy(1)) {
+      // Veto suspension if requested by bool returning await_suspend.
+      BasicBlock *RealSuspendBlock =
+          CGF.createBasicBlock(Prefix + Twine(".suspend.bool"));
+      CGF.Builder.CreateCondBr(SuspendRet, RealSuspendBlock, ReadyBlock);
+      CGF.EmitBlock(RealSuspendBlock);
+    } else if (SuspendRet->getType()->isPointerTy()) {
+      auto ResumeIntrinsic = CGF.CGM.getIntrinsic(llvm::Intrinsic::coro_resume);
+      Builder.CreateCall(ResumeIntrinsic, SuspendRet);
+    }
   }
 
   // Emit the suspend point.
@@ -338,6 +380,86 @@ static QualType getCoroutineSuspendExprReturnType(const ASTContext &Ctx,
 }
 #endif
 
+llvm::Function *
+CodeGenFunction::generateAwaitSuspendHelper(Twine const &CoroName,
+                                            Twine const &SuspendPointName,
+                                            CoroutineSuspendExpr const &S) {
+  std::string FuncName = "__await_suspend_helper_";
+  FuncName += CoroName.str();
+  FuncName += '_';
+  FuncName += SuspendPointName.str();
+
+  ASTContext &C = getContext();
+
+  FunctionArgList args;
+
+  ImplicitParamDecl AwaiterDecl(C, C.VoidPtrTy, ImplicitParamKind::Other);
+  ImplicitParamDecl FrameDecl(C, C.VoidPtrTy, ImplicitParamKind::Other);
+  QualType ReturnTy = S.getSuspendExpr()->getType();
+
+  if (ReturnTy->isBooleanType()) {
+    ReturnTy = C.BoolTy;
+  } else if (ReturnTy->isVoidPointerType()) {
+    ReturnTy = C.VoidPtrTy;
+  } else {
+    ReturnTy = C.VoidTy;
+  }
+
+  args.push_back(&AwaiterDecl);
+  args.push_back(&FrameDecl);
+
+  const CGFunctionInfo &FI =
+      CGM.getTypes().arrangeBuiltinFunctionDeclaration(ReturnTy, args);
+
+  llvm::FunctionType *LTy = CGM.getTypes().GetFunctionType(FI);
+
+  llvm::Function *Fn = llvm::Function::Create(
+      LTy, llvm::GlobalValue::PrivateLinkage, FuncName, &CGM.getModule());
+
+  Fn->addParamAttr(0, llvm::Attribute::AttrKind::NonNull);
+  Fn->addParamAttr(0, llvm::Attribute::AttrKind::NoUndef);
+
+  Fn->addParamAttr(1, llvm::Attribute::AttrKind::NoUndef);
+
+  Fn->setMustProgress();
+  Fn->addFnAttr(llvm::Attribute::AttrKind::AlwaysInline);
+
+  if (S.isSuspendNoThrow()) {
+    Fn->addFnAttr(llvm::Attribute::AttrKind::NoUnwind);
+  }
+
+  StartFunction(GlobalDecl(), ReturnTy, Fn, FI, args);
+
+  llvm::Value *AwaiterPtr = Builder.CreateLoad(GetAddrOfLocalVar(&AwaiterDecl));
+  auto AwaiterLValue =
+      MakeNaturalAlignAddrLValue(AwaiterPtr, AwaiterDecl.getType());
+
+  // FIXME: mark as aliasing with awaiter?
+  // FIXME: TBAA?
+  // FIXME: emit in a better way (maybe egenerate AST in SemaCoroutine)?
+  auto FramePtrRValue =
+      RValue::get(Builder.CreateLoad(GetAddrOfLocalVar(&FrameDecl)));
+
+  auto AwaiterBinder = CodeGenFunction::OpaqueValueMappingData::bind(
+      *this, S.getOpaqueValue(), AwaiterLValue);
+  auto FramePtrBinder = CodeGenFunction::OpaqueValueMappingData::bind(
+      *this, S.getOpaqueFramePtr(), FramePtrRValue);
+
+  auto *SuspendRet = EmitScalarExpr(S.getSuspendExpr());
+
+  auto UnbindCommonOnExit =
+      llvm::make_scope_exit([&] { AwaiterBinder.unbind(*this); });
+  auto UnbindFramePtrOnExit =
+      llvm::make_scope_exit([&] { FramePtrBinder.unbind(*this); });
+  if (SuspendRet != nullptr) {
+    Fn->addRetAttr(llvm::Attribute::AttrKind::NoUndef);
+    Builder.CreateStore(SuspendRet, ReturnValue);
+  }
+
+  FinishFunction();
+  return Fn;
+}
+
 LValue
 CodeGenFunction::EmitCoawaitLValue(const CoawaitExpr *E) {
   assert(getCoroutineSuspendExprReturnType(getContext(), E)->isReferenceType() &&
diff --git a/clang/lib/CodeGen/CodeGenFunction.h b/clang/lib/CodeGen/CodeGenFunction.h
index 06327a184717788..ed6d8b318301ba0 100644
--- a/clang/lib/CodeGen/CodeGenFunction.h
+++ b/clang/lib/CodeGen/CodeGenFunction.h
@@ -352,6 +352,10 @@ class CodeGenFunction : public CodeGenTypeCache {
     return isCoroutine() && CurCoro.InSuspendBlock;
   }
 
+  llvm::Function *generateAwaitSuspendHelper(Twine const &CoroName,
+                                             Twine const &SuspendPointName,
+                                             CoroutineSuspendExpr const &S);
+
   /// CurGD - The GlobalDecl for the current function being compiled.
   GlobalDecl CurGD;
 
diff --git a/clang/lib/Sema/SemaCoroutine.cpp b/clang/lib/Sema/SemaCoroutine.cpp
index a969b9383563b22..a46744205d0f298 100644
--- a/clang/lib/Sema/SemaCoroutine.cpp
+++ b/clang/lib/Sema/SemaCoroutine.cpp
@@ -31,6 +31,8 @@
 using namespace clang;
 using namespace sema;
 
+static bool isNoThrow(Sema &S, const Stmt *E);
+
 static LookupResult lookupMember(Sema &S, const char *Name, CXXRecordDecl *RD,
                                  SourceLocation Loc, bool &Res) {
   DeclarationName DN = S.PP.getIdentifierInfo(Name);
@@ -266,7 +268,7 @@ static ExprResult buildOperatorCoawaitCall(Sema &SemaRef, Scope *S,
 }
 
 static ExprResult buildCoroutineHandle(Sema &S, QualType PromiseType,
-                                       SourceLocation Loc) {
+                                       Expr *FramePtr, SourceLocation Loc) {
   QualType CoroHandleType = lookupCoroutineHandleType(S, PromiseType, Loc);
   if (CoroHandleType.isNull())
     return ExprError();
@@ -280,9 +282,6 @@ static ExprResult buildCoroutineHandle(Sema &S, QualType PromiseType,
     return ExprError();
   }
 
-  Expr *FramePtr =
-      S.BuildBuiltinCallExpr(Loc, Builtin::BI__builtin_coro_frame, {});
-
   CXXScopeSpec SS;
   ExprResult FromAddr =
       S.BuildDeclarationNameExpr(SS, Found, /*NeedsADL=*/false);
@@ -296,6 +295,8 @@ struct ReadySuspendResumeResult {
   enum AwaitCallType { ACT_Ready, ACT_Suspend, ACT_Resume };
   Expr *Results[3];
   OpaqueValueExpr *OpaqueValue;
+  OpaqueValueExpr *OpaqueFramePtr;
+  bool IsSuspendNoThrow;
   bool IsInvalid;
 };
 
@@ -380,66 +381,7 @@ static Expr *maybeTailCall(Sema &S, QualType RetType, Expr *E,
   // __builtin_coro_resume so that the cleanup code are not inserted in-between
   // the resume call and return instruction, which would interfere with the
   // musttail call contract.
-  JustAddress = S.MaybeCreateExprWithCleanups(JustAddress);
-  return S.BuildBuiltinCallExpr(Loc, Builtin::BI__builtin_coro_resume,
-                                JustAddress);
-}
-
-/// The await_suspend call performed by co_await is essentially asynchronous
-/// to the execution of the coroutine. Inlining it normally into an unsplit
-/// coroutine can cause miscompilation because the coroutine CFG misrepresents
-/// the true control flow of the program: things that happen in the
-/// await_suspend are not guaranteed to happen prior to the resumption of the
-/// coroutine, and things that happen after the resumption of the coroutine
-/// (including its exit and the potential deallocation of the coroutine frame)
-/// are not guaranteed to happen only after the end of await_suspend.
-///
-/// See https://github.com/llvm/llvm-project/issues/56301 and
-/// https://reviews.llvm.org/D157070 for the example and the full discussion.
-///
-/// The short-term solution to this problem is to mark the call as uninlinable.
-/// But we don't want to do this if the call is known to be trivial, which is
-/// very common.
-///
-/// The long-term solution may introduce patterns like:
-///
-///  call @llvm.coro.await_suspend(ptr %awaiter, ptr %handle,
-///                                ptr @awaitSuspendFn)
-///
-/// Then it is much easier to perform the safety analysis in the middle end.
-/// If it is safe to inline the call to awaitSuspend, we can replace it in the
-/// CoroEarly pass. Otherwise we could replace it in the CoroSplit pass.
-static void tryMarkAwaitSuspendNoInline(Sema &S, OpaqueValueExpr *Awaiter,
-                                        CallExpr *AwaitSuspend) {
-  // The method here to extract the awaiter decl is not precise.
-  // This is intentional. Since it is hard to perform the analysis in the
-  // frontend due to the complexity of C++'s type systems.
-  // And we prefer to perform such analysis in the middle end since it is
-  // easier to implement and more powerful.
-  CXXRecordDecl *AwaiterDecl =
-      Awaiter->getType().getNonReferenceType()->getAsCXXRecordDecl();
-
-  if (AwaiterDecl && AwaiterDecl->field_empty())
-    return;
-
-  FunctionDecl *FD = AwaitSuspend->getDirectCallee();
-
-  assert(FD);
-
-  // If the `await_suspend()` function is marked as `always_inline` explicitly,
-  // we should give the user the right to control the codegen.
-  if (FD->hasAttr<NoInlineAttr>() || FD->hasAttr<AlwaysInlineAttr>())
-    return;
-
-  // This is problematic if the user calls the await_suspend standalone. But on
-  // the on hand, it is not incorrect semantically since inlining is not part
-  // of the standard. On the other hand, it is relatively rare to call
-  // the await_suspend function standalone.
-  //
-  // And given we've already had the long-term plan, the current workaround
-  // looks relatively tolerant.
-  FD->addAttr(
-      NoInlineAttr::CreateImplicit(S.getASTContext(), FD->getLocation()));
+  return S.MaybeCreateExprWithCleanups(JustAddress);
 }
 
 /// Build calls to await_ready, await_suspend, and await_resume for a co_await
@@ -461,7 +403,11 @@ static ReadySuspendResumeResult buildCoawaitCalls(Sema &S, VarDecl *CoroPromise,
 
   // Assume valid until we see otherwise.
   // Further operations are responsible for setting IsInalid to true.
-  ReadySuspendResumeResult Calls = {{}, Operand, /*IsInvalid=*/false};
+  ReadySuspendResumeResult Calls = {{},
+                                    Operand,
+                                    /*OpaqueFramePtr=*/nullptr,
+                                    /*IsSuspendNoThrow=*/false,
+                                    /*IsInvalid=*/false};
 
   using ACT = ReadySuspendResumeResult::AwaitCallType;
 
@@ -495,8 +441,17 @@ static ReadySuspendResumeResult buildCoawaitCalls(Sema &S, VarDecl *CoroPromise,
       Calls.Results[ACT::ACT_Ready] = S.MaybeCreateExprWithCleanups(Conv.get());
   }
 
+  Expr *GetFramePtr =
+      S.BuildBuiltinCallExpr(Loc, Builtin::BI__builtin_coro_frame, {});
+
+  OpaqueValueExpr *FramePtr = new (S.Context)
+      OpaqueValueExpr(Loc, GetFramePtr->getType(), VK_PRValue,
+                      GetFramePtr->getObjectKind(), GetFramePtr);
+
+  Calls.OpaqueFramePtr = FramePtr;
+
   ExprResult CoroHandleRes =
-      buildCoroutineHandle(S, CoroPromise->getType(), Loc);
+      buildCoroutineHandle(S, CoroPromise->getType(), FramePtr, Loc);
   if (CoroHandleRes.isInvalid()) {
     Calls.IsInvalid = true;
     return Calls;
@@ -513,10 +468,6 @@ static ReadySuspendResumeResult buildCoawaitCalls(Sema &S, VarDecl *CoroPromise,
     //     type Z.
     QualType RetType = AwaitSuspend->getCallReturnType(S.Context);
 
-    // We need to mark await_suspend as noinline temporarily. See the comment
-    // of tryMarkAwaitSuspendNoInline for details.
-    tryMarkAwaitSuspendNoInline(S, Operand, AwaitSuspend);
-
     // Support for coroutine_handle returning await_suspend.
     if (Expr *TailCallSuspend =
             maybeTailCall(S, RetType, AwaitSuspend, Loc))
@@ -542,6 +493,10 @@ static ReadySuspendResumeResult buildCoawaitCalls(Sema &S, VarDecl *CoroPromise,
     }
   }
 
+  if (Calls.Results[ACT::ACT_Suspend]) {
+    Calls.IsSuspendNoThrow = isNoThrow(S, Calls.Results[ACT::ACT_Suspend]);
+  }
+
   BuildSubExpr(ACT::ACT_Resume, "await_resume", std::nullopt);
 
   // Make sure the awaiter object gets a chance to be cleaned up.
@@ -694,6 +649,59 @@ static FunctionScopeInfo *checkCoroutineContext(Sema &S, SourceLocation Loc,
   return ScopeInfo;
 }
 
+/// Recursively check \p E and all its children to see if any call target
+/// (including constructor call) is declared noexcept. Also any value returned
+/// from the call has a noexcept destructor.
+static bool isNoThrow(Sema &S, const Stmt *E) {
+  auto isDeclNoexcept = [&](const Decl *D, bool IsDtor = false) -> bool {
+    // In the case of dtor, the call to dtor is implicit and hence we should
+    // pass nullptr to canCalleeThrow.
+    if (Sema::canCalleeThrow(S, IsDtor ? nullptr : cast<Expr>(E), D)) {
+      return false;
+    }
+    return true;
+  };
+
+  if (auto *CE = dyn_cast<CXXConstructExpr>(E)) {
+    CXXConstructorDecl *Ctor = CE->getConstructor();
+    if (!isDeclNoexcept(Ctor)) {
+      return false;
+    }
+    // Check the corresponding destructor of the constructor.
+    if (!isDeclNoexcept(Ctor->getParent()->getDestructor(), /*IsDtor=*/true)) {
+      return false;
+    }
+  } else if (auto *CE = dyn_cast<CallExpr>(E)) {
+    if (CE->isTypeDependent())
+      return false;
+
+    if (!isDeclNoexcept(CE->getCalleeDecl())) {
+      return false;
+    }
+
+    QualType ReturnType = CE->getCallReturnType(S.getASTContext());
+    // Check the destructor of the call return type, if any.
+    if (ReturnType.isDestructedType() ==
+        QualType::DestructionKind::DK_cxx_destructor) {
+      const auto *T =
+          cast<RecordType>(ReturnType.getCanonicalType().getTypePtr());
+      if (!isDeclNoexcept(cast<CXXRecordDecl>(T->getDecl())->getDestructor(),
+                          /*IsDtor=*/true)) {
+        return false;
+      }
+    }
+  }
+  for (const auto *Child : E->children()) {
+    if (!Child)
+      continue;
+    if (!isNoThrow(S, Child)) {
+      return false;
+    }
+  }
+
+  return true;
+}
+
 /// Recursively check \p E and all its children to see if any call target
 /// (including constructor call) is declared noexcept. Also any value returned
 /// from the call has a noexcept destructor.
@@ -992,9 +1000,9 @@ ExprResult Sema::BuildResolvedCoawaitExpr(SourceLocation Loc, Expr *Operand,
   if (RSS.IsInvalid)
     return ExprError();
 
-  Expr *Res = new (Context)
-      CoawaitExpr(Loc, Operand, Awaiter, RSS.Results[0], RSS.Results[1],
-                  RSS.Results[2], RSS.OpaqueValue, IsImplicit);
+  Expr *Res = new (Context) CoawaitExpr(
+      Loc, Operand, Awaiter, RSS.Results[0], RSS.Results[1], RSS.Results[2],
+      RSS.IsSuspendNoThrow, RSS.OpaqueValue, RSS.OpaqueFramePtr, IsImplicit);
 
   return Res;
 }
@@ -1050,9 +1058,9 @@ ExprResult Sema::BuildCoyieldExpr(SourceLocation Loc, Expr *E) {
   if (RSS.IsInvalid)
     return ExprError();
 
-  Expr *Res =
-      new (Context) CoyieldExpr(Loc, Operand, E, RSS.Results[0], RSS.Results[1],
-                                RSS.Results[2], RSS.OpaqueValue);
+  Expr *Res = new (Context) CoyieldExpr(
+      Loc, Operand, E, RSS.Results[0], RSS.Results[1], RSS.Results[2],
+      RSS.IsSuspendNoThrow, RSS.OpaqueValue, RSS.OpaqueFramePtr);
 
   return Res;
 }
diff --git a/clang/test/AST/coroutine-locals-cleanup.cpp b/clang/test/AST/coroutine-locals-cleanup.cpp
index ce106b8e230a10e..6264df01fa2acb2 100644
--- a/clang/test/AST/coroutine-locals-cleanup.cpp
+++ b/clang/test/AST/coroutine-locals-cleanup.cpp
@@ -90,10 +90,7 @@ Task bar() {
 // CHECK:                 ExprWithCleanups {{.*}} 'bool'
 // CHECK-NEXT:              CXXMemberCallExpr {{.*}} 'bool'
 // CHECK-NEXT:                MemberExpr {{.*}} .await_ready
-// CHECK:                 CallExpr {{.*}} 'void'
-// CHECK-NEXT:              ImplicitCastExpr {{.*}} 'void (*)(void *)'
-// CHECK-NEXT:                DeclRefExpr {{.*}} '__builtin_coro_resume' 'void (void *)'
-// CHECK-NEXT:              ExprWithCleanups {{.*}} 'void *'
+// CHECK:                 ExprWithCleanups {{.*}} 'void *'
 
 // CHECK:           CaseStmt
 // CHECK:             ExprWithCleanups {{.*}} 'void'
@@ -103,7 +100,4 @@ Task bar() {
 // CHECK:                 ExprWithCleanups {{.*}} 'bool'
 // CHECK-NEXT:              CXXMemberCallExpr {{.*}} 'bool'
 // CHECK-NEXT:                MemberExpr {{.*}} .await_ready
-// CHECK:                 CallExpr {{.*}} 'void'
-// CHECK-NEXT:              ImplicitCastExpr {{.*}} 'void (*)(void *)'
-// CHECK-NEXT:                DeclRefExpr {{.*}} '__builtin_coro_resume' 'void (void *)'
-// CHECK-NEXT:              ExprWithCleanups {{.*}} 'void *'
+// CHECK:                 ExprWithCleanups {{.*}} 'void *'
diff --git a/clang/test/CodeGenCoroutines/coro-always-inline.cpp b/clang/test/CodeGenCoroutines/coro-always-inline.cpp
index 6e13a62fbd98659..d4f67a73f517268 100644
--- a/clang/test/CodeGenCoroutines/coro-always-inline.cpp
+++ b/clang/test/CodeGenCoroutines/coro-always-inline.cpp
@@ -34,7 +34,7 @@ struct coroutine_traits {
 // CHECK-LABEL: @_Z3foov
 // CHECK-LABEL: entry:
 // CHECK: %ref.tmp.reload.addr = getelementptr
-// CHECK: %ref.tmp4.reload.addr = getelementptr
+// CHECK: %ref.tmp3.reload.addr = getelementptr
 void foo() { co_return; }
 
 // Check that bar is not inlined even it's marked as always_inline.
diff --git a/clang/test/CodeGenCoroutines/coro-await.cpp b/clang/test/CodeGenCoroutines/coro-await.cpp
index dc5a765ccb83d78..701e1bf87b06f5d 100644
--- a/clang/test/CodeGenCoroutines/coro-await.cpp
+++ b/clang/test/CodeGenCoroutines/coro-await.cpp
@@ -71,16 +71,13 @@ extern "C" void f0() {
   // CHECK: [[SUSPEND_BB]]:
   // CHECK: %[[SUSPEND_ID:.+]] = call token @llvm.coro.save(
   // ---------------------------
-  // Build the coroutine handle and pass it to await_suspend
+  // Call coro.await.suspend
   // ---------------------------
-  // CHECK: call ptr @_ZNSt16coroutine_handleINSt16coroutine_traitsIJvEE12promise_typeEE12from_addressEPv(ptr %[[FRAME]])
-  //   ... many lines of code to coerce coroutine_handle into an ptr scalar
-  // CHECK: %[[CH:.+]] = load ptr, ptr %{{.+}}
-  // CHECK: call void @_ZN14suspend_always13await_suspendESt16coroutine_handleIvE(ptr {{[^,]*}} %[[AWAITABLE]], ptr %[[CH]])
+  // CHECK-NEXT: call void @llvm.coro.await.suspend(ptr %[[AWAITABLE]], ptr %[[FRAME]], ptr @__await_suspend_helper_f0_await)
   // -------------------------
   // Generate a suspend point:
   // -------------------------
-  // CHECK: %[[OUTCOME:.+]] = call i8 @llvm.coro.suspend(token %[[SUSPEND_ID]], i1 false)
+  // CHECK-NEXT: %[[OUTCOME:.+]] = call i8 @llvm.coro.suspend(token %[[SUSPEND_ID]], i1 false)
   // CHECK: switch i8 %[[OUTCOME]], label %[[RET_BB:.+]] [
   // CHECK:   i8 0, label %[[READY_BB]]
   // CHECK:   i8 1, label %[[CLEANUP_BB:.+]]
@@ -101,6 +98,17 @@ extern "C" void f0() {
   // CHECK-NEXT: call zeroext i1 @_ZN10final_susp11await_readyEv(ptr
   // CHECK: %[[FINALSP_ID:.+]] = call token @llvm.coro.save(
   // CHECK: call i8 @llvm.coro.suspend(token %[[FINALSP_ID]], i1 true)
+
+  // await suspend helper
+  // CHECK: define{{.*}} @__await_suspend_helper_f0_await(ptr {{[^,]*}} %[[AWAITABLE_ARG:.+]], ptr {{[^,]*}} %[[FRAME_ARG:.+]])
+  // CHECK: store ptr %[[AWAITABLE_ARG]], ptr %[[AWAITABLE_TMP:.+]],
+  // CHECK: store ptr %[[FRAME_ARG]], ptr %[[FRAME_TMP:.+]],
+  // CHECK: %[[AWAITABLE:.+]] = load ptr, ptr %[[AWAITABLE_TMP]]
+  // CHECK: %[[FRAME:.+]] = load ptr, ptr %[[FRAME_TMP]]
+  // CHECK: call ptr @_ZNSt16coroutine_handleINSt16coroutine_traitsIJvEE12promise_typeEE12from_addressEPv(ptr %[[FRAME]])
+  //   ... many lines of code to coerce coroutine_handle into an ptr scalar
+  // CHECK: %[[CH:.+]] = load ptr, ptr %{{.+}}
+  // CHECK: call void @_ZN14suspend_always13await_suspendESt16coroutine_handleIvE(ptr {{[^,]*}} %[[AWAITABLE]], ptr %[[CH]])
 }
 
 struct suspend_maybe {
@@ -131,7 +139,7 @@ extern "C" void f1(int) {
 
   // See if we need to suspend:
   // --------------------------
-  // CHECK: %[[READY:.+]] = call zeroext i1 @_ZN13suspend_maybe11await_readyEv(ptr {{[^,]*}} %[[AWAITABLE]])
+  // CHECK: %[[READY:.+]] = call zeroext i1 @_ZN13suspend_maybe11await_readyEv(ptr {{[^,]*}} %[[AWAITABLE:.+]])
   // CHECK: br i1 %[[READY]], label %[[READY_BB:.+]], label %[[SUSPEND_BB:.+]]
 
   // If we are suspending:
@@ -139,12 +147,9 @@ extern "C" void f1(int) {
   // CHECK: [[SUSPEND_BB]]:
   // CHECK: %[[SUSPEND_ID:.+]] = call token @llvm.coro.save(
   // ---------------------------
-  // Build the coroutine handle and pass it to await_suspend
+  // Call coro.await.suspend
   // ---------------------------
-  // CHECK: call ptr @_ZNSt16coroutine_handleINSt16coroutine_traitsIJviEE12promise_typeEE12from_addressEPv(ptr %[[FRAME]])
-  //   ... many lines of code to coerce coroutine_handle into an ptr scalar
-  // CHECK: %[[CH:.+]] = load ptr, ptr %{{.+}}
-  // CHECK: %[[YES:.+]] = call zeroext i1 @_ZN13suspend_maybe13await_suspendESt16coroutine_handleIvE(ptr {{[^,]*}} %[[AWAITABLE]], ptr %[[CH]])
+  // CHECK-NEXT: %[[YES:.+]] = call i1 @llvm.coro.await.suspend.bool(ptr %[[AWAITABLE]], ptr %[[FRAME]], ptr @__await_suspend_helper_f1_yield)
   // -------------------------------------------
   // See if await_suspend decided not to suspend
   // -------------------------------------------
@@ -155,6 +160,18 @@ extern "C" void f1(int) {
 
   // CHECK: [[READY_BB]]:
   // CHECK:     call void @_ZN13suspend_maybe12await_resumeEv(ptr {{[^,]*}} %[[AWAITABLE]])
+
+  // Await suspend helper
+  // CHECK: define {{.*}} i1 @__await_suspend_helper_f1_yield(ptr {{[^,]*}} %[[AWAITABLE_ARG:.+]], ptr {{[^,]*}} %[[FRAME_ARG:.+]])
+  // CHECK: store ptr %[[AWAITABLE_ARG]], ptr %[[AWAITABLE_TMP:.+]],
+  // CHECK: store ptr %[[FRAME_ARG]], ptr %[[FRAME_TMP:.+]],
+  // CHECK: %[[AWAITABLE:.+]] = load ptr, ptr %[[AWAITABLE_TMP]]
+  // CHECK: %[[FRAME:.+]] = load ptr, ptr %[[FRAME_TMP]]
+  // CHECK: call ptr @_ZNSt16coroutine_handleINSt16coroutine_traitsIJviEE12promise_typeEE12from_addressEPv(ptr %[[FRAME]])
+  //   ... many lines of code to coerce coroutine_handle into an ptr scalar
+  // CHECK: %[[CH:.+]] = load ptr, ptr %{{.+}}
+  // CHECK: %[[YES:.+]] = call zeroext i1 @_ZN13suspend_maybe13await_suspendESt16coroutine_handleIvE(ptr {{[^,]*}} %[[AWAITABLE]], ptr %[[CH]]) 
+  // CHECK-NEXT: ret i1 %[[YES]]
 }
 
 struct ComplexAwaiter {
@@ -340,11 +357,39 @@ struct TailCallAwait {
 
 // CHECK-LABEL: @TestTailcall(
 extern "C" void TestTailcall() {
+  // CHECK: %[[PROMISE:.+]] = alloca %"struct.std::coroutine_traits<void>::promise_type"
+  // CHECK: %[[FRAME:.+]] = call ptr @llvm.coro.begin(
   co_await TailCallAwait{};
+  // CHECK: %[[READY:.+]] = call zeroext i1 @_ZN13TailCallAwait11await_readyEv(ptr {{[^,]*}} %[[AWAITABLE:.+]])
+  // CHECK: br i1 %[[READY]], label %[[READY_BB:.+]], label %[[SUSPEND_BB:.+]]
 
-  // CHECK: %[[RESULT:.+]] = call ptr @_ZN13TailCallAwait13await_suspendESt16coroutine_handleIvE(ptr
-  // CHECK: %[[COERCE:.+]] = getelementptr inbounds %"struct.std::coroutine_handle", ptr %[[TMP:.+]], i32 0, i32 0
-  // CHECK: store ptr %[[RESULT]], ptr %[[COERCE]]
-  // CHECK: %[[ADDR:.+]] = call ptr @_ZNSt16coroutine_handleIvE7addressEv(ptr {{[^,]*}} %[[TMP]])
-  // CHECK: call void @llvm.coro.resume(ptr %[[ADDR]])
+  // If we are suspending:
+  // ---------------------
+  // CHECK: [[SUSPEND_BB]]:
+  // CHECK: %[[SUSPEND_ID:.+]] = call token @llvm.coro.save(
+  // ---------------------------
+  // Call coro.await.suspend
+  // ---------------------------
+  // CHECK-NEXT: %[[RESUMED:.+]] = call ptr @llvm.coro.await.suspend.handle(ptr %[[AWAITABLE]], ptr %[[FRAME]], ptr @__await_suspend_helper_TestTailcall_await)
+  // CHECK-NEXT: call void @llvm.coro.resume(ptr %[[RESUMED]])
+  // CHECK-NEXT: %[[OUTCOME:.+]] = call i8 @llvm.coro.suspend(token %[[SUSPEND_ID]], i1 false)
+  // CHECK-NEXT: switch i8 %[[OUTCOME]], label %[[RET_BB:.+]] [
+  // CHECK-NEXT:   i8 0, label %[[READY_BB]]
+  // CHECK-NEXT:   i8 1, label %{{.+}}
+  // CHECK-NEXT: ]
+
+  // Await suspend helper
+  // CHECK: define {{.*}} ptr @__await_suspend_helper_TestTailcall_await(ptr {{[^,]*}} %[[AWAITABLE_ARG:.+]], ptr {{[^,]*}} %[[FRAME_ARG:.+]])
+  // CHECK: store ptr %[[AWAITABLE_ARG]], ptr %[[AWAITABLE_TMP:.+]],
+  // CHECK: store ptr %[[FRAME_ARG]], ptr %[[FRAME_TMP:.+]],
+  // CHECK: %[[AWAITABLE:.+]] = load ptr, ptr %[[AWAITABLE_TMP]]
+  // CHECK: %[[FRAME:.+]] = load ptr, ptr %[[FRAME_TMP]]
+  // CHECK: call ptr  @_ZNSt16coroutine_handleINSt16coroutine_traitsIJvEE12promise_typeEE12from_addressEPv(ptr %[[FRAME]])
+  //   ... many lines of code to coerce coroutine_handle into an ptr scalar
+  // CHECK: %[[CH:.+]] = load ptr, ptr %{{.+}}
+  // CHECK-NEXT: %[[RESULT:.+]] = call ptr @_ZN13TailCallAwait13await_suspendESt16coroutine_handleIvE(ptr {{[^,]*}} %[[AWAITABLE]], ptr %[[CH]]) 
+  // CHECK-NEXT: %[[COERCE:.+]] = getelementptr inbounds %"struct.std::coroutine_handle", ptr %[[TMP:.+]], i32 0, i32 0
+  // CHECK-NEXT: store ptr %[[RESULT]], ptr %[[COERCE]]
+  // CHECK-NEXT: %[[ADDR:.+]] = call ptr @_ZNSt16coroutine_handleIvE7addressEv(ptr {{[^,]*}} %[[TMP]])
+  // CHECK-NEXT: ret ptr %[[ADDR]]
 }
diff --git a/clang/test/CodeGenCoroutines/coro-awaiter-noinline-suspend.cpp b/clang/test/CodeGenCoroutines/coro-awaiter-noinline-suspend.cpp
deleted file mode 100644
index f95286faf46ec8d..000000000000000
--- a/clang/test/CodeGenCoroutines/coro-awaiter-noinline-suspend.cpp
+++ /dev/null
@@ -1,168 +0,0 @@
-// Tests that we can mark await-suspend as noinline correctly.
-//
-// RUN: %clang_cc1 -std=c++20 -triple x86_64-unknown-linux-gnu -emit-llvm -o - %s \
-// RUN:     -O1 -disable-llvm-passes | FileCheck %s
-
-#include "Inputs/coroutine.h"
-
-struct Task {
-  struct promise_type {
-    struct FinalAwaiter {
-      bool await_ready() const noexcept { return false; }
-      template <typename PromiseType>
-      std::coroutine_handle<> await_suspend(std::coroutine_handle<PromiseType> h) noexcept {
-        return h.promise().continuation;
-      }
-      void await_resume() noexcept {}
-    };
-
-    Task get_return_object() noexcept {
-      return std::coroutine_handle<promise_type>::from_promise(*this);
-    }
-
-    std::suspend_always initial_suspend() noexcept { return {}; }
-    FinalAwaiter final_suspend() noexcept { return {}; }
-    void unhandled_exception() noexcept {}
-    void return_void() noexcept {}
-
-    std::coroutine_handle<> continuation;
-  };
-
-  Task(std::coroutine_handle<promise_type> handle);
-  ~Task();
-
-private:
-  std::coroutine_handle<promise_type> handle;
-};
-
-struct StatefulAwaiter {
-    int value;
-    bool await_ready() const noexcept { return false; }
-    template <typename PromiseType>
-    void await_suspend(std::coroutine_handle<PromiseType> h) noexcept {}
-    void await_resume() noexcept {}
-};
-
-typedef std::suspend_always NoStateAwaiter;
-using AnotherStatefulAwaiter = StatefulAwaiter;
-
-template <class T>
-struct TemplatedAwaiter {
-    T value;
-    bool await_ready() const noexcept { return false; }
-    template <typename PromiseType>
-    void await_suspend(std::coroutine_handle<PromiseType> h) noexcept {}
-    void await_resume() noexcept {}
-};
-
-
-class Awaitable {};
-StatefulAwaiter operator co_await(Awaitable) {
-  return StatefulAwaiter{};
-}
-
-StatefulAwaiter GlobalAwaiter;
-class Awaitable2 {};
-StatefulAwaiter& operator co_await(Awaitable2) {
-  return GlobalAwaiter;
-}
-
-struct AlwaysInlineStatefulAwaiter {
-    void* value;
-    bool await_ready() const noexcept { return false; }
-
-    template <typename PromiseType>
-    __attribute__((always_inline))
-    void await_suspend(std::coroutine_handle<PromiseType> h) noexcept {}
-
-    void await_resume() noexcept {}
-};
-
-Task testing() {
-    co_await std::suspend_always{};
-    co_await StatefulAwaiter{};
-    co_await AnotherStatefulAwaiter{};
-    
-    // Test lvalue case.
-    StatefulAwaiter awaiter;
-    co_await awaiter;
-
-    // The explicit call to await_suspend is not considered suspended.
-    awaiter.await_suspend(std::coroutine_handle<void>::from_address(nullptr));
-
-    co_await TemplatedAwaiter<int>{};
-    TemplatedAwaiter<int> TemplatedAwaiterInstace;
-    co_await TemplatedAwaiterInstace;
-
-    co_await Awaitable{};
-    co_await Awaitable2{};
-
-    co_await AlwaysInlineStatefulAwaiter{};
-}
-
-struct AwaitTransformTask {
-  struct promise_type {
-    struct FinalAwaiter {
-      bool await_ready() const noexcept { return false; }
-      template <typename PromiseType>
-      std::coroutine_handle<> await_suspend(std::coroutine_handle<PromiseType> h) noexcept {
-        return h.promise().continuation;
-      }
-      void await_resume() noexcept {}
-    };
-
-    AwaitTransformTask get_return_object() noexcept {
-      return std::coroutine_handle<promise_type>::from_promise(*this);
-    }
-
-    std::suspend_always initial_suspend() noexcept { return {}; }
-    FinalAwaiter final_suspend() noexcept { return {}; }
-    void unhandled_exception() noexcept {}
-    void return_void() noexcept {}
-
-    template <typename Awaitable>
-    auto await_transform(Awaitable &&awaitable) {
-      return awaitable;
-    }
-
-    std::coroutine_handle<> continuation;
-  };
-
-  AwaitTransformTask(std::coroutine_handle<promise_type> handle);
-  ~AwaitTransformTask();
-
-private:
-  std::coroutine_handle<promise_type> handle;
-};
-
-struct awaitableWithGetAwaiter {
-  bool await_ready() const noexcept { return false; }
-  template <typename PromiseType>
-  void await_suspend(std::coroutine_handle<PromiseType> h) noexcept {}
-  void await_resume() noexcept {}
-};
-
-AwaitTransformTask testingWithAwaitTransform() {
-  co_await awaitableWithGetAwaiter{};
-}
-
-// CHECK: define{{.*}}@_ZNSt14suspend_always13await_suspendESt16coroutine_handleIvE{{.*}}#[[NORMAL_ATTR:[0-9]+]]
-
-// CHECK: define{{.*}}@_ZN15StatefulAwaiter13await_suspendIN4Task12promise_typeEEEvSt16coroutine_handleIT_E{{.*}}#[[NOINLINE_ATTR:[0-9]+]]
-
-// CHECK: define{{.*}}@_ZN15StatefulAwaiter13await_suspendIvEEvSt16coroutine_handleIT_E{{.*}}#[[NORMAL_ATTR]]
-
-// CHECK: define{{.*}}@_ZN16TemplatedAwaiterIiE13await_suspendIN4Task12promise_typeEEEvSt16coroutine_handleIT_E{{.*}}#[[NOINLINE_ATTR]]
-
-// CHECK: define{{.*}}@_ZN27AlwaysInlineStatefulAwaiter13await_suspendIN4Task12promise_typeEEEvSt16coroutine_handleIT_E{{.*}}#[[ALWAYS_INLINE_ATTR:[0-9]+]]
-
-// CHECK: define{{.*}}@_ZN4Task12promise_type12FinalAwaiter13await_suspendIS0_EESt16coroutine_handleIvES3_IT_E{{.*}}#[[NORMAL_ATTR]]
-
-// CHECK: define{{.*}}@_ZN23awaitableWithGetAwaiter13await_suspendIN18AwaitTransformTask12promise_typeEEEvSt16coroutine_handleIT_E{{.*}}#[[NORMAL_ATTR]]
-
-// CHECK: define{{.*}}@_ZN18AwaitTransformTask12promise_type12FinalAwaiter13await_suspendIS0_EESt16coroutine_handleIvES3_IT_E{{.*}}#[[NORMAL_ATTR]]
-
-// CHECK-NOT: attributes #[[NORMAL_ATTR]] = noinline
-// CHECK: attributes #[[NOINLINE_ATTR]] = {{.*}}noinline
-// CHECK-NOT: attributes #[[ALWAYS_INLINE_ATTR]] = {{.*}}noinline
-// CHECK: attributes #[[ALWAYS_INLINE_ATTR]] = {{.*}}alwaysinline
diff --git a/clang/test/CodeGenCoroutines/coro-dwarf.cpp b/clang/test/CodeGenCoroutines/coro-dwarf.cpp
index 7914babe5483a45..7c44d1d4bef6568 100644
--- a/clang/test/CodeGenCoroutines/coro-dwarf.cpp
+++ b/clang/test/CodeGenCoroutines/coro-dwarf.cpp
@@ -70,3 +70,15 @@ void f_coro(int val, MoveOnly moParam, MoveAndCopy mcParam) {
 // CHECK: !{{[0-9]+}} = !DILocalVariable(name: "moParam", arg: 2, scope: ![[SP]], file: !{{[0-9]+}}, line: {{[0-9]+}}, type: !{{[0-9]+}})
 // CHECK: !{{[0-9]+}} = !DILocalVariable(name: "mcParam", arg: 3, scope: ![[SP]], file: !{{[0-9]+}}, line: {{[0-9]+}}, type: !{{[0-9]+}})
 // CHECK: !{{[0-9]+}} = !DILocalVariable(name: "__promise",
+
+// CHECK: !{{[0-9]+}} = distinct !DISubprogram(linkageName: "__await_suspend_helper__Z6f_coroi8MoveOnly11MoveAndCopy_init"
+// CHECK-NEXT: !{{[0-9]+}} = !DIFile
+// CHECK-NEXT: !{{[0-9]+}} = !DISubroutineType
+// CHECK-NEXT: !{{[0-9]+}} = !DILocalVariable(arg: 1,
+// CHECK-NEXT: !{{[0-9]+}} = !DILocation
+// CHECK-NEXT: !{{[0-9]+}} = !DILocalVariable(arg: 2,
+
+// CHECK: !{{[0-9]+}} = distinct !DISubprogram(linkageName: "__await_suspend_helper__Z6f_coroi8MoveOnly11MoveAndCopy_final"
+// CHECK-NEXT: !{{[0-9]+}} = !DILocalVariable(arg: 1,
+// CHECK-NEXT: !{{[0-9]+}} = !DILocation
+// CHECK-NEXT: !{{[0-9]+}} = !DILocalVariable(arg: 2,
diff --git a/clang/test/CodeGenCoroutines/coro-function-try-block.cpp b/clang/test/CodeGenCoroutines/coro-function-try-block.cpp
index f609eb55b8e771f..023a95870361fd2 100644
--- a/clang/test/CodeGenCoroutines/coro-function-try-block.cpp
+++ b/clang/test/CodeGenCoroutines/coro-function-try-block.cpp
@@ -19,5 +19,5 @@ task f() try {
 }
 
 // CHECK-LABEL: define{{.*}} void @_Z1fv(
-// CHECK: call void @_ZNSt13suspend_never13await_suspendESt16coroutine_handleIvE(
+// CHECK: call void @llvm.coro.await.suspend(
 // CHECK: call void @_ZN4task12promise_type11return_voidEv(
diff --git a/clang/test/CodeGenCoroutines/coro-simplify-suspend-point.cpp b/clang/test/CodeGenCoroutines/coro-simplify-suspend-point.cpp
new file mode 100644
index 000000000000000..32960311043c804
--- /dev/null
+++ b/clang/test/CodeGenCoroutines/coro-simplify-suspend-point.cpp
@@ -0,0 +1,66 @@
+// Test that we can perform suspend point simplification
+// RUN: %clang_cc1 -triple x86_64-unknown-linux-gnu -std=c++20 -O1 -emit-llvm %s -o - | FileCheck %s
+// RUN: %clang -std=c++20 -O1 -emit-llvm -c  %s -o %t && %clang -c %t
+
+#include "Inputs/coroutine.h"
+
+struct detached_task {
+  struct promise_type {
+    detached_task get_return_object() noexcept {
+      return detached_task{std::coroutine_handle<promise_type>::from_promise(*this)};
+    }
+
+    void return_void() noexcept {}
+
+    struct final_awaiter {
+      bool await_ready() noexcept { return false; }
+      std::coroutine_handle<> await_suspend(std::coroutine_handle<promise_type> h) noexcept {
+        h.destroy();
+        return std::noop_coroutine();
+      }
+      void await_resume() noexcept {}
+    };
+
+    void unhandled_exception() noexcept {}
+
+    final_awaiter final_suspend() noexcept { return {}; }
+
+    std::suspend_always initial_suspend() noexcept { return {}; }
+  };
+
+  ~detached_task() {
+    if (coro_) {
+      coro_.destroy();
+      coro_ = {};
+    }
+  }
+
+  void start() && {
+    auto tmp = coro_;
+    coro_ = {};
+    tmp.resume();
+  }
+
+  std::coroutine_handle<promise_type> coro_;
+};
+
+class SelfResumeAwaiter final
+{
+public:
+    bool await_ready() noexcept { return false; }
+    std::coroutine_handle<> await_suspend(std::coroutine_handle<> h) { 
+      return h; 
+    }
+    void await_resume() noexcept {}
+};
+
+// Check that there is only one call left: coroutine destroy
+// CHECK-LABEL: define{{.*}}void @_Z3foov.resume
+// CHECK-NOT: call{{.*}}
+// CHECK: tail call{{.*}}void %{{[0-9+]}}(
+// CHECK-NOT: call{{.*}}
+// CHECK: define
+detached_task foo() {
+  co_await SelfResumeAwaiter{};
+  co_return;
+}
diff --git a/clang/test/CodeGenCoroutines/coro-symmetric-transfer-01.cpp b/clang/test/CodeGenCoroutines/coro-symmetric-transfer-01.cpp
index c0b9e9ee2c55818..f1870b8eceff99a 100644
--- a/clang/test/CodeGenCoroutines/coro-symmetric-transfer-01.cpp
+++ b/clang/test/CodeGenCoroutines/coro-symmetric-transfer-01.cpp
@@ -47,13 +47,9 @@ detached_task foo() {
   co_return;
 }
 
+// FIXME: is this test needed now?
 // check that the lifetime of the coroutine handle used to obtain the address is contained within single basic block, and hence does not live across suspension points.
 // CHECK-LABEL: final.suspend:
 // CHECK:         %{{.+}} = call token @llvm.coro.save(ptr null)
-// CHECK:         call void @llvm.lifetime.start.p0(i64 8, ptr %[[HDL:.+]])
-// CHECK:         %[[CALL:.+]] = call ptr @_ZN13detached_task12promise_type13final_awaiter13await_suspendESt16coroutine_handleIS0_E(
-// CHECK:         %[[HDL_CAST2:.+]] = getelementptr inbounds %"struct.std::coroutine_handle.0", ptr %[[HDL]], i32 0, i32 0
-// CHECK:         store ptr %[[CALL]], ptr %[[HDL_CAST2]], align 8
-// CHECK:         %[[HDL_TRANSFER:.+]] = call noundef ptr @_ZNKSt16coroutine_handleIvE7addressEv(ptr noundef {{.*}}%[[HDL]])
-// CHECK:         call void @llvm.lifetime.end.p0(i64 8, ptr %[[HDL]])
+// CHECK:         %[[HDL_TRANSFER:.+]] = call ptr @llvm.coro.await.suspend.handle
 // CHECK:         call void @llvm.coro.resume(ptr %[[HDL_TRANSFER]])
diff --git a/clang/test/CodeGenCoroutines/coro-symmetric-transfer-02.cpp b/clang/test/CodeGenCoroutines/coro-symmetric-transfer-02.cpp
index 890d55e41de9538..ca6cf74115a3b1e 100644
--- a/clang/test/CodeGenCoroutines/coro-symmetric-transfer-02.cpp
+++ b/clang/test/CodeGenCoroutines/coro-symmetric-transfer-02.cpp
@@ -89,10 +89,8 @@ Task bar() {
 // CHECK:         br i1 %{{.+}}, label %[[CASE1_AWAIT_READY:.+]], label %[[CASE1_AWAIT_SUSPEND:.+]]
 // CHECK:       [[CASE1_AWAIT_SUSPEND]]:
 // CHECK-NEXT:    %{{.+}} = call token @llvm.coro.save(ptr null)
-// CHECK-NEXT:    call void @llvm.lifetime.start.p0(i64 8, ptr %[[TMP1:.+]])
-
-// CHECK:    call void @llvm.lifetime.end.p0(i64 8, ptr %[[TMP1]])
-// CHECK-NEXT:    call void @llvm.coro.resume
+// CHECK-NEXT:    %[[HANDLE1_PTR:.+]] = call ptr @llvm.coro.await.suspend.handle
+// CHECK-NEXT:    call void @llvm.coro.resume(ptr %[[HANDLE1_PTR]])
 // CHECK-NEXT:    %{{.+}} = call i8 @llvm.coro.suspend
 // CHECK-NEXT:    switch i8 %{{.+}}, label %coro.ret [
 // CHECK-NEXT:      i8 0, label %[[CASE1_AWAIT_READY]]
@@ -106,10 +104,8 @@ Task bar() {
 // CHECK:         br i1 %{{.+}}, label %[[CASE2_AWAIT_READY:.+]], label %[[CASE2_AWAIT_SUSPEND:.+]]
 // CHECK:       [[CASE2_AWAIT_SUSPEND]]:
 // CHECK-NEXT:    %{{.+}} = call token @llvm.coro.save(ptr null)
-// CHECK-NEXT:    call void @llvm.lifetime.start.p0(i64 8, ptr %[[TMP2:.+]])
-
-// CHECK:    call void @llvm.lifetime.end.p0(i64 8, ptr %[[TMP2]])
-// CHECK-NEXT:    call void @llvm.coro.resume
+// CHECK-NEXT:    %[[HANDLE2_PTR:.+]] = call ptr @llvm.coro.await.suspend.handle
+// CHECK-NEXT:    call void @llvm.coro.resume(ptr %[[HANDLE2_PTR]])
 // CHECK-NEXT:    %{{.+}} = call i8 @llvm.coro.suspend
 // CHECK-NEXT:    switch i8 %{{.+}}, label %coro.ret [
 // CHECK-NEXT:      i8 0, label %[[CASE2_AWAIT_READY]]
diff --git a/clang/test/CodeGenCoroutines/pr59181.cpp b/clang/test/CodeGenCoroutines/pr59181.cpp
index 80f4634db252146..ac732b6232e76e9 100644
--- a/clang/test/CodeGenCoroutines/pr59181.cpp
+++ b/clang/test/CodeGenCoroutines/pr59181.cpp
@@ -48,13 +48,14 @@ void foo() {
   bar(true);
 }
 
+// FIXME: the test doesn't seem to be relevant anymore,
+// because objects that require cleanup are no more emitted in the suspend block
 // CHECK: cleanup.cont:{{.*}}
 // CHECK-NEXT: load i8
 // CHECK-NEXT: trunc
 // CHECK-NEXT: store i1 false
-// CHECK-NEXT: call void @llvm.lifetime.start.p0(i64 8, ptr [[REF:%ref.tmp[0-9]+]])
 
 // CHECK: await.suspend:{{.*}}
-// CHECK-NOT: call void @llvm.lifetime.start.p0(i64 8, ptr [[REF]])
-// CHECK: call void @_ZZN4Task12promise_type15await_transformES_EN10Suspension13await_suspendESt16coroutine_handleIvE
-// CHECK-NEXT: call void @llvm.lifetime.end.p0(i64 8, ptr [[REF]])
+// CHECK-NOT: call void @llvm.lifetime
+// CHECK: call void @llvm.coro.await.suspend(
+// CHECK-NEXT: %{{[0-9]+}} = call i8 @llvm.coro.suspend(
diff --git a/clang/test/SemaCXX/co_await-ast.cpp b/clang/test/SemaCXX/co_await-ast.cpp
index 10cee21da0e87cd..ed3a988653eb29e 100644
--- a/clang/test/SemaCXX/co_await-ast.cpp
+++ b/clang/test/SemaCXX/co_await-ast.cpp
@@ -81,9 +81,10 @@ awaitable foo() {
 // CHECK:   |     |         `-CallExpr {{.*}} 'coroutine_handle<awaitable_frame>':'std::coroutine_handle<awaitable_frame>'
 // CHECK:   |     |           |-ImplicitCastExpr {{.*}} 'coroutine_handle<awaitable_frame> (*)(void *) noexcept' <FunctionToPointerDecay>
 // CHECK:   |     |           | `-DeclRefExpr {{.*}} 'coroutine_handle<awaitable_frame> (void *) noexcept' lvalue CXXMethod {{.*}} 'from_address' 'coroutine_handle<awaitable_frame> (void *) noexcept'
-// CHECK:   |     |           `-CallExpr {{.*}} 'void *'
-// CHECK:   |     |             `-ImplicitCastExpr {{.*}} 'void *(*)() noexcept' <FunctionToPointerDecay>
-// CHECK:   |     |               `-DeclRefExpr {{.*}} 'void *() noexcept' lvalue Function {{.*}} '__builtin_coro_frame' 'void *() noexcept'
+// CHECK:   |     |           `-OpaqueValueExpr {{.*}} 'void *'
+// CHECK:   |     |             `-CallExpr {{.*}} 'void *'
+// CHECK:   |     |               `-ImplicitCastExpr {{.*}} 'void *(*)() noexcept' <FunctionToPointerDecay>
+// CHECK:   |     |                 `-DeclRefExpr {{.*}} 'void *() noexcept' lvalue Function {{.*}} '__builtin_coro_frame' 'void *() noexcept'
 // CHECK:   |     `-CXXMemberCallExpr {{.*}} 'void'
 // CHECK:   |       `-MemberExpr {{.*}} '<bound member function type>' .await_resume {{.*}}
 // CHECK:   |         `-ImplicitCastExpr {{.*}} 'const awaitable_frame::result_t' lvalue <NoOp>
diff --git a/llvm/include/llvm/IR/Intrinsics.td b/llvm/include/llvm/IR/Intrinsics.td
index c2c0f74c315bab6..c8270a24c55d17b 100644
--- a/llvm/include/llvm/IR/Intrinsics.td
+++ b/llvm/include/llvm/IR/Intrinsics.td
@@ -1692,6 +1692,18 @@ def int_coro_promise : Intrinsic<[llvm_ptr_ty],
                                  [llvm_ptr_ty, llvm_i32_ty, llvm_i1_ty],
                                  [IntrNoMem, NoCapture<ArgIndex<0>>]>;
 
+def int_coro_await_suspend : Intrinsic<[],
+                                       [llvm_ptr_ty, llvm_ptr_ty, llvm_ptr_ty],
+                                       [Throws]>;
+
+def int_coro_await_suspend_bool : Intrinsic<[llvm_i1_ty],
+                                            [llvm_ptr_ty, llvm_ptr_ty, llvm_ptr_ty],
+                                            [Throws]>;
+
+def int_coro_await_suspend_handle : Intrinsic<[llvm_ptr_ty],
+                                              [llvm_ptr_ty, llvm_ptr_ty, llvm_ptr_ty],
+                                              [Throws]>;
+
 // Coroutine Lowering Intrinsics. Used internally by coroutine passes.
 
 def int_coro_subfn_addr : DefaultAttrsIntrinsic<
diff --git a/llvm/lib/IR/Verifier.cpp b/llvm/lib/IR/Verifier.cpp
index 027c6ccd1ab3045..59ee8c802b70034 100644
--- a/llvm/lib/IR/Verifier.cpp
+++ b/llvm/lib/IR/Verifier.cpp
@@ -4994,6 +4994,9 @@ void Verifier::visitInstruction(Instruction &I) {
                 F->getIntrinsicID() == Intrinsic::seh_scope_end ||
                 F->getIntrinsicID() == Intrinsic::coro_resume ||
                 F->getIntrinsicID() == Intrinsic::coro_destroy ||
+                F->getIntrinsicID() == Intrinsic::coro_await_suspend ||
+                F->getIntrinsicID() == Intrinsic::coro_await_suspend_bool ||
+                F->getIntrinsicID() == Intrinsic::coro_await_suspend_handle ||
                 F->getIntrinsicID() ==
                     Intrinsic::experimental_patchpoint_void ||
                 F->getIntrinsicID() == Intrinsic::experimental_patchpoint_i64 ||
diff --git a/llvm/lib/Transforms/Coroutines/CoroInstr.h b/llvm/lib/Transforms/Coroutines/CoroInstr.h
index f01aa58eb899961..ce4fe5a4e4de4e6 100644
--- a/llvm/lib/Transforms/Coroutines/CoroInstr.h
+++ b/llvm/lib/Transforms/Coroutines/CoroInstr.h
@@ -78,6 +78,36 @@ class LLVM_LIBRARY_VISIBILITY CoroAllocInst : public IntrinsicInst {
   }
 };
 
+/// This represents the llvm.coro.await.suspend instruction.
+class LLVM_LIBRARY_VISIBILITY CoroAwaitSuspendInst : public CallBase {
+  enum { AwaiterArg, FrameArg, HelperArg };
+
+public:
+  Value *getAwaiter() const { return getArgOperand(AwaiterArg); }
+
+  Value *getFrame() const { return getArgOperand(FrameArg); }
+
+  Function *getHelperFunction() const {
+    return cast<Function>(getArgOperand(HelperArg));
+  }
+
+  // Methods to support type inquiry through isa, cast, and dyn_cast:
+  static bool classof(const CallBase *CB) {
+    if (const Function *CF = CB->getCalledFunction()) {
+      auto IID = CF->getIntrinsicID();
+      return IID == Intrinsic::coro_await_suspend ||
+             IID == Intrinsic::coro_await_suspend_bool ||
+             IID == Intrinsic::coro_await_suspend_handle;
+    }
+
+    return false;
+  }
+
+  static bool classof(const Value *V) {
+    return isa<CallBase>(V) && classof(cast<CallBase>(V));
+  }
+};
+
 /// This represents a common base class for llvm.coro.id instructions.
 class LLVM_LIBRARY_VISIBILITY AnyCoroIdInst : public IntrinsicInst {
 public:
diff --git a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
index 90d40242ff2e444..00aeea2df2db039 100644
--- a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
@@ -79,6 +79,73 @@ using namespace llvm;
 
 namespace {
 
+// Created on demand if the coro-early pass has work to do.
+class Lowerer : public coro::LowererBase {
+  IRBuilder<> Builder;
+  void lowerAwaitSuspend(CoroAwaitSuspendInst *CB);
+
+public:
+  Lowerer(Module &M) : LowererBase(M), Builder(Context) {}
+
+  void lowerAwaitSuspends(Function &F);
+};
+
+void Lowerer::lowerAwaitSuspend(CoroAwaitSuspendInst *CB) {
+  auto Helper = CB->getHelperFunction();
+  auto Awaiter = CB->getAwaiter();
+  auto FramePtr = CB->getFrame();
+
+  Builder.SetInsertPoint(CB);
+
+  CallBase *NewCall = nullptr;
+  if (auto Invoke = dyn_cast<InvokeInst>(CB)) {
+    auto HelperInvoke =
+        Builder.CreateInvoke(Helper, Invoke->getNormalDest(),
+                             Invoke->getUnwindDest(), {Awaiter, FramePtr});
+
+    HelperInvoke->setCallingConv(Invoke->getCallingConv());
+    std::copy(Invoke->bundle_op_info_begin(), Invoke->bundle_op_info_end(),
+              HelperInvoke->bundle_op_info_begin());
+    AttributeList NewAttributes =
+        Invoke->getAttributes().removeParamAttributes(Context, 2);
+    HelperInvoke->setAttributes(NewAttributes);
+    HelperInvoke->setDebugLoc(Invoke->getDebugLoc());
+    NewCall = HelperInvoke;
+  } else if (auto Call = dyn_cast<CallInst>(CB)) {
+    auto HelperCall = Builder.CreateCall(Helper, {Awaiter, FramePtr});
+
+    AttributeList NewAttributes =
+        Call->getAttributes().removeParamAttributes(Context, 2);
+    HelperCall->setAttributes(NewAttributes);
+    HelperCall->setDebugLoc(Call->getDebugLoc());
+    NewCall = HelperCall;
+  }
+
+  CB->replaceAllUsesWith(NewCall);
+  CB->eraseFromParent();
+
+  InlineFunctionInfo FnInfo;
+  auto InlineRes = InlineFunction(*NewCall, FnInfo);
+  assert(InlineRes.isSuccess() && "Expected inlining to succeed");
+  (void)InlineRes;
+}
+
+void Lowerer::lowerAwaitSuspends(Function &F) {
+  SmallVector<CoroAwaitSuspendInst *, 4> AwaitSuspends;
+
+  for (Instruction &I : llvm::make_early_inc_range(instructions(F))) {
+    auto *CB = dyn_cast<CallBase>(&I);
+    if (!CB)
+      continue;
+
+    if (auto *AWS = dyn_cast<CoroAwaitSuspendInst>(CB))
+      AwaitSuspends.push_back(AWS);
+  }
+
+  for (auto *AWS : AwaitSuspends)
+    lowerAwaitSuspend(AWS);
+}
+
 /// A little helper class for building
 class CoroCloner {
 public:
@@ -1319,6 +1386,11 @@ static void handleNoSuspendCoroutine(coro::Shape &Shape) {
 // the coroutine and if that is the case we cannot eliminate the suspend point.
 static bool hasCallsInBlockBetween(Instruction *From, Instruction *To) {
   for (Instruction *I = From; I != To; I = I->getNextNode()) {
+    // This one could resume the coroutine,
+    // but additional analysis before the check should ensure,
+    // that it can't happen
+    if (isa<CoroAwaitSuspendInst>(I))
+      continue;
     // Assume that no intrinsic can resume the coroutine.
     if (isa<IntrinsicInst>(I))
       continue;
@@ -1359,6 +1431,9 @@ static bool hasCallsInBlocksBetween(BasicBlock *SaveBB, BasicBlock *ResDesBB) {
 }
 
 static bool hasCallsBetween(Instruction *Save, Instruction *ResumeOrDestroy) {
+  if (Save == ResumeOrDestroy)
+    return false;
+
   auto *SaveBB = Save->getParent();
   auto *ResumeOrDestroyBB = ResumeOrDestroy->getParent();
 
@@ -1381,6 +1456,36 @@ static bool hasCallsBetween(Instruction *Save, Instruction *ResumeOrDestroy) {
   return false;
 }
 
+// Check if await-suspend helper is "simple".
+// The conditions are:
+// 1. The return result is exactly coroutine frame parameter, passed to helper
+// 2. There are no calls between any of the returns and helper entry that could
+// resume or destroy it
+// FIXME: perform more sophisiticated analysis?
+static bool isSimpleHelper(CoroAwaitSuspendInst *AWS) {
+  auto Helper = AWS->getHelperFunction();
+
+  SmallVector<ReturnInst *, 4> Rets;
+
+  for (auto &BB : *Helper) {
+    if (BB.empty())
+      continue;
+    auto terminator = BB.getTerminator();
+    if (!terminator)
+      continue;
+    if (auto Ret = dyn_cast<ReturnInst>(terminator))
+      Rets.push_back(cast<ReturnInst>(terminator));
+  }
+
+  // FIXME: get rid of magical constant
+  for (auto Ret : Rets)
+    if (Ret->getReturnValue() != Helper->getArg(1) ||
+        hasCallsBetween(Helper->getEntryBlock().getFirstNonPHI(), Ret))
+      return false;
+
+  return true;
+}
+
 // If a SuspendIntrin is preceded by Resume or Destroy, we can eliminate the
 // suspend point and replace it with nornal control flow.
 static bool simplifySuspendPoint(CoroSuspendInst *Suspend,
@@ -1404,9 +1509,18 @@ static bool simplifySuspendPoint(CoroSuspendInst *Suspend,
   if (!SubFn)
     return false;
 
-  // Does not refer to the current coroutine, we cannot do anything with it.
-  if (SubFn->getFrame() != CoroBegin)
-    return false;
+  auto Frame = SubFn->getFrame();
+
+  // Check that frame directly always refers to the current coroutine,
+  // either directly or via helper
+  if (Frame != CoroBegin) {
+    auto *AWS = dyn_cast<CoroAwaitSuspendInst>(Frame);
+    if (!AWS)
+      return false;
+
+    if (AWS->getFrame() != CoroBegin || !isSimpleHelper(AWS))
+      return false;
+  }
 
   // See if the transformation is safe. Specifically, see if there are any
   // calls in between Save and CallInstr. They can potenitally resume the
@@ -1487,12 +1601,16 @@ static void simplifySuspendPoints(coro::Shape &Shape) {
 namespace {
 
 struct SwitchCoroutineSplitter {
-  static void split(Function &F, coro::Shape &Shape,
+  static void split(Module &M, Function &F, coro::Shape &Shape,
                     SmallVectorImpl<Function *> &Clones,
                     TargetTransformInfo &TTI) {
     assert(Shape.ABI == coro::ABI::Switch);
 
     createResumeEntryBlock(F, Shape);
+
+    Lowerer lowerer(M);
+    lowerer.lowerAwaitSuspends(F);
+
     auto *ResumeClone =
         createClone(F, ".resume", Shape, CoroCloner::Kind::SwitchResume);
     auto *DestroyClone =
@@ -2012,7 +2130,7 @@ class PrettyStackTraceFunction : public PrettyStackTraceEntry {
 } // namespace
 
 static coro::Shape
-splitCoroutine(Function &F, SmallVectorImpl<Function *> &Clones,
+splitCoroutine(Module &M, Function &F, SmallVectorImpl<Function *> &Clones,
                TargetTransformInfo &TTI, bool OptimizeFrame,
                std::function<bool(Instruction &)> MaterializableCallback) {
   PrettyStackTraceFunction prettyStackTrace(F);
@@ -2036,7 +2154,7 @@ splitCoroutine(Function &F, SmallVectorImpl<Function *> &Clones,
   } else {
     switch (Shape.ABI) {
     case coro::ABI::Switch:
-      SwitchCoroutineSplitter::split(F, Shape, Clones, TTI);
+      SwitchCoroutineSplitter::split(M, F, Shape, Clones, TTI);
       break;
     case coro::ABI::Async:
       splitAsyncCoroutine(F, Shape, Clones, TTI);
@@ -2221,7 +2339,7 @@ PreservedAnalyses CoroSplitPass::run(LazyCallGraph::SCC &C,
     SmallVector<Function *, 4> Clones;
     auto &ORE = FAM.getResult<OptimizationRemarkEmitterAnalysis>(F);
     const coro::Shape Shape =
-        splitCoroutine(F, Clones, FAM.getResult<TargetIRAnalysis>(F),
+        splitCoroutine(M, F, Clones, FAM.getResult<TargetIRAnalysis>(F),
                        OptimizeFrame, MaterializableCallback);
     updateCallGraphAfterCoroutineSplit(*N, Shape, Clones, C, CG, AM, UR, FAM);
 
diff --git a/llvm/lib/Transforms/Coroutines/Coroutines.cpp b/llvm/lib/Transforms/Coroutines/Coroutines.cpp
index eef5543bae24ab9..8ecdf019a2c84a4 100644
--- a/llvm/lib/Transforms/Coroutines/Coroutines.cpp
+++ b/llvm/lib/Transforms/Coroutines/Coroutines.cpp
@@ -67,6 +67,9 @@ static const char *const CoroIntrinsics[] = {
     "llvm.coro.async.resume",
     "llvm.coro.async.size.replace",
     "llvm.coro.async.store_resume",
+    "llvm.coro.await.suspend",
+    "llvm.coro.await.suspend.bool",
+    "llvm.coro.await.suspend.handle",
     "llvm.coro.begin",
     "llvm.coro.destroy",
     "llvm.coro.done",

>From c901979261d4280d3d398ce70f878fb901947712 Mon Sep 17 00:00:00 2001
From: fpasserby <125797601+fpasserby at users.noreply.github.com>
Date: Mon, 29 Jan 2024 22:17:02 +0100
Subject: [PATCH 02/17] Remove isAwaitSuspendNoThrow

---
 clang/include/clang/AST/ExprCXX.h   | 21 ++++-----
 clang/lib/CodeGen/CGCoroutine.cpp   | 27 +++++++-----
 clang/lib/CodeGen/CodeGenFunction.h |  3 +-
 clang/lib/Sema/SemaCoroutine.cpp    | 67 ++---------------------------
 4 files changed, 31 insertions(+), 87 deletions(-)

diff --git a/clang/include/clang/AST/ExprCXX.h b/clang/include/clang/AST/ExprCXX.h
index 57a505036def409..42cabe891ac0fb2 100644
--- a/clang/include/clang/AST/ExprCXX.h
+++ b/clang/include/clang/AST/ExprCXX.h
@@ -5035,19 +5035,18 @@ class CoroutineSuspendExpr : public Expr {
   enum SubExpr { Operand, Common, Ready, Suspend, Resume, Count };
 
   Stmt *SubExprs[SubExpr::Count];
-  bool IsSuspendNoThrow = false;
   OpaqueValueExpr *OpaqueValue = nullptr;
   OpaqueValueExpr *OpaqueFramePtr = nullptr;
 
 public:
   CoroutineSuspendExpr(StmtClass SC, SourceLocation KeywordLoc, Expr *Operand,
                        Expr *Common, Expr *Ready, Expr *Suspend, Expr *Resume,
-                       bool IsSuspendNoThrow, OpaqueValueExpr *OpaqueValue,
+                       OpaqueValueExpr *OpaqueValue,
                        OpaqueValueExpr *OpaqueFramePtr)
       : Expr(SC, Resume->getType(), Resume->getValueKind(),
              Resume->getObjectKind()),
-        KeywordLoc(KeywordLoc), IsSuspendNoThrow(IsSuspendNoThrow),
-        OpaqueValue(OpaqueValue), OpaqueFramePtr(OpaqueFramePtr) {
+        KeywordLoc(KeywordLoc), OpaqueValue(OpaqueValue),
+        OpaqueFramePtr(OpaqueFramePtr) {
     SubExprs[SubExpr::Operand] = Operand;
     SubExprs[SubExpr::Common] = Common;
     SubExprs[SubExpr::Ready] = Ready;
@@ -5104,8 +5103,6 @@ class CoroutineSuspendExpr : public Expr {
     return static_cast<Expr *>(SubExprs[SubExpr::Operand]);
   }
 
-  bool isSuspendNoThrow() const { return IsSuspendNoThrow; }
-
   SourceLocation getKeywordLoc() const { return KeywordLoc; }
 
   SourceLocation getBeginLoc() const LLVM_READONLY { return KeywordLoc; }
@@ -5134,12 +5131,12 @@ class CoawaitExpr : public CoroutineSuspendExpr {
 
 public:
   CoawaitExpr(SourceLocation CoawaitLoc, Expr *Operand, Expr *Common,
-              Expr *Ready, Expr *Suspend, Expr *Resume, bool IsSuspendNoThrow,
+              Expr *Ready, Expr *Suspend, Expr *Resume,
               OpaqueValueExpr *OpaqueValue, OpaqueValueExpr *OpaqueFramePtr,
               bool IsImplicit = false)
       : CoroutineSuspendExpr(CoawaitExprClass, CoawaitLoc, Operand, Common,
-                             Ready, Suspend, Resume, IsSuspendNoThrow,
-                             OpaqueValue, OpaqueFramePtr) {
+                             Ready, Suspend, Resume, OpaqueValue,
+                             OpaqueFramePtr) {
     CoawaitBits.IsImplicit = IsImplicit;
   }
 
@@ -5217,11 +5214,11 @@ class CoyieldExpr : public CoroutineSuspendExpr {
 
 public:
   CoyieldExpr(SourceLocation CoyieldLoc, Expr *Operand, Expr *Common,
-              Expr *Ready, Expr *Suspend, Expr *Resume, bool IsSuspendNoThrow,
+              Expr *Ready, Expr *Suspend, Expr *Resume,
               OpaqueValueExpr *OpaqueValue, OpaqueValueExpr *OpaqueFramePtr)
       : CoroutineSuspendExpr(CoyieldExprClass, CoyieldLoc, Operand, Common,
-                             Ready, Suspend, Resume, IsSuspendNoThrow,
-                             OpaqueValue, OpaqueFramePtr) {}
+                             Ready, Suspend, Resume, OpaqueValue,
+                             OpaqueFramePtr) {}
   CoyieldExpr(SourceLocation CoyieldLoc, QualType Ty, Expr *Operand,
               Expr *Common)
       : CoroutineSuspendExpr(CoyieldExprClass, CoyieldLoc, Ty, Operand,
diff --git a/clang/lib/CodeGen/CGCoroutine.cpp b/clang/lib/CodeGen/CGCoroutine.cpp
index 233a0137c751bf4..5bb4e7e950a74de 100644
--- a/clang/lib/CodeGen/CGCoroutine.cpp
+++ b/clang/lib/CodeGen/CGCoroutine.cpp
@@ -173,6 +173,10 @@ static bool ResumeStmtCanThrow(const Stmt *S) {
   return false;
 }
 
+static bool AwaitSuspendStmtCanThrow(const Stmt *S) {
+  return ResumeStmtCanThrow(S);
+}
+
 // Emit suspend expression which roughly looks like:
 //
 //   auto && x = CommonExpr();
@@ -233,8 +237,11 @@ static LValueOrRValue emitSuspendExpression(CodeGenFunction &CGF, CGCoroData &Co
   auto *NullPtr = llvm::ConstantPointerNull::get(CGF.CGM.Int8PtrTy);
   auto *SaveCall = Builder.CreateCall(CoroSave, {NullPtr});
 
+  const auto AwaitSuspendCanThrow =
+      AwaitSuspendStmtCanThrow(S.getSuspendExpr());
+
   auto SuspendHelper = CodeGenFunction(CGF.CGM).generateAwaitSuspendHelper(
-      CGF.CurFn->getName(), Prefix, S);
+      CGF.CurFn->getName(), Prefix, S, AwaitSuspendCanThrow);
 
   llvm::CallBase *SuspendRet = nullptr;
 
@@ -262,13 +269,12 @@ static LValueOrRValue emitSuspendExpression(CodeGenFunction &CGF, CGCoroData &Co
 
     llvm::Function *AwaitSuspendIntrinsic = CGF.CGM.getIntrinsic(IID);
     // FIXME: add call attributes?
-    if (S.isSuspendNoThrow()) {
-      SuspendRet = CGF.EmitNounwindRuntimeCall(AwaitSuspendIntrinsic,
-                                               SuspendHelperCallArgs);
-    } else {
+    if (AwaitSuspendCanThrow)
       SuspendRet =
           CGF.EmitCallOrInvoke(AwaitSuspendIntrinsic, SuspendHelperCallArgs);
-    }
+    else
+      SuspendRet = CGF.EmitNounwindRuntimeCall(AwaitSuspendIntrinsic,
+                                               SuspendHelperCallArgs);
 
     CGF.CurCoro.InSuspendBlock = false;
   }
@@ -380,10 +386,9 @@ static QualType getCoroutineSuspendExprReturnType(const ASTContext &Ctx,
 }
 #endif
 
-llvm::Function *
-CodeGenFunction::generateAwaitSuspendHelper(Twine const &CoroName,
-                                            Twine const &SuspendPointName,
-                                            CoroutineSuspendExpr const &S) {
+llvm::Function *CodeGenFunction::generateAwaitSuspendHelper(
+    Twine const &CoroName, Twine const &SuspendPointName,
+    CoroutineSuspendExpr const &S, bool CanThrow) {
   std::string FuncName = "__await_suspend_helper_";
   FuncName += CoroName.str();
   FuncName += '_';
@@ -424,7 +429,7 @@ CodeGenFunction::generateAwaitSuspendHelper(Twine const &CoroName,
   Fn->setMustProgress();
   Fn->addFnAttr(llvm::Attribute::AttrKind::AlwaysInline);
 
-  if (S.isSuspendNoThrow()) {
+  if (!CanThrow) {
     Fn->addFnAttr(llvm::Attribute::AttrKind::NoUnwind);
   }
 
diff --git a/clang/lib/CodeGen/CodeGenFunction.h b/clang/lib/CodeGen/CodeGenFunction.h
index ed6d8b318301ba0..d13be4e2eae69dc 100644
--- a/clang/lib/CodeGen/CodeGenFunction.h
+++ b/clang/lib/CodeGen/CodeGenFunction.h
@@ -354,7 +354,8 @@ class CodeGenFunction : public CodeGenTypeCache {
 
   llvm::Function *generateAwaitSuspendHelper(Twine const &CoroName,
                                              Twine const &SuspendPointName,
-                                             CoroutineSuspendExpr const &S);
+                                             CoroutineSuspendExpr const &S,
+                                             bool CanThrow);
 
   /// CurGD - The GlobalDecl for the current function being compiled.
   GlobalDecl CurGD;
diff --git a/clang/lib/Sema/SemaCoroutine.cpp b/clang/lib/Sema/SemaCoroutine.cpp
index a46744205d0f298..39004d711a02db0 100644
--- a/clang/lib/Sema/SemaCoroutine.cpp
+++ b/clang/lib/Sema/SemaCoroutine.cpp
@@ -31,8 +31,6 @@
 using namespace clang;
 using namespace sema;
 
-static bool isNoThrow(Sema &S, const Stmt *E);
-
 static LookupResult lookupMember(Sema &S, const char *Name, CXXRecordDecl *RD,
                                  SourceLocation Loc, bool &Res) {
   DeclarationName DN = S.PP.getIdentifierInfo(Name);
@@ -493,10 +491,6 @@ static ReadySuspendResumeResult buildCoawaitCalls(Sema &S, VarDecl *CoroPromise,
     }
   }
 
-  if (Calls.Results[ACT::ACT_Suspend]) {
-    Calls.IsSuspendNoThrow = isNoThrow(S, Calls.Results[ACT::ACT_Suspend]);
-  }
-
   BuildSubExpr(ACT::ACT_Resume, "await_resume", std::nullopt);
 
   // Make sure the awaiter object gets a chance to be cleaned up.
@@ -649,59 +643,6 @@ static FunctionScopeInfo *checkCoroutineContext(Sema &S, SourceLocation Loc,
   return ScopeInfo;
 }
 
-/// Recursively check \p E and all its children to see if any call target
-/// (including constructor call) is declared noexcept. Also any value returned
-/// from the call has a noexcept destructor.
-static bool isNoThrow(Sema &S, const Stmt *E) {
-  auto isDeclNoexcept = [&](const Decl *D, bool IsDtor = false) -> bool {
-    // In the case of dtor, the call to dtor is implicit and hence we should
-    // pass nullptr to canCalleeThrow.
-    if (Sema::canCalleeThrow(S, IsDtor ? nullptr : cast<Expr>(E), D)) {
-      return false;
-    }
-    return true;
-  };
-
-  if (auto *CE = dyn_cast<CXXConstructExpr>(E)) {
-    CXXConstructorDecl *Ctor = CE->getConstructor();
-    if (!isDeclNoexcept(Ctor)) {
-      return false;
-    }
-    // Check the corresponding destructor of the constructor.
-    if (!isDeclNoexcept(Ctor->getParent()->getDestructor(), /*IsDtor=*/true)) {
-      return false;
-    }
-  } else if (auto *CE = dyn_cast<CallExpr>(E)) {
-    if (CE->isTypeDependent())
-      return false;
-
-    if (!isDeclNoexcept(CE->getCalleeDecl())) {
-      return false;
-    }
-
-    QualType ReturnType = CE->getCallReturnType(S.getASTContext());
-    // Check the destructor of the call return type, if any.
-    if (ReturnType.isDestructedType() ==
-        QualType::DestructionKind::DK_cxx_destructor) {
-      const auto *T =
-          cast<RecordType>(ReturnType.getCanonicalType().getTypePtr());
-      if (!isDeclNoexcept(cast<CXXRecordDecl>(T->getDecl())->getDestructor(),
-                          /*IsDtor=*/true)) {
-        return false;
-      }
-    }
-  }
-  for (const auto *Child : E->children()) {
-    if (!Child)
-      continue;
-    if (!isNoThrow(S, Child)) {
-      return false;
-    }
-  }
-
-  return true;
-}
-
 /// Recursively check \p E and all its children to see if any call target
 /// (including constructor call) is declared noexcept. Also any value returned
 /// from the call has a noexcept destructor.
@@ -1002,7 +943,7 @@ ExprResult Sema::BuildResolvedCoawaitExpr(SourceLocation Loc, Expr *Operand,
 
   Expr *Res = new (Context) CoawaitExpr(
       Loc, Operand, Awaiter, RSS.Results[0], RSS.Results[1], RSS.Results[2],
-      RSS.IsSuspendNoThrow, RSS.OpaqueValue, RSS.OpaqueFramePtr, IsImplicit);
+      RSS.OpaqueValue, RSS.OpaqueFramePtr, IsImplicit);
 
   return Res;
 }
@@ -1058,9 +999,9 @@ ExprResult Sema::BuildCoyieldExpr(SourceLocation Loc, Expr *E) {
   if (RSS.IsInvalid)
     return ExprError();
 
-  Expr *Res = new (Context) CoyieldExpr(
-      Loc, Operand, E, RSS.Results[0], RSS.Results[1], RSS.Results[2],
-      RSS.IsSuspendNoThrow, RSS.OpaqueValue, RSS.OpaqueFramePtr);
+  Expr *Res = new (Context)
+      CoyieldExpr(Loc, Operand, E, RSS.Results[0], RSS.Results[1],
+                  RSS.Results[2], RSS.OpaqueValue, RSS.OpaqueFramePtr);
 
   return Res;
 }

>From b01bd4409a88ae10ce7986e8af069843d9f90ab4 Mon Sep 17 00:00:00 2001
From: fpasserby <125797601+fpasserby at users.noreply.github.com>
Date: Mon, 29 Jan 2024 22:18:24 +0100
Subject: [PATCH 03/17] Update AST serializers

---
 clang/lib/Serialization/ASTReaderStmt.cpp | 2 ++
 clang/lib/Serialization/ASTWriterStmt.cpp | 1 +
 2 files changed, 3 insertions(+)

diff --git a/clang/lib/Serialization/ASTReaderStmt.cpp b/clang/lib/Serialization/ASTReaderStmt.cpp
index 3da44ffccc38a28..2b6d113779559ae 100644
--- a/clang/lib/Serialization/ASTReaderStmt.cpp
+++ b/clang/lib/Serialization/ASTReaderStmt.cpp
@@ -484,6 +484,7 @@ void ASTStmtReader::VisitCoawaitExpr(CoawaitExpr *E) {
   for (auto &SubExpr: E->SubExprs)
     SubExpr = Record.readSubStmt();
   E->OpaqueValue = cast_or_null<OpaqueValueExpr>(Record.readSubStmt());
+  E->OpaqueFramePtr = cast_or_null<OpaqueValueExpr>(Record.readSubStmt());
   E->setIsImplicit(Record.readInt() != 0);
 }
 
@@ -493,6 +494,7 @@ void ASTStmtReader::VisitCoyieldExpr(CoyieldExpr *E) {
   for (auto &SubExpr: E->SubExprs)
     SubExpr = Record.readSubStmt();
   E->OpaqueValue = cast_or_null<OpaqueValueExpr>(Record.readSubStmt());
+  E->OpaqueFramePtr = cast_or_null<OpaqueValueExpr>(Record.readSubStmt());
 }
 
 void ASTStmtReader::VisitDependentCoawaitExpr(DependentCoawaitExpr *E) {
diff --git a/clang/lib/Serialization/ASTWriterStmt.cpp b/clang/lib/Serialization/ASTWriterStmt.cpp
index 484621ae813093e..232f0da4e0dd2e0 100644
--- a/clang/lib/Serialization/ASTWriterStmt.cpp
+++ b/clang/lib/Serialization/ASTWriterStmt.cpp
@@ -445,6 +445,7 @@ void ASTStmtWriter::VisitCoroutineSuspendExpr(CoroutineSuspendExpr *E) {
   for (Stmt *S : E->children())
     Record.AddStmt(S);
   Record.AddStmt(E->getOpaqueValue());
+  Record.AddStmt(E->getOpaqueFramePtr());
 }
 
 void ASTStmtWriter::VisitCoawaitExpr(CoawaitExpr *E) {

>From 710ca4eb7b7df77c92143fc4fdbc5faf7752a547 Mon Sep 17 00:00:00 2001
From: fpasserby <125797601+fpasserby at users.noreply.github.com>
Date: Tue, 30 Jan 2024 01:02:29 +0100
Subject: [PATCH 04/17] Document added intrinsics

---
 llvm/docs/Coroutines.rst | 271 ++++++++++++++++++++++++++++++++++++++-
 1 file changed, 270 insertions(+), 1 deletion(-)

diff --git a/llvm/docs/Coroutines.rst b/llvm/docs/Coroutines.rst
index d6219d264e4a9fb..fe892bbde67f639 100644
--- a/llvm/docs/Coroutines.rst
+++ b/llvm/docs/Coroutines.rst
@@ -1744,6 +1744,273 @@ a call to ``llvm.coro.suspend.retcon`` after resuming abnormally.
 In a yield-once coroutine, it is undefined behavior if the coroutine
 executes a call to ``llvm.coro.suspend.retcon`` after resuming in any way.
 
+.. _coro.await.suspend:
+
+'llvm.coro.await.suspend' Intrinsic
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+::
+
+  declare void @llvm.coro.await.suspend(
+                ptr <awaiter>,
+                ptr <handle>,
+                ptr <await_suspend_function>)
+
+Overview:
+"""""""""
+
+The '``llvm.coro.await.suspend``' intrinsic hides C++ `await-suspend`
+block code from optimizations on presplit coroutine body 
+to avoid miscompilations. This version of intrinsic corresponds to 
+'``void awaiter.await_suspend(...)``' variant.
+
+Arguments:
+""""""""""
+
+The first argument is a pointer to `awaiter` object.
+
+The second argument is a pointer to the current coroutine's frame.
+
+The third argument is a pointer to the helper function encapsulating
+`await-suspend` logic. Its signature must be
+
+.. code-block:: llvm
+
+    declare void @await_suspend_function(ptr %awaiter, ptr %hdl)
+
+Semantics:
+""""""""""
+
+The intrinsic must be used between corresponding `coro.save`_ and 
+`coro.suspend`_ calls. It is lowered to an inlined 
+`await_suspend_function` call during `CoroSplit`_ pass.
+
+Example:
+""""""""
+
+.. code-block:: llvm
+
+  ; before lowering
+  await.suspend:
+    %save = call token @llvm.coro.save(ptr %hdl)
+    call void @llvm.coro.await.suspend(
+                ptr %awaiter,
+                ptr %hdl,
+                ptr @await_suspend_function)
+    %suspend = call i8 @llvm.coro.suspend(token %save, i1 false)
+    ...
+
+  ; after lowering
+  await.suspend:
+    %save = call token @llvm.coro.save(ptr %hdl)
+    ; the call to await_suspend_function is inlined
+    call void @await_suspend_function(
+                ptr %awaiter,
+                ptr %hdl)
+    %suspend = call i8 @llvm.coro.suspend(token %save, i1 false)   
+    ...
+
+  ; helper function example
+  define void @await_suspend_function(ptr %awaiter, ptr %hdl)
+    entry:
+      %hdl.tmp = alloca %"struct.std::coroutine_handle"
+      %hdl.result.tmp = alloca %"struct.std::coroutine_handle"
+      %hdl.promise.tmp = alloca %"struct.std::coroutine_handle.0"
+      %hdl.promise = call ptr @"std::corouine_handle<promise_type>::from_address"(ptr %hdl)
+      %hdl.promise.tmp.dive = getelementptr inbounds %"struct.std::coroutine_handle.0",
+        ptr %hdl.promise.tmp, i32 0, i32 0
+      %hdl.promise.tmp.dive2 = getelementptr inbounds %"struct.std::coroutine_handle",
+        ptr %hdl.promise.tmp.dive, i32 0, i32 0
+      store ptr %hdl.promise, ptr %hdl.promise.tmp.dive2
+      call void @llvm.memcpy.p0.p0.i64(ptr %hdl.tmp, ptr %hdl.promise.tmp, i64 8, i1 false)
+      %hdl.tmp.dive = getelementptr inbounds %"struct.std::coroutine_handle",
+        ptr %hdl.tmp, i32 0, i32 0
+      %hdl.arg = load ptr, ptr %hdl.tmp.dive
+      call void @"Awaiter::await_suspend"(ptr %awaiter, ptr %hdl.arg)
+      ret void
+
+.. _coro.await.suspend.bool:
+
+'llvm.coro.await.suspend.bool' Intrinsic
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+::
+
+  declare i1 @llvm.coro.await.suspend.bool(
+                ptr <awaiter>,
+                ptr <handle>,
+                ptr <await_suspend_function>)
+
+Overview:
+"""""""""
+
+The '``llvm.coro.await.suspend.bool``' intrinsic hides C++ `await-suspend`
+block code from optimizations on presplit coroutine body 
+to avoid miscompilations. This version of intrinsic corresponds to 
+'``bool awaiter.await_suspend(...)``' variant.
+
+Arguments:
+""""""""""
+
+The first argument is a pointer to `awaiter` object.
+
+The second argument is a pointer to the current coroutine's frame.
+
+The third argument is a pointer to the helper function encapsulating
+`await-suspend` logic. Its signature must be
+
+.. code-block:: llvm
+
+    declare i1 @await_suspend_function(ptr %awaiter, ptr %hdl)
+
+Semantics:
+""""""""""
+
+The intrinsic must be used between corresponding `coro.save`_ and 
+`coro.suspend`_ calls. It is lowered to an inlined 
+`await_suspend_function` call during `CoroSplit`_ pass.
+
+If `await_suspend_function` call returns `true`, the current coroutine is
+immediately resumed.
+
+Example:
+""""""""
+
+.. code-block:: llvm
+
+  ; before lowering
+  await.suspend:
+    %save = call token @llvm.coro.save(ptr %hdl)
+    %resume = call i1 @llvm.coro.await.suspend(
+                ptr %awaiter,
+                ptr %hdl,
+                ptr @await_suspend_function)
+    br i1 %resume, %await.suspend.bool, %await.ready
+  await.suspend.bool:
+    %suspend = call i8 @llvm.coro.suspend(token %save, i1 false)
+    ...
+  await.ready:
+    call void @"Awaiter::await_ready"(ptr %awaiter)
+    ...
+
+  ; after lowering
+  await.suspend:
+    %save = call token @llvm.coro.save(ptr %hdl)
+    ; the call to await_suspend_function is inlined
+    %resume = call i1 @await_suspend_function(
+                ptr %awaiter,
+                ptr %hdl)
+    br i1 %resume, %await.suspend.bool, %await.ready
+    ...
+
+  ; helper function example
+  define i1 @await_suspend_function(ptr %awaiter, ptr %hdl)
+    entry:
+      %hdl.tmp = alloca %"struct.std::coroutine_handle"
+      %hdl.result.tmp = alloca %"struct.std::coroutine_handle"
+      %hdl.promise.tmp = alloca %"struct.std::coroutine_handle.0"
+      %hdl.promise = call ptr @"std::corouine_handle<promise_type>::from_address"(ptr %hdl)
+      %hdl.promise.tmp.dive = getelementptr inbounds %"struct.std::coroutine_handle.0",
+        ptr %hdl.promise.tmp, i32 0, i32 0
+      %hdl.promise.tmp.dive2 = getelementptr inbounds %"struct.std::coroutine_handle",
+        ptr %hdl.promise.tmp.dive, i32 0, i32 0
+      store ptr %hdl.promise, ptr %hdl.promise.tmp.dive2
+      call void @llvm.memcpy.p0.p0.i64(ptr %hdl.tmp, ptr %hdl.promise.tmp, i64 8, i1 false)
+      %hdl.tmp.dive = getelementptr inbounds %"struct.std::coroutine_handle",
+        ptr %hdl.tmp, i32 0, i32 0
+      %hdl.arg = load ptr, ptr %hdl.tmp.dive
+      %resume = call i1 @"Awaiter::await_suspend"(ptr %awaiter, ptr %hdl.arg)
+      ret i1 %resume
+
+.. _coro.await.suspend.handle:
+
+'llvm.coro.await.suspend.handle' Intrinsic
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+::
+
+  declare ptr @llvm.coro.await.suspend.handle(
+                ptr <awaiter>,
+                ptr <handle>,
+                ptr <await_suspend_function>)
+
+Overview:
+"""""""""
+
+The '``llvm.coro.await.suspend.handle``' intrinsic hides C++ `await-suspend`
+block code from optimizations on presplit coroutine body 
+to avoid miscompilations. This version of intrinsic corresponds to 
+'``std::corouine_handle<> awaiter.await_suspend(...)``' variant.
+
+Arguments:
+""""""""""
+
+The first argument is a pointer to `awaiter` object.
+
+The second argument is a pointer to the current coroutine's frame.
+
+The third argument is a pointer to the helper function encapsulating
+`await-suspend` logic. Its signature must be
+
+.. code-block:: llvm
+
+    declare ptr @await_suspend_function(ptr %awaiter, ptr %hdl)
+
+Semantics:
+""""""""""
+
+The intrinsic must be used between corresponding `coro.save`_ and 
+`coro.suspend`_ calls. It is lowered to an inlined 
+`await_suspend_function` call during `CoroSplit`_ pass.
+
+`await_suspend_function` must return a pointer to a valid
+coroutine frame, which is immediately resumed
+
+Example:
+""""""""
+
+.. code-block:: llvm
+
+  ; before lowering
+  await.suspend:
+    %save = call token @llvm.coro.save(ptr %hdl)
+    %next = call ptr @llvm.coro.await.suspend(
+                ptr %awaiter,
+                ptr %hdl,
+                ptr @await_suspend_function)
+    call void @llvm.coro.resume(%next)
+    ...
+
+  ; after lowering
+  await.suspend:
+    %save = call token @llvm.coro.save(ptr %hdl)
+    ; the call to await_suspend_function is inlined
+    %next = call ptr @await_suspend_function(
+                ptr %awaiter,
+                ptr %hdl)
+    call void @llvm.coro.resume(%next)
+    ...
+
+  ; helper function example
+  define ptr @await_suspend_function(ptr %awaiter, ptr %hdl)
+    entry:
+      %hdl.tmp = alloca %"struct.std::coroutine_handle"
+      %hdl.result.tmp = alloca %"struct.std::coroutine_handle"
+      %hdl.promise.tmp = alloca %"struct.std::coroutine_handle.0"
+      %hdl.promise = call ptr @"std::corouine_handle<promise_type>::from_address"(ptr %hdl)
+      %hdl.promise.tmp.dive = getelementptr inbounds %"struct.std::coroutine_handle.0",
+        ptr %hdl.promise.tmp, i32 0, i32 0
+      %hdl.promise.tmp.dive2 = getelementptr inbounds %"struct.std::coroutine_handle",
+        ptr %hdl.promise.tmp.dive, i32 0, i32 0
+      store ptr %hdl.promise, ptr %hdl.promise.tmp.dive2
+      call void @llvm.memcpy.p0.p0.i64(ptr %hdl.tmp, ptr %hdl.promise.tmp, i64 8, i1 false)
+      %hdl.tmp.dive = getelementptr inbounds %"struct.std::coroutine_handle",
+        ptr %hdl.tmp, i32 0, i32 0
+      %hdl.arg = load ptr, ptr %hdl.tmp.dive
+      %hdl.result = call ptr @"Awaiter::await_suspend"(ptr %awaiter, ptr %hdl.arg)
+      %hdl.result.tmp.dive = getelementptr inbounds %"struct.std::coroutine_handle",
+        ptr %hdl.result.tmp, i32 0, i32 0
+      store ptr %hdl.result, ptr %hdl.result.tmp.dive
+      %result.address = call ptr @"std::corouine_handle<>::address"(ptr %hdl.result.tmp)
+      ret ptr %result.address
+
 Coroutine Transformation Passes
 ===============================
 CoroEarly
@@ -1758,7 +2025,9 @@ and `coro.promise`_ intrinsics.
 CoroSplit
 ---------
 The pass CoroSplit builds coroutine frame and outlines resume and destroy parts
-into separate functions.
+into separate functions. This pass also lowers `coro.await.suspend`_,
+`coro.await.suspend.bool`_ and `coro.await.suspend.handle`_ intrinsics.
+
 
 CoroElide
 ---------

>From ccd27ca0725e0c03022411e8d0f0253762392ae3 Mon Sep 17 00:00:00 2001
From: fpasserby <125797601+fpasserby at users.noreply.github.com>
Date: Sun, 4 Feb 2024 14:13:30 +0100
Subject: [PATCH 05/17] Get rid of OpaqueFramePtr

---
 clang/include/clang/AST/ExprCXX.h         | 21 ++---
 clang/lib/CodeGen/CGCoroutine.cpp         | 94 +++++++++++++----------
 clang/lib/CodeGen/CodeGenFunction.h       | 15 ++++
 clang/lib/Sema/SemaCoroutine.cpp          | 22 ++----
 clang/lib/Serialization/ASTReaderStmt.cpp |  2 -
 clang/lib/Serialization/ASTWriterStmt.cpp |  1 -
 clang/test/SemaCXX/co_await-ast.cpp       |  7 +-
 7 files changed, 83 insertions(+), 79 deletions(-)

diff --git a/clang/include/clang/AST/ExprCXX.h b/clang/include/clang/AST/ExprCXX.h
index 42cabe891ac0fb2..a0e467b35778c5c 100644
--- a/clang/include/clang/AST/ExprCXX.h
+++ b/clang/include/clang/AST/ExprCXX.h
@@ -5036,17 +5036,14 @@ class CoroutineSuspendExpr : public Expr {
 
   Stmt *SubExprs[SubExpr::Count];
   OpaqueValueExpr *OpaqueValue = nullptr;
-  OpaqueValueExpr *OpaqueFramePtr = nullptr;
 
 public:
   CoroutineSuspendExpr(StmtClass SC, SourceLocation KeywordLoc, Expr *Operand,
                        Expr *Common, Expr *Ready, Expr *Suspend, Expr *Resume,
-                       OpaqueValueExpr *OpaqueValue,
-                       OpaqueValueExpr *OpaqueFramePtr)
+                       OpaqueValueExpr *OpaqueValue)
       : Expr(SC, Resume->getType(), Resume->getValueKind(),
              Resume->getObjectKind()),
-        KeywordLoc(KeywordLoc), OpaqueValue(OpaqueValue),
-        OpaqueFramePtr(OpaqueFramePtr) {
+        KeywordLoc(KeywordLoc), OpaqueValue(OpaqueValue) {
     SubExprs[SubExpr::Operand] = Operand;
     SubExprs[SubExpr::Common] = Common;
     SubExprs[SubExpr::Ready] = Ready;
@@ -5083,9 +5080,6 @@ class CoroutineSuspendExpr : public Expr {
   /// getOpaqueValue - Return the opaque value placeholder.
   OpaqueValueExpr *getOpaqueValue() const { return OpaqueValue; }
 
-  /// getOpaqueFramePtr - Return coroutine frame pointer placeholder.
-  OpaqueValueExpr *getOpaqueFramePtr() const { return OpaqueFramePtr; }
-
   Expr *getReadyExpr() const {
     return static_cast<Expr*>(SubExprs[SubExpr::Ready]);
   }
@@ -5132,11 +5126,9 @@ class CoawaitExpr : public CoroutineSuspendExpr {
 public:
   CoawaitExpr(SourceLocation CoawaitLoc, Expr *Operand, Expr *Common,
               Expr *Ready, Expr *Suspend, Expr *Resume,
-              OpaqueValueExpr *OpaqueValue, OpaqueValueExpr *OpaqueFramePtr,
-              bool IsImplicit = false)
+              OpaqueValueExpr *OpaqueValue, bool IsImplicit = false)
       : CoroutineSuspendExpr(CoawaitExprClass, CoawaitLoc, Operand, Common,
-                             Ready, Suspend, Resume, OpaqueValue,
-                             OpaqueFramePtr) {
+                             Ready, Suspend, Resume, OpaqueValue) {
     CoawaitBits.IsImplicit = IsImplicit;
   }
 
@@ -5215,10 +5207,9 @@ class CoyieldExpr : public CoroutineSuspendExpr {
 public:
   CoyieldExpr(SourceLocation CoyieldLoc, Expr *Operand, Expr *Common,
               Expr *Ready, Expr *Suspend, Expr *Resume,
-              OpaqueValueExpr *OpaqueValue, OpaqueValueExpr *OpaqueFramePtr)
+              OpaqueValueExpr *OpaqueValue)
       : CoroutineSuspendExpr(CoyieldExprClass, CoyieldLoc, Operand, Common,
-                             Ready, Suspend, Resume, OpaqueValue,
-                             OpaqueFramePtr) {}
+                             Ready, Suspend, Resume, OpaqueValue) {}
   CoyieldExpr(SourceLocation CoyieldLoc, QualType Ty, Expr *Operand,
               Expr *Common)
       : CoroutineSuspendExpr(CoyieldExprClass, CoyieldLoc, Ty, Operand,
diff --git a/clang/lib/CodeGen/CGCoroutine.cpp b/clang/lib/CodeGen/CGCoroutine.cpp
index 5bb4e7e950a74de..d6eefddcd24c21f 100644
--- a/clang/lib/CodeGen/CGCoroutine.cpp
+++ b/clang/lib/CodeGen/CGCoroutine.cpp
@@ -182,14 +182,27 @@ static bool AwaitSuspendStmtCanThrow(const Stmt *S) {
 //   auto && x = CommonExpr();
 //   if (!x.await_ready()) {
 //      llvm_coro_save();
-//      x.await_suspend(...);     (*)
+//      llvm_coro_await_suspend(&x, frame, helper) (*)
 //      llvm_coro_suspend(); (**)
 //   }
 //   x.await_resume();
 //
 // where the result of the entire expression is the result of x.await_resume()
 //
-//   (*) If x.await_suspend return type is bool, it allows to veto a suspend:
+//   (*) llvm_coro_await_suspend_{void, bool, handle} is lowered to
+//      helper(&x, frame) when it's certain not to interfere with
+//      coroutine frame generation. await_suspend expression is
+//      asynchronous to the coroutine body and not all analyses
+//      and transformations can currently handle it correctly.
+//
+//      Helper function encapsulates x.await_suspend(...) call and looks like:
+//
+//      auto __await_suspend_helper(auto& awaiter, void* frame) {
+//        std::coroutine_handle<> handle(frame);
+//        return awaiter.await_suspend(handle);
+//      }
+//
+//      If x.await_suspend return type is bool, it allows to veto a suspend:
 //      if (x.await_suspend(...))
 //        llvm_coro_suspend();
 //
@@ -245,39 +258,34 @@ static LValueOrRValue emitSuspendExpression(CodeGenFunction &CGF, CGCoroData &Co
 
   llvm::CallBase *SuspendRet = nullptr;
 
-  {
-    CGF.CurCoro.InSuspendBlock = true;
-
-    auto FramePtrBinder = CodeGenFunction::OpaqueValueMappingData::bind(
-        CGF, S.getOpaqueFramePtr(), S.getOpaqueFramePtr()->getSourceExpr());
-    auto UnbindFramePtrOnExit =
-        llvm::make_scope_exit([&] { FramePtrBinder.unbind(CGF); });
-
-    SmallVector<llvm::Value *, 3> SuspendHelperCallArgs;
-    SuspendHelperCallArgs.push_back(
-        CGF.getOrCreateOpaqueLValueMapping(S.getOpaqueValue()).getPointer(CGF));
-    SuspendHelperCallArgs.push_back(
-        CGF.getOrCreateOpaqueRValueMapping(S.getOpaqueFramePtr())
-            .getScalarVal());
-    SuspendHelperCallArgs.push_back(SuspendHelper);
-
-    auto IID = llvm::Intrinsic::coro_await_suspend;
-    if (S.getSuspendExpr()->getType()->isBooleanType())
-      IID = llvm::Intrinsic::coro_await_suspend_bool;
-    else if (S.getSuspendExpr()->getType()->isVoidPointerType())
-      IID = llvm::Intrinsic::coro_await_suspend_handle;
-
-    llvm::Function *AwaitSuspendIntrinsic = CGF.CGM.getIntrinsic(IID);
-    // FIXME: add call attributes?
-    if (AwaitSuspendCanThrow)
-      SuspendRet =
-          CGF.EmitCallOrInvoke(AwaitSuspendIntrinsic, SuspendHelperCallArgs);
-    else
-      SuspendRet = CGF.EmitNounwindRuntimeCall(AwaitSuspendIntrinsic,
-                                               SuspendHelperCallArgs);
+  CGF.CurCoro.InSuspendBlock = true;
 
-    CGF.CurCoro.InSuspendBlock = false;
-  }
+  assert(CGF.CurCoro.Data && CGF.CurCoro.Data->CoroBegin &&
+         "expected to be called in coroutine context");
+
+  SmallVector<llvm::Value *, 3> SuspendHelperCallArgs;
+  SuspendHelperCallArgs.push_back(
+      CGF.getOrCreateOpaqueLValueMapping(S.getOpaqueValue()).getPointer(CGF));
+
+  SuspendHelperCallArgs.push_back(CGF.CurCoro.Data->CoroBegin);
+  SuspendHelperCallArgs.push_back(SuspendHelper);
+
+  auto IID = llvm::Intrinsic::coro_await_suspend;
+  if (S.getSuspendExpr()->getType()->isBooleanType())
+    IID = llvm::Intrinsic::coro_await_suspend_bool;
+  else if (S.getSuspendExpr()->getType()->isVoidPointerType())
+    IID = llvm::Intrinsic::coro_await_suspend_handle;
+
+  llvm::Function *AwaitSuspendIntrinsic = CGF.CGM.getIntrinsic(IID);
+  // FIXME: add call attributes?
+  if (AwaitSuspendCanThrow)
+    SuspendRet =
+        CGF.EmitCallOrInvoke(AwaitSuspendIntrinsic, SuspendHelperCallArgs);
+  else
+    SuspendRet = CGF.EmitNounwindRuntimeCall(AwaitSuspendIntrinsic,
+                                             SuspendHelperCallArgs);
+
+  CGF.CurCoro.InSuspendBlock = false;
 
   if (SuspendRet != nullptr) {
     if (SuspendRet->getType()->isIntegerTy(1)) {
@@ -439,28 +447,25 @@ llvm::Function *CodeGenFunction::generateAwaitSuspendHelper(
   auto AwaiterLValue =
       MakeNaturalAlignAddrLValue(AwaiterPtr, AwaiterDecl.getType());
 
-  // FIXME: mark as aliasing with awaiter?
+  CurAwaitSuspendHelper.FramePtr =
+      Builder.CreateLoad(GetAddrOfLocalVar(&FrameDecl));
+
+  // FIXME: mark as aliasing with frame?
   // FIXME: TBAA?
   // FIXME: emit in a better way (maybe egenerate AST in SemaCoroutine)?
-  auto FramePtrRValue =
-      RValue::get(Builder.CreateLoad(GetAddrOfLocalVar(&FrameDecl)));
-
   auto AwaiterBinder = CodeGenFunction::OpaqueValueMappingData::bind(
       *this, S.getOpaqueValue(), AwaiterLValue);
-  auto FramePtrBinder = CodeGenFunction::OpaqueValueMappingData::bind(
-      *this, S.getOpaqueFramePtr(), FramePtrRValue);
 
   auto *SuspendRet = EmitScalarExpr(S.getSuspendExpr());
 
   auto UnbindCommonOnExit =
       llvm::make_scope_exit([&] { AwaiterBinder.unbind(*this); });
-  auto UnbindFramePtrOnExit =
-      llvm::make_scope_exit([&] { FramePtrBinder.unbind(*this); });
   if (SuspendRet != nullptr) {
     Fn->addRetAttr(llvm::Attribute::AttrKind::NoUndef);
     Builder.CreateStore(SuspendRet, ReturnValue);
   }
 
+  CurAwaitSuspendHelper.FramePtr = nullptr;
   FinishFunction();
   return Fn;
 }
@@ -961,6 +966,11 @@ RValue CodeGenFunction::EmitCoroutineIntrinsic(const CallExpr *E,
     if (CurCoro.Data && CurCoro.Data->CoroBegin) {
       return RValue::get(CurCoro.Data->CoroBegin);
     }
+
+    if (CurAwaitSuspendHelper.FramePtr) {
+      return RValue::get(CurAwaitSuspendHelper.FramePtr);
+    }
+
     CGM.Error(E->getBeginLoc(), "this builtin expect that __builtin_coro_begin "
                                 "has been used earlier in this function");
     auto *NullPtr = llvm::ConstantPointerNull::get(Builder.getPtrTy());
diff --git a/clang/lib/CodeGen/CodeGenFunction.h b/clang/lib/CodeGen/CodeGenFunction.h
index d13be4e2eae69dc..815211480c59e9b 100644
--- a/clang/lib/CodeGen/CodeGenFunction.h
+++ b/clang/lib/CodeGen/CodeGenFunction.h
@@ -352,6 +352,21 @@ class CodeGenFunction : public CodeGenTypeCache {
     return isCoroutine() && CurCoro.InSuspendBlock;
   }
 
+  // Holds FramePtr for await_suspend helper generation,
+  // so that __builtin_coro_frame call can be lowered
+  // directly to value of its second argument
+  struct AwaitSuspendHelperInfo {
+    llvm::Value *FramePtr = nullptr;
+  };
+  AwaitSuspendHelperInfo CurAwaitSuspendHelper;
+
+  // Generates helper function for `llvm.coro.await.suspend.*` intrinisics.
+  // It encapsulates SuspendExpr in a function, to separate it's body
+  // from the main coroutine to avoid miscompilations. Intrinisic
+  // is lowered to this function call in CoroSplit pass
+  // Function signature is:
+  // <type> __await_suspend_helper_<name>(ptr %awaiter, ptr %hdl)
+  // where type is one of (void, i1, ptr)
   llvm::Function *generateAwaitSuspendHelper(Twine const &CoroName,
                                              Twine const &SuspendPointName,
                                              CoroutineSuspendExpr const &S,
diff --git a/clang/lib/Sema/SemaCoroutine.cpp b/clang/lib/Sema/SemaCoroutine.cpp
index 39004d711a02db0..f250c3f07b7836f 100644
--- a/clang/lib/Sema/SemaCoroutine.cpp
+++ b/clang/lib/Sema/SemaCoroutine.cpp
@@ -293,7 +293,6 @@ struct ReadySuspendResumeResult {
   enum AwaitCallType { ACT_Ready, ACT_Suspend, ACT_Resume };
   Expr *Results[3];
   OpaqueValueExpr *OpaqueValue;
-  OpaqueValueExpr *OpaqueFramePtr;
   bool IsSuspendNoThrow;
   bool IsInvalid;
 };
@@ -403,7 +402,6 @@ static ReadySuspendResumeResult buildCoawaitCalls(Sema &S, VarDecl *CoroPromise,
   // Further operations are responsible for setting IsInalid to true.
   ReadySuspendResumeResult Calls = {{},
                                     Operand,
-                                    /*OpaqueFramePtr=*/nullptr,
                                     /*IsSuspendNoThrow=*/false,
                                     /*IsInvalid=*/false};
 
@@ -439,15 +437,9 @@ static ReadySuspendResumeResult buildCoawaitCalls(Sema &S, VarDecl *CoroPromise,
       Calls.Results[ACT::ACT_Ready] = S.MaybeCreateExprWithCleanups(Conv.get());
   }
 
-  Expr *GetFramePtr =
+  Expr *FramePtr =
       S.BuildBuiltinCallExpr(Loc, Builtin::BI__builtin_coro_frame, {});
 
-  OpaqueValueExpr *FramePtr = new (S.Context)
-      OpaqueValueExpr(Loc, GetFramePtr->getType(), VK_PRValue,
-                      GetFramePtr->getObjectKind(), GetFramePtr);
-
-  Calls.OpaqueFramePtr = FramePtr;
-
   ExprResult CoroHandleRes =
       buildCoroutineHandle(S, CoroPromise->getType(), FramePtr, Loc);
   if (CoroHandleRes.isInvalid()) {
@@ -941,9 +933,9 @@ ExprResult Sema::BuildResolvedCoawaitExpr(SourceLocation Loc, Expr *Operand,
   if (RSS.IsInvalid)
     return ExprError();
 
-  Expr *Res = new (Context) CoawaitExpr(
-      Loc, Operand, Awaiter, RSS.Results[0], RSS.Results[1], RSS.Results[2],
-      RSS.OpaqueValue, RSS.OpaqueFramePtr, IsImplicit);
+  Expr *Res = new (Context)
+      CoawaitExpr(Loc, Operand, Awaiter, RSS.Results[0], RSS.Results[1],
+                  RSS.Results[2], RSS.OpaqueValue, IsImplicit);
 
   return Res;
 }
@@ -999,9 +991,9 @@ ExprResult Sema::BuildCoyieldExpr(SourceLocation Loc, Expr *E) {
   if (RSS.IsInvalid)
     return ExprError();
 
-  Expr *Res = new (Context)
-      CoyieldExpr(Loc, Operand, E, RSS.Results[0], RSS.Results[1],
-                  RSS.Results[2], RSS.OpaqueValue, RSS.OpaqueFramePtr);
+  Expr *Res =
+      new (Context) CoyieldExpr(Loc, Operand, E, RSS.Results[0], RSS.Results[1],
+                                RSS.Results[2], RSS.OpaqueValue);
 
   return Res;
 }
diff --git a/clang/lib/Serialization/ASTReaderStmt.cpp b/clang/lib/Serialization/ASTReaderStmt.cpp
index 2b6d113779559ae..3da44ffccc38a28 100644
--- a/clang/lib/Serialization/ASTReaderStmt.cpp
+++ b/clang/lib/Serialization/ASTReaderStmt.cpp
@@ -484,7 +484,6 @@ void ASTStmtReader::VisitCoawaitExpr(CoawaitExpr *E) {
   for (auto &SubExpr: E->SubExprs)
     SubExpr = Record.readSubStmt();
   E->OpaqueValue = cast_or_null<OpaqueValueExpr>(Record.readSubStmt());
-  E->OpaqueFramePtr = cast_or_null<OpaqueValueExpr>(Record.readSubStmt());
   E->setIsImplicit(Record.readInt() != 0);
 }
 
@@ -494,7 +493,6 @@ void ASTStmtReader::VisitCoyieldExpr(CoyieldExpr *E) {
   for (auto &SubExpr: E->SubExprs)
     SubExpr = Record.readSubStmt();
   E->OpaqueValue = cast_or_null<OpaqueValueExpr>(Record.readSubStmt());
-  E->OpaqueFramePtr = cast_or_null<OpaqueValueExpr>(Record.readSubStmt());
 }
 
 void ASTStmtReader::VisitDependentCoawaitExpr(DependentCoawaitExpr *E) {
diff --git a/clang/lib/Serialization/ASTWriterStmt.cpp b/clang/lib/Serialization/ASTWriterStmt.cpp
index 232f0da4e0dd2e0..484621ae813093e 100644
--- a/clang/lib/Serialization/ASTWriterStmt.cpp
+++ b/clang/lib/Serialization/ASTWriterStmt.cpp
@@ -445,7 +445,6 @@ void ASTStmtWriter::VisitCoroutineSuspendExpr(CoroutineSuspendExpr *E) {
   for (Stmt *S : E->children())
     Record.AddStmt(S);
   Record.AddStmt(E->getOpaqueValue());
-  Record.AddStmt(E->getOpaqueFramePtr());
 }
 
 void ASTStmtWriter::VisitCoawaitExpr(CoawaitExpr *E) {
diff --git a/clang/test/SemaCXX/co_await-ast.cpp b/clang/test/SemaCXX/co_await-ast.cpp
index ed3a988653eb29e..10cee21da0e87cd 100644
--- a/clang/test/SemaCXX/co_await-ast.cpp
+++ b/clang/test/SemaCXX/co_await-ast.cpp
@@ -81,10 +81,9 @@ awaitable foo() {
 // CHECK:   |     |         `-CallExpr {{.*}} 'coroutine_handle<awaitable_frame>':'std::coroutine_handle<awaitable_frame>'
 // CHECK:   |     |           |-ImplicitCastExpr {{.*}} 'coroutine_handle<awaitable_frame> (*)(void *) noexcept' <FunctionToPointerDecay>
 // CHECK:   |     |           | `-DeclRefExpr {{.*}} 'coroutine_handle<awaitable_frame> (void *) noexcept' lvalue CXXMethod {{.*}} 'from_address' 'coroutine_handle<awaitable_frame> (void *) noexcept'
-// CHECK:   |     |           `-OpaqueValueExpr {{.*}} 'void *'
-// CHECK:   |     |             `-CallExpr {{.*}} 'void *'
-// CHECK:   |     |               `-ImplicitCastExpr {{.*}} 'void *(*)() noexcept' <FunctionToPointerDecay>
-// CHECK:   |     |                 `-DeclRefExpr {{.*}} 'void *() noexcept' lvalue Function {{.*}} '__builtin_coro_frame' 'void *() noexcept'
+// CHECK:   |     |           `-CallExpr {{.*}} 'void *'
+// CHECK:   |     |             `-ImplicitCastExpr {{.*}} 'void *(*)() noexcept' <FunctionToPointerDecay>
+// CHECK:   |     |               `-DeclRefExpr {{.*}} 'void *() noexcept' lvalue Function {{.*}} '__builtin_coro_frame' 'void *() noexcept'
 // CHECK:   |     `-CXXMemberCallExpr {{.*}} 'void'
 // CHECK:   |       `-MemberExpr {{.*}} '<bound member function type>' .await_resume {{.*}}
 // CHECK:   |         `-ImplicitCastExpr {{.*}} 'const awaitable_frame::result_t' lvalue <NoOp>

>From fcafc7c9eb7aaa11d533d24dd35fc2a4dfd023be Mon Sep 17 00:00:00 2001
From: fpasserby <125797601+fpasserby at users.noreply.github.com>
Date: Sun, 4 Feb 2024 15:03:21 +0100
Subject: [PATCH 06/17] Rename coro.await.suspend.void

---
 clang/lib/CodeGen/CGCoroutine.cpp                  |  2 +-
 clang/test/CodeGenCoroutines/coro-await.cpp        |  2 +-
 .../CodeGenCoroutines/coro-function-try-block.cpp  |  2 +-
 clang/test/CodeGenCoroutines/pr59181.cpp           |  2 +-
 llvm/docs/Coroutines.rst                           | 14 ++++++++------
 llvm/include/llvm/IR/Intrinsics.td                 |  2 +-
 llvm/lib/IR/Verifier.cpp                           |  2 +-
 llvm/lib/Transforms/Coroutines/CoroInstr.h         |  2 +-
 8 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/clang/lib/CodeGen/CGCoroutine.cpp b/clang/lib/CodeGen/CGCoroutine.cpp
index d6eefddcd24c21f..d64e64dec521d7f 100644
--- a/clang/lib/CodeGen/CGCoroutine.cpp
+++ b/clang/lib/CodeGen/CGCoroutine.cpp
@@ -270,7 +270,7 @@ static LValueOrRValue emitSuspendExpression(CodeGenFunction &CGF, CGCoroData &Co
   SuspendHelperCallArgs.push_back(CGF.CurCoro.Data->CoroBegin);
   SuspendHelperCallArgs.push_back(SuspendHelper);
 
-  auto IID = llvm::Intrinsic::coro_await_suspend;
+  auto IID = llvm::Intrinsic::coro_await_suspend_void;
   if (S.getSuspendExpr()->getType()->isBooleanType())
     IID = llvm::Intrinsic::coro_await_suspend_bool;
   else if (S.getSuspendExpr()->getType()->isVoidPointerType())
diff --git a/clang/test/CodeGenCoroutines/coro-await.cpp b/clang/test/CodeGenCoroutines/coro-await.cpp
index 701e1bf87b06f5d..377f14e8933cdc4 100644
--- a/clang/test/CodeGenCoroutines/coro-await.cpp
+++ b/clang/test/CodeGenCoroutines/coro-await.cpp
@@ -73,7 +73,7 @@ extern "C" void f0() {
   // ---------------------------
   // Call coro.await.suspend
   // ---------------------------
-  // CHECK-NEXT: call void @llvm.coro.await.suspend(ptr %[[AWAITABLE]], ptr %[[FRAME]], ptr @__await_suspend_helper_f0_await)
+  // CHECK-NEXT: call void @llvm.coro.await.suspend.void(ptr %[[AWAITABLE]], ptr %[[FRAME]], ptr @__await_suspend_helper_f0_await)
   // -------------------------
   // Generate a suspend point:
   // -------------------------
diff --git a/clang/test/CodeGenCoroutines/coro-function-try-block.cpp b/clang/test/CodeGenCoroutines/coro-function-try-block.cpp
index 023a95870361fd2..b7a796cc241af4b 100644
--- a/clang/test/CodeGenCoroutines/coro-function-try-block.cpp
+++ b/clang/test/CodeGenCoroutines/coro-function-try-block.cpp
@@ -19,5 +19,5 @@ task f() try {
 }
 
 // CHECK-LABEL: define{{.*}} void @_Z1fv(
-// CHECK: call void @llvm.coro.await.suspend(
+// CHECK: call void @llvm.coro.await.suspend.void(
 // CHECK: call void @_ZN4task12promise_type11return_voidEv(
diff --git a/clang/test/CodeGenCoroutines/pr59181.cpp b/clang/test/CodeGenCoroutines/pr59181.cpp
index ac732b6232e76e9..a3b050dacf0ba0e 100644
--- a/clang/test/CodeGenCoroutines/pr59181.cpp
+++ b/clang/test/CodeGenCoroutines/pr59181.cpp
@@ -57,5 +57,5 @@ void foo() {
 
 // CHECK: await.suspend:{{.*}}
 // CHECK-NOT: call void @llvm.lifetime
-// CHECK: call void @llvm.coro.await.suspend(
+// CHECK: call void @llvm.coro.await.suspend.void(
 // CHECK-NEXT: %{{[0-9]+}} = call i8 @llvm.coro.suspend(
diff --git a/llvm/docs/Coroutines.rst b/llvm/docs/Coroutines.rst
index fe892bbde67f639..2cd1ade7911e030 100644
--- a/llvm/docs/Coroutines.rst
+++ b/llvm/docs/Coroutines.rst
@@ -1746,11 +1746,11 @@ executes a call to ``llvm.coro.suspend.retcon`` after resuming in any way.
 
 .. _coro.await.suspend:
 
-'llvm.coro.await.suspend' Intrinsic
+'llvm.coro.await.suspend.void' Intrinsic
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 ::
 
-  declare void @llvm.coro.await.suspend(
+  declare void @llvm.coro.await.suspend.void(
                 ptr <awaiter>,
                 ptr <handle>,
                 ptr <await_suspend_function>)
@@ -1758,7 +1758,7 @@ executes a call to ``llvm.coro.suspend.retcon`` after resuming in any way.
 Overview:
 """""""""
 
-The '``llvm.coro.await.suspend``' intrinsic hides C++ `await-suspend`
+The '``llvm.coro.await.suspend.void``' intrinsic hides C++ `await-suspend`
 block code from optimizations on presplit coroutine body 
 to avoid miscompilations. This version of intrinsic corresponds to 
 '``void awaiter.await_suspend(...)``' variant.
@@ -1792,7 +1792,7 @@ Example:
   ; before lowering
   await.suspend:
     %save = call token @llvm.coro.save(ptr %hdl)
-    call void @llvm.coro.await.suspend(
+    call void @llvm.coro.await.suspend.void(
                 ptr %awaiter,
                 ptr %hdl,
                 ptr @await_suspend_function)
@@ -1879,7 +1879,7 @@ Example:
   ; before lowering
   await.suspend:
     %save = call token @llvm.coro.save(ptr %hdl)
-    %resume = call i1 @llvm.coro.await.suspend(
+    %resume = call i1 @llvm.coro.await.suspend.bool(
                 ptr %awaiter,
                 ptr %hdl,
                 ptr @await_suspend_function)
@@ -1971,11 +1971,12 @@ Example:
   ; before lowering
   await.suspend:
     %save = call token @llvm.coro.save(ptr %hdl)
-    %next = call ptr @llvm.coro.await.suspend(
+    %next = call ptr @llvm.coro.await.suspend.handle(
                 ptr %awaiter,
                 ptr %hdl,
                 ptr @await_suspend_function)
     call void @llvm.coro.resume(%next)
+    %suspend = call i8 @llvm.coro.suspend(token %save, i1 false)
     ...
 
   ; after lowering
@@ -1986,6 +1987,7 @@ Example:
                 ptr %awaiter,
                 ptr %hdl)
     call void @llvm.coro.resume(%next)
+    %suspend = call i8 @llvm.coro.suspend(token %save, i1 false)
     ...
 
   ; helper function example
diff --git a/llvm/include/llvm/IR/Intrinsics.td b/llvm/include/llvm/IR/Intrinsics.td
index c8270a24c55d17b..144298fd7c0162e 100644
--- a/llvm/include/llvm/IR/Intrinsics.td
+++ b/llvm/include/llvm/IR/Intrinsics.td
@@ -1692,7 +1692,7 @@ def int_coro_promise : Intrinsic<[llvm_ptr_ty],
                                  [llvm_ptr_ty, llvm_i32_ty, llvm_i1_ty],
                                  [IntrNoMem, NoCapture<ArgIndex<0>>]>;
 
-def int_coro_await_suspend : Intrinsic<[],
+def int_coro_await_suspend_void : Intrinsic<[],
                                        [llvm_ptr_ty, llvm_ptr_ty, llvm_ptr_ty],
                                        [Throws]>;
 
diff --git a/llvm/lib/IR/Verifier.cpp b/llvm/lib/IR/Verifier.cpp
index 59ee8c802b70034..1a5583e56aa8a2f 100644
--- a/llvm/lib/IR/Verifier.cpp
+++ b/llvm/lib/IR/Verifier.cpp
@@ -4994,7 +4994,7 @@ void Verifier::visitInstruction(Instruction &I) {
                 F->getIntrinsicID() == Intrinsic::seh_scope_end ||
                 F->getIntrinsicID() == Intrinsic::coro_resume ||
                 F->getIntrinsicID() == Intrinsic::coro_destroy ||
-                F->getIntrinsicID() == Intrinsic::coro_await_suspend ||
+                F->getIntrinsicID() == Intrinsic::coro_await_suspend_void ||
                 F->getIntrinsicID() == Intrinsic::coro_await_suspend_bool ||
                 F->getIntrinsicID() == Intrinsic::coro_await_suspend_handle ||
                 F->getIntrinsicID() ==
diff --git a/llvm/lib/Transforms/Coroutines/CoroInstr.h b/llvm/lib/Transforms/Coroutines/CoroInstr.h
index ce4fe5a4e4de4e6..63146cf6742aa6b 100644
--- a/llvm/lib/Transforms/Coroutines/CoroInstr.h
+++ b/llvm/lib/Transforms/Coroutines/CoroInstr.h
@@ -95,7 +95,7 @@ class LLVM_LIBRARY_VISIBILITY CoroAwaitSuspendInst : public CallBase {
   static bool classof(const CallBase *CB) {
     if (const Function *CF = CB->getCalledFunction()) {
       auto IID = CF->getIntrinsicID();
-      return IID == Intrinsic::coro_await_suspend ||
+      return IID == Intrinsic::coro_await_suspend_void ||
              IID == Intrinsic::coro_await_suspend_bool ||
              IID == Intrinsic::coro_await_suspend_handle;
     }

>From 87db125db5d3bbaf8e6e3fe42b6d255f040e092e Mon Sep 17 00:00:00 2001
From: fpasserby <125797601+fpasserby at users.noreply.github.com>
Date: Sun, 4 Feb 2024 15:52:54 +0100
Subject: [PATCH 07/17] Switch on await_suspend return type

---
 clang/include/clang/AST/ExprCXX.h | 18 ++++++++++
 clang/lib/CodeGen/CGCoroutine.cpp | 58 +++++++++++++++++++------------
 2 files changed, 53 insertions(+), 23 deletions(-)

diff --git a/clang/include/clang/AST/ExprCXX.h b/clang/include/clang/AST/ExprCXX.h
index a0e467b35778c5c..82bebd8c781bf31 100644
--- a/clang/include/clang/AST/ExprCXX.h
+++ b/clang/include/clang/AST/ExprCXX.h
@@ -5038,6 +5038,8 @@ class CoroutineSuspendExpr : public Expr {
   OpaqueValueExpr *OpaqueValue = nullptr;
 
 public:
+  enum SuspendReturnType { SuspendVoid, SuspendBool, SuspendHandle };
+
   CoroutineSuspendExpr(StmtClass SC, SourceLocation KeywordLoc, Expr *Operand,
                        Expr *Common, Expr *Ready, Expr *Suspend, Expr *Resume,
                        OpaqueValueExpr *OpaqueValue)
@@ -5097,6 +5099,22 @@ class CoroutineSuspendExpr : public Expr {
     return static_cast<Expr *>(SubExprs[SubExpr::Operand]);
   }
 
+  SuspendReturnType getSuspendReturnType() const {
+    auto *SuspendExpr = getSuspendExpr();
+    assert(SuspendExpr);
+
+    auto SuspendType = SuspendExpr->getType();
+
+    if (SuspendType->isVoidType())
+      return SuspendReturnType::SuspendVoid;
+    if (SuspendType->isBooleanType())
+      return SuspendReturnType::SuspendBool;
+    if (SuspendType->isVoidPointerType())
+      return SuspendReturnType::SuspendHandle;
+
+    llvm_unreachable("Unexpected await_suspend expression return type");
+  }
+
   SourceLocation getKeywordLoc() const { return KeywordLoc; }
 
   SourceLocation getBeginLoc() const LLVM_READONLY { return KeywordLoc; }
diff --git a/clang/lib/CodeGen/CGCoroutine.cpp b/clang/lib/CodeGen/CGCoroutine.cpp
index d64e64dec521d7f..3d14d425e32e647 100644
--- a/clang/lib/CodeGen/CGCoroutine.cpp
+++ b/clang/lib/CodeGen/CGCoroutine.cpp
@@ -193,7 +193,7 @@ static bool AwaitSuspendStmtCanThrow(const Stmt *S) {
 //      helper(&x, frame) when it's certain not to interfere with
 //      coroutine frame generation. await_suspend expression is
 //      asynchronous to the coroutine body and not all analyses
-//      and transformations can currently handle it correctly.
+//      and transformations can handle it correctly at the moment.
 //
 //      Helper function encapsulates x.await_suspend(...) call and looks like:
 //
@@ -270,11 +270,20 @@ static LValueOrRValue emitSuspendExpression(CodeGenFunction &CGF, CGCoroData &Co
   SuspendHelperCallArgs.push_back(CGF.CurCoro.Data->CoroBegin);
   SuspendHelperCallArgs.push_back(SuspendHelper);
 
-  auto IID = llvm::Intrinsic::coro_await_suspend_void;
-  if (S.getSuspendExpr()->getType()->isBooleanType())
+  const auto SuspendReturnType = S.getSuspendReturnType();
+  llvm::Intrinsic::ID IID;
+
+  switch (SuspendReturnType) {
+  case CoroutineSuspendExpr::SuspendVoid:
+    IID = llvm::Intrinsic::coro_await_suspend_void;
+    break;
+  case CoroutineSuspendExpr::SuspendBool:
     IID = llvm::Intrinsic::coro_await_suspend_bool;
-  else if (S.getSuspendExpr()->getType()->isVoidPointerType())
+    break;
+  case CoroutineSuspendExpr::SuspendHandle:
     IID = llvm::Intrinsic::coro_await_suspend_handle;
+    break;
+  }
 
   llvm::Function *AwaitSuspendIntrinsic = CGF.CGM.getIntrinsic(IID);
   // FIXME: add call attributes?
@@ -285,19 +294,30 @@ static LValueOrRValue emitSuspendExpression(CodeGenFunction &CGF, CGCoroData &Co
     SuspendRet = CGF.EmitNounwindRuntimeCall(AwaitSuspendIntrinsic,
                                              SuspendHelperCallArgs);
 
+  assert(SuspendRet);
   CGF.CurCoro.InSuspendBlock = false;
 
-  if (SuspendRet != nullptr) {
-    if (SuspendRet->getType()->isIntegerTy(1)) {
-      // Veto suspension if requested by bool returning await_suspend.
-      BasicBlock *RealSuspendBlock =
-          CGF.createBasicBlock(Prefix + Twine(".suspend.bool"));
-      CGF.Builder.CreateCondBr(SuspendRet, RealSuspendBlock, ReadyBlock);
-      CGF.EmitBlock(RealSuspendBlock);
-    } else if (SuspendRet->getType()->isPointerTy()) {
-      auto ResumeIntrinsic = CGF.CGM.getIntrinsic(llvm::Intrinsic::coro_resume);
-      Builder.CreateCall(ResumeIntrinsic, SuspendRet);
-    }
+  switch (SuspendReturnType) {
+  case CoroutineSuspendExpr::SuspendVoid:
+    assert(SuspendRet->getType()->isVoidTy());
+    break;
+  case CoroutineSuspendExpr::SuspendBool: {
+    assert(SuspendRet->getType()->isIntegerTy());
+
+    // Veto suspension if requested by bool returning await_suspend.
+    BasicBlock *RealSuspendBlock =
+        CGF.createBasicBlock(Prefix + Twine(".suspend.bool"));
+    CGF.Builder.CreateCondBr(SuspendRet, RealSuspendBlock, ReadyBlock);
+    CGF.EmitBlock(RealSuspendBlock);
+    break;
+  }
+  case CoroutineSuspendExpr::SuspendHandle: {
+    assert(SuspendRet->getType()->isPointerTy());
+
+    auto ResumeIntrinsic = CGF.CGM.getIntrinsic(llvm::Intrinsic::coro_resume);
+    Builder.CreateCall(ResumeIntrinsic, SuspendRet);
+    break;
+  }
   }
 
   // Emit the suspend point.
@@ -410,14 +430,6 @@ llvm::Function *CodeGenFunction::generateAwaitSuspendHelper(
   ImplicitParamDecl FrameDecl(C, C.VoidPtrTy, ImplicitParamKind::Other);
   QualType ReturnTy = S.getSuspendExpr()->getType();
 
-  if (ReturnTy->isBooleanType()) {
-    ReturnTy = C.BoolTy;
-  } else if (ReturnTy->isVoidPointerType()) {
-    ReturnTy = C.VoidPtrTy;
-  } else {
-    ReturnTy = C.VoidTy;
-  }
-
   args.push_back(&AwaiterDecl);
   args.push_back(&FrameDecl);
 

>From 7c918a1284f481b09fed069c48dfc12941beebec Mon Sep 17 00:00:00 2001
From: fpasserby <125797601+fpasserby at users.noreply.github.com>
Date: Sun, 4 Feb 2024 15:58:25 +0100
Subject: [PATCH 08/17] Don't mark await_suspend_helper as nounwind

---
 clang/lib/CodeGen/CGCoroutine.cpp   | 13 +++++--------
 clang/lib/CodeGen/CodeGenFunction.h |  3 +--
 2 files changed, 6 insertions(+), 10 deletions(-)

diff --git a/clang/lib/CodeGen/CGCoroutine.cpp b/clang/lib/CodeGen/CGCoroutine.cpp
index 3d14d425e32e647..ca06a927de5c977 100644
--- a/clang/lib/CodeGen/CGCoroutine.cpp
+++ b/clang/lib/CodeGen/CGCoroutine.cpp
@@ -254,7 +254,7 @@ static LValueOrRValue emitSuspendExpression(CodeGenFunction &CGF, CGCoroData &Co
       AwaitSuspendStmtCanThrow(S.getSuspendExpr());
 
   auto SuspendHelper = CodeGenFunction(CGF.CGM).generateAwaitSuspendHelper(
-      CGF.CurFn->getName(), Prefix, S, AwaitSuspendCanThrow);
+      CGF.CurFn->getName(), Prefix, S);
 
   llvm::CallBase *SuspendRet = nullptr;
 
@@ -414,9 +414,10 @@ static QualType getCoroutineSuspendExprReturnType(const ASTContext &Ctx,
 }
 #endif
 
-llvm::Function *CodeGenFunction::generateAwaitSuspendHelper(
-    Twine const &CoroName, Twine const &SuspendPointName,
-    CoroutineSuspendExpr const &S, bool CanThrow) {
+llvm::Function *
+CodeGenFunction::generateAwaitSuspendHelper(Twine const &CoroName,
+                                            Twine const &SuspendPointName,
+                                            CoroutineSuspendExpr const &S) {
   std::string FuncName = "__await_suspend_helper_";
   FuncName += CoroName.str();
   FuncName += '_';
@@ -449,10 +450,6 @@ llvm::Function *CodeGenFunction::generateAwaitSuspendHelper(
   Fn->setMustProgress();
   Fn->addFnAttr(llvm::Attribute::AttrKind::AlwaysInline);
 
-  if (!CanThrow) {
-    Fn->addFnAttr(llvm::Attribute::AttrKind::NoUnwind);
-  }
-
   StartFunction(GlobalDecl(), ReturnTy, Fn, FI, args);
 
   llvm::Value *AwaiterPtr = Builder.CreateLoad(GetAddrOfLocalVar(&AwaiterDecl));
diff --git a/clang/lib/CodeGen/CodeGenFunction.h b/clang/lib/CodeGen/CodeGenFunction.h
index 815211480c59e9b..243edd58a7e3541 100644
--- a/clang/lib/CodeGen/CodeGenFunction.h
+++ b/clang/lib/CodeGen/CodeGenFunction.h
@@ -369,8 +369,7 @@ class CodeGenFunction : public CodeGenTypeCache {
   // where type is one of (void, i1, ptr)
   llvm::Function *generateAwaitSuspendHelper(Twine const &CoroName,
                                              Twine const &SuspendPointName,
-                                             CoroutineSuspendExpr const &S,
-                                             bool CanThrow);
+                                             CoroutineSuspendExpr const &S);
 
   /// CurGD - The GlobalDecl for the current function being compiled.
   GlobalDecl CurGD;

>From 221ec9cedb538da919fc688389c832abfcfbad35 Mon Sep 17 00:00:00 2001
From: fpasserby <125797601+fpasserby at users.noreply.github.com>
Date: Sun, 4 Feb 2024 16:05:06 +0100
Subject: [PATCH 09/17] Don't inline in CoroSplit

---
 clang/test/CodeGenCoroutines/pr65054.cpp     | 2 +-
 llvm/lib/Transforms/Coroutines/CoroSplit.cpp | 5 -----
 2 files changed, 1 insertion(+), 6 deletions(-)

diff --git a/clang/test/CodeGenCoroutines/pr65054.cpp b/clang/test/CodeGenCoroutines/pr65054.cpp
index 834b71050f59ffa..55bb52233394b1f 100644
--- a/clang/test/CodeGenCoroutines/pr65054.cpp
+++ b/clang/test/CodeGenCoroutines/pr65054.cpp
@@ -55,6 +55,6 @@ MyTask FooBar() {
 // FRONTEND: attributes #[[address_attr]] = {{.*}}alwaysinline
 
 // CHECK-O0: define{{.*}}@_Z6FooBarv.resume
-// CHECK-O0: call{{.*}}@_ZN7Awaiter13await_suspendESt16coroutine_handleIvE
+// CHECK-O0: call{{.*}}@__await_suspend_helper__Z6FooBarv_await(
 // CHECK-O0-NOT: store
 // CHECK-O0: ret void
diff --git a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
index 00aeea2df2db039..4147c8864935813 100644
--- a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
@@ -123,11 +123,6 @@ void Lowerer::lowerAwaitSuspend(CoroAwaitSuspendInst *CB) {
 
   CB->replaceAllUsesWith(NewCall);
   CB->eraseFromParent();
-
-  InlineFunctionInfo FnInfo;
-  auto InlineRes = InlineFunction(*NewCall, FnInfo);
-  assert(InlineRes.isSuccess() && "Expected inlining to succeed");
-  (void)InlineRes;
 }
 
 void Lowerer::lowerAwaitSuspends(Function &F) {

>From 832be2b14e30f34217df42f642e2acb7f07d2665 Mon Sep 17 00:00:00 2001
From: fpasserby <125797601+fpasserby at users.noreply.github.com>
Date: Sun, 4 Feb 2024 17:04:01 +0100
Subject: [PATCH 10/17] Improve docs

---
 clang/lib/CodeGen/CGCoroutine.cpp |   2 +-
 llvm/docs/Coroutines.rst          | 117 ++++++++++++++----------------
 2 files changed, 55 insertions(+), 64 deletions(-)

diff --git a/clang/lib/CodeGen/CGCoroutine.cpp b/clang/lib/CodeGen/CGCoroutine.cpp
index ca06a927de5c977..6f5a19af0b37a11 100644
--- a/clang/lib/CodeGen/CGCoroutine.cpp
+++ b/clang/lib/CodeGen/CGCoroutine.cpp
@@ -191,7 +191,7 @@ static bool AwaitSuspendStmtCanThrow(const Stmt *S) {
 //
 //   (*) llvm_coro_await_suspend_{void, bool, handle} is lowered to
 //      helper(&x, frame) when it's certain not to interfere with
-//      coroutine frame generation. await_suspend expression is
+//      coroutine transform. await_suspend expression is
 //      asynchronous to the coroutine body and not all analyses
 //      and transformations can handle it correctly at the moment.
 //
diff --git a/llvm/docs/Coroutines.rst b/llvm/docs/Coroutines.rst
index 2cd1ade7911e030..a175c7de24ed601 100644
--- a/llvm/docs/Coroutines.rst
+++ b/llvm/docs/Coroutines.rst
@@ -1744,7 +1744,7 @@ a call to ``llvm.coro.suspend.retcon`` after resuming abnormally.
 In a yield-once coroutine, it is undefined behavior if the coroutine
 executes a call to ``llvm.coro.suspend.retcon`` after resuming in any way.
 
-.. _coro.await.suspend:
+.. _coro.await.suspend.void:
 
 'llvm.coro.await.suspend.void' Intrinsic
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
@@ -1758,9 +1758,19 @@ executes a call to ``llvm.coro.suspend.retcon`` after resuming in any way.
 Overview:
 """""""""
 
-The '``llvm.coro.await.suspend.void``' intrinsic hides C++ `await-suspend`
-block code from optimizations on presplit coroutine body 
-to avoid miscompilations. This version of intrinsic corresponds to 
+The '``llvm.coro.await.suspend.void``' intrinsic encapsulates C++ 
+`await-suspend` block until it can't interfere with coroutine transform.
+
+The `await_suspend` block of `co_await` is essentially asynchronous
+to the execution of the coroutine. Inlining it normally into an unsplit
+coroutine can cause miscompilation because the coroutine CFG misrepresents
+the true control flow of the program: things that happen in the
+await_suspend are not guaranteed to happen prior to the resumption of the
+coroutine, and things that happen after the resumption of the coroutine
+(including its exit and the potential deallocation of the coroutine frame)
+are not guaranteed to happen only after the end of `await_suspend`.
+
+This version of intrinsic corresponds to 
 '``void awaiter.await_suspend(...)``' variant.
 
 Arguments:
@@ -1781,7 +1791,7 @@ Semantics:
 """"""""""
 
 The intrinsic must be used between corresponding `coro.save`_ and 
-`coro.suspend`_ calls. It is lowered to an inlined 
+`coro.suspend`_ calls. It is lowered to a direct 
 `await_suspend_function` call during `CoroSplit`_ pass.
 
 Example:
@@ -1802,7 +1812,7 @@ Example:
   ; after lowering
   await.suspend:
     %save = call token @llvm.coro.save(ptr %hdl)
-    ; the call to await_suspend_function is inlined
+    ; the call to await_suspend_function can be inlined
     call void @await_suspend_function(
                 ptr %awaiter,
                 ptr %hdl)
@@ -1812,19 +1822,7 @@ Example:
   ; helper function example
   define void @await_suspend_function(ptr %awaiter, ptr %hdl)
     entry:
-      %hdl.tmp = alloca %"struct.std::coroutine_handle"
-      %hdl.result.tmp = alloca %"struct.std::coroutine_handle"
-      %hdl.promise.tmp = alloca %"struct.std::coroutine_handle.0"
-      %hdl.promise = call ptr @"std::corouine_handle<promise_type>::from_address"(ptr %hdl)
-      %hdl.promise.tmp.dive = getelementptr inbounds %"struct.std::coroutine_handle.0",
-        ptr %hdl.promise.tmp, i32 0, i32 0
-      %hdl.promise.tmp.dive2 = getelementptr inbounds %"struct.std::coroutine_handle",
-        ptr %hdl.promise.tmp.dive, i32 0, i32 0
-      store ptr %hdl.promise, ptr %hdl.promise.tmp.dive2
-      call void @llvm.memcpy.p0.p0.i64(ptr %hdl.tmp, ptr %hdl.promise.tmp, i64 8, i1 false)
-      %hdl.tmp.dive = getelementptr inbounds %"struct.std::coroutine_handle",
-        ptr %hdl.tmp, i32 0, i32 0
-      %hdl.arg = load ptr, ptr %hdl.tmp.dive
+      %hdl.arg = ... ; construct std::coroutine_handle from %hdl
       call void @"Awaiter::await_suspend"(ptr %awaiter, ptr %hdl.arg)
       ret void
 
@@ -1842,9 +1840,19 @@ Example:
 Overview:
 """""""""
 
-The '``llvm.coro.await.suspend.bool``' intrinsic hides C++ `await-suspend`
-block code from optimizations on presplit coroutine body 
-to avoid miscompilations. This version of intrinsic corresponds to 
+The '``llvm.coro.await.suspend.bool``' intrinsic encapsulates C++
+`await-suspend` block until it can't interfere with coroutine transform.
+
+The `await_suspend` block of `co_await` is essentially asynchronous
+to the execution of the coroutine. Inlining it normally into an unsplit
+coroutine can cause miscompilation because the coroutine CFG misrepresents
+the true control flow of the program: things that happen in the
+await_suspend are not guaranteed to happen prior to the resumption of the
+coroutine, and things that happen after the resumption of the coroutine
+(including its exit and the potential deallocation of the coroutine frame)
+are not guaranteed to happen only after the end of `await_suspend`.
+
+This version of intrinsic corresponds to 
 '``bool awaiter.await_suspend(...)``' variant.
 
 Arguments:
@@ -1865,7 +1873,7 @@ Semantics:
 """"""""""
 
 The intrinsic must be used between corresponding `coro.save`_ and 
-`coro.suspend`_ calls. It is lowered to an inlined 
+`coro.suspend`_ calls. It is lowered to a direct 
 `await_suspend_function` call during `CoroSplit`_ pass.
 
 If `await_suspend_function` call returns `true`, the current coroutine is
@@ -1888,13 +1896,13 @@ Example:
     %suspend = call i8 @llvm.coro.suspend(token %save, i1 false)
     ...
   await.ready:
-    call void @"Awaiter::await_ready"(ptr %awaiter)
+    call void @"Awaiter::await_resume"(ptr %awaiter)
     ...
 
   ; after lowering
   await.suspend:
     %save = call token @llvm.coro.save(ptr %hdl)
-    ; the call to await_suspend_function is inlined
+    ; the call to await_suspend_function can inlined
     %resume = call i1 @await_suspend_function(
                 ptr %awaiter,
                 ptr %hdl)
@@ -1904,19 +1912,7 @@ Example:
   ; helper function example
   define i1 @await_suspend_function(ptr %awaiter, ptr %hdl)
     entry:
-      %hdl.tmp = alloca %"struct.std::coroutine_handle"
-      %hdl.result.tmp = alloca %"struct.std::coroutine_handle"
-      %hdl.promise.tmp = alloca %"struct.std::coroutine_handle.0"
-      %hdl.promise = call ptr @"std::corouine_handle<promise_type>::from_address"(ptr %hdl)
-      %hdl.promise.tmp.dive = getelementptr inbounds %"struct.std::coroutine_handle.0",
-        ptr %hdl.promise.tmp, i32 0, i32 0
-      %hdl.promise.tmp.dive2 = getelementptr inbounds %"struct.std::coroutine_handle",
-        ptr %hdl.promise.tmp.dive, i32 0, i32 0
-      store ptr %hdl.promise, ptr %hdl.promise.tmp.dive2
-      call void @llvm.memcpy.p0.p0.i64(ptr %hdl.tmp, ptr %hdl.promise.tmp, i64 8, i1 false)
-      %hdl.tmp.dive = getelementptr inbounds %"struct.std::coroutine_handle",
-        ptr %hdl.tmp, i32 0, i32 0
-      %hdl.arg = load ptr, ptr %hdl.tmp.dive
+      %hdl.arg = ... ; construct std::coroutine_handle from %hdl
       %resume = call i1 @"Awaiter::await_suspend"(ptr %awaiter, ptr %hdl.arg)
       ret i1 %resume
 
@@ -1934,9 +1930,19 @@ Example:
 Overview:
 """""""""
 
-The '``llvm.coro.await.suspend.handle``' intrinsic hides C++ `await-suspend`
-block code from optimizations on presplit coroutine body 
-to avoid miscompilations. This version of intrinsic corresponds to 
+The '``llvm.coro.await.suspend.handle``' intrinsic encapsulates C++
+`await-suspend` block until it can't interfere with coroutine transform.
+
+The `await_suspend` block of `co_await` is essentially asynchronous
+to the execution of the coroutine. Inlining it normally into an unsplit
+coroutine can cause miscompilation because the coroutine CFG misrepresents
+the true control flow of the program: things that happen in the
+await_suspend are not guaranteed to happen prior to the resumption of the
+coroutine, and things that happen after the resumption of the coroutine
+(including its exit and the potential deallocation of the coroutine frame)
+are not guaranteed to happen only after the end of `await_suspend`.
+
+This version of intrinsic corresponds to 
 '``std::corouine_handle<> awaiter.await_suspend(...)``' variant.
 
 Arguments:
@@ -1957,7 +1963,7 @@ Semantics:
 """"""""""
 
 The intrinsic must be used between corresponding `coro.save`_ and 
-`coro.suspend`_ calls. It is lowered to an inlined 
+`coro.suspend`_ calls. It is lowered to a direct 
 `await_suspend_function` call during `CoroSplit`_ pass.
 
 `await_suspend_function` must return a pointer to a valid
@@ -1982,7 +1988,7 @@ Example:
   ; after lowering
   await.suspend:
     %save = call token @llvm.coro.save(ptr %hdl)
-    ; the call to await_suspend_function is inlined
+    ; the call to await_suspend_function can be inlined
     %next = call ptr @await_suspend_function(
                 ptr %awaiter,
                 ptr %hdl)
@@ -1993,25 +1999,10 @@ Example:
   ; helper function example
   define ptr @await_suspend_function(ptr %awaiter, ptr %hdl)
     entry:
-      %hdl.tmp = alloca %"struct.std::coroutine_handle"
-      %hdl.result.tmp = alloca %"struct.std::coroutine_handle"
-      %hdl.promise.tmp = alloca %"struct.std::coroutine_handle.0"
-      %hdl.promise = call ptr @"std::corouine_handle<promise_type>::from_address"(ptr %hdl)
-      %hdl.promise.tmp.dive = getelementptr inbounds %"struct.std::coroutine_handle.0",
-        ptr %hdl.promise.tmp, i32 0, i32 0
-      %hdl.promise.tmp.dive2 = getelementptr inbounds %"struct.std::coroutine_handle",
-        ptr %hdl.promise.tmp.dive, i32 0, i32 0
-      store ptr %hdl.promise, ptr %hdl.promise.tmp.dive2
-      call void @llvm.memcpy.p0.p0.i64(ptr %hdl.tmp, ptr %hdl.promise.tmp, i64 8, i1 false)
-      %hdl.tmp.dive = getelementptr inbounds %"struct.std::coroutine_handle",
-        ptr %hdl.tmp, i32 0, i32 0
-      %hdl.arg = load ptr, ptr %hdl.tmp.dive
-      %hdl.result = call ptr @"Awaiter::await_suspend"(ptr %awaiter, ptr %hdl.arg)
-      %hdl.result.tmp.dive = getelementptr inbounds %"struct.std::coroutine_handle",
-        ptr %hdl.result.tmp, i32 0, i32 0
-      store ptr %hdl.result, ptr %hdl.result.tmp.dive
-      %result.address = call ptr @"std::corouine_handle<>::address"(ptr %hdl.result.tmp)
-      ret ptr %result.address
+      %hdl.arg = ... ; construct std::coroutine_handle from %hdl
+      %hdl.raw = call ptr @"Awaiter::await_suspend"(ptr %awaiter, ptr %hdl.arg)
+      %hdl.result = ... ; get address of returned coroutine handle
+      ret ptr %hdl.result
 
 Coroutine Transformation Passes
 ===============================
@@ -2027,7 +2018,7 @@ and `coro.promise`_ intrinsics.
 CoroSplit
 ---------
 The pass CoroSplit builds coroutine frame and outlines resume and destroy parts
-into separate functions. This pass also lowers `coro.await.suspend`_,
+into separate functions. This pass also lowers `coro.await.suspend.void`_,
 `coro.await.suspend.bool`_ and `coro.await.suspend.handle`_ intrinsics.
 
 

>From b3361282ba0279679600eb18b9d90d1590ab222a Mon Sep 17 00:00:00 2001
From: fpasserby <125797601+fpasserby at users.noreply.github.com>
Date: Sun, 4 Feb 2024 17:12:22 +0100
Subject: [PATCH 11/17] Rename to await.suspend.void in one more place

---
 llvm/lib/Transforms/Coroutines/Coroutines.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/lib/Transforms/Coroutines/Coroutines.cpp b/llvm/lib/Transforms/Coroutines/Coroutines.cpp
index 8ecdf019a2c84a4..f8f3edd9292a15d 100644
--- a/llvm/lib/Transforms/Coroutines/Coroutines.cpp
+++ b/llvm/lib/Transforms/Coroutines/Coroutines.cpp
@@ -67,7 +67,7 @@ static const char *const CoroIntrinsics[] = {
     "llvm.coro.async.resume",
     "llvm.coro.async.size.replace",
     "llvm.coro.async.store_resume",
-    "llvm.coro.await.suspend",
+    "llvm.coro.await.suspend.void",
     "llvm.coro.await.suspend.bool",
     "llvm.coro.await.suspend.handle",
     "llvm.coro.begin",

>From ca78ac0fcf2fcf752b0c009c60412376d76d8a5d Mon Sep 17 00:00:00 2001
From: fpasserby <125797601+fpasserby at users.noreply.github.com>
Date: Sun, 4 Feb 2024 17:56:57 +0100
Subject: [PATCH 12/17] Improve lowering

---
 llvm/lib/Transforms/Coroutines/CoroInternal.h |   1 +
 llvm/lib/Transforms/Coroutines/CoroSplit.cpp  | 106 +++++++-----------
 llvm/lib/Transforms/Coroutines/Coroutines.cpp |   7 +-
 3 files changed, 49 insertions(+), 65 deletions(-)

diff --git a/llvm/lib/Transforms/Coroutines/CoroInternal.h b/llvm/lib/Transforms/Coroutines/CoroInternal.h
index 388cf8d2aee71c8..09d1430b4c559e7 100644
--- a/llvm/lib/Transforms/Coroutines/CoroInternal.h
+++ b/llvm/lib/Transforms/Coroutines/CoroInternal.h
@@ -84,6 +84,7 @@ struct LLVM_LIBRARY_VISIBILITY Shape {
   SmallVector<CoroAlignInst *, 2> CoroAligns;
   SmallVector<AnyCoroSuspendInst *, 4> CoroSuspends;
   SmallVector<CallInst*, 2> SwiftErrorOps;
+  SmallVector<CoroAwaitSuspendInst *, 4> CoroAwaitSuspends;
 
   // Field indexes for special fields in the switch lowering.
   struct SwitchFieldIndex {
diff --git a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
index 4147c8864935813..8eef64658ed7214 100644
--- a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
@@ -79,68 +79,6 @@ using namespace llvm;
 
 namespace {
 
-// Created on demand if the coro-early pass has work to do.
-class Lowerer : public coro::LowererBase {
-  IRBuilder<> Builder;
-  void lowerAwaitSuspend(CoroAwaitSuspendInst *CB);
-
-public:
-  Lowerer(Module &M) : LowererBase(M), Builder(Context) {}
-
-  void lowerAwaitSuspends(Function &F);
-};
-
-void Lowerer::lowerAwaitSuspend(CoroAwaitSuspendInst *CB) {
-  auto Helper = CB->getHelperFunction();
-  auto Awaiter = CB->getAwaiter();
-  auto FramePtr = CB->getFrame();
-
-  Builder.SetInsertPoint(CB);
-
-  CallBase *NewCall = nullptr;
-  if (auto Invoke = dyn_cast<InvokeInst>(CB)) {
-    auto HelperInvoke =
-        Builder.CreateInvoke(Helper, Invoke->getNormalDest(),
-                             Invoke->getUnwindDest(), {Awaiter, FramePtr});
-
-    HelperInvoke->setCallingConv(Invoke->getCallingConv());
-    std::copy(Invoke->bundle_op_info_begin(), Invoke->bundle_op_info_end(),
-              HelperInvoke->bundle_op_info_begin());
-    AttributeList NewAttributes =
-        Invoke->getAttributes().removeParamAttributes(Context, 2);
-    HelperInvoke->setAttributes(NewAttributes);
-    HelperInvoke->setDebugLoc(Invoke->getDebugLoc());
-    NewCall = HelperInvoke;
-  } else if (auto Call = dyn_cast<CallInst>(CB)) {
-    auto HelperCall = Builder.CreateCall(Helper, {Awaiter, FramePtr});
-
-    AttributeList NewAttributes =
-        Call->getAttributes().removeParamAttributes(Context, 2);
-    HelperCall->setAttributes(NewAttributes);
-    HelperCall->setDebugLoc(Call->getDebugLoc());
-    NewCall = HelperCall;
-  }
-
-  CB->replaceAllUsesWith(NewCall);
-  CB->eraseFromParent();
-}
-
-void Lowerer::lowerAwaitSuspends(Function &F) {
-  SmallVector<CoroAwaitSuspendInst *, 4> AwaitSuspends;
-
-  for (Instruction &I : llvm::make_early_inc_range(instructions(F))) {
-    auto *CB = dyn_cast<CallBase>(&I);
-    if (!CB)
-      continue;
-
-    if (auto *AWS = dyn_cast<CoroAwaitSuspendInst>(CB))
-      AwaitSuspends.push_back(AWS);
-  }
-
-  for (auto *AWS : AwaitSuspends)
-    lowerAwaitSuspend(AWS);
-}
-
 /// A little helper class for building
 class CoroCloner {
 public:
@@ -229,6 +167,45 @@ class CoroCloner {
 
 } // end anonymous namespace
 
+// FIXME:
+// Lower the intrinisc earlier if coroutine frame doesn't escape
+static void lowerAwaitSuspend(IRBuilder<> &Builder, CoroAwaitSuspendInst *CB) {
+  auto Helper = CB->getHelperFunction();
+  auto Awaiter = CB->getAwaiter();
+  auto FramePtr = CB->getFrame();
+
+  Builder.SetInsertPoint(CB);
+
+  CallBase *NewCall = nullptr;
+  if (auto Invoke = dyn_cast<InvokeInst>(CB)) {
+    auto HelperInvoke =
+        Builder.CreateInvoke(Helper, Invoke->getNormalDest(),
+                             Invoke->getUnwindDest(), {Awaiter, FramePtr});
+
+    HelperInvoke->setCallingConv(Invoke->getCallingConv());
+    std::copy(Invoke->bundle_op_info_begin(), Invoke->bundle_op_info_end(),
+              HelperInvoke->bundle_op_info_begin());
+    AttributeList NewAttributes =
+        Invoke->getAttributes().removeParamAttributes(Invoke->getContext(), 2);
+    HelperInvoke->setAttributes(NewAttributes);
+    HelperInvoke->setDebugLoc(Invoke->getDebugLoc());
+    NewCall = HelperInvoke;
+  } else if (auto Call = dyn_cast<CallInst>(CB)) {
+    auto HelperCall = Builder.CreateCall(Helper, {Awaiter, FramePtr});
+
+    AttributeList NewAttributes =
+        Call->getAttributes().removeParamAttributes(Call->getContext(), 2);
+    HelperCall->setAttributes(NewAttributes);
+    HelperCall->setDebugLoc(Call->getDebugLoc());
+    NewCall = HelperCall;
+  } else {
+    llvm_unreachable("Unexpected coro_await_suspend invocation method");
+  }
+
+  CB->replaceAllUsesWith(NewCall);
+  CB->eraseFromParent();
+}
+
 static void maybeFreeRetconStorage(IRBuilder<> &Builder,
                                    const coro::Shape &Shape, Value *FramePtr,
                                    CallGraph *CG) {
@@ -1603,8 +1580,9 @@ struct SwitchCoroutineSplitter {
 
     createResumeEntryBlock(F, Shape);
 
-    Lowerer lowerer(M);
-    lowerer.lowerAwaitSuspends(F);
+    IRBuilder<> Builder(M.getContext());
+    for (auto *AWS : Shape.CoroAwaitSuspends)
+      lowerAwaitSuspend(Builder, AWS);
 
     auto *ResumeClone =
         createClone(F, ".resume", Shape, CoroCloner::Kind::SwitchResume);
diff --git a/llvm/lib/Transforms/Coroutines/Coroutines.cpp b/llvm/lib/Transforms/Coroutines/Coroutines.cpp
index f8f3edd9292a15d..539354506c443e0 100644
--- a/llvm/lib/Transforms/Coroutines/Coroutines.cpp
+++ b/llvm/lib/Transforms/Coroutines/Coroutines.cpp
@@ -67,9 +67,9 @@ static const char *const CoroIntrinsics[] = {
     "llvm.coro.async.resume",
     "llvm.coro.async.size.replace",
     "llvm.coro.async.store_resume",
-    "llvm.coro.await.suspend.void",
     "llvm.coro.await.suspend.bool",
     "llvm.coro.await.suspend.handle",
+    "llvm.coro.await.suspend.void",
     "llvm.coro.begin",
     "llvm.coro.destroy",
     "llvm.coro.done",
@@ -258,6 +258,11 @@ void coro::Shape::buildFrom(Function &F) {
           }
         }
         break;
+      case Intrinsic::coro_await_suspend_void:
+      case Intrinsic::coro_await_suspend_bool:
+      case Intrinsic::coro_await_suspend_handle:
+        CoroAwaitSuspends.push_back(cast<CoroAwaitSuspendInst>(II));
+        break;
       }
     }
   }

>From edac452b85db4328a78c0b16df3e5a16086a8041 Mon Sep 17 00:00:00 2001
From: fpasserby <125797601+fpasserby at users.noreply.github.com>
Date: Sun, 4 Feb 2024 18:14:50 +0100
Subject: [PATCH 13/17] Rename helper -> wrapper

---
 clang/lib/CodeGen/CGCoroutine.cpp            | 38 ++++++++---------
 clang/lib/CodeGen/CodeGenFunction.h          | 16 +++----
 clang/test/CodeGenCoroutines/coro-await.cpp  | 18 ++++----
 clang/test/CodeGenCoroutines/coro-dwarf.cpp  |  4 +-
 clang/test/CodeGenCoroutines/pr65054.cpp     |  2 +-
 llvm/docs/Coroutines.rst                     | 12 +++---
 llvm/lib/Transforms/Coroutines/CoroInstr.h   |  6 +--
 llvm/lib/Transforms/Coroutines/CoroSplit.cpp | 44 ++++++++++----------
 8 files changed, 70 insertions(+), 70 deletions(-)

diff --git a/clang/lib/CodeGen/CGCoroutine.cpp b/clang/lib/CodeGen/CGCoroutine.cpp
index 6f5a19af0b37a11..6c8c8553b87ed6a 100644
--- a/clang/lib/CodeGen/CGCoroutine.cpp
+++ b/clang/lib/CodeGen/CGCoroutine.cpp
@@ -182,7 +182,7 @@ static bool AwaitSuspendStmtCanThrow(const Stmt *S) {
 //   auto && x = CommonExpr();
 //   if (!x.await_ready()) {
 //      llvm_coro_save();
-//      llvm_coro_await_suspend(&x, frame, helper) (*)
+//      llvm_coro_await_suspend(&x, frame, wrapper) (*)
 //      llvm_coro_suspend(); (**)
 //   }
 //   x.await_resume();
@@ -190,14 +190,14 @@ static bool AwaitSuspendStmtCanThrow(const Stmt *S) {
 // where the result of the entire expression is the result of x.await_resume()
 //
 //   (*) llvm_coro_await_suspend_{void, bool, handle} is lowered to
-//      helper(&x, frame) when it's certain not to interfere with
+//      wrapper(&x, frame) when it's certain not to interfere with
 //      coroutine transform. await_suspend expression is
 //      asynchronous to the coroutine body and not all analyses
 //      and transformations can handle it correctly at the moment.
 //
-//      Helper function encapsulates x.await_suspend(...) call and looks like:
+//      Wrapper function encapsulates x.await_suspend(...) call and looks like:
 //
-//      auto __await_suspend_helper(auto& awaiter, void* frame) {
+//      auto __await_suspend_wrapper(auto& awaiter, void* frame) {
 //        std::coroutine_handle<> handle(frame);
 //        return awaiter.await_suspend(handle);
 //      }
@@ -253,7 +253,7 @@ static LValueOrRValue emitSuspendExpression(CodeGenFunction &CGF, CGCoroData &Co
   const auto AwaitSuspendCanThrow =
       AwaitSuspendStmtCanThrow(S.getSuspendExpr());
 
-  auto SuspendHelper = CodeGenFunction(CGF.CGM).generateAwaitSuspendHelper(
+  auto SuspendWrapper = CodeGenFunction(CGF.CGM).generateAwaitSuspendWrapper(
       CGF.CurFn->getName(), Prefix, S);
 
   llvm::CallBase *SuspendRet = nullptr;
@@ -263,12 +263,12 @@ static LValueOrRValue emitSuspendExpression(CodeGenFunction &CGF, CGCoroData &Co
   assert(CGF.CurCoro.Data && CGF.CurCoro.Data->CoroBegin &&
          "expected to be called in coroutine context");
 
-  SmallVector<llvm::Value *, 3> SuspendHelperCallArgs;
-  SuspendHelperCallArgs.push_back(
+  SmallVector<llvm::Value *, 3> SuspendIntrinsicCallArgs;
+  SuspendIntrinsicCallArgs.push_back(
       CGF.getOrCreateOpaqueLValueMapping(S.getOpaqueValue()).getPointer(CGF));
 
-  SuspendHelperCallArgs.push_back(CGF.CurCoro.Data->CoroBegin);
-  SuspendHelperCallArgs.push_back(SuspendHelper);
+  SuspendIntrinsicCallArgs.push_back(CGF.CurCoro.Data->CoroBegin);
+  SuspendIntrinsicCallArgs.push_back(SuspendWrapper);
 
   const auto SuspendReturnType = S.getSuspendReturnType();
   llvm::Intrinsic::ID IID;
@@ -289,10 +289,10 @@ static LValueOrRValue emitSuspendExpression(CodeGenFunction &CGF, CGCoroData &Co
   // FIXME: add call attributes?
   if (AwaitSuspendCanThrow)
     SuspendRet =
-        CGF.EmitCallOrInvoke(AwaitSuspendIntrinsic, SuspendHelperCallArgs);
+        CGF.EmitCallOrInvoke(AwaitSuspendIntrinsic, SuspendIntrinsicCallArgs);
   else
     SuspendRet = CGF.EmitNounwindRuntimeCall(AwaitSuspendIntrinsic,
-                                             SuspendHelperCallArgs);
+                                             SuspendIntrinsicCallArgs);
 
   assert(SuspendRet);
   CGF.CurCoro.InSuspendBlock = false;
@@ -415,10 +415,10 @@ static QualType getCoroutineSuspendExprReturnType(const ASTContext &Ctx,
 #endif
 
 llvm::Function *
-CodeGenFunction::generateAwaitSuspendHelper(Twine const &CoroName,
-                                            Twine const &SuspendPointName,
-                                            CoroutineSuspendExpr const &S) {
-  std::string FuncName = "__await_suspend_helper_";
+CodeGenFunction::generateAwaitSuspendWrapper(Twine const &CoroName,
+                                             Twine const &SuspendPointName,
+                                             CoroutineSuspendExpr const &S) {
+  std::string FuncName = "__await_suspend_wrapper_";
   FuncName += CoroName.str();
   FuncName += '_';
   FuncName += SuspendPointName.str();
@@ -456,7 +456,7 @@ CodeGenFunction::generateAwaitSuspendHelper(Twine const &CoroName,
   auto AwaiterLValue =
       MakeNaturalAlignAddrLValue(AwaiterPtr, AwaiterDecl.getType());
 
-  CurAwaitSuspendHelper.FramePtr =
+  CurAwaitSuspendWrapper.FramePtr =
       Builder.CreateLoad(GetAddrOfLocalVar(&FrameDecl));
 
   // FIXME: mark as aliasing with frame?
@@ -474,7 +474,7 @@ CodeGenFunction::generateAwaitSuspendHelper(Twine const &CoroName,
     Builder.CreateStore(SuspendRet, ReturnValue);
   }
 
-  CurAwaitSuspendHelper.FramePtr = nullptr;
+  CurAwaitSuspendWrapper.FramePtr = nullptr;
   FinishFunction();
   return Fn;
 }
@@ -976,8 +976,8 @@ RValue CodeGenFunction::EmitCoroutineIntrinsic(const CallExpr *E,
       return RValue::get(CurCoro.Data->CoroBegin);
     }
 
-    if (CurAwaitSuspendHelper.FramePtr) {
-      return RValue::get(CurAwaitSuspendHelper.FramePtr);
+    if (CurAwaitSuspendWrapper.FramePtr) {
+      return RValue::get(CurAwaitSuspendWrapper.FramePtr);
     }
 
     CGM.Error(E->getBeginLoc(), "this builtin expect that __builtin_coro_begin "
diff --git a/clang/lib/CodeGen/CodeGenFunction.h b/clang/lib/CodeGen/CodeGenFunction.h
index 243edd58a7e3541..6c825a302913df7 100644
--- a/clang/lib/CodeGen/CodeGenFunction.h
+++ b/clang/lib/CodeGen/CodeGenFunction.h
@@ -352,24 +352,24 @@ class CodeGenFunction : public CodeGenTypeCache {
     return isCoroutine() && CurCoro.InSuspendBlock;
   }
 
-  // Holds FramePtr for await_suspend helper generation,
+  // Holds FramePtr for await_suspend wrapper generation,
   // so that __builtin_coro_frame call can be lowered
   // directly to value of its second argument
-  struct AwaitSuspendHelperInfo {
+  struct AwaitSuspendWrapperInfo {
     llvm::Value *FramePtr = nullptr;
   };
-  AwaitSuspendHelperInfo CurAwaitSuspendHelper;
+  AwaitSuspendWrapperInfo CurAwaitSuspendWrapper;
 
-  // Generates helper function for `llvm.coro.await.suspend.*` intrinisics.
+  // Generates wrapper function for `llvm.coro.await.suspend.*` intrinisics.
   // It encapsulates SuspendExpr in a function, to separate it's body
   // from the main coroutine to avoid miscompilations. Intrinisic
   // is lowered to this function call in CoroSplit pass
   // Function signature is:
-  // <type> __await_suspend_helper_<name>(ptr %awaiter, ptr %hdl)
+  // <type> __await_suspend_wrapper_<name>(ptr %awaiter, ptr %hdl)
   // where type is one of (void, i1, ptr)
-  llvm::Function *generateAwaitSuspendHelper(Twine const &CoroName,
-                                             Twine const &SuspendPointName,
-                                             CoroutineSuspendExpr const &S);
+  llvm::Function *generateAwaitSuspendWrapper(Twine const &CoroName,
+                                              Twine const &SuspendPointName,
+                                              CoroutineSuspendExpr const &S);
 
   /// CurGD - The GlobalDecl for the current function being compiled.
   GlobalDecl CurGD;
diff --git a/clang/test/CodeGenCoroutines/coro-await.cpp b/clang/test/CodeGenCoroutines/coro-await.cpp
index 377f14e8933cdc4..75851d8805bb6e3 100644
--- a/clang/test/CodeGenCoroutines/coro-await.cpp
+++ b/clang/test/CodeGenCoroutines/coro-await.cpp
@@ -73,7 +73,7 @@ extern "C" void f0() {
   // ---------------------------
   // Call coro.await.suspend
   // ---------------------------
-  // CHECK-NEXT: call void @llvm.coro.await.suspend.void(ptr %[[AWAITABLE]], ptr %[[FRAME]], ptr @__await_suspend_helper_f0_await)
+  // CHECK-NEXT: call void @llvm.coro.await.suspend.void(ptr %[[AWAITABLE]], ptr %[[FRAME]], ptr @__await_suspend_wrapper_f0_await)
   // -------------------------
   // Generate a suspend point:
   // -------------------------
@@ -99,8 +99,8 @@ extern "C" void f0() {
   // CHECK: %[[FINALSP_ID:.+]] = call token @llvm.coro.save(
   // CHECK: call i8 @llvm.coro.suspend(token %[[FINALSP_ID]], i1 true)
 
-  // await suspend helper
-  // CHECK: define{{.*}} @__await_suspend_helper_f0_await(ptr {{[^,]*}} %[[AWAITABLE_ARG:.+]], ptr {{[^,]*}} %[[FRAME_ARG:.+]])
+  // Await suspend wrapper
+  // CHECK: define{{.*}} @__await_suspend_wrapper_f0_await(ptr {{[^,]*}} %[[AWAITABLE_ARG:.+]], ptr {{[^,]*}} %[[FRAME_ARG:.+]])
   // CHECK: store ptr %[[AWAITABLE_ARG]], ptr %[[AWAITABLE_TMP:.+]],
   // CHECK: store ptr %[[FRAME_ARG]], ptr %[[FRAME_TMP:.+]],
   // CHECK: %[[AWAITABLE:.+]] = load ptr, ptr %[[AWAITABLE_TMP]]
@@ -149,7 +149,7 @@ extern "C" void f1(int) {
   // ---------------------------
   // Call coro.await.suspend
   // ---------------------------
-  // CHECK-NEXT: %[[YES:.+]] = call i1 @llvm.coro.await.suspend.bool(ptr %[[AWAITABLE]], ptr %[[FRAME]], ptr @__await_suspend_helper_f1_yield)
+  // CHECK-NEXT: %[[YES:.+]] = call i1 @llvm.coro.await.suspend.bool(ptr %[[AWAITABLE]], ptr %[[FRAME]], ptr @__await_suspend_wrapper_f1_yield)
   // -------------------------------------------
   // See if await_suspend decided not to suspend
   // -------------------------------------------
@@ -161,8 +161,8 @@ extern "C" void f1(int) {
   // CHECK: [[READY_BB]]:
   // CHECK:     call void @_ZN13suspend_maybe12await_resumeEv(ptr {{[^,]*}} %[[AWAITABLE]])
 
-  // Await suspend helper
-  // CHECK: define {{.*}} i1 @__await_suspend_helper_f1_yield(ptr {{[^,]*}} %[[AWAITABLE_ARG:.+]], ptr {{[^,]*}} %[[FRAME_ARG:.+]])
+  // Await suspend wrapper
+  // CHECK: define {{.*}} i1 @__await_suspend_wrapper_f1_yield(ptr {{[^,]*}} %[[AWAITABLE_ARG:.+]], ptr {{[^,]*}} %[[FRAME_ARG:.+]])
   // CHECK: store ptr %[[AWAITABLE_ARG]], ptr %[[AWAITABLE_TMP:.+]],
   // CHECK: store ptr %[[FRAME_ARG]], ptr %[[FRAME_TMP:.+]],
   // CHECK: %[[AWAITABLE:.+]] = load ptr, ptr %[[AWAITABLE_TMP]]
@@ -370,7 +370,7 @@ extern "C" void TestTailcall() {
   // ---------------------------
   // Call coro.await.suspend
   // ---------------------------
-  // CHECK-NEXT: %[[RESUMED:.+]] = call ptr @llvm.coro.await.suspend.handle(ptr %[[AWAITABLE]], ptr %[[FRAME]], ptr @__await_suspend_helper_TestTailcall_await)
+  // CHECK-NEXT: %[[RESUMED:.+]] = call ptr @llvm.coro.await.suspend.handle(ptr %[[AWAITABLE]], ptr %[[FRAME]], ptr @__await_suspend_wrapper_TestTailcall_await)
   // CHECK-NEXT: call void @llvm.coro.resume(ptr %[[RESUMED]])
   // CHECK-NEXT: %[[OUTCOME:.+]] = call i8 @llvm.coro.suspend(token %[[SUSPEND_ID]], i1 false)
   // CHECK-NEXT: switch i8 %[[OUTCOME]], label %[[RET_BB:.+]] [
@@ -378,8 +378,8 @@ extern "C" void TestTailcall() {
   // CHECK-NEXT:   i8 1, label %{{.+}}
   // CHECK-NEXT: ]
 
-  // Await suspend helper
-  // CHECK: define {{.*}} ptr @__await_suspend_helper_TestTailcall_await(ptr {{[^,]*}} %[[AWAITABLE_ARG:.+]], ptr {{[^,]*}} %[[FRAME_ARG:.+]])
+  // Await suspend wrapper
+  // CHECK: define {{.*}} ptr @__await_suspend_wrapper_TestTailcall_await(ptr {{[^,]*}} %[[AWAITABLE_ARG:.+]], ptr {{[^,]*}} %[[FRAME_ARG:.+]])
   // CHECK: store ptr %[[AWAITABLE_ARG]], ptr %[[AWAITABLE_TMP:.+]],
   // CHECK: store ptr %[[FRAME_ARG]], ptr %[[FRAME_TMP:.+]],
   // CHECK: %[[AWAITABLE:.+]] = load ptr, ptr %[[AWAITABLE_TMP]]
diff --git a/clang/test/CodeGenCoroutines/coro-dwarf.cpp b/clang/test/CodeGenCoroutines/coro-dwarf.cpp
index 7c44d1d4bef6568..2c9c827e6753d68 100644
--- a/clang/test/CodeGenCoroutines/coro-dwarf.cpp
+++ b/clang/test/CodeGenCoroutines/coro-dwarf.cpp
@@ -71,14 +71,14 @@ void f_coro(int val, MoveOnly moParam, MoveAndCopy mcParam) {
 // CHECK: !{{[0-9]+}} = !DILocalVariable(name: "mcParam", arg: 3, scope: ![[SP]], file: !{{[0-9]+}}, line: {{[0-9]+}}, type: !{{[0-9]+}})
 // CHECK: !{{[0-9]+}} = !DILocalVariable(name: "__promise",
 
-// CHECK: !{{[0-9]+}} = distinct !DISubprogram(linkageName: "__await_suspend_helper__Z6f_coroi8MoveOnly11MoveAndCopy_init"
+// CHECK: !{{[0-9]+}} = distinct !DISubprogram(linkageName: "__await_suspend_wrapper__Z6f_coroi8MoveOnly11MoveAndCopy_init"
 // CHECK-NEXT: !{{[0-9]+}} = !DIFile
 // CHECK-NEXT: !{{[0-9]+}} = !DISubroutineType
 // CHECK-NEXT: !{{[0-9]+}} = !DILocalVariable(arg: 1,
 // CHECK-NEXT: !{{[0-9]+}} = !DILocation
 // CHECK-NEXT: !{{[0-9]+}} = !DILocalVariable(arg: 2,
 
-// CHECK: !{{[0-9]+}} = distinct !DISubprogram(linkageName: "__await_suspend_helper__Z6f_coroi8MoveOnly11MoveAndCopy_final"
+// CHECK: !{{[0-9]+}} = distinct !DISubprogram(linkageName: "__await_suspend_wrapper__Z6f_coroi8MoveOnly11MoveAndCopy_final"
 // CHECK-NEXT: !{{[0-9]+}} = !DILocalVariable(arg: 1,
 // CHECK-NEXT: !{{[0-9]+}} = !DILocation
 // CHECK-NEXT: !{{[0-9]+}} = !DILocalVariable(arg: 2,
diff --git a/clang/test/CodeGenCoroutines/pr65054.cpp b/clang/test/CodeGenCoroutines/pr65054.cpp
index 55bb52233394b1f..506990163e9b8d1 100644
--- a/clang/test/CodeGenCoroutines/pr65054.cpp
+++ b/clang/test/CodeGenCoroutines/pr65054.cpp
@@ -55,6 +55,6 @@ MyTask FooBar() {
 // FRONTEND: attributes #[[address_attr]] = {{.*}}alwaysinline
 
 // CHECK-O0: define{{.*}}@_Z6FooBarv.resume
-// CHECK-O0: call{{.*}}@__await_suspend_helper__Z6FooBarv_await(
+// CHECK-O0: call{{.*}}@__await_suspend_wrapper__Z6FooBarv_await(
 // CHECK-O0-NOT: store
 // CHECK-O0: ret void
diff --git a/llvm/docs/Coroutines.rst b/llvm/docs/Coroutines.rst
index a175c7de24ed601..83369d93c309a79 100644
--- a/llvm/docs/Coroutines.rst
+++ b/llvm/docs/Coroutines.rst
@@ -1780,7 +1780,7 @@ The first argument is a pointer to `awaiter` object.
 
 The second argument is a pointer to the current coroutine's frame.
 
-The third argument is a pointer to the helper function encapsulating
+The third argument is a pointer to the wrapper function encapsulating
 `await-suspend` logic. Its signature must be
 
 .. code-block:: llvm
@@ -1819,7 +1819,7 @@ Example:
     %suspend = call i8 @llvm.coro.suspend(token %save, i1 false)   
     ...
 
-  ; helper function example
+  ; wrapper function example
   define void @await_suspend_function(ptr %awaiter, ptr %hdl)
     entry:
       %hdl.arg = ... ; construct std::coroutine_handle from %hdl
@@ -1862,7 +1862,7 @@ The first argument is a pointer to `awaiter` object.
 
 The second argument is a pointer to the current coroutine's frame.
 
-The third argument is a pointer to the helper function encapsulating
+The third argument is a pointer to the wrapper function encapsulating
 `await-suspend` logic. Its signature must be
 
 .. code-block:: llvm
@@ -1909,7 +1909,7 @@ Example:
     br i1 %resume, %await.suspend.bool, %await.ready
     ...
 
-  ; helper function example
+  ; wrapper function example
   define i1 @await_suspend_function(ptr %awaiter, ptr %hdl)
     entry:
       %hdl.arg = ... ; construct std::coroutine_handle from %hdl
@@ -1952,7 +1952,7 @@ The first argument is a pointer to `awaiter` object.
 
 The second argument is a pointer to the current coroutine's frame.
 
-The third argument is a pointer to the helper function encapsulating
+The third argument is a pointer to the wrapper function encapsulating
 `await-suspend` logic. Its signature must be
 
 .. code-block:: llvm
@@ -1996,7 +1996,7 @@ Example:
     %suspend = call i8 @llvm.coro.suspend(token %save, i1 false)
     ...
 
-  ; helper function example
+  ; wrapper function example
   define ptr @await_suspend_function(ptr %awaiter, ptr %hdl)
     entry:
       %hdl.arg = ... ; construct std::coroutine_handle from %hdl
diff --git a/llvm/lib/Transforms/Coroutines/CoroInstr.h b/llvm/lib/Transforms/Coroutines/CoroInstr.h
index 63146cf6742aa6b..8d7f1b19d9ed0cb 100644
--- a/llvm/lib/Transforms/Coroutines/CoroInstr.h
+++ b/llvm/lib/Transforms/Coroutines/CoroInstr.h
@@ -80,15 +80,15 @@ class LLVM_LIBRARY_VISIBILITY CoroAllocInst : public IntrinsicInst {
 
 /// This represents the llvm.coro.await.suspend instruction.
 class LLVM_LIBRARY_VISIBILITY CoroAwaitSuspendInst : public CallBase {
-  enum { AwaiterArg, FrameArg, HelperArg };
+  enum { AwaiterArg, FrameArg, WrapperArg };
 
 public:
   Value *getAwaiter() const { return getArgOperand(AwaiterArg); }
 
   Value *getFrame() const { return getArgOperand(FrameArg); }
 
-  Function *getHelperFunction() const {
-    return cast<Function>(getArgOperand(HelperArg));
+  Function *getWrapperFunction() const {
+    return cast<Function>(getArgOperand(WrapperArg));
   }
 
   // Methods to support type inquiry through isa, cast, and dyn_cast:
diff --git a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
index 8eef64658ed7214..c9ab2385f92dc46 100644
--- a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
@@ -170,7 +170,7 @@ class CoroCloner {
 // FIXME:
 // Lower the intrinisc earlier if coroutine frame doesn't escape
 static void lowerAwaitSuspend(IRBuilder<> &Builder, CoroAwaitSuspendInst *CB) {
-  auto Helper = CB->getHelperFunction();
+  auto Wrapper = CB->getWrapperFunction();
   auto Awaiter = CB->getAwaiter();
   auto FramePtr = CB->getFrame();
 
@@ -178,26 +178,26 @@ static void lowerAwaitSuspend(IRBuilder<> &Builder, CoroAwaitSuspendInst *CB) {
 
   CallBase *NewCall = nullptr;
   if (auto Invoke = dyn_cast<InvokeInst>(CB)) {
-    auto HelperInvoke =
-        Builder.CreateInvoke(Helper, Invoke->getNormalDest(),
+    auto WrapperInvoke =
+        Builder.CreateInvoke(Wrapper, Invoke->getNormalDest(),
                              Invoke->getUnwindDest(), {Awaiter, FramePtr});
 
-    HelperInvoke->setCallingConv(Invoke->getCallingConv());
+    WrapperInvoke->setCallingConv(Invoke->getCallingConv());
     std::copy(Invoke->bundle_op_info_begin(), Invoke->bundle_op_info_end(),
-              HelperInvoke->bundle_op_info_begin());
+              WrapperInvoke->bundle_op_info_begin());
     AttributeList NewAttributes =
         Invoke->getAttributes().removeParamAttributes(Invoke->getContext(), 2);
-    HelperInvoke->setAttributes(NewAttributes);
-    HelperInvoke->setDebugLoc(Invoke->getDebugLoc());
-    NewCall = HelperInvoke;
+    WrapperInvoke->setAttributes(NewAttributes);
+    WrapperInvoke->setDebugLoc(Invoke->getDebugLoc());
+    NewCall = WrapperInvoke;
   } else if (auto Call = dyn_cast<CallInst>(CB)) {
-    auto HelperCall = Builder.CreateCall(Helper, {Awaiter, FramePtr});
+    auto WrapperCall = Builder.CreateCall(Wrapper, {Awaiter, FramePtr});
 
     AttributeList NewAttributes =
         Call->getAttributes().removeParamAttributes(Call->getContext(), 2);
-    HelperCall->setAttributes(NewAttributes);
-    HelperCall->setDebugLoc(Call->getDebugLoc());
-    NewCall = HelperCall;
+    WrapperCall->setAttributes(NewAttributes);
+    WrapperCall->setDebugLoc(Call->getDebugLoc());
+    NewCall = WrapperCall;
   } else {
     llvm_unreachable("Unexpected coro_await_suspend invocation method");
   }
@@ -1428,18 +1428,18 @@ static bool hasCallsBetween(Instruction *Save, Instruction *ResumeOrDestroy) {
   return false;
 }
 
-// Check if await-suspend helper is "simple".
+// Check if await-suspend wrapper is "simple".
 // The conditions are:
-// 1. The return result is exactly coroutine frame parameter, passed to helper
-// 2. There are no calls between any of the returns and helper entry that could
+// 1. The return result is exactly coroutine frame parameter, passed to wrapper
+// 2. There are no calls between any of the returns and wrapper entry that could
 // resume or destroy it
 // FIXME: perform more sophisiticated analysis?
-static bool isSimpleHelper(CoroAwaitSuspendInst *AWS) {
-  auto Helper = AWS->getHelperFunction();
+static bool isSimpleWrapper(CoroAwaitSuspendInst *AWS) {
+  auto Wrapper = AWS->getWrapperFunction();
 
   SmallVector<ReturnInst *, 4> Rets;
 
-  for (auto &BB : *Helper) {
+  for (auto &BB : *Wrapper) {
     if (BB.empty())
       continue;
     auto terminator = BB.getTerminator();
@@ -1451,8 +1451,8 @@ static bool isSimpleHelper(CoroAwaitSuspendInst *AWS) {
 
   // FIXME: get rid of magical constant
   for (auto Ret : Rets)
-    if (Ret->getReturnValue() != Helper->getArg(1) ||
-        hasCallsBetween(Helper->getEntryBlock().getFirstNonPHI(), Ret))
+    if (Ret->getReturnValue() != Wrapper->getArg(1) ||
+        hasCallsBetween(Wrapper->getEntryBlock().getFirstNonPHI(), Ret))
       return false;
 
   return true;
@@ -1484,13 +1484,13 @@ static bool simplifySuspendPoint(CoroSuspendInst *Suspend,
   auto Frame = SubFn->getFrame();
 
   // Check that frame directly always refers to the current coroutine,
-  // either directly or via helper
+  // either directly or via wrapper
   if (Frame != CoroBegin) {
     auto *AWS = dyn_cast<CoroAwaitSuspendInst>(Frame);
     if (!AWS)
       return false;
 
-    if (AWS->getFrame() != CoroBegin || !isSimpleHelper(AWS))
+    if (AWS->getFrame() != CoroBegin || !isSimpleWrapper(AWS))
       return false;
   }
 

>From 8261a969ce6f6b1a23637b8bc6df376593851d5b Mon Sep 17 00:00:00 2001
From: fpasserby <125797601+fpasserby at users.noreply.github.com>
Date: Sun, 4 Feb 2024 18:17:18 +0100
Subject: [PATCH 14/17] Remove unnecessary comments

---
 clang/test/CodeGenCoroutines/coro-symmetric-transfer-01.cpp | 1 -
 clang/test/CodeGenCoroutines/pr59181.cpp                    | 2 --
 llvm/lib/Transforms/Coroutines/CoroSplit.cpp                | 2 ++
 3 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/clang/test/CodeGenCoroutines/coro-symmetric-transfer-01.cpp b/clang/test/CodeGenCoroutines/coro-symmetric-transfer-01.cpp
index f1870b8eceff99a..da30e12c63cffb2 100644
--- a/clang/test/CodeGenCoroutines/coro-symmetric-transfer-01.cpp
+++ b/clang/test/CodeGenCoroutines/coro-symmetric-transfer-01.cpp
@@ -47,7 +47,6 @@ detached_task foo() {
   co_return;
 }
 
-// FIXME: is this test needed now?
 // check that the lifetime of the coroutine handle used to obtain the address is contained within single basic block, and hence does not live across suspension points.
 // CHECK-LABEL: final.suspend:
 // CHECK:         %{{.+}} = call token @llvm.coro.save(ptr null)
diff --git a/clang/test/CodeGenCoroutines/pr59181.cpp b/clang/test/CodeGenCoroutines/pr59181.cpp
index a3b050dacf0ba0e..21e784e0031de33 100644
--- a/clang/test/CodeGenCoroutines/pr59181.cpp
+++ b/clang/test/CodeGenCoroutines/pr59181.cpp
@@ -48,8 +48,6 @@ void foo() {
   bar(true);
 }
 
-// FIXME: the test doesn't seem to be relevant anymore,
-// because objects that require cleanup are no more emitted in the suspend block
 // CHECK: cleanup.cont:{{.*}}
 // CHECK-NEXT: load i8
 // CHECK-NEXT: trunc
diff --git a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
index c9ab2385f92dc46..89dca96f4ca2bec 100644
--- a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
@@ -169,6 +169,8 @@ class CoroCloner {
 
 // FIXME:
 // Lower the intrinisc earlier if coroutine frame doesn't escape
+// and it is known that other transformations, for example, sanitizers
+// won't lead to incorrect code.
 static void lowerAwaitSuspend(IRBuilder<> &Builder, CoroAwaitSuspendInst *CB) {
   auto Wrapper = CB->getWrapperFunction();
   auto Awaiter = CB->getAwaiter();

>From 20a7dfb5e61d5986fadacc1cc1b52fd72facc533 Mon Sep 17 00:00:00 2001
From: fpasserby <125797601+fpasserby at users.noreply.github.com>
Date: Sun, 4 Feb 2024 22:57:24 +0100
Subject: [PATCH 15/17] Add await suspend lowering tests

---
 llvm/lib/Transforms/Coroutines/CoroSplit.cpp  |  22 ++--
 llvm/lib/Transforms/Coroutines/Coroutines.cpp |  11 +-
 .../coro-await-suspend-lower-invoke.ll        | 123 ++++++++++++++++++
 .../Coroutines/coro-await-suspend-lower.ll    |  96 ++++++++++++++
 .../Coroutines/coro-await-suspend-simplify.ll |  92 +++++++++++++
 5 files changed, 328 insertions(+), 16 deletions(-)
 create mode 100644 llvm/test/Transforms/Coroutines/coro-await-suspend-lower-invoke.ll
 create mode 100644 llvm/test/Transforms/Coroutines/coro-await-suspend-lower.ll
 create mode 100644 llvm/test/Transforms/Coroutines/coro-await-suspend-simplify.ll

diff --git a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
index 89dca96f4ca2bec..3019dee47b94ef8 100644
--- a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
@@ -1412,10 +1412,10 @@ static bool hasCallsBetween(Instruction *Save, Instruction *ResumeOrDestroy) {
   auto *ResumeOrDestroyBB = ResumeOrDestroy->getParent();
 
   if (SaveBB == ResumeOrDestroyBB)
-    return hasCallsInBlockBetween(Save->getNextNode(), ResumeOrDestroy);
+    return hasCallsInBlockBetween(Save, ResumeOrDestroy);
 
   // Any calls from Save to the end of the block?
-  if (hasCallsInBlockBetween(Save->getNextNode(), nullptr))
+  if (hasCallsInBlockBetween(Save, nullptr))
     return true;
 
   // Any calls from begging of the block up to ResumeOrDestroy?
@@ -1437,7 +1437,10 @@ static bool hasCallsBetween(Instruction *Save, Instruction *ResumeOrDestroy) {
 // resume or destroy it
 // FIXME: perform more sophisiticated analysis?
 static bool isSimpleWrapper(CoroAwaitSuspendInst *AWS) {
-  auto Wrapper = AWS->getWrapperFunction();
+  auto *Wrapper = AWS->getWrapperFunction();
+
+  if (Wrapper->empty())
+    return false;
 
   SmallVector<ReturnInst *, 4> Rets;
 
@@ -1575,17 +1578,12 @@ static void simplifySuspendPoints(coro::Shape &Shape) {
 namespace {
 
 struct SwitchCoroutineSplitter {
-  static void split(Module &M, Function &F, coro::Shape &Shape,
+  static void split(Function &F, coro::Shape &Shape,
                     SmallVectorImpl<Function *> &Clones,
                     TargetTransformInfo &TTI) {
     assert(Shape.ABI == coro::ABI::Switch);
 
     createResumeEntryBlock(F, Shape);
-
-    IRBuilder<> Builder(M.getContext());
-    for (auto *AWS : Shape.CoroAwaitSuspends)
-      lowerAwaitSuspend(Builder, AWS);
-
     auto *ResumeClone =
         createClone(F, ".resume", Shape, CoroCloner::Kind::SwitchResume);
     auto *DestroyClone =
@@ -2122,6 +2120,10 @@ splitCoroutine(Module &M, Function &F, SmallVectorImpl<Function *> &Clones,
   buildCoroutineFrame(F, Shape, TTI, MaterializableCallback);
   replaceFrameSizeAndAlignment(Shape);
 
+  IRBuilder<> Builder(M.getContext());
+  for (auto *AWS : Shape.CoroAwaitSuspends)
+    lowerAwaitSuspend(Builder, AWS);
+
   // If there are no suspend points, no split required, just remove
   // the allocation and deallocation blocks, they are not needed.
   if (Shape.CoroSuspends.empty()) {
@@ -2129,7 +2131,7 @@ splitCoroutine(Module &M, Function &F, SmallVectorImpl<Function *> &Clones,
   } else {
     switch (Shape.ABI) {
     case coro::ABI::Switch:
-      SwitchCoroutineSplitter::split(M, F, Shape, Clones, TTI);
+      SwitchCoroutineSplitter::split(F, Shape, Clones, TTI);
       break;
     case coro::ABI::Async:
       splitAsyncCoroutine(F, Shape, Clones, TTI);
diff --git a/llvm/lib/Transforms/Coroutines/Coroutines.cpp b/llvm/lib/Transforms/Coroutines/Coroutines.cpp
index 539354506c443e0..ec7423c5c126f03 100644
--- a/llvm/lib/Transforms/Coroutines/Coroutines.cpp
+++ b/llvm/lib/Transforms/Coroutines/Coroutines.cpp
@@ -177,7 +177,11 @@ void coro::Shape::buildFrom(Function &F) {
   SmallVector<CoroSaveInst *, 2> UnusedCoroSaves;
 
   for (Instruction &I : instructions(F)) {
-    if (auto II = dyn_cast<IntrinsicInst>(&I)) {
+    // FIXME: coro_await_suspend_* are not proper `IntrinisicInst`s
+    // because they might be invoked
+    if (auto AWS = dyn_cast<CoroAwaitSuspendInst>(&I)) {
+      CoroAwaitSuspends.push_back(AWS);
+    } else if (auto II = dyn_cast<IntrinsicInst>(&I)) {
       switch (II->getIntrinsicID()) {
       default:
         continue;
@@ -258,11 +262,6 @@ void coro::Shape::buildFrom(Function &F) {
           }
         }
         break;
-      case Intrinsic::coro_await_suspend_void:
-      case Intrinsic::coro_await_suspend_bool:
-      case Intrinsic::coro_await_suspend_handle:
-        CoroAwaitSuspends.push_back(cast<CoroAwaitSuspendInst>(II));
-        break;
       }
     }
   }
diff --git a/llvm/test/Transforms/Coroutines/coro-await-suspend-lower-invoke.ll b/llvm/test/Transforms/Coroutines/coro-await-suspend-lower-invoke.ll
new file mode 100644
index 000000000000000..fbc4a2c006f84e5
--- /dev/null
+++ b/llvm/test/Transforms/Coroutines/coro-await-suspend-lower-invoke.ll
@@ -0,0 +1,123 @@
+; Tests that invoke <type> @llvm.coro.await.suspend lowers to invoke @helper
+; RUN: opt < %s -passes='module(coro-early),cgscc(coro-split),simplifycfg' -S | FileCheck %s
+
+%Awaiter = type {}
+
+; CHECK:     define {{[^@]*}} @f.resume(ptr {{[^%]*}} %[[HDL:.+]])
+; CHECK:       %[[AWAITER:.+]] = getelementptr inbounds %f.Frame, ptr %[[HDL]], i32 0, i32 0
+define void @f() presplitcoroutine personality i32 0 {
+entry:
+  %awaiter = alloca %Awaiter
+  %id = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null)
+  %size = call i32 @llvm.coro.size.i32()
+  %alloc = call ptr @malloc(i32 %size)
+  %hdl = call ptr @llvm.coro.begin(token %id, ptr %alloc)
+  ; Initial suspend so that all 3 await_suspend invocations are inside f.resume
+  %suspend.init = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %suspend.init, label %ret [
+    i8 0, label %step
+    i8 1, label %cleanup
+  ]
+
+; CHECK:        invoke void @await_suspend_wrapper_void(ptr %[[AWAITER]], ptr %[[HDL]])
+; CHECK-NEXT:       to label %[[STEP_CONT:[^ ]+]] unwind label %[[PAD:[^ ]+]]
+step:
+  %save = call token @llvm.coro.save(ptr null)
+  invoke void @llvm.coro.await.suspend.void(ptr %awaiter, ptr %hdl, ptr @await_suspend_wrapper_void)
+    to label %step.continue unwind label %pad
+
+; CHECK       [[STEP_CONT]]:
+step.continue:
+  %suspend = call i8 @llvm.coro.suspend(token %save, i1 false)
+  switch i8 %suspend, label %ret [
+    i8 0, label %step1
+    i8 1, label %cleanup
+  ]
+
+; CHECK:        %[[RESUME:.+]] = invoke i1 @await_suspend_wrapper_bool(ptr %[[AWAITER]], ptr %[[HDL]])
+; CHECK-NEXT:       to label %[[STEP1_CONT:[^ ]+]] unwind label %[[PAD]]
+step1:
+  %save1 = call token @llvm.coro.save(ptr null)
+  %resume.bool = invoke i1 @llvm.coro.await.suspend.bool(ptr %awaiter, ptr %hdl, ptr @await_suspend_wrapper_bool)
+    to label %step1.continue unwind label %pad
+
+; CHECK:      [[STEP1_CONT]]:
+; CHECK-NEXT:   br i1 %[[RESUME]], label %{{[^,]+}}, label %[[STEP2:.+]]
+step1.continue:
+  br i1 %resume.bool, label %suspend.cond, label %step2
+
+suspend.cond:
+  %suspend1 = call i8 @llvm.coro.suspend(token %save1, i1 false)
+  switch i8 %suspend1, label %ret [
+    i8 0, label %step2
+    i8 1, label %cleanup
+  ]
+
+; CHECK:      [[STEP2]]:
+; CHECK:        %[[NEXT_HDL:.+]] = invoke ptr @await_suspend_wrapper_handle(ptr %[[AWAITER]], ptr %[[HDL]])
+; CHECK-NEXT:       to label %[[STEP2_CONT:[^ ]+]] unwind label %[[PAD]]
+step2:
+  %save2 = call token @llvm.coro.save(ptr null)
+  %resume.handle = invoke ptr @llvm.coro.await.suspend.handle(ptr %awaiter, ptr %hdl, ptr @await_suspend_wrapper_handle)
+    to label %step2.continue unwind label %pad
+
+; CHECK:      [[STEP2_CONT]]:
+; CHECK-NEXT:   %[[NEXT_RESUME:.+]] = call ptr @llvm.coro.subfn.addr(ptr %[[NEXT_HDL]], i8 0)
+; CHECK-NEXT:   musttail call {{.*}} void %[[NEXT_RESUME]](ptr %[[NEXT_HDL]])
+step2.continue:
+  call void @llvm.coro.resume(ptr %resume.handle)
+  %suspend2 = call i8 @llvm.coro.suspend(token %save2, i1 false)
+  switch i8 %suspend2, label %ret [
+    i8 0, label %step3
+    i8 1, label %cleanup
+  ]
+
+step3:
+  br label %cleanup
+
+pad:
+  %lp = landingpad { ptr, i32 }
+          catch ptr null
+  %exn = extractvalue { ptr, i32 } %lp, 0
+  call ptr @__cxa_begin_catch(ptr %exn)
+  call void @__cxa_end_catch()
+  br label %cleanup
+
+cleanup:
+  %mem = call ptr @llvm.coro.free(token %id, ptr %hdl)
+  call void @free(ptr %mem)
+  br label %ret
+
+ret:
+  call i1 @llvm.coro.end(ptr %hdl, i1 0, token none)
+  ret void
+}
+
+; check that we were haven't accidentally went out of @f.resume body
+; CHECK-LABEL: @f.destroy(
+; CHECK-LABEL: @f.cleanup(
+
+declare void @await_suspend_wrapper_void(ptr, ptr)
+declare i1 @await_suspend_wrapper_bool(ptr, ptr)
+declare ptr @await_suspend_wrapper_handle(ptr, ptr)
+
+declare ptr @llvm.coro.free(token, ptr)
+declare i32 @llvm.coro.size.i32()
+declare i8  @llvm.coro.suspend(token, i1)
+declare void @llvm.coro.resume(ptr)
+declare void @llvm.coro.destroy(ptr)
+
+declare token @llvm.coro.id(i32, ptr, ptr, ptr)
+declare i1 @llvm.coro.alloc(token)
+declare ptr @llvm.coro.begin(token, ptr)
+declare void @llvm.coro.await.suspend.void(ptr, ptr, ptr)
+declare i1 @llvm.coro.await.suspend.bool(ptr, ptr, ptr)
+declare ptr @llvm.coro.await.suspend.handle(ptr, ptr, ptr)
+declare i1 @llvm.coro.end(ptr, i1, token)
+
+declare ptr @__cxa_begin_catch(ptr)
+declare void @use_val(i32)
+declare void @__cxa_end_catch()
+
+declare noalias ptr @malloc(i32)
+declare void @free(ptr)
diff --git a/llvm/test/Transforms/Coroutines/coro-await-suspend-lower.ll b/llvm/test/Transforms/Coroutines/coro-await-suspend-lower.ll
new file mode 100644
index 000000000000000..0f574c4acc26e7b
--- /dev/null
+++ b/llvm/test/Transforms/Coroutines/coro-await-suspend-lower.ll
@@ -0,0 +1,96 @@
+; Tests lowerings of different versions of coro.await.suspend
+; RUN: opt < %s -passes='module(coro-early),cgscc(coro-split),simplifycfg' -S | FileCheck %s
+
+%Awaiter = type {}
+
+; CHECK:     define {{[^@]*}} @f.resume(ptr {{[^%]*}} %[[HDL:.+]])
+; CHECK:       %[[AWAITER:.+]] = getelementptr inbounds %f.Frame, ptr %[[HDL]], i32 0, i32 0
+define void @f() presplitcoroutine {
+entry:
+  %awaiter = alloca %Awaiter
+  %id = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null)
+  %size = call i32 @llvm.coro.size.i32()
+  %alloc = call ptr @malloc(i32 %size)
+  %hdl = call ptr @llvm.coro.begin(token %id, ptr %alloc)
+  %suspend.init = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %suspend.init, label %ret [
+    i8 0, label %step
+    i8 1, label %cleanup
+  ]
+
+; CHECK:        call void @await_suspend_wrapper_void(ptr %[[AWAITER]], ptr %[[HDL]])
+; CHECK-NEXT:   br label %{{.*}}
+step:
+  %save = call token @llvm.coro.save(ptr null)
+  call void @llvm.coro.await.suspend.void(ptr %awaiter, ptr %hdl, ptr @await_suspend_wrapper_void)
+  %suspend = call i8 @llvm.coro.suspend(token %save, i1 false)
+  switch i8 %suspend, label %ret [
+    i8 0, label %step1
+    i8 1, label %cleanup
+  ]
+
+; CHECK:        %[[RESUME:.+]] = call i1 @await_suspend_wrapper_bool(ptr %[[AWAITER]], ptr %[[HDL]])
+; CHECK-NEXT:   br i1 %[[RESUME]], label %{{[^,]+}}, label %[[STEP2:.+]]
+step1:
+  %save1 = call token @llvm.coro.save(ptr null)
+  %resume.bool = call i1 @llvm.coro.await.suspend.bool(ptr %awaiter, ptr %hdl, ptr @await_suspend_wrapper_bool)
+  br i1 %resume.bool, label %suspend.cond, label %step2
+
+suspend.cond:
+  %suspend1 = call i8 @llvm.coro.suspend(token %save1, i1 false)
+  switch i8 %suspend1, label %ret [
+    i8 0, label %step2
+    i8 1, label %cleanup
+  ]
+
+; CHECK:      [[STEP2]]:
+; CHECK:        %[[NEXT_HDL:.+]] = call ptr @await_suspend_wrapper_handle(ptr %[[AWAITER]], ptr %[[HDL]])
+; CHECK-NEXT:   %[[CONT:.+]] = call ptr @llvm.coro.subfn.addr(ptr %[[NEXT_HDL]], i8 0)
+; CHECK-NEXT:   musttail call {{.*}} void %[[CONT]](ptr %[[NEXT_HDL]])
+step2:
+  %save2 = call token @llvm.coro.save(ptr null)
+  %resume.handle = call ptr @llvm.coro.await.suspend.handle(ptr %awaiter, ptr %hdl, ptr @await_suspend_wrapper_handle)
+  call void @llvm.coro.resume(ptr %resume.handle)
+  %suspend2 = call i8 @llvm.coro.suspend(token %save2, i1 false)
+  switch i8 %suspend2, label %ret [
+    i8 0, label %step3
+    i8 1, label %cleanup
+  ]
+
+step3:
+  br label %cleanup
+
+cleanup:
+  %mem = call ptr @llvm.coro.free(token %id, ptr %hdl)
+  call void @free(ptr %mem)
+  br label %ret
+
+ret:
+  call i1 @llvm.coro.end(ptr %hdl, i1 0, token none)
+  ret void
+}
+
+; check that we were haven't accidentally went out of @f.resume body
+; CHECK-LABEL: @f.destroy(
+; CHECK-LABEL: @f.cleanup(
+
+declare void @await_suspend_wrapper_void(ptr, ptr)
+declare i1 @await_suspend_wrapper_bool(ptr, ptr)
+declare ptr @await_suspend_wrapper_handle(ptr, ptr)
+
+declare ptr @llvm.coro.free(token, ptr)
+declare i32 @llvm.coro.size.i32()
+declare i8  @llvm.coro.suspend(token, i1)
+declare void @llvm.coro.resume(ptr)
+declare void @llvm.coro.destroy(ptr)
+
+declare token @llvm.coro.id(i32, ptr, ptr, ptr)
+declare i1 @llvm.coro.alloc(token)
+declare ptr @llvm.coro.begin(token, ptr)
+declare void @llvm.coro.await.suspend.void(ptr, ptr, ptr)
+declare i1 @llvm.coro.await.suspend.bool(ptr, ptr, ptr)
+declare ptr @llvm.coro.await.suspend.handle(ptr, ptr, ptr)
+declare i1 @llvm.coro.end(ptr, i1, token)
+
+declare noalias ptr @malloc(i32)
+declare void @free(ptr)
diff --git a/llvm/test/Transforms/Coroutines/coro-await-suspend-simplify.ll b/llvm/test/Transforms/Coroutines/coro-await-suspend-simplify.ll
new file mode 100644
index 000000000000000..8bd7ba6ceb96641
--- /dev/null
+++ b/llvm/test/Transforms/Coroutines/coro-await-suspend-simplify.ll
@@ -0,0 +1,92 @@
+; Tests that CoroSplit can succesfully see through await suspend wrapper
+; and remove redundant suspend points only if there are no calls in the wrapper
+; RUN: opt < %s -passes='module(coro-early),cgscc(coro-split),simplifycfg' -S | FileCheck %s
+
+; CHECK-LABEL: @f(
+; CHECK:        call ptr @await_suspend_wrapper(
+; CHECK-NEXT:   call void @free(
+define void @f() presplitcoroutine {
+entry:
+  %id = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null)
+  %size = call i32 @llvm.coro.size.i32()
+  %alloc = call ptr @malloc(i32 %size)
+  %hdl = call ptr @llvm.coro.begin(token %id, ptr %alloc)
+  %save = call token @llvm.coro.save(ptr null)
+  %hdl.resume = call ptr @llvm.coro.await.suspend.handle(ptr null, ptr %hdl, ptr @await_suspend_wrapper)
+  call void @llvm.coro.resume(ptr %hdl.resume)
+  %suspend = call i8 @llvm.coro.suspend(token %save, i1 false)
+  switch i8 %suspend, label %ret [
+    i8 0, label %resume
+    i8 1, label %cleanup
+  ]
+
+resume:
+  br label %cleanup
+
+cleanup:
+  %mem = call ptr @llvm.coro.free(token %id, ptr %hdl)
+  call void @free(ptr %mem)
+  br label %ret
+
+ret:
+  call i1 @llvm.coro.end(ptr %hdl, i1 0, token none)
+  ret void
+}
+
+define ptr @await_suspend_wrapper(ptr, ptr %hdl) {
+  ret ptr %hdl
+}
+
+; CHECK-LABEL: @f1(
+; CHECK:        %[[HDL:.+]] = call ptr @await_suspend_wrapper1(
+; CHECK-NEXT:   %[[FN:.+]] = call ptr @llvm.coro.subfn.addr(ptr %[[HDL]], i8 0)
+; CHECK-NEXT:   call {{.*}} %[[FN]](ptr %[[HDL]])
+define void @f1() presplitcoroutine {
+entry:
+  %id = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null)
+  %size = call i32 @llvm.coro.size.i32()
+  %alloc = call ptr @malloc(i32 %size)
+  %hdl = call ptr @llvm.coro.begin(token %id, ptr %alloc)
+  %save = call token @llvm.coro.save(ptr null)
+  %hdl.resume = call ptr @llvm.coro.await.suspend.handle(ptr null, ptr %hdl, ptr @await_suspend_wrapper1)
+  call void @llvm.coro.resume(ptr %hdl.resume)
+  %suspend = call i8 @llvm.coro.suspend(token %save, i1 false)
+  switch i8 %suspend, label %ret [
+    i8 0, label %resume
+    i8 1, label %cleanup
+  ]
+
+resume:
+  br label %cleanup
+
+cleanup:
+  %mem = call ptr @llvm.coro.free(token %id, ptr %hdl)
+  call void @free(ptr %mem)
+  br label %ret
+
+ret:
+  call i1 @llvm.coro.end(ptr %hdl, i1 0, token none)
+  ret void
+}
+
+define ptr @await_suspend_wrapper1(ptr, ptr %hdl) {
+  call void @external()
+  ret ptr %hdl
+}
+
+declare ptr @llvm.coro.free(token, ptr)
+declare i32 @llvm.coro.size.i32()
+declare i8  @llvm.coro.suspend(token, i1)
+declare void @llvm.coro.resume(ptr)
+declare void @llvm.coro.destroy(ptr)
+
+declare token @llvm.coro.id(i32, ptr, ptr, ptr)
+declare i1 @llvm.coro.alloc(token)
+declare ptr @llvm.coro.begin(token, ptr)
+declare ptr @llvm.coro.await.suspend.handle(ptr, ptr, ptr)
+declare i1 @llvm.coro.end(ptr, i1, token)
+
+declare void @external()
+
+declare noalias ptr @malloc(i32)
+declare void @free(ptr)

>From 1b6e5a56d4be1d174bc4f5c1568a0244ffccb8e0 Mon Sep 17 00:00:00 2001
From: fpasserby <125797601+fpasserby at users.noreply.github.com>
Date: Mon, 5 Feb 2024 00:05:16 +0100
Subject: [PATCH 16/17] Remove unnecessary changes from SemaCoroutine

---
 clang/lib/Sema/SemaCoroutine.cpp | 16 ++++++----------
 1 file changed, 6 insertions(+), 10 deletions(-)

diff --git a/clang/lib/Sema/SemaCoroutine.cpp b/clang/lib/Sema/SemaCoroutine.cpp
index f250c3f07b7836f..cfbf5791a3b999e 100644
--- a/clang/lib/Sema/SemaCoroutine.cpp
+++ b/clang/lib/Sema/SemaCoroutine.cpp
@@ -266,7 +266,7 @@ static ExprResult buildOperatorCoawaitCall(Sema &SemaRef, Scope *S,
 }
 
 static ExprResult buildCoroutineHandle(Sema &S, QualType PromiseType,
-                                       Expr *FramePtr, SourceLocation Loc) {
+                                       SourceLocation Loc) {
   QualType CoroHandleType = lookupCoroutineHandleType(S, PromiseType, Loc);
   if (CoroHandleType.isNull())
     return ExprError();
@@ -280,6 +280,9 @@ static ExprResult buildCoroutineHandle(Sema &S, QualType PromiseType,
     return ExprError();
   }
 
+  Expr *FramePtr =
+      S.BuildBuiltinCallExpr(Loc, Builtin::BI__builtin_coro_frame, {});
+
   CXXScopeSpec SS;
   ExprResult FromAddr =
       S.BuildDeclarationNameExpr(SS, Found, /*NeedsADL=*/false);
@@ -293,7 +296,6 @@ struct ReadySuspendResumeResult {
   enum AwaitCallType { ACT_Ready, ACT_Suspend, ACT_Resume };
   Expr *Results[3];
   OpaqueValueExpr *OpaqueValue;
-  bool IsSuspendNoThrow;
   bool IsInvalid;
 };
 
@@ -400,10 +402,7 @@ static ReadySuspendResumeResult buildCoawaitCalls(Sema &S, VarDecl *CoroPromise,
 
   // Assume valid until we see otherwise.
   // Further operations are responsible for setting IsInalid to true.
-  ReadySuspendResumeResult Calls = {{},
-                                    Operand,
-                                    /*IsSuspendNoThrow=*/false,
-                                    /*IsInvalid=*/false};
+  ReadySuspendResumeResult Calls = {{}, Operand, /*IsInvalid=*/false};
 
   using ACT = ReadySuspendResumeResult::AwaitCallType;
 
@@ -437,11 +436,8 @@ static ReadySuspendResumeResult buildCoawaitCalls(Sema &S, VarDecl *CoroPromise,
       Calls.Results[ACT::ACT_Ready] = S.MaybeCreateExprWithCleanups(Conv.get());
   }
 
-  Expr *FramePtr =
-      S.BuildBuiltinCallExpr(Loc, Builtin::BI__builtin_coro_frame, {});
-
   ExprResult CoroHandleRes =
-      buildCoroutineHandle(S, CoroPromise->getType(), FramePtr, Loc);
+      buildCoroutineHandle(S, CoroPromise->getType(), Loc);
   if (CoroHandleRes.isInvalid()) {
     Calls.IsInvalid = true;
     return Calls;

>From 507bc1700e9fe7c8f7921579bc1b05407fa693a7 Mon Sep 17 00:00:00 2001
From: fpasserby <125797601+fpasserby at users.noreply.github.com>
Date: Mon, 4 Mar 2024 22:58:34 +0100
Subject: [PATCH 17/17] Review fixes

---
 clang/include/clang/AST/ExprCXX.h             | 11 ++-
 clang/lib/CodeGen/CGCoroutine.cpp             | 40 ++++----
 clang/lib/Sema/SemaCoroutine.cpp              | 29 +-----
 .../coro-simplify-suspend-point.cpp           | 66 -------------
 clang/test/CodeGenCoroutines/pr56329.cpp      |  4 +
 clang/test/CodeGenCoroutines/pr65054.cpp      |  7 --
 llvm/lib/Transforms/Coroutines/CoroSplit.cpp  | 78 ++++------------
 .../Coroutines/coro-await-suspend-simplify.ll | 92 -------------------
 8 files changed, 49 insertions(+), 278 deletions(-)
 delete mode 100644 clang/test/CodeGenCoroutines/coro-simplify-suspend-point.cpp
 delete mode 100644 llvm/test/Transforms/Coroutines/coro-await-suspend-simplify.ll

diff --git a/clang/include/clang/AST/ExprCXX.h b/clang/include/clang/AST/ExprCXX.h
index 82bebd8c781bf31..6003b866c9f5645 100644
--- a/clang/include/clang/AST/ExprCXX.h
+++ b/clang/include/clang/AST/ExprCXX.h
@@ -5038,7 +5038,8 @@ class CoroutineSuspendExpr : public Expr {
   OpaqueValueExpr *OpaqueValue = nullptr;
 
 public:
-  enum SuspendReturnType { SuspendVoid, SuspendBool, SuspendHandle };
+  // These types correspond to the three C++ 'await_suspend' return variants
+  enum class SuspendReturnType { SuspendVoid, SuspendBool, SuspendHandle };
 
   CoroutineSuspendExpr(StmtClass SC, SourceLocation KeywordLoc, Expr *Operand,
                        Expr *Common, Expr *Ready, Expr *Suspend, Expr *Resume,
@@ -5109,10 +5110,12 @@ class CoroutineSuspendExpr : public Expr {
       return SuspendReturnType::SuspendVoid;
     if (SuspendType->isBooleanType())
       return SuspendReturnType::SuspendBool;
-    if (SuspendType->isVoidPointerType())
-      return SuspendReturnType::SuspendHandle;
 
-    llvm_unreachable("Unexpected await_suspend expression return type");
+    // Void pointer is the type of handle.address(), which is returned
+    // from the await suspend wrapper so that the temporary coroutine handle
+    // value won't go to the frame by mistake
+    assert(SuspendType->isVoidPointerType());
+    return SuspendReturnType::SuspendHandle;
   }
 
   SourceLocation getKeywordLoc() const { return KeywordLoc; }
diff --git a/clang/lib/CodeGen/CGCoroutine.cpp b/clang/lib/CodeGen/CGCoroutine.cpp
index 6c8c8553b87ed6a..d537a9cc93f840a 100644
--- a/clang/lib/CodeGen/CGCoroutine.cpp
+++ b/clang/lib/CodeGen/CGCoroutine.cpp
@@ -182,8 +182,8 @@ static bool AwaitSuspendStmtCanThrow(const Stmt *S) {
 //   auto && x = CommonExpr();
 //   if (!x.await_ready()) {
 //      llvm_coro_save();
-//      llvm_coro_await_suspend(&x, frame, wrapper) (*)
-//      llvm_coro_suspend(); (**)
+//      llvm_coro_await_suspend(&x, frame, wrapper) (*) (**)
+//      llvm_coro_suspend(); (***)
 //   }
 //   x.await_resume();
 //
@@ -202,11 +202,11 @@ static bool AwaitSuspendStmtCanThrow(const Stmt *S) {
 //        return awaiter.await_suspend(handle);
 //      }
 //
-//      If x.await_suspend return type is bool, it allows to veto a suspend:
+//  (**) If x.await_suspend return type is bool, it allows to veto a suspend:
 //      if (x.await_suspend(...))
 //        llvm_coro_suspend();
 //
-//  (**) llvm_coro_suspend() encodes three possible continuations as
+//  (***) llvm_coro_suspend() encodes three possible continuations as
 //       a switch instruction:
 //
 //  %where-to = call i8 @llvm.coro.suspend(...)
@@ -250,14 +250,9 @@ static LValueOrRValue emitSuspendExpression(CodeGenFunction &CGF, CGCoroData &Co
   auto *NullPtr = llvm::ConstantPointerNull::get(CGF.CGM.Int8PtrTy);
   auto *SaveCall = Builder.CreateCall(CoroSave, {NullPtr});
 
-  const auto AwaitSuspendCanThrow =
-      AwaitSuspendStmtCanThrow(S.getSuspendExpr());
-
   auto SuspendWrapper = CodeGenFunction(CGF.CGM).generateAwaitSuspendWrapper(
       CGF.CurFn->getName(), Prefix, S);
 
-  llvm::CallBase *SuspendRet = nullptr;
-
   CGF.CurCoro.InSuspendBlock = true;
 
   assert(CGF.CurCoro.Data && CGF.CurCoro.Data->CoroBegin &&
@@ -271,21 +266,26 @@ static LValueOrRValue emitSuspendExpression(CodeGenFunction &CGF, CGCoroData &Co
   SuspendIntrinsicCallArgs.push_back(SuspendWrapper);
 
   const auto SuspendReturnType = S.getSuspendReturnType();
-  llvm::Intrinsic::ID IID;
+  llvm::Intrinsic::ID AwaitSuspendIID;
 
   switch (SuspendReturnType) {
-  case CoroutineSuspendExpr::SuspendVoid:
-    IID = llvm::Intrinsic::coro_await_suspend_void;
+  case CoroutineSuspendExpr::SuspendReturnType::SuspendVoid:
+    AwaitSuspendIID = llvm::Intrinsic::coro_await_suspend_void;
     break;
-  case CoroutineSuspendExpr::SuspendBool:
-    IID = llvm::Intrinsic::coro_await_suspend_bool;
+  case CoroutineSuspendExpr::SuspendReturnType::SuspendBool:
+    AwaitSuspendIID = llvm::Intrinsic::coro_await_suspend_bool;
     break;
-  case CoroutineSuspendExpr::SuspendHandle:
-    IID = llvm::Intrinsic::coro_await_suspend_handle;
+  case CoroutineSuspendExpr::SuspendReturnType::SuspendHandle:
+    AwaitSuspendIID = llvm::Intrinsic::coro_await_suspend_handle;
     break;
   }
 
-  llvm::Function *AwaitSuspendIntrinsic = CGF.CGM.getIntrinsic(IID);
+  llvm::Function *AwaitSuspendIntrinsic = CGF.CGM.getIntrinsic(AwaitSuspendIID);
+
+  const auto AwaitSuspendCanThrow =
+      AwaitSuspendStmtCanThrow(S.getSuspendExpr());
+
+  llvm::CallBase *SuspendRet = nullptr;
   // FIXME: add call attributes?
   if (AwaitSuspendCanThrow)
     SuspendRet =
@@ -298,10 +298,10 @@ static LValueOrRValue emitSuspendExpression(CodeGenFunction &CGF, CGCoroData &Co
   CGF.CurCoro.InSuspendBlock = false;
 
   switch (SuspendReturnType) {
-  case CoroutineSuspendExpr::SuspendVoid:
+  case CoroutineSuspendExpr::SuspendReturnType::SuspendVoid:
     assert(SuspendRet->getType()->isVoidTy());
     break;
-  case CoroutineSuspendExpr::SuspendBool: {
+  case CoroutineSuspendExpr::SuspendReturnType::SuspendBool: {
     assert(SuspendRet->getType()->isIntegerTy());
 
     // Veto suspension if requested by bool returning await_suspend.
@@ -311,7 +311,7 @@ static LValueOrRValue emitSuspendExpression(CodeGenFunction &CGF, CGCoroData &Co
     CGF.EmitBlock(RealSuspendBlock);
     break;
   }
-  case CoroutineSuspendExpr::SuspendHandle: {
+  case CoroutineSuspendExpr::SuspendReturnType::SuspendHandle: {
     assert(SuspendRet->getType()->isPointerTy());
 
     auto ResumeIntrinsic = CGF.CGM.getIntrinsic(llvm::Intrinsic::coro_resume);
diff --git a/clang/lib/Sema/SemaCoroutine.cpp b/clang/lib/Sema/SemaCoroutine.cpp
index cfbf5791a3b999e..736632857efc369 100644
--- a/clang/lib/Sema/SemaCoroutine.cpp
+++ b/clang/lib/Sema/SemaCoroutine.cpp
@@ -347,39 +347,14 @@ static Expr *maybeTailCall(Sema &S, QualType RetType, Expr *E,
 
   Expr *JustAddress = AddressExpr.get();
 
-  // FIXME: Without optimizations, the temporary result from `await_suspend()`
-  // may be put on the coroutine frame since the coroutine frame constructor
-  // will think the temporary variable will escape from the
-  // `coroutine_handle<>::address()` call. This is problematic since the
-  // coroutine should be considered to be suspended after it enters
-  // `await_suspend` so it shouldn't access/update the coroutine frame after
-  // that.
-  //
-  // See https://github.com/llvm/llvm-project/issues/65054 for the report.
-  //
-  // The long term solution may wrap the whole logic about `await-suspend`
-  // into a standalone function. This is similar to the proposed solution
-  // in tryMarkAwaitSuspendNoInline. See the comments there for details.
-  //
-  // The short term solution here is to mark `coroutine_handle<>::address()`
-  // function as always-inline so that the coroutine frame constructor won't
-  // think the temporary result is escaped incorrectly.
-  if (auto *FD = cast<CallExpr>(JustAddress)->getDirectCallee())
-    if (!FD->hasAttr<AlwaysInlineAttr>() && !FD->hasAttr<NoInlineAttr>())
-      FD->addAttr(AlwaysInlineAttr::CreateImplicit(S.getASTContext(),
-                                                   FD->getLocation()));
-
   // Check that the type of AddressExpr is void*
   if (!JustAddress->getType().getTypePtr()->isVoidPointerType())
     S.Diag(cast<CallExpr>(JustAddress)->getCalleeDecl()->getLocation(),
            diag::warn_coroutine_handle_address_invalid_return_type)
         << JustAddress->getType();
 
-  // Clean up temporary objects so that they don't live across suspension points
-  // unnecessarily. We choose to clean up before the call to
-  // __builtin_coro_resume so that the cleanup code are not inserted in-between
-  // the resume call and return instruction, which would interfere with the
-  // musttail call contract.
+  // Clean up temporary objects, because the resulting expression
+  // will become the body of await_suspend wrapper.
   return S.MaybeCreateExprWithCleanups(JustAddress);
 }
 
diff --git a/clang/test/CodeGenCoroutines/coro-simplify-suspend-point.cpp b/clang/test/CodeGenCoroutines/coro-simplify-suspend-point.cpp
deleted file mode 100644
index 32960311043c804..000000000000000
--- a/clang/test/CodeGenCoroutines/coro-simplify-suspend-point.cpp
+++ /dev/null
@@ -1,66 +0,0 @@
-// Test that we can perform suspend point simplification
-// RUN: %clang_cc1 -triple x86_64-unknown-linux-gnu -std=c++20 -O1 -emit-llvm %s -o - | FileCheck %s
-// RUN: %clang -std=c++20 -O1 -emit-llvm -c  %s -o %t && %clang -c %t
-
-#include "Inputs/coroutine.h"
-
-struct detached_task {
-  struct promise_type {
-    detached_task get_return_object() noexcept {
-      return detached_task{std::coroutine_handle<promise_type>::from_promise(*this)};
-    }
-
-    void return_void() noexcept {}
-
-    struct final_awaiter {
-      bool await_ready() noexcept { return false; }
-      std::coroutine_handle<> await_suspend(std::coroutine_handle<promise_type> h) noexcept {
-        h.destroy();
-        return std::noop_coroutine();
-      }
-      void await_resume() noexcept {}
-    };
-
-    void unhandled_exception() noexcept {}
-
-    final_awaiter final_suspend() noexcept { return {}; }
-
-    std::suspend_always initial_suspend() noexcept { return {}; }
-  };
-
-  ~detached_task() {
-    if (coro_) {
-      coro_.destroy();
-      coro_ = {};
-    }
-  }
-
-  void start() && {
-    auto tmp = coro_;
-    coro_ = {};
-    tmp.resume();
-  }
-
-  std::coroutine_handle<promise_type> coro_;
-};
-
-class SelfResumeAwaiter final
-{
-public:
-    bool await_ready() noexcept { return false; }
-    std::coroutine_handle<> await_suspend(std::coroutine_handle<> h) { 
-      return h; 
-    }
-    void await_resume() noexcept {}
-};
-
-// Check that there is only one call left: coroutine destroy
-// CHECK-LABEL: define{{.*}}void @_Z3foov.resume
-// CHECK-NOT: call{{.*}}
-// CHECK: tail call{{.*}}void %{{[0-9+]}}(
-// CHECK-NOT: call{{.*}}
-// CHECK: define
-detached_task foo() {
-  co_await SelfResumeAwaiter{};
-  co_return;
-}
diff --git a/clang/test/CodeGenCoroutines/pr56329.cpp b/clang/test/CodeGenCoroutines/pr56329.cpp
index 31d4849af4e71ef..ad8b1b990179c94 100644
--- a/clang/test/CodeGenCoroutines/pr56329.cpp
+++ b/clang/test/CodeGenCoroutines/pr56329.cpp
@@ -115,5 +115,9 @@ Task Outer() {
 // CHECK-NOT: _exit
 // CHECK: musttail call
 // CHECK: musttail call
+// CHECK: musttail call
 // CHECK-NEXT: ret void
+// CHECK-EMPTY:
+// CHECK-NEXT: unreachable:
+// CHECK-NEXT: unreachable
 // CHECK-NEXT: }
diff --git a/clang/test/CodeGenCoroutines/pr65054.cpp b/clang/test/CodeGenCoroutines/pr65054.cpp
index 506990163e9b8d1..7af9c04fca180e7 100644
--- a/clang/test/CodeGenCoroutines/pr65054.cpp
+++ b/clang/test/CodeGenCoroutines/pr65054.cpp
@@ -1,7 +1,3 @@
-// RUN: %clang_cc1 -triple x86_64-unknown-linux-gnu -std=c++20 \
-// RUN:      -O0 -disable-llvm-passes -emit-llvm %s -o - \
-// RUN:      | FileCheck %s --check-prefix=FRONTEND
-
 // The output of O0 is highly redundant and hard to test. Also it is not good
 // limit the output of O0. So we test the optimized output from O0. The idea
 // is the optimizations shouldn't change the semantics of the program.
@@ -51,9 +47,6 @@ MyTask FooBar() {
   }
 }
 
-// FRONTEND: define{{.*}}@_ZNKSt16coroutine_handleIvE7addressEv{{.*}}#[[address_attr:[0-9]+]]
-// FRONTEND: attributes #[[address_attr]] = {{.*}}alwaysinline
-
 // CHECK-O0: define{{.*}}@_Z6FooBarv.resume
 // CHECK-O0: call{{.*}}@__await_suspend_wrapper__Z6FooBarv_await(
 // CHECK-O0-NOT: store
diff --git a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
index 3019dee47b94ef8..7993aaf6925d684 100644
--- a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
@@ -168,7 +168,7 @@ class CoroCloner {
 } // end anonymous namespace
 
 // FIXME:
-// Lower the intrinisc earlier if coroutine frame doesn't escape
+// Lower the intrinisc in CoroEarly phase if coroutine frame doesn't escape
 // and it is known that other transformations, for example, sanitizers
 // won't lead to incorrect code.
 static void lowerAwaitSuspend(IRBuilder<> &Builder, CoroAwaitSuspendInst *CB) {
@@ -208,6 +208,12 @@ static void lowerAwaitSuspend(IRBuilder<> &Builder, CoroAwaitSuspendInst *CB) {
   CB->eraseFromParent();
 }
 
+static void lowerAwaitSuspends(Function &F, coro::Shape &Shape) {
+  IRBuilder<> Builder(F.getContext());
+  for (auto *AWS : Shape.CoroAwaitSuspends)
+    lowerAwaitSuspend(Builder, AWS);
+}
+
 static void maybeFreeRetconStorage(IRBuilder<> &Builder,
                                    const coro::Shape &Shape, Value *FramePtr,
                                    CallGraph *CG) {
@@ -1360,11 +1366,6 @@ static void handleNoSuspendCoroutine(coro::Shape &Shape) {
 // the coroutine and if that is the case we cannot eliminate the suspend point.
 static bool hasCallsInBlockBetween(Instruction *From, Instruction *To) {
   for (Instruction *I = From; I != To; I = I->getNextNode()) {
-    // This one could resume the coroutine,
-    // but additional analysis before the check should ensure,
-    // that it can't happen
-    if (isa<CoroAwaitSuspendInst>(I))
-      continue;
     // Assume that no intrinsic can resume the coroutine.
     if (isa<IntrinsicInst>(I))
       continue;
@@ -1405,17 +1406,14 @@ static bool hasCallsInBlocksBetween(BasicBlock *SaveBB, BasicBlock *ResDesBB) {
 }
 
 static bool hasCallsBetween(Instruction *Save, Instruction *ResumeOrDestroy) {
-  if (Save == ResumeOrDestroy)
-    return false;
-
   auto *SaveBB = Save->getParent();
   auto *ResumeOrDestroyBB = ResumeOrDestroy->getParent();
 
   if (SaveBB == ResumeOrDestroyBB)
-    return hasCallsInBlockBetween(Save, ResumeOrDestroy);
+    return hasCallsInBlockBetween(Save->getNextNode(), ResumeOrDestroy);
 
   // Any calls from Save to the end of the block?
-  if (hasCallsInBlockBetween(Save, nullptr))
+  if (hasCallsInBlockBetween(Save->getNextNode(), nullptr))
     return true;
 
   // Any calls from begging of the block up to ResumeOrDestroy?
@@ -1430,39 +1428,6 @@ static bool hasCallsBetween(Instruction *Save, Instruction *ResumeOrDestroy) {
   return false;
 }
 
-// Check if await-suspend wrapper is "simple".
-// The conditions are:
-// 1. The return result is exactly coroutine frame parameter, passed to wrapper
-// 2. There are no calls between any of the returns and wrapper entry that could
-// resume or destroy it
-// FIXME: perform more sophisiticated analysis?
-static bool isSimpleWrapper(CoroAwaitSuspendInst *AWS) {
-  auto *Wrapper = AWS->getWrapperFunction();
-
-  if (Wrapper->empty())
-    return false;
-
-  SmallVector<ReturnInst *, 4> Rets;
-
-  for (auto &BB : *Wrapper) {
-    if (BB.empty())
-      continue;
-    auto terminator = BB.getTerminator();
-    if (!terminator)
-      continue;
-    if (auto Ret = dyn_cast<ReturnInst>(terminator))
-      Rets.push_back(cast<ReturnInst>(terminator));
-  }
-
-  // FIXME: get rid of magical constant
-  for (auto Ret : Rets)
-    if (Ret->getReturnValue() != Wrapper->getArg(1) ||
-        hasCallsBetween(Wrapper->getEntryBlock().getFirstNonPHI(), Ret))
-      return false;
-
-  return true;
-}
-
 // If a SuspendIntrin is preceded by Resume or Destroy, we can eliminate the
 // suspend point and replace it with nornal control flow.
 static bool simplifySuspendPoint(CoroSuspendInst *Suspend,
@@ -1486,18 +1451,9 @@ static bool simplifySuspendPoint(CoroSuspendInst *Suspend,
   if (!SubFn)
     return false;
 
-  auto Frame = SubFn->getFrame();
-
-  // Check that frame directly always refers to the current coroutine,
-  // either directly or via wrapper
-  if (Frame != CoroBegin) {
-    auto *AWS = dyn_cast<CoroAwaitSuspendInst>(Frame);
-    if (!AWS)
-      return false;
-
-    if (AWS->getFrame() != CoroBegin || !isSimpleWrapper(AWS))
-      return false;
-  }
+  // Does not refer to the current coroutine, we cannot do anything with it.
+  if (SubFn->getFrame() != CoroBegin)
+    return false;
 
   // See if the transformation is safe. Specifically, see if there are any
   // calls in between Save and CallInstr. They can potenitally resume the
@@ -2103,7 +2059,7 @@ class PrettyStackTraceFunction : public PrettyStackTraceEntry {
 } // namespace
 
 static coro::Shape
-splitCoroutine(Module &M, Function &F, SmallVectorImpl<Function *> &Clones,
+splitCoroutine(Function &F, SmallVectorImpl<Function *> &Clones,
                TargetTransformInfo &TTI, bool OptimizeFrame,
                std::function<bool(Instruction &)> MaterializableCallback) {
   PrettyStackTraceFunction prettyStackTrace(F);
@@ -2116,14 +2072,12 @@ splitCoroutine(Module &M, Function &F, SmallVectorImpl<Function *> &Clones,
   if (!Shape.CoroBegin)
     return Shape;
 
+  lowerAwaitSuspends(F, Shape);
+
   simplifySuspendPoints(Shape);
   buildCoroutineFrame(F, Shape, TTI, MaterializableCallback);
   replaceFrameSizeAndAlignment(Shape);
 
-  IRBuilder<> Builder(M.getContext());
-  for (auto *AWS : Shape.CoroAwaitSuspends)
-    lowerAwaitSuspend(Builder, AWS);
-
   // If there are no suspend points, no split required, just remove
   // the allocation and deallocation blocks, they are not needed.
   if (Shape.CoroSuspends.empty()) {
@@ -2316,7 +2270,7 @@ PreservedAnalyses CoroSplitPass::run(LazyCallGraph::SCC &C,
     SmallVector<Function *, 4> Clones;
     auto &ORE = FAM.getResult<OptimizationRemarkEmitterAnalysis>(F);
     const coro::Shape Shape =
-        splitCoroutine(M, F, Clones, FAM.getResult<TargetIRAnalysis>(F),
+        splitCoroutine(F, Clones, FAM.getResult<TargetIRAnalysis>(F),
                        OptimizeFrame, MaterializableCallback);
     updateCallGraphAfterCoroutineSplit(*N, Shape, Clones, C, CG, AM, UR, FAM);
 
diff --git a/llvm/test/Transforms/Coroutines/coro-await-suspend-simplify.ll b/llvm/test/Transforms/Coroutines/coro-await-suspend-simplify.ll
deleted file mode 100644
index 8bd7ba6ceb96641..000000000000000
--- a/llvm/test/Transforms/Coroutines/coro-await-suspend-simplify.ll
+++ /dev/null
@@ -1,92 +0,0 @@
-; Tests that CoroSplit can succesfully see through await suspend wrapper
-; and remove redundant suspend points only if there are no calls in the wrapper
-; RUN: opt < %s -passes='module(coro-early),cgscc(coro-split),simplifycfg' -S | FileCheck %s
-
-; CHECK-LABEL: @f(
-; CHECK:        call ptr @await_suspend_wrapper(
-; CHECK-NEXT:   call void @free(
-define void @f() presplitcoroutine {
-entry:
-  %id = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null)
-  %size = call i32 @llvm.coro.size.i32()
-  %alloc = call ptr @malloc(i32 %size)
-  %hdl = call ptr @llvm.coro.begin(token %id, ptr %alloc)
-  %save = call token @llvm.coro.save(ptr null)
-  %hdl.resume = call ptr @llvm.coro.await.suspend.handle(ptr null, ptr %hdl, ptr @await_suspend_wrapper)
-  call void @llvm.coro.resume(ptr %hdl.resume)
-  %suspend = call i8 @llvm.coro.suspend(token %save, i1 false)
-  switch i8 %suspend, label %ret [
-    i8 0, label %resume
-    i8 1, label %cleanup
-  ]
-
-resume:
-  br label %cleanup
-
-cleanup:
-  %mem = call ptr @llvm.coro.free(token %id, ptr %hdl)
-  call void @free(ptr %mem)
-  br label %ret
-
-ret:
-  call i1 @llvm.coro.end(ptr %hdl, i1 0, token none)
-  ret void
-}
-
-define ptr @await_suspend_wrapper(ptr, ptr %hdl) {
-  ret ptr %hdl
-}
-
-; CHECK-LABEL: @f1(
-; CHECK:        %[[HDL:.+]] = call ptr @await_suspend_wrapper1(
-; CHECK-NEXT:   %[[FN:.+]] = call ptr @llvm.coro.subfn.addr(ptr %[[HDL]], i8 0)
-; CHECK-NEXT:   call {{.*}} %[[FN]](ptr %[[HDL]])
-define void @f1() presplitcoroutine {
-entry:
-  %id = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null)
-  %size = call i32 @llvm.coro.size.i32()
-  %alloc = call ptr @malloc(i32 %size)
-  %hdl = call ptr @llvm.coro.begin(token %id, ptr %alloc)
-  %save = call token @llvm.coro.save(ptr null)
-  %hdl.resume = call ptr @llvm.coro.await.suspend.handle(ptr null, ptr %hdl, ptr @await_suspend_wrapper1)
-  call void @llvm.coro.resume(ptr %hdl.resume)
-  %suspend = call i8 @llvm.coro.suspend(token %save, i1 false)
-  switch i8 %suspend, label %ret [
-    i8 0, label %resume
-    i8 1, label %cleanup
-  ]
-
-resume:
-  br label %cleanup
-
-cleanup:
-  %mem = call ptr @llvm.coro.free(token %id, ptr %hdl)
-  call void @free(ptr %mem)
-  br label %ret
-
-ret:
-  call i1 @llvm.coro.end(ptr %hdl, i1 0, token none)
-  ret void
-}
-
-define ptr @await_suspend_wrapper1(ptr, ptr %hdl) {
-  call void @external()
-  ret ptr %hdl
-}
-
-declare ptr @llvm.coro.free(token, ptr)
-declare i32 @llvm.coro.size.i32()
-declare i8  @llvm.coro.suspend(token, i1)
-declare void @llvm.coro.resume(ptr)
-declare void @llvm.coro.destroy(ptr)
-
-declare token @llvm.coro.id(i32, ptr, ptr, ptr)
-declare i1 @llvm.coro.alloc(token)
-declare ptr @llvm.coro.begin(token, ptr)
-declare ptr @llvm.coro.await.suspend.handle(ptr, ptr, ptr)
-declare i1 @llvm.coro.end(ptr, i1, token)
-
-declare void @external()
-
-declare noalias ptr @malloc(i32)
-declare void @free(ptr)



More information about the cfe-commits mailing list