[unroll] Prune all but first copy of invariant exit

If we have an exit which is controlled by a loop invariant condition and which dominates the latch, we know only the copy in the first unrolled iteration can be taken. All other copies are dead.

The change itself is pretty straight forward, but let me add two points of context:
* I'd have expected other transform passes to catch this after unrolling, but I'm seeing multiple examples where we get to the end of O2/O3 without simplifying.
* I'd like to do a stronger change which did CSE during unroll and accounted for invariant expressions (as defined by SCEV instead of trivial ones from LoopInfo), but that doesn't fit cleanly into the current code structure.

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

1 files changed, 5 insertions, 0 deletions

diff --git a/llvm/lib/Transforms/Utils/LoopUnroll.cpp b/llvm/lib/Transforms/Utils/LoopUnroll.cpp
index b0c622b..0a530f5 100644
--- a/llvm/lib/Transforms/Utils/LoopUnroll.cpp
+++ b/llvm/lib/Transforms/Utils/LoopUnroll.cpp
@@ -310,6 +310,7 @@ LoopUnrollResult llvm::UnrollLoop(Loop *L, UnrollLoopOptions ULO, LoopInfo *LI,
     unsigned TripMultiple;
     unsigned BreakoutTrip;
     bool ExitOnTrue;
+    bool InvariantExit;
     SmallVector<BasicBlock *> ExitingBlocks;
   };
   DenseMap<BasicBlock *, ExitInfo> ExitInfos;
@@ -333,6 +334,8 @@ LoopUnrollResult llvm::UnrollLoop(Loop *L, UnrollLoopOptions ULO, LoopInfo *LI,
           (unsigned)GreatestCommonDivisor64(ULO.Count, Info.TripMultiple);
     }
     Info.ExitOnTrue = !L->contains(BI->getSuccessor(0));
+    Info.InvariantExit = L->isLoopInvariant(BI->getCondition()) &&
+      DT->dominates(ExitingBlock, LatchBlock);
     Info.ExitingBlocks.push_back(ExitingBlock);
     LLVM_DEBUG(dbgs() << "  Exiting block %" << ExitingBlock->getName()
                       << ": TripCount=" << Info.TripCount
@@ -685,6 +688,8 @@ LoopUnrollResult llvm::UnrollLoop(Loop *L, UnrollLoopOptions ULO, LoopInfo *LI,
 
   auto WillExit = [&](const ExitInfo &Info, unsigned i, unsigned j,
                       bool IsLatch) -> Optional<bool> {
+    if (Info.InvariantExit && i != 0)
+      return false;
     if (CompletelyUnroll) {
       if (PreserveOnlyFirst) {
         if (i == 0)