[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/include/llvm/IR/Instructions.h

void setSuccessor(unsigned idx, BasicBlock *NewSucc) {
     *(&Op<-1>() - idx) = NewSucc;
   }
 
   /// Swap the successors of this branch instruction.
   ///
   /// Swaps the successors of the branch instruction. This also swaps any
   /// branch weight metadata associated with the instruction so that it
   /// continues to map correctly to each operand.
-  void swapSuccessors();
+  void swapSuccessors(BranchProbabilityInfo *BPI = nullptr);
 
   iterator_range<succ_op_iterator> successors() {
     return make_range(
         succ_op_iterator(std::next(value_op_begin(), isConditional() ? 1 : 0)),
         succ_op_iterator(value_op_end()));
   }
 
   iterator_range<const_succ_op_iterator> successors() const {

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/lib/IR/Instructions.cpp

 //===- Instructions.cpp - Implement the LLVM instructions -----------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements all of the non-inline methods for the LLVM instruction
 // classes.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/IR/Instructions.h"
 #include "LLVMContextImpl.h"
+#include "llvm/Analysis/BranchProbabilityInfo.h"
 #include "llvm/ADT/SmallBitVector.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Constant.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DataLayout.h"
if (BI.getNumOperands() != 1) {
     assert(BI.getNumOperands() == 3 && "BR can have 1 or 3 operands!");
     Op<-3>() = BI.Op<-3>();
     Op<-2>() = BI.Op<-2>();
   }
   Op<-1>() = BI.Op<-1>();
   SubclassOptionalData = BI.SubclassOptionalData;
 }
 
-void BranchInst::swapSuccessors() {
+void BranchInst::swapSuccessors(BranchProbabilityInfo *BPI) {
   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();
+  if (BPI && getParent())
+    BPI->swapSuccEdgesProbabilities(getParent());
 }
 
 //===----------------------------------------------------------------------===//
 //                        AllocaInst Implementation
 //===----------------------------------------------------------------------===//
 
 static Value *getAISize(LLVMContext &Context, Value *Amt) {
   if (!Amt)

llvm/unittests/Analysis/BranchProbabilityInfoTest.cpp

TEST_F(BranchProbabilityInfoTest, SwapProbabilities) {
   BranchInst *Branch = dyn_cast<BranchInst>(LoopHeaderBB->getTerminator());
   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();
+  Branch->swapSuccessors(BPI);
   // 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