[Coroutines] Refactor CoroShape::buildFrom for future use by ABI objects

[TylerNowicki]

• Refactor buildFrom to separate the analysis, abi related operations,
  tidying and bailout.
• In a follow-up PR the code in initABI will be moved to an ABI object
  init method. And the Shape constructor will no longer perform any
  lowering, instead it will just call analysis. This will make the Shape
  object a bit more useful because it can be constructed and used
  anywhere. It may even be useful to make it an analysis pass.
• In a follow-up PR the OptimizeFrame flag will also be removed from the
  Shape and instead will be passed directly to buildCoroutineFrame
  (although it would be nice to find another way to trigger this
  optimization). This is the only thing that Shape cannot determine from
  the Function/Coroutine, but it is only needed within
  buildCoroutineFrame.
• Note, that it was necessary to introduce two new SmallVectors to
  ShapeInfo, these were previously used internally in buildFrom. They
  are CoroFrames and UnusedCoroSaves. The cleanCoroutine and
  invalidateCoroutine (bailout) methods will erase the insts they
  manage.
• Note, code for moving the final suspend to the last element in the
  CoroSuspends vector was placed in the analysis method.

See RFC for more info: https://discourse.llvm.org/t/rfc-abi-objects-for-coroutines/81057

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

@llvm/pr-subscribers-coroutines

Author: Tyler Nowicki (TylerNowicki)

Changes

• Refactor buildFrom to separate the analysis, abi related operations, tidying and bailout.
• In a follow-up PR the code in initABI will be moved to an ABI object init method and the Shape constructor will no longer perform any lowering. This will make the Shape object a bit more useful because it can be constructed and used anywhere. It may even be useful to make it an analysis pass. This future patch will also simplify the constructor down to a single call to analysis().
• In another follow-up PR the OptimizeFrame flag will also be removed from the Shape and instead will be passed directly to buildCoroutineFrame (although it would be nice to find another way to trigger this optimization). This is the only thing that Shape cannot determine from the Function/Coroutine, but it is only needed within buildCoroutineFrame.
• Note, that it was necessary to introduce two new SmallVectors, one to track CoroFrames and the other for UnusedCoroSaves. The tidyCoroutine method requires both, while invalidateCoroutine (bailout) method just requires the former.
• Note, code for moving the final suspend to the last element in the CoroSuspends vector was placed in the analysis method.

See RFC for more info: https://discourse.llvm.org/t/rfc-abi-objects-for-coroutines/81057

---

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

2 Files Affected:

• (modified) llvm/lib/Transforms/Coroutines/CoroShape.h (+47-8)
• (modified) llvm/lib/Transforms/Coroutines/Coroutines.cpp (+70-67)

diff --git a/llvm/lib/Transforms/Coroutines/CoroShape.h b/llvm/lib/Transforms/Coroutines/CoroShape.h
index f5798b63bf7325..3d1b38082173d1 100644
--- a/llvm/lib/Transforms/Coroutines/CoroShape.h
+++ b/llvm/lib/Transforms/Coroutines/CoroShape.h
@@ -50,15 +50,49 @@ enum class ABI {
 // Holds structural Coroutine Intrinsics for a particular function and other
 // values used during CoroSplit pass.
 struct LLVM_LIBRARY_VISIBILITY Shape {
-  CoroBeginInst *CoroBegin;
+  CoroBeginInst *CoroBegin = nullptr;
   SmallVector<AnyCoroEndInst *, 4> CoroEnds;
   SmallVector<CoroSizeInst *, 2> CoroSizes;
   SmallVector<CoroAlignInst *, 2> CoroAligns;
   SmallVector<AnyCoroSuspendInst *, 4> CoroSuspends;
-  SmallVector<CallInst *, 2> SwiftErrorOps;
   SmallVector<CoroAwaitSuspendInst *, 4> CoroAwaitSuspends;
   SmallVector<CallInst *, 2> SymmetricTransfers;
 
+  // Values invalidated by invalidateCoroutine() and tidyCoroutine()
+  SmallVector<CoroFrameInst *, 8> CoroFrames;
+  SmallVector<CoroSaveInst *, 2> UnusedCoroSaves;
+
+  // Values invalidated by replaceSwiftErrorOps()
+  SmallVector<CallInst *, 2> SwiftErrorOps;
+
+  void clear() {
+    CoroBegin = nullptr;
+    CoroEnds.clear();
+    CoroSizes.clear();
+    CoroAligns.clear();
+    CoroSuspends.clear();
+    CoroAwaitSuspends.clear();
+    SymmetricTransfers.clear();
+
+    CoroFrames.clear();
+    UnusedCoroSaves.clear();
+
+    SwiftErrorOps.clear();
+
+    FrameTy = nullptr;
+    FramePtr = nullptr;
+    AllocaSpillBlock = nullptr;
+  }
+
+  // Scan the function and collect the above intrinsics for later processing
+  void analyze(Function &F);
+  // If for some reason, we were not able to find coro.begin, bailout.
+  void invalidateCoroutine(Function &F);
+  // Perform ABI related initial transformation
+  void initABI();
+  // Remove orphaned and unnecessary intrinsics
+  void tidyCoroutine();
+
   // Field indexes for special fields in the switch lowering.
   struct SwitchFieldIndex {
     enum {
@@ -76,11 +110,11 @@ struct LLVM_LIBRARY_VISIBILITY Shape {
 
   coro::ABI ABI;
 
-  StructType *FrameTy;
+  StructType *FrameTy = nullptr;
   Align FrameAlign;
-  uint64_t FrameSize;
-  Value *FramePtr;
-  BasicBlock *AllocaSpillBlock;
+  uint64_t FrameSize = 0;
+  Value *FramePtr = nullptr;
+  BasicBlock *AllocaSpillBlock = nullptr;
 
   /// This would only be true if optimization are enabled.
   bool OptimizeFrame;
@@ -237,9 +271,14 @@ struct LLVM_LIBRARY_VISIBILITY Shape {
   Shape() = default;
   explicit Shape(Function &F, bool OptimizeFrame = false)
       : OptimizeFrame(OptimizeFrame) {
-    buildFrom(F);
+    analyze(F);
+    if (!CoroBegin) {
+      invalidateCoroutine(F);
+      return;
+    }
+    initABI();
+    tidyCoroutine();
   }
-  void buildFrom(Function &F);
 };
 
 } // end namespace coro
diff --git a/llvm/lib/Transforms/Coroutines/Coroutines.cpp b/llvm/lib/Transforms/Coroutines/Coroutines.cpp
index 5cc13a584aef32..c1042b21883f67 100644
--- a/llvm/lib/Transforms/Coroutines/Coroutines.cpp
+++ b/llvm/lib/Transforms/Coroutines/Coroutines.cpp
@@ -177,17 +177,6 @@ void coro::suppressCoroAllocs(LLVMContext &Context,
   }
 }
 
-static void clear(coro::Shape &Shape) {
-  Shape.CoroBegin = nullptr;
-  Shape.CoroEnds.clear();
-  Shape.CoroSizes.clear();
-  Shape.CoroSuspends.clear();
-
-  Shape.FrameTy = nullptr;
-  Shape.FramePtr = nullptr;
-  Shape.AllocaSpillBlock = nullptr;
-}
-
 static CoroSaveInst *createCoroSave(CoroBeginInst *CoroBegin,
                                     CoroSuspendInst *SuspendInst) {
   Module *M = SuspendInst->getModule();
@@ -200,13 +189,12 @@ static CoroSaveInst *createCoroSave(CoroBeginInst *CoroBegin,
 }
 
 // Collect "interesting" coroutine intrinsics.
-void coro::Shape::buildFrom(Function &F) {
+void coro::Shape::analyze(Function &F) {
+  clear();
+
   bool HasFinalSuspend = false;
   bool HasUnwindCoroEnd = false;
   size_t FinalSuspendIndex = 0;
-  clear(*this);
-  SmallVector<CoroFrameInst *, 8> CoroFrames;
-  SmallVector<CoroSaveInst *, 2> UnusedCoroSaves;
 
   for (Instruction &I : instructions(F)) {
     // FIXME: coro_await_suspend_* are not proper `IntrinisicInst`s
@@ -298,8 +286,58 @@ void coro::Shape::buildFrom(Function &F) {
     }
   }
 
-  // If for some reason, we were not able to find coro.begin, bailout.
-  if (!CoroBegin) {
+  // If there is no CoroBegin then this is not a coroutine.
+  if (!CoroBegin)
+    return;
+
+  // Determination of ABI and initializing lowering info
+  auto Id = CoroBegin->getId();
+  auto IntrID = Id->getIntrinsicID();
+  if (IntrID == Intrinsic::coro_id) {
+    ABI = coro::ABI::Switch;
+    SwitchLowering.HasFinalSuspend = HasFinalSuspend;
+    SwitchLowering.HasUnwindCoroEnd = HasUnwindCoroEnd;
+
+    auto SwitchId = getSwitchCoroId();
+    SwitchLowering.ResumeSwitch = nullptr;
+    SwitchLowering.PromiseAlloca = SwitchId->getPromise();
+    SwitchLowering.ResumeEntryBlock = nullptr;
+
+    // Move final suspend to the last element in the CoroSuspends vector.
+    if (SwitchLowering.HasFinalSuspend &&
+        FinalSuspendIndex != CoroSuspends.size() - 1)
+      std::swap(CoroSuspends[FinalSuspendIndex], CoroSuspends.back());
+  } else if (IntrID == Intrinsic::coro_id_async) {
+    ABI = coro::ABI::Async;
+    auto *AsyncId = getAsyncCoroId();
+    AsyncId->checkWellFormed();
+    AsyncLowering.Context = AsyncId->getStorage();
+    AsyncLowering.ContextArgNo = AsyncId->getStorageArgumentIndex();
+    AsyncLowering.ContextHeaderSize = AsyncId->getStorageSize();
+    AsyncLowering.ContextAlignment = AsyncId->getStorageAlignment().value();
+    AsyncLowering.AsyncFuncPointer = AsyncId->getAsyncFunctionPointer();
+    AsyncLowering.AsyncCC = F.getCallingConv();
+  } else if (IntrID == Intrinsic::coro_id_retcon ||
+             IntrID == Intrinsic::coro_id_retcon_once) {
+    ABI = IntrID == Intrinsic::coro_id_retcon ? coro::ABI::Retcon
+                                              : coro::ABI::RetconOnce;
+    auto ContinuationId = getRetconCoroId();
+    ContinuationId->checkWellFormed();
+    auto Prototype = ContinuationId->getPrototype();
+    RetconLowering.ResumePrototype = Prototype;
+    RetconLowering.Alloc = ContinuationId->getAllocFunction();
+    RetconLowering.Dealloc = ContinuationId->getDeallocFunction();
+    RetconLowering.ReturnBlock = nullptr;
+    RetconLowering.IsFrameInlineInStorage = false;
+  } else {
+    llvm_unreachable("coro.begin is not dependent on a coro.id call");
+  }
+}
+
+// If for some reason, we were not able to find coro.begin, bailout.
+void coro::Shape::invalidateCoroutine(Function &F) {
+  assert(!CoroBegin);
+  {
     // Replace coro.frame which are supposed to be lowered to the result of
     // coro.begin with undef.
     auto *Undef = UndefValue::get(PointerType::get(F.getContext(), 0));
@@ -320,21 +358,13 @@ void coro::Shape::buildFrom(Function &F) {
     // Replace all coro.ends with unreachable instruction.
     for (AnyCoroEndInst *CE : CoroEnds)
       changeToUnreachable(CE);
-
-    return;
   }
+}
 
-  auto Id = CoroBegin->getId();
-  switch (auto IdIntrinsic = Id->getIntrinsicID()) {
-  case Intrinsic::coro_id: {
-    auto SwitchId = cast<CoroIdInst>(Id);
-    this->ABI = coro::ABI::Switch;
-    this->SwitchLowering.HasFinalSuspend = HasFinalSuspend;
-    this->SwitchLowering.HasUnwindCoroEnd = HasUnwindCoroEnd;
-    this->SwitchLowering.ResumeSwitch = nullptr;
-    this->SwitchLowering.PromiseAlloca = SwitchId->getPromise();
-    this->SwitchLowering.ResumeEntryBlock = nullptr;
-
+// Perform semantic checking and initialization of the ABI
+void coro::Shape::initABI() {
+  switch (ABI) {
+  case coro::ABI::Switch: {
     for (auto *AnySuspend : CoroSuspends) {
       auto Suspend = dyn_cast<CoroSuspendInst>(AnySuspend);
       if (!Suspend) {
@@ -349,33 +379,11 @@ void coro::Shape::buildFrom(Function &F) {
     }
     break;
   }
-  case Intrinsic::coro_id_async: {
-    auto *AsyncId = cast<CoroIdAsyncInst>(Id);
-    AsyncId->checkWellFormed();
-    this->ABI = coro::ABI::Async;
-    this->AsyncLowering.Context = AsyncId->getStorage();
-    this->AsyncLowering.ContextArgNo = AsyncId->getStorageArgumentIndex();
-    this->AsyncLowering.ContextHeaderSize = AsyncId->getStorageSize();
-    this->AsyncLowering.ContextAlignment =
-        AsyncId->getStorageAlignment().value();
-    this->AsyncLowering.AsyncFuncPointer = AsyncId->getAsyncFunctionPointer();
-    this->AsyncLowering.AsyncCC = F.getCallingConv();
+  case coro::ABI::Async: {
     break;
   };
-  case Intrinsic::coro_id_retcon:
-  case Intrinsic::coro_id_retcon_once: {
-    auto ContinuationId = cast<AnyCoroIdRetconInst>(Id);
-    ContinuationId->checkWellFormed();
-    this->ABI = (IdIntrinsic == Intrinsic::coro_id_retcon
-                  ? coro::ABI::Retcon
-                  : coro::ABI::RetconOnce);
-    auto Prototype = ContinuationId->getPrototype();
-    this->RetconLowering.ResumePrototype = Prototype;
-    this->RetconLowering.Alloc = ContinuationId->getAllocFunction();
-    this->RetconLowering.Dealloc = ContinuationId->getDeallocFunction();
-    this->RetconLowering.ReturnBlock = nullptr;
-    this->RetconLowering.IsFrameInlineInStorage = false;
-
+  case coro::ABI::Retcon:
+  case coro::ABI::RetconOnce: {
     // Determine the result value types, and make sure they match up with
     // the values passed to the suspends.
     auto ResultTys = getRetconResultTypes();
@@ -408,7 +416,7 @@ void coro::Shape::buildFrom(Function &F) {
 
 #ifndef NDEBUG
           Suspend->dump();
-          Prototype->getFunctionType()->dump();
+          RetconLowering.ResumePrototype->getFunctionType()->dump();
 #endif
           report_fatal_error("argument to coro.suspend.retcon does not "
                              "match corresponding prototype function result");
@@ -417,14 +425,14 @@ void coro::Shape::buildFrom(Function &F) {
       if (SI != SE || RI != RE) {
 #ifndef NDEBUG
         Suspend->dump();
-        Prototype->getFunctionType()->dump();
+        RetconLowering.ResumePrototype->getFunctionType()->dump();
 #endif
         report_fatal_error("wrong number of arguments to coro.suspend.retcon");
       }
 
       // Check that the result type of the suspend matches the resume types.
       Type *SResultTy = Suspend->getType();
-      ArrayRef<Type*> SuspendResultTys;
+      ArrayRef<Type *> SuspendResultTys;
       if (SResultTy->isVoidTy()) {
         // leave as empty array
       } else if (auto SResultStructTy = dyn_cast<StructType>(SResultTy)) {
@@ -436,7 +444,7 @@ void coro::Shape::buildFrom(Function &F) {
       if (SuspendResultTys.size() != ResumeTys.size()) {
 #ifndef NDEBUG
         Suspend->dump();
-        Prototype->getFunctionType()->dump();
+        RetconLowering.ResumePrototype->getFunctionType()->dump();
 #endif
         report_fatal_error("wrong number of results from coro.suspend.retcon");
       }
@@ -444,7 +452,7 @@ void coro::Shape::buildFrom(Function &F) {
         if (SuspendResultTys[I] != ResumeTys[I]) {
 #ifndef NDEBUG
           Suspend->dump();
-          Prototype->getFunctionType()->dump();
+          RetconLowering.ResumePrototype->getFunctionType()->dump();
 #endif
           report_fatal_error("result from coro.suspend.retcon does not "
                              "match corresponding prototype function param");
@@ -453,23 +461,18 @@ void coro::Shape::buildFrom(Function &F) {
     }
     break;
   }
-
   default:
     llvm_unreachable("coro.begin is not dependent on a coro.id call");
   }
+}
 
+void coro::Shape::tidyCoroutine() {
   // The coro.free intrinsic is always lowered to the result of coro.begin.
   for (CoroFrameInst *CF : CoroFrames) {
     CF->replaceAllUsesWith(CoroBegin);
     CF->eraseFromParent();
   }
 
-  // Move final suspend to be the last element in the CoroSuspends vector.
-  if (ABI == coro::ABI::Switch &&
-      SwitchLowering.HasFinalSuspend &&
-      FinalSuspendIndex != CoroSuspends.size() - 1)
-    std::swap(CoroSuspends[FinalSuspendIndex], CoroSuspends.back());
-
   // Remove orphaned coro.saves.
   for (CoroSaveInst *CoroSave : UnusedCoroSaves)
     CoroSave->eraseFromParent();

[ChuanqiXu9]

Suggested change

	void tidyCoroutine();
	void cleanCoroutine();

nit

[TylerNowicki]

Done

[ChuanqiXu9]

It comes from the original buildFrom. But as far as I know, they are not invalided?

[TylerNowicki]

Right, they were on the stack in buildFrame. In my commit message I called them out because they are now part of the Shape struct itself.
 
The methods both will erase the Insts in the SmallVectors. For example, invalidateCoroutine() erases CoroFrameInsts within the CoroFrames SmallVector using eraseFromParent. The method does not then clear CoroFrames. Using an inst after eraseFromParent() can be a bit dangerous, for example, calling eraseFromParent again will trigger an exception. I realized that I should clear the SmallVectors after they are handled by invalidateCoroutine() and clearCoroutine(). I updated the patch to do this (sorry about the accidental force push).

[ChuanqiXu9]

Let's following the original style to use switch.

[TylerNowicki]

Done

[ChuanqiXu9]

nit: The name looks slightly odd to me. It looks like they will be cleared in invalidateCoroutine(). This is conflicting with the comment.

And also I like to rename them to ToBeInvalidatedCoroFrames and so on. Then in analyze, remain the original CoroFrames and UnusedCoroSaves, let's only assign values to ToBeInvalidatedCoroFrames (and so on) when it is invalid actually.

[TylerNowicki]

I see you approved the change but I don't see a lgtm? Do you want me to make these nit changes in another patch?

About the nit changes, let me see if I have understood what you want.

1. Keep CoroShape::CoroFrames and CoroShape::UnusedCoroSaves
2. Declare CoroShape::ToBeInvalidatedCoroFrames and CoroShape::ToBeInvalidatedUnusedCoroSaves
3. In CoroShape::invalidateCoroutine() and CoroShape::cleanCoroutine() instead of erasing the Instructions, move them to CoroShape::ToBeInvalidated{CoroFrames,UnusedCoroSaves}
4. Add code to CoroShape::clear() or CoroShape::~CoroShape() to erase the Instructions in CoroShape::ToBeInvalidated{CoroFrames,UnusedCoroSaves}

Is that right?

Should we erase instructions in CoroShape::clear() or CoroShape::~CoroShape()?

Also I noticed that CoroShape::invalidateCoroutine() also erases CoroSuspends and CoroSaves (those that are used). Should we also delay erasing these until CoroShape::clear() or CoroShape::~CoroShape()?

[TylerNowicki]

My original intent to the changes was to ensure that we could create CoroShape objects from the Function without changing the Function's instructions in anyway. This makes CoroShape a little more useful because right now if a method needs the info we have to rely on it being passed into the method, and it cannot be used by other passes.

[ChuanqiXu9]

I see you approved the change but I don't see a lgtm? Do you want me to make these nit changes in another patch?

Approve with a nit means, you can land it as if now, but I'll be more happier if you address the nit comments.

1. Keep CoroShape::CoroFrames and CoroShape::UnusedCoroSaves

I mean, keep them as is, a local variables, IIUC.

My original intent to the changes was to ensure that we could create CoroShape objects from the Function without changing the Function's instructions in anyway.

The idea sounds good. But it will erase the instructions, right? I don't feel bad with this since this is the current behavior of the building process of CoroShape. But it will be good to make it a pure analysis process.

[TylerNowicki]

Okay, I can see I made an unintentional functional change to CoroShape (changing the building process). I will make CoroFrames and UnusedCoroSaves local variables again. I have some ideas about how to make a pure analysis. I will pursue that after landing the ABI object patches.