4 files changed, 45 insertions, 5 deletions

diff --git a/include/llvm/Analysis/TargetTransformInfo.h b/include/llvm/Analysis/TargetTransformInfo.h
index 06810a7..908612c 100644
--- a/include/llvm/Analysis/TargetTransformInfo.h
+++ b/include/llvm/Analysis/TargetTransformInfo.h
@@ -191,6 +191,20 @@ public:
   /// incurs significant execution cost.
   virtual bool isLoweredToCall(const Function *F) const;
 
+  /// Parameters that control the generic loop unrolling transformation.
+  struct UnrollingPreferences {
+    unsigned Threshold; ///< The cost threshold for the unrolled loop.
+    unsigned OptSizeThreshold; ///< The cost threshold for the unrolled loop
+                               ///< when optimizing for size.
+    bool     Partial;   ///< Allow partial loop unrolling.
+    bool     Runtime;   ///< Perform runtime unrolling.
+  };
+
+  /// \brief Get target-customized preferences for the generic loop unrolling
+  /// transformation. Returns true if the UnrollingPreferences struct has been
+  /// initialized.
+  virtual bool getUnrollingPreferences(UnrollingPreferences &UP) const;
+
   /// @}
 
   /// \name Scalar Target Information
diff --git a/lib/Analysis/TargetTransformInfo.cpp b/lib/Analysis/TargetTransformInfo.cpp
index 8c6005b..b587ff2 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 d5340e6..e1380b7 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 80d060b..f8ff275 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;