[PGO] Add a threshold for number of critical edges in PGO

For some auto-generated sources, we have a huge number of critical
edges (like from switch statements). We have seen instance of 183777
critical edges in one function.

After we split the critical edges in PGO instrumentation/profile-use
pass, the CFG is so large that we have compiler time issues in
downstream passes (like in machine CSE and block placement). Here I
add a threshold to skip PGO if the number of critical edges are too
large.

The threshold is large enough so that it will not affect the majority
of PGO compilation.

Also sync the logic for skipping instrumentation and profile-use. I
think this is the correct thing to do.

Differential Revision: https://reviews.llvm.org/D137184

Author: xur-llvm

llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp

@@ -302,12 +302,17 @@ static cl::opt<std::string> PGOTraceFuncHash(
 
 static cl::opt<unsigned> PGOFunctionSizeThreshold(
     "pgo-function-size-threshold", cl::Hidden,
-    cl::desc("Do not instrument functions smaller than this threshold"));
+    cl::desc("Do not instrument functions smaller than this threshold."));
 
 static cl::opt<bool> MatchMemProf(
     "pgo-match-memprof", cl::init(true), cl::Hidden,
     cl::desc("Perform matching and annotation of memprof profiles."));
 
+static cl::opt<unsigned> PGOFunctionCriticalEdgeThreshold(
+    "pgo-critical-edge-threshold", cl::init(20000), cl::Hidden,
+    cl::desc("Do not instrument functions with the number of critical edges "
+             " greater than this threshold."));
+
 namespace llvm {
 // Command line option to turn on CFG dot dump after profile annotation.
 // Defined in Analysis/BlockFrequencyInfo.cpp:  -pgo-view-counts
@@ -1846,6 +1851,38 @@ static void collectComdatMembers(
       ComdatMembers.insert(std::make_pair(C, &GA));
 }
 
+// Don't perform PGO instrumeatnion / profile-use.
+static bool skipPGO(const Function &F) {
+  if (F.isDeclaration())
+    return true;
+  if (F.hasFnAttribute(llvm::Attribute::NoProfile))
+    return true;
+  if (F.hasFnAttribute(llvm::Attribute::SkipProfile))
+    return true;
+  if (F.getInstructionCount() < PGOFunctionSizeThreshold)
+    return true;
+
+  // If there are too many critical edges, PGO might cause
+  // compiler time problem. Skip PGO if the number of
+  // critical edges execeed the threshold.
+  unsigned NumCriticalEdges = 0;
+  for (auto &BB : F) {
+    const Instruction *TI = BB.getTerminator();
+    for (unsigned I = 0, E = TI->getNumSuccessors(); I != E; ++I) {
+      if (isCriticalEdge(TI, I))
+        NumCriticalEdges++;
+    }
+  }
+  if (NumCriticalEdges > PGOFunctionCriticalEdgeThreshold) {
+    LLVM_DEBUG(dbgs() << "In func " << F.getName()
+                      << ", NumCriticalEdges=" << NumCriticalEdges
+                      << " exceed the threshold. Skip PGO.\n");
+    return true;
+  }
+
+  return false;
+}
+
 static bool InstrumentAllFunctions(
     Module &M, function_ref<TargetLibraryInfo &(Function &)> LookupTLI,
     function_ref<BranchProbabilityInfo *(Function &)> LookupBPI,
@@ -1858,13 +1895,7 @@ static bool InstrumentAllFunctions(
   collectComdatMembers(M, ComdatMembers);
 
   for (auto &F : M) {
-    if (F.isDeclaration())
-      continue;
-    if (F.hasFnAttribute(llvm::Attribute::NoProfile))
-      continue;
-    if (F.hasFnAttribute(llvm::Attribute::SkipProfile))
-      continue;
-    if (F.getInstructionCount() < PGOFunctionSizeThreshold)
+    if (skipPGO(F))
       continue;
     auto &TLI = LookupTLI(F);
     auto *BPI = LookupBPI(F);
@@ -2092,7 +2123,7 @@ static bool annotateAllFunctions(
   if (PGOInstrumentEntry.getNumOccurrences() > 0)
     InstrumentFuncEntry = PGOInstrumentEntry;
   for (auto &F : M) {
-    if (F.isDeclaration())
+    if (skipPGO(F))
       continue;
     auto &TLI = LookupTLI(F);
     auto *BPI = LookupBPI(F);

llvm/test/Transforms/PGOProfile/critical-edge-threshold.ll

@@ -0,0 +1,27 @@
+; Test the critical edge threahold
+; RUN: opt < %s -passes=pgo-instr-gen -pgo-critical-edge-threshold=1 -pgo-instrument-entry=true -S | FileCheck %s
+
+@sum = dso_local global i32 0, align 4
+
+define void @foo(i32 %a, i32 %b) {
+entry:
+  %tobool.not = icmp eq i32 %a, 0
+  br i1 %tobool.not, label %if.end4, label %if.then
+
+if.then:
+  %0 = load i32, ptr @sum, align 4
+  %inc = add nsw i32 %0, 1
+  store i32 %inc, ptr @sum, align 4
+  %tobool1.not = icmp eq i32 %b, 0
+  br i1 %tobool1.not, label %if.end4, label %if.then2
+
+if.then2:
+  %inc3 = add nsw i32 %0, 2
+  store i32 %inc3, ptr @sum, align 4
+  br label %if.end4
+
+if.end4:
+  ret void
+}
+
+; CHECK-NOT: call void @llvm.instrprof.increment(ptr @__profn_foo