diff --git a/llvm/include/llvm/CodeGen/MachineBasicBlock.h b/llvm/include/llvm/CodeGen/MachineBasicBlock.h --- a/llvm/include/llvm/CodeGen/MachineBasicBlock.h +++ b/llvm/include/llvm/CodeGen/MachineBasicBlock.h @@ -492,6 +492,10 @@ void splitSuccessor(MachineBasicBlock *Old, MachineBasicBlock *New, bool NormalizeSuccProbs = false); + /// Update PHI nodes in this MBB to refer to ToMBB instead of FromMBB. + void updatePHIs(MachineBasicBlock *FromMBB, + MachineBasicBlock *ToMBB); + /// 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). diff --git a/llvm/lib/CodeGen/MachineBasicBlock.cpp b/llvm/lib/CodeGen/MachineBasicBlock.cpp --- a/llvm/lib/CodeGen/MachineBasicBlock.cpp +++ b/llvm/lib/CodeGen/MachineBasicBlock.cpp @@ -771,6 +771,16 @@ Predecessors.erase(I); } +void MachineBasicBlock::updatePHIs(MachineBasicBlock *FromMBB, + MachineBasicBlock *ToMBB) { + for (MachineInstr &MI : phis()) + for (unsigned i = 2, e = MI.getNumOperands() + 1; i != e; i += 2) { + MachineOperand &MO = MI.getOperand(i); + if (MO.getMBB() == FromMBB) + MO.setMBB(ToMBB); + } +} + void MachineBasicBlock::transferSuccessors(MachineBasicBlock *FromMBB) { if (this == FromMBB) return; @@ -778,7 +788,8 @@ while (!FromMBB->succ_empty()) { MachineBasicBlock *Succ = *FromMBB->succ_begin(); - // If probability list is empty it means we don't use it (disabled optimization). + // If probability list is empty it means we don't use it (disabled + // optimization). if (!FromMBB->Probs.empty()) { auto Prob = *FromMBB->Probs.begin(); addSuccessor(Succ, Prob); @@ -804,13 +815,7 @@ FromMBB->removeSuccessor(Succ); // Fix up any PHI nodes in the successor. - 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); - } + Succ->updatePHIs(FromMBB, this); } normalizeSuccProbs(); } @@ -985,13 +990,8 @@ } } - // Fix PHI nodes in Succ so they refer to NMBB instead of this - for (MachineBasicBlock::instr_iterator - i = Succ->instr_begin(),e = Succ->instr_end(); - i != e && i->isPHI(); ++i) - for (unsigned ni = 1, ne = i->getNumOperands(); ni != ne; ni += 2) - if (i->getOperand(ni+1).getMBB() == this) - i->getOperand(ni+1).setMBB(NMBB); + // Fix PHI nodes in Succ so they refer to NMBB instead of this. + Succ->updatePHIs(this, NMBB); // Inherit live-ins from the successor for (const auto &LI : Succ->liveins()) diff --git a/llvm/lib/CodeGen/MachinePipeliner.cpp b/llvm/lib/CodeGen/MachinePipeliner.cpp --- a/llvm/lib/CodeGen/MachinePipeliner.cpp +++ b/llvm/lib/CodeGen/MachinePipeliner.cpp @@ -2234,15 +2234,7 @@ } // Fix any Phi nodes in the loop exit block. - for (MachineInstr &MI : *LoopExitBB) { - if (!MI.isPHI()) - break; - for (unsigned i = 2, e = MI.getNumOperands() + 1; i != e; i += 2) { - MachineOperand &MO = MI.getOperand(i); - if (MO.getMBB() == BB) - MO.setMBB(PredBB); - } - } + LoopExitBB->updatePHIs(BB, PredBB); // Create a branch to the new epilog from the kernel. // Remove the original branch and add a new branch to the epilog.