[LoopUnroll] Use the upper bound of the loop trip count to completely unroll loops

[haicheng]

Summary

This patch tries to fully unroll loops having break statement like this

for (int i = 0; i < 8; i++) {                                                                     
    if (a[i] == value) {                                                                               
        found = true; 
        break;
    }                                                                                             
}

GCC can fully unroll such loops, but currently LLVM cannot because LLVM only supports loops having exact constant trip counts.

The upper bound of the trip count can be obtained from calling ScalarEvolution::getMaxBackedgeTakenCount(). Part of the patch is the refactoring work in SCEV to prevent duplicating code.

The feature of using the upper bound is enabled under the same circumstance when runtime unrolling is enabled since both are used to unroll loops without knowing the exact constant trip count.

The modified test/CodeGen/AMDGPU/tti-unroll-prefs.ll can be used as the test case of this patch.

[mzolotukhin]

Hi,

This makes total sense to me. But I think there is no need in introducing yet another knob in loop-unroller: have you measured the effect of the change on benchmarks (LLVM testsuite/SPEC/others)? I think it should be correct to unroll a loop with a small trip count even if it's not exactly known: the effect of this transformation would be not worse then from unrolling a similar loop with a constant, but equal to upper bound, trip-count. The benchmark results would help confirming or rejecting this assumption.

Thanks,
Michael

lib/Transforms/Scalar/LoopUnrollPass.cpp

934 ↗

(On Diff #72006)

Unnecessary change.

[ABataev]

lib/Analysis/ScalarEvolution.cpp

5349 ↗

(On Diff #72006)

const auto *MaxExitCount

lib/Transforms/Utils/LoopUnroll.cpp

543 ↗

(On Diff #72006)

Maybe just NeedConditional = (UseUpperBound && j)?

[haicheng]

In D24790#550261, @mzolotukhin wrote:

Hi,

This makes total sense to me. But I think there is no need in introducing yet another knob in loop-unroller: have you measured the effect of the change on benchmarks (LLVM testsuite/SPEC/others)? I think it should be correct to unroll a loop with a small trip count even if it's not exactly known: the effect of this transformation would be not worse then from unrolling a similar loop with a constant, but equal to upper bound, trip-count. The benchmark results would help confirming or rejecting this assumption.

Thanks,
Michael

Thank you for reviewing my patch, Michael.

My initial design was to completely unroll all loops having small trip count upper bounds whenever the exact trip counts are unknown, but I saw several regressions in spec200x and internal benchmarks (e.g. spec2000/eon -1.8%, spec2000/gcc -1.2%) running on a AArch64 device. One of the major reason was that I unrolled many loops with calls. As you may already know, the cost model of call is not that awesome. BasicTTIImplBase::getUnrollingPreferences() can help me check the call IRs in the loop and I need a boolean anyway to pass the check result so that I create a new entry in UnrollingPreferences.

If using exact trip count to unroll, the unrolled loop usually becomes a giant basic block which is preferable. However, if using the upper bound to unroll, the unrolled loop usually become a sequence of small basic blocks because it is not safe to merge loop blocks belonging to different iterations. Some of these blocks may not be executed during runtime. This is another reason that I think we may need to be more conservative to use upper bound to unroll loops.

I tried several different configurations and the patch I uploaded was the best I found. No noticeable regressions in the benchmarks I tested and there are several small wins in spec200x (e.g. spec2000/perlbmk +0.8%, spec2006/xalancbmk +0.7%) and many larger wins in our internal benchmarks which are much larger than spec2006.

Haicheng

[haicheng]

Address the comments from Michael and Alexey. Thank you very much.

[mzolotukhin]

Hi Haicheng,

Thanks for working on this, please find my answers below and some more remarks/nit-picks inline.

One of the major reason was that I unrolled many loops with calls. As you may already know, the cost model of call is not that awesome.

Maybe we just fix the cost model for calls instead :-) But I know, it might not be actually possible at all.

If using exact trip count to unroll, the unrolled loop usually becomes a giant basic block which is preferable. However, if using the upper bound to unroll, the unrolled loop usually become a sequence of small basic blocks because it is not safe to merge loop blocks belonging to different iterations. Some of these blocks may not be executed during runtime. This is another reason that I think we may need to be more conservative to use upper bound to unroll loops.

I see, this makes perfect sense to me. Indeed, having separate thresholds might be reasonable.

I tried several different configurations and the patch I uploaded was the best I found.

Have you tested it on any other architecture except AArch64?

Thanks,
Michael

lib/Analysis/ScalarEvolution.cpp

5316 ↗

(On Diff #72293)

The first letter shouldn't be capitalized. I think the routine name also might be improved, but I don't have any better suggestions at the moment.

lib/Transforms/Scalar/LoopUnrollPass.cpp

766 ↗

(On Diff #72293)

s/make/makes/

767 ↗

(On Diff #72293)

s/staticaly/statically/

771 ↗

(On Diff #72293)

This gets confusing. Can we somehow rename these variables so that it's clearer what's the difference between them?

779 ↗

(On Diff #72293)

What if TripCount matches MaxTripCount? Will we think that we're using upper bound instead of the exact trip count in this case?

1007 ↗

(On Diff #72293)

I'm not a big fan of reusing this variable. Initially it was supposed to come from UP and show if we're allowed to use upper-bound instead of exactly-known trip-count, but now we're also using it to communicate between various routines (and the interface was not specified). Can it be refactored somehow?

lib/Transforms/Utils/LoopUnroll.cpp

206 ↗

(On Diff #72293)

Please add a comment about UseUpperBound. Actually, I don't think that's a good name for this argument, because this functions doesn't use upper bound for anything. What it needs is a flag indicating that conditional branches must be preserved - I'd suggest to reflect that in the argument name.

543 ↗

(On Diff #72293)

Is it intentional that we can make NeedConditional true after it was set to false before that? From semantics it looks like this never happens (in this case we can assert it), but I'd like to make sure it was not overlooked. Did you intend to replace this if with else if as well?

test/CodeGen/AMDGPU/tti-unroll-prefs.ll

14–16 ↗

(On Diff #72293)

We could've just passed -unroll-threshold to overcome this. It might make sense to do it even now, so that the test doesn't break if UnrollPreferences are changed.

[haicheng]

Address Michael's comments. Thank you.

[haicheng]

lib/Transforms/Scalar/LoopUnrollPass.cpp

779 ↗

(On Diff #72293)

TripCount and MaxTripCount cannot both be non zero because MaxTripCount is computed only if TripCount is Zero. If one is non zero, the other one must be zero. If they are both zero, FullUnrollTripCount is zero and then we cannot enter here.

lib/Transforms/Utils/LoopUnroll.cpp

543 ↗

(On Diff #72293)

I intended to replace if with else if. Using complete unroll should not enter the else if part. I think the change improves the readability.

test/CodeGen/AMDGPU/tti-unroll-prefs.ll

14–16 ↗

(On Diff #72293)

This loop has a break statement so that we cannot get the exact trip cont. I think we cannot fully unroll the loop without my change. Passing -unroll-threshold cannot overcome this.

[haicheng]

lib/Transforms/Scalar/LoopUnrollPass.cpp

779 ↗

(On Diff #72293)

I also added a sentence in the comment to describe this.

lib/Transforms/Utils/LoopUnroll.cpp

543 ↗

(On Diff #72293)

Without the change of else if, the logic of my patch is also wrong.

[haicheng]

Hi Michael,

Please see my inlined response. Thank you.

In D24790#554730, @mzolotukhin wrote:

Hi Haicheng,

Thanks for working on this, please find my answers below and some more remarks/nit-picks inline.

One of the major reason was that I unrolled many loops with calls. As you may already know, the cost model of call is not that awesome.

Maybe we just fix the cost model for calls instead :-) But I know, it might not be actually possible at all.

It should be fixed, but I have no clue how to do that now. Can I leave it for the future?

If using exact trip count to unroll, the unrolled loop usually becomes a giant basic block which is preferable. However, if using the upper bound to unroll, the unrolled loop usually become a sequence of small basic blocks because it is not safe to merge loop blocks belonging to different iterations. Some of these blocks may not be executed during runtime. This is another reason that I think we may need to be more conservative to use upper bound to unroll loops.

I see, this makes perfect sense to me. Indeed, having separate thresholds might be reasonable.

I tried threshold 100 for using upper bound to unroll instead of the current default threshold 150, but I saw some regressions.

I have some other random thoughts about unroll threshold that not directly related to this patch and not tested with any benchmarks yet. I think we may want to encourage the unroller to unroll loops with larger trip count to reduce more loop overhead. For example, if we have a loop whose size is 75 and can be unrolled twice and we have another loop whose size is 21 but can be unrolled 8 times, we may prefer unrolling the latter more. It occurs to me because I noticed that GCC unrolls more small loops with high trip count than LLVM does and LLVM unrolls more large loops with small trip count than GCC does. Maybe we can use some equations like 100+10*trip_count to calculate the threshold. What do you think?

I tried several different configurations and the patch I uploaded was the best I found.

Have you tested it on any other architecture except AArch64?

I ran spec2000/2006 on x86 last week, no noticeable regressions. Some small improvements: spec2000/bzip2 +0.7, spec2006/dealII +2.1% gcc +0.8%.

Thanks,
Michael

[haicheng]

Hi Michael,

Would you please take another look at my modified patch? Thank you in advance.

Haicheng

[mzolotukhin]

Hi Haicheng,

Please find a couple of minor comments inline.

Thanks,
Michael

lib/Transforms/Scalar/LoopUnrollPass.cpp

779 ↗

(On Diff #72293)

Could you add an assert for this please?

95–102 ↗

(On Diff #73277)

Do we really need two options here? Can the first one be replaced with (UnrollMaxUpperBound == 0)?

[haicheng]

Address Michael's comments. Thank you.

[mzolotukhin]

Hi Haicheng,

The patch looks mostly good to me modulo a couple of small remarks, please feel free to commit once they are addressed.

Thanks,
Michael

lib/Transforms/Scalar/LoopUnrollPass.cpp

768–769 ↗

(On Diff #74193)

The assert might fail to catch a bug if the multiplication result overflows and wraps around to 0.

test/Transforms/LoopUnroll/AArch64/full-unroll-trip-count-upper-bound.ll

1–2 ↗

(On Diff #74193)

Shouldn't we have -unroll-upper here? Actually, I think we need to check that a) a loop is not unrolled without this flag, and b) it's unrolled with it.

Also, what happened to the changes in the other (existing) test?

[haicheng]

Addressed more comments from Michael. I will commit tomorrow morning if no new comments coming today.

[evstupac]

llvm/trunk/lib/Transforms/Scalar/LoopUnrollPass.cpp

1027

Sorry for comment after approve.
What I see here is that TripCount could implicitly (as TripCount was passed to computeUnrollCount as a reference) become MaxTripCount. Here it is passed to UnrollLoop as TripCount.
Could you please add a comment to UnrollLoop that TripCount parameter is not always a real loop trip count, but sometimes MaxTripCount. Or rename TripCount to MaxTripCount in UnrollLoop parameters. There could be future errors if someone assumes that TripCount is exact loop trip count.

[haicheng]

llvm/trunk/lib/Transforms/Scalar/LoopUnrollPass.cpp

1027

I will modify the comments in a separate patch to make it clearer. However, you may need to know that even the trip counter is calculated by SCEV::getSmallConstantTripCount(), it may not be the exact times the loop header get executed. Please see the comments before llvm::UnrollLoop() in LoopUnroll.cpp.

llvm/trunk/include/llvm/Analysis/ScalarEvolution.h

public:
   /// value. Returns 0 if the trip count is unknown or not constant. This
   /// "trip count" assumes that control exits via ExitingBlock. More
   /// precisely, it is the number of times that control may reach ExitingBlock
   /// before taking the branch. For loops with multiple exits, it may not be
   /// the number times that the loop header executes if the loop exits
   /// prematurely via another branch.
   unsigned getSmallConstantTripCount(Loop *L, BasicBlock *ExitingBlock);
 
+  /// Returns the upper bound of the loop trip count as a normal unsigned
+  /// value.
+  /// Returns 0 if the trip count is unknown or not constant.
+  unsigned getSmallConstantMaxTripCount(Loop *L);
+
   /// Returns the largest constant divisor of the trip count of the
   /// loop if it is a single-exit loop and we can compute a small maximum for
   /// that loop.
   ///
   /// Implemented in terms of the \c getSmallConstantTripMultiple overload with
   /// the single exiting block passed to it. See that routine for details.
   unsigned getSmallConstantTripMultiple(Loop *L);
 

llvm/trunk/include/llvm/Analysis/TargetTransformInfo.h

struct UnrollingPreferences {
     /// Allow generation of a loop remainder (extra iterations after unroll).
     bool AllowRemainder;
     /// Allow emitting expensive instructions (such as divisions) when computing
     /// the trip count of a loop for runtime unrolling.
     bool AllowExpensiveTripCount;
     /// Apply loop unroll on any kind of loop
     /// (mainly to loops that fail runtime unrolling).
     bool Force;
+    /// Allow using trip count upper bound to unroll loops.
+    bool UpperBound;
   };
 
   /// \brief Get target-customized preferences for the generic loop unrolling
   /// transformation. The caller will initialize UP with the current
   /// target-independent defaults.
   void getUnrollingPreferences(Loop *L, UnrollingPreferences &UP) const;
 
   /// @}

llvm/trunk/include/llvm/CodeGen/BasicTTIImpl.h

for (Loop::block_iterator I = L->block_begin(), E = L->block_end(); I != E;
               continue;
           }
 
           return;
         }
     }
 
     // Enable runtime and partial unrolling up to the specified size.
-    UP.Partial = UP.Runtime = true;
+    // Enable using trip count upper bound to unroll loops.
+    UP.Partial = UP.Runtime = UP.UpperBound = true;
     UP.PartialThreshold = MaxOps;
 
     // Avoid unrolling when optimizing for size.
     UP.OptSizeThreshold = 0;
     UP.PartialOptSizeThreshold = 0;
   }
 
   /// @}

llvm/trunk/include/llvm/Transforms/Scalar.h

 //
 Pass *createLoopInstSimplifyPass();
 
 //===----------------------------------------------------------------------===//
 //
 // LoopUnroll - This pass is a simple loop unrolling pass.
 //
 Pass *createLoopUnrollPass(int Threshold = -1, int Count = -1,
-                           int AllowPartial = -1, int Runtime = -1);
-// Create an unrolling pass for full unrolling only.
+                           int AllowPartial = -1, int Runtime = -1,
+                           int UpperBound = -1);
+// Create an unrolling pass for full unrolling that uses exact trip count only.
 Pass *createSimpleLoopUnrollPass();
 
 //===----------------------------------------------------------------------===//
 //
 // LoopReroll - This pass is a simple loop rerolling pass.
 //
 Pass *createLoopRerollPass();
 

llvm/trunk/include/llvm/Transforms/Scalar/LoopUnrollPass.h

 
 namespace llvm {
 
 struct LoopUnrollPass : public PassInfoMixin<LoopUnrollPass> {
   Optional<unsigned> ProvidedCount;
   Optional<unsigned> ProvidedThreshold;
   Optional<bool> ProvidedAllowPartial;
   Optional<bool> ProvidedRuntime;
+  Optional<bool> ProvidedUpperBound;
 
   PreservedAnalyses run(Loop &L, LoopAnalysisManager &AM);
 };
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_SCALAR_LOOPUNROLLPASS_H

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

 class LPPassManager;
 class MDNode;
 class Pass;
 class OptimizationRemarkEmitter;
 class ScalarEvolution;
 
 bool UnrollLoop(Loop *L, unsigned Count, unsigned TripCount, bool Force,
                 bool AllowRuntime, bool AllowExpensiveTripCount,
-                unsigned TripMultiple, LoopInfo *LI, ScalarEvolution *SE,
-                DominatorTree *DT, AssumptionCache *AC,
+                bool UseUpperBound, unsigned TripMultiple, LoopInfo *LI,
+                ScalarEvolution *SE, DominatorTree *DT, AssumptionCache *AC,
                 OptimizationRemarkEmitter *ORE, bool PreserveLCSSA);
 
 bool UnrollRuntimeLoopRemainder(Loop *L, unsigned Count,
                                 bool AllowExpensiveTripCount,
                                 bool UseEpilogRemainder, LoopInfo *LI,
                                 ScalarEvolution *SE, DominatorTree *DT,
                                 bool PreserveLCSSA);
 
 MDNode *GetUnrollMetadata(MDNode *LoopID, StringRef Name);
 }
 
 #endif

llvm/trunk/lib/Analysis/ScalarEvolution.cpp

 }
 
 
 
 //===----------------------------------------------------------------------===//
 //                   Iteration Count Computation Code
 //
 
+static unsigned getConstantTripCount(const SCEVConstant *ExitCount) {
+  if (!ExitCount)
+    return 0;
+
+  ConstantInt *ExitConst = ExitCount->getValue();
+
+  // Guard against huge trip counts.
+  if (ExitConst->getValue().getActiveBits() > 32)
+    return 0;
+
+  // In case of integer overflow, this returns 0, which is correct.
+  return ((unsigned)ExitConst->getZExtValue()) + 1;
+}
+
 unsigned ScalarEvolution::getSmallConstantTripCount(Loop *L) {
   if (BasicBlock *ExitingBB = L->getExitingBlock())
     return getSmallConstantTripCount(L, ExitingBB);
 
   // No trip count information for multiple exits.
   return 0;
 }
 
 unsigned ScalarEvolution::getSmallConstantTripCount(Loop *L,
                                                     BasicBlock *ExitingBlock) {
   assert(ExitingBlock && "Must pass a non-null exiting block!");
   assert(L->isLoopExiting(ExitingBlock) &&
          "Exiting block must actually branch out of the loop!");
   const SCEVConstant *ExitCount =
       dyn_cast<SCEVConstant>(getExitCount(L, ExitingBlock));
-  if (!ExitCount)
-    return 0;
-
-  ConstantInt *ExitConst = ExitCount->getValue();
-
-  // Guard against huge trip counts.
-  if (ExitConst->getValue().getActiveBits() > 32)
-    return 0;
+  return getConstantTripCount(ExitCount);
+}
 
-  // In case of integer overflow, this returns 0, which is correct.
-  return ((unsigned)ExitConst->getZExtValue()) + 1;
+unsigned ScalarEvolution::getSmallConstantMaxTripCount(Loop *L) {
+  const auto *MaxExitCount =
+      dyn_cast<SCEVConstant>(getMaxBackedgeTakenCount(L));
+  return getConstantTripCount(MaxExitCount);
 }
 
 unsigned ScalarEvolution::getSmallConstantTripMultiple(Loop *L) {
   if (BasicBlock *ExitingBB = L->getExitingBlock())
     return getSmallConstantTripMultiple(L, ExitingBB);
 
   // No trip multiple information for multiple exits.
   return 0;

llvm/trunk/lib/Transforms/Scalar/LoopUnrollPass.cpp

static cl::opt<bool> UnrollAllowRemainder(
     "unroll-allow-remainder", cl::Hidden,
     cl::desc("Allow generation of a loop remainder (extra iterations) "
              "when unrolling a loop."));
 
 static cl::opt<bool>
     UnrollRuntime("unroll-runtime", cl::ZeroOrMore, cl::Hidden,
                   cl::desc("Unroll loops with run-time trip counts"));
 
+static cl::opt<unsigned> UnrollMaxUpperBound(
+    "unroll-max-upperbound", cl::init(8), cl::Hidden,
+    cl::desc(
+        "The max of trip count upper bound that is considered in unrolling"));
+
 static cl::opt<unsigned> PragmaUnrollThreshold(
     "pragma-unroll-threshold", cl::init(16 * 1024), cl::Hidden,
     cl::desc("Unrolled size limit for loops with an unroll(full) or "
              "unroll_count pragma."));
 
 /// A magic value for use with the Threshold parameter to indicate
 /// that the loop unroll should be performed regardless of how much
 /// code expansion would result.
 static const unsigned NoThreshold = UINT_MAX;
 
 /// Gather the various unrolling parameters based on the defaults, compiler
 /// flags, TTI overrides and user specified parameters.
 static TargetTransformInfo::UnrollingPreferences gatherUnrollingPreferences(
     Loop *L, const TargetTransformInfo &TTI, Optional<unsigned> UserThreshold,
     Optional<unsigned> UserCount, Optional<bool> UserAllowPartial,
-    Optional<bool> UserRuntime) {
+    Optional<bool> UserRuntime, Optional<bool> UserUpperBound) {
   TargetTransformInfo::UnrollingPreferences UP;
 
   // Set up the defaults
   UP.Threshold = 150;
   UP.PercentDynamicCostSavedThreshold = 50;
   UP.DynamicCostSavingsDiscount = 100;
   UP.OptSizeThreshold = 0;
   UP.PartialThreshold = UP.Threshold;
   UP.PartialOptSizeThreshold = 0;
   UP.Count = 0;
   UP.DefaultUnrollRuntimeCount = 8;
   UP.MaxCount = UINT_MAX;
   UP.FullUnrollMaxCount = UINT_MAX;
   UP.Partial = false;
   UP.Runtime = false;
   UP.AllowRemainder = true;
   UP.AllowExpensiveTripCount = false;
   UP.Force = false;
+  UP.UpperBound = false;
 
   // Override with any target specific settings
   TTI.getUnrollingPreferences(L, UP);
 
   // Apply size attributes
   if (L->getHeader()->getParent()->optForSize()) {
     UP.Threshold = UP.OptSizeThreshold;
     UP.PartialThreshold = UP.PartialOptSizeThreshold;
static TargetTransformInfo::UnrollingPreferences gatherUnrollingPreferences(
   if (UnrollFullMaxCount.getNumOccurrences() > 0)
     UP.FullUnrollMaxCount = UnrollFullMaxCount;
   if (UnrollAllowPartial.getNumOccurrences() > 0)
     UP.Partial = UnrollAllowPartial;
   if (UnrollAllowRemainder.getNumOccurrences() > 0)
     UP.AllowRemainder = UnrollAllowRemainder;
   if (UnrollRuntime.getNumOccurrences() > 0)
     UP.Runtime = UnrollRuntime;
+  if (UnrollMaxUpperBound == 0)
+    UP.UpperBound = false;
 
   // Apply user values provided by argument
   if (UserThreshold.hasValue()) {
     UP.Threshold = *UserThreshold;
     UP.PartialThreshold = *UserThreshold;
   }
   if (UserCount.hasValue())
     UP.Count = *UserCount;
   if (UserAllowPartial.hasValue())
     UP.Partial = *UserAllowPartial;
   if (UserRuntime.hasValue())
     UP.Runtime = *UserRuntime;
+  if (UserUpperBound.hasValue())
+    UP.UpperBound = *UserUpperBound;
 
   return UP;
 }
 
 namespace {
 /// A struct to densely store the state of an instruction after unrolling at
 /// each iteration.
 ///
static bool canUnrollCompletely(Loop *L, unsigned Threshold,
   DEBUG(dbgs() << "    Rolled dynamic cost: " << RolledDynamicCost << "\n");
   DEBUG(dbgs() << "    Percent cost saved: " << PercentDynamicCostSaved
                << "\n");
   return false;
 }
 
 // Returns true if unroll count was set explicitly.
 // Calculates unroll count and writes it to UP.Count.
-static bool computeUnrollCount(Loop *L, const TargetTransformInfo &TTI,
-                               DominatorTree &DT, LoopInfo *LI,
-                               ScalarEvolution *SE,
-                               OptimizationRemarkEmitter *ORE,
-                               unsigned TripCount, unsigned TripMultiple,
-                               unsigned LoopSize,
-                               TargetTransformInfo::UnrollingPreferences &UP) {
+static bool computeUnrollCount(
+    Loop *L, const TargetTransformInfo &TTI, DominatorTree &DT, LoopInfo *LI,
+    ScalarEvolution *SE, OptimizationRemarkEmitter *ORE, unsigned &TripCount,
+    unsigned MaxTripCount, unsigned &TripMultiple, unsigned LoopSize,
+    TargetTransformInfo::UnrollingPreferences &UP, bool &UseUpperBound) {
   // BEInsns represents number of instructions optimized when "back edge"
   // becomes "fall through" in unrolled loop.
   // For now we count a conditional branch on a backedge and a comparison
   // feeding it.
   unsigned BEInsns = 2;
   // Check for explicit Count.
   // 1st priority is unroll count set by "unroll-count" option.
   bool UserUnrollCount = UnrollCount.getNumOccurrences() > 0;
if (ExplicitUnroll && TripCount != 0) {
     // unrolling limits. Set thresholds to at least the PragmaThreshold value
     // which is larger than the default limits.
     UP.Threshold = std::max<unsigned>(UP.Threshold, PragmaUnrollThreshold);
     UP.PartialThreshold =
         std::max<unsigned>(UP.PartialThreshold, PragmaUnrollThreshold);
   }
 
   // 3rd priority is full unroll count.
-  // Full unroll make sense only when TripCount could be staticaly calculated.
+  // Full unroll makes sense only when TripCount or its upper bound could be
+  // statically calculated.
   // Also we need to check if we exceed FullUnrollMaxCount.
-  if (TripCount && TripCount <= UP.FullUnrollMaxCount) {
+  // If using the upper bound to unroll, TripMultiple should be set to 1 because
+  // we do not know when loop may exit.
+  // MaxTripCount and ExactTripCount cannot both be non zero since we only
+  // compute the former when the latter is zero.
+  unsigned ExactTripCount = TripCount;
+  assert((ExactTripCount == 0 || MaxTripCount == 0) &&
+         "ExtractTripCound and MaxTripCount cannot both be non zero.");
+  unsigned FullUnrollTripCount = ExactTripCount ? ExactTripCount : MaxTripCount;
+  if (FullUnrollTripCount && FullUnrollTripCount <= UP.FullUnrollMaxCount) {
     // When computing the unrolled size, note that BEInsns are not replicated
     // like the rest of the loop body.
-    UnrolledSize = (uint64_t)(LoopSize - BEInsns) * TripCount + BEInsns;
+    UnrolledSize =
+        (uint64_t)(LoopSize - BEInsns) * FullUnrollTripCount + BEInsns;
     if (canUnrollCompletely(L, UP.Threshold, 100, UP.DynamicCostSavingsDiscount,
                             UnrolledSize, UnrolledSize)) {
+      UseUpperBound = (MaxTripCount == FullUnrollTripCount);
+      TripCount = FullUnrollTripCount;
+      TripMultiple = UP.UpperBound ? 1 : TripMultiple;
       UP.Count = TripCount;
       return ExplicitUnroll;
     } else {
       // The loop isn't that small, but we still can fully unroll it if that
       // helps to remove a significant number of instructions.
       // To check that, run additional analysis on the loop.
       if (Optional<EstimatedUnrollCost> Cost = analyzeLoopUnrollCost(
-              L, TripCount, DT, *SE, TTI,
+              L, FullUnrollTripCount, DT, *SE, TTI,
               UP.Threshold + UP.DynamicCostSavingsDiscount))
         if (canUnrollCompletely(L, UP.Threshold,
                                 UP.PercentDynamicCostSavedThreshold,
                                 UP.DynamicCostSavingsDiscount,
                                 Cost->UnrolledCost, Cost->RolledDynamicCost)) {
+          UseUpperBound = (MaxTripCount == FullUnrollTripCount);
+          TripCount = FullUnrollTripCount;
+          TripMultiple = UP.UpperBound ? 1 : TripMultiple;
           UP.Count = TripCount;
           return ExplicitUnroll;
         }
     }
   }
 
   // 4rd priority is partial unrolling.
   // Try partial unroll only when TripCount could be staticaly calculated.

 
 static bool tryToUnrollLoop(Loop *L, DominatorTree &DT, LoopInfo *LI,
                             ScalarEvolution *SE, const TargetTransformInfo &TTI,
                             AssumptionCache &AC, OptimizationRemarkEmitter &ORE,
                             bool PreserveLCSSA,
                             Optional<unsigned> ProvidedCount,
                             Optional<unsigned> ProvidedThreshold,
                             Optional<bool> ProvidedAllowPartial,
-                            Optional<bool> ProvidedRuntime) {
+                            Optional<bool> ProvidedRuntime,
+                            Optional<bool> ProvidedUpperBound) {
   DEBUG(dbgs() << "Loop Unroll: F[" << L->getHeader()->getParent()->getName()
                << "] Loop %" << L->getHeader()->getName() << "\n");
   if (HasUnrollDisablePragma(L)) {
     return false;
   }
 
   unsigned NumInlineCandidates;
   bool NotDuplicatable;
static bool tryToUnrollLoop(Loop *L, DominatorTree &DT, LoopInfo *LI,
   if (!L->isLoopSimplifyForm()) {
     DEBUG(
         dbgs() << "  Not unrolling loop which is not in loop-simplify form.\n");
     return false;
   }
 
   // Find trip count and trip multiple if count is not available
   unsigned TripCount = 0;
+  unsigned MaxTripCount = 0;
   unsigned TripMultiple = 1;
   // If there are multiple exiting blocks but one of them is the latch, use the
   // latch for the trip count estimation. Otherwise insist on a single exiting
   // block for the trip count estimation.
   BasicBlock *ExitingBlock = L->getLoopLatch();
   if (!ExitingBlock || !L->isLoopExiting(ExitingBlock))
     ExitingBlock = L->getExitingBlock();
   if (ExitingBlock) {
     TripCount = SE->getSmallConstantTripCount(L, ExitingBlock);
     TripMultiple = SE->getSmallConstantTripMultiple(L, ExitingBlock);
   }
 
   TargetTransformInfo::UnrollingPreferences UP = gatherUnrollingPreferences(
       L, TTI, ProvidedThreshold, ProvidedCount, ProvidedAllowPartial,
-      ProvidedRuntime);
+      ProvidedRuntime, ProvidedUpperBound);
 
   // Exit early if unrolling is disabled.
   if (UP.Threshold == 0 && (!UP.Partial || UP.PartialThreshold == 0))
     return false;
 
   // If the loop contains a convergent operation, the prelude we'd add
   // to do the first few instructions before we hit the unrolled loop
   // is unsafe -- it adds a control-flow dependency to the convergent
   // operation.  Therefore restrict remainder loop (try unrollig without).
   //
   // TODO: This is quite conservative.  In practice, convergent_op()
   // is likely to be called unconditionally in the loop.  In this
   // case, the program would be ill-formed (on most architectures)
   // unless n were the same on all threads in a thread group.
   // Assuming n is the same on all threads, any kind of unrolling is
   // safe.  But currently llvm's notion of convergence isn't powerful
   // enough to express this.
   if (Convergent)
     UP.AllowRemainder = false;
 
-  bool IsCountSetExplicitly = computeUnrollCount(
-      L, TTI, DT, LI, SE, &ORE, TripCount, TripMultiple, LoopSize, UP);
+  // Try to find the trip count upper bound if it is allowed and we cannot find
+  // exact trip count.
+  if (UP.UpperBound) {
+    if (!TripCount) {
+      MaxTripCount = SE->getSmallConstantMaxTripCount(L);
+      // Only unroll with small upper bound.
+      if (MaxTripCount > UnrollMaxUpperBound)
+        MaxTripCount = 0;
+    }
+  }
+
+  // computeUnrollCount() decides whether it is beneficial to use upper bound to
+  // fully unroll the loop.
+  bool UseUpperBound = false;
+  bool IsCountSetExplicitly =
+      computeUnrollCount(L, TTI, DT, LI, SE, &ORE, TripCount, MaxTripCount,
+                         TripMultiple, LoopSize, UP, UseUpperBound);
   if (!UP.Count)
     return false;
   // Unroll factor (Count) must be less or equal to TripCount.
   if (TripCount && UP.Count > TripCount)
     UP.Count = TripCount;
 
   // Unroll the loop.
   if (!UnrollLoop(L, UP.Count, TripCount, UP.Force, UP.Runtime,
-                  UP.AllowExpensiveTripCount, TripMultiple, LI, SE, &DT, &AC,
-                  &ORE, PreserveLCSSA))
+                  UP.AllowExpensiveTripCount, UseUpperBound, TripMultiple, LI,
+                  SE, &DT, &AC, &ORE, PreserveLCSSA))
     return false;
 
   // If loop has an unroll count pragma or unrolled by explicitly set count
   // mark loop as unrolled to prevent unrolling beyond that requested.
   if (IsCountSetExplicitly)
     SetLoopAlreadyUnrolled(L);
   return true;
 }
 
 namespace {
 class LoopUnroll : public LoopPass {
 public:
   static char ID; // Pass ID, replacement for typeid
   LoopUnroll(Optional<unsigned> Threshold = None,
              Optional<unsigned> Count = None,
-             Optional<bool> AllowPartial = None, Optional<bool> Runtime = None)
+             Optional<bool> AllowPartial = None, Optional<bool> Runtime = None,
+             Optional<bool> UpperBound = None)
       : LoopPass(ID), ProvidedCount(std::move(Count)),
         ProvidedThreshold(Threshold), ProvidedAllowPartial(AllowPartial),
-        ProvidedRuntime(Runtime) {
+        ProvidedRuntime(Runtime), ProvidedUpperBound(UpperBound) {
     initializeLoopUnrollPass(*PassRegistry::getPassRegistry());
   }
 
   Optional<unsigned> ProvidedCount;
   Optional<unsigned> ProvidedThreshold;
   Optional<bool> ProvidedAllowPartial;
   Optional<bool> ProvidedRuntime;
+  Optional<bool> ProvidedUpperBound;
 
   bool runOnLoop(Loop *L, LPPassManager &) override {
     if (skipLoop(L))
       return false;
 
     Function &F = *L->getHeader()->getParent();
 
     auto &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
     LoopInfo *LI = &getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
     ScalarEvolution *SE = &getAnalysis<ScalarEvolutionWrapperPass>().getSE();
     const TargetTransformInfo &TTI =
         getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
     auto &AC = getAnalysis<AssumptionCacheTracker>().getAssumptionCache(F);
     // For the old PM, we can't use OptimizationRemarkEmitter as an analysis
     // pass.  Function analyses need to be preserved across loop transformations
     // but ORE cannot be preserved (see comment before the pass definition).
     OptimizationRemarkEmitter ORE(&F);
     bool PreserveLCSSA = mustPreserveAnalysisID(LCSSAID);
 
     return tryToUnrollLoop(L, DT, LI, SE, TTI, AC, ORE, PreserveLCSSA,
                            ProvidedCount, ProvidedThreshold,
-                           ProvidedAllowPartial, ProvidedRuntime);
+                           ProvidedAllowPartial, ProvidedRuntime,
+                           ProvidedUpperBound);
   }
 
   /// This transformation requires natural loop information & requires that
   /// loop preheaders be inserted into the CFG...
   ///
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<AssumptionCacheTracker>();
     AU.addRequired<TargetTransformInfoWrapperPass>();
     // FIXME: Loop passes are required to preserve domtree, and for now we just
     // recreate dom info if anything gets unrolled.
     getLoopAnalysisUsage(AU);
   }
 };
 }
 
 char LoopUnroll::ID = 0;
 INITIALIZE_PASS_BEGIN(LoopUnroll, "loop-unroll", "Unroll loops", false, false)
 INITIALIZE_PASS_DEPENDENCY(AssumptionCacheTracker)
 INITIALIZE_PASS_DEPENDENCY(LoopPass)
 INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
 INITIALIZE_PASS_END(LoopUnroll, "loop-unroll", "Unroll loops", false, false)
 
 Pass *llvm::createLoopUnrollPass(int Threshold, int Count, int AllowPartial,
-                                 int Runtime) {
+                                 int Runtime, int UpperBound) {
   // TODO: It would make more sense for this function to take the optionals
   // directly, but that's dangerous since it would silently break out of tree
   // callers.
   return new LoopUnroll(Threshold == -1 ? None : Optional<unsigned>(Threshold),
                         Count == -1 ? None : Optional<unsigned>(Count),
                         AllowPartial == -1 ? None
                                            : Optional<bool>(AllowPartial),
-                        Runtime == -1 ? None : Optional<bool>(Runtime));
+                        Runtime == -1 ? None : Optional<bool>(Runtime),
+                        UpperBound == -1 ? None : Optional<bool>(UpperBound));
 }
 
 Pass *llvm::createSimpleLoopUnrollPass() {
-  return llvm::createLoopUnrollPass(-1, -1, 0, 0);
+  return llvm::createLoopUnrollPass(-1, -1, 0, 0, 0);
 }
 
 PreservedAnalyses LoopUnrollPass::run(Loop &L, LoopAnalysisManager &AM) {
   const auto &FAM =
       AM.getResult<FunctionAnalysisManagerLoopProxy>(L).getManager();
   Function *F = L.getHeader()->getParent();
 
 
PreservedAnalyses LoopUnrollPass::run(Loop &L, LoopAnalysisManager &AM) {
   if (!TTI)
     report_fatal_error("LoopUnrollPass: TargetIRAnalysis not cached at a higher level");
   if (!AC)
     report_fatal_error("LoopUnrollPass: AssumptionAnalysis not cached at a higher level");
   if (!ORE)
     report_fatal_error("LoopUnrollPass: OptimizationRemarkEmitterAnalysis not "
                        "cached at a higher level");
 
-  bool Changed = tryToUnrollLoop(
-      &L, *DT, LI, SE, *TTI, *AC, *ORE, /*PreserveLCSSA*/ true, ProvidedCount,
-      ProvidedThreshold, ProvidedAllowPartial, ProvidedRuntime);
+  bool Changed =
+      tryToUnrollLoop(&L, *DT, LI, SE, *TTI, *AC, *ORE, /*PreserveLCSSA*/ true,
+                      ProvidedCount, ProvidedThreshold, ProvidedAllowPartial,
+                      ProvidedRuntime, ProvidedUpperBound);
 
   if (!Changed)
     return PreservedAnalyses::all();
   return getLoopPassPreservedAnalyses();
 }

llvm/trunk/lib/Transforms/Utils/LoopUnroll.cpp

 /// executes. UnrollLoop relaxes the definition to permit early exits: here
 /// TripCount is the iteration on which control exits LatchBlock if no early
 /// exits were taken. Note that UnrollLoop assumes that the loop counter test
 /// terminates LatchBlock in order to remove unnecesssary instances of the
 /// test. In other words, control may exit the loop prior to TripCount
 /// iterations via an early branch, but control may not exit the loop from the
 /// LatchBlock's terminator prior to TripCount iterations.
 ///
+/// PreserveCondBr indicates whether the conditional branch of the LatchBlock
+/// needs to be preserved.  It is needed when we use trip count upper bound to
+/// fully unroll the loop.
+///
 /// Similarly, TripMultiple divides the number of times that the LatchBlock may
 /// execute without exiting the loop.
 ///
 /// If AllowRuntime is true then UnrollLoop will consider unrolling loops that
 /// have a runtime (i.e. not compile time constant) trip count.  Unrolling these
 /// loops require a unroll "prologue" that runs "RuntimeTripCount % Count"
 /// iterations before branching into the unrolled loop.  UnrollLoop will not
 /// runtime-unroll the loop if computing RuntimeTripCount will be expensive and
 /// AllowExpensiveTripCount is false.
 ///
 /// The LoopInfo Analysis that is passed will be kept consistent.
 ///
 /// This utility preserves LoopInfo. It will also preserve ScalarEvolution and
 /// DominatorTree if they are non-null.
 bool llvm::UnrollLoop(Loop *L, unsigned Count, unsigned TripCount, bool Force,
                       bool AllowRuntime, bool AllowExpensiveTripCount,
-                      unsigned TripMultiple, LoopInfo *LI, ScalarEvolution *SE,
-                      DominatorTree *DT, AssumptionCache *AC,
-                      OptimizationRemarkEmitter *ORE, bool PreserveLCSSA) {
+                      bool PreserveCondBr, unsigned TripMultiple, LoopInfo *LI,
+                      ScalarEvolution *SE, DominatorTree *DT,
+                      AssumptionCache *AC, OptimizationRemarkEmitter *ORE,
+                      bool PreserveLCSSA) {
   BasicBlock *Preheader = L->getLoopPreheader();
   if (!Preheader) {
     DEBUG(dbgs() << "  Can't unroll; loop preheader-insertion failed.\n");
     return false;
   }
 
   BasicBlock *LatchBlock = L->getLoopLatch();
   if (!LatchBlock) {
if (RuntimeTripCount && j != 0) {
       NeedConditional = false;
     }
 
     // For a complete unroll, make the last iteration end with a branch
     // to the exit block.
     if (CompletelyUnroll) {
       if (j == 0)
         Dest = LoopExit;
-      NeedConditional = false;
-    }
-
+      // If using trip count upper bound to completely unroll, we need to keep
+      // the conditional branch except the last one because the loop may exit
+      // after any iteration.
+      assert(NeedConditional &&
+             "NeedCondition cannot be modified by both complete "
+             "unrolling and runtime unrolling");
+      NeedConditional = (PreserveCondBr && j);
+    } else if (j != BreakoutTrip && (TripMultiple == 0 || j % TripMultiple != 0)) {
       // If we know the trip count or a multiple of it, we can safely use an
       // unconditional branch for some iterations.
-    if (j != BreakoutTrip && (TripMultiple == 0 || j % TripMultiple != 0)) {
       NeedConditional = false;
     }
 
     if (NeedConditional) {
       // Update the conditional branch's successor for the following
       // iteration.
       Term->setSuccessor(!ContinueOnTrue, Dest);
     } else {

llvm/trunk/test/Transforms/LoopUnroll/AArch64/full-unroll-trip-count-upper-bound.ll

+; RUN: opt -loop-unroll -S -mtriple aarch64 -mcpu=cortex-a57 %s | FileCheck %s -check-prefix=UNROLL
+; RUN: opt -loop-unroll -unroll-max-upperbound=0 -S -mtriple aarch64 -mcpu=cortex-a57 %s | FileCheck %s -check-prefix=NOUNROLL
+
+; This IR comes from this C code:
+;
+;   for (int i = 0; i < 4; i++) {
+;     if (src[i] == 1) {
+;       *dst = i;
+;       break;
+;     }
+;   }
+;
+; This test is meant to check that this loop is unrolled into four iterations.
+
+; UNROLL-LABEL: @test
+; UNROLL: load i32, i32*
+; UNROLL: load i32, i32*
+; UNROLL: load i32, i32*
+; UNROLL: load i32, i32*
+; UNROLL-NOT: load i32, i32*
+; NOUNROLL-LABEL: @test
+; NOUNROLL: load i32, i32*
+; NOUNROLL-NOT: load i32, i32*
+
+define void @test(i32* %dst, i32* %src) {
+entry:
+  br label %for.body
+
+for.body:                                         ; preds = %entry, %for.body
+  %i = phi i32 [ 0, %entry ], [ %inc, %for.body ]
+  %0 = sext i32 %i to i64
+  %1 = getelementptr inbounds i32, i32* %src, i64 %0
+  %2 = load i32, i32* %1
+  %inc = add nsw i32 %i, 1
+  %cmp1 = icmp slt i32 %inc, 4
+  %cmp3 = icmp eq i32 %2, 1 
+  %or.cond = and i1 %cmp3, %cmp1
+  br i1 %or.cond, label %for.body, label %exit
+
+exit:                                          ; preds = %for.body
+  store i32 %i, i32* %dst
+  ret void
+}