[llvm] [llvm][Coro] Lower coro.frame in CoroEarly instead of CoroSplit (PR #137976)

[Weibo He via llvm-commits]

https://github.com/NewSigma updated https://github.com/llvm/llvm-project/pull/137976

>From 85c437a132098d4a157a5822fe621d0337a4210e Mon Sep 17 00:00:00 2001
From: NewSigma <NewSigma at 163.com>
Date: Fri, 16 May 2025 17:34:34 +0800
Subject: [PATCH] [llvm][Coro] Lower coro.frame in CoroEarly instead of
 CoroSplit

---
 .../llvm/Transforms/Coroutines/CoroShape.h    | 16 +++++--------
 llvm/lib/Transforms/Coroutines/CoroEarly.cpp  | 20 ++++++++++++++--
 llvm/lib/Transforms/Coroutines/Coroutines.cpp | 24 +------------------
 3 files changed, 25 insertions(+), 35 deletions(-)

diff --git a/llvm/include/llvm/Transforms/Coroutines/CoroShape.h b/llvm/include/llvm/Transforms/Coroutines/CoroShape.h
index 891774b446571..038fc28903981 100644
--- a/llvm/include/llvm/Transforms/Coroutines/CoroShape.h
+++ b/llvm/include/llvm/Transforms/Coroutines/CoroShape.h
@@ -78,17 +78,14 @@ struct Shape {
   }
 
   // Scan the function and collect the above intrinsics for later processing
-  void analyze(Function &F, SmallVectorImpl<CoroFrameInst *> &CoroFrames,
-               SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves,
+  void analyze(Function &F, SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves,
                CoroPromiseInst *&CoroPromise);
   // If for some reason, we were not able to find coro.begin, bailout.
-  void invalidateCoroutine(Function &F,
-                           SmallVectorImpl<CoroFrameInst *> &CoroFrames);
+  void invalidateCoroutine(Function &F);
   // Perform ABI related initial transformation
   void initABI();
   // Remove orphaned and unnecessary intrinsics
-  void cleanCoroutine(SmallVectorImpl<CoroFrameInst *> &CoroFrames,
-                      SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves,
+  void cleanCoroutine(SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves,
                       CoroPromiseInst *CoroPromise);
 
   // Field indexes for special fields in the switch lowering.
@@ -265,16 +262,15 @@ struct Shape {
 
   Shape() = default;
   explicit Shape(Function &F) {
-    SmallVector<CoroFrameInst *, 8> CoroFrames;
     SmallVector<CoroSaveInst *, 2> UnusedCoroSaves;
     CoroPromiseInst *CoroPromise = nullptr;
 
-    analyze(F, CoroFrames, UnusedCoroSaves, CoroPromise);
+    analyze(F, UnusedCoroSaves, CoroPromise);
     if (!CoroBegin) {
-      invalidateCoroutine(F, CoroFrames);
+      invalidateCoroutine(F);
       return;
     }
-    cleanCoroutine(CoroFrames, UnusedCoroSaves, CoroPromise);
+    cleanCoroutine(UnusedCoroSaves, CoroPromise);
   }
 };
 
diff --git a/llvm/lib/Transforms/Coroutines/CoroEarly.cpp b/llvm/lib/Transforms/Coroutines/CoroEarly.cpp
index eea6dfba14e37..717e1bfc1a074 100644
--- a/llvm/lib/Transforms/Coroutines/CoroEarly.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroEarly.cpp
@@ -25,8 +25,10 @@ class Lowerer : public coro::LowererBase {
   IRBuilder<> Builder;
   PointerType *const AnyResumeFnPtrTy;
   Constant *NoopCoro = nullptr;
+  SmallVector<CoroFrameInst *, 8> CoroFrames;
 
   void lowerResumeOrDestroy(CallBase &CB, CoroSubFnInst::ResumeKind);
+  void lowerCoroFrames(Function &F, Value *CoroBegin);
   void lowerCoroPromise(CoroPromiseInst *Intrin);
   void lowerCoroDone(IntrinsicInst *II);
   void lowerCoroNoop(IntrinsicInst *II);
@@ -51,6 +53,18 @@ void Lowerer::lowerResumeOrDestroy(CallBase &CB,
   CB.setCallingConv(CallingConv::Fast);
 }
 
+void Lowerer::lowerCoroFrames(Function &F, Value *CoroBegin) {
+  // Lower with poison if we cannot find coro.begin
+  if (CoroBegin == nullptr)
+    CoroBegin = PoisonValue::get(PointerType::get(F.getContext(), 0));
+
+  for (CoroFrameInst *CF : CoroFrames) {
+    CF->replaceAllUsesWith(CoroBegin);
+    CF->eraseFromParent();
+  }
+  CoroFrames.clear();
+}
+
 // Coroutine promise field is always at the fixed offset from the beginning of
 // the coroutine frame. i8* coro.promise(i8*, i1 from) intrinsic adds an offset
 // to a passed pointer to move from coroutine frame to coroutine promise and
@@ -257,6 +271,8 @@ void Lowerer::lowerEarlyIntrinsics(Function &F) {
         break;
     }
   }
+  // The coro.frame intrinsic is always lowered to the result of coro.begin.
+  lowerCoroFrames(F, CoroBegin);
 
   if (CoroId) {
     // Make sure that all CoroFree reference the coro.id intrinsic.
@@ -282,8 +298,8 @@ static bool declaresCoroEarlyIntrinsics(const Module &M) {
       M, {"llvm.coro.id", "llvm.coro.id.retcon", "llvm.coro.id.retcon.once",
           "llvm.coro.id.async", "llvm.coro.destroy", "llvm.coro.done",
           "llvm.coro.end", "llvm.coro.end.async", "llvm.coro.noop",
-          "llvm.coro.free", "llvm.coro.promise", "llvm.coro.resume",
-          "llvm.coro.suspend"});
+          "llvm.coro.frame", "llvm.coro.free", "llvm.coro.promise",
+          "llvm.coro.resume", "llvm.coro.suspend"});
 }
 
 PreservedAnalyses CoroEarlyPass::run(Module &M, ModuleAnalysisManager &) {
diff --git a/llvm/lib/Transforms/Coroutines/Coroutines.cpp b/llvm/lib/Transforms/Coroutines/Coroutines.cpp
index 02500ff778b80..b9b0d135e37b3 100644
--- a/llvm/lib/Transforms/Coroutines/Coroutines.cpp
+++ b/llvm/lib/Transforms/Coroutines/Coroutines.cpp
@@ -191,7 +191,6 @@ static CoroSaveInst *createCoroSave(CoroBeginInst *CoroBegin,
 
 // Collect "interesting" coroutine intrinsics.
 void coro::Shape::analyze(Function &F,
-                          SmallVectorImpl<CoroFrameInst *> &CoroFrames,
                           SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves,
                           CoroPromiseInst *&CoroPromise) {
   clear();
@@ -215,9 +214,6 @@ void coro::Shape::analyze(Function &F,
       case Intrinsic::coro_align:
         CoroAligns.push_back(cast<CoroAlignInst>(II));
         break;
-      case Intrinsic::coro_frame:
-        CoroFrames.push_back(cast<CoroFrameInst>(II));
-        break;
       case Intrinsic::coro_save:
         // After optimizations, coro_suspends using this coro_save might have
         // been removed, remember orphaned coro_saves to remove them later.
@@ -351,19 +347,9 @@ void coro::Shape::analyze(Function &F,
 }
 
 // If for some reason, we were not able to find coro.begin, bailout.
-void coro::Shape::invalidateCoroutine(
-    Function &F, SmallVectorImpl<CoroFrameInst *> &CoroFrames) {
+void coro::Shape::invalidateCoroutine(Function &F) {
   assert(!CoroBegin);
   {
-    // Replace coro.frame which are supposed to be lowered to the result of
-    // coro.begin with poison.
-    auto *Poison = PoisonValue::get(PointerType::get(F.getContext(), 0));
-    for (CoroFrameInst *CF : CoroFrames) {
-      CF->replaceAllUsesWith(Poison);
-      CF->eraseFromParent();
-    }
-    CoroFrames.clear();
-
     // Replace all coro.suspend with poison and remove related coro.saves if
     // present.
     for (AnyCoroSuspendInst *CS : CoroSuspends) {
@@ -482,15 +468,7 @@ void coro::AnyRetconABI::init() {
 }
 
 void coro::Shape::cleanCoroutine(
-    SmallVectorImpl<CoroFrameInst *> &CoroFrames,
     SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves, CoroPromiseInst *PI) {
-  // The coro.frame intrinsic is always lowered to the result of coro.begin.
-  for (CoroFrameInst *CF : CoroFrames) {
-    CF->replaceAllUsesWith(CoroBegin);
-    CF->eraseFromParent();
-  }
-  CoroFrames.clear();
-
   // Remove orphaned coro.saves.
   for (CoroSaveInst *CoroSave : UnusedCoroSaves)
     CoroSave->eraseFromParent();