Try reminder unrolling, when checks fail, rotate loop and retry

Author: Marek Sedlacek

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

@@ -72,18 +72,6 @@ enum class LoopUnrollResult {
   FullyUnrolled
 };
 
-/// Represents the result of a \c UnrollRuntimeLoopRemainder invocation.
-enum class LoopReminderUnrollResult {
-  /// The loop reminder was not modified.
-  Unmodified,
-
-  /// The loop was rotated, but not unrolled.
-  Rotated,
-
-  /// The loop reminder was unrolled.
-  Unrolled
-};
-
 struct UnrollLoopOptions {
   unsigned Count;
   bool Force;
@@ -105,13 +93,13 @@ LLVM_ABI LoopUnrollResult UnrollLoop(Loop *L, UnrollLoopOptions ULO,
                                      Loop **RemainderLoop = nullptr,
                                      AAResults *AA = nullptr);
 
-LLVM_ABI LoopReminderUnrollResult UnrollRuntimeLoopRemainder(
+LLVM_ABI bool UnrollRuntimeLoopRemainder(
     Loop *L, unsigned Count, bool AllowExpensiveTripCount,
     bool UseEpilogRemainder, bool UnrollRemainder, bool ForgetAllSCEV,
     LoopInfo *LI, ScalarEvolution *SE, DominatorTree *DT, AssumptionCache *AC,
     const TargetTransformInfo *TTI, bool PreserveLCSSA,
     unsigned SCEVExpansionBudget, bool RuntimeUnrollMultiExit,
-    bool AllowLoopRotation, Loop **ResultLoop = nullptr);
+    Loop **ResultLoop = nullptr);
 
 LLVM_ABI LoopUnrollResult UnrollAndJamLoop(
     Loop *L, unsigned Count, unsigned TripCount, unsigned TripMultiple,

llvm/lib/Transforms/Utils/LoopUnroll.cpp

@@ -58,6 +58,7 @@
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/Cloning.h"
 #include "llvm/Transforms/Utils/Local.h"
+#include "llvm/Transforms/Utils/LoopRotationUtils.h"
 #include "llvm/Transforms/Utils/LoopSimplify.h"
 #include "llvm/Transforms/Utils/LoopUtils.h"
 #include "llvm/Transforms/Utils/SimplifyIndVar.h"
@@ -535,31 +536,60 @@ llvm::UnrollLoop(Loop *L, UnrollLoopOptions ULO, LoopInfo *LI,
       UnrollRuntimeEpilog.getNumOccurrences() ? UnrollRuntimeEpilog
                                               : isEpilogProfitable(L);
 
-  LoopReminderUnrollResult UnrollReminderResult =
-      LoopReminderUnrollResult::Unmodified;
+  bool LoopRotated = false;
+  bool ReminderUnrolled = false;
   if (ULO.Runtime) {
-    UnrollReminderResult = UnrollRuntimeLoopRemainder(
+    // Call unroll with disabled rotation, to see if it is possible without it.
+    ReminderUnrolled = UnrollRuntimeLoopRemainder(
         L, ULO.Count, ULO.AllowExpensiveTripCount, EpilogProfitability,
         ULO.UnrollRemainder, ULO.ForgetAllSCEV, LI, SE, DT, AC, TTI,
         PreserveLCSSA, ULO.SCEVExpansionBudget, ULO.RuntimeUnrollMultiExit,
-        /*AllowLoopRotation*/ true, RemainderLoop);
+        RemainderLoop);
+
+    // If unroll is not possible, then try with loop rotation.
+    if (!ReminderUnrolled) {
+      BasicBlock *OrigHeader = L->getHeader();
+      BranchInst *BI = dyn_cast<BranchInst>(OrigHeader->getTerminator());
+      if (BI && !BI->isUnconditional() &&
+          isa<SCEVCouldNotCompute>(SE->getExitCount(L, L->getLoopLatch())) &&
+          !isa<SCEVCouldNotCompute>(SE->getExitCount(L, OrigHeader))) {
+        LLVM_DEBUG(
+            dbgs() << "  Rotating loop to make the exit count computable.\n");
+        SimplifyQuery SQ{OrigHeader->getDataLayout()};
+        SQ.TLI = nullptr;
+        SQ.DT = DT;
+        SQ.AC = AC;
+        LoopRotated =
+            llvm::LoopRotation(L, LI, TTI, AC, DT, SE,
+                               /*MemorySSAUpdater*/ nullptr, SQ,
+                               /*RotationOnly*/ false, /*Threshold*/ 16,
+                               /*IsUtilMode*/ false, /*PrepareForLTO*/ false,
+                               [](Loop *, ScalarEvolution *) { return true; });
+      }
+      if (LoopRotated) {
+        // Loop was rotated, try unrolling.
+        ReminderUnrolled = UnrollRuntimeLoopRemainder(
+            L, ULO.Count, ULO.AllowExpensiveTripCount, EpilogProfitability,
+            ULO.UnrollRemainder, ULO.ForgetAllSCEV, LI, SE, DT, AC, TTI,
+            PreserveLCSSA, ULO.SCEVExpansionBudget, ULO.RuntimeUnrollMultiExit,
+            RemainderLoop);
+      }
+    }
+    // Latch block needs to be updated.
     LatchBlock = L->getLoopLatch();
     LatchIsExiting = L->isLoopExiting(LatchBlock);
   }
 
-  if (ULO.Runtime &&
-      UnrollReminderResult != LoopReminderUnrollResult::Unrolled) {
+  if (ULO.Runtime && !ReminderUnrolled) {
     if (ULO.Force)
       ULO.Runtime = false;
     else {
       LLVM_DEBUG(dbgs() << "Won't unroll; remainder loop could not be "
                            "generated when assuming runtime trip count\n");
       // Loop might have been rotated inside of UnrollRuntimeLoopRemainder and
       // this needs to be propagated.
-      return UnrollReminderResult == LoopReminderUnrollResult::Rotated
-                 ? LoopUnrollResult::Modified
-                 : LoopUnrollResult::Unmodified;
-      ;
+      return LoopRotated ? LoopUnrollResult::Modified
+                         : LoopUnrollResult::Unmodified;
     }
   }
 

llvm/lib/Transforms/Utils/LoopUnrollAndJam.cpp

@@ -237,18 +237,16 @@ llvm::UnrollAndJamLoop(Loop *L, unsigned Count, unsigned TripCount,
   // 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 % Count != 0) {
-    auto UnrollReminderResult = UnrollRuntimeLoopRemainder(
-        L, Count, /*AllowExpensiveTripCount*/ false,
-        /*UseEpilogRemainder*/ true, UnrollRemainder, /*ForgetAllSCEV*/ false,
-        LI, SE, DT, AC, TTI, /*PreserveLCSSA*/ true, SCEVCheapExpansionBudget,
-        /*AllowLoopRotation*/ false, EpilogueLoop);
-    if (UnrollReminderResult != LoopReminderUnrollResult::Unrolled) {
+    if (!UnrollRuntimeLoopRemainder(L, Count, /*AllowExpensiveTripCount*/ false,
+                                    /*UseEpilogRemainder*/ true,
+                                    UnrollRemainder, /*ForgetAllSCEV*/ false,
+                                    LI, SE, DT, AC, TTI, true,
+                                    SCEVCheapExpansionBudget, EpilogueLoop)) {
       LLVM_DEBUG(dbgs() << "Won't unroll-and-jam; remainder loop could not be "
                            "generated when assuming runtime trip count\n");
-      return UnrollReminderResult == LoopReminderUnrollResult::Rotated
-                 ? LoopUnrollResult::Modified
-                 : LoopUnrollResult::Unmodified;
+      return LoopUnrollResult::Unmodified;
     }
   }
 

llvm/lib/Transforms/Utils/LoopUnrollRuntime.cpp

@@ -37,7 +37,6 @@
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/Cloning.h"
 #include "llvm/Transforms/Utils/Local.h"
-#include "llvm/Transforms/Utils/LoopRotationUtils.h"
 #include "llvm/Transforms/Utils/LoopUtils.h"
 #include "llvm/Transforms/Utils/ScalarEvolutionExpander.h"
 #include "llvm/Transforms/Utils/UnrollLoop.h"
@@ -575,46 +574,22 @@ static Value *CreateTripRemainder(IRBuilder<> &B, Value *BECount,
 ///        if (extraiters != 0) jump Epil: // Omitted if unroll factor is 2.
 /// EpilExit:
 
-LoopReminderUnrollResult llvm::UnrollRuntimeLoopRemainder(
+bool llvm::UnrollRuntimeLoopRemainder(
     Loop *L, unsigned Count, bool AllowExpensiveTripCount,
     bool UseEpilogRemainder, bool UnrollRemainder, bool ForgetAllSCEV,
     LoopInfo *LI, ScalarEvolution *SE, DominatorTree *DT, AssumptionCache *AC,
     const TargetTransformInfo *TTI, bool PreserveLCSSA,
     unsigned SCEVExpansionBudget, bool RuntimeUnrollMultiExit,
-    bool AllowLoopRotation, Loop **ResultLoop) {
+    Loop **ResultLoop) {
   LLVM_DEBUG(dbgs() << "Trying runtime unrolling on Loop: \n");
   LLVM_DEBUG(L->dump());
   LLVM_DEBUG(UseEpilogRemainder ? dbgs() << "Using epilog remainder.\n"
                                 : dbgs() << "Using prolog remainder.\n");
 
-  LoopReminderUnrollResult Result = LoopReminderUnrollResult::Unmodified;
-
-  // Rotate loop if it makes the exit count from the latch computable.
-  if (AllowLoopRotation) {
-    BasicBlock *OrigHeader = L->getHeader();
-    BranchInst *BI = dyn_cast<BranchInst>(OrigHeader->getTerminator());
-    if (BI && !BI->isUnconditional() &&
-        isa<SCEVCouldNotCompute>(SE->getExitCount(L, L->getLoopLatch())) &&
-        !isa<SCEVCouldNotCompute>(SE->getExitCount(L, OrigHeader))) {
-      LLVM_DEBUG(
-          dbgs() << "  Rotating loop to make the exit count computable.\n");
-      SimplifyQuery SQ{OrigHeader->getDataLayout()};
-      SQ.TLI = nullptr;
-      SQ.DT = DT;
-      SQ.AC = AC;
-      if (llvm::LoopRotation(L, LI, TTI, AC, DT, SE,
-                             /*MemorySSAUpdater*/ nullptr, SQ,
-                             /*RotationOnly*/ false, /*Threshold*/ 16,
-                             /*IsUtilMode*/ false, /*PrepareForLTO*/ false,
-                             [](Loop *, ScalarEvolution *) { return true; }))
-        Result = LoopReminderUnrollResult::Rotated;
-    }
-  }
-
   // Make sure the loop is in canonical form.
   if (!L->isLoopSimplifyForm()) {
     LLVM_DEBUG(dbgs() << "Not in simplify form!\n");
-    return Result;
+    return false;
   }
 
   // Guaranteed by LoopSimplifyForm.
@@ -628,7 +603,7 @@ LoopReminderUnrollResult llvm::UnrollRuntimeLoopRemainder(
     LLVM_DEBUG(
         dbgs()
         << "Loop latch not terminated by a conditional branch.\n");
-    return Result;
+    return false;
   }
 
   unsigned ExitIndex = LatchBR->getSuccessor(0) == Header ? 1 : 0;
@@ -640,7 +615,7 @@ LoopReminderUnrollResult llvm::UnrollRuntimeLoopRemainder(
     LLVM_DEBUG(
         dbgs()
         << "One of the loop latch successors must be the exit block.\n");
-    return Result;
+    return false;
   }
 
   // These are exit blocks other than the target of the latch exiting block.
@@ -652,12 +627,12 @@ LoopReminderUnrollResult llvm::UnrollRuntimeLoopRemainder(
     // We rely on LCSSA form being preserved when the exit blocks are transformed.
     // (Note that only an off-by-default mode of the old PM disables PreserveLCCA.)
     if (!PreserveLCSSA)
-      return Result;
+      return false;
 
     // Priority goes to UnrollRuntimeMultiExit if it's supplied.
     if (UnrollRuntimeMultiExit.getNumOccurrences()) {
       if (!UnrollRuntimeMultiExit)
-        return Result;
+        return false;
     } else {
       // Otherwise perform multi-exit unrolling, if either the target indicates
       // it is profitable or the general profitability heuristics apply.
@@ -666,14 +641,14 @@ LoopReminderUnrollResult llvm::UnrollRuntimeLoopRemainder(
                                                    UseEpilogRemainder)) {
         LLVM_DEBUG(dbgs() << "Multiple exit/exiting blocks in loop and "
                              "multi-exit unrolling not enabled!\n");
-        return Result;
+        return false;
       }
     }
   }
   // Use Scalar Evolution to compute the trip count. This allows more loops to
   // be unrolled than relying on induction var simplification.
   if (!SE)
-    return Result;
+    return false;
 
   // Only unroll loops with a computable trip count.
   // We calculate the backedge count by using getExitCount on the Latch block,
@@ -683,7 +658,7 @@ LoopReminderUnrollResult llvm::UnrollRuntimeLoopRemainder(
   const SCEV *BECountSC = SE->getExitCount(L, Latch);
   if (isa<SCEVCouldNotCompute>(BECountSC)) {
     LLVM_DEBUG(dbgs() << "Could not compute exit block SCEV\n");
-    return Result;
+    return false;
   }
 
   unsigned BEWidth = cast<IntegerType>(BECountSC->getType())->getBitWidth();
@@ -694,7 +669,7 @@ LoopReminderUnrollResult llvm::UnrollRuntimeLoopRemainder(
       SE->getAddExpr(BECountSC, SE->getConstant(BECountSC->getType(), 1));
   if (isa<SCEVCouldNotCompute>(TripCountSC)) {
     LLVM_DEBUG(dbgs() << "Could not compute trip count SCEV.\n");
-    return Result;
+    return false;
   }
 
   BasicBlock *PreHeader = L->getLoopPreheader();
@@ -705,7 +680,7 @@ LoopReminderUnrollResult llvm::UnrollRuntimeLoopRemainder(
       Expander.isHighCostExpansion(TripCountSC, L, SCEVExpansionBudget, TTI,
                                    PreHeaderBR)) {
     LLVM_DEBUG(dbgs() << "High cost for expanding trip count scev!\n");
-    return Result;
+    return false;
   }
 
   // This constraint lets us deal with an overflowing trip count easily; see the
@@ -714,7 +689,7 @@ LoopReminderUnrollResult llvm::UnrollRuntimeLoopRemainder(
     LLVM_DEBUG(
         dbgs()
         << "Count failed constraint on overflow trip count calculation.\n");
-    return Result;
+    return false;
   }
 
   // Loop structure is the following:
@@ -1060,5 +1035,5 @@ LoopReminderUnrollResult llvm::UnrollRuntimeLoopRemainder(
   if (ResultLoop && UnrollResult != LoopUnrollResult::FullyUnrolled)
     *ResultLoop = remainderLoop;
   NumRuntimeUnrolled++;
-  return LoopReminderUnrollResult::Unrolled;
+  return true;
 }