Skip to content

Commit c616f19

Browse files
[MemProf] Refactor context node creation into a new helper (NFC) (#108408)
Simplify code by refactoring some common handling for node creation into a helper function.
1 parent 68ddd6c commit c616f19

File tree

1 file changed

+17
-20
lines changed

1 file changed

+17
-20
lines changed

llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp

Lines changed: 17 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -619,6 +619,16 @@ class CallsiteContextGraph {
619619
return static_cast<const DerivedCCG *>(this)->getLabel(Func, Call, CloneNo);
620620
}
621621

622+
// Create and return a new ContextNode.
623+
ContextNode *createNewNode(bool IsAllocation, const FuncTy *F = nullptr,
624+
CallInfo C = CallInfo()) {
625+
NodeOwner.push_back(std::make_unique<ContextNode>(IsAllocation, C));
626+
auto *NewNode = NodeOwner.back().get();
627+
if (F)
628+
NodeToCallingFunc[NewNode] = F;
629+
return NewNode;
630+
}
631+
622632
/// Helpers to find the node corresponding to the given call or stackid.
623633
ContextNode *getNodeForInst(const CallInfo &C);
624634
ContextNode *getNodeForAlloc(const CallInfo &C);
@@ -1082,11 +1092,8 @@ typename CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::ContextNode *
10821092
CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::addAllocNode(
10831093
CallInfo Call, const FuncTy *F) {
10841094
assert(!getNodeForAlloc(Call));
1085-
NodeOwner.push_back(
1086-
std::make_unique<ContextNode>(/*IsAllocation=*/true, Call));
1087-
ContextNode *AllocNode = NodeOwner.back().get();
1095+
ContextNode *AllocNode = createNewNode(/*IsAllocation=*/true, F, Call);
10881096
AllocationCallToContextNodeMap[Call] = AllocNode;
1089-
NodeToCallingFunc[AllocNode] = F;
10901097
// Use LastContextId as a uniq id for MIB allocation nodes.
10911098
AllocNode->OrigStackOrAllocId = LastContextId;
10921099
// Alloc type should be updated as we add in the MIBs. We should assert
@@ -1143,9 +1150,7 @@ void CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::addStackNodesForMIB(
11431150
auto StackId = getStackId(*ContextIter);
11441151
ContextNode *StackNode = getNodeForStackId(StackId);
11451152
if (!StackNode) {
1146-
NodeOwner.push_back(
1147-
std::make_unique<ContextNode>(/*IsAllocation=*/false));
1148-
StackNode = NodeOwner.back().get();
1153+
StackNode = createNewNode(/*IsAllocation=*/false);
11491154
StackEntryIdToContextNodeMap[StackId] = StackNode;
11501155
StackNode->OrigStackOrAllocId = StackId;
11511156
}
@@ -1448,10 +1453,7 @@ void CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::
14481453
continue;
14491454

14501455
// Create new context node.
1451-
NodeOwner.push_back(
1452-
std::make_unique<ContextNode>(/*IsAllocation=*/false, Call));
1453-
ContextNode *NewNode = NodeOwner.back().get();
1454-
NodeToCallingFunc[NewNode] = Func;
1456+
ContextNode *NewNode = createNewNode(/*IsAllocation=*/false, Func, Call);
14551457
NonAllocationCallToContextNodeMap[Call] = NewNode;
14561458
CreatedNode = true;
14571459
NewNode->AllocTypes = computeAllocType(SavedContextIds);
@@ -2164,10 +2166,7 @@ bool CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::calleesMatch(
21642166
} else {
21652167
FuncToCallsWithMetadata[Func].push_back({NewCall});
21662168
// Create Node and record node info.
2167-
NodeOwner.push_back(
2168-
std::make_unique<ContextNode>(/*IsAllocation=*/false, NewCall));
2169-
NewNode = NodeOwner.back().get();
2170-
NodeToCallingFunc[NewNode] = Func;
2169+
NewNode = createNewNode(/*IsAllocation=*/false, Func, NewCall);
21712170
TailCallToContextNodeMap[NewCall] = NewNode;
21722171
NewNode->AllocTypes = Edge->AllocTypes;
21732172
}
@@ -2740,13 +2739,11 @@ CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::moveEdgeToNewCalleeClone(
27402739
const std::shared_ptr<ContextEdge> &Edge, EdgeIter *CallerEdgeI,
27412740
DenseSet<uint32_t> ContextIdsToMove) {
27422741
ContextNode *Node = Edge->Callee;
2743-
NodeOwner.push_back(
2744-
std::make_unique<ContextNode>(Node->IsAllocation, Node->Call));
2745-
ContextNode *Clone = NodeOwner.back().get();
2742+
assert(NodeToCallingFunc.count(Node));
2743+
ContextNode *Clone =
2744+
createNewNode(Node->IsAllocation, NodeToCallingFunc[Node], Node->Call);
27462745
Node->addClone(Clone);
27472746
Clone->MatchingCalls = Node->MatchingCalls;
2748-
assert(NodeToCallingFunc.count(Node));
2749-
NodeToCallingFunc[Clone] = NodeToCallingFunc[Node];
27502747
moveEdgeToExistingCalleeClone(Edge, Clone, CallerEdgeI, /*NewClone=*/true,
27512748
ContextIdsToMove);
27522749
return Clone;

0 commit comments

Comments
 (0)