Skip to content

Commit 5899a93

Browse files
tzuralontzuralon
committed
[Coroutines] fixed the bugs introduced in the last 3 tests, improved split performance
1 parent 7f00eeb commit 5899a93

File tree

1 file changed

+148
-128
lines changed

1 file changed

+148
-128
lines changed

‎llvm/lib/Transforms/Coroutines/CoroFrame.cpp‎

Lines changed: 148 additions & 128 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717

1818
#include "CoroInternal.h"
1919
#include "llvm/ADT/ScopeExit.h"
20+
#include "llvm/ADT/SmallSet.h"
2021
#include "llvm/ADT/SmallString.h"
2122
#include "llvm/Analysis/StackLifetime.h"
2223
#include "llvm/IR/DIBuilder.h"
@@ -975,157 +976,176 @@ static StructType *buildFrameType(Function &F, coro::Shape &Shape,
975976
return FrameTy;
976977
}
977978

978-
// Fixer for the "Instruction does not dominate all uses!" bug
979-
// The fix consists of mapping problematic paths (where CoroBegin does not
980-
// dominate cleanup BBs) and clones them to 2 flows - the one that insertSpills
981-
// intended to create (using the spill) and another one, preserving the logics
982-
// of pre-splitting, which would be triggered if unwinding happened before
983-
// CoroBegin
984-
static void
985-
splitBasicBlocksNotDominatedByCoroBegin(const FrameDataInfo &FrameData,
986-
coro::Shape &Shape, Function *F,
987-
DominatorTree &DT) {
988-
ValueToValueMapTy VMap;
989-
DenseMap<BasicBlock *, BasicBlock *>
990-
ProcessedSpillBlockToAlternativeUnspilledBlockMap;
991-
bool FunctionHasSomeBlockNotDominatedByCoroBegin;
992-
SmallVector<BasicBlock *> SpillUserBlocks;
993-
994-
for (const auto &E : FrameData.Spills) {
995-
for (auto *U : E.second) {
996-
if (std::find(SpillUserBlocks.begin(), SpillUserBlocks.end(),
997-
U->getParent()) == SpillUserBlocks.end()) {
998-
SpillUserBlocks.push_back(U->getParent());
999-
}
1000-
}
1001-
}
1002-
SpillUserBlocks.push_back(Shape.AllocaSpillBlock);
1003-
1004-
do {
1005-
FunctionHasSomeBlockNotDominatedByCoroBegin = false;
1006-
// We want to traverse the function post-order (predecessors first),
1007-
// and check dominance starting CoroBegin
1008-
bool HaveTraversedCoroBegin = false;
1009-
for (BasicBlock *CurrentBlock : ReversePostOrderTraversal<Function *>(F)) {
1010-
if (!HaveTraversedCoroBegin &&
1011-
CurrentBlock != Shape.CoroBegin->getParent()) {
1012-
continue;
1013-
}
1014-
HaveTraversedCoroBegin = true;
1015-
1016-
// Focus on 2 types of users that produce errors - those in
1017-
// FrameData.Spills, and decendants of Shape.AllocaSpillBlocks
1018-
if (!DT.dominates(Shape.CoroBegin, CurrentBlock) &&
1019-
std::find(SpillUserBlocks.begin(), SpillUserBlocks.end(),
1020-
CurrentBlock) != SpillUserBlocks.end()) {
1021-
// Mark another iteration of the loop is needed, to verify that no more
1022-
// dominance issues after current run
1023-
FunctionHasSomeBlockNotDominatedByCoroBegin = true;
1024-
1025-
// Clone (preserve) the current basic block, before it will be modified
1026-
// by insertSpills
1027-
auto UnspilledAlternativeBlock =
1028-
CloneBasicBlock(CurrentBlock, VMap, ".unspilled_alternative", F);
979+
// This function assumes that it is being called on basic block in reversed
980+
// post-order, meaning predecessors are visited before successors failing to do
981+
// so will cause VMap to be non-valid and will cause instructions to fail
982+
// mapping to their corresponding clones
983+
static void finalizeBasicBlockCloneAndTrackSuccessors(
984+
BasicBlock *InitialBlock, BasicBlock *ClonedBlock, ValueToValueMapTy &VMap,
985+
SmallSet<BasicBlock *, 20> &SuccessorBlocksSet) {
986+
// This code will examine the basic block, fix issues caused by clones
987+
// for example - tailor cleanupret to the corresponding cleanuppad
988+
// it will use VMap to do so
989+
// in addition, it will add the node successors to SuccessorBlocksSet
1029990

1030991
// Remap the instructions, VMap here aggregates instructions across
1031-
// multiple BasicBlocks, and we assume that traversal is post-order,
992+
// multiple BasicBlocks, and we assume that traversal is reversed post-order,
1032993
// therefore successor blocks (for example instructions having funclet
1033994
// tags) will be mapped correctly to the new cloned cleanuppad
1034-
for (Instruction &I : *UnspilledAlternativeBlock) {
1035-
RemapInstruction(&I, VMap,
995+
for (Instruction &ClonedBlockInstruction : *ClonedBlock) {
996+
RemapInstruction(&ClonedBlockInstruction, VMap,
1036997
RF_NoModuleLevelChanges | RF_IgnoreMissingLocals);
1037998
}
1038999

1039-
// Keep track between the processed spill basic block and the cloned
1040-
// alternative unspilled basic block Will help us fix instructions that
1041-
// their context is complex (for example cleanuppad of funclet is
1042-
// defined in another BB)
1043-
ProcessedSpillBlockToAlternativeUnspilledBlockMap[CurrentBlock] =
1044-
UnspilledAlternativeBlock;
1045-
1046-
SmallVector<Instruction *> FixUpPredTerminators;
1047-
1048-
// Find the specific predecessors that does not dominated by
1049-
// Shape.CoroBegin We don't fix them here but later because it's not
1050-
// safe while using predecessors as iterator function
1051-
for (BasicBlock *Pred : predecessors(CurrentBlock)) {
1052-
if (!DT.dominates(Shape.CoroBegin, Pred)) {
1053-
FixUpPredTerminators.push_back(Pred->getTerminator());
1054-
}
1055-
}
1000+
const auto &InitialBlockTerminator = InitialBlock->getTerminator();
10561001

1057-
// Fixups for current block terminator
1058-
const auto &CurrentBlockTerminator = CurrentBlock->getTerminator();
1059-
1060-
// If it's cleanupret, find the correspondant cleanuppad (use the map to
1061-
// find it)
1062-
if (auto CurrentBlockCleanupReturnTerminator =
1063-
dyn_cast<CleanupReturnInst>(CurrentBlockTerminator)) {
1064-
BasicBlock *CBCPBB =
1065-
CurrentBlockCleanupReturnTerminator->getCleanupPad()->getParent();
1066-
CleanupReturnInst *DBT = dyn_cast<CleanupReturnInst>(
1067-
UnspilledAlternativeBlock->getTerminator());
1068-
1069-
// Again assuming post-order traversal - if we mapped the predecessing
1070-
// cleanuppad block before, we should find it here If not, do nothing
1071-
if (ProcessedSpillBlockToAlternativeUnspilledBlockMap.contains(
1072-
CBCPBB)) {
1073-
Instruction *DCPr =
1074-
&ProcessedSpillBlockToAlternativeUnspilledBlockMap[CBCPBB]
1075-
->front();
1076-
CleanupPadInst *DCP = cast<CleanupPadInst>(DCPr);
1077-
DBT->setCleanupPad(DCP);
1078-
}
1002+
// If it's cleanupret, find the correspondant cleanuppad (use the VMap to
1003+
// find it).
1004+
if (auto *InitialBlockTerminatorCleanupReturn =
1005+
dyn_cast<CleanupReturnInst>(InitialBlockTerminator)) {
1006+
// if none found do nothing
1007+
if (VMap.find(InitialBlockTerminatorCleanupReturn->getCleanupPad()) ==
1008+
VMap.end()) {
1009+
return;
1010+
}
10791011

1012+
auto *ClonedBlockTerminatorCleanupReturn =
1013+
cast<CleanupReturnInst>(ClonedBlock->getTerminator());
1014+
1015+
// Assuming reversed post-order traversal
1016+
llvm::Value *ClonedBlockCleanupPadValue =
1017+
VMap[InitialBlockTerminatorCleanupReturn->getCleanupPad()];
1018+
auto *ClonedBlockCleanupPad =
1019+
cast<CleanupPadInst>(ClonedBlockCleanupPadValue);
1020+
ClonedBlockTerminatorCleanupReturn->setCleanupPad(ClonedBlockCleanupPad);
1021+
10801022
// If it's a branch/invoke, keep track of its successors, we want to
1081-
// calculate dominance between CoroBegin and them also They might need
1082-
// clone as well
1083-
} else if (auto CurrentBlockBranchTerminator =
1084-
dyn_cast<BranchInst>(CurrentBlockTerminator)) {
1023+
// calculate dominance between CoroBegin and them also
1024+
} else if (auto *InitialBlockTerminatorBranch =
1025+
dyn_cast<BranchInst>(InitialBlockTerminator)) {
10851026
for (unsigned int successorIdx = 0;
1086-
successorIdx < CurrentBlockBranchTerminator->getNumSuccessors();
1027+
successorIdx < InitialBlockTerminatorBranch->getNumSuccessors();
10871028
++successorIdx) {
1088-
SpillUserBlocks.push_back(
1089-
CurrentBlockBranchTerminator->getSuccessor(successorIdx));
1029+
SuccessorBlocksSet.insert(
1030+
InitialBlockTerminatorBranch->getSuccessor(successorIdx));
10901031
}
1091-
} else if (auto CurrentBlockInvokeTerminator =
1092-
dyn_cast<InvokeInst>(CurrentBlockTerminator)) {
1093-
SpillUserBlocks.push_back(
1094-
CurrentBlockInvokeTerminator->getUnwindDest());
1032+
} else if (auto *InitialBlockTerminatorInvoke =
1033+
dyn_cast<InvokeInst>(InitialBlockTerminator)) {
1034+
SuccessorBlocksSet.insert(InitialBlockTerminatorInvoke->getUnwindDest());
1035+
} else if (isa<ReturnInst>(InitialBlockTerminator)) {
1036+
// No action needed
10951037
} else {
1096-
report_fatal_error("Not implemented terminator for this specific "
1097-
"instruction fixup in current block fixups");
1098-
}
1038+
InitialBlockTerminator->print(dbgs());
1039+
report_fatal_error("Terminator is not implemented in "
1040+
"finalizeBasicBlockCloneAndTrackSuccessors");
1041+
}
1042+
}
10991043

1100-
// Fixups on the predecessors terminator - direct them to out untouched
1101-
// alternative block to break dominance error.
1102-
for (auto FixUpPredTerminator : FixUpPredTerminators) {
1103-
if (auto FixUpPredTerminatorInvoke =
1104-
dyn_cast<InvokeInst>(FixUpPredTerminator)) {
1105-
FixUpPredTerminatorInvoke->setUnwindDest(UnspilledAlternativeBlock);
1106-
} else if (auto FixUpPredTerminatorBranch =
1107-
dyn_cast<BranchInst>(FixUpPredTerminator)) {
1044+
// Dominance issue fixer for each predecessor satisfying predicate function
1045+
void splitIfBasicBlockPredecessors(
1046+
BasicBlock *InitialBlock, BasicBlock *ReplacementBlock,
1047+
std::function<bool(BasicBlock *)> Predicate) {
1048+
1049+
SmallVector<BasicBlock *> InitialBlockPredecessors;
1050+
1051+
auto Predecessors = predecessors(InitialBlock);
1052+
std::copy_if(Predecessors.begin(), Predecessors.end(),
1053+
std::back_inserter(InitialBlockPredecessors), Predicate);
1054+
1055+
// Fixups on the predecessors terminator - point them to ReplacementBlock.
1056+
for (auto *InitialBlockPredecessor : InitialBlockPredecessors) {
1057+
auto *InitialBlockPredecessorTerminator =
1058+
InitialBlockPredecessor->getTerminator();
1059+
if (auto *InitialBlockPredecessorTerminatorInvoke =
1060+
dyn_cast<InvokeInst>(InitialBlockPredecessorTerminator)) {
1061+
if (InitialBlock ==
1062+
InitialBlockPredecessorTerminatorInvoke->getUnwindDest()) {
1063+
InitialBlockPredecessorTerminatorInvoke->setUnwindDest(
1064+
ReplacementBlock);
1065+
} else {
1066+
InitialBlockPredecessorTerminatorInvoke->setNormalDest(
1067+
ReplacementBlock);
1068+
}
1069+
} else if (auto *InitialBlockPredecessorTerminatorBranch =
1070+
dyn_cast<BranchInst>(InitialBlockPredecessorTerminator)) {
11081071
for (unsigned int successorIdx = 0;
1109-
successorIdx < FixUpPredTerminatorBranch->getNumSuccessors();
1072+
successorIdx <
1073+
InitialBlockPredecessorTerminatorBranch->getNumSuccessors();
11101074
++successorIdx) {
1111-
if (CurrentBlock ==
1112-
FixUpPredTerminatorBranch->getSuccessor(successorIdx)) {
1113-
FixUpPredTerminatorBranch->setSuccessor(
1114-
successorIdx, UnspilledAlternativeBlock);
1115-
}
1116-
}
1075+
if (InitialBlock ==
1076+
InitialBlockPredecessorTerminatorBranch->getSuccessor(
1077+
successorIdx)) {
1078+
InitialBlockPredecessorTerminatorBranch->setSuccessor(
1079+
successorIdx, ReplacementBlock);
1080+
}
1081+
}
1082+
} else if (auto *InitialBlockPredecessorTerminatorCleanupReturn =
1083+
dyn_cast<CleanupReturnInst>(
1084+
InitialBlockPredecessorTerminator)) {
1085+
InitialBlockPredecessorTerminatorCleanupReturn->setUnwindDest(
1086+
ReplacementBlock);
11171087
} else {
1118-
report_fatal_error("Not implemented terminator for this specific "
1119-
"instruction in pred fixups");
1088+
InitialBlockPredecessorTerminator->print(dbgs());
1089+
report_fatal_error(
1090+
"Terminator is not implemented in splitIfBasicBlockPredecessors");
1091+
}
1092+
}
1093+
}
1094+
1095+
// Fixer for the "Instruction does not dominate all uses!" bug
1096+
// The fix consists of mapping problematic paths (where CoroBegin does not
1097+
// dominate cleanup nodes) and duplicates them to 2 flows - the one that
1098+
// insertSpills intended to create (using the spill) and another one, preserving
1099+
// the logics of pre-splitting, which would be triggered if unwinding happened
1100+
// before CoroBegin
1101+
static void
1102+
splitBasicBlocksNotDominatedByCoroBegin(const FrameDataInfo &FrameData,
1103+
const coro::Shape &Shape, Function *F,
1104+
DominatorTree &DT) {
1105+
ValueToValueMapTy VMap;
1106+
SmallSet<BasicBlock *, 20> SpillUserBlocksSet;
1107+
1108+
// Prepare the node set, logics will be run only on those nodes
1109+
for (const auto &E : FrameData.Spills) {
1110+
for (auto *U : E.second) {
1111+
auto CurrentBlock = U->getParent();
1112+
if (!DT.dominates(Shape.CoroBegin, CurrentBlock)) {
1113+
SpillUserBlocksSet.insert(CurrentBlock);
1114+
}
11201115
}
11211116
}
11221117

1123-
// We changed dominance tree, so recalculate.
1124-
DT.recalculate(*F);
1118+
// Run is in reversed post order, to enforce visiting predecessors before
1119+
// successors
1120+
for (BasicBlock *CurrentBlock : ReversePostOrderTraversal<Function *>(F)) {
1121+
if (!SpillUserBlocksSet.contains(CurrentBlock)) {
11251122
continue;
11261123
}
1127-
}
1128-
} while (FunctionHasSomeBlockNotDominatedByCoroBegin);
1124+
SpillUserBlocksSet.erase(CurrentBlock);
1125+
1126+
// Preserve the current node. the duplicate will become the unspilled
1127+
// alternative
1128+
auto UnspilledAlternativeBlock =
1129+
CloneBasicBlock(CurrentBlock, VMap, ".unspilled_alternative", F);
1130+
1131+
// Remap node instructions, keep track of successors to visit them in next
1132+
// iterations
1133+
finalizeBasicBlockCloneAndTrackSuccessors(
1134+
CurrentBlock, UnspilledAlternativeBlock, VMap, SpillUserBlocksSet);
1135+
1136+
// Split only if predecessor breaks dominance against CoroBegin
1137+
splitIfBasicBlockPredecessors(CurrentBlock, UnspilledAlternativeBlock,
1138+
[&DT, &Shape](BasicBlock *PredecessorNode) {
1139+
return !DT.dominates(Shape.CoroBegin,
1140+
PredecessorNode);
1141+
});
1142+
1143+
// We changed dominance tree, so recalculate
1144+
DT.recalculate(*F);
1145+
}
1146+
1147+
assert(SpillUserBlocksSet.empty() &&
1148+
"Graph is corrupted by SpillUserBlocksSet");
11291149
}
11301150

11311151
// Replace all alloca and SSA values that are accessed across suspend points

0 commit comments

Comments
 (0)