Index: llvm/trunk/lib/Transforms/Utils/LoopUnroll.cpp =================================================================== --- llvm/trunk/lib/Transforms/Utils/LoopUnroll.cpp +++ llvm/trunk/lib/Transforms/Utils/LoopUnroll.cpp @@ -401,8 +401,11 @@ // 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(); + // According to our guards and profitability checks the only + // meaningful exit should be latch block. + BasicBlock *ExitingBlock = L->getLoopLatch(); assert(ExitingBlock && "Loop without exiting block?"); + assert(L->isLoopExiting(ExitingBlock) && "Latch is not exiting?"); Preheader = L->getLoopPreheader(); ULO.TripCount = SE->getSmallConstantTripCount(L, ExitingBlock); ULO.TripMultiple = SE->getSmallConstantTripMultiple(L, ExitingBlock); Index: llvm/trunk/lib/Transforms/Utils/LoopUnrollPeel.cpp =================================================================== --- llvm/trunk/lib/Transforms/Utils/LoopUnrollPeel.cpp +++ llvm/trunk/lib/Transforms/Utils/LoopUnrollPeel.cpp @@ -432,18 +432,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(); @@ -489,9 +489,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, e = LatchBR->getNumSuccessors(); idx < e; ++idx) + if (LatchBR->getSuccessor(idx) == Header) { + LatchBR->setSuccessor(idx, InsertBot); + break; + } if (DT) DT->changeImmediateDominator(InsertBot, NewLatch); @@ -522,14 +524,14 @@ // 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 Edge : ExitEdges) + for (PHINode &PHI : Edge.second->phis()) { + Value *LatchVal = PHI.getIncomingValueForBlock(Edge.first); + Instruction *LatchInst = dyn_cast(LatchVal); + if (LatchInst && L->contains(LatchInst)) + LatchVal = VMap[LatchVal]; + PHI.addIncoming(LatchVal, cast(VMap[Edge.first])); + } // LastValueMap is updated with the values for the current loop // which are used the next time this function is called. @@ -558,7 +560,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(); @@ -640,8 +643,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. @@ -652,7 +655,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 Edge : ExitEdges) + DT->changeImmediateDominator(Edge.second, + cast(LVMap[Edge.first])); #ifdef EXPENSIVE_CHECKS assert(DT->verify(DominatorTree::VerificationLevel::Fast)); #endif