Index: include/llvm/Analysis/BranchProbabilityInfo.h
===================================================================
--- include/llvm/Analysis/BranchProbabilityInfo.h
+++ include/llvm/Analysis/BranchProbabilityInfo.h
@@ -17,6 +17,7 @@
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/IndexedMap.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/IR/ValueHandle.h"
@@ -141,10 +142,6 @@
   };
   DenseSet<BasicBlockCallbackVH, DenseMapInfo<Value*>> Handles;
 
-  // Since we allow duplicate edges from one basic block to another, we use
-  // a pair (PredBlock and an index in the successors) to specify an edge.
-  typedef std::pair<const BasicBlock *, unsigned> Edge;
-
   // Default weight value. Used when we don't have information about the edge.
   // TODO: DEFAULT_WEIGHT makes sense during static predication, when none of
   // the successors have a weight yet. But it doesn't make sense when providing
@@ -153,7 +150,11 @@
   // weight to just "inherit" the non-zero weight of an adjacent successor.
   static const uint32_t DEFAULT_WEIGHT = 16;
 
-  DenseMap<Edge, BranchProbability> Probs;
+  // Since we allow duplicate edges from one basic block to another, we need to
+  // use successor index to determine edge destination. This results in a
+  // following data structure:
+  //   BB -> (SuccessorIdx -> BranchProbability)
+  DenseMap<const BasicBlock *, IndexedMap<BranchProbability>> Probs;
 
   /// \brief Track the last function we run over for printing.
   const Function *LastF;
Index: lib/Analysis/BranchProbabilityInfo.cpp
===================================================================
--- lib/Analysis/BranchProbabilityInfo.cpp
+++ lib/Analysis/BranchProbabilityInfo.cpp
@@ -577,6 +577,12 @@
   return nullptr;
 }
 
+static bool
+EdgeProbabilityRecorded(const IndexedMap<BranchProbability> &SuccMap,
+                        unsigned Index) {
+  return SuccMap.inBounds(Index) && !SuccMap[Index].isUnknown();
+}
+
 /// Get the raw edge probability for the edge. If can't find it, return a
 /// default probability 1/N where N is the number of successors. Here an edge is
 /// specified using PredBlock and an
@@ -584,10 +590,11 @@
 BranchProbability
 BranchProbabilityInfo::getEdgeProbability(const BasicBlock *Src,
                                           unsigned IndexInSuccessors) const {
-  auto I = Probs.find(std::make_pair(Src, IndexInSuccessors));
+  const auto I = Probs.find(Src);
 
-  if (I != Probs.end())
-    return I->second;
+  if (I != Probs.end() &&
+      EdgeProbabilityRecorded(I->second, IndexInSuccessors))
+    return I->second[IndexInSuccessors];
 
   return {1,
           static_cast<uint32_t>(std::distance(succ_begin(Src), succ_end(Src)))};
@@ -605,16 +612,24 @@
 BranchProbabilityInfo::getEdgeProbability(const BasicBlock *Src,
                                           const BasicBlock *Dst) const {
   auto Prob = BranchProbability::getZero();
+  uint32_t succ_num = std::distance(succ_begin(Src), succ_end(Src));
+
+  const auto BBProbsIt = Probs.find(Src);
+  if (BBProbsIt == Probs.end())
+    return {1, succ_num};
+  const IndexedMap<BranchProbability> &BBSuccProbs = BBProbsIt->second;
+
   bool FoundProb = false;
-  for (succ_const_iterator I = succ_begin(Src), E = succ_end(Src); I != E; ++I)
-    if (*I == Dst) {
-      auto MapI = Probs.find(std::make_pair(Src, I.getSuccessorIndex()));
-      if (MapI != Probs.end()) {
-        FoundProb = true;
-        Prob += MapI->second;
-      }
+  for (succ_const_iterator I = succ_begin(Src), E = succ_end(Src); I != E;
+       ++I) {
+    unsigned SuccIdx = I.getSuccessorIndex();
+
+    if (*I == Dst && EdgeProbabilityRecorded(BBSuccProbs, SuccIdx)) {
+      FoundProb = true;
+      Prob += BBSuccProbs[SuccIdx];
     }
-  uint32_t succ_num = std::distance(succ_begin(Src), succ_end(Src));
+  }
+
   return FoundProb ? Prob : BranchProbability(1, succ_num);
 }
 
@@ -623,8 +638,16 @@
 void BranchProbabilityInfo::setEdgeProbability(const BasicBlock *Src,
                                                unsigned IndexInSuccessors,
                                                BranchProbability Prob) {
-  Probs[std::make_pair(Src, IndexInSuccessors)] = Prob;
+  IndexedMap<BranchProbability> &BBSuccProbs = Probs[Src];
+
+  // This is nop if IndexInSuccessors is already present in the map.
+  // If map was resized new elements will be filled with the 'Unknown' branch
+  // probability.
+  BBSuccProbs.grow(IndexInSuccessors);
+  BBSuccProbs[IndexInSuccessors] = Prob;
+
   Handles.insert(BasicBlockCallbackVH(Src, this));
+
   DEBUG(dbgs() << "set edge " << Src->getName() << " -> " << IndexInSuccessors
                << " successor probability to " << Prob << "\n");
 }
@@ -643,11 +666,7 @@
 }
 
 void BranchProbabilityInfo::eraseBlock(const BasicBlock *BB) {
-  for (auto I = Probs.begin(), E = Probs.end(); I != E; ++I) {
-    auto Key = I->first;
-    if (Key.first == BB)
-      Probs.erase(Key);
-  }
+  Probs.erase(BB);
 }
 
 void BranchProbabilityInfo::calculate(const Function &F, const LoopInfo &LI) {