[LoopUnroll] Push runtime unrolling decision up into tryToUnrollLoop()

[nikic]

Summary

Currently, UnrollLoop() is passed an AllowRuntime flag and decides itself whether runtime unrolling should be used or not. This patch pushes the decision into the caller and allows us to eliminate the ULO.TripCount and ULO.TripMultiple parameters.

[nikic]

Fix comment typos.

[nikic]

llvm/lib/Transforms/Utils/LoopUnroll.cpp

391

I'm not particularly familiar with convergent operations, but looking back at D17526 the concern here seems to have been specifically about the prologue introduced by runtime unrolling, not anything else.

llvm/lib/Transforms/Utils/LoopUnrollAndJam.cpp

251

This change deserves a comment: If TripMultiple is 1 and Count is 1 and UnrollRemainder was set, this tried to unroll the remainder loop with a Count of 0, which is not legal. It wasn't caught previously due to the check

// Don't enter the unroll code if there is nothing to do.
if (ULO.TripCount == 0 && ULO.Count < 2) {

which happened before the assert(Count > 0) assertion. I've removed that check and the assertion is triggered now.

The change here adjust things to not try to create a remainder loop for an unroll count of 1, as this doesn't make sense in the first place.

[reames]

LGTM. Nice cleanup.

llvm/include/llvm/Transforms/Utils/UnrollLoop.h

enum class LoopUnrollResult {
 
   /// The loop was fully unrolled into straight-line code.  We no longer have
   /// any back-edges.
   FullyUnrolled
 };
 
 struct UnrollLoopOptions {
   unsigned Count;
-  unsigned TripCount;
   bool Force;
-  bool AllowRuntime;
+  bool Runtime;
   bool AllowExpensiveTripCount;
-  unsigned TripMultiple;
   bool UnrollRemainder;
   bool ForgetAllSCEV;
 };
 
 LoopUnrollResult UnrollLoop(Loop *L, UnrollLoopOptions ULO, LoopInfo *LI,
                             ScalarEvolution *SE, DominatorTree *DT,
                             AssumptionCache *AC,
                             const llvm::TargetTransformInfo *TTI,

llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp

if (peelLoop(L, PP.PeelCount, LI, &SE, &DT, &AC, PreserveLCSSA)) {
       // we had, so we don't want to unroll or peel again.
       if (PP.PeelProfiledIterations)
         L->setLoopAlreadyUnrolled();
       return LoopUnrollResult::PartiallyUnrolled;
     }
     return LoopUnrollResult::Unmodified;
   }
 
+  // At this point, UP.Runtime indicates that run-time unrolling is allowed.
+  // However, we only want to actually perform it if we don't know the trip
+  // count and the unroll count doesn't divide the known trip multiple.
+  // TODO: This decision should probably be pushed up into
+  // computeUnrollCount().
+  UP.Runtime &= TripCount == 0 && TripMultiple % UP.Count != 0;
+
   // Save loop properties before it is transformed.
   MDNode *OrigLoopID = L->getLoopID();
 
   // Unroll the loop.
   Loop *RemainderLoop = nullptr;
   LoopUnrollResult UnrollResult = UnrollLoop(
       L,
-      {UP.Count, TripCount, UP.Force, UP.Runtime, UP.AllowExpensiveTripCount,
-       TripMultiple, UP.UnrollRemainder, ForgetAllSCEV},
+      {UP.Count, UP.Force, UP.Runtime, UP.AllowExpensiveTripCount,
+       UP.UnrollRemainder, ForgetAllSCEV},
       LI, &SE, &DT, &AC, &TTI, &ORE, PreserveLCSSA, &RemainderLoop);
   if (UnrollResult == LoopUnrollResult::Unmodified)
     return LoopUnrollResult::Unmodified;
 
   if (RemainderLoop) {
     Optional<MDNode *> RemainderLoopID =
         makeFollowupLoopID(OrigLoopID, {LLVMLoopUnrollFollowupAll,
                                         LLVMLoopUnrollFollowupRemainder});

llvm/lib/Transforms/Utils/LoopUnroll.cpp

void llvm::simplifyLoopAfterUnroll(Loop *L, bool SimplifyIVs, LoopInfo *LI,
   }
 }
 
 /// Unroll the given loop by Count. The loop must be in LCSSA form.  Unrolling
 /// can only fail when the loop's latch block is not terminated by a conditional
 /// branch instruction. However, if the trip count (and multiple) are not known,
 /// loop unrolling will mostly produce more code that is no faster.
 ///
-/// TripCount is an upper bound on the number of times the loop header runs.
-/// Note that the trip count does not need to be exact, it can be any upper
-/// bound on the true trip count.
-///
-/// Similarly, TripMultiple divides the number of times that the LatchBlock may
-/// execute without exiting the loop.
-///
-/// If AllowRuntime is true then UnrollLoop will consider unrolling loops that
-/// have a runtime (i.e. not compile time constant) trip count.  Unrolling these
-/// loops require a unroll "prologue" that runs "RuntimeTripCount % Count"
-/// iterations before branching into the unrolled loop.  UnrollLoop will not
-/// runtime-unroll the loop if computing RuntimeTripCount will be expensive and
-/// AllowExpensiveTripCount is false.
+/// If Runtime is true then UnrollLoop will try to insert a prologue or
+/// epilogue that ensures the latch has a trip multiple of Count. UnrollLoop
+/// will not runtime-unroll the loop if computing the run-time trip count will
+/// be expensive and AllowExpensiveTripCount is false.
 ///
 /// The LoopInfo Analysis that is passed will be kept consistent.
 ///
 /// This utility preserves LoopInfo. It will also preserve ScalarEvolution and
 /// DominatorTree if they are non-null.
 ///
 /// If RemainderLoop is non-null, it will receive the remainder loop (if
 /// required and not fully unrolled).
LoopUnrollResult llvm::UnrollLoop(Loop *L, UnrollLoopOptions ULO, LoopInfo *LI,
 
   if (L->getHeader()->hasAddressTaken()) {
     // The loop-rotate pass can be helpful to avoid this in many cases.
     LLVM_DEBUG(
         dbgs() << "  Won't unroll loop: address of header block is taken.\n");
     return LoopUnrollResult::Unmodified;
   }
 
-  if (ULO.TripCount != 0)
-    LLVM_DEBUG(dbgs() << "  Trip Count = " << ULO.TripCount << "\n");
-  if (ULO.TripMultiple != 1)
-    LLVM_DEBUG(dbgs() << "  Trip Multiple = " << ULO.TripMultiple << "\n");
-
-  // Don't enter the unroll code if there is nothing to do.
-  if (ULO.TripCount == 0 && ULO.Count < 2) {
-    LLVM_DEBUG(dbgs() << "Won't unroll; almost nothing to do\n");
-    return LoopUnrollResult::Unmodified;
-  }
-
   assert(ULO.Count > 0);
-  assert(ULO.TripMultiple > 0);
-  assert(ULO.TripCount == 0 || ULO.TripCount % ULO.TripMultiple == 0);
 
   // All these values should be taken only after peeling because they might have
   // changed.
   BasicBlock *Preheader = L->getLoopPreheader();
   BasicBlock *Header = L->getHeader();
   BasicBlock *LatchBlock = L->getLoopLatch();
   SmallVector<BasicBlock *, 4> ExitBlocks;
   L->getExitBlocks(ExitBlocks);
LoopUnrollResult llvm::UnrollLoop(Loop *L, UnrollLoopOptions ULO, LoopInfo *LI,
 
   // Are we eliminating the loop control altogether?  Note that we can know
   // we're eliminating the backedge without knowing exactly which iteration
   // of the unrolled body exits.
   const bool CompletelyUnroll = ULO.Count == MaxTripCount;
 
   const bool PreserveOnlyFirst = CompletelyUnroll && MaxOrZero;
 
-  // We assume a run-time trip count if the compiler cannot
-  // figure out the loop trip count and the unroll-runtime
-  // flag is specified.
-  bool RuntimeTripCount = !CompletelyUnroll && ULO.TripCount == 0 &&
-                          ULO.TripMultiple % ULO.Count != 0 && ULO.AllowRuntime;
+  // There's no point in performing runtime unrolling if this unroll count
+  // results in a full unroll.
+  if (CompletelyUnroll)
+    ULO.Runtime = false;
 
   // Go through all exits of L and see if there are any phi-nodes there. We just
   // conservatively assume that they're inserted to preserve LCSSA form, which
   // means that complete unrolling might break this form. We need to either fix
   // it in-place after the transformation, or entirely rebuild LCSSA. TODO: For
   // now we just recompute LCSSA for the outer loop, but it should be possible
   // to fix it in-place.
   bool NeedToFixLCSSA =
LoopUnrollResult llvm::UnrollLoop(Loop *L, UnrollLoopOptions ULO, LoopInfo *LI,
   // unconditional branch in the unrolled loop in some cases.
   bool LatchIsExiting = L->isLoopExiting(LatchBlock);
   if (!LatchBI || (LatchBI->isConditional() && !LatchIsExiting)) {
     LLVM_DEBUG(
         dbgs() << "Can't unroll; a conditional latch must exit the loop");
     return LoopUnrollResult::Unmodified;
   }
 
-  // Loops containing convergent instructions must have a count that divides
-  // their TripMultiple.
+  // Loops containing convergent instructions cannot use runtime unrolling,
+  // as the prologue/epilogue may add additional control-dependencies to
+  // convergent operations.
   LLVM_DEBUG(
       {
         bool HasConvergent = false;
         for (auto &BB : L->blocks())
           for (auto &I : *BB)
             if (auto *CB = dyn_cast<CallBase>(&I))
               HasConvergent |= CB->isConvergent();
-        assert((!HasConvergent || ULO.TripMultiple % ULO.Count == 0) &&
-               "Unroll count must divide trip multiple if loop contains a "
-               "convergent operation.");
+        assert((!HasConvergent || !ULO.Runtime) &&
+               "Can't runtime unroll if loop contains a convergent operation.");
       });
 
   bool EpilogProfitability =
       UnrollRuntimeEpilog.getNumOccurrences() ? UnrollRuntimeEpilog
                                               : isEpilogProfitable(L);
 
-  if (RuntimeTripCount &&
+  if (ULO.Runtime &&
       !UnrollRuntimeLoopRemainder(L, ULO.Count, ULO.AllowExpensiveTripCount,
                                   EpilogProfitability, ULO.UnrollRemainder,
                                   ULO.ForgetAllSCEV, LI, SE, DT, AC, TTI,
                                   PreserveLCSSA, RemainderLoop)) {
     if (ULO.Force)
-      RuntimeTripCount = false;
+      ULO.Runtime = false;
     else {
       LLVM_DEBUG(dbgs() << "Won't unroll; remainder loop could not be "
                            "generated when assuming runtime trip count\n");
       return LoopUnrollResult::Unmodified;
     }
   }
 
   using namespace ore;
   // Report the unrolling decision.
   if (CompletelyUnroll) {
     LLVM_DEBUG(dbgs() << "COMPLETELY UNROLLING loop %" << Header->getName()
                       << " with trip count " << ULO.Count << "!\n");
     if (ORE)
       ORE->emit([&]() {
         return OptimizationRemark(DEBUG_TYPE, "FullyUnrolled", L->getStartLoc(),
                                   L->getHeader())
                << "completely unrolled loop with "
                << NV("UnrollCount", ULO.Count) << " iterations";
       });
   } else {
     LLVM_DEBUG(dbgs() << "UNROLLING loop %" << Header->getName() << " by "
                       << ULO.Count);
-    if (RuntimeTripCount)
+    if (ULO.Runtime)
       LLVM_DEBUG(dbgs() << " with run-time trip count");
     LLVM_DEBUG(dbgs() << "!\n");
 
     if (ORE)
       ORE->emit([&]() {
         OptimizationRemark Diag(DEBUG_TYPE, "PartialUnrolled", L->getStartLoc(),
                                 L->getHeader());
         Diag << "unrolled loop by a factor of " << NV("UnrollCount", ULO.Count);
-        if (RuntimeTripCount)
+        if (ULO.Runtime)
           Diag << " with run-time trip count";
         return Diag;
       });
   }
 
   // We are going to make changes to this loop. SCEV may be keeping cached info
   // about it, in particular about backedge taken count. The changes we make
   // are guaranteed to invalidate this information for our loop. It is tempting
if (CompletelyUnroll) {
       // Complete (but possibly inexact) unrolling
       if (j == 0)
         return true;
       if (Info.TripCount && j != Info.TripCount)
         return false;
       return None;
     }
 
-    if (RuntimeTripCount) {
+    if (ULO.Runtime) {
       // If runtime unrolling inserts a prologue, information about non-latch
       // exits may be stale.
       if (IsLatch && j != 0)
         return false;
       return None;
     }
 
     if (j != Info.BreakoutTrip &&

llvm/lib/Transforms/Utils/LoopUnrollAndJam.cpp

llvm::UnrollAndJamLoop(Loop *L, unsigned Count, unsigned TripCount,
   assert(Count > 0);
   assert(TripMultiple > 0);
   assert(TripCount == 0 || TripCount % TripMultiple == 0);
 
   // Are we eliminating the loop control altogether?
   bool CompletelyUnroll = (Count == TripCount);
 
   // We use the runtime remainder in cases where we don't know trip multiple
-  if (TripMultiple == 1 || TripMultiple % Count != 0) {
+  if (TripMultiple % Count != 0) {
     if (!UnrollRuntimeLoopRemainder(L, Count, /*AllowExpensiveTripCount*/ false,
                                     /*UseEpilogRemainder*/ true,
                                     UnrollRemainder, /*ForgetAllSCEV*/ false,
                                     LI, SE, DT, AC, TTI, true, EpilogueLoop)) {
       LLVM_DEBUG(dbgs() << "Won't unroll-and-jam; remainder loop could not be "
                            "generated when assuming runtime trip count\n");
       return LoopUnrollResult::Unmodified;
     }

llvm/lib/Transforms/Utils/LoopUnrollRuntime.cpp

if (remainderLoop)
       formDedicatedExitBlocks(remainderLoop, DT, LI, nullptr, PreserveLCSSA);
   }
 
   auto UnrollResult = LoopUnrollResult::Unmodified;
   if (remainderLoop && UnrollRemainder) {
     LLVM_DEBUG(dbgs() << "Unrolling remainder loop\n");
     UnrollResult =
         UnrollLoop(remainderLoop,
-                   {/*Count*/ Count - 1, /*TripCount*/ Count - 1,
-                    /*Force*/ false, /*AllowRuntime*/ false,
-                    /*AllowExpensiveTripCount*/ false, /*TripMultiple*/ 1,
+                   {/*Count*/ Count - 1, /*Force*/ false, /*Runtime*/ false,
+                    /*AllowExpensiveTripCount*/ false,
                     /*UnrollRemainder*/ false, ForgetAllSCEV},
                    LI, SE, DT, AC, TTI, /*ORE*/ nullptr, PreserveLCSSA);
   }
 
   if (ResultLoop && UnrollResult != LoopUnrollResult::FullyUnrolled)
     *ResultLoop = remainderLoop;
   NumRuntimeUnrolled++;
   return true;
 }