[llvm] Fix stack overflow in allPathsGoThroughCold past 6b11573b8c5e (PR #106384)

[Danial Klimkin via llvm-commits]

https://github.com/dklimkin updated https://github.com/llvm/llvm-project/pull/106384

>From 2200dddd855d4fb46b7630d67384bd425cea6a7d Mon Sep 17 00:00:00 2001
From: Danial Klimkin <dklimkin at google.com>
Date: Wed, 28 Aug 2024 15:05:10 +0200
Subject: [PATCH] Fix stack overflow in allPathsGoThroughCold

---
 llvm/lib/Transforms/IPO/FunctionAttrs.cpp | 82 +++++++++++------------
 1 file changed, 41 insertions(+), 41 deletions(-)

diff --git a/llvm/lib/Transforms/IPO/FunctionAttrs.cpp b/llvm/lib/Transforms/IPO/FunctionAttrs.cpp
index 603a1565e48c45..1c1ca3d3795f31 100644
--- a/llvm/lib/Transforms/IPO/FunctionAttrs.cpp
+++ b/llvm/lib/Transforms/IPO/FunctionAttrs.cpp
@@ -62,6 +62,7 @@
 #include <iterator>
 #include <map>
 #include <optional>
+#include <stack>
 #include <vector>
 
 using namespace llvm;
@@ -1762,54 +1763,53 @@ 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;
+  std::stack<BasicBlock *> Jobs;
+  Jobs.push(&F.front());
+
+  while (!Jobs.empty()) {
+    BasicBlock *BB = Jobs.top();
+    Jobs.pop();
+
+    // 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 (llvm::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(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) {