[Coroutines] Verify normalization was not missed

[TylerNowicki]

• Add asserts to verify normalization of the coroutine happened.
• This will be important when normalization becomes part of an ABI object.

--- From a previous discussion here #108076

Normalization performs these important steps:

split around each suspend, this adds BBs before/after each suspend so each suspend now exists in its own block.
split around coro.end (similar to above)
break critical edges and add single-edge phis in the new blocks (also removing other single-edge phis).
Each of these things can individually be tested
A) Check that each suspend is the only inst in its BB
B) Check that coro.end is the only inst in its BB
C) Check that each edge of a multi-edge phis is preceded by single-edge phi in an immediate pred

For 1) and 2) I believe the purpose of the transform is in part for suspend crossing info's analysis so it can specifically 'mark' the suspend blocks and identify the end of the coroutine. There are some existing places within suspend crossing info that visit the CoroSuspends and CoroEnds so we could check A) and B) there.

For 3) I believe the purpose of this transform is for insertSpills to work properly. Infact there is already a check for the result of this transform!

assert(PN->getNumIncomingValues() == 1 &&
"unexpected number of incoming "
"values in the PHINode");

I think to verify the result of normalization we just need to add checks A) and B) to suspend crossing info.

[llvmbot]

@llvm/pr-subscribers-coroutines

@llvm/pr-subscribers-llvm-transforms

Author: Tyler Nowicki (TylerNowicki)

Changes

• Add asserts to verify normalization of the coroutine happened.
• This will be important if/when normalization becomes part of an ABI object.

---

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

4 Files Affected:

• (modified) llvm/lib/Transforms/Coroutines/CoroFrame.cpp (+9-7)
• (modified) llvm/lib/Transforms/Coroutines/CoroInternal.h (+3-1)
• (modified) llvm/lib/Transforms/Coroutines/CoroSplit.cpp (+2-1)
• (modified) llvm/lib/Transforms/Coroutines/SuspendCrossingInfo.cpp (+11-1)

diff --git a/llvm/lib/Transforms/Coroutines/CoroFrame.cpp b/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
index 4b76fc79361008..b792835159a8a3 100644
--- a/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
@@ -1284,7 +1284,7 @@ static void insertSpills(const FrameDataInfo &FrameData, coro::Shape &Shape) {
 
       // If we have a single edge PHINode, remove it and replace it with a
       // reload from the coroutine frame. (We already took care of multi edge
-      // PHINodes by rewriting them in the rewritePHIs function).
+      // PHINodes by normalizing them in the rewritePHIs function).
       if (auto *PN = dyn_cast<PHINode>(U)) {
         assert(PN->getNumIncomingValues() == 1 &&
                "unexpected number of incoming "
@@ -1754,7 +1754,8 @@ static bool willLeaveFunctionImmediatelyAfter(BasicBlock *BB,
   if (depth == 0) return false;
 
   // If this is a suspend block, we're about to exit the resumption function.
-  if (isSuspendBlock(BB)) return true;
+  if (isSuspendBlock(BB))
+    return true;
 
   // Recurse into the successors.
   for (auto *Succ : successors(BB)) {
@@ -2288,9 +2289,8 @@ static void doRematerializations(
   rewriteMaterializableInstructions(AllRemats);
 }
 
-void coro::buildCoroutineFrame(
-    Function &F, Shape &Shape, TargetTransformInfo &TTI,
-    const std::function<bool(Instruction &)> &MaterializableCallback) {
+void coro::normalizeCoroutine(Function &F, coro::Shape &Shape,
+                              TargetTransformInfo &TTI) {
   // Don't eliminate swifterror in async functions that won't be split.
   if (Shape.ABI != coro::ABI::Async || !Shape.CoroSuspends.empty())
     eliminateSwiftError(F, Shape);
@@ -2337,10 +2337,12 @@ void coro::buildCoroutineFrame(
   // Transforms multi-edge PHI Nodes, so that any value feeding into a PHI will
   // never have its definition separated from the PHI by the suspend point.
   rewritePHIs(F);
+}
 
-  // Build suspend crossing info.
+void coro::buildCoroutineFrame(
+    Function &F, Shape &Shape,
+    const std::function<bool(Instruction &)> &MaterializableCallback) {
   SuspendCrossingInfo Checker(F, Shape.CoroSuspends, Shape.CoroEnds);
-
   doRematerializations(F, Checker, MaterializableCallback);
 
   const DominatorTree DT(F);
diff --git a/llvm/lib/Transforms/Coroutines/CoroInternal.h b/llvm/lib/Transforms/Coroutines/CoroInternal.h
index be86f96525b677..698c21a797420a 100644
--- a/llvm/lib/Transforms/Coroutines/CoroInternal.h
+++ b/llvm/lib/Transforms/Coroutines/CoroInternal.h
@@ -281,8 +281,10 @@ struct LLVM_LIBRARY_VISIBILITY Shape {
 };
 
 bool defaultMaterializable(Instruction &V);
+void normalizeCoroutine(Function &F, coro::Shape &Shape,
+                        TargetTransformInfo &TTI);
 void buildCoroutineFrame(
-    Function &F, Shape &Shape, TargetTransformInfo &TTI,
+    Function &F, Shape &Shape,
     const std::function<bool(Instruction &)> &MaterializableCallback);
 CallInst *createMustTailCall(DebugLoc Loc, Function *MustTailCallFn,
                              TargetTransformInfo &TTI,
diff --git a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
index 494c4d632de95f..dc3829d7f28eb1 100644
--- a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
@@ -2030,7 +2030,8 @@ splitCoroutine(Function &F, SmallVectorImpl<Function *> &Clones,
   lowerAwaitSuspends(F, Shape);
 
   simplifySuspendPoints(Shape);
-  buildCoroutineFrame(F, Shape, TTI, MaterializableCallback);
+  normalizeCoroutine(F, Shape, TTI);
+  buildCoroutineFrame(F, Shape, MaterializableCallback);
   replaceFrameSizeAndAlignment(Shape);
   bool isNoSuspendCoroutine = Shape.CoroSuspends.empty();
 
diff --git a/llvm/lib/Transforms/Coroutines/SuspendCrossingInfo.cpp b/llvm/lib/Transforms/Coroutines/SuspendCrossingInfo.cpp
index 6b0dc126d5471a..84699e653db602 100644
--- a/llvm/lib/Transforms/Coroutines/SuspendCrossingInfo.cpp
+++ b/llvm/lib/Transforms/Coroutines/SuspendCrossingInfo.cpp
@@ -165,8 +165,13 @@ SuspendCrossingInfo::SuspendCrossingInfo(
   // Mark all CoroEnd Blocks. We do not propagate Kills beyond coro.ends as
   // the code beyond coro.end is reachable during initial invocation of the
   // coroutine.
-  for (auto *CE : CoroEnds)
+  for (auto *CE : CoroEnds) {
+    // Verify CoroEnd was normalized
+    assert(CE->getParent()->getFirstInsertionPt() == CE->getIterator() &&
+           CE->getParent()->size() <= 2 && "CoroEnd must be in its own BB");
+
     getBlockData(CE->getParent()).End = true;
+  }
 
   // Mark all suspend blocks and indicate that they kill everything they
   // consume. Note, that crossing coro.save also requires a spill, as any code
@@ -179,6 +184,11 @@ SuspendCrossingInfo::SuspendCrossingInfo(
     B.Kills |= B.Consumes;
   };
   for (auto *CSI : CoroSuspends) {
+    // Verify CoroSuspend was normalized
+    assert(CSI->getParent()->getFirstInsertionPt() == CSI->getIterator() &&
+           CSI->getParent()->size() <= 2 &&
+           "CoroSuspend must be in its own BB");
+
     markSuspendBlock(CSI);
     if (auto *Save = CSI->getCoroSave())
       markSuspendBlock(Save);

[TylerNowicki]

This PR is ready for review now

#108076