[llvm-branch-commits] [llvm] 68106bd - [Sample Profile Loader] Fix potential invalidated reference (#73181)

[via llvm-branch-commits]

Author: William Junda Huang
Date: 2023-11-29T21:48:55Z
New Revision: 68106bd492e294ecf0a7e2829dd9edf6cd72f3ef

URL: https://github.com/llvm/llvm-project/commit/68106bd492e294ecf0a7e2829dd9edf6cd72f3ef
DIFF: https://github.com/llvm/llvm-project/commit/68106bd492e294ecf0a7e2829dd9edf6cd72f3ef.diff

LOG: [Sample Profile Loader]  Fix potential invalidated reference (#73181)

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.

Added: 
    

Modified: 
    llvm/include/llvm/Transforms/IPO/ProfiledCallGraph.h

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/Transforms/IPO/ProfiledCallGraph.h b/llvm/include/llvm/Transforms/IPO/ProfiledCallGraph.h
index 3f986caeb547287..5381ada37fe27e9 100644
--- a/llvm/include/llvm/Transforms/IPO/ProfiledCallGraph.h
+++ b/llvm/include/llvm/Transforms/IPO/ProfiledCallGraph.h
@@ -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);
     }
   }
 
@@ -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);
@@ -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)
@@ -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;
 };