Skip to content

[Coroutines] Move materialization code into its own utils #108240

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Sep 12, 2024
+351 −297
Merged

[Coroutines] Move materialization code into its own utils #108240

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions llvm/lib/Transforms/Coroutines/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
296 changes: 6 additions & 290 deletions llvm/lib/Transforms/Coroutines/CoroFrame.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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 {
Expand Down Expand Up @@ -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";
Expand Down Expand Up @@ -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) {
Expand All @@ -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,
Expand All @@ -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.
Expand Down Expand Up @@ -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());
}
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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");
}
}
Expand Down
1 change: 1 addition & 0 deletions llvm/lib/Transforms/Coroutines/CoroInternal.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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>);
Expand Down
4 changes: 4 additions & 0 deletions llvm/lib/Transforms/Coroutines/Coroutines.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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");
Expand Down
Loading
Loading