[llvm] [Sample Profile Loader] Fix potential invalidated reference (PR #73181)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-pgo

Author: William Junda Huang (huangjd)

<details>
<summary>Changes</summary>

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. 

---
Full diff: https://github.com/llvm/llvm-project/pull/73181.diff


1 Files Affected:

- (modified) llvm/include/llvm/Transforms/IPO/ProfiledCallGraph.h (+12-8) 


``````````diff
diff --git a/llvm/include/llvm/Transforms/IPO/ProfiledCallGraph.h b/llvm/include/llvm/Transforms/IPO/ProfiledCallGraph.h
index 3f986caeb547287..3d844cc92aa0a8e 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,8 +208,9 @@ class ProfiledCallGraph {
   }
 
   ProfiledCallGraphNode Root;
-  HashKeyMap<std::unordered_map, FunctionId, ProfiledCallGraphNode>
-      ProfiledFunctions;
+  // backing buffer for ProfiledCallGraphNodes.
+  std::list<ProfiledCallGraphNode> ProfiledCallGraphNodeList;
+  llvm::DenseMap<FunctionId, ProfiledCallGraphNode*> ProfiledFunctions;
 };
 
 } // end namespace sampleprof

``````````

</details>


https://github.com/llvm/llvm-project/pull/73181