[Coroutines] Split buildCoroutineFrame into normalization and frame building #108076

TylerNowicki · 2024-09-10T19:55:01Z

Split buildCoroutineFrame into code related to normalization and code related to actually building the coroutine frame.
This will enable future specialization of buildCoroutineFrame for different ABIs while the normalization can be done by splitCoroutine prior to calling buildCoroutineFrame.

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

* Split buildCoroutineFrame into code related to normalization and code related to actually building the coroutine frame. * This will enable future specialization of buildCoroutineFrame for different ABIs.

llvmbot · 2024-09-10T19:55:20Z

@llvm/pr-subscribers-llvm-transforms

Author: Tyler Nowicki (TylerNowicki)

Changes

Split buildCoroutineFrame into code related to normalization and code related to actually building the coroutine frame.
This will enable future specialization of buildCoroutineFrame for different ABIs while the normalization can be done by splitCoroutine prior to calling buildCoroutineFrame.

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/108076.diff

3 Files Affected:

(modified) llvm/lib/Transforms/Coroutines/CoroFrame.cpp (+8-6)
(modified) llvm/lib/Transforms/Coroutines/CoroInternal.h (+3-1)
(modified) llvm/lib/Transforms/Coroutines/CoroSplit.cpp (+2-1)

diff --git a/llvm/lib/Transforms/Coroutines/CoroFrame.cpp b/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
index 4b76fc79361008..8ee4bfa3b888df 100644
--- a/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
@@ -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();

llvmbot · 2024-09-10T19:55:20Z

@llvm/pr-subscribers-coroutines

Author: Tyler Nowicki (TylerNowicki)

Changes

Split buildCoroutineFrame into code related to normalization and code related to actually building the coroutine frame.
This will enable future specialization of buildCoroutineFrame for different ABIs while the normalization can be done by splitCoroutine prior to calling buildCoroutineFrame.

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/108076.diff

3 Files Affected:

(modified) llvm/lib/Transforms/Coroutines/CoroFrame.cpp (+8-6)
(modified) llvm/lib/Transforms/Coroutines/CoroInternal.h (+3-1)
(modified) llvm/lib/Transforms/Coroutines/CoroSplit.cpp (+2-1)

diff --git a/llvm/lib/Transforms/Coroutines/CoroFrame.cpp b/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
index 4b76fc79361008..8ee4bfa3b888df 100644
--- a/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
@@ -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();

yuxuanchen1997

Overall I think a neutral/positive change. One thing I am worried was that how we are going to distinguish a "normalized" vs a "unnormalized" coroutine? i.e. How do we find out that we missed a normalization here?

TylerNowicki · 2024-09-10T20:21:12Z

Overall I think a neutral/positive change. One thing I am worried was that how we are going to distinguish a "normalized" vs a "unnormalized" coroutine? i.e. How do we find out that we missed a normalization here?

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.

TylerNowicki · 2024-09-10T21:47:00Z

Overall I think a neutral/positive change. One thing I am worried was that how we are going to distinguish a "normalized" vs a "unnormalized" coroutine? i.e. How do we find out that we missed a normalization here?

I posted a PR to add a couple of asserts that check for normalization #108096. The PR includes this change #108076 (I posted it too quickly), I will rebase after this change merges.

ChuanqiXu9

LGTM

ChuanqiXu9 · 2024-09-11T01:49:15Z

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

+  if (isSuspendBlock(BB))
+    return true;


nit: generally we don't unrelated changes.

ChuanqiXu9 · 2024-09-11T01:51:17Z

I think the concept normalization is pretty narrow in CoroFrame itself (or even buildCoroutineFrame itself). It is simply some transformations to make it easier to construct and use the frame. So I didn't take it as a formal concept.

* 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.

[Coroutines] Split buildCoroutineFrame

34559ad

* Split buildCoroutineFrame into code related to normalization and code related to actually building the coroutine frame. * This will enable future specialization of buildCoroutineFrame for different ABIs.

TylerNowicki added coroutines C++20 coroutines llvm:transforms labels Sep 10, 2024

TylerNowicki requested review from ChuanqiXu9 and yuxuanchen1997 September 10, 2024 19:55

TylerNowicki self-assigned this Sep 10, 2024

yuxuanchen1997 approved these changes Sep 10, 2024

View reviewed changes

ChuanqiXu9 approved these changes Sep 11, 2024

View reviewed changes

TylerNowicki merged commit 9a9f155 into llvm:main Sep 11, 2024

TylerNowicki mentioned this pull request Sep 11, 2024

[Coroutines] Verify normalization was not missed #108096

Merged

TylerNowicki deleted the users/tylernowicki/coro-refactor3 branch September 11, 2024 14:38

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[Coroutines] Split buildCoroutineFrame into normalization and frame building #108076

[Coroutines] Split buildCoroutineFrame into normalization and frame building #108076

Uh oh!

TylerNowicki commented Sep 10, 2024

Uh oh!

llvmbot commented Sep 10, 2024

Uh oh!

llvmbot commented Sep 10, 2024

Uh oh!

yuxuanchen1997 left a comment

Uh oh!

TylerNowicki commented Sep 10, 2024

Uh oh!

TylerNowicki commented Sep 10, 2024 •

edited

Loading

Uh oh!

ChuanqiXu9 left a comment

Uh oh!

ChuanqiXu9 Sep 11, 2024

Uh oh!

ChuanqiXu9 commented Sep 11, 2024

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

[Coroutines] Split buildCoroutineFrame into normalization and frame building #108076

[Coroutines] Split buildCoroutineFrame into normalization and frame building #108076

Uh oh!

Conversation

TylerNowicki commented Sep 10, 2024

Uh oh!

llvmbot commented Sep 10, 2024

Uh oh!

llvmbot commented Sep 10, 2024

Uh oh!

yuxuanchen1997 left a comment

Choose a reason for hiding this comment

Uh oh!

TylerNowicki commented Sep 10, 2024

Uh oh!

TylerNowicki commented Sep 10, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ChuanqiXu9 left a comment

Choose a reason for hiding this comment

Uh oh!

ChuanqiXu9 Sep 11, 2024

Choose a reason for hiding this comment

Uh oh!

ChuanqiXu9 commented Sep 11, 2024

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

TylerNowicki commented Sep 10, 2024 •

edited

Loading