Generalize the reading of probability metadata to work for both branches

and switches, with arbitrary numbers of successors. Still optimized for
the common case of 2 successors for a conditional branch.

Add a test case for switch metadata showing up in the BlockFrequencyInfo pass.

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@142493 91177308-0d34-0410-b5e6-96231b3b80d8

2 files changed, 67 insertions, 16 deletions

diff --git a/lib/Analysis/BranchProbabilityInfo.cpp b/lib/Analysis/BranchProbabilityInfo.cpp
index 70de3d1c49..52090c9fc1 100644
--- a/lib/Analysis/BranchProbabilityInfo.cpp
+++ b/lib/Analysis/BranchProbabilityInfo.cpp
@@ -140,30 +140,38 @@ public:
 
 // Propagate existing explicit probabilities from either profile data or
 // 'expect' intrinsic processing.
-// FIXME: This doesn't correctly extract probabilities for switches.
 bool BranchProbabilityAnalysis::calcMetadataWeights(BasicBlock *BB) {
-  BranchInst *BI = dyn_cast<BranchInst>(BB->getTerminator());
-  if (!BI || !BI->isConditional())
+  TerminatorInst *TI = BB->getTerminator();
+  if (TI->getNumSuccessors() == 1)
+    return false;
+  if (!isa<BranchInst>(TI) && !isa<SwitchInst>(TI))
     return false;
 
-  MDNode *WeightsNode = BI->getMetadata(LLVMContext::MD_prof);
-  if (!WeightsNode || WeightsNode->getNumOperands() < 3)
+  MDNode *WeightsNode = TI->getMetadata(LLVMContext::MD_prof);
+  if (!WeightsNode)
     return false;
 
-  // Pull the weights out of the metadata. Note that the zero operand is the
-  // name.
-  ConstantInt *Weights[] = {
-    dyn_cast<ConstantInt>(WeightsNode->getOperand(1)),
-    dyn_cast<ConstantInt>(WeightsNode->getOperand(2))
-  };
-  if (!Weights[0] || !Weights[1])
+  // Ensure there are weights for all of the successors. Note that the first
+  // operand to the metadata node is a name, not a weight.
+  if (WeightsNode->getNumOperands() != TI->getNumSuccessors() + 1)
     return false;
 
+  // Build up the final weights that will be used in a temporary buffer, but
+  // don't add them until all weihts are present. Each weight value is clamped
+  // to [1, getMaxWeightFor(BB)].
   uint32_t WeightLimit = getMaxWeightFor(BB);
-  BP->setEdgeWeight(BB, BI->getSuccessor(0),
-                    Weights[0]->getLimitedValue(WeightLimit));
-  BP->setEdgeWeight(BB, BI->getSuccessor(1),
-                    Weights[1]->getLimitedValue(WeightLimit));
+  SmallVector<uint32_t, 2> Weights;
+  Weights.reserve(TI->getNumSuccessors());
+  for (unsigned i = 1, e = WeightsNode->getNumOperands(); i != e; ++i) {
+    ConstantInt *Weight = dyn_cast<ConstantInt>(WeightsNode->getOperand(i));
+    if (!Weight)
+      return false;
+    Weights.push_back(
+      std::max<uint32_t>(1, Weight->getLimitedValue(WeightLimit)));
+  }
+  assert(Weights.size() == TI->getNumSuccessors() && "Checked above");
+  for (unsigned i = 0, e = TI->getNumSuccessors(); i != e; ++i)
+    BP->setEdgeWeight(BB, TI->getSuccessor(i), Weights[i]);
 
   return true;
 }
diff --git a/test/Analysis/BlockFrequencyInfo/basic.ll b/test/Analysis/BlockFrequencyInfo/basic.ll
index f043dc08e4..540d06b1f5 100644
--- a/test/Analysis/BlockFrequencyInfo/basic.ll
+++ b/test/Analysis/BlockFrequencyInfo/basic.ll
@@ -47,3 +47,46 @@ exit:
 }
 
 !0 = metadata !{metadata !"branch_weights", i32 64, i32 4}
+
+define i32 @test3(i32 %i, i32 %a, i32 %b, i32 %c, i32 %d, i32 %e) {
+; CHECK: Printing analysis {{.*}} for function 'test3'
+; CHECK: entry = 1024
+entry:
+  switch i32 %i, label %case_a [ i32 1, label %case_b
+                                 i32 2, label %case_c
+                                 i32 3, label %case_d
+                                 i32 4, label %case_e ], !prof !1
+
+; CHECK: case_a = 51
+case_a:
+  br label %exit
+
+; CHECK: case_b = 51
+case_b:
+  br label %exit
+
+; The 'case_c' branch is predicted more likely via branch weight metadata.
+; CHECK: case_c = 819
+case_c:
+  br label %exit
+
+; CHECK: case_d = 51
+case_d:
+  br label %exit
+
+; CHECK: case_e = 51
+case_e:
+  br label %exit
+
+; FIXME: It may be a bug that we don't sum back to 1024.
+; CHECK: exit = 1023
+exit:
+  %result = phi i32 [ %a, %case_a ],
+                    [ %b, %case_b ],
+                    [ %c, %case_c ],
+                    [ %d, %case_d ],
+                    [ %e, %case_e ]
+  ret i32 %result
+}
+
+!1 = metadata !{metadata !"branch_weights", i32 4, i32 4, i32 64, i32 4, i32 4}