[llvm][Coro] Lower coro.frame in CoroEarly instead of CoroSplit

[NewSigma]

Documentation of CoroEarly states:

This pass lowers coro.frame, coro.done, and coro.promise intrinsics.

However, currently coro.frame lowering is handled by CoroSplit. This commit moves this operation to CoroEarly.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

@llvm/pr-subscribers-coroutines

Author: None (NewSigma)

Changes

Documentation of CoroEarly states:

> This pass lowers coro.frame, coro.done, and coro.promise intrinsics.

However, currently coro.frame lowering is handled by CoroSplit. This commit moves this operation to CoroEarly.

---

Full diff: https://github.com/llvm/llvm-project/pull/137976.diff

3 Files Affected:

• (modified) llvm/include/llvm/Transforms/Coroutines/CoroShape.h (+6-10)
• (modified) llvm/lib/Transforms/Coroutines/CoroEarly.cpp (+41-4)
• (modified) llvm/lib/Transforms/Coroutines/Coroutines.cpp (+1-23)

diff --git a/llvm/include/llvm/Transforms/Coroutines/CoroShape.h b/llvm/include/llvm/Transforms/Coroutines/CoroShape.h
index ea93ced1ce29e..c3b14fbc36e66 100644
--- a/llvm/include/llvm/Transforms/Coroutines/CoroShape.h
+++ b/llvm/include/llvm/Transforms/Coroutines/CoroShape.h
@@ -78,16 +78,13 @@ 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);
   // 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);
 
   // Field indexes for special fields in the switch lowering.
   struct SwitchFieldIndex {
@@ -263,15 +260,14 @@ struct Shape {
 
   Shape() = default;
   explicit Shape(Function &F) {
-    SmallVector<CoroFrameInst *, 8> CoroFrames;
     SmallVector<CoroSaveInst *, 2> UnusedCoroSaves;
 
-    analyze(F, CoroFrames, UnusedCoroSaves);
+    analyze(F, UnusedCoroSaves);
     if (!CoroBegin) {
-      invalidateCoroutine(F, CoroFrames);
+      invalidateCoroutine(F);
       return;
     }
-    cleanCoroutine(CoroFrames, UnusedCoroSaves);
+    cleanCoroutine(UnusedCoroSaves);
   }
 };
 
diff --git a/llvm/lib/Transforms/Coroutines/CoroEarly.cpp b/llvm/lib/Transforms/Coroutines/CoroEarly.cpp
index 5375448d2d2e2..df62cc6c51953 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);
@@ -50,6 +52,18 @@ void Lowerer::lowerResumeOrDestroy(CallBase &CB,
   CB.setCallingConv(CallingConv::Fast);
 }
 
+void Lowerer::lowerCoroFrames(Function &F, Value *CoroBegin) {
+  // Lower with poison if we cannot func 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
@@ -165,6 +179,7 @@ static void setCannotDuplicate(CoroIdInst *CoroId) {
 
 void Lowerer::lowerEarlyIntrinsics(Function &F) {
   CoroIdInst *CoroId = nullptr;
+  CoroBeginInst *CoroBegin = nullptr;
   SmallVector<CoroFreeInst *, 4> CoroFrees;
   bool HasCoroSuspend = false;
   for (Instruction &I : llvm::make_early_inc_range(instructions(F))) {
@@ -175,6 +190,26 @@ void Lowerer::lowerEarlyIntrinsics(Function &F) {
     switch (CB->getIntrinsicID()) {
       default:
         continue;
+      case Intrinsic::coro_begin:
+      case Intrinsic::coro_begin_custom_abi: {
+        auto CBI = cast<CoroBeginInst>(&I);
+
+        // Ignore coro id's that aren't pre-split.
+        auto Id = dyn_cast<CoroIdInst>(CBI->getId());
+        if (Id && !Id->getInfo().isPreSplit())
+          break;
+
+        if (CoroBegin)
+          report_fatal_error(
+              "coroutine should have exactly one defining @llvm.coro.begin");
+        CBI->addRetAttr(Attribute::NonNull);
+        CBI->addRetAttr(Attribute::NoAlias);
+        CoroBegin = CBI;
+        break;
+      }
+      case Intrinsic::coro_frame:
+        CoroFrames.push_back(cast<CoroFrameInst>(&I));
+        break;
       case Intrinsic::coro_free:
         CoroFrees.push_back(cast<CoroFreeInst>(&I));
         break;
@@ -226,6 +261,8 @@ void Lowerer::lowerEarlyIntrinsics(Function &F) {
         break;
     }
   }
+  // The coro.frame intrinsic is always lowered to the result of coro.begin.
+  lowerCoroFrames(F, CoroBegin);
 
   // Make sure that all CoroFree reference the coro.id intrinsic.
   // Token type is not exposed through coroutine C/C++ builtins to plain C, so
@@ -246,10 +283,10 @@ void Lowerer::lowerEarlyIntrinsics(Function &F) {
 static bool declaresCoroEarlyIntrinsics(const Module &M) {
   return coro::declaresIntrinsics(
       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.id.async", "llvm.coro.begin", "llvm.coro.destroy",
+          "llvm.coro.done", "llvm.coro.end", "llvm.coro.end.async",
+          "llvm.coro.noop", "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 7b59c39283ded..7dee7a6f571d7 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) {
   clear();
 
@@ -214,9 +213,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.
@@ -345,19 +341,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) {
@@ -476,15 +462,7 @@ void coro::AnyRetconABI::init() {
 }
 
 void coro::Shape::cleanCoroutine(
-    SmallVectorImpl<CoroFrameInst *> &CoroFrames,
     SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves) {
-  // 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();

[NewSigma]

Invite @TylerNowicki and @ChuanqiXu9 for code review