[IndVars] Extend reasoning about loop invariant exits to non-header blocks
Noticed while glancing through the code for other reasons. The extension is trivial enough, decided to just do it.
llvm-svn: 360694
diff --git a/llvm/lib/Transforms/Scalar/IndVarSimplify.cpp b/llvm/lib/Transforms/Scalar/IndVarSimplify.cpp
index ecf0bef..333de75 100644
--- a/llvm/lib/Transforms/Scalar/IndVarSimplify.cpp
+++ b/llvm/lib/Transforms/Scalar/IndVarSimplify.cpp
@@ -728,11 +728,13 @@
IncomingValIdx != E; ++IncomingValIdx) {
auto *IncomingBB = PN.getIncomingBlock(IncomingValIdx);
- // We currently only support loop exits from loop header. If the
- // incoming block is not loop header, we need to recursively check
- // all conditions starting from loop header are loop invariants.
- // Additional support might be added in the future.
- if (IncomingBB != LoopHeader)
+ // Can we prove that the exit must run on the first iteration if it
+ // runs at all? (i.e. early exits are fine for our purposes, but
+ // traces which lead to this exit being taken on the 2nd iteration
+ // aren't.) Note that this is about whether the exit branch is
+ // executed, not about whether it is taken.
+ if (!L->getLoopLatch() ||
+ !DT->dominates(IncomingBB, L->getLoopLatch()))
continue;
// Get condition that leads to the exit path.
@@ -753,8 +755,8 @@
auto *ExitVal = dyn_cast<PHINode>(PN.getIncomingValue(IncomingValIdx));
- // Only deal with PHIs.
- if (!ExitVal)
+ // Only deal with PHIs in the loop header.
+ if (!ExitVal || ExitVal->getParent() != L->getHeader())
continue;
// If ExitVal is a PHI on the loop header, then we know its