Expand cloneLoopWithPreheader() to support cloning loop nest

[Whitney]

Summary

cloneLoopWithPreheader() currently only support innermost loop, and assert otherwise.

[Meinersbur]

llvm/lib/Transforms/Utils/CloneFunction.cpp

744

Consider using DenseMap.

771

[style] I like getting a reference that can be updated instead of another map lookup. I.e.

Loop *&NewLoop = LMap[CurLoop];
...
NewLoop = LI->AllocateLoop();

776–777

[serious] Since the iteration order of OrigLoop->getBlocks() is undefined, we might get a block that is nested in two loops without having visited the parent first, i.e. NewParentLoop would be null in this case.

Consider iterating over LoopInfo's loop tree before iterating over all block to ensure that the new loop structure has been created.

[Whitney]

Addressed all review comments from Michael.

[Whitney]

llvm/lib/Transforms/Utils/CloneFunction.cpp

776–777

Good catch! Modified as suggested.

[Whitney]

@Meinersbur Thanks for the review! What do you think of the updated changeset?

[Meinersbur]

LGTM

llvm/include/llvm/Analysis/LoopInfo.h

#endif
   void getLoopLatches(SmallVectorImpl<BlockT *> &LoopLatches) const {
     assert(!isInvalid() && "Loop not in a valid state!");
     BlockT *H = getHeader();
     for (const auto Pred : children<Inverse<BlockT *>>(H))
       if (contains(Pred))
         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
   //
 
   /// This method is used by other analyses to update loop information.
   /// NewBB is set to be a new member of the current loop.
   /// Because of this, it is added as a member of all parent loops, and is added
   /// to the specified LoopInfo object as being in the current basic block.  It

llvm/include/llvm/Analysis/LoopInfoImpl.h

 SmallVector<LoopT *, 4> LoopInfoBase<BlockT, LoopT>::getLoopsInPreorder() {
   SmallVector<LoopT *, 4> PreOrderLoops, PreOrderWorklist;
   // The outer-most loop actually goes into the result in the same relative
   // order as we walk it. But LoopInfo stores the top level loops in reverse
   // program order so for here we reverse it to get forward program order.
   // 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;
 }
 
 template <class BlockT, class LoopT>
 SmallVector<LoopT *, 4>
 LoopInfoBase<BlockT, LoopT>::getLoopsInReverseSiblingPreorder() {

llvm/lib/Transforms/Utils/CloneFunction.cpp

 ///
 /// Updates LoopInfo and DominatorTree assuming the loop is dominated by block
 /// \p LoopDomBB.  Insert the new blocks before block specified in \p Before.
 Loop *llvm::cloneLoopWithPreheader(BasicBlock *Before, BasicBlock *LoopDomBB,
                                    Loop *OrigLoop, ValueToValueMapTy &VMap,
                                    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
     LI->addTopLevelLoop(NewLoop);
 
   BasicBlock *OrigPH = OrigLoop->getLoopPreheader();
   assert(OrigPH && "No preheader");
   BasicBlock *NewPH = CloneBasicBlock(OrigPH, VMap, NameSuffix, F);
   // To rename the loop PHIs.
   VMap[OrigPH] = NewPH;
   Blocks.push_back(NewPH);
 
   // Update LoopInfo.
   if (ParentLoop)
     ParentLoop->addBasicBlockToLoop(NewPH, *LI);
 
   // Update DominatorTree.
   DT->addNewBlock(NewPH, LoopDomBB);
 
-  for (BasicBlock *BB : OrigLoop->getBlocks()) {
+  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();
+
+        // 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);
+      }
+
       BasicBlock *NewBB = CloneBasicBlock(BB, VMap, NameSuffix, F);
       VMap[BB] = 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.
+      // Add DominatorTree node. After seeing all blocks, update to correct
+      // IDom.
       DT->addNewBlock(NewBB, NewPH);
 
       Blocks.push_back(NewBB);
     }
+  }
 
   for (BasicBlock *BB : OrigLoop->getBlocks()) {
     // Update DominatorTree.
     BasicBlock *IDomBB = DT->getNode(BB)->getIDom()->getBlock();
     DT->changeImmediateDominator(cast<BasicBlock>(VMap[BB]),
                                  cast<BasicBlock>(VMap[IDomBB]));
   }