From 3432f06be35b1cdf14345e6ee39ef03d14dec77a Mon Sep 17 00:00:00 2001
From: Serguei Katkov <serguei.katkov@azul.com>
Date: Mon, 15 Jul 2019 09:13:11 +0000
Subject: [PATCH] [Loop Peeling] Fix the bug with IDom setting for exit loops

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.

Reviewers: reames, fhahn
Reviewed By: reames
Subscribers: hiraditya, zzheng, llvm-commits
Differential Revision: https://reviews.llvm.org/D64618


git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@366050 91177308-0d34-0410-b5e6-96231b3b80d8
---
 lib/Transforms/Utils/LoopUnrollPeel.cpp       | 21 ++++++-
 .../LoopUnroll/peel-loop-pgo-deopt-idom.ll    | 55 +++++++++++++++++++
 2 files changed, 73 insertions(+), 3 deletions(-)
 create mode 100644 test/Transforms/LoopUnroll/peel-loop-pgo-deopt-idom.ll

diff --git a/lib/Transforms/Utils/LoopUnrollPeel.cpp b/lib/Transforms/Utils/LoopUnrollPeel.cpp
index 6394d74f316..deb38df4420 100644
--- a/lib/Transforms/Utils/LoopUnrollPeel.cpp
+++ b/lib/Transforms/Utils/LoopUnrollPeel.cpp
@@ -583,6 +583,18 @@ bool llvm::peelLoop(Loop *L, unsigned PeelCount, LoopInfo *LI,
   SmallVector<std::pair<BasicBlock *, BasicBlock *>, 4> ExitEdges;
   L->getExitEdges(ExitEdges);
 
+  DenseMap<BasicBlock *, BasicBlock *> ExitIDom;
+  if (DT) {
+    assert(L->hasDedicatedExits() && "No dedicated exits?");
+    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
@@ -675,9 +687,9 @@ bool llvm::peelLoop(Loop *L, unsigned PeelCount, LoopInfo *LI,
       // 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
@@ -719,6 +731,9 @@ bool llvm::peelLoop(Loop *L, unsigned PeelCount, LoopInfo *LI,
   // 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);
 
diff --git a/test/Transforms/LoopUnroll/peel-loop-pgo-deopt-idom.ll b/test/Transforms/LoopUnroll/peel-loop-pgo-deopt-idom.ll
new file mode 100644
index 00000000000..ab3488c8110
--- /dev/null
+++ b/test/Transforms/LoopUnroll/peel-loop-pgo-deopt-idom.ll
@@ -0,0 +1,55 @@
+; 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}
-- 
2.40.0