Loop Utils: Cloning any level loop nest

[mivnay]

Summary

Currently, only inner-most loops are allowed to be cloned. This changes the cloneLoopWithPreheader() function to clone any level loop. This also updates the DominatorTree and LoopInfo.

Thanks,
Vinay Madhusudan,
Compiler Tree Technologies

lib/Transforms/Utils/CloneFunction.cpp

void llvm::remapInstructionsInBlocks(
     const SmallVectorImpl<BasicBlock *> &Blocks, ValueToValueMapTy &VMap) {
   // Rewrite the code to refer to itself.
   for (auto *BB : Blocks)
     for (auto &Inst : *BB)
       RemapInstruction(&Inst, VMap,
                        RF_NoModuleLevelChanges | RF_IgnoreMissingLocals);
 }
 
+/// \brief Clones the original loop \p OrigLoop structure
+/// and keeps it ready to add the basic blocks.
+static void createNewLoops(Loop *OrigLoop, LoopInfo *LI, Loop *ParentLoop,
+   std::map<Loop*, Loop*>  &ClonedLoopMap) {
+  if (OrigLoop->empty()) return;
+
+  for (auto CurrLoop :  OrigLoop->getSubLoops()) {
+    Loop *NewLoop = new Loop();
+    ParentLoop->addChildLoop(NewLoop);
+    ClonedLoopMap[CurrLoop] = NewLoop;
+
+    // Recursively add the new loops.
+    createNewLoops(CurrLoop, LI, NewLoop, ClonedLoopMap);
+  }
+}
+
 /// \brief Clones a loop \p OrigLoop.  Returns the loop and the blocks in \p
 /// Blocks.
 ///
 /// 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();
 
   Loop *NewLoop = new Loop();
   if (ParentLoop)
     ParentLoop->addChildLoop(NewLoop);
   else
     LI->addTopLevelLoop(NewLoop);
 
+  // Map each old Loop with new one.
+  std::map<Loop*, Loop*> ClonedLoopMap;
+  // Add the top level loop provided for cloning.
+  ClonedLoopMap[OrigLoop] = NewLoop;
+
+  // Recursively clone the loop structure.
+  createNewLoops(OrigLoop, LI, NewLoop, ClonedLoopMap);
+
   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()) {
     BasicBlock *NewBB = CloneBasicBlock(BB, VMap, NameSuffix, F);
     VMap[BB] = NewBB;
 
+    // Get the innermost loop for the BB.
+    Loop* L = LI->getLoopFor(BB);
+    // Get the corresponding cloned loop.
+    Loop* NewClonedLoop = ClonedLoopMap[L];
+    assert(NewClonedLoop && "Could not find the corresponding cloned loop");
     // Update LoopInfo.
-    NewLoop->addBasicBlockToLoop(NewBB, *LI);
+    NewClonedLoop->addBasicBlockToLoop(NewBB, *LI);
 
     // Add DominatorTree node. After seeing all blocks, update to correct IDom.
     DT->addNewBlock(NewBB, NewPH);
 
     Blocks.push_back(NewBB);
   }
 
   for (BasicBlock *BB : OrigLoop->getBlocks()) {