[LoopSimplifyCFG] Teach LoopSimplifyCFG to constant-fold branches and switches

[mkazantsev]

Summary

This patch introduces infrastructure and the simplest case for constant-folding
of branch and switch instructions within loop into unconditional branches.
It is useful as a cleanup for such passes as loop unswitching that sometimes
produce such branches.

Only the simplest case supported in this patch: after the folding, no block
should become dead or stop being part of the loop. Support for more
sophisticated cases will go separately in follow-up patches.

[mkazantsev]

Fix typo in dump()

[reames]

Drive by comments.

It looks like you forgot to include your test cases?

lib/Transforms/Scalar/LoopSimplifyCFG.cpp

182 ↗

(On Diff #172301)

Why bother to keep around both sets when given an exit block you can answer liveness questions with only one copy? (Same question on other parallel dead/live structures BTW)

196 ↗

(On Diff #172301)

I think you need a better name for this variable.

198 ↗

(On Diff #172301)

Not clear to me why you need this bit of code in this patch.

248 ↗

(On Diff #172301)

Where's the code for LI updating?

[dmgreen]

lib/Transforms/Scalar/LoopSimplifyCFG.cpp

64 ↗

(On Diff #172301)

Should be CI?

[anna]

lib/Transforms/Scalar/LoopSimplifyCFG.cpp

248 ↗

(On Diff #172301)

From the summary of this patch, it looks like we don't need LI update in this patch:

Only the simplest case supported in this patch: after the folding, no block
should become dead or stop being part of the loop. Support for more
sophisticated cases will go separately in follow-up patches.

If that's the case, I'd really like an assert to make sure that the LI is still correct after this transform is completed.

I'll review the code to see if what's stated here is actually true (for example, we have a number of prechecks before performing the actual transform). See run function below.

[anna]

Couple of comments inline. Pls add test cases.

lib/Transforms/Scalar/LoopSimplifyCFG.cpp

64 ↗

(On Diff #172301)

yup, looks like a typo.

136 ↗

(On Diff #172301)

if constant folding will be done. At this point CF is not done.

159 ↗

(On Diff #172301)

Pls add test cases where:

1. Block's only incoming edge if constant folded will make the block dead - this transform shouldn't kick in.
2. Block has multiple incoming edges - even if all except one is constant folded, the block is still live.

198 ↗

(On Diff #172301)

If LoopIsLive is false, we should just return at this point now, right? For this patch, I believe this chunk of code is not needed. Am I missing something?

261 ↗

(On Diff #172301)

Pls add test cases for each of the scenarios described below. The constantfold xform should kick in with this patch.

288 ↗

(On Diff #172301)

We need a bunch of asserts here:

1. assert that L.getHeader() is reachableFromEntry.
2. assert that LI is correct (mentioned in earlier comment).

[anna]

based on above comments.

[mkazantsev]

As for the test: I have commited all related tests I am planning to support in this pass in the file test/Transforms/LoopSimplifyCFG/constant-fold-branch.ll. It has all situations with dead blocks, live blocks, dead loops etc, please take a look and tell me if some important case is missing. My plan is to support folding in all these cases incrementally.

As for comments in code, I'll address them shortly.

[mkazantsev]

lib/Transforms/Scalar/LoopSimplifyCFG.cpp

64 ↗

(On Diff #172301)

Good catch!

159 ↗

(On Diff #172301)

Already in test/Transforms/LoopSimplifyCFG/constant-fold-branch.ll, it has a variety of such cases with inner/exit blocks and loops becoming dead.

182 ↗

(On Diff #172301)

As we support more cases in follow-up patches, sets of live and dead blocks have different purposes and are all used. See:
https://reviews.llvm.org/D54023
https://reviews.llvm.org/D54025

This patch just collects all required information and enables the simplest case (that doesn't need all of it). However, as we add more logic, these sets are used more actively.

All this logic is only split into two patches to make it easier to review.

198 ↗

(On Diff #172301)

You are right. But this patch is just collecting all required data for follow-up patches https://reviews.llvm.org/D54023 https://reviews.llvm.org/D54025

My idea was to collect *all* required data in this patch and then add transforms that use it one by one, without adding more analysis. If you think it would be easier to review if the data collection is also added incrementally, I can rework the change.

248 ↗

(On Diff #172301)

No LI update is needed because no block gets deleted or leaves the loop. This patch only handles the situation when it is so. You can take a look on follow-up dependent patches, they have LI updates.

I will add the assert that LI is still correct after the transform.

261 ↗

(On Diff #172301)

All in test/Transforms/LoopSimplifyCFG/constant-fold-branch.ll. Currently they are all negative tests. Patches to this pass enable groups of them.

288 ↗

(On Diff #172301)

AFAIK reachibility of header is implied by correctness of LI. Will do.

[skatkov]

lib/Transforms/Scalar/LoopSimplifyCFG.cpp

48 ↗

(On Diff #172301)

Please update the comment that you support only branch and switch terminators and for others you just return nullptr.

49 ↗

(On Diff #172301)

In terms of compile time improvement I see that in analysis you use this function a lot.
So it makes sense to cache the result. So you need a set of live pairs <Src, Dest>.

70 ↗

(On Diff #172301)

May be just return this LiveSucc here and after the loop just return the default dest?

103 ↗

(On Diff #172301)

According to code how you fill in it, it seems that FoldCandidates contains only blocks from L, not form sub loops...

160 ↗

(On Diff #172301)

With this check you want to process only blocks from the current loop, not in sub loops. What is the reason for that?
Simplification?

196 ↗

(On Diff #172301)

In terms of analysis I think the restriction that there is an only one loop latch is redundant.

279 ↗

(On Diff #172301)

It seems that this check should be the first one.
If there is no to fold so no one of the check above will be true.

345 ↗

(On Diff #172301)

Be careful here in the future, if constantFoldTerminators will eliminate the current loop you should not run merge blocks...

[skatkov]

Also Max, please consider split this patch into two ones. One just NFC patch with analysis and update tests to check the output of dump and then patches which does modifications.

[mkazantsev]

Making a separate NFC which checks analyzes behavior doesn't sound because I am using sets to store blocks status, and printing out the contents of sets is non-deterministic.

[mkazantsev]

lib/Transforms/Scalar/LoopSimplifyCFG.cpp

49 ↗

(On Diff #172301)

This check is free for branches, for switches it's O(num_succ) which is negligible. I think it's better keep it that way than add a cache which can be a source of potential bugs.

160 ↗

(On Diff #172301)

If we modify some condition in the inner loop, we may end up needing to update its loop info, which we don't want to do. Besides, we expect that all branches in inner loops will have been folded by the moment when we will process the outer loop.

196 ↗

(On Diff #172301)

It makes things easier, and we canonicalize most loops to this form, so I think we shouldn't bother about multi-latch case. We can support them later if needed.

345 ↗

(On Diff #172301)

Yes, it will be handled when we support deletion of the current loop (I plan to make it as the last step).

[mkazantsev]

Fixed typo bug, added LI and DT verification, changed FoldingCandidates to vector to keep deterministic order of updates, updated tests with decline messages to make sure that we exercise all possible scenarios in our test file.

[anna]

comments inline. Algo looks fine.

lib/Transforms/Scalar/LoopSimplifyCFG.cpp

346 ↗

(On Diff #172898)

You're not running this to a fixed point, so why not just do Changed = BranchFolder.run()

198 ↗

(On Diff #172301)

pls add this with the patch that deletes the loop, instead of in this patch. I'd like to reduce the complexity of this patch.

288 ↗

(On Diff #172301)

I checked, but LI.verify() doesn't contain that check for header is reachableFromEntry. We recompute a new LI and then verify original LI and the recomputed one is the same.

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

8 ↗

(On Diff #172898)

I understand why you added these, but it's Really a lot of information in the debug output. I think "Give up constant terminator folding in loop header" is enough.

[apilipenko]

lib/Transforms/Scalar/LoopSimplifyCFG.cpp

114–115 ↗

(On Diff #172898)

Minor. Extract a local lambda?

[mkazantsev]

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

8 ↗

(On Diff #172898)

Specifying the reason lets us be sure that we exercise all scenarios in this file. Besides, all these messages will be deleted one by one as the optimization is generalized. It's more a temporal thing for current development than something that will stay for long. Let's leave it as is, we're going to delete them anyways.

[mkazantsev]

lib/Transforms/Scalar/LoopSimplifyCFG.cpp

198 ↗

(On Diff #172301)

Sorry for misleading you earlier. This logic works if we *do not* delete the loop. I cannot remove it. It is required to detect the cases we are not supporting yet.

[mkazantsev]

lib/Transforms/Scalar/LoopSimplifyCFG.cpp

288 ↗

(On Diff #172301)

Both things are asserted now.

[mkazantsev]

Added lambda for printing;
Reworked to make it clear that we are doing it for non-dead loops;
Simplified logic in constantFoldTerminators

[anna]

lib/Transforms/Scalar/LoopSimplifyCFG.cpp

192 ↗

(On Diff #173121)

I find this name misleading because the current loop can be deleted in 2 ways:

1. latch to header edge is dead
2. edge to header from outside this loop is dead

Second case does not happen now, because you only consider blocks in current loop as candidate for folding.

However, I think you need to add both cases for completeness here.

198 ↗

(On Diff #172301)

okay, makes sense.

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

8 ↗

(On Diff #172898)

okay, pls correct typo: "blocks that are not dead".

[anna]

Forgot to mark request changes. See previous inline comments.

[mkazantsev]

lib/Transforms/Scalar/LoopSimplifyCFG.cpp

192 ↗

(On Diff #173121)

Technically, we only modify the current loop and cannot change anything outside this loop, so 2nd case is just out of scope.

[mkazantsev]

lib/Transforms/Scalar/LoopSimplifyCFG.cpp

192 ↗

(On Diff #173121)

I've added an explanatory comment why the second case is impossible in the field's comment.

[mkazantsev]

Fixed typo, added comment that explains why the loop can be dead in one way only.

[anna]

LGTM. Thanks for working through the comments.

lib/Transforms/Scalar/LoopSimplifyCFG.cpp

90 ↗

(On Diff #174351)

Minor: can remove the "in practice" here.

[dmgreen]

Hello. I have an error from this, which I think may well be a knock-on affect for the code later in the pass. Something like this:

define void @repo(i32 %size) {
entry:
  br label %loop

loop:
  %i = phi i32 [ 0, %entry ], [ %inc, %loop2 ]
  br i1 true, label %loop1, label %loop2

loop1:
  call void @something()
  br label %loop2

loop2:
  %phi = phi i32 [ 1, %loop ], [ 2, %loop1 ]
  %inc = add i32 %i, %phi
  %cmp = icmp ult i32 %inc, 3900
  br i1 %cmp, label %loop, label %end

end:
  ret void
}

declare void @something()

The br i1 true is const-folded, but the %phi gets the wrong value, taking the 1 where it should be taking the 2.

[mkazantsev]

Thanks for pointing out, the patch has been reverted due to very same problem. It is now fixed by recommit patch rL347289.

[dmgreen]

Hello, Another one for you, this time to do with LCSSA not being preserved, run with "opt -loop-simplifycfg -indvars simplified.ll -S":

target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
target triple = "aarch64-arm-none-eabi"

define void @c() {
entry:
  br label %d

d.loopexit:                                       ; preds = %if.end.7, %for.body
  %.lcssa = phi i32 [ %1, %for.body ], [ 0, %if.end.7 ]
  br label %d

d:                                                ; preds = %d.loopexit, %entry
  %0 = phi i32 [ undef, %entry ], [ %.lcssa, %d.loopexit ]
  br label %for.cond

for.cond:                                         ; preds = %if.end.8, %d
  %1 = phi i32 [ %0, %d ], [ 0, %if.end.8 ]
  br label %for.body

for.body:                                         ; preds = %for.cond
  %tobool2 = icmp eq i32 %1, 0
  br i1 %tobool2, label %if.end, label %d.loopexit

if.end:                                           ; preds = %for.body
  br label %if.end.7

if.end.7:                                         ; preds = %if.end
  br i1 true, label %if.end.8, label %d.loopexit

if.end.8:                                         ; preds = %if.end.7
  br label %for.cond
}

[mkazantsev]

Thanks for the test, I will fix it.

[mkazantsev]

Fix for it is D54841

llvm/trunk/lib/Transforms/Scalar/LoopSimplifyCFG.cpp

 #include "llvm/Transforms/Utils.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/Local.h"
 #include "llvm/Transforms/Utils/LoopUtils.h"
 using namespace llvm;
 
 #define DEBUG_TYPE "loop-simplifycfg"
 
+STATISTIC(NumTerminatorsFolded,
+          "Number of terminators folded to unconditional branches");
+
+/// If \p BB is a switch or a conditional branch, but only one of its successors
+/// can be reached from this block in runtime, return this successor. Otherwise,
+/// return nullptr.
+static BasicBlock *getOnlyLiveSuccessor(BasicBlock *BB) {
+  Instruction *TI = BB->getTerminator();
+  if (BranchInst *BI = dyn_cast<BranchInst>(TI)) {
+    if (BI->isUnconditional())
+      return nullptr;
+    if (BI->getSuccessor(0) == BI->getSuccessor(1))
+      return BI->getSuccessor(0);
+    ConstantInt *Cond = dyn_cast<ConstantInt>(BI->getCondition());
+    if (!Cond)
+      return nullptr;
+    return Cond->isZero() ? BI->getSuccessor(1) : BI->getSuccessor(0);
+  }
+
+  if (SwitchInst *SI = dyn_cast<SwitchInst>(TI)) {
+    auto *CI = dyn_cast<ConstantInt>(SI->getCondition());
+    if (!CI)
+      return nullptr;
+    for (auto Case : SI->cases())
+      if (Case.getCaseValue() == CI)
+        return Case.getCaseSuccessor();
+    return SI->getDefaultDest();
+  }
+
+  return nullptr;
+}
+
+/// Helper class that can turn branches and switches with constant conditions
+/// into unconditional branches.
+class ConstantTerminatorFoldingImpl {
+private:
+  Loop &L;
+  LoopInfo &LI;
+  DominatorTree &DT;
+
+  // 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
+  // has become unreachable from loop header.
+  bool DeleteCurrentLoop = false;
+
+  // The blocks of the original loop that will still be reachable from entry
+  // after the constant folding.
+  SmallPtrSet<BasicBlock *, 8> LiveLoopBlocks;
+  // The blocks of the original loop that will become unreachable from entry
+  // after the constant folding.
+  SmallPtrSet<BasicBlock *, 8> DeadLoopBlocks;
+  // The exits of the original loop that will still be reachable from entry
+  // after the constant folding.
+  SmallPtrSet<BasicBlock *, 8> LiveExitBlocks;
+  // The exits of the original loop that will become unreachable from entry
+  // after the constant folding.
+  SmallPtrSet<BasicBlock *, 8> DeadExitBlocks;
+  // The blocks that will still be a part of the current loop after folding.
+  SmallPtrSet<BasicBlock *, 8> BlocksInLoopAfterFolding;
+  // The blocks that have terminators with constant condition that can be
+  // folded. Note: fold candidates should be in L but not in any of its
+  // subloops to avoid complex LI updates.
+  SmallVector<BasicBlock *, 8> FoldCandidates;
+
+  void dump() const {
+    dbgs() << "Constant terminator folding for loop " << L << "\n";
+    dbgs() << "After terminator constant-folding, the loop will";
+    if (!DeleteCurrentLoop)
+      dbgs() << " not";
+    dbgs() << " be destroyed\n";
+    dbgs() << "Blocks in which we can constant-fold terminator:\n";
+    for (const BasicBlock *BB : FoldCandidates)
+      dbgs() << "\t" << BB->getName() << "\n";
+    auto PrintOutSet = [&](const char *Message,
+                           const SmallPtrSetImpl<BasicBlock *> &S) {
+      dbgs() << Message << "\n";
+      for (const BasicBlock *BB : S)
+        dbgs() << "\t" << BB->getName() << "\n";
+    };
+    PrintOutSet("Live blocks from the original loop:", LiveLoopBlocks);
+    PrintOutSet("Dead blocks from the original loop:", DeadLoopBlocks);
+    PrintOutSet("Live exit blocks:", LiveExitBlocks);
+    PrintOutSet("Dead exit blocks:", DeadExitBlocks);
+    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;
+
+      // If a loop block wasn't marked as live so far, then it's dead.
+      if (!LiveLoopBlocks.count(BB)) {
+        DeadLoopBlocks.insert(BB);
+        continue;
+      }
+
+      BasicBlock *TheOnlySucc = getOnlyLiveSuccessor(BB);
+
+      // If a block has only one live successor, it's a candidate on constant
+      // folding. Only handle blocks from current loop: branches in child loops
+      // are skipped because if they can be folded, they should be folded during
+      // the processing of child loops.
+      if (TheOnlySucc && LI.getLoopFor(BB) == &L)
+        FoldCandidates.push_back(BB);
+
+      // Handle successors.
+      auto ProcessSuccessor = [&](BasicBlock *Succ, bool IsLive) {
+        if (!L.contains(Succ)) {
+          if (IsLive)
+            LiveExitBlocks.insert(Succ);
+          ExitBlocks.insert(Succ);
+        } else if (IsLive)
+          LiveLoopBlocks.insert(Succ);
+      };
+      for (BasicBlock *Succ : successors(BB))
+        ProcessSuccessor(Succ, !TheOnlySucc || TheOnlySucc == Succ);
+    }
+
+    // Sanity check: amount of dead and live loop blocks should match the total
+    // number of blocks in loop.
+    assert(L.getNumBlocks() == LiveLoopBlocks.size() + DeadLoopBlocks.size() &&
+           "Malformed block sets?");
+
+    // Now, all exit blocks that are not marked as live are dead.
+    for (auto *ExitBlock : ExitBlocks)
+      if (!LiveExitBlocks.count(ExitBlock))
+        DeadExitBlocks.insert(ExitBlock);
+
+    // Whether or not the edge From->To will still be present in graph after the
+    // folding.
+    auto IsEdgeLive = [&](BasicBlock *From, BasicBlock *To) {
+      if (!LiveLoopBlocks.count(From))
+        return false;
+      BasicBlock *TheOnlySucc = getOnlyLiveSuccessor(From);
+      return !TheOnlySucc || TheOnlySucc == To;
+    };
+
+    // The loop will not be destroyed if its latch is live.
+    DeleteCurrentLoop = !IsEdgeLive(L.getLoopLatch(), L.getHeader());
+
+    // If we are going to delete the current loop completely, no extra analysis
+    // is needed.
+    if (DeleteCurrentLoop)
+      return;
+
+    // Otherwise, we should check which blocks will still be a part of the
+    // current loop after the transform.
+    BlocksInLoopAfterFolding.insert(L.getLoopLatch());
+    // If the loop is live, then we should compute what blocks are still in
+    // loop after all branch folding has been done. A block is in loop if
+    // it has a live edge to another block that is in the loop; by definition,
+    // latch is in the loop.
+    auto BlockIsInLoop = [&](BasicBlock *BB) {
+      return any_of(successors(BB), [&](BasicBlock *Succ) {
+        return BlocksInLoopAfterFolding.count(Succ) && IsEdgeLive(BB, Succ);
+      });
+    };
+    for (auto I = DFS.beginPostorder(), E = DFS.endPostorder(); I != E; ++I) {
+      BasicBlock *BB = *I;
+      if (BlockIsInLoop(BB))
+        BlocksInLoopAfterFolding.insert(BB);
+    }
+
+    // Sanity check: header must be in loop.
+    assert(BlocksInLoopAfterFolding.count(L.getHeader()) &&
+           "Header not in loop?");
+  }
+
+  /// Constant-fold terminators of blocks acculumated in FoldCandidates into the
+  /// unconditional branches.
+  void foldTerminators() {
+    DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Eager);
+
+    for (BasicBlock *BB : FoldCandidates) {
+      assert(LI.getLoopFor(BB) == &L && "Should be a loop block!");
+      BasicBlock *TheOnlySucc = getOnlyLiveSuccessor(BB);
+      assert(TheOnlySucc && "Should have one live successor!");
+
+      LLVM_DEBUG(dbgs() << "Replacing terminator of " << BB->getName()
+                        << " with an unconditional branch to the block "
+                        << TheOnlySucc->getName() << "\n");
+
+      SmallPtrSet<BasicBlock *, 2> DeadSuccessors;
+      // Remove all BB's successors except for the live one.
+      for (auto *Succ : successors(BB))
+        if (Succ != TheOnlySucc)
+          DeadSuccessors.insert(Succ);
+
+      IRBuilder<> Builder(BB->getContext());
+      Instruction *Term = BB->getTerminator();
+      Builder.SetInsertPoint(Term);
+      Builder.CreateBr(TheOnlySucc);
+      Term->eraseFromParent();
+
+      for (auto *DeadSucc : DeadSuccessors)
+        DTU.deleteEdge(BB, DeadSucc);
+
+      ++NumTerminatorsFolded;
+    }
+  }
+
+public:
+  ConstantTerminatorFoldingImpl(Loop &L, LoopInfo &LI, DominatorTree &DT)
+      : L(L), LI(LI), DT(DT) {}
+  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()
+                      << ": ");
+
+    // Nothing to constant-fold.
+    if (FoldCandidates.empty()) {
+      LLVM_DEBUG(
+          dbgs() << "No constant terminator folding candidates found in loop "
+                 << L.getHeader()->getName() << "\n");
+      return false;
+    }
+
+    // TODO: Support deletion of the current loop.
+    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 deletion of dead loop blocks.
+    if (!DeadLoopBlocks.empty()) {
+      LLVM_DEBUG(dbgs() << "Give up constant terminator folding in loop "
+                        << L.getHeader()->getName()
+                        << ": we don't currently"
+                           " support deletion of dead in-loop blocks.\n");
+      return false;
+    }
+
+    // TODO: Support dead loop exits.
+    if (!DeadExitBlocks.empty()) {
+      LLVM_DEBUG(dbgs() << "Give up constant terminator folding in loop "
+                        << L.getHeader()->getName()
+                        << ": we don't currently support dead loop exits.\n");
+      return false;
+    }
+
+    // TODO: Support blocks that are not dead, but also not in loop after the
+    // folding.
+    if (BlocksInLoopAfterFolding.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.
+    foldTerminators();
+
+#ifndef NDEBUG
+    // Make sure that we have preserved all data structures after the transform.
+    DT.verify();
+    assert(DT.isReachableFromEntry(L.getHeader()));
+    LI.verify(DT);
+#endif
+
+    return true;
+  }
+};
+
+/// Turn branches and switches with known constant conditions into unconditional
+/// branches.
+static bool constantFoldTerminators(Loop &L, DominatorTree &DT, LoopInfo &LI) {
+  // To keep things simple, only process loops with single latch. We
+  // canonicalize most loops to this form. We can support multi-latch if needed.
+  if (!L.getLoopLatch())
+    return false;
+
+  ConstantTerminatorFoldingImpl BranchFolder(L, LI, DT);
+  return BranchFolder.run();
+}
+
 static bool mergeBlocksIntoPredecessors(Loop &L, DominatorTree &DT,
                                         LoopInfo &LI, MemorySSAUpdater *MSSAU) {
   bool Changed = false;
   DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Eager);
   // Copy blocks into a temporary array to avoid iterator invalidation issues
   // as we remove them.
   SmallVector<WeakTrackingVH, 16> Blocks(L.blocks());
 
static bool mergeBlocksIntoPredecessors(Loop &L, DominatorTree &DT,
 
   return Changed;
 }
 
 static bool simplifyLoopCFG(Loop &L, DominatorTree &DT, LoopInfo &LI,
                             ScalarEvolution &SE, MemorySSAUpdater *MSSAU) {
   bool Changed = false;
 
+  // Constant-fold terminators with known constant conditions.
+  Changed |= constantFoldTerminators(L, DT, LI);
+
   // Eliminate unconditional branches by merging blocks into their predecessors.
   Changed |= mergeBlocksIntoPredecessors(L, DT, LI, MSSAU);
 
   if (Changed)
     SE.forgetTopmostLoop(&L);
 
   return Changed;
 }

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

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
-; RUN: opt -S -loop-simplifycfg < %s | FileCheck %s
-; RUN: opt -S -passes='require<domtree>,loop(simplify-cfg)' < %s | FileCheck %s
-; RUN: opt -S -loop-simplifycfg -enable-mssa-loop-dependency=true -verify-memoryssa < %s | FileCheck %s
+; 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_block_test_branch_loop: Give up constant terminator folding in loop header: we don't currently support deletion of dead in-loop blocks.
+; CHECK-LABEL: In function dead_block_test_switch_loop: Give up constant terminator folding in loop header: we don't currently support deletion of dead in-loop blocks.
+; CHECK-LABEL: In function dead_block_propogate_test_branch_loop: Give up constant terminator folding in loop header: we don't currently support deletion of dead in-loop blocks.
+; CHECK-LABEL: In function dead_block_propogate_test_switch_loop: Give up constant terminator folding in loop header: we don't currently support deletion of dead in-loop blocks.
+; CHECK-LABEL: In function dead_exit_test_branch_loop: Give up constant terminator folding in loop header: we don't currently support dead loop exits.
+; CHECK-LABEL: In function dead_exit_test_switch_loop: Give up constant terminator folding in loop header: we don't currently support dead loop exits.
+; CHECK-LABEL: In function dead_loop_test_branch_loop: Give up constant terminator folding in loop header: we don't currently support deletion of the current loop.
+; CHECK-LABEL: In function dead_loop_test_switch_loop: Give up constant terminator folding in loop header: we don't currently support deletion of the current loop.
+; CHECK-LABEL: In function dead_sub_loop_test_branch_loop: No constant terminator folding candidates found in loop dead_loop
+; CHECK-LABEL: In function dead_sub_loop_test_branch_loop: No constant terminator folding candidates found in loop live_loop
+; CHECK-LABEL: In function dead_sub_loop_test_branch_loop: Give up constant terminator folding in loop header: we don't currently support deletion of dead in-loop blocks.
+; CHECK-LABEL: In function dead_sub_loop_test_switch_loop: No constant terminator folding candidates found in loop live_loop
+; CHECK-LABEL: In function dead_sub_loop_test_switch_loop: No constant terminator folding candidates found in loop dead_loop
+; CHECK-LABEL: In function dead_sub_loop_test_switch_loop: Give up constant terminator folding in loop header: we don't currently support deletion of dead in-loop blocks.
+; CHECK-LABEL: In function inf_loop_test_branch_loop: Give up constant terminator folding in loop header: we don't currently support deletion of dead in-loop blocks.
+; CHECK-LABEL: In function inf_loop_test_switch_loop: Give up constant terminator folding in loop header: we don't currently support deletion of dead in-loop blocks.
+; CHECK-LABEL: In function live_block_test_branch_loop: Constant terminator folding for loop Loop at depth 1 containing: %header<header>,%check,%live,%backedge<latch><exiting>
+; CHECK:         Replacing terminator of check with an unconditional branch to the block backedge
+; CHECK-LABEL: In function live_block_test_switch_loop: Constant terminator folding for loop Loop at depth 1 containing: %header<header>,%check,%live,%backedge<latch><exiting>
+; CHECK:         Replacing terminator of check with an unconditional branch to the block backedge
+; CHECK-LABEL: In function partial_sub_loop_test_branch_loop: Give up constant terminator folding in loop header: we don't currently support deletion of dead in-loop blocks.
+; CHECK-LABEL: In function partial_sub_loop_test_branch_loop: No constant terminator folding candidates found in loop outer_header
+; CHECK-LABEL: In function partial_sub_loop_test_switch_loop: Give up constant terminator folding in loop header: we don't currently support deletion of dead in-loop blocks.
+; CHECK-LABEL: In function partial_sub_loop_test_switch_loop: No constant terminator folding candidates found in loop outer_header
+; CHECK-LABEL: In function full_sub_loop_test_branch_loop: Give up constant terminator folding in loop header: we don't currently support deletion of the current loop.
+; CHECK-LABEL: In function full_sub_loop_test_branch_loop: No constant terminator folding candidates found in loop outer_header
+; CHECK-LABEL: In function full_sub_loop_test_switch_loop: Give up constant terminator folding in loop header: we don't currently support deletion of the current loop.
+; CHECK-LABEL: In function full_sub_loop_test_switch_loop: No constant terminator folding candidates found in loop outer_header
+; CHECK-LABEL: In function full_sub_loop_test_branch_loop_inverse_1: Give up constant terminator folding in loop header: we don't currently support deletion of the current loop.
+; CHECK-LABEL: In function full_sub_loop_test_branch_loop_inverse_1: No constant terminator folding candidates found in loop outer_header
+; CHECK-LABEL: In function full_sub_loop_test_switch_loop_inverse_1: Give up constant terminator folding in loop header: we don't currently support deletion of the current loop.
+; CHECK-LABEL: In function full_sub_loop_test_switch_loop_inverse_1: No constant terminator folding candidates found in loop outer_header
+; CHECK-LABEL: In function full_sub_loop_test_branch_loop_inverse_2: Give up constant terminator folding in loop header: we don't currently support dead loop exits.
+; CHECK-LABEL: In function full_sub_loop_test_branch_loop_inverse_2: No constant terminator folding candidates found in loop outer_header
+; CHECK-LABEL: In function full_sub_loop_test_switch_loop_inverse_2: Give up constant terminator folding in loop header: we don't currently support dead loop exits.
+; CHECK-LABEL: In function full_sub_loop_test_switch_loop_inverse_2: No constant terminator folding candidates found in loop outer_header
+; CHECK-LABEL: In function full_sub_loop_test_branch_loop_inverse_3: Give up constant terminator folding in loop header: we don't currently support deletion of dead in-loop blocks.
+; CHECK-LABEL: In function full_sub_loop_test_branch_loop_inverse_3: No constant terminator folding candidates found in loop outer_header
+; CHECK-LABEL: In function full_sub_loop_test_switch_loop_inverse_3: Give up constant terminator folding in loop header: we don't currently support deletion of dead in-loop blocks.
+; CHECK-LABEL: In function full_sub_loop_test_switch_loop_inverse_3: No constant terminator folding candidates found in loop outer_header
+
 ; 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

 define i32 @live_block_test_branch_loop(i1 %c, i32 %end) {
 ; CHECK-LABEL: @live_block_test_branch_loop(
 ; CHECK-NEXT:  preheader:
 ; CHECK-NEXT:    br label [[HEADER:%.*]]
 ; CHECK:       header:
 ; CHECK-NEXT:    [[I:%.*]] = phi i32 [ 0, [[PREHEADER:%.*]] ], [ [[I_INC:%.*]], [[BACKEDGE:%.*]] ]
 ; CHECK-NEXT:    br i1 [[C:%.*]], label [[CHECK:%.*]], label [[LIVE:%.*]]
 ; CHECK:       check:
-; CHECK-NEXT:    br i1 true, label [[BACKEDGE]], label [[LIVE]]
+; CHECK-NEXT:    br label [[BACKEDGE]]
 ; CHECK:       live:
 ; CHECK-NEXT:    [[I_2:%.*]] = add i32 [[I]], 1
 ; CHECK-NEXT:    br label [[BACKEDGE]]
 ; CHECK:       backedge:
 ; CHECK-NEXT:    [[I_1:%.*]] = phi i32 [ [[I]], [[CHECK]] ], [ [[I_2]], [[LIVE]] ]
 ; CHECK-NEXT:    [[I_INC]] = add i32 [[I_1]], 1
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp slt i32 [[I_INC]], [[END:%.*]]
 ; CHECK-NEXT:    br i1 [[CMP]], label [[HEADER]], label [[EXIT:%.*]]

 define i32 @live_block_test_switch_loop(i1 %c, i32 %end) {
 ; CHECK-LABEL: @live_block_test_switch_loop(
 ; CHECK-NEXT:  preheader:
 ; CHECK-NEXT:    br label [[HEADER:%.*]]
 ; CHECK:       header:
 ; CHECK-NEXT:    [[I:%.*]] = phi i32 [ 0, [[PREHEADER:%.*]] ], [ [[I_INC:%.*]], [[BACKEDGE:%.*]] ]
 ; CHECK-NEXT:    br i1 [[C:%.*]], label [[CHECK:%.*]], label [[LIVE:%.*]]
 ; CHECK:       check:
-; CHECK-NEXT:    switch i32 1, label [[LIVE]] [
-; CHECK-NEXT:    i32 0, label [[LIVE]]
-; CHECK-NEXT:    i32 1, label [[BACKEDGE]]
-; CHECK-NEXT:    i32 2, label [[LIVE]]
-; CHECK-NEXT:    ]
+; CHECK-NEXT:    br label [[BACKEDGE]]
 ; CHECK:       live:
 ; CHECK-NEXT:    [[I_2:%.*]] = add i32 [[I]], 1
 ; CHECK-NEXT:    br label [[BACKEDGE]]
 ; CHECK:       backedge:
 ; CHECK-NEXT:    [[I_1:%.*]] = phi i32 [ [[I]], [[CHECK]] ], [ [[I_2]], [[LIVE]] ]
 ; CHECK-NEXT:    [[I_INC]] = add i32 [[I_1]], 1
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp slt i32 [[I_INC]], [[END:%.*]]
 ; CHECK-NEXT:    br i1 [[CMP]], label [[HEADER]], label [[EXIT:%.*]]