Expand cloneLoopWithPreheader() to support cloning loop nest

Summary: cloneLoopWithPreheader() currently only support innermost loop,
and assert otherwise.
Reviewers: Meinersbur, fhahn, kbarton
Reviewed By: Meinersbur
Subscribers: hiraditya, jsji, llvm-commits
Tag: LLVM
Differential Revision: https://reviews.llvm.org/D63446

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@364310 91177308-0d34-0410-b5e6-96231b3b80d8

include/llvm/Analysis/LoopInfo.h

@@ -311,6 +311,40 @@ template <class BlockT, class LoopT> class LoopBase {
         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 <class Type>
+  static void getInnerLoopsInPreorder(const LoopT &L,
+                                      SmallVectorImpl<Type> &PreOrderLoops) {
+    SmallVector<LoopT *, 4> 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<const LoopT *, 4> getLoopsInPreorder() const {
+    SmallVector<const LoopT *, 4> PreOrderLoops;
+    const LoopT *CurLoop = static_cast<const LoopT *>(this);
+    PreOrderLoops.push_back(CurLoop);
+    getInnerLoopsInPreorder(*CurLoop, PreOrderLoops);
+    return PreOrderLoops;
+  }
+  SmallVector<LoopT *, 4> getLoopsInPreorder() {
+    SmallVector<LoopT *, 4> PreOrderLoops;
+    LoopT *CurLoop = static_cast<LoopT *>(this);
+    PreOrderLoops.push_back(CurLoop);
+    getInnerLoopsInPreorder(*CurLoop, PreOrderLoops);
+    return PreOrderLoops;
+  }
+
   //===--------------------------------------------------------------------===//
   // APIs for updating loop information after changing the CFG
   //

include/llvm/Analysis/LoopInfoImpl.h

@@ -587,16 +587,9 @@ SmallVector<LoopT *, 4> LoopInfoBase<BlockT, LoopT>::getLoopsInPreorder() {
   // 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;

lib/Transforms/Utils/CloneFunction.cpp

@@ -739,12 +739,12 @@ Loop *llvm::cloneLoopWithPreheader(BasicBlock *Before, BasicBlock *LoopDomBB,
                                    const Twine &NameSuffix, LoopInfo *LI,
                                    DominatorTree *DT,
                                    SmallVectorImpl<BasicBlock *> &Blocks) {
-  assert(OrigLoop->getSubLoops().empty() &&
-         "Loop to be cloned cannot have inner loop");
   Function *F = OrigLoop->getHeader()->getParent();
   Loop *ParentLoop = OrigLoop->getParentLoop();
+  DenseMap<Loop *, Loop *> LMap;
 
   Loop *NewLoop = LI->AllocateLoop();
+  LMap[OrigLoop] = NewLoop;
   if (ParentLoop)
     ParentLoop->addChildLoop(NewLoop);
   else
@@ -764,17 +764,38 @@ Loop *llvm::cloneLoopWithPreheader(BasicBlock *Before, BasicBlock *LoopDomBB,
   // 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()) {