3 files changed, 31 insertions, 5 deletions

diff --git a/lib/Analysis/TargetTransformInfo.cpp b/lib/Analysis/TargetTransformInfo.cpp
index 8c6005b31cbe..b587ff27ea67 100644
--- a/lib/Analysis/TargetTransformInfo.cpp
+++ b/lib/Analysis/TargetTransformInfo.cpp
@@ -96,6 +96,11 @@ bool TargetTransformInfo::isLoweredToCall(const Function *F) const {
   return PrevTTI->isLoweredToCall(F);
 }
 
+bool TargetTransformInfo::getUnrollingPreferences(
+                            UnrollingPreferences &UP) const {
+  return PrevTTI->getUnrollingPreferences(UP);
+}
+
 bool TargetTransformInfo::isLegalAddImmediate(int64_t Imm) const {
   return PrevTTI->isLegalAddImmediate(Imm);
 }
@@ -469,6 +474,10 @@ struct NoTTI : ImmutablePass, TargetTransformInfo {
     return true;
   }
 
+  virtual bool getUnrollingPreferences(UnrollingPreferences &) const {
+    return false;
+  }
+
   bool isLegalAddImmediate(int64_t Imm) const {
     return false;
   }
diff --git a/lib/CodeGen/BasicTargetTransformInfo.cpp b/lib/CodeGen/BasicTargetTransformInfo.cpp
index d5340e60231e..e1380b73e9ed 100644
--- a/lib/CodeGen/BasicTargetTransformInfo.cpp
+++ b/lib/CodeGen/BasicTargetTransformInfo.cpp
@@ -84,6 +84,7 @@ public:
   virtual unsigned getJumpBufSize() const;
   virtual bool shouldBuildLookupTables() const;
   virtual bool haveFastSqrt(Type *Ty) const;
+  virtual bool getUnrollingPreferences(UnrollingPreferences &UP) const;
 
   /// @}
 
@@ -189,6 +190,10 @@ bool BasicTTI::haveFastSqrt(Type *Ty) const {
   return TLI->isTypeLegal(VT) && TLI->isOperationLegalOrCustom(ISD::FSQRT, VT);
 }
 
+bool BasicTTI::getUnrollingPreferences(UnrollingPreferences &) const {
+  return false;
+}
+
 //===----------------------------------------------------------------------===//
 //
 // Calls used by the vectorizers.
diff --git a/lib/Transforms/Scalar/LoopUnrollPass.cpp b/lib/Transforms/Scalar/LoopUnrollPass.cpp
index 80d060b926ea..f8ff275e9699 100644
--- a/lib/Transforms/Scalar/LoopUnrollPass.cpp
+++ b/lib/Transforms/Scalar/LoopUnrollPass.cpp
@@ -55,6 +55,8 @@ namespace {
       CurrentAllowPartial = (P == -1) ? UnrollAllowPartial : (bool)P;
 
       UserThreshold = (T != -1) || (UnrollThreshold.getNumOccurrences() > 0);
+      UserAllowPartial = (P != -1) ||
+                         (UnrollAllowPartial.getNumOccurrences() > 0);
 
       initializeLoopUnrollPass(*PassRegistry::getPassRegistry());
     }
@@ -76,6 +78,7 @@ namespace {
     unsigned CurrentThreshold;
     bool     CurrentAllowPartial;
     bool     UserThreshold;        // CurrentThreshold is user-specified.
+    bool     UserAllowPartial;     // CurrentAllowPartial is user-specified.
 
     bool runOnLoop(Loop *L, LPPassManager &LPM);
 
@@ -145,16 +148,20 @@ bool LoopUnroll::runOnLoop(Loop *L, LPPassManager &LPM) {
         << "] Loop %" << Header->getName() << "\n");
   (void)Header;
 
+  TargetTransformInfo::UnrollingPreferences UP;
+  bool HasUP = TTI.getUnrollingPreferences(UP);
+
   // Determine the current unrolling threshold.  While this is normally set
   // from UnrollThreshold, it is overridden to a smaller value if the current
   // function is marked as optimize-for-size, and the unroll threshold was
   // not user specified.
-  unsigned Threshold = CurrentThreshold;
+  unsigned Threshold = (HasUP && !UserThreshold) ? UP.Threshold :
+                                                   CurrentThreshold;
   if (!UserThreshold &&
       Header->getParent()->getAttributes().
         hasAttribute(AttributeSet::FunctionIndex,
                      Attribute::OptimizeForSize))
-    Threshold = OptSizeUnrollThreshold;
+    Threshold = HasUP ? UP.OptSizeThreshold : OptSizeUnrollThreshold;
 
   // Find trip count and trip multiple if count is not available
   unsigned TripCount = 0;
@@ -184,6 +191,9 @@ bool LoopUnroll::runOnLoop(Loop *L, LPPassManager &LPM) {
     Count = TripCount;
   }
 
+  bool Runtime = (HasUP && UnrollRuntime.getNumOccurrences() == 0) ?
+                 UP.Runtime : UnrollRuntime;
+
   // Enforce the threshold.
   if (Threshold != NoThreshold) {
     unsigned NumInlineCandidates;
@@ -204,7 +214,9 @@ bool LoopUnroll::runOnLoop(Loop *L, LPPassManager &LPM) {
     if (TripCount != 1 && Size > Threshold) {
       DEBUG(dbgs() << "  Too large to fully unroll with count: " << Count
             << " because size: " << Size << ">" << Threshold << "\n");
-      if (!CurrentAllowPartial && !(UnrollRuntime && TripCount == 0)) {
+      bool AllowPartial = (HasUP && !UserAllowPartial) ? UP.Partial :
+                                                         CurrentAllowPartial;
+      if (!AllowPartial && !(Runtime && TripCount == 0)) {
         DEBUG(dbgs() << "  will not try to unroll partially because "
               << "-unroll-allow-partial not given\n");
         return false;
@@ -215,7 +227,7 @@ bool LoopUnroll::runOnLoop(Loop *L, LPPassManager &LPM) {
         while (Count != 0 && TripCount%Count != 0)
           Count--;
       }
-      else if (UnrollRuntime) {
+      else if (Runtime) {
         // Reduce unroll count to be a lower power-of-two value
         while (Count != 0 && Size > Threshold) {
           Count >>= 1;
@@ -231,7 +243,7 @@ bool LoopUnroll::runOnLoop(Loop *L, LPPassManager &LPM) {
   }
 
   // Unroll the loop.
-  if (!UnrollLoop(L, Count, TripCount, UnrollRuntime, TripMultiple, LI, &LPM))
+  if (!UnrollLoop(L, Count, TripCount, Runtime, TripMultiple, LI, &LPM))
     return false;
 
   return true;