Added flag to enable loop rotation in UnrollRuntimeLoopRemainder

Author: Marek Sedlacek

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

@@ -111,7 +111,7 @@ LLVM_ABI LoopReminderUnrollResult UnrollRuntimeLoopRemainder(
     LoopInfo *LI, ScalarEvolution *SE, DominatorTree *DT, AssumptionCache *AC,
     const TargetTransformInfo *TTI, bool PreserveLCSSA,
     unsigned SCEVExpansionBudget, bool RuntimeUnrollMultiExit,
-    Loop **ResultLoop = nullptr);
+    bool AllowLoopRotation, Loop **ResultLoop = nullptr);
 
 LLVM_ABI LoopUnrollResult UnrollAndJamLoop(
     Loop *L, unsigned Count, unsigned TripCount, unsigned TripMultiple,

llvm/lib/Transforms/Utils/LoopUnroll.cpp

@@ -542,7 +542,7 @@ llvm::UnrollLoop(Loop *L, UnrollLoopOptions ULO, LoopInfo *LI,
         L, ULO.Count, ULO.AllowExpensiveTripCount, EpilogProfitability,
         ULO.UnrollRemainder, ULO.ForgetAllSCEV, LI, SE, DT, AC, TTI,
         PreserveLCSSA, ULO.SCEVExpansionBudget, ULO.RuntimeUnrollMultiExit,
-        RemainderLoop);
+        /*AllowLoopRotation*/ true, RemainderLoop);
     LatchBlock = L->getLoopLatch();
     LatchIsExiting = L->isLoopExiting(LatchBlock);
   }

llvm/lib/Transforms/Utils/LoopUnrollAndJam.cpp

@@ -237,12 +237,12 @@ 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, true, SCEVCheapExpansionBudget, EpilogueLoop);
+        LI, SE, DT, AC, TTI, /*PreserveLCSSA*/ true, SCEVCheapExpansionBudget,
+        /*AllowLoopRotation*/ false, EpilogueLoop);
     if (UnrollReminderResult != LoopReminderUnrollResult::Unrolled) {
       LLVM_DEBUG(dbgs() << "Won't unroll-and-jam; remainder loop could not be "
                            "generated when assuming runtime trip count\n");

llvm/lib/Transforms/Utils/LoopUnrollRuntime.cpp

@@ -581,7 +581,7 @@ LoopReminderUnrollResult llvm::UnrollRuntimeLoopRemainder(
     LoopInfo *LI, ScalarEvolution *SE, DominatorTree *DT, AssumptionCache *AC,
     const TargetTransformInfo *TTI, bool PreserveLCSSA,
     unsigned SCEVExpansionBudget, bool RuntimeUnrollMultiExit,
-    Loop **ResultLoop) {
+    bool AllowLoopRotation, Loop **ResultLoop) {
   LLVM_DEBUG(dbgs() << "Trying runtime unrolling on Loop: \n");
   LLVM_DEBUG(L->dump());
   LLVM_DEBUG(UseEpilogRemainder ? dbgs() << "Using epilog remainder.\n"
@@ -590,22 +590,25 @@ LoopReminderUnrollResult llvm::UnrollRuntimeLoopRemainder(
   LoopReminderUnrollResult Result = LoopReminderUnrollResult::Unmodified;
 
   // Rotate loop if it makes the exit count from the latch computable.
-  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, nullptr /*MemorySSAUpdater*/,
-                           SQ, false /*RotationOnly*/, 16 /*Threshold*/,
-                           false /*IsUtilMode*/, false /*PrepareForLTO*/,
-                           [](Loop *, ScalarEvolution *) { return true; }))
-      Result = LoopReminderUnrollResult::Rotated;
+  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.

llvm/unittests/Transforms/Utils/UnrollLoopTest.cpp

@@ -73,6 +73,7 @@ while.end:                                        ; preds = %while.cond
 
   bool ret =
       UnrollRuntimeLoopRemainder(L, 4, true, false, false, false, &LI, &SE, &DT,
-                                 &AC, /*TTI=*/nullptr, PreserveLCSSA, 4, false);
+                                 &AC, /*TTI=*/nullptr, PreserveLCSSA, 4, false,
+                                 false);
   EXPECT_FALSE(ret);
 }