[LoopUnroll] Fix bug in computeUnrollCount causing it to not honor MaxCount

[anna]

lib/Transforms/Scalar/LoopUnrollPass.cpp

813

This check should actually be present along all paths from this function?

839

What happens if we early return from this function with result as true, but UP.Count > UP.MaxCount?

Perhaps we should be setting the max value at the caller of this function :

tryToUnrollLoop() {
....
 // 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;
  // Add the check after this?

}

[gberry]

lib/Transforms/Scalar/LoopUnrollPass.cpp

839

I think this is the only place that MaxCount wasn't being honored where it should be. The earlier checks that return early are:

• explicit -unroll-count settings and pragma unroll counts, which I don't think should be affected by MaxCount
• full unroll count, which has a separate max: FullUnrollMaxCount which seems to be honored
• loop peeling: sets Count to 1, so MaxCount should not have any effect

[anna]

LGTM.
Just clarifying: this just honours the user defined max count right (if user hasn't defined the value, it would be UINT_MAX).

[gberry]

In D34532#791464, @anna wrote:

LGTM.
Just clarifying: this just honours the user defined max count right (if user hasn't defined the value, it would be UINT_MAX).

That, or the value set by the target hook getUnrollingPreferences(), which is how I originally ran into the bug, when implementing D34533

lib/Transforms/Scalar/LoopUnrollPass.cpp

if (TripCount) {
     if (UP.Count == 0)
       UP.Count = TripCount;
     if (UP.PartialThreshold != NoThreshold) {
       // Reduce unroll count to be modulo of TripCount for partial unrolling.
       if (getUnrolledLoopSize(LoopSize, UP) > UP.PartialThreshold)
         UP.Count =
             (std::max(UP.PartialThreshold, UP.BEInsns + 1) - UP.BEInsns) /
             (LoopSize - UP.BEInsns);
       if (UP.Count > UP.MaxCount)
         UP.Count = UP.MaxCount;
       while (UP.Count != 0 && TripCount % UP.Count != 0)
         UP.Count--;
       if (UP.AllowRemainder && UP.Count <= 1) {
         // If there is no Count that is modulo of TripCount, set Count to
         // largest power-of-two factor that satisfies the threshold limit.
         // As we'll create fixup loop, do the type of unrolling only if
         // remainder loop is allowed.
if (UP.PartialThreshold != NoThreshold) {
                                        L->getStartLoc(), L->getHeader())
               << "Unable to unroll loop as directed by unroll(enable) pragma "
                  "because unrolled size is too large.");
         UP.Count = 0;
       }
     } else {
       UP.Count = TripCount;
     }
+    if (UP.Count > UP.MaxCount)
+      UP.Count = UP.MaxCount;
     if ((PragmaFullUnroll || PragmaEnableUnroll) && TripCount &&
         UP.Count != TripCount)
       ORE->emit(
           OptimizationRemarkMissed(DEBUG_TYPE, "FullUnrollAsDirectedTooLarge",
                                    L->getStartLoc(), L->getHeader())
           << "Unable to fully unroll loop as directed by unroll pragma because "
              "unrolled size is too large.");
     return ExplicitUnroll;

test/Transforms/LoopUnroll/unroll-maxcount.ll

+; RUN: opt < %s -S -loop-unroll -unroll-allow-partial -unroll-max-count=1 | FileCheck %s
+; Checks that unroll MaxCount is honored.
+;
+; CHECK-LABEL: @foo(
+; CHECK-LABEL: for.body:
+; CHECK-NEXT: phi
+; CHECK-NEXT: getelementptr
+; CHECK-NEXT: load
+; CHECK-NEXT: add
+; CHECK-NEXT: store
+; CHECK-NEXT: add
+; CHECK-NEXT: icmp
+; CHECK-NEXT: br
+define void @foo(i32* nocapture %a) {
+entry:
+  br label %for.body
+
+for.body:                                         ; preds = %for.body, %entry
+  %indvars.iv = phi i64 [ 0, %entry ], [ %indvars.iv.next, %for.body ]
+  %arrayidx = getelementptr inbounds i32, i32* %a, i64 %indvars.iv
+  %0 = load i32, i32* %arrayidx, align 4
+  %inc = add nsw i32 %0, 1
+  store i32 %inc, i32* %arrayidx, align 4
+  %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
+  %exitcond = icmp eq i64 %indvars.iv.next, 1024
+  br i1 %exitcond, label %for.end, label %for.body
+
+for.end:                                          ; preds = %for.body
+  ret void
+}
+