[SimpleLoopUnswitch] After unswitch delete dead blocks in parent loops

[fedor.sergeev]

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.

[fedor.sergeev]

Original unreduced testcase I was using when filing PR38737 also passes with this change.

[asbirlea]

This LGTM.
FWIW, I think we're having an underlying problem and this and the previous PRs may be related.
In unswitchTrivialSwitch, at the very beginning we're removing switch cases and querying LoopInfo to get the outer loop for each switch successor without updating it. LoopInfo is updated only at the very end in hoistLoopToNewParent. Adding an artificial assert to check if LI should be updated after removing switch cases, triggers on 2/3 tests from D51375 and existing trivial-unswitch tests. Looking into it further.

lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

 }
 
 static void
 deleteDeadBlocksFromLoop(Loop &L,
                          SmallVectorImpl<BasicBlock *> &ExitBlocks,
                          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());
 
   // Filter out the dead blocks from the exit blocks list so that it can be
   // used in the caller.
   llvm::erase_if(ExitBlocks,
                  [&](BasicBlock *BB) { return DeadBlockSet.count(BB); });
for (auto *BB : DeadBlocks) {
     // Check that the dominator tree has already been updated.
     assert(!DT.getNode(BB) && "Should already have cleared domtree!");
     LI.changeLoopFor(BB, nullptr);
     BB->dropAllReferences();
   }
 
   // Actually delete the blocks now that they've been fully unhooked from the
   // IR.
-  for (auto *BB : DeadBlocks)
+  for (auto *BB : DeadBlockSet)
     BB->eraseFromParent();
 }
 
 /// Recompute the set of blocks in a loop after unswitching.
 ///
 /// This walks from the original headers predecessors to rebuild the loop. We
 /// take advantage of the fact that new blocks can't have been added, and so we
 /// filter by the original loop's blocks. This also handles potentially

test/Transforms/SimpleLoopUnswitch/delete-dead-blocks.ll

+; 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
+}