Index: include/llvm/Analysis/LoopInfo.h =================================================================== --- include/llvm/Analysis/LoopInfo.h +++ include/llvm/Analysis/LoopInfo.h @@ -277,11 +277,22 @@ /// Otherwise return null. BlockT *getUniqueExitBlock() const; - /// Edge type. - typedef std::pair Edge; - /// Return all pairs of (_inside_block_,_outside_block_). - void getExitEdges(SmallVectorImpl &ExitEdges) const; + template + void getExitEdges(SmallVectorImpl > & + ExitEdges) const { + // Assert that ParamBlockT and BlockT are the same +/- const qualifier. + static_assert( + std::is_same::type, + typename std::remove_const::type>::value, + "not same!"); + assert(!isInvalid() && "Loop not in a valid state!"); + for (const auto BB : blocks()) + for (const auto &Succ : children(BB)) + if (!contains(Succ)) + // Not in current loop? It must be an exit block. + ExitEdges.emplace_back(BB, Succ); + } /// If there is a preheader for this loop, return it. A loop has a preheader /// if there is only one edge to the header of the loop from outside of the Index: include/llvm/Analysis/LoopInfoImpl.h =================================================================== --- include/llvm/Analysis/LoopInfoImpl.h +++ include/llvm/Analysis/LoopInfoImpl.h @@ -149,18 +149,6 @@ return nullptr; } -/// getExitEdges - Return all pairs of (_inside_block_,_outside_block_). -template -void LoopBase::getExitEdges( - SmallVectorImpl &ExitEdges) const { - assert(!isInvalid() && "Loop not in a valid state!"); - for (const auto BB : blocks()) - for (const auto &Succ : children(BB)) - if (!contains(Succ)) - // Not in current loop? It must be an exit block. - ExitEdges.emplace_back(BB, Succ); -} - /// getLoopPreheader - If there is a preheader for this loop, return it. A /// loop has a preheader if there is only one edge to the header of the loop /// from outside of the loop and it is legal to hoist instructions into the Index: lib/Transforms/Utils/LoopUnroll.cpp =================================================================== --- lib/Transforms/Utils/LoopUnroll.cpp +++ lib/Transforms/Utils/LoopUnroll.cpp @@ -401,7 +401,12 @@ // Successful peeling may result in a change in the loop preheader/trip // counts. If we later unroll the loop, we want these to be updated. if (Peeled) { - BasicBlock *ExitingBlock = L->getExitingBlock(); + BasicBlock *ExitingBlock = nullptr; + for (auto *Succ : successors(L->getLoopLatch())) + if (!L->contains(Succ)) { + ExitingBlock = Succ; + break; + } assert(ExitingBlock && "Loop without exiting block?"); Preheader = L->getLoopPreheader(); ULO.TripCount = SE->getSmallConstantTripCount(L, ExitingBlock); Index: lib/Transforms/Utils/LoopUnrollPeel.cpp =================================================================== --- lib/Transforms/Utils/LoopUnrollPeel.cpp +++ lib/Transforms/Utils/LoopUnrollPeel.cpp @@ -461,18 +461,18 @@ /// InsertBot. /// \param IterNumber The serial number of the iteration currently being /// peeled off. -/// \param Exit The exit block of the original loop. +/// \param ExitEdges The exit edges of the original loop. /// \param[out] NewBlocks A list of the blocks in the newly created clone /// \param[out] VMap The value map between the loop and the new clone. /// \param LoopBlocks A helper for DFS-traversal of the loop. /// \param LVMap A value-map that maps instructions from the original loop to /// instructions in the last peeled-off iteration. -static void cloneLoopBlocks(Loop *L, unsigned IterNumber, BasicBlock *InsertTop, - BasicBlock *InsertBot, BasicBlock *Exit, - SmallVectorImpl &NewBlocks, - LoopBlocksDFS &LoopBlocks, ValueToValueMapTy &VMap, - ValueToValueMapTy &LVMap, DominatorTree *DT, - LoopInfo *LI) { +static void cloneLoopBlocks( + Loop *L, unsigned IterNumber, BasicBlock *InsertTop, BasicBlock *InsertBot, + SmallVectorImpl > &ExitEdges, + SmallVectorImpl &NewBlocks, LoopBlocksDFS &LoopBlocks, + ValueToValueMapTy &VMap, ValueToValueMapTy &LVMap, DominatorTree *DT, + LoopInfo *LI) { BasicBlock *Header = L->getHeader(); BasicBlock *Latch = L->getLoopLatch(); BasicBlock *PreHeader = L->getLoopPreheader(); @@ -518,9 +518,11 @@ // iteration (for every other iteration) BasicBlock *NewLatch = cast(VMap[Latch]); BranchInst *LatchBR = cast(NewLatch->getTerminator()); - unsigned HeaderIdx = (LatchBR->getSuccessor(0) == Header ? 0 : 1); - LatchBR->setSuccessor(HeaderIdx, InsertBot); - LatchBR->setSuccessor(1 - HeaderIdx, Exit); + for (unsigned idx = 0; idx < LatchBR->getNumSuccessors(); ++idx) + if (LatchBR->getSuccessor(idx) == Header) { + LatchBR->setSuccessor(idx, InsertBot); + break; + } if (DT) DT->changeImmediateDominator(InsertBot, NewLatch); @@ -551,14 +553,15 @@ // we've just created. Note that this must happen *after* the incoming // values are adjusted, since the value going out of the latch may also be // a value coming into the header. - for (BasicBlock::iterator I = Exit->begin(); isa(I); ++I) { - PHINode *PHI = cast(I); - Value *LatchVal = PHI->getIncomingValueForBlock(Latch); - Instruction *LatchInst = dyn_cast(LatchVal); - if (LatchInst && L->contains(LatchInst)) - LatchVal = VMap[LatchVal]; - PHI->addIncoming(LatchVal, cast(VMap[Latch])); - } + for (auto It : ExitEdges) + for (BasicBlock::iterator I = It.second->begin(); isa(I); ++I) { + PHINode *PHI = cast(I); + Value *LatchVal = PHI->getIncomingValueForBlock(It.first); + Instruction *LatchInst = dyn_cast(LatchVal); + if (LatchInst && L->contains(LatchInst)) + LatchVal = VMap[LatchVal]; + PHI->addIncoming(LatchVal, cast(VMap[It.first])); + } // LastValueMap is updated with the values for the current loop // which are used the next time this function is called. @@ -587,7 +590,8 @@ BasicBlock *Header = L->getHeader(); BasicBlock *PreHeader = L->getLoopPreheader(); BasicBlock *Latch = L->getLoopLatch(); - BasicBlock *Exit = L->getUniqueExitBlock(); + SmallVector, 4> ExitEdges; + L->getExitEdges(ExitEdges); Function *F = Header->getParent(); @@ -669,8 +673,8 @@ else CurHeaderWeight = 1; - cloneLoopBlocks(L, Iter, InsertTop, InsertBot, Exit, - NewBlocks, LoopBlocks, VMap, LVMap, DT, LI); + cloneLoopBlocks(L, Iter, InsertTop, InsertBot, ExitEdges, NewBlocks, + LoopBlocks, VMap, LVMap, DT, LI); // Remap to use values from the current iteration instead of the // previous one. @@ -681,7 +685,9 @@ // latter is the first cloned loop body, as original PreHeader dominates // the original loop body. if (Iter == 0) - DT->changeImmediateDominator(Exit, cast(LVMap[Latch])); + for (auto It : ExitEdges) + DT->changeImmediateDominator(It.second, + cast(LVMap[It.first])); #ifdef EXPENSIVE_CHECKS assert(DT->verify(DominatorTree::VerificationLevel::Fast)); #endif