Normalize MBB's successors' probabilities in several locations.

[congh]

Summary

This patch adds some missing calls to MBB::normalizeSuccProbs() in several locations where it should be called. Those places are found by checking if the sum of successors' probabilities is approximate one in MachineBlockPlacement pass with some instrumented code (not in this patch).

[davidxl]

include/llvm/CodeGen/MachineBasicBlock.h

467–470

Can you add some comment about the guidelines when to use the normalize arg? For instance when the removal is the final one.

lib/CodeGen/IfConversion.cpp

1262–1263

Is it better to add the same 'normalize' parameter to the addSuccessor method?

A related -- for addSuccessor, what is the guideline to normalize? How can we tell if the normalization actually masks a bug in the updater?

lib/CodeGen/TailDuplication.cpp

754 ↗

(On Diff #41988)

This one seems unnecessary.

862 ↗

(On Diff #41988)

The normalization seems redundant -- See assertion before: assert(PredB->succ_empty()) ..

If the assertion is not true, the way prob is updated is not correct. It should do this:
OldProb = getEdgeProb(PredBB, PredBB->succ_begin());

for (.... ) {

PredBB->addSuccessor (*I, getEdgeProb(TailBB,I)*OldProb)

}

Normalization is not needed either.

[congh]

include/llvm/CodeGen/MachineBasicBlock.h

467–470

I have added some comments on normalizeSuccProbs(). This is usually done when the current update on the MBB is done. There are three effects of using this interface: 1. Incorrect. 2. Necessary. 3. Redundant. We should forbid the first case and try to always do it if it is necessary.

lib/CodeGen/IfConversion.cpp

1262–1263

As the user already provided a probability when calling addSuccessor(), I think adding the 'normalize' parameter to it is unnecessary. Actually in most cases we don't have to do the normalization.

The time to do normalization is the update is done and no more update is coming to this MBB. It is difficult to detect if a normalization is correct, or may be expensive to do this. It is the user's responsibility to keep its usage correct.

lib/CodeGen/TailDuplication.cpp

754 ↗

(On Diff #41988)

Agree. Thanks for pointing it out!.

862 ↗

(On Diff #41988)

You are right. The normalization is redundant here.

[congh]

Update the patch according to David's comments.

[davidxl]

Are there any test case changes?

lib/CodeGen/IfConversion.cpp

1263

Add a brief comment about why normalization is needed (i.e., why the input BPs do not sum up to 1)?

1719

The comment needs to be fixed (probably in a different patch).

1753

It is unclear why this is needed.

1754

Will its successor edge be discarded later?

[congh]

In D15259#305550, @davidxl wrote:

Are there any test case changes?

There is no test changes.

lib/CodeGen/IfConversion.cpp

1263

I tried removing this normalization and there isn't new test failures. So I think this is redundant. I added a runtime check instead, which is a new interface of MBB that checks if all successors' probabilities sum up to one and produces assertion failure if not.

1719

As this is a very small change, I think it is ok to do it in this patch. Done.

1753

This is needed as here we have adjusted its successors's probabilities which cannot guarantee the sum is one.

1754

I checked all uses of this function and found after each call FromBBI.BB's successors are all normalized somehow. I commented this line and didn't see any new failures in my test. So we don't have to normalize FromBBI.BB's successors.

[congh]

lib/CodeGen/IfConversion.cpp

1263

Unfortunately the validation fails several times here, which means sum != 1 for BBI.BB's successors. However, they are going to be normalized later in RemoveExtraEdges(BBI) so it seems the normalization is still redundant.

The if converter allows many invalid temporary states which are later fixed by some clean-up functions. This makes updating and validating probabilities quite difficult.

[congh]

Update the patch according to David's comment.

[davidxl]

LGTM

include/llvm/CodeGen/MachineBasicBlock.h

public:
   /// BPI is not available (e.g. -O0 is used), in which case edge probabilities
   /// won't be used. Using this interface can save some space.
   void addSuccessorWithoutProb(MachineBasicBlock *Succ);
 
   /// Set successor probability of a given iterator.
   void setSuccProbability(succ_iterator I, BranchProbability Prob);
 
   /// Normalize probabilities of all successors so that the sum of them becomes
-  /// one.
+  /// one. This is usually done when the current update on this MBB is done, and
+  /// the sum of its successors' probabilities is not guaranteed to be one. The
+  /// user is responsible for the correct use of this function.
+  /// MBB::removeSuccessor() has an option to do this automatically.
   void normalizeSuccProbs() {
     BranchProbability::normalizeProbabilities(Probs.begin(), Probs.end());
   }
 
   /// Remove successor from the successors list of this MachineBasicBlock. The
   /// Predecessors list of Succ is automatically updated.
-  void removeSuccessor(MachineBasicBlock *Succ);
+  /// If NormalizeSuccProbs is true, then normalize successors' probabilities
+  /// after the successor is removed.
+  void removeSuccessor(MachineBasicBlock *Succ,
+                       bool NormalizeSuccProbs = false);
 
   /// Remove specified successor from the successors list of this
   /// MachineBasicBlock. The Predecessors list of Succ is automatically updated.
+  /// If NormalizeSuccProbs is true, then normalize successors' probabilities
+  /// after the successor is removed.
   /// Return the iterator to the element after the one removed.
-  succ_iterator removeSuccessor(succ_iterator I);
+  succ_iterator removeSuccessor(succ_iterator I,
+                                bool NormalizeSuccProbs = false);
 
   /// Replace successor OLD with NEW and update probability info.
   void replaceSuccessor(MachineBasicBlock *Old, MachineBasicBlock *New);
 
   /// Transfers all the successors from MBB to this machine basic block (i.e.,
   /// copies all the successors FromMBB and remove all the successors from
   /// FromMBB).
   void transferSuccessors(MachineBasicBlock *FromMBB);

lib/CodeGen/EarlyIfConversion.cpp

void SSAIfConv::convertIf(SmallVectorImpl<MachineBasicBlock*> &RemovedBlocks) {
   bool ExtraPreds = Tail->pred_size() != 2;
   if (ExtraPreds)
     rewritePHIOperands();
   else
     replacePHIInstrs();
 
   // Fix up the CFG, temporarily leave Head without any successors.
   Head->removeSuccessor(TBB);
-  Head->removeSuccessor(FBB);
+  Head->removeSuccessor(FBB, true);
   if (TBB != Tail)
-    TBB->removeSuccessor(Tail);
+    TBB->removeSuccessor(Tail, true);
   if (FBB != Tail)
-    FBB->removeSuccessor(Tail);
+    FBB->removeSuccessor(Tail, true);
 
   // Fix up Head's terminators.
   // It should become a single branch or a fallthrough.
   DebugLoc HeadDL = Head->getFirstTerminator()->getDebugLoc();
   TII->RemoveBranch(*Head);
 
   // Erase the now empty conditional blocks. It is likely that Head can fall
   // through to Tail, and we can join the two blocks.

lib/CodeGen/IfConversion.cpp

bool IfConverter::IfConvertSimple(BBInfo &BBI, IfcvtKind Kind) {
   if (CvtBBI->BB->pred_size() > 1) {
     BBI.NonPredSize -= TII->RemoveBranch(*BBI.BB);
     // Copy instructions in the true block, predicate them, and add them to
     // the entry block.
     CopyAndPredicateBlock(BBI, *CvtBBI, Cond);
 
     // RemoveExtraEdges won't work if the block has an unanalyzable branch, so
     // explicitly remove CvtBBI as a successor.
-    BBI.BB->removeSuccessor(CvtBBI->BB);
+    BBI.BB->removeSuccessor(CvtBBI->BB, true);
   } else {
     RemoveKills(CvtBBI->BB->begin(), CvtBBI->BB->end(), DontKill, *TRI);
     PredicateBlock(*CvtBBI, CvtBBI->BB->end(), Cond);
 
     // Merge converted block into entry block.
     BBI.NonPredSize -= TII->RemoveBranch(*BBI.BB);
     MergeBlocks(BBI, *CvtBBI);
   }
bool IfConverter::IfConvertTriangle(BBInfo &BBI, IfcvtKind Kind) {
   if (CvtBBI->BB->pred_size() > 1) {
     BBI.NonPredSize -= TII->RemoveBranch(*BBI.BB);
     // Copy instructions in the true block, predicate them, and add them to
     // the entry block.
     CopyAndPredicateBlock(BBI, *CvtBBI, Cond, true);
 
     // RemoveExtraEdges won't work if the block has an unanalyzable branch, so
     // explicitly remove CvtBBI as a successor.
-    BBI.BB->removeSuccessor(CvtBBI->BB);
+    BBI.BB->removeSuccessor(CvtBBI->BB, true);
   } else {
     // Predicate the 'true' block after removing its branch.
     CvtBBI->NonPredSize -= TII->RemoveBranch(*CvtBBI->BB);
     PredicateBlock(*CvtBBI, CvtBBI->BB->end(), Cond);
 
     // Now merge the entry of the triangle with the true block.
     BBI.NonPredSize -= TII->RemoveBranch(*BBI.BB);
     MergeBlocks(BBI, *CvtBBI, false);
if (HasEarlyExit) {
     auto NewNext = BBNext + BBCvt * CvtNext;
     auto NewTrueBBIter =
         std::find(BBI.BB->succ_begin(), BBI.BB->succ_end(), NewTrueBB);
     if (NewTrueBBIter != BBI.BB->succ_end())
       BBI.BB->setSuccProbability(NewTrueBBIter, NewNext);
 
     auto NewFalse = BBCvt * CvtFalse;
     TII->InsertBranch(*BBI.BB, CvtBBI->FalseBB, nullptr, RevCond, dl);
     BBI.BB->addSuccessor(CvtBBI->FalseBB, NewFalse);
+    BBI.BB->normalizeSuccProbs();
   }
 
   // Merge in the 'false' block if the 'false' block has no other
   // predecessors. Otherwise, add an unconditional branch to 'false'.
   bool FalseBBDead = false;
   bool IterIfcvt = true;
   bool isFallThrough = canFallThroughTo(BBI.BB, NextBBI->BB);
   if (!isFallThrough) {
if (CanMergeTail) {
       BBI.HasFallThrough = false;
     }
   }
 
   // RemoveExtraEdges won't work if the block has an unanalyzable branch,
   // which can happen here if TailBB is unanalyzable and is merged, so
   // explicitly remove BBI1 and BBI2 as successors.
   BBI.BB->removeSuccessor(BBI1->BB);
-  BBI.BB->removeSuccessor(BBI2->BB);
+  BBI.BB->removeSuccessor(BBI2->BB, true);
   RemoveExtraEdges(BBI);
 
   // Update block info.
   BBI.IsDone = TrueBBI.IsDone = FalseBBI.IsDone = true;
   InvalidatePreds(BBI.BB);
 
   // FIXME: Must maintain LiveIns.
   return true;
if (AddEdges) {
       if (!To2FromProb.isZero())
         NewProb *= To2FromProb;
     }
 
     FromBBI.BB->removeSuccessor(Succ);
 
     if (AddEdges) {
       // If the edge from ToBBI.BB to Succ already exists, update the
       // probability of this edge by adding NewWeight to it. An example is shown
       // below, in which A is ToBBI.BB and B is FromBBI.BB. In this case we
       // don't have to set C as A's successor as it already is. We only need to
       // update the edge probability on A->C. Note that B will not be
       // immediately removed from A's successors. It is possible that B->D is
       // not removed either if D is a fallthrough of B. Later the edge A->D
       // (generated here) and B->D will be combined into one edge. To maintain
       // correct edge probability of this combined edge, we need to set the edge
       // probability of A->B to zero, which is already done above. The edge
if (AddEdges) {
         ToBBI.BB->addSuccessor(Succ, NewProb);
     }
   }
 
   // Now FromBBI always falls through to the next block!
   if (NBB && !FromBBI.BB->isSuccessor(NBB))
     FromBBI.BB->addSuccessor(NBB);
 
+  ToBBI.BB->normalizeSuccProbs();
+  FromBBI.BB->normalizeSuccProbs();
+
   ToBBI.Predicate.append(FromBBI.Predicate.begin(), FromBBI.Predicate.end());
   FromBBI.Predicate.clear();
 
   ToBBI.NonPredSize += FromBBI.NonPredSize;
   ToBBI.ExtraCost += FromBBI.ExtraCost;
   ToBBI.ExtraCost2 += FromBBI.ExtraCost2;
   FromBBI.NonPredSize = 0;
   FromBBI.ExtraCost = 0;

lib/CodeGen/MachineBasicBlock.cpp

void MachineBasicBlock::addSuccessorWithoutProb(MachineBasicBlock *Succ) {
   // We need to make sure probability list is either empty or has the same size
   // of successor list. When this function is called, we can safely delete all
   // probability in the list.
   Probs.clear();
   Successors.push_back(Succ);
   Succ->addPredecessor(this);
 }
 
-void MachineBasicBlock::removeSuccessor(MachineBasicBlock *Succ) {
+void MachineBasicBlock::removeSuccessor(MachineBasicBlock *Succ,
+                                        bool NormalizeSuccProbs) {
   succ_iterator I = std::find(Successors.begin(), Successors.end(), Succ);
-  removeSuccessor(I);
+  removeSuccessor(I, NormalizeSuccProbs);
 }
 
 MachineBasicBlock::succ_iterator
-MachineBasicBlock::removeSuccessor(succ_iterator I) {
+MachineBasicBlock::removeSuccessor(succ_iterator I, bool NormalizeSuccProbs) {
   assert(I != Successors.end() && "Not a current successor!");
 
   // If probability list is empty it means we don't use it (disabled
   // optimization).
   if (!Probs.empty()) {
     probability_iterator WI = getProbabilityIterator(I);
     Probs.erase(WI);
+    if (NormalizeSuccProbs)
+      normalizeSuccProbs();
   }
 
   (*I)->removePredecessor(this);
   return Successors.erase(I);
 }
 
 void MachineBasicBlock::replaceSuccessor(MachineBasicBlock *Old,
                                          MachineBasicBlock *New) {
while (!FromMBB->succ_empty()) {
     for (MachineBasicBlock::instr_iterator MI = Succ->instr_begin(),
            ME = Succ->instr_end(); MI != ME && MI->isPHI(); ++MI)
       for (unsigned i = 2, e = MI->getNumOperands()+1; i != e; i += 2) {
         MachineOperand &MO = MI->getOperand(i);
         if (MO.getMBB() == FromMBB)
           MO.setMBB(this);
       }
   }
+  normalizeSuccProbs();
 }
 
 bool MachineBasicBlock::isPredecessor(const MachineBasicBlock *MBB) const {
   return std::find(pred_begin(), pred_end(), MBB) != pred_end();
 }
 
 bool MachineBasicBlock::isSuccessor(const MachineBasicBlock *MBB) const {
   return std::find(succ_begin(), succ_end(), MBB) != succ_end();
if (!SeenMBBs.insert(MBB).second ||
       // This is a superfluous edge, remove it.
       SI = removeSuccessor(SI);
       Changed = true;
     } else {
       ++SI;
     }
   }
 
+  if (Changed)
+    normalizeSuccProbs();
   return Changed;
 }
 
 /// Find the next valid DebugLoc starting at MBBI, skipping any DBG_VALUE
 /// instructions.  Return UnknownLoc if there is none.
 DebugLoc
 MachineBasicBlock::findDebugLoc(instr_iterator MBBI) {
   DebugLoc DL;

lib/Target/AArch64/AArch64ConditionalCompares.cpp

 
 void SSACCmpConv::convert(SmallVectorImpl<MachineBasicBlock *> &RemovedBlocks) {
   DEBUG(dbgs() << "Merging BB#" << CmpBB->getNumber() << " into BB#"
                << Head->getNumber() << ":\n" << *CmpBB);
 
   // All CmpBB instructions are moved into Head, and CmpBB is deleted.
   // Update the CFG first.
   updateTailPHIs();
-  Head->removeSuccessor(CmpBB);
-  CmpBB->removeSuccessor(Tail);
+  Head->removeSuccessor(CmpBB, true);
+  CmpBB->removeSuccessor(Tail, true);
   Head->transferSuccessorsAndUpdatePHIs(CmpBB);
   DebugLoc TermDL = Head->getFirstTerminator()->getDebugLoc();
   TII->RemoveBranch(*Head);
 
   // If the Head terminator was one of the cbz / tbz branches with built-in
   // compare, we need to insert an explicit compare instruction in its place.
   if (HeadCond[0].getImm() == -1) {
     ++NumCompBranches;

lib/Target/AMDGPU/AMDILCFGStructurizer.cpp

if (LoopRep->contains(MBB)) {
                           LoopHeader, LoopRep);
       ContMBB.push_back(MBB);
       ++NumCont;
     }
   }
 
   for (SmallVectorImpl<MachineBasicBlock *>::iterator It = ContMBB.begin(),
       E = ContMBB.end(); It != E; ++It) {
-    (*It)->removeSuccessor(LoopHeader);
+    (*It)->removeSuccessor(LoopHeader, true);
   }
 
   numLoopcontPatternMatch += NumCont;
 
   return NumCont;
 }
 
 

 void AMDGPUCFGStructurizer::mergeSerialBlock(MachineBasicBlock *DstMBB,
     MachineBasicBlock *SrcMBB) {
   DEBUG(
     dbgs() << "serialPattern BB" << DstMBB->getNumber()
            << " <= BB" << SrcMBB->getNumber() << "\n";
   );
   DstMBB->splice(DstMBB->end(), SrcMBB, SrcMBB->begin(), SrcMBB->end());
 
-  DstMBB->removeSuccessor(SrcMBB);
+  DstMBB->removeSuccessor(SrcMBB, true);
   cloneSuccessorList(DstMBB, SrcMBB);
 
   removeSuccessor(SrcMBB);
   MLI->removeBlock(SrcMBB);
   retireBlock(SrcMBB);
 }
 
 void AMDGPUCFGStructurizer::mergeIfthenelseBlock(MachineInstr *BranchMI,

 //    landBlk
 
   MachineBasicBlock::iterator I = BranchMI;
   insertCondBranchBefore(I, getBranchNzeroOpcode(OldOpcode),
       BranchDL);
 
   if (TrueMBB) {
     MBB->splice(I, TrueMBB, TrueMBB->begin(), TrueMBB->end());
-    MBB->removeSuccessor(TrueMBB);
+    MBB->removeSuccessor(TrueMBB, true);
     if (LandMBB && TrueMBB->succ_size()!=0)
-      TrueMBB->removeSuccessor(LandMBB);
+      TrueMBB->removeSuccessor(LandMBB, true);
     retireBlock(TrueMBB);
     MLI->removeBlock(TrueMBB);
   }
 
   if (FalseMBB) {
     insertInstrBefore(I, AMDGPU::ELSE);
     MBB->splice(I, FalseMBB, FalseMBB->begin(),
                    FalseMBB->end());
-    MBB->removeSuccessor(FalseMBB);
+    MBB->removeSuccessor(FalseMBB, true);
     if (LandMBB && FalseMBB->succ_size() != 0)
-      FalseMBB->removeSuccessor(LandMBB);
+      FalseMBB->removeSuccessor(LandMBB, true);
     retireBlock(FalseMBB);
     MLI->removeBlock(FalseMBB);
   }
   insertInstrBefore(I, AMDGPU::ENDIF);
 
   BranchMI->eraseFromParent();
 
   if (LandMBB && TrueMBB && FalseMBB)
void AMDGPUCFGStructurizer::mergeLoopbreakBlock(MachineBasicBlock *ExitingMBB,
   if (TrueBranch != LandMBB)
     reversePredicateSetter(I);
   insertCondBranchBefore(ExitingMBB, I, AMDGPU::IF_PREDICATE_SET, AMDGPU::PREDICATE_BIT, DL);
   insertInstrBefore(I, AMDGPU::BREAK);
   insertInstrBefore(I, AMDGPU::ENDIF);
   //now branchInst can be erase safely
   BranchMI->eraseFromParent();
   //now take care of successors, retire blocks
-  ExitingMBB->removeSuccessor(LandMBB);
+  ExitingMBB->removeSuccessor(LandMBB, true);
 }
 
 void AMDGPUCFGStructurizer::settleLoopcontBlock(MachineBasicBlock *ContingMBB,
     MachineBasicBlock *ContMBB) {
   DEBUG(dbgs() << "settleLoopcontBlock conting = BB"
                << ContingMBB->getNumber()
                << ", cont = BB" << ContMBB->getNumber() << "\n";);
 
void AMDGPUCFGStructurizer::removeRedundantConditionalBranch(
   if (MBB1 != MBB2)
     return;
 
   MachineInstr *BranchMI = getNormalBlockBranchInstr(MBB);
   assert(BranchMI && isCondBranch(BranchMI));
   DEBUG(dbgs() << "Removing unneeded cond branch instr: " << *BranchMI);
   BranchMI->eraseFromParent();
   SHOWNEWBLK(MBB1, "Removing redundant successor");
-  MBB->removeSuccessor(MBB1);
+  MBB->removeSuccessor(MBB1, true);
 }
 
 void AMDGPUCFGStructurizer::addDummyExitBlock(
     SmallVectorImpl<MachineBasicBlock*> &RetMBB) {
   MachineBasicBlock *DummyExitBlk = FuncRep->CreateMachineBasicBlock();
   FuncRep->push_back(DummyExitBlk);  //insert to function
   insertInstrEnd(DummyExitBlk, AMDGPU::RETURN);
 

lib/Target/ARM/ARMISelLowering.cpp

while (!Successors.empty()) {
       MachineBasicBlock *SMBB = Successors.pop_back_val();
       if (SMBB->isEHPad()) {
         BB->removeSuccessor(SMBB);
         MBBLPads.push_back(SMBB);
       }
     }
 
     BB->addSuccessor(DispatchBB, BranchProbability::getZero());
+    BB->normalizeSuccProbs();
 
     // Find the invoke call and mark all of the callee-saved registers as
     // 'implicit defined' so that they're spilled. This prevents code from
     // moving instructions to before the EH block, where they will never be
     // executed.
     for (MachineBasicBlock::reverse_iterator
            II = BB->rbegin(), IE = BB->rend(); II != IE; ++II) {
       if (!II->isCall()) continue;

lib/Target/Hexagon/HexagonEarlyIfConv.cpp

for (DTNodeVectType::iterator I = Cn.begin(), E = Cn.end(); I != E; ++I) {
       MDT->changeImmediateDominator(SB, IDB);
     }
   }
 
   while (B->succ_size() > 0)
     B->removeSuccessor(B->succ_begin());
 
   for (auto I = B->pred_begin(), E = B->pred_end(); I != E; ++I)
-    (*I)->removeSuccessor(B);
+    (*I)->removeSuccessor(B, true);
 
   Deleted.insert(B);
   MDT->eraseNode(B);
   MFN->erase(B->getIterator());
 }
 
 
 void HexagonEarlyIfConversion::eliminatePhis(MachineBasicBlock *B) {
void HexagonEarlyIfConversion::mergeBlocks(MachineBasicBlock *PredB,
   bool TermOk = hasUncondBranch(SuccB);
   eliminatePhis(SuccB);
   TII->RemoveBranch(*PredB);
   PredB->removeSuccessor(SuccB);
   PredB->splice(PredB->end(), SuccB, SuccB->begin(), SuccB->end());
   MachineBasicBlock::succ_iterator I, E = SuccB->succ_end();
   for (I = SuccB->succ_begin(); I != E; ++I)
     PredB->addSuccessor(*I);
+  PredB->normalizeSuccProbs();
   replacePhiEdges(SuccB, PredB);
   removeBlock(SuccB);
   if (!TermOk)
     PredB->updateTerminator();
 }
 
 
 void HexagonEarlyIfConversion::simplifyFlowGraph(const FlowPattern &FP) {

lib/Target/Mips/MipsLongBranch.cpp

void MipsLongBranch::splitMBB(MachineBasicBlock *MBB) {
 
   // Create a new MBB. Move instructions in MBB to the newly created MBB.
   MachineBasicBlock *NewMBB =
     MF->CreateMachineBasicBlock(MBB->getBasicBlock());
 
   // Insert NewMBB and fix control flow.
   MachineBasicBlock *Tgt = getTargetMBB(*FirstBr);
   NewMBB->transferSuccessors(MBB);
-  NewMBB->removeSuccessor(Tgt);
+  NewMBB->removeSuccessor(Tgt, true);
   MBB->addSuccessor(NewMBB);
   MBB->addSuccessor(Tgt);
   MF->insert(std::next(MachineFunction::iterator(MBB)), NewMBB);
 
   NewMBB->splice(NewMBB->end(), MBB, (++LastBr).base(), MBB->end());
 }
 
 // Fill MBBInfos.

lib/Target/PowerPC/PPCEarlyReturn.cpp

bool processBlock(MachineBasicBlock &ReturnMBB) {
           PredToRemove.push_back(*PI);
         }
 
         if (BlockChanged)
           Changed = true;
       }
 
       for (unsigned i = 0, ie = PredToRemove.size(); i != ie; ++i)
-        PredToRemove[i]->removeSuccessor(&ReturnMBB);
+        PredToRemove[i]->removeSuccessor(&ReturnMBB, true);
 
       if (Changed && !ReturnMBB.hasAddressTaken()) {
         // We now might be able to merge this blr-only block into its
         // by-layout predecessor.
         if (ReturnMBB.pred_size() == 1) {
           MachineBasicBlock &PrevMBB = **ReturnMBB.pred_begin();
           if (PrevMBB.isLayoutSuccessor(&ReturnMBB) && PrevMBB.canFallThrough()) {
             // Move the blr into the preceding block.
             PrevMBB.splice(PrevMBB.end(), &ReturnMBB, I);
-            PrevMBB.removeSuccessor(&ReturnMBB);
+            PrevMBB.removeSuccessor(&ReturnMBB, true);
           }
         }
 
         if (ReturnMBB.pred_empty())
           ReturnMBB.eraseFromParent();
       }
 
       return Changed;