[Coroutines] Move spill related methods to a Spill utils #107884
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 code related to spilling into SpillUtils to clean up CoroFrame
|
@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 58.58 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/107884.diff 4 Files Affected:
diff --git a/llvm/lib/Transforms/Coroutines/CMakeLists.txt b/llvm/lib/Transforms/Coroutines/CMakeLists.txt
index ddc4caa82b0e98..c6508174a7f109 100644
--- a/llvm/lib/Transforms/Coroutines/CMakeLists.txt
+++ b/llvm/lib/Transforms/Coroutines/CMakeLists.txt
@@ -8,6 +8,7 @@ add_llvm_component_library(LLVMCoroutines
CoroFrame.cpp
CoroSplit.cpp
SuspendCrossingInfo.cpp
+ SpillUtils.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 7a631a646f3f8b..c6027a7de2f806 100644
--- a/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
@@ -16,27 +16,22 @@
//===----------------------------------------------------------------------===//
#include "CoroInternal.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/CFG.h"
-#include "llvm/Analysis/PtrUseVisitor.h"
#include "llvm/Analysis/StackLifetime.h"
-#include "llvm/Config/llvm-config.h"
-#include "llvm/IR/CFG.h"
#include "llvm/IR/DIBuilder.h"
#include "llvm/IR/DebugInfo.h"
#include "llvm/IR/Dominators.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/InstIterator.h"
#include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/Module.h"
#include "llvm/Support/Debug.h"
-#include "llvm/Support/MathExtras.h"
#include "llvm/Support/OptimizedStructLayout.h"
-#include "llvm/Support/circular_raw_ostream.h"
-#include "llvm/Support/raw_ostream.h"
#include "llvm/Transforms/Utils/BasicBlockUtils.h"
#include "llvm/Transforms/Utils/Local.h"
#include "llvm/Transforms/Utils/PromoteMemToReg.h"
@@ -175,24 +170,17 @@ template <> struct GraphTraits<RematGraph *> {
namespace {
class FrameTypeBuilder;
// Mapping from the to-be-spilled value to all the users that need reload.
-using SpillInfo = SmallMapVector<Value *, SmallVector<Instruction *, 2>, 8>;
-struct AllocaInfo {
- AllocaInst *Alloca;
- DenseMap<Instruction *, std::optional<APInt>> Aliases;
- bool MayWriteBeforeCoroBegin;
- AllocaInfo(AllocaInst *Alloca,
- DenseMap<Instruction *, std::optional<APInt>> Aliases,
- bool MayWriteBeforeCoroBegin)
- : Alloca(Alloca), Aliases(std::move(Aliases)),
- MayWriteBeforeCoroBegin(MayWriteBeforeCoroBegin) {}
-};
struct FrameDataInfo {
// All the values (that are not allocas) that needs to be spilled to the
// frame.
- SpillInfo Spills;
+ coro::SpillInfo &Spills;
// Allocas contains all values defined as allocas that need to live in the
// frame.
- SmallVector<AllocaInfo, 8> Allocas;
+ SmallVectorImpl<coro::AllocaInfo> &Allocas;
+
+ FrameDataInfo(coro::SpillInfo &Spills,
+ SmallVectorImpl<coro::AllocaInfo> &Allocas)
+ : Spills(Spills), Allocas(Allocas) {}
SmallVector<Value *, 8> getAllDefs() const {
SmallVector<Value *, 8> Defs;
@@ -271,7 +259,7 @@ struct FrameDataInfo {
} // namespace
#ifndef NDEBUG
-static void dumpSpills(StringRef Title, const SpillInfo &Spills) {
+static void dumpSpills(StringRef Title, const coro::SpillInfo &Spills) {
dbgs() << "------------- " << Title << " --------------\n";
for (const auto &E : Spills) {
E.first->dump();
@@ -290,7 +278,7 @@ static void dumpRemats(
}
}
-static void dumpAllocas(const SmallVectorImpl<AllocaInfo> &Allocas) {
+static void dumpAllocas(const SmallVectorImpl<coro::AllocaInfo> &Allocas) {
dbgs() << "------------- Allocas --------------\n";
for (const auto &A : Allocas) {
A.Alloca->dump();
@@ -538,7 +526,7 @@ void FrameTypeBuilder::addFieldForAllocas(const Function &F,
return StackLifetimeAnalyzer.getLiveRange(AI1).overlaps(
StackLifetimeAnalyzer.getLiveRange(AI2));
};
- auto GetAllocaSize = [&](const AllocaInfo &A) {
+ auto GetAllocaSize = [&](const coro::AllocaInfo &A) {
std::optional<TypeSize> RetSize = A.Alloca->getAllocationSize(DL);
assert(RetSize && "Variable Length Arrays (VLA) are not supported.\n");
assert(!RetSize->isScalable() && "Scalable vectors are not yet supported");
@@ -1119,367 +1107,6 @@ static StructType *buildFrameType(Function &F, coro::Shape &Shape,
return FrameTy;
}
-// We use a pointer use visitor to track how an alloca is being used.
-// The goal is to be able to answer the following three questions:
-// 1. Should this alloca be allocated on the frame instead.
-// 2. Could the content of the alloca be modified prior to CoroBegn, which would
-// require copying the data from alloca to the frame after CoroBegin.
-// 3. Is there any alias created for this alloca prior to CoroBegin, but used
-// after CoroBegin. In that case, we will need to recreate the alias after
-// CoroBegin based off the frame. To answer question 1, we track two things:
-// a. List of all BasicBlocks that use this alloca or any of the aliases of
-// the alloca. In the end, we check if there exists any two basic blocks that
-// cross suspension points. If so, this alloca must be put on the frame. b.
-// Whether the alloca or any alias of the alloca is escaped at some point,
-// either by storing the address somewhere, or the address is used in a
-// function call that might capture. If it's ever escaped, this alloca must be
-// put on the frame conservatively.
-// To answer quetion 2, we track through the variable MayWriteBeforeCoroBegin.
-// Whenever a potential write happens, either through a store instruction, a
-// function call or any of the memory intrinsics, we check whether this
-// instruction is prior to CoroBegin. To answer question 3, we track the offsets
-// of all aliases created for the alloca prior to CoroBegin but used after
-// CoroBegin. std::optional is used to be able to represent the case when the
-// offset is unknown (e.g. when you have a PHINode that takes in different
-// offset values). We cannot handle unknown offsets and will assert. This is the
-// potential issue left out. An ideal solution would likely require a
-// significant redesign.
-namespace {
-struct AllocaUseVisitor : PtrUseVisitor<AllocaUseVisitor> {
- using Base = PtrUseVisitor<AllocaUseVisitor>;
- AllocaUseVisitor(const DataLayout &DL, const DominatorTree &DT,
- const coro::Shape &CoroShape,
- const SuspendCrossingInfo &Checker,
- bool ShouldUseLifetimeStartInfo)
- : PtrUseVisitor(DL), DT(DT), CoroShape(CoroShape), Checker(Checker),
- ShouldUseLifetimeStartInfo(ShouldUseLifetimeStartInfo) {
- for (AnyCoroSuspendInst *SuspendInst : CoroShape.CoroSuspends)
- CoroSuspendBBs.insert(SuspendInst->getParent());
- }
-
- void visit(Instruction &I) {
- Users.insert(&I);
- Base::visit(I);
- // If the pointer is escaped prior to CoroBegin, we have to assume it would
- // be written into before CoroBegin as well.
- if (PI.isEscaped() &&
- !DT.dominates(CoroShape.CoroBegin, PI.getEscapingInst())) {
- MayWriteBeforeCoroBegin = true;
- }
- }
- // We need to provide this overload as PtrUseVisitor uses a pointer based
- // visiting function.
- void visit(Instruction *I) { return visit(*I); }
-
- void visitPHINode(PHINode &I) {
- enqueueUsers(I);
- handleAlias(I);
- }
-
- void visitSelectInst(SelectInst &I) {
- enqueueUsers(I);
- handleAlias(I);
- }
-
- void visitStoreInst(StoreInst &SI) {
- // Regardless whether the alias of the alloca is the value operand or the
- // pointer operand, we need to assume the alloca is been written.
- handleMayWrite(SI);
-
- if (SI.getValueOperand() != U->get())
- return;
-
- // We are storing the pointer into a memory location, potentially escaping.
- // As an optimization, we try to detect simple cases where it doesn't
- // actually escape, for example:
- // %ptr = alloca ..
- // %addr = alloca ..
- // store %ptr, %addr
- // %x = load %addr
- // ..
- // If %addr is only used by loading from it, we could simply treat %x as
- // another alias of %ptr, and not considering %ptr being escaped.
- auto IsSimpleStoreThenLoad = [&]() {
- auto *AI = dyn_cast<AllocaInst>(SI.getPointerOperand());
- // If the memory location we are storing to is not an alloca, it
- // could be an alias of some other memory locations, which is difficult
- // to analyze.
- if (!AI)
- return false;
- // StoreAliases contains aliases of the memory location stored into.
- SmallVector<Instruction *, 4> StoreAliases = {AI};
- while (!StoreAliases.empty()) {
- Instruction *I = StoreAliases.pop_back_val();
- for (User *U : I->users()) {
- // If we are loading from the memory location, we are creating an
- // alias of the original pointer.
- if (auto *LI = dyn_cast<LoadInst>(U)) {
- enqueueUsers(*LI);
- handleAlias(*LI);
- continue;
- }
- // If we are overriding the memory location, the pointer certainly
- // won't escape.
- if (auto *S = dyn_cast<StoreInst>(U))
- if (S->getPointerOperand() == I)
- continue;
- if (auto *II = dyn_cast<IntrinsicInst>(U))
- if (II->isLifetimeStartOrEnd())
- continue;
- // BitCastInst creats aliases of the memory location being stored
- // into.
- if (auto *BI = dyn_cast<BitCastInst>(U)) {
- StoreAliases.push_back(BI);
- continue;
- }
- return false;
- }
- }
-
- return true;
- };
-
- if (!IsSimpleStoreThenLoad())
- PI.setEscaped(&SI);
- }
-
- // All mem intrinsics modify the data.
- void visitMemIntrinsic(MemIntrinsic &MI) { handleMayWrite(MI); }
-
- void visitBitCastInst(BitCastInst &BC) {
- Base::visitBitCastInst(BC);
- handleAlias(BC);
- }
-
- void visitAddrSpaceCastInst(AddrSpaceCastInst &ASC) {
- Base::visitAddrSpaceCastInst(ASC);
- handleAlias(ASC);
- }
-
- void visitGetElementPtrInst(GetElementPtrInst &GEPI) {
- // The base visitor will adjust Offset accordingly.
- Base::visitGetElementPtrInst(GEPI);
- handleAlias(GEPI);
- }
-
- void visitIntrinsicInst(IntrinsicInst &II) {
- // When we found the lifetime markers refers to a
- // subrange of the original alloca, ignore the lifetime
- // markers to avoid misleading the analysis.
- if (!IsOffsetKnown || !Offset.isZero())
- return Base::visitIntrinsicInst(II);
- switch (II.getIntrinsicID()) {
- default:
- return Base::visitIntrinsicInst(II);
- case Intrinsic::lifetime_start:
- LifetimeStarts.insert(&II);
- LifetimeStartBBs.push_back(II.getParent());
- break;
- case Intrinsic::lifetime_end:
- LifetimeEndBBs.insert(II.getParent());
- break;
- }
- }
-
- void visitCallBase(CallBase &CB) {
- for (unsigned Op = 0, OpCount = CB.arg_size(); Op < OpCount; ++Op)
- if (U->get() == CB.getArgOperand(Op) && !CB.doesNotCapture(Op))
- PI.setEscaped(&CB);
- handleMayWrite(CB);
- }
-
- bool getShouldLiveOnFrame() const {
- if (!ShouldLiveOnFrame)
- ShouldLiveOnFrame = computeShouldLiveOnFrame();
- return *ShouldLiveOnFrame;
- }
-
- bool getMayWriteBeforeCoroBegin() const { return MayWriteBeforeCoroBegin; }
-
- DenseMap<Instruction *, std::optional<APInt>> getAliasesCopy() const {
- assert(getShouldLiveOnFrame() && "This method should only be called if the "
- "alloca needs to live on the frame.");
- for (const auto &P : AliasOffetMap)
- if (!P.second)
- report_fatal_error("Unable to handle an alias with unknown offset "
- "created before CoroBegin.");
- return AliasOffetMap;
- }
-
-private:
- const DominatorTree &DT;
- const coro::Shape &CoroShape;
- const SuspendCrossingInfo &Checker;
- // All alias to the original AllocaInst, created before CoroBegin and used
- // after CoroBegin. Each entry contains the instruction and the offset in the
- // original Alloca. They need to be recreated after CoroBegin off the frame.
- DenseMap<Instruction *, std::optional<APInt>> AliasOffetMap{};
- SmallPtrSet<Instruction *, 4> Users{};
- SmallPtrSet<IntrinsicInst *, 2> LifetimeStarts{};
- SmallVector<BasicBlock *> LifetimeStartBBs{};
- SmallPtrSet<BasicBlock *, 2> LifetimeEndBBs{};
- SmallPtrSet<const BasicBlock *, 2> CoroSuspendBBs{};
- bool MayWriteBeforeCoroBegin{false};
- bool ShouldUseLifetimeStartInfo{true};
-
- mutable std::optional<bool> ShouldLiveOnFrame{};
-
- bool computeShouldLiveOnFrame() const {
- // If lifetime information is available, we check it first since it's
- // more precise. We look at every pair of lifetime.start intrinsic and
- // every basic block that uses the pointer to see if they cross suspension
- // points. The uses cover both direct uses as well as indirect uses.
- if (ShouldUseLifetimeStartInfo && !LifetimeStarts.empty()) {
- // If there is no explicit lifetime.end, then assume the address can
- // cross suspension points.
- if (LifetimeEndBBs.empty())
- return true;
-
- // If there is a path from a lifetime.start to a suspend without a
- // corresponding lifetime.end, then the alloca's lifetime persists
- // beyond that suspension point and the alloca must go on the frame.
- llvm::SmallVector<BasicBlock *> Worklist(LifetimeStartBBs);
- if (isManyPotentiallyReachableFromMany(Worklist, CoroSuspendBBs,
- &LifetimeEndBBs, &DT))
- return true;
-
- // Addresses are guaranteed to be identical after every lifetime.start so
- // we cannot use the local stack if the address escaped and there is a
- // suspend point between lifetime markers. This should also cover the
- // case of a single lifetime.start intrinsic in a loop with suspend point.
- if (PI.isEscaped()) {
- for (auto *A : LifetimeStarts) {
- for (auto *B : LifetimeStarts) {
- if (Checker.hasPathOrLoopCrossingSuspendPoint(A->getParent(),
- B->getParent()))
- return true;
- }
- }
- }
- return false;
- }
- // FIXME: Ideally the isEscaped check should come at the beginning.
- // However there are a few loose ends that need to be fixed first before
- // we can do that. We need to make sure we are not over-conservative, so
- // that the data accessed in-between await_suspend and symmetric transfer
- // is always put on the stack, and also data accessed after coro.end is
- // always put on the stack (esp the return object). To fix that, we need
- // to:
- // 1) Potentially treat sret as nocapture in calls
- // 2) Special handle the return object and put it on the stack
- // 3) Utilize lifetime.end intrinsic
- if (PI.isEscaped())
- return true;
-
- for (auto *U1 : Users)
- for (auto *U2 : Users)
- if (Checker.isDefinitionAcrossSuspend(*U1, U2))
- return true;
-
- return false;
- }
-
- void handleMayWrite(const Instruction &I) {
- if (!DT.dominates(CoroShape.CoroBegin, &I))
- MayWriteBeforeCoroBegin = true;
- }
-
- bool usedAfterCoroBegin(Instruction &I) {
- for (auto &U : I.uses())
- if (DT.dominates(CoroShape.CoroBegin, U))
- return true;
- return false;
- }
-
- void handleAlias(Instruction &I) {
- // We track all aliases created prior to CoroBegin but used after.
- // These aliases may need to be recreated after CoroBegin if the alloca
- // need to live on the frame.
- if (DT.dominates(CoroShape.CoroBegin, &I) || !usedAfterCoroBegin(I))
- return;
-
- if (!IsOffsetKnown) {
- AliasOffetMap[&I].reset();
- } else {
- auto Itr = AliasOffetMap.find(&I);
- if (Itr == AliasOffetMap.end()) {
- AliasOffetMap[&I] = Offset;
- } else if (Itr->second && *Itr->second != Offset) {
- // If we have seen two different possible values for this alias, we set
- // it to empty.
- AliasOffetMap[&I].reset();
- }
- }
- }
-};
-} // namespace
-
-// We need to make room to insert a spill after initial PHIs, but before
-// catchswitch instruction. Placing it before violates the requirement that
-// catchswitch, like all other EHPads must be the first nonPHI in a block.
-//
-// Split away catchswitch into a separate block and insert in its place:
-//
-// cleanuppad <InsertPt> cleanupret.
-//
-// cleanupret instruction will act as an insert point for the spill.
-static Instruction *splitBeforeCatchSwitch(CatchSwitchInst *CatchSwitch) {
- BasicBlock *CurrentBlock = CatchSwitch->getParent();
- BasicBlock *NewBlock = CurrentBlock->splitBasicBlock(CatchSwitch);
- CurrentBlock->getTerminator()->eraseFromParent();
-
- auto *CleanupPad =
- CleanupPadInst::Create(CatchSwitch->getParentPad(), {}, "", CurrentBlock);
- auto *CleanupRet =
- CleanupReturnInst::Create(CleanupPad, NewBlock, CurrentBlock);
- return CleanupRet;
-}
-
-static BasicBlock::iterator getSpillInsertionPt(const coro::Shape &Shape,
- Value *Def,
- const DominatorTree &DT) {
- BasicBlock::iterator InsertPt;
- if (auto *Arg = dyn_cast<Argument>(Def)) {
- // For arguments, we will place the store instruction right after
- // the coroutine frame pointer instruction, i.e. coro.begin.
- InsertPt = Shape.getInsertPtAfterFramePtr();
-
- // If we're spilling an Argument, make sure we clear 'nocapture'
- // from the coroutine function.
- Arg->getParent()->removeParamAttr(Arg->getArgNo(), Attribute::NoCapture);
- } else if (auto *CSI = dyn_cast<AnyCoroSuspendInst>(Def)) {
- // Don't spill immediately after a suspend; splitting assumes
- // that the suspend will be followed by a branch.
- InsertPt = CSI->getParent()->getSingleSuccessor()->getFirstNonPHIIt();
- } else {
- auto *I = cast<Instruction>(Def);
- if (!DT.dominates(Shape.CoroBegin, I)) {
- // If it is not dominated by CoroBegin, then spill should be
- // inserted immediately after CoroFrame is computed.
- InsertPt = Shape.getInsertPtAfterFramePtr();
- } else if (auto *II = dyn_cast<InvokeInst>(I)) {
- // If we are spilling the result of the invoke instruction, split
- // the normal edge and insert the spill in the new block.
- auto *NewBB = SplitEdge(II->getParent(), II->getNormalDest());
- InsertPt = NewBB->getTerminator()->getIterator();
- } else if (isa<PHINode>(I)) {
- // Skip the PHINodes and EH pads instructions.
- BasicBlock *DefBlock = I->getParent();
- if (auto *CSI = dyn_cast<CatchSwitchInst>(DefBlock->getTerminator()))
- InsertPt = splitBeforeCatchSwitch(CSI)->getIterator();
- else
- InsertPt = DefBlock->getFirstInsertionPt();
- } else {
- assert(!I->isTerminator() && "unexpected terminator");
- // For all other values, the spill is placed immediately after
- // the definition.
- InsertPt = I->getNextNode()->getIterator();
- }
- }
-
- return InsertPt;
-}
-
// Replace all alloca and SSA values that are accessed across suspend points
// with GetElementPointer from coroutine frame + loads and stores. Create an
// AllocaSpillBB that will become the new entry block for the resume parts of
@@ -1564,7 +1191,7 @@ static void insertSpills(const FrameDataInfo &FrameData, coro::Shape &Shape) {
auto SpillAlignment = Align(FrameData.getAlign(Def));
// Create a store instruction storing the value into the
// coroutine frame.
- BasicBlock::iterator InsertPt = getSpillInsertionPt(Shape, Def, DT);
+ BasicBlock::iterator InsertPt = coro::getSpillInsertionPt(Shape, Def, DT);
Type *ByValTy = nullptr;
if (auto *Arg = dyn_cast<Argument>(Def)) {
@@ -2015,13 +1642,6 @@ bool coro::defaultMaterializable(Instruction &V) {
isa<BinaryOperator>(&V) || isa<CmpInst>(&V) || isa<SelectInst>(&V));
}
-// Check for structural coroutine intrinsics that should not be spilled into
-// the coroutine frame.
-static bool isCoroutineStructureIntrinsic(Instruction &I) {
- return isa<CoroIdInst>(&I) ||...
[truncated]
|
ChuanqiXu9
approved these changes
Sep 10, 2024
There was a problem hiding this comment.
I didn't check if they are strictly the same line by line. But I believe it is trivial code movement. So LGTM with removing the prefix [llvm/llvm-project].
But it looks like there were some regression reports for the first issue. Please take a look before landing.
Looks like the test failure is due to a permissions issue. "C:\a\lld-x86_64-win\llvm-project\llvm\unittests\Support\ProgramTest.cpp(163): error: fs::remove(Twine(LongPath)): did not return errc::success. |
TylerNowicki
changed the title
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