[BranchProbabilityInfo] Handle irreducible loops.

Summary:
Compute the strongly connected components of the CFG and fall back to
use these for blocks that are in loops that are not detected by
LoopInfo when computing loop back-edge and exit branch probabilities.

Reviewers: dexonsmith, davidxl

Subscribers: mcrosier, llvm-commits

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

llvm-svn: 317094

include/llvm/Analysis/BranchProbabilityInfo.h

@@ -137,6 +137,15 @@ public:
   /// Forget analysis results for the given basic block.
   void eraseBlock(const BasicBlock *BB);
 
+  // Use to track SCCs for handling irreducible loops.
+  using SccMap = DenseMap<const BasicBlock *, int>;
+  using SccHeaderMap = DenseMap<const BasicBlock *, bool>;
+  using SccHeaderMaps = std::vector<SccHeaderMap>;
+  struct SccInfo {
+    SccMap SccNums;
+    SccHeaderMaps SccHeaders;
+  };
+
 private:
   // We need to store CallbackVH's in order to correctly handle basic block
   // removal.
@@ -185,7 +194,8 @@ private:
   bool calcMetadataWeights(const BasicBlock *BB);
   bool calcColdCallHeuristics(const BasicBlock *BB);
   bool calcPointerHeuristics(const BasicBlock *BB);
-  bool calcLoopBranchHeuristics(const BasicBlock *BB, const LoopInfo &LI);
+  bool calcLoopBranchHeuristics(const BasicBlock *BB, const LoopInfo &LI,
+                                SccInfo &SccI);
   bool calcZeroHeuristics(const BasicBlock *BB, const TargetLibraryInfo *TLI);
   bool calcFloatingPointHeuristics(const BasicBlock *BB);
   bool calcInvokeHeuristics(const BasicBlock *BB);

lib/Analysis/BranchProbabilityInfo.cpp

@@ -13,6 +13,7 @@
 
 #include "llvm/Analysis/BranchProbabilityInfo.h"
 #include "llvm/ADT/PostOrderIterator.h"
+#include "llvm/ADT/SCCIterator.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Analysis/LoopInfo.h"
@@ -424,25 +425,73 @@ bool BranchProbabilityInfo::calcPointerHeuristics(const BasicBlock *BB) {
   return true;
 }
 
+static int getSCCNum(const BasicBlock *BB,
+                     const BranchProbabilityInfo::SccInfo &SccI) {
+  auto SccIt = SccI.SccNums.find(BB);
+  if (SccIt == SccI.SccNums.end())
+    return -1;
+  return SccIt->second;
+}
+
+// Consider any block that is an entry point to the SCC as a header.
+static bool isSCCHeader(const BasicBlock *BB, int SccNum,
+                        BranchProbabilityInfo::SccInfo &SccI) {
+  assert(getSCCNum(BB, SccI) == SccNum);
+
+  // Lazily compute the set of headers for a given SCC and cache the results
+  // in the SccHeaderMap.
+  if (SccI.SccHeaders.size() <= static_cast<unsigned>(SccNum))
+    SccI.SccHeaders.resize(SccNum + 1);
+  auto &HeaderMap = SccI.SccHeaders[SccNum];
+  bool Inserted;
+  BranchProbabilityInfo::SccHeaderMap::iterator HeaderMapIt;
+  std::tie(HeaderMapIt, Inserted) = HeaderMap.insert(std::make_pair(BB, false));
+  if (Inserted) {
+    bool IsHeader = llvm::any_of(make_range(pred_begin(BB), pred_end(BB)),
+                                 [&](const BasicBlock *Pred) {
+                                   return getSCCNum(Pred, SccI) != SccNum;
+                                 });
+    HeaderMapIt->second = IsHeader;
+    return IsHeader;
+  } else
+    return HeaderMapIt->second;
+}
+
 // Calculate Edge Weights using "Loop Branch Heuristics". Predict backedges
 // as taken, exiting edges as not-taken.
 bool BranchProbabilityInfo::calcLoopBranchHeuristics(const BasicBlock *BB,
-                                                     const LoopInfo &LI) {
+                                                     const LoopInfo &LI,
+                                                     SccInfo &SccI) {
+  int SccNum;
   Loop *L = LI.getLoopFor(BB);
-  if (!L)
-    return false;
+  if (!L) {
+    SccNum = getSCCNum(BB, SccI);
+    if (SccNum < 0)
+      return false;
+  }
 
   SmallVector<unsigned, 8> BackEdges;
   SmallVector<unsigned, 8> ExitingEdges;
   SmallVector<unsigned, 8> InEdges; // Edges from header to the loop.
 
   for (succ_const_iterator I = succ_begin(BB), E = succ_end(BB); I != E; ++I) {
-    if (!L->contains(*I))
-      ExitingEdges.push_back(I.getSuccessorIndex());
-    else if (L->getHeader() == *I)
-      BackEdges.push_back(I.getSuccessorIndex());
-    else
-      InEdges.push_back(I.getSuccessorIndex());
+    // Use LoopInfo if we have it, otherwise fall-back to SCC info to catch
+    // irreducible loops.
+    if (L) {
+      if (!L->contains(*I))
+        ExitingEdges.push_back(I.getSuccessorIndex());
+      else if (L->getHeader() == *I)
+        BackEdges.push_back(I.getSuccessorIndex());
+      else
+        InEdges.push_back(I.getSuccessorIndex());
+    } else {
+      if (getSCCNum(*I, SccI) != SccNum)
+        ExitingEdges.push_back(I.getSuccessorIndex());
+      else if (isSCCHeader(*I, SccNum, SccI))
+        BackEdges.push_back(I.getSuccessorIndex());
+      else
+        InEdges.push_back(I.getSuccessorIndex());
+    }
   }
 
   if (BackEdges.empty() && ExitingEdges.empty())
@@ -771,6 +820,27 @@ void BranchProbabilityInfo::calculate(const Function &F, const LoopInfo &LI,
   assert(PostDominatedByUnreachable.empty());
   assert(PostDominatedByColdCall.empty());
 
+  // Record SCC numbers of blocks in the CFG to identify irreducible loops.
+  // FIXME: We could only calculate this if the CFG is known to be irreducible
+  // (perhaps cache this info in LoopInfo if we can easily calculate it there?).
+  int SccNum = 0;
+  SccInfo SccI;
+  for (scc_iterator<const Function *> It = scc_begin(&F); !It.isAtEnd();
+       ++It, ++SccNum) {
+    // Ignore single-block SCCs since they either aren't loops or LoopInfo will
+    // catch them.
+    const std::vector<const BasicBlock *> &Scc = *It;
+    if (Scc.size() == 1)
+      continue;
+
+    DEBUG(dbgs() << "BPI: SCC " << SccNum << ":");
+    for (auto *BB : Scc) {
+      DEBUG(dbgs() << " " << BB->getName());
+      SccI.SccNums[BB] = SccNum;
+    }
+    DEBUG(dbgs() << "\n");
+  }
+
   // Walk the basic blocks in post-order so that we can build up state about
   // the successors of a block iteratively.
   for (auto BB : post_order(&F.getEntryBlock())) {
@@ -786,7 +856,7 @@ void BranchProbabilityInfo::calculate(const Function &F, const LoopInfo &LI,
       continue;
     if (calcColdCallHeuristics(BB))
       continue;
-    if (calcLoopBranchHeuristics(BB, LI))
+    if (calcLoopBranchHeuristics(BB, LI, SccI))
       continue;
     if (calcPointerHeuristics(BB))
       continue;

test/Analysis/BranchProbabilityInfo/loop.ll

@@ -6,6 +6,7 @@ declare void @g1()
 declare void @g2()
 declare void @g3()
 declare void @g4()
+declare i32 @g5()
 
 define void @test1(i32 %a, i32 %b) {
 entry:
@@ -364,3 +365,39 @@ for.end15:
   call void @g4()
   ret void
 }
+
+; Test that an irreducible loop gets heavily weighted back-edges.
+define void @test9(i32 %i, i32 %x, i32 %c) {
+entry:
+  %tobool = icmp eq i32 %c, 0
+  br i1 %tobool, label %if.end, label %midloop
+; CHECK:  edge entry -> if.end probability is 0x30000000 / 0x80000000 = 37.50%
+; CHECK:  edge entry -> midloop probability is 0x50000000 / 0x80000000 = 62.50%
+
+if.end:
+  br label %for.cond
+; CHECK:  edge if.end -> for.cond probability is 0x80000000 / 0x80000000 = 100.00% [HOT edge]
+
+for.cond:
+  %i.addr.0 = phi i32 [ %inc, %for.inc ], [ 0, %if.end ]
+  %cmp = icmp slt i32 %i.addr.0, %x
+  br i1 %cmp, label %midloop, label %end
+; CHECK:  edge for.cond -> midloop probability is 0x7c000000 / 0x80000000 = 96.88% [HOT edge]
+; CHECK:  edge for.cond -> end probability is 0x04000000 / 0x80000000 = 3.12%
+
+midloop:
+  %i.addr.1 = phi i32 [ %i, %entry ], [ %i.addr.0, %for.cond ]
+  %call1 = call i32 @g5()
+  %tobool2 = icmp eq i32 %call1, 0
+  br i1 %tobool2, label %for.inc, label %end
+; CHECK:  edge midloop -> for.inc probability is 0x7c000000 / 0x80000000 = 96.88% [HOT edge]
+; CHECK:  edge midloop -> end probability is 0x04000000 / 0x80000000 = 3.12%
+
+for.inc:
+  %inc = add nsw i32 %i.addr.1, 1
+  br label %for.cond
+; CHECK:  edge for.inc -> for.cond probability is 0x80000000 / 0x80000000 = 100.00% [HOT edge]
+
+end:
+  ret void
+}