[BPI] Add method to swap outgoing edges probabilities

[aleksandr.popov]

Summary

The motivation is need to update branch probability info after
swapping successors of branch instruction

[fhahn]

Could you also link a the patch that actually makes use of this?

[aleksandr.popov]

Could you also link a the patch that actually makes use of this?

Sure, https://reviews.llvm.org/D148244

[mkazantsev]

You could add a unit test here: https://github.com/llvm/llvm-project/blob/main/llvm/unittests/Analysis/BranchProbabilityInfoTest.cpp

[mkazantsev]

llvm/lib/Analysis/BranchProbabilityInfo.cpp

1184

Why do you need this, exactly? I see that swapSuccessors is already updating the metadata:

void BranchInst::swapSuccessors() {
  assert(isConditional() &&
         "Cannot swap successors of an unconditional branch");
  Op<-1>().swap(Op<-2>());

  // Update profile metadata if present and it matches our structural
  // expectations.
  swapProfMetadata();
}

Is analysis returning invalid result after that?

[mkazantsev]

llvm/lib/Analysis/BranchProbabilityInfo.cpp

1184

Yeah, seems it doesn't. Why don't we just add BPI an optional parameter to swapSuccessors and make this update there in this case?

[aleksandr.popov]

Yes, good point, I've move update into swapSuccessors.
But do you suggest to use existing set/getEdgeProbability to swap BPI or to add this new method swapSuccEdgesProbabilities?

[mkazantsev]

I think we should add a parameter to existing method. There are already precedents of that, e.g. SplitBlock at https://llvm.org/doxygen/BasicBlockUtils_8h.html

[nikic]

This is a layering violation: IR cannot depend on Analysis.

[aleksandr.popov]

This is a layering violation: IR cannot depend on Analysis.

Got it, thanks! Returned previous version

[nikic]

Looks reasonable to me.

llvm/include/llvm/Analysis/BranchProbabilityInfo.h

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);
 
+  /// Swap outgoing edges probabilities for \p Src with branch terminator
+  void swapSuccEdgesProbabilities(const BasicBlock *Src);
+
   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, DominatorTree *DT,
                  PostDominatorTree *PDT);

llvm/lib/Analysis/BranchProbabilityInfo.cpp

void BranchProbabilityInfo::copyEdgeProbabilities(BasicBlock *Src,
   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");
   }
 }
 
+void BranchProbabilityInfo::swapSuccEdgesProbabilities(const BasicBlock *Src) {
+  assert(Src->getTerminator()->getNumSuccessors() == 2);
+  if (!Probs.contains(std::make_pair(Src, 0)))
+    return; // No probability is set for edges from Src
+  assert(Probs.contains(std::make_pair(Src, 1)));
+  std::swap(Probs[std::make_pair(Src, 0)], Probs[std::make_pair(Src, 1)]);
+}
+
 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");

llvm/unittests/Analysis/BranchProbabilityInfoTest.cpp

TEST_F(BranchProbabilityInfoTest, SwapProbabilities) {
   ASSERT_TRUE(Branch != nullptr);
   // Save the probabilities before successors swapping
   BranchProbabilityInfo *BPI = &buildBPI(*F);
   auto ProbEdge0 = BPI->getEdgeProbability(LoopHeaderBB, 0U);
   auto ProbEdge1 = BPI->getEdgeProbability(LoopHeaderBB, 1U);
   EXPECT_LT(ProbEdge0, ProbEdge1);
 
   Branch->swapSuccessors();
+  BPI->swapSuccEdgesProbabilities(LoopHeaderBB);
   // TODO: Check the probabilities are swapped as well as the edges
-  EXPECT_EQ(ProbEdge0, BPI->getEdgeProbability(LoopHeaderBB, 0U));
-  EXPECT_EQ(ProbEdge1, BPI->getEdgeProbability(LoopHeaderBB, 1U));
+  EXPECT_EQ(ProbEdge0, BPI->getEdgeProbability(LoopHeaderBB, 1U));
+  EXPECT_EQ(ProbEdge1, BPI->getEdgeProbability(LoopHeaderBB, 0U));
 }
 
 } // end anonymous namespace
 } // end namespace llvm