Index: llvm/include/llvm/Analysis/LoopInfo.h =================================================================== --- llvm/include/llvm/Analysis/LoopInfo.h +++ llvm/include/llvm/Analysis/LoopInfo.h @@ -311,6 +311,40 @@ LoopLatches.push_back(Pred); } + /// Return all inner loops in the loop nest rooted by the loop in preorder, + /// with siblings in forward program order. + template + static void getInnerLoopsInPreorder(const LoopT &L, + SmallVectorImpl &PreOrderLoops) { + SmallVector PreOrderWorklist; + PreOrderWorklist.append(L.rbegin(), L.rend()); + + while (!PreOrderWorklist.empty()) { + LoopT *L = PreOrderWorklist.pop_back_val(); + // Sub-loops are stored in forward program order, but will process the + // worklist backwards so append them in reverse order. + PreOrderWorklist.append(L->rbegin(), L->rend()); + PreOrderLoops.push_back(L); + } + } + + /// Return all loops in the loop nest rooted by the loop in preorder, with + /// siblings in forward program order. + SmallVector getLoopsInPreorder() const { + SmallVector PreOrderLoops; + const LoopT *CurLoop = static_cast(this); + PreOrderLoops.push_back(CurLoop); + getInnerLoopsInPreorder(*CurLoop, PreOrderLoops); + return PreOrderLoops; + } + SmallVector getLoopsInPreorder() { + SmallVector PreOrderLoops; + LoopT *CurLoop = static_cast(this); + PreOrderLoops.push_back(CurLoop); + getInnerLoopsInPreorder(*CurLoop, PreOrderLoops); + return PreOrderLoops; + } + //===--------------------------------------------------------------------===// // APIs for updating loop information after changing the CFG // Index: llvm/include/llvm/Analysis/LoopInfoImpl.h =================================================================== --- llvm/include/llvm/Analysis/LoopInfoImpl.h +++ llvm/include/llvm/Analysis/LoopInfoImpl.h @@ -587,16 +587,9 @@ // FIXME: If we change the order of LoopInfo we will want to remove the // reverse here. for (LoopT *RootL : reverse(*this)) { - assert(PreOrderWorklist.empty() && - "Must start with an empty preorder walk worklist."); - PreOrderWorklist.push_back(RootL); - do { - LoopT *L = PreOrderWorklist.pop_back_val(); - // Sub-loops are stored in forward program order, but will process the - // worklist backwards so append them in reverse order. - PreOrderWorklist.append(L->rbegin(), L->rend()); - PreOrderLoops.push_back(L); - } while (!PreOrderWorklist.empty()); + auto PreOrderLoopsInRootL = RootL->getLoopsInPreorder(); + PreOrderLoops.append(PreOrderLoopsInRootL.begin(), + PreOrderLoopsInRootL.end()); } return PreOrderLoops; Index: llvm/lib/Transforms/Utils/CloneFunction.cpp =================================================================== --- llvm/lib/Transforms/Utils/CloneFunction.cpp +++ llvm/lib/Transforms/Utils/CloneFunction.cpp @@ -739,12 +739,12 @@ const Twine &NameSuffix, LoopInfo *LI, DominatorTree *DT, SmallVectorImpl &Blocks) { - assert(OrigLoop->getSubLoops().empty() && - "Loop to be cloned cannot have inner loop"); Function *F = OrigLoop->getHeader()->getParent(); Loop *ParentLoop = OrigLoop->getParentLoop(); + DenseMap LMap; Loop *NewLoop = LI->AllocateLoop(); + LMap[OrigLoop] = NewLoop; if (ParentLoop) ParentLoop->addChildLoop(NewLoop); else @@ -764,17 +764,38 @@ // Update DominatorTree. DT->addNewBlock(NewPH, LoopDomBB); - for (BasicBlock *BB : OrigLoop->getBlocks()) { - BasicBlock *NewBB = CloneBasicBlock(BB, VMap, NameSuffix, F); - VMap[BB] = NewBB; + for (Loop *CurLoop : OrigLoop->getLoopsInPreorder()) { + for (BasicBlock *BB : CurLoop->getBlocks()) { + if (CurLoop != LI->getLoopFor(BB)) + continue; + + Loop *&NewLoop = LMap[CurLoop]; + if (!NewLoop) { + NewLoop = LI->AllocateLoop(); - // Update LoopInfo. - NewLoop->addBasicBlockToLoop(NewBB, *LI); + // Establish the parent/child relationship. + Loop *OrigParent = CurLoop->getParentLoop(); + assert(OrigParent && "Could not find the original parent loop"); + Loop *NewParentLoop = LMap[OrigParent]; + assert(NewParentLoop && "Could not find the new parent loop"); + + NewParentLoop->addChildLoop(NewLoop); + } - // Add DominatorTree node. After seeing all blocks, update to correct IDom. - DT->addNewBlock(NewBB, NewPH); + BasicBlock *NewBB = CloneBasicBlock(BB, VMap, NameSuffix, F); + VMap[BB] = NewBB; - Blocks.push_back(NewBB); + // Update LoopInfo. + NewLoop->addBasicBlockToLoop(NewBB, *LI); + if (BB == CurLoop->getHeader()) + NewLoop->moveToHeader(NewBB); + + // Add DominatorTree node. After seeing all blocks, update to correct + // IDom. + DT->addNewBlock(NewBB, NewPH); + + Blocks.push_back(NewBB); + } } for (BasicBlock *BB : OrigLoop->getBlocks()) {