Recommit "[FunctionAttrs] deduce attr cold on functions if all CG paths call a cold function"

Fixed up the uar test that was failing. It seems with the new `cold`
attribute the order of the functions is different. As far as I can
tell this is not a concern.

Closes #105559

Author: goldsteinn

compiler-rt/test/metadata/uar.cpp

@@ -2,7 +2,7 @@
 // RUN: %clangxx %s -O1 -o %t -fexperimental-sanitize-metadata=covered,uar -fsanitize=address,signed-integer-overflow,alignment && %t | FileCheck %s
 // RUN: %clangxx %s -O1 -o %t -mcmodel=large -fexperimental-sanitize-metadata=covered,uar -fsanitize=address,signed-integer-overflow,alignment && %t | FileCheck %s
 
-// CHECK: metadata add version 2
+// CHECK-DAG: metadata add version 2
 
 __attribute__((noinline, not_tail_called)) void escape(const volatile void *p) {
   [[maybe_unused]] static const volatile void *sink;
@@ -14,51 +14,51 @@ __attribute__((noinline, not_tail_called)) void use(int x) {
   sink += x;
 }
 
-// CHECK: empty: features=0 stack_args=0
+// CHECK-DAG: empty: features=0 stack_args=0
 void empty() {}
 
-// CHECK: simple: features=0 stack_args=0
+// CHECK-DAG: simple: features=0 stack_args=0
 int simple(int *data, int index) { return data[index + 1]; }
 
-// CHECK: builtins: features=0 stack_args=0
+// CHECK-DAG: builtins: features=0 stack_args=0
 int builtins() {
   int x = 0;
   __builtin_prefetch(&x);
   return x;
 }
 
-// CHECK: ellipsis: features=0 stack_args=0
+// CHECK-DAG: ellipsis: features=0 stack_args=0
 void ellipsis(const char *fmt, ...) {
   int x;
   escape(&x);
 }
 
-// CHECK: non_empty_function: features=2 stack_args=0
+// CHECK-DAG: non_empty_function: features=2 stack_args=0
 void non_empty_function() {
   int x;
   escape(&x);
 }
 
-// CHECK: no_stack_args: features=2 stack_args=0
+// CHECK-DAG: no_stack_args: features=2 stack_args=0
 void no_stack_args(long a0, long a1, long a2, long a3, long a4, long a5) {
   int x;
   escape(&x);
 }
 
-// CHECK: stack_args: features=6 stack_args=16
+// CHECK-DAG: stack_args: features=6 stack_args=16
 void stack_args(long a0, long a1, long a2, long a3, long a4, long a5, long a6) {
   int x;
   escape(&x);
 }
 
-// CHECK: more_stack_args: features=6 stack_args=32
+// CHECK-DAG: more_stack_args: features=6 stack_args=32
 void more_stack_args(long a0, long a1, long a2, long a3, long a4, long a5,
                      long a6, long a7, long a8) {
   int x;
   escape(&x);
 }
 
-// CHECK: struct_stack_args: features=6 stack_args=144
+// CHECK-DAG: struct_stack_args: features=6 stack_args=144
 struct large {
   char x[131];
 };
@@ -69,28 +69,28 @@ void struct_stack_args(large a) {
 
 __attribute__((noinline)) int tail_called(int x) { return x; }
 
-// CHECK: with_tail_call: features=2
+// CHECK-DAG: with_tail_call: features=2
 int with_tail_call(int x) { [[clang::musttail]] return tail_called(x); }
 
 __attribute__((noinline, noreturn)) int noreturn(int x) { __builtin_trap(); }
 
-// CHECK: with_noreturn_tail_call: features=0
+// CHECK-DAG: with_noreturn_tail_call: features=0
 int with_noreturn_tail_call(int x) { return noreturn(x); }
 
-// CHECK: local_array: features=0
+// CHECK-DAG: local_array: features=0
 void local_array(int x) {
   int data[10];
   use(data[x]);
 }
 
-// CHECK: local_alloca: features=0
+// CHECK-DAG: local_alloca: features=0
 void local_alloca(int size, int i, int j) {
   volatile int *p = static_cast<int *>(__builtin_alloca(size));
   p[i] = 0;
   use(p[j]);
 }
 
-// CHECK: escaping_alloca: features=2
+// CHECK-DAG: escaping_alloca: features=2
 void escaping_alloca(int size, int i) {
   volatile int *p = static_cast<int *>(__builtin_alloca(size));
   escape(&p[i]);

llvm/lib/Transforms/IPO/FunctionAttrs.cpp

@@ -82,6 +82,7 @@ STATISTIC(NumNoUnwind, "Number of functions marked as nounwind");
 STATISTIC(NumNoFree, "Number of functions marked as nofree");
 STATISTIC(NumWillReturn, "Number of functions marked as willreturn");
 STATISTIC(NumNoSync, "Number of functions marked as nosync");
+STATISTIC(NumCold, "Number of functions marked as cold");
 
 STATISTIC(NumThinLinkNoRecurse,
           "Number of functions marked as norecurse during thinlink");
@@ -1745,6 +1746,7 @@ static bool canReturn(Function &F) {
   return false;
 }
 
+
 // Set the noreturn function attribute if possible.
 static void addNoReturnAttrs(const SCCNodeSet &SCCNodes,
                              SmallSet<Function *, 8> &Changed) {
@@ -1760,6 +1762,72 @@ 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;
+
+  // 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;
+      return false;
+    }
+    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) {
+  for (Function *F : SCCNodes) {
+    if (!F || !F->hasExactDefinition() || F->hasFnAttribute(Attribute::Naked) ||
+        F->hasFnAttribute(Attribute::Cold) || F->hasFnAttribute(Attribute::Hot))
+      continue;
+
+    // Potential TODO: We could add attribute `cold` on functions with `coldcc`.
+    if (allPathsGoThroughCold(*F)) {
+      F->addFnAttr(Attribute::Cold);
+      ++NumCold;
+      Changed.insert(F);
+      continue;
+    }
+  }
+}
+
 static bool functionWillReturn(const Function &F) {
   // We can infer and propagate function attributes only when we know that the
   // definition we'll get at link time is *exactly* the definition we see now.
@@ -1853,6 +1921,7 @@ deriveAttrsInPostOrder(ArrayRef<Function *> Functions, AARGetterT &&AARGetter,
   addArgumentAttrs(Nodes.SCCNodes, Changed);
   inferConvergent(Nodes.SCCNodes, Changed);
   addNoReturnAttrs(Nodes.SCCNodes, Changed);
+  addColdAttrs(Nodes.SCCNodes, Changed);
   addWillReturn(Nodes.SCCNodes, Changed);
   addNoUndefAttrs(Nodes.SCCNodes, Changed);