LoopVectorize: MaxVF should not be larger than the loop trip count

[zvi]

Summary

Improve how MaxVF is computed while taking into account that MaxVF should not be larger than the loop's trip count.

Other than saving on compile-time by pruning the possible MaxVF candidates, this patch fixes pr34438 which exposed the following flow:

1. Short trip count identified -> Don't bail out, set OptForSize:=True to avoid tail-loop and runtime checks.
2. Compute MaxVF returned 16 on a target supporting AVX512.
3. OptForSize -> choose VF:=MaxVF.
4. Bail out because TripCount = 8, VF = 16, TripCount % VF !=0 means we need a tail loop.

With this patch step 2. will choose MaxVF=8 based on TripCount.

[hfinkel]

test/Transforms/LoopVectorize/X86/pr34438.ll

7 ↗

(On Diff #113698)

Can you please test the output directly instead of testing the debug output? We generally test debug output when there's no other effective way to write a test. It's better to have non-asserts-requiring tests.

[zvi]

Test the contents of the transformed function instead of using debug prints. This addresses Hal's comment. Thanks.

[hfinkel]

LGTM

[Ayal]

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

6165

Would be good to restrict MaxVF to TC even if we're not OptForSize, although it's probably less likely to have a known TC that is larger than Tiny but smaller than current MaxVF.

6243

Better set MaxVectorSize to PowerOf2Floor(ConstTripCount), to also handle the case where ConstTripCount is not a power of 2.

Also, while we're at it: computeFeasibleMaxVF() is currently not required/documented to return a power of 2. We should either check that MaxVF is a power of 2 when checking if (TC % MaxVF != 0), and then can simply set MaxVectorSize to ConstTripCount here. Otherwise we should document/assert that MaxVF is a power of 2.

[zvi]

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

6165

Thanks for pointing that out! Created pr34468 showing a missed opportunity for TinyTripCount < ConstTripCount < MaxVL.

6243

Better set MaxVectorSize to PowerOf2Floor(ConstTripCount), to also handle the case where ConstTripCount is not a power of 2.

What would there be a benefit of this suggestion given that we currently bail out if there is a need for a tail-loop in the case of short-trip-count? Should we tie this with the work on vectorizing short trip count and allowing runtime checks/tail loops?

Also, while we're at it: computeFeasibleMaxVF() is currently not required/documented to return a power of 2. We should either check that MaxVF is a power of 2 when checking if (TC % MaxVF != 0), and then can simply set MaxVectorSize to ConstTripCount here. Otherwise we should document/assert that MaxVF is a power of 2.

It did not occur to me that MaxVF may not be a power of two. Can we label this as a follow-up work?

[Ayal]

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

6243

Better set MaxVectorSize to PowerOf2Floor(ConstTripCount), to also handle the case where ConstTripCount is not a power of 2.

What would there be a benefit of this suggestion given that we currently bail out if there is a need for a tail-loop in the case of short-trip-count? Should we tie this with the work on vectorizing short trip count and allowing runtime checks/tail loops?

Well, your pr34468 shows where there would be a benefit, namely when TinyTripCount < ConstTripCount < MaxVL which allows runtime checks, i.e., does not bail-out if a tail-loop is needed.

It did not occur to me that MaxVF may not be a power of two. Can we label this as a follow-up work?

Sure, suggest to simply document and assert that MaxVF is expected to be a power of two.

[anna]

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

6165

This is exactly the case we are seeing. Currently, even with this patch, for loops with constant trip count = 16, we will have OptForSize as false (because TC=16 is same as TinyTripCountVectorThreshold).

So, if the MaxVF=32, we still won't vectorize the loop. I'm working on the fix.

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 computeFeasibleMaxVF(bool OptForSize, unsigned ConstTripCount = 0);
 
   /// 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;
   }
 
   if (!OptForSize) // Remaining checks deal with scalar loop when OptForSize.
     return computeFeasibleMaxVF(OptForSize);
 
   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");
     return None;
   }
 
-  unsigned MaxVF = computeFeasibleMaxVF(OptForSize);
+  unsigned MaxVF = computeFeasibleMaxVF(OptForSize, TC);
 
   if (TC % MaxVF != 0) {
     // If the trip count that we found modulo the vectorization factor is not
     // zero then we require a tail.
     // FIXME: look for a smaller MaxVF that does divide TC rather than give up.
     // FIXME: return None if loop requiresScalarEpilog(<MaxVF>), or look for a
     //        smaller MaxVF that does not require a scalar epilog.
 
     ORE->emit(createMissedAnalysis("NoTailLoopWithOptForSize")
               << "cannot optimize for size and vectorize at the "
                  "same time. Enable vectorization of this loop "
                  "with '#pragma clang loop vectorize(enable)' "
                  "when compiling with -Os/-Oz");
     DEBUG(dbgs() << "LV: Aborting. A tail loop is required with -Os/-Oz.\n");
     return None;
   }
 
   return MaxVF;
 }
 
-unsigned LoopVectorizationCostModel::computeFeasibleMaxVF(bool OptForSize) {
+unsigned
+LoopVectorizationCostModel::computeFeasibleMaxVF(bool OptForSize,
+                                                 unsigned ConstTripCount) {
   MinBWs = computeMinimumValueSizes(TheLoop->getBlocks(), *DB, &TTI);
   unsigned SmallestType, WidestType;
   std::tie(SmallestType, WidestType) = getSmallestAndWidestTypes();
   unsigned WidestRegister = TTI.getRegisterBitWidth(true);
   unsigned MaxSafeDepDist = -1U;
 
   // Get the maximum safe dependence distance in bits computed by LAA. If the
   // loop contains any interleaved accesses, we divide the dependence distance
LoopVectorizationCostModel::computeFeasibleMaxVF(bool OptForSize,
   DEBUG(dbgs() << "LV: The Smallest and Widest types: " << SmallestType << " / "
                << WidestType << " bits.\n");
   DEBUG(dbgs() << "LV: The Widest register is: " << WidestRegister
                << " bits.\n");
 
   if (MaxVectorSize == 0) {
     DEBUG(dbgs() << "LV: The target has no vector registers.\n");
     MaxVectorSize = 1;
-  }
+  } else if (ConstTripCount && ConstTripCount < MaxVectorSize &&
+             isPowerOf2_32(ConstTripCount))
+    MaxVectorSize = ConstTripCount;
 
   assert(MaxVectorSize <= 64 && "Did not expect to pack so many elements"
                                 " into one vector!");
 
   unsigned MaxVF = MaxVectorSize;
   if (MaximizeBandwidth && !OptForSize) {
     // Collect all viable vectorization factors.
     SmallVector<unsigned, 8> VFs;

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

+; PR34438
+; Loop has a short trip count of 8 iterations. It should be vectorized because no runtime checks or tail loop are necessary.
+; Two cases tested AVX (MaxVF=8 = TripCount) and AVX512 (MaxVF=16 > TripCount)
+
+; RUN: opt < %s -loop-vectorize -mtriple=x86_64-apple-macosx10.8.0 -mcpu=corei7-avx -S | FileCheck %s
+; RUN: opt < %s -loop-vectorize -mtriple=x86_64-apple-macosx10.8.0 -mcpu=skylake-avx512 -S | FileCheck %s
+
+target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128-n8:16:32:64-S128"
+target triple = "x86_64-apple-macosx10.8.0"
+
+define void @small_tc(float* noalias nocapture %A, float* noalias nocapture readonly %B) {
+; CHECK-LABEL: @small_tc
+; CHECK:    load <8 x float>, <8 x float>*
+; CHECK:    fadd fast <8 x float>
+entry:
+  br label %for.body
+
+for.body:
+  %indvars.iv = phi i64 [ 0, %entry ], [ %indvars.iv.next, %for.body ]
+  %arrayidx = getelementptr inbounds float, float* %B, i64 %indvars.iv
+  %0 = load float, float* %arrayidx, align 4, !llvm.mem.parallel_loop_access !3
+  %arrayidx2 = getelementptr inbounds float, float* %A, i64 %indvars.iv
+  %1 = load float, float* %arrayidx2, align 4, !llvm.mem.parallel_loop_access !3
+  %add = fadd fast float %0, %1
+  store float %add, float* %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, 8
+  br i1 %exitcond, label %for.end, label %for.body, !llvm.loop !4
+
+for.end:
+  ret void
+}
+
+!3 = !{!3}
+!4 = !{!4}