[BranchProbabilityInfo] Introduce method copyEdgeProbabilities(). NFC

[yrouban]

Summary

A new method is introduced to allow bulk copy of outgoing edge probabilities from one block to another. This can be useful when a block is a cloned from another one and we do not know if there are edge probabilities set for the original block or not. Copying outside of the BranchProbabilityInfo class makes the user unconditionally set the cloned block's edge probabilities even if they are unset for the original block.
This new method will be used in JumpThreading.

[MaskRay]

Thanks!

llvm/include/llvm/Analysis/BranchProbabilityInfo.h

public:
   /// Set the raw probabilities for all edges from the given block.
   ///
   /// This allows a pass to explicitly set edge probabilities for a block. It
   /// can be used when updating the CFG to update the branch probability
   /// information.
   void setEdgeProbability(const BasicBlock *Src,
                           const SmallVectorImpl<BranchProbability> &Probs);
 
+  /// Copy outgoing edge probabilities from \p Src to \p Dst.
+  ///
+  /// This allows to keep probabilities unset for the destination if they were
+  /// unset for source.
+  void copyEdgeProbabilities(BasicBlock *Src, BasicBlock *Dst);
+
   static BranchProbability getBranchProbStackProtector(bool IsLikely) {
     static const BranchProbability LikelyProb((1u << 20) - 1, 1u << 20);
     return IsLikely ? LikelyProb : LikelyProb.getCompl();
   }
 
   void calculate(const Function &F, const LoopInfo &LI,
                  const TargetLibraryInfo *TLI, PostDominatorTree *PDT);
 

llvm/lib/Analysis/BranchProbabilityInfo.cpp

void BranchProbabilityInfo::setEdgeProbability(
   // 1.0 exactly. That is TotalNumerator == BranchProbability::getDenominator.
   // Instead, every single probability in Probs must be as accurate as possible.
   // This results in error 1/denominator at most, thus the total absolute error
   // should be within Probs.size / BranchProbability::getDenominator.
   assert(TotalNumerator <= BranchProbability::getDenominator() + Probs.size());
   assert(TotalNumerator >= BranchProbability::getDenominator() - Probs.size());
 }
 
+void BranchProbabilityInfo::copyEdgeProbabilities(BasicBlock *Src, BasicBlock *Dst) {
+  eraseBlock(Dst); // Erase stale data if any.
+  unsigned NumSuccessors = Src->getTerminator()->getNumSuccessors();
+  assert(NumSuccessors == Dst->getTerminator()->getNumSuccessors());
+  if (NumSuccessors == 0)
+    return; // Nothing to set.
+  if (this->Probs.find(std::make_pair(Src, 0)) == this->Probs.end())
+    return; // No probability is set for edges from Src. Keep the same for Dst.
+
+  Handles.insert(BasicBlockCallbackVH(Dst, this));
+  for (unsigned SuccIdx = 0; SuccIdx < NumSuccessors; ++SuccIdx) {
+    auto Prob = this->Probs[std::make_pair(Src, SuccIdx)];
+    this->Probs[std::make_pair(Dst, SuccIdx)] = Prob;
+    LLVM_DEBUG(dbgs() << "set edge " << Dst->getName() << " -> " << SuccIdx
+                      << " successor probability to " << Prob << "\n");
+  }
+}
+
 raw_ostream &
 BranchProbabilityInfo::printEdgeProbability(raw_ostream &OS,
                                             const BasicBlock *Src,
                                             const BasicBlock *Dst) const {
   const BranchProbability Prob = getEdgeProbability(Src, Dst);
   OS << "edge " << Src->getName() << " -> " << Dst->getName()
      << " probability is " << Prob
      << (isEdgeHot(Src, Dst) ? " [HOT edge]\n" : "\n");