[LoopUnroll] Freeze tripcount rather than condition

This is a followup to D125754. We introduce two branches, one
before the unrolled loop and one before the epilogue (and similar
for the prologue case). The previous patch only froze the
condition on the first branch.

Rather than independently freezing the second condition, this patch
instead freezes TripCount and bases BECount on it. These are the
two quantities involved in the conditions, and this ensures that
both work on a consistent, non-poisonous trip count.

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

1 files changed, 20 insertions, 8 deletions

diff --git a/llvm/lib/Transforms/Utils/LoopUnrollRuntime.cpp b/llvm/lib/Transforms/Utils/LoopUnrollRuntime.cpp
index 71c15f9..37a013c 100644
--- a/llvm/lib/Transforms/Utils/LoopUnrollRuntime.cpp
+++ b/llvm/lib/Transforms/Utils/LoopUnrollRuntime.cpp
@@ -738,11 +738,28 @@ bool llvm::UnrollRuntimeLoopRemainder(
   // Compute the number of extra iterations required, which is:
   //  extra iterations = run-time trip count % loop unroll factor
   PreHeaderBR = cast<BranchInst>(PreHeader->getTerminator());
+  IRBuilder<> B(PreHeaderBR);
   Value *TripCount = Expander.expandCodeFor(TripCountSC, TripCountSC->getType(),
                                             PreHeaderBR);
-  Value *BECount = Expander.expandCodeFor(BECountSC, BECountSC->getType(),
-                                          PreHeaderBR);
-  IRBuilder<> B(PreHeaderBR);
+  Value *BECount;
+  // If there are other exits before the latch, that may cause the latch exit
+  // branch to never be executed, and the latch exit count may be poison.
+  // In this case, freeze the TripCount and base BECount on the frozen
+  // TripCount. We will introduce two branches using these values, and it's
+  // important that they see a consistent value (which would not be guaranteed
+  // if were frozen independently.)
+  if ((!OtherExits.empty() || !SE->loopHasNoAbnormalExits(L)) &&
+      !isGuaranteedNotToBeUndefOrPoison(TripCount, AC, PreHeaderBR, DT)) {
+    TripCount = B.CreateFreeze(TripCount);
+    BECount =
+        B.CreateAdd(TripCount, ConstantInt::get(TripCount->getType(), -1));
+  } else {
+    // If we don't need to freeze, use SCEVExpander for BECount as well, to
+    // allow slightly better value reuse.
+    BECount =
+        Expander.expandCodeFor(BECountSC, BECountSC->getType(), PreHeaderBR);
+  }
+
   Value * const ModVal = CreateTripRemainder(B, BECount, TripCount, Count);
 
   Value *BranchVal =
@@ -752,11 +769,6 @@ bool llvm::UnrollRuntimeLoopRemainder(
                            B.CreateIsNotNull(ModVal, "lcmp.mod");
   BasicBlock *RemainderLoop = UseEpilogRemainder ? NewExit : PrologPreHeader;
   BasicBlock *UnrollingLoop = UseEpilogRemainder ? NewPreHeader : PrologExit;
-  // Freeze the condition if there are other exits before the latch that may
-  // cause the latch exit branch to never be executed.
-  if ((!OtherExits.empty() || !SE->loopHasNoAbnormalExits(L)) &&
-      !isGuaranteedNotToBeUndefOrPoison(BranchVal, AC, PreHeaderBR, DT))
-    BranchVal = B.CreateFreeze(BranchVal);
   // Branch to either remainder (extra iterations) loop or unrolling loop.
   B.CreateCondBr(BranchVal, RemainderLoop, UnrollingLoop);
   PreHeaderBR->eraseFromParent();