[SimpleLoopUnswitch] After unswitch delete dead blocks in parent loops

Summary:
Assert from PR38737 happens on the dead block inside the parent loop
after unswitching nontrivial switch in the inner loop.

deleteDeadBlocksFromLoop now takes extra care to detect/remove dead
blocks in all the parent loops in addition to the blocks from original
loop being unswitched.

Reviewers: asbirlea, chandlerc

Reviewed By: asbirlea

Subscribers: llvm-commits

Differential Revision: https://reviews.llvm.org/D51415

llvm-svn: 340955

2 files changed, 55 insertions, 2 deletions

diff --git a/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp b/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
index 8b67c470a35..3cc2c651f6c 100644
--- a/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
+++ b/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
@@ -1380,13 +1380,21 @@ deleteDeadBlocksFromLoop(Loop &L,
                          DominatorTree &DT, LoopInfo &LI) {
   // Find all the dead blocks, and remove them from their successors.
   SmallVector<BasicBlock *, 16> DeadBlocks;
-  for (BasicBlock *BB : llvm::concat<BasicBlock *const>(L.blocks(), ExitBlocks))
+  for (BasicBlock *BB : ExitBlocks)
     if (!DT.isReachableFromEntry(BB)) {
       for (BasicBlock *SuccBB : successors(BB))
         SuccBB->removePredecessor(BB);
       DeadBlocks.push_back(BB);
     }
 
+  for (Loop *ParentL = &L; ParentL; ParentL = ParentL->getParentLoop())
+    for (BasicBlock *BB : ParentL->blocks())
+      if (!DT.isReachableFromEntry(BB)) {
+        for (BasicBlock *SuccBB : successors(BB))
+          SuccBB->removePredecessor(BB);
+        DeadBlocks.push_back(BB);
+      }
+
   SmallPtrSet<BasicBlock *, 16> DeadBlockSet(DeadBlocks.begin(),
                                              DeadBlocks.end());
 
@@ -1431,7 +1439,7 @@ deleteDeadBlocksFromLoop(Loop &L,
 
   // Actually delete the blocks now that they've been fully unhooked from the
   // IR.
-  for (auto *BB : DeadBlocks)
+  for (auto *BB : DeadBlockSet)
     BB->eraseFromParent();
 }
 
diff --git a/llvm/test/Transforms/SimpleLoopUnswitch/delete-dead-blocks.ll b/llvm/test/Transforms/SimpleLoopUnswitch/delete-dead-blocks.ll
new file mode 100644
index 00000000000..0be1c8549ed
--- /dev/null
+++ b/llvm/test/Transforms/SimpleLoopUnswitch/delete-dead-blocks.ll
@@ -0,0 +1,45 @@
+; RUN: opt < %s -simple-loop-unswitch -enable-nontrivial-unswitch -S 2>&1 | FileCheck %s
+; RUN: opt < %s -passes=unswitch -enable-nontrivial-unswitch -S 2>&1 | FileCheck %s
+;
+; Checking that (dead) blocks from inner loop are deleted after unswitch.
+;
+declare void @foo()
+
+; CHECK-LABEL: @Test
+define void @Test(i32) {
+entry:
+  br label %outer
+outer:
+  %oi = phi i32 [ 0, %entry ], [ %oinc, %outer_continue]
+  br label %inner
+inner:
+  %ii = phi i32 [ 0, %outer ], [ %iinc, %continue]
+  call void @foo() 
+  switch i32 %0, label %get_out2 [
+    i32 0, label %continue
+    i32 1, label %case1
+    i32 2, label %get_out
+  ]
+;
+; since we unswitch on the above switch, %case1 and %continue blocks
+; become dead in the original loop
+;
+; CHECK-NOT: case1:
+case1:
+  br label %continue
+; CHECK-NOT: {{^}}continue:
+continue:
+  %iinc = add i32 %ii, 1
+  %icmp = icmp eq i32 %ii, 100
+  br i1 %icmp, label %inner, label %outer_continue
+
+outer_continue:
+  %oinc = add i32 %oi, 1
+  %ocmp = icmp eq i32 %oi, 100
+  br i1 %ocmp, label %outer, label %get_out
+
+get_out:
+  ret void
+get_out2:
+  unreachable
+}