3 files changed, 31 insertions, 17 deletions

diff --git a/llvm/include/llvm/Transforms/Utils/LoopRotationUtils.h b/llvm/include/llvm/Transforms/Utils/LoopRotationUtils.h
index ea4d2cd2ace..2dbeb1149b8 100644
--- a/llvm/include/llvm/Transforms/Utils/LoopRotationUtils.h
+++ b/llvm/include/llvm/Transforms/Utils/LoopRotationUtils.h
@@ -24,11 +24,16 @@ class ScalarEvolution;
 struct SimplifyQuery;
 class TargetTransformInfo;
 
-/// \brief Convert a loop into a loop with bottom test.
-bool LoopRotation(Loop *L, unsigned MaxHeaderSize, LoopInfo *LI,
-                  const TargetTransformInfo *TTI, AssumptionCache *AC,
-                  DominatorTree *DT, ScalarEvolution *SE,
-                  const SimplifyQuery &SQ);
+/// \brief Convert a loop into a loop with bottom test. It may
+/// perform loop latch simplication as well if the flag RotationOnly
+/// is false. The flag Threshold represents the size threshold of the loop
+/// header. If the loop header's size exceeds the threshold, the loop rotation
+/// will give up. The flag IsUtilMode controls the heuristic used in the
+/// LoopRotation. If it is true, the profitability heuristic will be ignored.
+bool LoopRotation(Loop *L, LoopInfo *LI, const TargetTransformInfo *TTI,
+                  AssumptionCache *AC, DominatorTree *DT, ScalarEvolution *SE,
+                  const SimplifyQuery &SQ, bool RotationOnly,
+                  unsigned Threshold, bool IsUtilMode);
 
 } // namespace llvm
 
diff --git a/llvm/lib/Transforms/Scalar/LoopRotation.cpp b/llvm/lib/Transforms/Scalar/LoopRotation.cpp
index 6ea9c2bce4f..eeaad39dc1d 100644
--- a/llvm/lib/Transforms/Scalar/LoopRotation.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopRotation.cpp
@@ -40,8 +40,8 @@ PreservedAnalyses LoopRotatePass::run(Loop &L, LoopAnalysisManager &AM,
   const DataLayout &DL = L.getHeader()->getModule()->getDataLayout();
   const SimplifyQuery SQ = getBestSimplifyQuery(AR, DL);
 
-  bool Changed =
-      LoopRotation(&L, Threshold, &AR.LI, &AR.TTI, &AR.AC, &AR.DT, &AR.SE, SQ);
+  bool Changed = LoopRotation(&L, &AR.LI, &AR.TTI, &AR.AC, &AR.DT, &AR.SE, SQ,
+                              false, Threshold, false);
 
   if (!Changed)
     return PreservedAnalyses::all();
@@ -84,7 +84,8 @@ public:
     auto *SEWP = getAnalysisIfAvailable<ScalarEvolutionWrapperPass>();
     auto *SE = SEWP ? &SEWP->getSE() : nullptr;
     const SimplifyQuery SQ = getBestSimplifyQuery(*this, F);
-    return LoopRotation(L, MaxHeaderSize, LI, TTI, AC, DT, SE, SQ);
+    return LoopRotation(L, LI, TTI, AC, DT, SE, SQ, false, MaxHeaderSize,
+                        false);
   }
 };
 }
diff --git a/llvm/lib/Transforms/Utils/LoopRotationUtils.cpp b/llvm/lib/Transforms/Utils/LoopRotationUtils.cpp
index 529a0ec3525..d96c1014877 100644
--- a/llvm/lib/Transforms/Utils/LoopRotationUtils.cpp
+++ b/llvm/lib/Transforms/Utils/LoopRotationUtils.cpp
@@ -54,13 +54,16 @@ class LoopRotate {
   DominatorTree *DT;
   ScalarEvolution *SE;
   const SimplifyQuery &SQ;
+  bool RotationOnly;
+  bool IsUtilMode;
 
 public:
   LoopRotate(unsigned MaxHeaderSize, LoopInfo *LI,
              const TargetTransformInfo *TTI, AssumptionCache *AC,
-             DominatorTree *DT, ScalarEvolution *SE, const SimplifyQuery &SQ)
+             DominatorTree *DT, ScalarEvolution *SE, const SimplifyQuery &SQ,
+             bool RotationOnly, bool IsUtilMode)
       : MaxHeaderSize(MaxHeaderSize), LI(LI), TTI(TTI), AC(AC), DT(DT), SE(SE),
-        SQ(SQ) {}
+        SQ(SQ), RotationOnly(RotationOnly), IsUtilMode(IsUtilMode) {}
   bool processLoop(Loop *L);
 
 private:
@@ -219,7 +222,7 @@ bool LoopRotate::rotateLoop(Loop *L, bool SimplifiedLatch) {
 
   // Rotate if either the loop latch does *not* exit the loop, or if the loop
   // latch was just simplified. Or if we think it will be profitable.
-  if (L->isLoopExiting(OrigLatch) && !SimplifiedLatch &&
+  if (L->isLoopExiting(OrigLatch) && !SimplifiedLatch && IsUtilMode == false &&
       !shouldRotateLoopExitingLatch(L))
     return false;
 
@@ -604,10 +607,13 @@ bool LoopRotate::processLoop(Loop *L) {
   // Save the loop metadata.
   MDNode *LoopMD = L->getLoopID();
 
+  bool SimplifiedLatch = false;
+
   // Simplify the loop latch before attempting to rotate the header
   // upward. Rotation may not be needed if the loop tail can be folded into the
   // loop exit.
-  bool SimplifiedLatch = simplifyLoopLatch(L);
+  if (!RotationOnly)
+    SimplifiedLatch = simplifyLoopLatch(L);
 
   bool MadeChange = rotateLoop(L, SimplifiedLatch);
   assert((!MadeChange || L->isLoopExiting(L->getLoopLatch())) &&
@@ -623,11 +629,13 @@ bool LoopRotate::processLoop(Loop *L) {
 
 
 /// The utility to convert a loop into a loop with bottom test.
-bool llvm::LoopRotation(Loop *L, unsigned MaxHeaderSize, LoopInfo *LI,
-                        const TargetTransformInfo *TTI, AssumptionCache *AC,
-                        DominatorTree *DT, ScalarEvolution *SE,
-                        const SimplifyQuery &SQ) {
-  LoopRotate LR(MaxHeaderSize, LI, TTI, AC, DT, SE, SQ);
+bool llvm::LoopRotation(Loop *L, LoopInfo *LI, const TargetTransformInfo *TTI,
+                        AssumptionCache *AC, DominatorTree *DT,
+                        ScalarEvolution *SE, const SimplifyQuery &SQ,
+                        bool RotationOnly = true,
+                        unsigned Threshold = unsigned(-1),
+                        bool IsUtilMode = true) {
+  LoopRotate LR(Threshold, LI, TTI, AC, DT, SE, SQ, RotationOnly, IsUtilMode);
 
   return LR.processLoop(L);
 }