From 243a05e2edafe99b82a31c9e93139be2b19afdd6 Mon Sep 17 00:00:00 2001
From: Michael Zolotukhin <mzolotukhin@apple.com>
Date: Wed, 6 Apr 2016 21:47:12 +0000
Subject: [PATCH] [LoopUnroll] Fix the way we update DT after complete
 unrolling.

Updating dominators for exit-blocks of the unrolled loops is not enough,
as shown in PR27157. The proper way is to update dominators for all
dominance-children of original loop blocks.

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@265605 91177308-0d34-0410-b5e6-96231b3b80d8
---
 lib/Transforms/Utils/LoopUnroll.cpp   | 25 ++++++++-----
 test/Transforms/LoopUnroll/pr27157.ll | 53 +++++++++++++++++++++++++++
 2 files changed, 68 insertions(+), 10 deletions(-)
 create mode 100644 test/Transforms/LoopUnroll/pr27157.ll

diff --git a/lib/Transforms/Utils/LoopUnroll.cpp b/lib/Transforms/Utils/LoopUnroll.cpp
index 20822df819d..e9840424354 100644
--- a/lib/Transforms/Utils/LoopUnroll.cpp
+++ b/lib/Transforms/Utils/LoopUnroll.cpp
@@ -262,6 +262,7 @@ bool llvm::UnrollLoop(Loop *L, unsigned Count, unsigned TripCount,
   bool CompletelyUnroll = Count == TripCount;
   SmallVector<BasicBlock *, 4> ExitBlocks;
   L->getExitBlocks(ExitBlocks);
+  std::vector<BasicBlock*> OriginalLoopBlocks = L->getBlocks();
 
   // Go through all exits of L and see if there are any phi-nodes there. We just
   // conservatively assume that they're inserted to preserve LCSSA form, which
@@ -551,20 +552,24 @@ bool llvm::UnrollLoop(Loop *L, unsigned Count, unsigned TripCount,
       Term->eraseFromParent();
     }
   }
-  // Update dominators of loop exit blocks.
-  // Immediate dominator of an exit block might change, because we add more
+  // Update dominators of blocks we might reach through exits.
+  // Immediate dominator of such block might change, because we add more
   // routes which can lead to the exit: we can now reach it from the copied
-  // iterations too. Thus, the new idom of the exit block will be the nearest
+  // iterations too. Thus, the new idom of the block will be the nearest
   // common dominator of the previous idom and common dominator of all copies of
-  // the exiting block. This is equivalent to the nearest common dominator of
+  // the previous idom. This is equivalent to the nearest common dominator of
   // the previous idom and the first latch, which dominates all copies of the
-  // exiting block.
+  // previous idom.
   if (DT && Count > 1) {
-    for (auto Exit : ExitBlocks) {
-      BasicBlock *PrevIDom = DT->getNode(Exit)->getIDom()->getBlock();
-      BasicBlock *NewIDom =
-          DT->findNearestCommonDominator(PrevIDom, Latches[0]);
-      DT->changeImmediateDominator(Exit, NewIDom);
+    for (auto *BB : OriginalLoopBlocks) {
+      auto *BBDomNode = DT->getNode(BB);
+      for (auto *ChildDomNode : BBDomNode->getChildren()) {
+        auto *ChildBB = ChildDomNode->getBlock();
+        if (L->contains(ChildBB))
+          continue;
+        BasicBlock *NewIDom = DT->findNearestCommonDominator(BB, Latches[0]);
+        DT->changeImmediateDominator(ChildBB, NewIDom);
+      }
     }
   }
 
diff --git a/test/Transforms/LoopUnroll/pr27157.ll b/test/Transforms/LoopUnroll/pr27157.ll
new file mode 100644
index 00000000000..917bcf1f944
--- /dev/null
+++ b/test/Transforms/LoopUnroll/pr27157.ll
@@ -0,0 +1,53 @@
+; RUN: opt -loop-unroll -debug-only=loop-unroll -disable-output < %s
+; REQUIRES: asserts
+; Compile this test with debug flag on to verify domtree right after loop unrolling.
+target datalayout = "E-m:e-i1:8:16-i8:8:16-i64:64-f128:64-v128:64-a:8:16-n32:64"
+
+; PR27157
+define void @foo() {
+entry:
+  br label %loop_header
+loop_header:
+  %iv = phi i64 [ 0, %entry ], [ %iv_next, %loop_latch ]
+  br i1 undef, label %loop_latch, label %loop_exiting_bb1
+loop_exiting_bb1:
+  br i1 false, label %loop_exiting_bb2, label %exit1.loopexit
+loop_exiting_bb2:
+  br i1 false, label %loop_latch, label %bb
+bb:
+  br label %exit1
+loop_latch:
+  %iv_next = add nuw nsw i64 %iv, 1
+  %cmp = icmp ne i64 %iv_next, 2
+  br i1 %cmp, label %loop_header, label %exit2
+exit1.loopexit:
+  br label %exit1
+exit1:
+  ret void
+exit2:
+  ret void
+}
+
+define void @foo2() {
+entry:
+  br label %loop.header
+loop.header:
+  %iv = phi i32 [ 0, %entry ], [ %iv.inc, %latch ]
+  %iv.inc = add i32 %iv, 1
+  br i1 undef, label %diamond, label %latch
+diamond:
+  br i1 undef, label %left, label %right
+left:
+  br i1 undef, label %exit, label %merge
+right:
+  br i1 undef, label %exit, label %merge
+merge:
+  br label %latch
+latch:
+  %end.cond = icmp eq i32 %iv, 1
+  br i1 %end.cond, label %exit1, label %loop.header
+exit:
+  ret void
+exit1:
+  ret void
+}