[Loop Peeling] Fix the bug with IDom setting for exit loops

[skatkov]

Summary

It is possible that loop exit has two predecessors in a loop body.
In this case after the peeling the iDom of the exit should be a clone of
iDom of original exit but no a clone of a block coming to this exit.

[reames]

LGTM. An alternate implementation would be to require that each exit has one predecessor in the legality for peeling check. Either would be fine, and if you want to do the alternate, you can land that without review..

llvm/lib/Transforms/Utils/LoopUnrollPeel.cpp

588

Please add an assert that L has dedicated exits, because that's the precondition for your dominance fact.

[skatkov]

llvm/lib/Transforms/Utils/LoopUnrollPeel.cpp

588

In the beginning of this method we have an assert that we canPeel which returns true only if loop is in simplify form what includes the check that all exits are dedicated.

It is to far so I'm ok to add an assert here as well.

llvm/lib/Transforms/Utils/LoopUnrollPeel.cpp

bool llvm::peelLoop(Loop *L, unsigned PeelCount, LoopInfo *LI,
   LoopBlocks.perform(LI);
 
   BasicBlock *Header = L->getHeader();
   BasicBlock *PreHeader = L->getLoopPreheader();
   BasicBlock *Latch = L->getLoopLatch();
   SmallVector<std::pair<BasicBlock *, BasicBlock *>, 4> ExitEdges;
   L->getExitEdges(ExitEdges);
 
+  DenseMap<BasicBlock *, BasicBlock *> ExitIDom;
+  if (DT)
+    for (auto Edge : ExitEdges) {
+      if (ExitIDom.count(Edge.second))
+        continue;
+      BasicBlock *BB = DT->getNode(Edge.second)->getIDom()->getBlock();
+      assert(L->contains(BB) && "IDom is not in a loop");
+      ExitIDom[Edge.second] = BB;
+    }
+
   Function *F = Header->getParent();
 
   // Set up all the necessary basic blocks. It is convenient to split the
   // preheader into 3 parts - two blocks to anchor the peeled copy of the loop
   // body, and a new preheader for the "real" loop.
 
   // Peeling the first iteration transforms.
   //
for (unsigned Iter = 0; Iter < PeelCount; ++Iter) {
     // previous one.
     remapInstructionsInBlocks(NewBlocks, VMap);
 
     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)
-        for (auto Edge : ExitEdges)
-          DT->changeImmediateDominator(Edge.second,
-                                       cast<BasicBlock>(LVMap[Edge.first]));
+        for (auto Exit : ExitIDom)
+          DT->changeImmediateDominator(Exit.first,
+                                       cast<BasicBlock>(LVMap[Exit.second]));
 #ifdef EXPENSIVE_CHECKS
       assert(DT->verify(DominatorTree::VerificationLevel::Fast));
 #endif
     }
 
     auto *LatchBRCopy = cast<BranchInst>(VMap[LatchBR]);
     updateBranchWeights(InsertBot, LatchBRCopy, Iter,
                         PeelCount, ExitWeight);
#endif
   fixupBranchWeights(Header, LatchBR, ExitWeight, CurHeaderWeight);
 
   if (Loop *ParentLoop = L->getParentLoop())
     L = ParentLoop;
 
   // We modified the loop, update SE.
   SE->forgetTopmostLoop(L);
 
+  // Finally DomtTree must be correct.
+  assert(DT->verify(DominatorTree::VerificationLevel::Fast));
+
   // FIXME: Incrementally update loop-simplify
   simplifyLoop(L, DT, LI, SE, AC, nullptr, PreserveLCSSA);
 
   NumPeeled++;
 
   return true;
 }

llvm/test/Transforms/LoopUnroll/peel-loop-pgo-deopt-idom.ll

+; REQUIRES: asserts
+; RUN: opt < %s -S -debug-only=loop-unroll -loop-unroll -unroll-runtime -unroll-peel-multi-deopt-exit 2>&1 | FileCheck %s
+; RUN: opt < %s -S -debug-only=loop-unroll -unroll-peel-multi-deopt-exit -passes='require<profile-summary>,function(require<opt-remark-emit>,unroll)' 2>&1 | FileCheck %s
+
+; Regression test for setting the correct idom for exit blocks.
+
+; CHECK: Loop Unroll: F[basic]
+; CHECK: PEELING loop %for.body with iteration count 1!
+
+define i32 @basic(i32* %p, i32 %k, i1 %c1, i1 %c2) #0 !prof !3 {
+entry:
+  %cmp3 = icmp slt i32 0, %k
+  br i1 %cmp3, label %for.body.lr.ph, label %for.end
+
+for.body.lr.ph:                                   ; preds = %entry
+  br label %for.body
+
+for.body:                                         ; preds = %for.body.lr.ph, %for.body
+  %i.05 = phi i32 [ 0, %for.body.lr.ph ], [ %inc, %latch ]
+  %p.addr.04 = phi i32* [ %p, %for.body.lr.ph ], [ %incdec.ptr, %latch ]
+  %incdec.ptr = getelementptr inbounds i32, i32* %p.addr.04, i32 1
+  store i32 %i.05, i32* %p.addr.04, align 4
+  %inc = add nsw i32 %i.05, 1
+  %cmp = icmp slt i32 %inc, %k
+  br i1 %c1, label %continue, label %to_side_exit
+
+continue:
+  br i1 %c2, label %latch, label %side_exit, !prof !2
+
+latch:
+  br i1 %cmp, label %for.body, label %for.cond.for.end_crit_edge, !prof !1
+
+for.cond.for.end_crit_edge:                       ; preds = %for.body
+  br label %for.end
+
+to_side_exit:
+  br i1 %c2, label %continue, label %side_exit, !prof !2
+
+
+for.end:                                          ; preds = %for.cond.for.end_crit_edge, %entry
+  %res = phi i32 [ 0, %entry ], [ %inc, %for.cond.for.end_crit_edge ]
+  ret i32 %res
+
+side_exit:
+  %rval = call i32(...) @llvm.experimental.deoptimize.i32() [ "deopt"(i32 %inc) ]
+  ret i32 %rval
+}
+
+declare i32 @llvm.experimental.deoptimize.i32(...)
+
+attributes #0 = { nounwind }
+
+!1 = !{!"branch_weights", i32 1, i32 1}
+!2 = !{!"branch_weights", i32 1, i32 0}
+!3 = !{!"function_entry_count", i64 1}