[LV] Clamp the VF to the trip count

[anna]

Summary

When the MaxVectorSize > ConstantTripCount, we should just clamp the
vectorization factor to be the ConstantTripCount.
This vectorizes loops where the TinyTripCountThreshold >= TripCount < MaxVF.

Earlier we were finding the maximum vector width, which could be greater than
the trip count itself. The Loop vectorizer does all the work for generating a
vectorizable loop, but in the end we would always choose the scalar loop (since
the VF > trip count). This allows us to choose the VF keeping in mind the trip
count if available.

This is a fix on top of rL312472.

[Ayal]

lib/Transforms/Vectorize/LoopVectorize.cpp

6166 ↗

(On Diff #114643)

Better remove the default setting of TC = 0 in the declaration of computeFeasibleMaxVF(), which becomes dead.

6179 ↗

(On Diff #114643)

Probably good to hoist this DEBUG line along with computing TC earlier.

6258 ↗

(On Diff #114643)

While we're at it, in this or a follow-up patch:

• Suggest to also clamp NewMaxVectorSize above to TC.
• VS below should start from MaxVectorSize*2 instead of MaxVectorSize; it's pointless and wasteful to calculateRegisterUsage for MaxVectorSize, as the latter will be returned by default anyway.
• Optimize calculateRegisterUsage() for the case where VFs is empty.

Doing this will also take care of your early exit above.

test/Transforms/LoopVectorize/X86/vector_max_bandwidth.ll

3 ↗

(On Diff #114643)

You could reuse CHECK-AVX2 instead of adding CHECK-CLAMP, right?

[anna]

lib/Transforms/Vectorize/LoopVectorize.cpp

6179 ↗

(On Diff #114643)

When OptForSize is true (i.e. at this point in the code), we know for sure that a small trip count value returns a constant trip count or value of 1, if we didn't know the trip count.

I left this here rather than move up because there maybe *more* spurious cases where we do not have a *small* trip count, and the DEBUG statement wouldn't be useful.

6258 ↗

(On Diff #114643)

To me, it seems rather than clamping the NewMaxVectorSize, explicitly showing the early return is clearer.
The second and third suggestions I'll do in a follow-up patch.

Also, an additional point to follow up: I think we can early return in even earlier case - when VF is set to 1 because the target has no vector registers.

[anna]

Addressed review comments - no longer need the default argument,
updated test to reuse the existing RUN.

[anna]

updated previous diff to CHECK for the correct VF in the first test.

[Ayal]

lib/Transforms/Vectorize/LoopVectorize.cpp

6179 ↗

(On Diff #114643)

At this point I think TC can also be zero if we don't know it, which is why we check if TC<2 below (i.e., if (TC==0 || TC == 1)).

Printing out the value of TC when computed above can only aid debugging, afaics.

In any case, I'm also ok leaving it here; you're also printing TC in computeFeasibleMaxVF() where it clamps MaxVF.

6258 ↗

(On Diff #114643)

Agreed!

test/Transforms/LoopVectorize/X86/vector_max_bandwidth.ll

54 ↗

(On Diff #114798)

The max possible vector width for this test on AVX1 is 16, so we're missing the point in checking its selected VF here. Suffice to CHECK-AVX2 only, whose max possible vector width is 32, as noted.

[anna]

@Ayal, any other comments or does this look good to go? Thanks.

lib/Transforms/Vectorize/LoopVectorize.cpp

6179 ↗

(On Diff #114643)

yes, the clamped value is printed in the debug statement.

test/Transforms/LoopVectorize/X86/vector_max_bandwidth.ll

54 ↗

(On Diff #114798)

agreed, I'll remove the CHECK-AVX1 check.

[anna]

minor review comment: removed unnecessary CHECK.

[Ayal]

@Ayal, any other comments or does this look good to go? Thanks.

This looks good to go!

llvm/trunk/lib/Transforms/Vectorize/LoopVectorize.cpp

void collectUniformsAndScalars(unsigned VF) {
     setCostBasedWideningDecision(VF);
     collectLoopUniforms(VF);
     collectLoopScalars(VF);
   }
 
 private:
   /// \return An upper bound for the vectorization factor, larger than zero.
   /// One is returned if vectorization should best be avoided due to cost.
-  unsigned computeFeasibleMaxVF(bool OptForSize, unsigned ConstTripCount = 0);
+  unsigned computeFeasibleMaxVF(bool OptForSize, unsigned ConstTripCount);
 
   /// The vectorization cost is a combination of the cost itself and a boolean
   /// indicating whether any of the contributing operations will actually
   /// operate on
   /// vector values after type legalization in the backend. If this latter value
   /// is
   /// false, then all operations will be scalarized (i.e. no vectorization has
   /// actually taken place).
if (Legal->getRuntimePointerChecking()->Need && TTI.hasBranchDivergence()) {
 
     ORE->emit(
       createMissedAnalysis("CantVersionLoopWithDivergentTarget")
       << "runtime pointer checks needed. Not enabled for divergent target");
 
     return None;
   }
 
+  unsigned TC = PSE.getSE()->getSmallConstantTripCount(TheLoop);
   if (!OptForSize) // Remaining checks deal with scalar loop when OptForSize.
-    return computeFeasibleMaxVF(OptForSize);
+    return computeFeasibleMaxVF(OptForSize, TC);
 
   if (Legal->getRuntimePointerChecking()->Need) {
     ORE->emit(createMissedAnalysis("CantVersionLoopWithOptForSize")
               << "runtime pointer checks needed. Enable vectorization of this "
                  "loop with '#pragma clang loop vectorize(enable)' when "
                  "compiling with -Os/-Oz");
     DEBUG(dbgs()
           << "LV: Aborting. Runtime ptr check is required with -Os/-Oz.\n");
     return None;
   }
 
   // If we optimize the program for size, avoid creating the tail loop.
-  unsigned TC = PSE.getSE()->getSmallConstantTripCount(TheLoop);
   DEBUG(dbgs() << "LV: Found trip count: " << TC << '\n');
 
   // If we don't know the precise trip count, don't try to vectorize.
   if (TC < 2) {
     ORE->emit(
         createMissedAnalysis("UnknownLoopCountComplexCFG")
         << "unable to calculate the loop count due to complex control flow");
     DEBUG(dbgs() << "LV: Aborting. A tail loop is required with -Os/-Oz.\n");
WidestRegister =
       ((WidestRegister < MaxSafeDepDist) ? WidestRegister : MaxSafeDepDist);
   unsigned MaxVectorSize = WidestRegister / WidestType;
 
   DEBUG(dbgs() << "LV: The Smallest and Widest types: " << SmallestType << " / "
                << WidestType << " bits.\n");
   DEBUG(dbgs() << "LV: The Widest register is: " << WidestRegister
                << " bits.\n");
 
+  assert(MaxVectorSize <= 64 && "Did not expect to pack so many elements"
+                                " into one vector!");
   if (MaxVectorSize == 0) {
     DEBUG(dbgs() << "LV: The target has no vector registers.\n");
     MaxVectorSize = 1;
   } else if (ConstTripCount && ConstTripCount < MaxVectorSize &&
-             isPowerOf2_32(ConstTripCount))
+             isPowerOf2_32(ConstTripCount)) {
+    // We need to clamp the VF to be the ConstTripCount. There is no point in
+    // choosing a higher viable VF as done in the loop below.
+    DEBUG(dbgs() << "LV: Clamping the MaxVF to the constant trip count: "
+                 << ConstTripCount << "\n");
     MaxVectorSize = ConstTripCount;
-
-  assert(MaxVectorSize <= 64 && "Did not expect to pack so many elements"
-                                " into one vector!");
+    return MaxVectorSize;
+  }
 
   unsigned MaxVF = MaxVectorSize;
   if (MaximizeBandwidth && !OptForSize) {
     // Collect all viable vectorization factors.
     SmallVector<unsigned, 8> VFs;
     unsigned NewMaxVectorSize = WidestRegister / SmallestType;
     for (unsigned VS = MaxVectorSize; VS <= NewMaxVectorSize; VS *= 2)
       VFs.push_back(VS);

llvm/trunk/test/Transforms/LoopVectorize/X86/vector_max_bandwidth.ll

for.body:
   %3 = load i32, i32* %arrayidx10, align 4
   %add11 = add nsw i32 %3, %2
   %arrayidx13 = getelementptr inbounds [1000 x i32], [1000 x i32]* @u, i64 0, i64 %indvars.iv
   store i32 %add11, i32* %arrayidx13, align 4
   %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
   %exitcond = icmp eq i64 %indvars.iv.next, 1000
   br i1 %exitcond, label %for.cond.cleanup, label %for.body
 }
+
+; We should not choose a VF larger than the constant TC.
+; VF chosen should be atmost 16 (not the max possible vector width = 32 for AVX2)
+define void @not_too_small_tc(i8* noalias nocapture %A, i8* noalias nocapture readonly %B) {
+; CHECK-LABEL: not_too_small_tc
+; CHECK-AVX2: LV: Selecting VF: 16.
+entry:
+  br label %for.body
+
+for.body:
+  %indvars.iv = phi i64 [ 0, %entry ], [ %indvars.iv.next, %for.body ]
+  %arrayidx = getelementptr inbounds i8, i8* %B, i64 %indvars.iv
+  %l1 = load i8, i8* %arrayidx, align 4, !llvm.mem.parallel_loop_access !3
+  %arrayidx2 = getelementptr inbounds i8, i8* %A, i64 %indvars.iv
+  %l2 = load i8, i8* %arrayidx2, align 4, !llvm.mem.parallel_loop_access !3
+  %add = add i8 %l1, %l2
+  store i8 %add, i8* %arrayidx2, align 4, !llvm.mem.parallel_loop_access !3
+  %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
+  %exitcond = icmp eq i64 %indvars.iv.next, 16
+  br i1 %exitcond, label %for.end, label %for.body, !llvm.loop !4
+
+for.end:
+  ret void
+}
+!3 = !{!3}
+!4 = !{!4}