Update domtree incrementally in loop peeling.

[sepavloff]

Summary

With this change dominator tree remains in sync after each step of loop
peeling.

[mkuper]

LGTM with a couple of nits.

lib/Transforms/Utils/LoopUnrollPeel.cpp

191 ↗

(On Diff #85390)

Can you move this to immediately below line 207? I think it makes more sense to have this after the comment lines 203-206.

384 ↗

(On Diff #85390)

Can you decouple this from D29028?
I'd suggest just doing this under UnrollVerifyDomtree for now.

[mkuper]

lib/Transforms/Utils/LoopUnrollPeel.cpp

384 ↗

(On Diff #85390)

I see Michael Z approved D29028, so feel free to ignore this.

llvm/trunk/lib/Transforms/Utils/LoopUnrollPeel.cpp

 /// \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<BasicBlock *> &NewBlocks,
                             LoopBlocksDFS &LoopBlocks, ValueToValueMapTy &VMap,
-                            ValueToValueMapTy &LVMap, LoopInfo *LI) {
+                            ValueToValueMapTy &LVMap, DominatorTree *DT,
+                            LoopInfo *LI) {
 
   BasicBlock *Header = L->getHeader();
   BasicBlock *Latch = L->getLoopLatch();
   BasicBlock *PreHeader = L->getLoopPreheader();
 
   Function *F = Header->getParent();
   LoopBlocksDFS::RPOIterator BlockBegin = LoopBlocks.beginRPO();
   LoopBlocksDFS::RPOIterator BlockEnd = LoopBlocks.endRPO();
   Loop *ParentLoop = L->getParentLoop();
 
   // For each block in the original loop, create a new copy,
   // and update the value map with the newly created values.
   for (LoopBlocksDFS::RPOIterator BB = BlockBegin; BB != BlockEnd; ++BB) {
     BasicBlock *NewBB = CloneBasicBlock(*BB, VMap, ".peel", F);
     NewBlocks.push_back(NewBB);
 
     if (ParentLoop)
       ParentLoop->addBasicBlockToLoop(NewBB, *LI);
 
     VMap[*BB] = NewBB;
+
+    // If dominator tree is available, insert nodes to represent cloned blocks.
+    if (DT) {
+      if (Header == *BB)
+        DT->addNewBlock(NewBB, InsertTop);
+      else {
+        DomTreeNode *IDom = DT->getNode(*BB)->getIDom();
+        // VMap must contain entry for IDom, as the iteration order is RPO.
+        DT->addNewBlock(NewBB, cast<BasicBlock>(VMap[IDom->getBlock()]));
+      }
+    }
   }
 
   // Hook-up the control flow for the newly inserted blocks.
   // The new header is hooked up directly to the "top", which is either
   // the original loop preheader (for the first iteration) or the previous
   // iteration's exiting block (for every other iteration)
   InsertTop->getTerminator()->setSuccessor(0, cast<BasicBlock>(VMap[Header]));
 
   // Similarly, for the latch:
   // The original exiting edge is still hooked up to the loop exit.
   // The backedge now goes to the "bottom", which is either the loop's real
   // header (for the last peeled iteration) or the copied header of the next
   // iteration (for every other iteration)
-  BranchInst *LatchBR =
-      cast<BranchInst>(cast<BasicBlock>(VMap[Latch])->getTerminator());
+  BasicBlock *NewLatch = cast<BasicBlock>(VMap[Latch]);
+  BranchInst *LatchBR = cast<BranchInst>(NewLatch->getTerminator());
   unsigned HeaderIdx = (LatchBR->getSuccessor(0) == Header ? 0 : 1);
   LatchBR->setSuccessor(HeaderIdx, InsertBot);
   LatchBR->setSuccessor(1 - HeaderIdx, Exit);
+  if (DT)
+    DT->changeImmediateDominator(InsertBot, NewLatch);
 
   // The new copy of the loop body starts with a bunch of PHI nodes
   // that pick an incoming value from either the preheader, or the previous
   // loop iteration. Since this copy is no longer part of the loop, we
   // resolve this statically:
   // For the first iteration, we use the value from the preheader directly.
   // For any other iteration, we replace the phi with the value generated by
   // the immediately preceding clone of the loop body (which represents
for (unsigned Iter = 0; Iter < PeelCount; ++Iter) {
     // FIXME: due to the way the distribution is constructed, we need a
     // guard here to make sure we don't end up with non-positive weights.
     if (ExitWeight < CurHeaderWeight)
       CurHeaderWeight -= ExitWeight;
     else
       CurHeaderWeight = 1;
 
     cloneLoopBlocks(L, Iter, InsertTop, InsertBot, Exit,
-                    NewBlocks, LoopBlocks, VMap, LVMap, LI);
+                    NewBlocks, LoopBlocks, VMap, LVMap, DT, LI);
+    if (DT) {
+      // Latches of the cloned loops dominate over the loop exit, so idom of the
+      // latter is the first cloned loop body, as original PreHeader dominates
+      // the original loop body.
+      if (Iter == 0)
+        DT->changeImmediateDominator(Exit, cast<BasicBlock>(LVMap[Latch]));
+#ifndef NDEBUG
+      if (VerifyDomInfo)
+        DT->verifyDomTree();
+#endif
+    }
+
     updateBranchWeights(InsertBot, cast<BranchInst>(VMap[LatchBR]), Iter,
                         PeelCount, ExitWeight);
 
     InsertTop = InsertBot;
     InsertBot = SplitBlock(InsertBot, InsertBot->getTerminator(), DT, LI);
     InsertBot->setName(Header->getName() + ".peel.next");
 
     F->getBasicBlockList().splice(InsertTop->getIterator(),
else
       BackEdgeWeight = 1;
     MDBuilder MDB(LatchBR->getContext());
     MDNode *WeightNode =
         HeaderIdx ? MDB.createBranchWeights(ExitWeight, BackEdgeWeight)
                   : MDB.createBranchWeights(BackEdgeWeight, ExitWeight);
     LatchBR->setMetadata(LLVMContext::MD_prof, WeightNode);
   }
 
-  // FIXME: Incrementally update domtree.
-  DT->recalculate(*L->getHeader()->getParent());
-
   // If the loop is nested, we changed the parent loop, update SE.
   if (Loop *ParentLoop = L->getParentLoop()) {
     SE->forgetLoop(ParentLoop);
 
     // FIXME: Incrementally update loop-simplify
     simplifyLoop(ParentLoop, DT, LI, SE, AC, PreserveLCSSA);
   } else {
     // FIXME: Incrementally update loop-simplify
     simplifyLoop(L, DT, LI, SE, AC, PreserveLCSSA);
   }
 
   NumPeeled++;
 
   return true;
 }

llvm/trunk/test/Transforms/LoopUnroll/peel-loop.ll

-; RUN: opt < %s -S -loop-unroll -unroll-force-peel-count=3 -simplifycfg -instcombine | FileCheck %s
+; RUN: opt < %s -S -loop-unroll -unroll-force-peel-count=3 -verify-dom-info -simplifycfg -instcombine | FileCheck %s
 
 ; Basic loop peeling - check that we can peel-off the first 3 loop iterations
 ; when explicitly requested.
 ; CHECK-LABEL: @basic
 ; CHECK: %[[CMP0:.*]] = icmp sgt i32 %k, 0
 ; CHECK: br i1 %[[CMP0]], label %[[NEXT0:.*]], label %for.end
 ; CHECK: [[NEXT0]]:
 ; CHECK: store i32 0, i32* %p, align 4