Skip to content

Commit

Permalink
[Sample Profile Loader] Fix potential invalidated reference (#73181)
Browse files Browse the repository at this point in the history
There is a potential issue in ProfiledCallGraph where pointers to
ProfiledCallGraphNode are used to construct edges, while
ProfiledCallGraphNode instances are being added to a hash map
ProfiledFunctions simultaneously. If rehash happens, those pointers are
invalidated, resulting in undefined behavior/crash. Previously (before
md5phase2) ProfiledFunctions is a llvm::StringMap, which also have the
same issue theoretically when rehashing but was not observed. This patch
fixes this potential issue by using a backing buffer for
ProrfiledCallGraphNode that does not relocate.
  • Loading branch information
huangjd committed Nov 29, 2023
1 parent f150ecc commit 68106bd
Showing 1 changed file with 12 additions and 7 deletions.
19 changes: 12 additions & 7 deletions llvm/include/llvm/Transforms/IPO/ProfiledCallGraph.h
Original file line number Diff line number Diff line change
Expand Up @@ -140,8 +140,11 @@ class ProfiledCallGraph {
if (!ProfiledFunctions.count(Name)) {
// Link to synthetic root to make sure every node is reachable
// from root. This does not affect SCC order.
ProfiledFunctions[Name] = ProfiledCallGraphNode(Name);
Root.Edges.emplace(&Root, &ProfiledFunctions[Name], 0);
// Store the pointer of the node because the map can be rehashed.
auto &Node =
ProfiledCallGraphNodeList.emplace_back(ProfiledCallGraphNode(Name));
ProfiledFunctions[Name] = &Node;
Root.Edges.emplace(&Root, ProfiledFunctions[Name], 0);
}
}

Expand All @@ -152,9 +155,9 @@ class ProfiledCallGraph {
auto CalleeIt = ProfiledFunctions.find(CalleeName);
if (CalleeIt == ProfiledFunctions.end())
return;
ProfiledCallGraphEdge Edge(&ProfiledFunctions[CallerName],
&CalleeIt->second, Weight);
auto &Edges = ProfiledFunctions[CallerName].Edges;
ProfiledCallGraphEdge Edge(ProfiledFunctions[CallerName],
CalleeIt->second, Weight);
auto &Edges = ProfiledFunctions[CallerName]->Edges;
auto EdgeIt = Edges.find(Edge);
if (EdgeIt == Edges.end()) {
Edges.insert(Edge);
Expand Down Expand Up @@ -193,7 +196,7 @@ class ProfiledCallGraph {
return;

for (auto &Node : ProfiledFunctions) {
auto &Edges = Node.second.Edges;
auto &Edges = Node.second->Edges;
auto I = Edges.begin();
while (I != Edges.end()) {
if (I->Weight <= Threshold)
Expand All @@ -205,7 +208,9 @@ class ProfiledCallGraph {
}

ProfiledCallGraphNode Root;
HashKeyMap<std::unordered_map, FunctionId, ProfiledCallGraphNode>
// backing buffer for ProfiledCallGraphNodes.
std::list<ProfiledCallGraphNode> ProfiledCallGraphNodeList;
HashKeyMap<llvm::DenseMap, FunctionId, ProfiledCallGraphNode*>
ProfiledFunctions;
};

Expand Down

0 comments on commit 68106bd

Please sign in to comment.