[llvm] 64f1995 - Fix stack overflow in allPathsGoThroughCold past 6b11573b8c5e (#106384)

[via llvm-commits]

Author: Danial Klimkin
Date: 2024-08-30T14:59:00+02:00
New Revision: 64f19951718075fdd2d2b6d072e8e5ca15a1c6c4

URL: https://github.com/llvm/llvm-project/commit/64f19951718075fdd2d2b6d072e8e5ca15a1c6c4
DIFF: https://github.com/llvm/llvm-project/commit/64f19951718075fdd2d2b6d072e8e5ca15a1c6c4.diff

LOG: Fix stack overflow in allPathsGoThroughCold past 6b11573b8c5e (#106384)

Recursion here causes stack overflow on large inputs. Fixing by
unrolling via a stack.

Added: 
    

Modified: 
    llvm/lib/Transforms/IPO/FunctionAttrs.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/IPO/FunctionAttrs.cpp b/llvm/lib/Transforms/IPO/FunctionAttrs.cpp
index 603a1565e48c45..79746201133bdd 100644
--- a/llvm/lib/Transforms/IPO/FunctionAttrs.cpp
+++ b/llvm/lib/Transforms/IPO/FunctionAttrs.cpp
@@ -1762,54 +1762,52 @@ static void addNoReturnAttrs(const SCCNodeSet &SCCNodes,
   }
 }
 
-static bool
-allBBPathsGoThroughCold(BasicBlock *BB,
-                        SmallDenseMap<BasicBlock *, bool, 16> &Visited) {
-  // If BB contains a cold callsite this path through the CG is cold.
-  // Ignore whether the instructions actually are guranteed to transfer
-  // execution. Divergent behavior is considered unlikely.
-  if (any_of(*BB, [](Instruction &I) {
-        if (auto *CB = dyn_cast<CallBase>(&I))
-          return CB->hasFnAttr(Attribute::Cold);
-        return false;
-      })) {
-    Visited[BB] = true;
-    return true;
-  }
-
-  auto Succs = successors(BB);
-  // We found a path that doesn't go through any cold callsite.
-  if (Succs.empty())
-    return false;
+static bool allPathsGoThroughCold(Function &F) {
+  SmallDenseMap<BasicBlock *, bool, 16> ColdPaths;
+  ColdPaths[&F.front()] = false;
+  SmallVector<BasicBlock *> Jobs;
+  Jobs.push_back(&F.front());
+
+  while (!Jobs.empty()) {
+    BasicBlock *BB = Jobs.pop_back_val();
+
+    // If block contains a cold callsite this path through the CG is cold.
+    // Ignore whether the instructions actually are guaranteed to transfer
+    // execution. Divergent behavior is considered unlikely.
+    if (any_of(*BB, [](Instruction &I) {
+          if (auto *CB = dyn_cast<CallBase>(&I))
+            return CB->hasFnAttr(Attribute::Cold);
+          return false;
+        })) {
+      ColdPaths[BB] = true;
+      continue;
+    }
 
-  // We didn't find a cold callsite in this BB, so check that all successors
-  // contain a cold callsite (or that their successors do).
-  // Potential TODO: We could use static branch hints to assume certain
-  // successor paths are inherently cold, irrespective of if they contain a cold
-  // callsite.
-  for (auto *Succ : Succs) {
-    // Start with false, this is necessary to ensure we don't turn loops into
-    // cold.
-    auto R = Visited.try_emplace(Succ, false);
-    if (!R.second) {
-      if (R.first->second)
-        continue;
+    auto Succs = successors(BB);
+    // We found a path that doesn't go through any cold callsite.
+    if (Succs.empty())
       return false;
+
+    // We didn't find a cold callsite in this BB, so check that all successors
+    // contain a cold callsite (or that their successors do).
+    // Potential TODO: We could use static branch hints to assume certain
+    // successor paths are inherently cold, irrespective of if they contain a
+    // cold callsite.
+    for (BasicBlock *Succ : Succs) {
+      // Start with false, this is necessary to ensure we don't turn loops into
+      // cold.
+      auto [Iter, Inserted] = ColdPaths.try_emplace(Succ, false);
+      if (!Inserted) {
+        if (Iter->second)
+          continue;
+        return false;
+      }
+      Jobs.push_back(Succ);
     }
-    if (!allBBPathsGoThroughCold(Succ, Visited))
-      return false;
-    Visited[Succ] = true;
   }
-
   return true;
 }
 
-static bool allPathsGoThroughCold(Function &F) {
-  SmallDenseMap<BasicBlock *, bool, 16> Visited;
-  Visited[&F.front()] = false;
-  return allBBPathsGoThroughCold(&F.front(), Visited);
-}
-
 // Set the cold function attribute if possible.
 static void addColdAttrs(const SCCNodeSet &SCCNodes,
                          SmallSet<Function *, 8> &Changed) {