[Coroutines] Move materialization code into its own utils #108240
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
* Move materialization out of CoroFrame to MaterializationUtils.h * Move spill related utilities that were used by materialization to SpillUtils * Move isSuspendBlock (needed by materialization) to CoroInternal
|
@llvm/pr-subscribers-llvm-transforms @llvm/pr-subscribers-coroutines Author: Tyler Nowicki (TylerNowicki) Changes
See RFC for more info: https://discourse.llvm.org/t/rfc-abi-objects-for-coroutines/81057 Patch is 29.97 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/108240.diff 8 Files Affected:
diff --git a/llvm/lib/Transforms/Coroutines/CMakeLists.txt b/llvm/lib/Transforms/Coroutines/CMakeLists.txt
index c6508174a7f109..46ef5cd4e8cfe8 100644
--- a/llvm/lib/Transforms/Coroutines/CMakeLists.txt
+++ b/llvm/lib/Transforms/Coroutines/CMakeLists.txt
@@ -9,6 +9,7 @@ add_llvm_component_library(LLVMCoroutines
CoroSplit.cpp
SuspendCrossingInfo.cpp
SpillUtils.cpp
+ MaterializationUtils.cpp
ADDITIONAL_HEADER_DIRS
${LLVM_MAIN_INCLUDE_DIR}/llvm/Transforms/Coroutines
diff --git a/llvm/lib/Transforms/Coroutines/CoroFrame.cpp b/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
index 8ee4bfa3b888df..b74c9f01cd2395 100644
--- a/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
@@ -16,10 +16,10 @@
//===----------------------------------------------------------------------===//
#include "CoroInternal.h"
+#include "MaterializationUtils.h"
#include "SpillUtils.h"
#include "SuspendCrossingInfo.h"
#include "llvm/ADT/BitVector.h"
-#include "llvm/ADT/PostOrderIterator.h"
#include "llvm/ADT/ScopeExit.h"
#include "llvm/ADT/SmallString.h"
#include "llvm/Analysis/StackLifetime.h"
@@ -36,135 +36,12 @@
#include "llvm/Transforms/Utils/Local.h"
#include "llvm/Transforms/Utils/PromoteMemToReg.h"
#include <algorithm>
-#include <deque>
#include <optional>
using namespace llvm;
extern cl::opt<bool> UseNewDbgInfoFormat;
-// The "coro-suspend-crossing" flag is very noisy. There is another debug type,
-// "coro-frame", which results in leaner debug spew.
-#define DEBUG_TYPE "coro-suspend-crossing"
-
-namespace {
-
-// RematGraph is used to construct a DAG for rematerializable instructions
-// When the constructor is invoked with a candidate instruction (which is
-// materializable) it builds a DAG of materializable instructions from that
-// point.
-// Typically, for each instruction identified as re-materializable across a
-// suspend point, a RematGraph will be created.
-struct RematGraph {
- // Each RematNode in the graph contains the edges to instructions providing
- // operands in the current node.
- struct RematNode {
- Instruction *Node;
- SmallVector<RematNode *> Operands;
- RematNode() = default;
- RematNode(Instruction *V) : Node(V) {}
- };
-
- RematNode *EntryNode;
- using RematNodeMap =
- SmallMapVector<Instruction *, std::unique_ptr<RematNode>, 8>;
- RematNodeMap Remats;
- const std::function<bool(Instruction &)> &MaterializableCallback;
- SuspendCrossingInfo &Checker;
-
- RematGraph(const std::function<bool(Instruction &)> &MaterializableCallback,
- Instruction *I, SuspendCrossingInfo &Checker)
- : MaterializableCallback(MaterializableCallback), Checker(Checker) {
- std::unique_ptr<RematNode> FirstNode = std::make_unique<RematNode>(I);
- EntryNode = FirstNode.get();
- std::deque<std::unique_ptr<RematNode>> WorkList;
- addNode(std::move(FirstNode), WorkList, cast<User>(I));
- while (WorkList.size()) {
- std::unique_ptr<RematNode> N = std::move(WorkList.front());
- WorkList.pop_front();
- addNode(std::move(N), WorkList, cast<User>(I));
- }
- }
-
- void addNode(std::unique_ptr<RematNode> NUPtr,
- std::deque<std::unique_ptr<RematNode>> &WorkList,
- User *FirstUse) {
- RematNode *N = NUPtr.get();
- if (Remats.count(N->Node))
- return;
-
- // We haven't see this node yet - add to the list
- Remats[N->Node] = std::move(NUPtr);
- for (auto &Def : N->Node->operands()) {
- Instruction *D = dyn_cast<Instruction>(Def.get());
- if (!D || !MaterializableCallback(*D) ||
- !Checker.isDefinitionAcrossSuspend(*D, FirstUse))
- continue;
-
- if (Remats.count(D)) {
- // Already have this in the graph
- N->Operands.push_back(Remats[D].get());
- continue;
- }
-
- bool NoMatch = true;
- for (auto &I : WorkList) {
- if (I->Node == D) {
- NoMatch = false;
- N->Operands.push_back(I.get());
- break;
- }
- }
- if (NoMatch) {
- // Create a new node
- std::unique_ptr<RematNode> ChildNode = std::make_unique<RematNode>(D);
- N->Operands.push_back(ChildNode.get());
- WorkList.push_back(std::move(ChildNode));
- }
- }
- }
-
-#if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
- static std::string getBasicBlockLabel(const BasicBlock *BB) {
- if (BB->hasName())
- return BB->getName().str();
-
- std::string S;
- raw_string_ostream OS(S);
- BB->printAsOperand(OS, false);
- return OS.str().substr(1);
- }
-
- void dump() const {
- dbgs() << "Entry (";
- dbgs() << getBasicBlockLabel(EntryNode->Node->getParent());
- dbgs() << ") : " << *EntryNode->Node << "\n";
- for (auto &E : Remats) {
- dbgs() << *(E.first) << "\n";
- for (RematNode *U : E.second->Operands)
- dbgs() << " " << *U->Node << "\n";
- }
- }
-#endif
-};
-} // end anonymous namespace
-
-namespace llvm {
-
-template <> struct GraphTraits<RematGraph *> {
- using NodeRef = RematGraph::RematNode *;
- using ChildIteratorType = RematGraph::RematNode **;
-
- static NodeRef getEntryNode(RematGraph *G) { return G->EntryNode; }
- static ChildIteratorType child_begin(NodeRef N) {
- return N->Operands.begin();
- }
- static ChildIteratorType child_end(NodeRef N) { return N->Operands.end(); }
-};
-
-} // end namespace llvm
-
-#undef DEBUG_TYPE // "coro-suspend-crossing"
#define DEBUG_TYPE "coro-frame"
namespace {
@@ -268,15 +145,6 @@ static void dumpSpills(StringRef Title, const coro::SpillInfo &Spills) {
I->dump();
}
}
-static void dumpRemats(
- StringRef Title,
- const SmallMapVector<Instruction *, std::unique_ptr<RematGraph>, 8> &RM) {
- dbgs() << "------------- " << Title << "--------------\n";
- for (const auto &E : RM) {
- E.second->dump();
- dbgs() << "--\n";
- }
-}
static void dumpAllocas(const SmallVectorImpl<coro::AllocaInfo> &Allocas) {
dbgs() << "------------- Allocas --------------\n";
@@ -1634,93 +1502,6 @@ static void rewritePHIs(Function &F) {
rewritePHIs(*BB);
}
-/// Default materializable callback
-// Check for instructions that we can recreate on resume as opposed to spill
-// the result into a coroutine frame.
-bool coro::defaultMaterializable(Instruction &V) {
- return (isa<CastInst>(&V) || isa<GetElementPtrInst>(&V) ||
- isa<BinaryOperator>(&V) || isa<CmpInst>(&V) || isa<SelectInst>(&V));
-}
-
-// For each instruction identified as materializable across the suspend point,
-// and its associated DAG of other rematerializable instructions,
-// recreate the DAG of instructions after the suspend point.
-static void rewriteMaterializableInstructions(
- const SmallMapVector<Instruction *, std::unique_ptr<RematGraph>, 8>
- &AllRemats) {
- // This has to be done in 2 phases
- // Do the remats and record the required defs to be replaced in the
- // original use instructions
- // Once all the remats are complete, replace the uses in the final
- // instructions with the new defs
- typedef struct {
- Instruction *Use;
- Instruction *Def;
- Instruction *Remat;
- } ProcessNode;
-
- SmallVector<ProcessNode> FinalInstructionsToProcess;
-
- for (const auto &E : AllRemats) {
- Instruction *Use = E.first;
- Instruction *CurrentMaterialization = nullptr;
- RematGraph *RG = E.second.get();
- ReversePostOrderTraversal<RematGraph *> RPOT(RG);
- SmallVector<Instruction *> InstructionsToProcess;
-
- // If the target use is actually a suspend instruction then we have to
- // insert the remats into the end of the predecessor (there should only be
- // one). This is so that suspend blocks always have the suspend instruction
- // as the first instruction.
- auto InsertPoint = &*Use->getParent()->getFirstInsertionPt();
- if (isa<AnyCoroSuspendInst>(Use)) {
- BasicBlock *SuspendPredecessorBlock =
- Use->getParent()->getSinglePredecessor();
- assert(SuspendPredecessorBlock && "malformed coro suspend instruction");
- InsertPoint = SuspendPredecessorBlock->getTerminator();
- }
-
- // Note: skip the first instruction as this is the actual use that we're
- // rematerializing everything for.
- auto I = RPOT.begin();
- ++I;
- for (; I != RPOT.end(); ++I) {
- Instruction *D = (*I)->Node;
- CurrentMaterialization = D->clone();
- CurrentMaterialization->setName(D->getName());
- CurrentMaterialization->insertBefore(InsertPoint);
- InsertPoint = CurrentMaterialization;
-
- // Replace all uses of Def in the instructions being added as part of this
- // rematerialization group
- for (auto &I : InstructionsToProcess)
- I->replaceUsesOfWith(D, CurrentMaterialization);
-
- // Don't replace the final use at this point as this can cause problems
- // for other materializations. Instead, for any final use that uses a
- // define that's being rematerialized, record the replace values
- for (unsigned i = 0, E = Use->getNumOperands(); i != E; ++i)
- if (Use->getOperand(i) == D) // Is this operand pointing to oldval?
- FinalInstructionsToProcess.push_back(
- {Use, D, CurrentMaterialization});
-
- InstructionsToProcess.push_back(CurrentMaterialization);
- }
- }
-
- // Finally, replace the uses with the defines that we've just rematerialized
- for (auto &R : FinalInstructionsToProcess) {
- if (auto *PN = dyn_cast<PHINode>(R.Use)) {
- assert(PN->getNumIncomingValues() == 1 && "unexpected number of incoming "
- "values in the PHINode");
- PN->replaceAllUsesWith(R.Remat);
- PN->eraseFromParent();
- continue;
- }
- R.Use->replaceUsesOfWith(R.Def, R.Remat);
- }
-}
-
// Splits the block at a particular instruction unless it is the first
// instruction in the block with a single predecessor.
static BasicBlock *splitBlockIfNotFirst(Instruction *I, const Twine &Name) {
@@ -1741,10 +1522,6 @@ static void splitAround(Instruction *I, const Twine &Name) {
splitBlockIfNotFirst(I->getNextNode(), "After" + Name);
}
-static bool isSuspendBlock(BasicBlock *BB) {
- return isa<AnyCoroSuspendInst>(BB->front());
-}
-
/// After we split the coroutine, will the given basic block be along
/// an obvious exit path for the resumption function?
static bool willLeaveFunctionImmediatelyAfter(BasicBlock *BB,
@@ -1754,7 +1531,7 @@ 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))
+ if (coro::isSuspendBlock(BB))
return true;
// Recurse into the successors.
@@ -1995,7 +1772,8 @@ static void sinkLifetimeStartMarkers(Function &F, coro::Shape &Shape,
DomSet.insert(&F.getEntryBlock());
for (auto *CSI : Shape.CoroSuspends) {
BasicBlock *SuspendBlock = CSI->getParent();
- assert(isSuspendBlock(SuspendBlock) && SuspendBlock->getSingleSuccessor() &&
+ assert(coro::isSuspendBlock(SuspendBlock) &&
+ SuspendBlock->getSingleSuccessor() &&
"should have split coro.suspend into its own block");
DomSet.insert(SuspendBlock->getSingleSuccessor());
}
@@ -2227,68 +2005,6 @@ void coro::salvageDebugInfo(
}
}
-static void doRematerializations(
- Function &F, SuspendCrossingInfo &Checker,
- const std::function<bool(Instruction &)> &MaterializableCallback) {
- if (F.hasOptNone())
- return;
-
- coro::SpillInfo Spills;
-
- // See if there are materializable instructions across suspend points
- // We record these as the starting point to also identify materializable
- // defs of uses in these operations
- for (Instruction &I : instructions(F)) {
- if (!MaterializableCallback(I))
- continue;
- for (User *U : I.users())
- if (Checker.isDefinitionAcrossSuspend(I, U))
- Spills[&I].push_back(cast<Instruction>(U));
- }
-
- // Process each of the identified rematerializable instructions
- // and add predecessor instructions that can also be rematerialized.
- // This is actually a graph of instructions since we could potentially
- // have multiple uses of a def in the set of predecessor instructions.
- // The approach here is to maintain a graph of instructions for each bottom
- // level instruction - where we have a unique set of instructions (nodes)
- // and edges between them. We then walk the graph in reverse post-dominator
- // order to insert them past the suspend point, but ensure that ordering is
- // correct. We also rely on CSE removing duplicate defs for remats of
- // different instructions with a def in common (rather than maintaining more
- // complex graphs for each suspend point)
-
- // We can do this by adding new nodes to the list for each suspend
- // point. Then using standard GraphTraits to give a reverse post-order
- // traversal when we insert the nodes after the suspend
- SmallMapVector<Instruction *, std::unique_ptr<RematGraph>, 8> AllRemats;
- for (auto &E : Spills) {
- for (Instruction *U : E.second) {
- // Don't process a user twice (this can happen if the instruction uses
- // more than one rematerializable def)
- if (AllRemats.count(U))
- continue;
-
- // Constructor creates the whole RematGraph for the given Use
- auto RematUPtr =
- std::make_unique<RematGraph>(MaterializableCallback, U, Checker);
-
- LLVM_DEBUG(dbgs() << "***** Next remat group *****\n";
- ReversePostOrderTraversal<RematGraph *> RPOT(RematUPtr.get());
- for (auto I = RPOT.begin(); I != RPOT.end();
- ++I) { (*I)->Node->dump(); } dbgs()
- << "\n";);
-
- AllRemats[U] = std::move(RematUPtr);
- }
- }
-
- // Rewrite materializable instructions to be materialized at the use
- // point.
- LLVM_DEBUG(dumpRemats("Materializations", AllRemats));
- rewriteMaterializableInstructions(AllRemats);
-}
-
void coro::normalizeCoroutine(Function &F, coro::Shape &Shape,
TargetTransformInfo &TTI) {
// Don't eliminate swifterror in async functions that won't be split.
@@ -2324,8 +2040,8 @@ void coro::normalizeCoroutine(Function &F, coro::Shape &Shape,
IRBuilder<> Builder(AsyncEnd);
SmallVector<Value *, 8> Args(AsyncEnd->args());
auto Arguments = ArrayRef<Value *>(Args).drop_front(3);
- auto *Call = createMustTailCall(AsyncEnd->getDebugLoc(), MustTailCallFn,
- TTI, Arguments, Builder);
+ auto *Call = coro::createMustTailCall(
+ AsyncEnd->getDebugLoc(), MustTailCallFn, TTI, Arguments, Builder);
splitAround(Call, "MustTailCall.Before.CoroEnd");
}
}
diff --git a/llvm/lib/Transforms/Coroutines/CoroInternal.h b/llvm/lib/Transforms/Coroutines/CoroInternal.h
index 698c21a797420a..891798f53b2d00 100644
--- a/llvm/lib/Transforms/Coroutines/CoroInternal.h
+++ b/llvm/lib/Transforms/Coroutines/CoroInternal.h
@@ -21,6 +21,7 @@ class CallGraph;
namespace coro {
+bool isSuspendBlock(BasicBlock *BB);
bool declaresAnyIntrinsic(const Module &M);
bool declaresIntrinsics(const Module &M,
const std::initializer_list<StringRef>);
diff --git a/llvm/lib/Transforms/Coroutines/Coroutines.cpp b/llvm/lib/Transforms/Coroutines/Coroutines.cpp
index be257339e0ac49..cdc442bc819c37 100644
--- a/llvm/lib/Transforms/Coroutines/Coroutines.cpp
+++ b/llvm/lib/Transforms/Coroutines/Coroutines.cpp
@@ -100,6 +100,10 @@ static bool isCoroutineIntrinsicName(StringRef Name) {
}
#endif
+bool coro::isSuspendBlock(BasicBlock *BB) {
+ return isa<AnyCoroSuspendInst>(BB->front());
+}
+
bool coro::declaresAnyIntrinsic(const Module &M) {
for (StringRef Name : CoroIntrinsics) {
assert(isCoroutineIntrinsicName(Name) && "not a coroutine intrinsic");
diff --git a/llvm/lib/Transforms/Coroutines/MaterializationUtils.cpp b/llvm/lib/Transforms/Coroutines/MaterializationUtils.cpp
new file mode 100644
index 00000000000000..708e8734175f93
--- /dev/null
+++ b/llvm/lib/Transforms/Coroutines/MaterializationUtils.cpp
@@ -0,0 +1,308 @@
+//===- MaterializationUtils.cpp - Builds and manipulates coroutine frame
+//-------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+// This file contains classes used to materialize insts after suspends points.
+//===----------------------------------------------------------------------===//
+
+#include "MaterializationUtils.h"
+#include "SpillUtils.h"
+#include "llvm/ADT/PostOrderIterator.h"
+#include "llvm/IR/Dominators.h"
+#include "llvm/IR/InstIterator.h"
+#include "llvm/IR/Instruction.h"
+#include <deque>
+
+using namespace llvm;
+
+using namespace coro;
+
+// The "coro-suspend-crossing" flag is very noisy. There is another debug type,
+// "coro-frame", which results in leaner debug spew.
+#define DEBUG_TYPE "coro-suspend-crossing"
+
+namespace {
+
+// RematGraph is used to construct a DAG for rematerializable instructions
+// When the constructor is invoked with a candidate instruction (which is
+// materializable) it builds a DAG of materializable instructions from that
+// point.
+// Typically, for each instruction identified as re-materializable across a
+// suspend point, a RematGraph will be created.
+struct RematGraph {
+ // Each RematNode in the graph contains the edges to instructions providing
+ // operands in the current node.
+ struct RematNode {
+ Instruction *Node;
+ SmallVector<RematNode *> Operands;
+ RematNode() = default;
+ RematNode(Instruction *V) : Node(V) {}
+ };
+
+ RematNode *EntryNode;
+ using RematNodeMap =
+ SmallMapVector<Instruction *, std::unique_ptr<RematNode>, 8>;
+ RematNodeMap Remats;
+ const std::function<bool(Instruction &)> &MaterializableCallback;
+ SuspendCrossingInfo &Checker;
+
+ RematGraph(const std::function<bool(Instruction &)> &MaterializableCallback,
+ Instruction *I, SuspendCrossingInfo &Checker)
+ : MaterializableCallback(MaterializableCallback), Checker(Checker) {
+ std::unique_ptr<RematNode> FirstNode = std::make_unique<RematNode>(I);
+ EntryNode = FirstNode.get();
+ std::deque<std::unique_ptr<RematNode>> WorkList;
+ addNode(std::move(FirstNode), WorkList, cast<User>(I));
+ while (WorkList.size()) {
+ std::unique_ptr<RematNode> N = std::move(WorkList.front());
+ WorkList.pop_front();
+ addNode(std::move(N), WorkList, cast<User>(I));
+ }
+ }
+
+ void addNode(std::unique_ptr<RematNode> NUPtr,
+ std::deque<std::unique_ptr<RematNode>> &WorkList,
+ User *FirstUse) {
+ RematNode *N = NUPtr.get();
+ if (Remats.count(N->Node))
+ return;
+
+ // We haven't see this node yet - add to the list
+ Remats[N->Node] = std::move(NUPtr);
+ for (auto &Def : N->Node->operands()) {
+ Instruction *D = dyn_cast<Instruction>(Def.get());
+ if (!D || !MaterializableCallback(*D) ||
+ !Checker.isDefinitionAcrossSuspend(*D, FirstUse))
+ continue;
+
+ if (Remats.count(D)) {
+ // Already have this in the graph
+ N->Operands.push_back(Remats[D].get());
+ continue;
+ }
+
+ bool NoMatch = true;
+ for (auto &I : WorkList) {
+ if (I->Node == D) {
+ NoMatch = false;
+ N->Operands.push_back(I.get());
+ break;
+ }
+ }
+ if (NoMatch) {
+ // Create a new node
+ std::unique_ptr<RematNode> ChildNode = std::make_unique<RematNode>(D);
+ N->Operands.push_back(ChildNode.get());
+ WorkList.push_back(std::move(ChildNode));
+ }
+ }
+ }
+
+#if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
+ static std::string getBasicBlockLabel(const BasicBlock *BB) {
+ if (BB->hasName())
+ return BB->getName().str();
+
+ std::string S;
+ raw_string_ostream OS(S);
+ BB->printAsOperand(OS, false);
+ return OS.str().substr(1);
+ }
+
+ void dump() const {
+ dbgs() << "Entry (";
+ dbgs() << getBasicBlockLabel(EntryNode->Node->getParent());
+ dbgs() << ") : " << *EntryNode->Node << "\n"...
[truncated]
|
ChuanqiXu9
approved these changes
Sep 12, 2024
|
LLVM Buildbot has detected a new failure on builder Full details are available at: https://lab.llvm.org/buildbot/#/builders/128/builds/636 Here is the relevant piece of the build log for the reference |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
See RFC for more info: https://discourse.llvm.org/t/rfc-abi-objects-for-coroutines/81057