1 files changed, 26 insertions, 9 deletions

diff --git a/lib/Transforms/Scalar/LoopUnrollPass.cpp b/lib/Transforms/Scalar/LoopUnrollPass.cpp
index 08ac38dec5..ecd350b258 100644
--- a/lib/Transforms/Scalar/LoopUnrollPass.cpp
+++ b/lib/Transforms/Scalar/LoopUnrollPass.cpp
@@ -19,6 +19,7 @@
 #include "llvm/Analysis/ScalarEvolution.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/IR/DataLayout.h"
+#include "llvm/IR/Dominators.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
@@ -86,12 +87,12 @@ namespace {
     bool     UserAllowPartial;     // CurrentAllowPartial is user-specified.
     bool     UserRuntime;          // CurrentRuntime is user-specified.
 
-    bool runOnLoop(Loop *L, LPPassManager &LPM);
+    bool runOnLoop(Loop *L, LPPassManager &LPM) override;
 
     /// This transformation requires natural loop information & requires that
     /// loop preheaders be inserted into the CFG...
     ///
-    virtual void getAnalysisUsage(AnalysisUsage &AU) const {
+    void getAnalysisUsage(AnalysisUsage &AU) const override {
       AU.addRequired<LoopInfo>();
       AU.addPreserved<LoopInfo>();
       AU.addRequiredID(LoopSimplifyID);
@@ -105,7 +106,7 @@ namespace {
       // If loop unroll does not preserve dom info then LCSSA pass on next
       // loop will receive invalid dom info.
       // For now, recreate dom info, if loop is unrolled.
-      AU.addPreserved<DominatorTree>();
+      AU.addPreserved<DominatorTreeWrapperPass>();
     }
   };
 }
@@ -124,6 +125,10 @@ Pass *llvm::createLoopUnrollPass(int Threshold, int Count, int AllowPartial,
   return new LoopUnroll(Threshold, Count, AllowPartial, Runtime);
 }
 
+Pass *llvm::createSimpleLoopUnrollPass() {
+  return llvm::createLoopUnrollPass(-1, -1, 0, 0);
+}
+
 /// ApproximateLoopSize - Approximate the size of the loop.
 static unsigned ApproximateLoopSize(const Loop *L, unsigned &NumCalls,
                                     bool &NotDuplicatable,
@@ -146,6 +151,9 @@ static unsigned ApproximateLoopSize(const Loop *L, unsigned &NumCalls,
 }
 
 bool LoopUnroll::runOnLoop(Loop *L, LPPassManager &LPM) {
+  if (skipOptnoneFunction(L))
+    return false;
+
   LoopInfo *LI = &getAnalysis<LoopInfo>();
   ScalarEvolution *SE = &getAnalysis<ScalarEvolution>();
   const TargetTransformInfo &TTI = getAnalysis<TargetTransformInfo>();
@@ -158,7 +166,10 @@ bool LoopUnroll::runOnLoop(Loop *L, LPPassManager &LPM) {
   TargetTransformInfo::UnrollingPreferences UP;
   UP.Threshold = CurrentThreshold;
   UP.OptSizeThreshold = OptSizeUnrollThreshold;
+  UP.PartialThreshold = CurrentThreshold;
+  UP.PartialOptSizeThreshold = OptSizeUnrollThreshold;
   UP.Count = CurrentCount;
+  UP.MaxCount = UINT_MAX;
   UP.Partial = CurrentAllowPartial;
   UP.Runtime = CurrentRuntime;
   TTI.getUnrollingPreferences(L, UP);
@@ -168,11 +179,15 @@ bool LoopUnroll::runOnLoop(Loop *L, LPPassManager &LPM) {
   // function is marked as optimize-for-size, and the unroll threshold was
   // not user specified.
   unsigned Threshold = UserThreshold ? CurrentThreshold : UP.Threshold;
+  unsigned PartialThreshold =
+    UserThreshold ? CurrentThreshold : UP.PartialThreshold;
   if (!UserThreshold &&
       Header->getParent()->getAttributes().
         hasAttribute(AttributeSet::FunctionIndex,
-                     Attribute::OptimizeForSize))
+                     Attribute::OptimizeForSize)) {
     Threshold = UP.OptSizeThreshold;
+    PartialThreshold = UP.PartialOptSizeThreshold;
+  }
 
   // Find trip count and trip multiple if count is not available
   unsigned TripCount = 0;
@@ -206,14 +221,14 @@ bool LoopUnroll::runOnLoop(Loop *L, LPPassManager &LPM) {
   }
 
   // Enforce the threshold.
-  if (Threshold != NoThreshold) {
+  if (Threshold != NoThreshold && PartialThreshold != NoThreshold) {
     unsigned NumInlineCandidates;
     bool notDuplicatable;
     unsigned LoopSize = ApproximateLoopSize(L, NumInlineCandidates,
                                             notDuplicatable, TTI);
     DEBUG(dbgs() << "  Loop Size = " << LoopSize << "\n");
     if (notDuplicatable) {
-      DEBUG(dbgs() << "  Not unrolling loop which contains non duplicatable"
+      DEBUG(dbgs() << "  Not unrolling loop which contains non-duplicatable"
             << " instructions.\n");
       return false;
     }
@@ -233,17 +248,19 @@ bool LoopUnroll::runOnLoop(Loop *L, LPPassManager &LPM) {
       }
       if (TripCount) {
         // Reduce unroll count to be modulo of TripCount for partial unrolling
-        Count = Threshold / LoopSize;
+        Count = PartialThreshold / LoopSize;
         while (Count != 0 && TripCount%Count != 0)
           Count--;
       }
       else if (Runtime) {
         // Reduce unroll count to be a lower power-of-two value
-        while (Count != 0 && Size > Threshold) {
+        while (Count != 0 && Size > PartialThreshold) {
           Count >>= 1;
           Size = LoopSize*Count;
         }
       }
+      if (Count > UP.MaxCount)
+        Count = UP.MaxCount;
       if (Count < 2) {
         DEBUG(dbgs() << "  could not unroll partially\n");
         return false;
@@ -253,7 +270,7 @@ bool LoopUnroll::runOnLoop(Loop *L, LPPassManager &LPM) {
   }
 
   // Unroll the loop.
-  if (!UnrollLoop(L, Count, TripCount, Runtime, TripMultiple, LI, &LPM))
+  if (!UnrollLoop(L, Count, TripCount, Runtime, TripMultiple, LI, this, &LPM))
     return false;
 
   return true;