[LoopSimplifyCFG] Support blocks that are live, but no longer a part of current loop

[mkazantsev]

Summary

Adds support for situations when block don't become dead after the
folding, but stop being a part of the current loop. In this case, we
should do proper upates to Loop Info to keep it consistent.

[mkazantsev]

Still WIP, some bug with subloops.

[mkazantsev]

Just setting "Plan changes" to not confuse reviewers too much with amount of patches in progress.

[mkazantsev]

Rebased, ready for review since underlying patches are now merged.

[mkazantsev]

I plan to hold it off before current implementation is enabled and works without issues for some time.

[mkazantsev]

Rebased.

lib/Transforms/Scalar/LoopSimplifyCFG.cpp

 
 /// Helper class that can turn branches and switches with constant conditions
 /// into unconditional branches.
 class ConstantTerminatorFoldingImpl {
 private:
   Loop &L;
   LoopInfo &LI;
   DominatorTree &DT;
+  LoopBlocksDFS DFS;
 
   // Whether or not the current loop will still exist after terminator constant
   // folding will be done. In theory, there are two ways how it can happen:
   // 1. Loop's latch(es) become unreachable from loop header;
   // 2. Loop's header becomes unreachable from method entry.
   // In practice, the second situation is impossible because we only modify the
   // current loop and its preheader and do not affect preheader's reachibility
   // from any other block. So this variable set to true means that loop's latch
void dump() const {
     if (!DeleteCurrentLoop)
       PrintOutSet("The following blocks will still be part of the loop:",
                   BlocksInLoopAfterFolding);
   }
 
   /// Fill all information about status of blocks and exits of the current loop
   /// if constant folding of all branches will be done.
   void analyze() {
-    LoopBlocksDFS DFS(&L);
     DFS.perform(&LI);
     assert(DFS.isComplete() && "DFS is expected to be finished");
 
     // Collect live and dead loop blocks and exits.
     SmallPtrSet<BasicBlock *, 8> ExitBlocks;
     LiveLoopBlocks.insert(L.getHeader());
     for (auto I = DFS.beginRPO(), E = DFS.endRPO(); I != E; ++I) {
       BasicBlock *BB = *I;
for (BasicBlock *BB : FoldCandidates) {
 
       for (auto *DeadSucc : DeadSuccessors)
         DTU.deleteEdge(BB, DeadSucc);
 
       ++NumTerminatorsFolded;
     }
   }
 
+  bool fixupCurrentLoop() {
+    // If there is no blocks that are no longer in current loop, nothing to do.
+    if (BlocksInLoopAfterFolding.size() + DeadLoopBlocks.size() ==
+        L.getNumBlocks())
+      return false;
+
+    // Traverse blocks in Postorder to preserve the following invariant: by the
+    // moment when we process a block, all its successors already have correct
+    // loops set.
+    for (auto I = DFS.beginPostorder(), E = DFS.endPostorder(); I != E; ++I) {
+      BasicBlock *BB = *I;
+      // If this block stayed in current loop, skip it.
+      if (BlocksInLoopAfterFolding.count(BB))
+        continue;
+
+      // If the block being processed is a header of a child loop, detatch it.
+      Loop *ChildLoop = nullptr;
+      if (LI.isLoopHeader(BB)) {
+        ChildLoop = LI.getLoopFor(BB);
+        if (ChildLoop->getParentLoop() != &L)
+          ChildLoop = nullptr;
+        else
+          L.removeChildLoop(ChildLoop);
+      }
+
+      // Find the most nested ancestor loop that is still BB's successor. BB
+      // also belongs to this loop.
+      Loop *StillReachable = nullptr;
+      for (BasicBlock *Succ : successors(BB)) {
+        Loop *BBL = LI.getLoopFor(Succ);
+        if (BBL && BBL->contains(L.getHeader()))
+          if (!StillReachable || StillReachable->getLoopDepth() >
+              BBL->getLoopDepth())
+            StillReachable = BBL;
+      }
+
+      assert(StillReachable != &L && "Current loop should not be reachable!");
+      // Reattach the child loop is needed.
+      if (ChildLoop) {
+        if (StillReachable)
+          StillReachable->addChildLoop(ChildLoop);
+        else
+          LI.addTopLevelLoop(ChildLoop);
+      }
+      // Finally, set the correct loop for BB.
+      LI.changeLoopFor(BB, StillReachable);
+      for (Loop *Current = &L; Current != StillReachable;
+           Current = Current->getParentLoop())
+        Current->removeBlockFromLoop(BB);
+    }
+
+    return true;
+  }
+
+  void fixupLCSSA() {
+    formLCSSA(L, DT, &LI, nullptr);
+  }
+
 public:
   ConstantTerminatorFoldingImpl(Loop &L, LoopInfo &LI, DominatorTree &DT)
-      : L(L), LI(LI), DT(DT) {}
+      : L(L), LI(LI), DT(DT), DFS(&L) {}
   bool run() {
     assert(L.getLoopLatch() && "Should be single latch!");
 
     // Collect all available information about status of blocks after constant
     // folding.
     analyze();
 
     LLVM_DEBUG(dbgs() << "In function " << L.getHeader()->getParent()->getName()
if (DeleteCurrentLoop) {
       LLVM_DEBUG(
           dbgs()
           << "Give up constant terminator folding in loop "
           << L.getHeader()->getName()
           << ": we don't currently support deletion of the current loop.\n");
       return false;
     }
 
-    // TODO: Support blocks that are not dead, but also not in loop after the
-    // folding.
-    if (BlocksInLoopAfterFolding.size() + DeadLoopBlocks.size() !=
-        L.getNumBlocks()) {
-      LLVM_DEBUG(
-          dbgs() << "Give up constant terminator folding in loop "
-                 << L.getHeader()->getName()
-                 << ": we don't currently"
-                    " support blocks that are not dead, but will stop "
-                    "being a part of the loop after constant-folding.\n");
-      return false;
-    }
-
     // Dump analysis results.
     LLVM_DEBUG(dump());
 
     LLVM_DEBUG(dbgs() << "Constant-folding " << FoldCandidates.size()
                       << " terminators in loop " << L.getHeader()->getName()
                       << "\n");
 
     // Make the actual transforms.
     handleDeadExits();
     foldTerminators();
+    bool NewExitsAppeared = fixupCurrentLoop();
 
     if (!DeadLoopBlocks.empty()) {
       LLVM_DEBUG(dbgs() << "Deleting " << DeadLoopBlocks.size()
                     << " dead blocks in loop " << L.getHeader()->getName()
                     << "\n");
       deleteDeadLoopBlocks();
     }
 
+    if (NewExitsAppeared)
+      fixupLCSSA();
+
 #ifndef NDEBUG
     // Make sure that we have preserved all data structures after the transform.
     DT.verify();
     assert(DT.isReachableFromEntry(L.getHeader()));
     LI.verify(DT);
+    assert(L.isRecursivelyLCSSAForm(DT, LI));
 #endif
 
     return true;
   }
 };
 
 /// Turn branches and switches with known constant conditions into unconditional
 /// branches.

test/Transforms/LoopSimplifyCFG/constant-fold-branch.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; REQUIRES: asserts
 ; RUN: opt -S -loop-simplifycfg -debug-only=loop-simplifycfg 2>&1 < %s | FileCheck %s
 ; RUN: opt -S -passes='require<domtree>,loop(simplify-cfg)' -debug-only=loop-simplifycfg 2>&1 < %s | FileCheck %s
 ; RUN: opt -S -loop-simplifycfg -enable-mssa-loop-dependency=true -verify-memoryssa -debug-only=loop-simplifycfg 2>&1 < %s | FileCheck %s
 
 target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128-ni:1"
 
-; CHECK-LABEL: In function dead_backedge_test_branch_loop: Give up constant terminator folding in loop header: we don't currently support blocks that are not dead, but will stop being a part of the loop after constant-folding.
-; CHECK-LABEL: In function dead_backedge_test_switch_loop: Give up constant terminator folding in loop header: we don't currently support blocks that are not dead, but will stop being a part of the loop after constant-folding.
+; CHECK-LABEL: In function dead_backedge_test_branch_loop: Constant terminator folding for loop Loop at depth 1 containing: %header<header>,%dead_backedge<exiting>,%header.backedge<latch>
+; CHECK:         Constant-folding 1 terminators in loop header
+; CHECK-LABEL: In function dead_backedge_test_switch_loop: Constant terminator folding for loop Loop at depth 1 containing: %header<header>,%dead_backedge<exiting>,%header.backedge<latch>
+; CHECK:         Constant-folding 1 terminators in loop header
 ; CHECK-LABEL: In function dead_block_test_branch_loop: Constant terminator folding for loop Loop at depth 1 containing: %header<header>,%dead,%backedge<latch><exiting>
 ; CHECK:         Deleting 1 dead blocks in loop header
 ; CHECK-LABEL: In function dead_block_test_switch_loop: Constant terminator folding for loop Loop at depth 1 containing: %header<header>,%dead,%backedge<latch><exiting>
 ; CHECK:         Deleting 1 dead blocks in loop header
 ; CHECK-LABEL: In function dead_block_propogate_test_branch_loop: Constant terminator folding for loop Loop at depth 1 containing: %header<header>,%dead,%dummy,%backedge<latch><exiting>
 ; CHECK:         Deleting 2 dead blocks in loop header
 ; CHECK-LABEL: In function dead_block_propogate_test_switch_loop: Constant terminator folding for loop Loop at depth 1 containing: %header<header>,%dead,%dummy,%backedge<latch><exiting>
 ; CHECK:         Deleting 2 dead blocks in loop header

 ; CHECK-LABEL: In function exit_switch_from_inner_to_grandparent: Constant terminator folding for loop Loop at depth 3 containing: %loop_3<header><exiting>,%loop_3_backedge<latch><exiting>
 
 ; Make sure that we can eliminate a provably dead backedge.
 define i32 @dead_backedge_test_branch_loop(i32 %end) {
 ; CHECK-LABEL: @dead_backedge_test_branch_loop(
 ; CHECK-NEXT:  preheader:
 ; CHECK-NEXT:    br label [[HEADER:%.*]]
 ; CHECK:       header:
-; CHECK-NEXT:    [[I:%.*]] = phi i32 [ 0, [[PREHEADER:%.*]] ], [ [[I_BE:%.*]], [[HEADER_BACKEDGE:%.*]] ]
-; CHECK-NEXT:    [[I_1:%.*]] = add i32 [[I]], 1
+; CHECK-NEXT:    [[I:%.*]] = phi i32 [ 0, [[PREHEADER:%.*]] ], [ [[I_1:%.*]], [[HEADER_BACKEDGE:%.*]] ]
+; CHECK-NEXT:    [[I_1]] = add i32 [[I]], 1
 ; CHECK-NEXT:    [[CMP1:%.*]] = icmp slt i32 [[I_1]], 100
 ; CHECK-NEXT:    br i1 [[CMP1]], label [[HEADER_BACKEDGE]], label [[DEAD_BACKEDGE:%.*]]
 ; CHECK:       header.backedge:
-; CHECK-NEXT:    [[I_BE]] = phi i32 [ [[I_1]], [[HEADER]] ], [ [[I_2:%.*]], [[DEAD_BACKEDGE]] ]
 ; CHECK-NEXT:    br label [[HEADER]]
 ; CHECK:       dead_backedge:
-; CHECK-NEXT:    [[I_2]] = add i32 [[I_1]], 10
-; CHECK-NEXT:    br i1 false, label [[HEADER_BACKEDGE]], label [[EXIT:%.*]]
+; CHECK-NEXT:    [[I_1_LCSSA:%.*]] = phi i32 [ [[I_1]], [[HEADER]] ]
+; CHECK-NEXT:    [[I_2:%.*]] = add i32 [[I_1_LCSSA]], 10
+; CHECK-NEXT:    br label [[EXIT:%.*]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    [[I_2_LCSSA:%.*]] = phi i32 [ [[I_2]], [[DEAD_BACKEDGE]] ]
 ; CHECK-NEXT:    ret i32 [[I_2_LCSSA]]
 ;
 preheader:
   br label %header
 
 header:

 }
 
 ; Make sure that we can eliminate a provably dead backedge with switch.
 define i32 @dead_backedge_test_switch_loop(i32 %end) {
 ; CHECK-LABEL: @dead_backedge_test_switch_loop(
 ; CHECK-NEXT:  preheader:
 ; CHECK-NEXT:    br label [[HEADER:%.*]]
 ; CHECK:       header:
-; CHECK-NEXT:    [[I:%.*]] = phi i32 [ 0, [[PREHEADER:%.*]] ], [ [[I_BE:%.*]], [[HEADER_BACKEDGE:%.*]] ]
-; CHECK-NEXT:    [[I_1:%.*]] = add i32 [[I]], 1
+; CHECK-NEXT:    [[I:%.*]] = phi i32 [ 0, [[PREHEADER:%.*]] ], [ [[I_1:%.*]], [[HEADER_BACKEDGE:%.*]] ]
+; CHECK-NEXT:    [[I_1]] = add i32 [[I]], 1
 ; CHECK-NEXT:    [[CMP1:%.*]] = icmp slt i32 [[I_1]], 100
 ; CHECK-NEXT:    br i1 [[CMP1]], label [[HEADER_BACKEDGE]], label [[DEAD_BACKEDGE:%.*]]
 ; CHECK:       header.backedge:
-; CHECK-NEXT:    [[I_BE]] = phi i32 [ [[I_1]], [[HEADER]] ], [ [[I_2:%.*]], [[DEAD_BACKEDGE]] ]
 ; CHECK-NEXT:    br label [[HEADER]]
 ; CHECK:       dead_backedge:
-; CHECK-NEXT:    [[I_2]] = add i32 [[I_1]], 10
-; CHECK-NEXT:    switch i32 1, label [[EXIT:%.*]] [
-; CHECK-NEXT:    i32 0, label [[HEADER_BACKEDGE]]
-; CHECK-NEXT:    ]
+; CHECK-NEXT:    [[I_1_LCSSA:%.*]] = phi i32 [ [[I_1]], [[HEADER]] ]
+; CHECK-NEXT:    [[I_2:%.*]] = add i32 [[I_1_LCSSA]], 10
+; CHECK-NEXT:    br label [[EXIT:%.*]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    [[I_2_LCSSA:%.*]] = phi i32 [ [[I_2]], [[DEAD_BACKEDGE]] ]
 ; CHECK-NEXT:    ret i32 [[I_2_LCSSA]]
 ;
 preheader:
   br label %header
 
 header: