[LoopVectorize] Let selectVectorizationFactor choose a scalable VF.

[sdesmalen]

Summary

This patch changes the algorithm of selectVectorizationFactor to
iterate all power of 2 VFs between Min/MaxFactor, for both fixed
and scalable Min/MaxFactor VFs. It then invokes the cost-model as
normal, and selects the most cost-effective VF to vectorize the
loop.

[sdesmalen]

Rebased

[sdesmalen]

Rebased

[sdesmalen]

Rebased

llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h

if (LHSPtr->hasValue())
       *LHSPtr = ElementCount::isKnownGT(**LHSPtr, VF) ? **LHSPtr : VF;
     else
       *LHSPtr = VF;
   }
 
   explicit operator bool() const { return FixedVF || ScalableVF; }
   bool hasFixedVF() const { return FixedVF.hasValue(); }
   bool hasScalableVF() const { return ScalableVF.hasValue(); }
+  bool isScalar() const {
+    return FixedVF && FixedVF->isScalar() && !ScalableVF;
+  }
   Optional<ElementCount> getFixedVF() const { return FixedVF; }
   Optional<ElementCount> getScalableVF() const { return ScalableVF; }
 };
 
 
 /// Planner drives the vectorization process after having passed
 /// Legality checks.
 class LoopVectorizationPlanner {

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

public:
   /// otherwise.
   bool runtimeChecksRequired();
 
   /// \return The most profitable vectorization factor and the cost of that VF.
   /// This method checks every power of two up to MaxVF. If UserVF is not ZERO
   /// then this vectorization factor will be selected if vectorization is
   /// possible.
   VectorizationFactor
-  selectVectorizationFactor(const OptionalVFCandidates &MaxVF);
+  selectVectorizationFactor(const OptionalVFCandidates &MinVFs,
+                            const OptionalVFCandidates &MaxVFs);
   VectorizationFactor
   selectEpilogueVectorizationFactor(const ElementCount MaxVF,
                                     const LoopVectorizationPlanner &LVP);
 
   /// Setup cost-based decisions for user vectorization factor.
   void selectUserVectorizationFactor(ElementCount UserVF) {
     collectUniformsAndScalars(UserVF);
     collectInstsToScalarize(UserVF);
if (ElementCount MinVF =
         MaxVF = MinVF;
       }
     }
   }
   return MaxVF;
 }
 
 VectorizationFactor LoopVectorizationCostModel::selectVectorizationFactor(
+    const OptionalVFCandidates &MinFactors,
     const OptionalVFCandidates &MaxFactors) {
-  // Ignore the Scalable MaxVF in decision making process for now.
-  ElementCount MaxVF = *MaxFactors.getFixedVF();
-
-  // FIXME: This can be fixed for scalable vectors later, because at this stage
-  // the LoopVectorizer will only consider vectorizing a loop with scalable
-  // vectors when the loop has a hint to enable vectorization for a given VF.
-  assert(!MaxVF.isScalable() && "scalable vectors not yet supported");
-
   InstructionCost ExpectedCost = expectedCost(ElementCount::getFixed(1)).first;
   LLVM_DEBUG(dbgs() << "LV: Scalar loop costs: " << ExpectedCost << ".\n");
   assert(ExpectedCost.isValid() && "Unexpected invalid cost for scalar loop");
 
   auto Width = ElementCount::getFixed(1);
   const float ScalarCost = *ExpectedCost.getValue();
   float Cost = ScalarCost;
 
   bool ForceVectorization = Hints->getForce() == LoopVectorizeHints::FK_Enabled;
-  if (ForceVectorization && MaxVF.isVector()) {
+  if (ForceVectorization && !MaxFactors.isScalar()) {
     // Ignore scalar width, because the user explicitly wants vectorization.
     // Initialize cost to max so that VF = 2 is, at least, chosen during cost
     // evaluation.
     Cost = std::numeric_limits<float>::max();
   }
 
-  for (auto i = ElementCount::getFixed(2); ElementCount::isKnownLE(i, MaxVF);
-       i *= 2) {
+  SmallVector<ElementCount> PowerOf2VFs;
+  genAllPowerOfTwoFactorsUptoMaxFactors(PowerOf2VFs, MinFactors, MaxFactors);
+  for (const auto &i : PowerOf2VFs) {
     // Notice that the vector loop needs to be executed less times, so
     // we need to divide the cost of the vector loops by the width of
     // the vector elements.
     VectorizationCostTy C = expectedCost(i);
     assert(C.first.isValid() && "Unexpected invalid cost for vector loop");
-    float VectorCost = *C.first.getValue() / (float)i.getFixedValue();
+    float VectorCost = *C.first.getValue() / (float)i.getKnownMinValue();
     LLVM_DEBUG(dbgs() << "LV: Vector loop of width " << i
                       << " costs: " << (int)VectorCost << ".\n");
     if (!C.second && !ForceVectorization) {
       LLVM_DEBUG(
           dbgs() << "LV: Not considering vector loop of width " << i
                  << " because it will not generate any vector instructions.\n");
       continue;
     }
if (canTruncateToMinimalBitwidth(I, VF)) {
             largestIntegerVectorType(ToVectorTy(I->getType(), VF), MinVecTy);
       } else if (Opcode == Instruction::ZExt || Opcode == Instruction::SExt) {
         SrcVecTy = largestIntegerVectorType(SrcVecTy, MinVecTy);
         VectorTy =
             smallestIntegerVectorType(ToVectorTy(I->getType(), VF), MinVecTy);
       }
     }
 
-    unsigned N;
-    if (isScalarAfterVectorization(I, VF)) {
-      assert(!VF.isScalable() && "VF is assumed to be non scalable");
-      N = VF.getKnownMinValue();
-    } else
-      N = 1;
+    unsigned N = isScalarAfterVectorization(I, VF) ? VF.getKnownMinValue() : 1;
     return N *
            TTI.getCastInstrCost(Opcode, VectorTy, SrcVecTy, CCH, CostKind, I);
   }
   case Instruction::Call: {
     bool NeedToScalarize;
     CallInst *CI = cast<CallInst>(I);
     InstructionCost CallCost = getVectorCallCost(CI, VF, NeedToScalarize);
     if (getVectorIntrinsicIDForCall(CI, TLI)) {
LoopVectorizationPlanner::plan(ElementCount UserVF, unsigned UserIC) {
   }
 
   ElementCount MaxUserVF = (UserVF.isScalable() && MaxFactors.hasScalableVF())
                                ? *MaxFactors.getScalableVF()
                                : *MaxFactors.getFixedVF();
   assert(MaxUserVF.isNonZero() && "MaxUserVF is zero.");
 
   bool UserVFIsLegal = ElementCount::isKnownLE(UserVF, MaxUserVF);
-  if (!UserVF.isZero() &&
-      (UserVFIsLegal || (UserVF.isScalable() && MaxUserVF.isScalable()))) {
-    // FIXME: MaxUserVF is temporarily used inplace of UserVF for illegal
-    // scalable VFs here, this should be reverted to only use legal UserVFs once
-    // the loop below supports scalable VFs.
-    ElementCount VF = UserVFIsLegal ? UserVF : MaxUserVF;
+  if (!UserVF.isZero() && UserVFIsLegal) {
     LLVM_DEBUG(dbgs() << "LV: Using " << (UserVFIsLegal ? "user" : "max")
-                      << " VF " << VF << ".\n");
-    assert(isPowerOf2_32(VF.getKnownMinValue()) &&
+                      << " VF " << UserVF << ".\n");
+    assert(isPowerOf2_32(UserVF.getKnownMinValue()) &&
            "VF needs to be a power of two");
     // Collect the instructions (and their associated costs) that will be more
     // profitable to scalarize.
-    CM.selectUserVectorizationFactor(VF);
+    CM.selectUserVectorizationFactor(UserVF);
     CM.collectInLoopReductions();
-    buildVPlansWithVPRecipes({VF}, {VF});
+    buildVPlansWithVPRecipes({UserVF}, {UserVF});
     LLVM_DEBUG(printPlans(dbgs()));
-    return {{VF, 0}};
+    return {{UserVF, 0}};
   }
 
   OptionalVFCandidates MinFactors(ElementCount::getFixed(1),
                                   ElementCount::getScalable(1));
 
   SmallVector<ElementCount> PowerOf2VFs;
   genAllPowerOfTwoFactorsUptoMaxFactors(PowerOf2VFs, MinFactors, MaxFactors);
   for (auto &VF : PowerOf2VFs) {
if (Optional<ElementCount> FixedMaxVF = MaxFactors.getFixedVF())
     if (FixedMaxVF->isScalar()) {
       assert((!MaxFactors.hasScalableVF() ||
               MaxFactors.getScalableVF()->isScalar()) &&
              "Unexpected max scalable VF");
       return VectorizationFactor::Disabled();
     }
 
   // Select the optimal vectorization factor.
-  return CM.selectVectorizationFactor(MaxFactors);
+  MinFactors = OptionalVFCandidates(ElementCount::getFixed(2),
+                                    ElementCount::getScalable(1));
+  return CM.selectVectorizationFactor(MinFactors, MaxFactors);
 }
 
 void LoopVectorizationPlanner::setBestPlan(ElementCount VF, unsigned UF) {
   LLVM_DEBUG(dbgs() << "Setting best plan to VF=" << VF << ", UF=" << UF
                     << '\n');
   BestVF = VF;
   BestUF = UF;
 

llvm/test/Transforms/LoopVectorize/AArch64/scalable-vf-analysis.ll

 ; REQUIRES: asserts
 ; RUN: opt -mtriple=aarch64-none-linux-gnu -mattr=+sve -enable-scalable-vectorization -loop-vectorize -S -debug-only=loop-vectorize < %s 2>&1 | FileCheck %s
 ; RUN: opt -mtriple=aarch64-none-linux-gnu -mattr=+sve -enable-scalable-vectorization -loop-vectorize -S -debug-only=loop-vectorize -vectorizer-maximize-bandwidth < %s 2>&1 | FileCheck %s --check-prefix=CHECKMAXBW
 
 ; Test that the MaxVF for the following loop, that has no dependence distances,
 ; is calculated as vscale x 4 (max legal SVE vector size) or vscale x 16
 ; (maximized bandwidth for i8 in the loop).
 define void @test0(i32* %a, i8* %b, i32* %c) {
 ; CHECK: LV: Checking a loop in "test0"
 ; CHECK: LV: Found feasible scalable VF = vscale x 4
+; CHECK: LV: Selecting VF: vscale x 4.
 ; CHECKMAXBW: LV: Checking a loop in "test0"
 ; CHECKMAXBW: LV: Found feasible scalable VF = vscale x 16
+; CHECKMAXBW: LV: Selecting VF: vscale x 16.
 entry:
   br label %loop
 
 loop:
   %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop ]
   %arrayidx = getelementptr inbounds i32, i32* %c, i64 %iv
   %0 = load i32, i32* %arrayidx, align 4
   %arrayidx2 = getelementptr inbounds i8, i8* %b, i64 %iv

 exit:
   ret void
 }
 ; Test that the MaxVF for the following loop, with a dependence distance
 ; of 64 elements, is calculated as (maxvscale = 16) * 4.
 define void @test1(i32* %a, i8* %b) {
 ; CHECK: LV: Checking a loop in "test1"
 ; CHECK: LV: Found feasible scalable VF = vscale x 4
+; CHECK: LV: Selecting VF: vscale x 4.
 ; CHECKMAXBW: LV: Checking a loop in "test1"
 ; CHECKMAXBW: LV: Found feasible scalable VF = vscale x 4
+; CHECKMAXBW: LV: Selecting VF: 16.
 entry:
   br label %loop
 
 loop:
   %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop ]
   %arrayidx = getelementptr inbounds i32, i32* %a, i64 %iv
   %0 = load i32, i32* %arrayidx, align 4
   %arrayidx2 = getelementptr inbounds i8, i8* %b, i64 %iv
exit:
   ret void
 }
 
 ; Test that the MaxVF for the following loop, with a dependence distance
 ; of 32 elements, is calculated as (maxvscale = 16) * 2.
 define void @test2(i32* %a, i8* %b) {
 ; CHECK: LV: Checking a loop in "test2"
 ; CHECK: LV: Found feasible scalable VF = vscale x 2
+; CHECK: LV: Selecting VF: vscale x 2.
 ; CHECKMAXBW: LV: Checking a loop in "test2"
 ; CHECKMAXBW: LV: Found feasible scalable VF = vscale x 2
+; CHECKMAXBW: LV: Selecting VF: 16.
 entry:
   br label %loop
 
 loop:
   %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop ]
   %arrayidx = getelementptr inbounds i32, i32* %a, i64 %iv
   %0 = load i32, i32* %arrayidx, align 4
   %arrayidx2 = getelementptr inbounds i8, i8* %b, i64 %iv
exit:
   ret void
 }
 
 ; Test that the MaxVF for the following loop, with a dependence distance
 ; of 16 elements, is calculated as (maxvscale = 16) * 1.
 define void @test3(i32* %a, i8* %b) {
 ; CHECK: LV: Checking a loop in "test3"
 ; CHECK: LV: Found feasible scalable VF = vscale x 1
+; CHECK: LV: Selecting VF: 1.
 ; CHECKMAXBW: LV: Checking a loop in "test3"
 ; CHECKMAXBW: LV: Found feasible scalable VF = vscale x 2
+; CHECKMAXBW: LV: Selecting VF: 16.
 entry:
   br label %loop
 
 loop:
   %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop ]
   %arrayidx = getelementptr inbounds i32, i32* %a, i64 %iv
   %0 = load i32, i32* %arrayidx, align 4
   %arrayidx2 = getelementptr inbounds i8, i8* %b, i64 %iv

llvm/test/Transforms/LoopVectorize/AArch64/scalable-vf-hint.ll

 ; }
 ;
 ; Max fixed VF=32, Max scalable VF=2, unsafe to vectorize. Should clamp to 2.
 
 ; CHECK-DBG-LABEL: LV: Checking a loop in "test4"
 ; CHECK-DBG: LV: The max safe VF is: vscale x 2.
 ; CHECK-DBG: LV: User VF=vscale x 4 is unsafe, clamping to max safe VF=vscale x 2.
 ; CHECK-DBG: remark: <unknown>:0:0: User-specified vectorization factor vscale x 4 is unsafe, clamping to maximum safe vectorization factor vscale x 2
-; CHECK-DBG: LV: Using max VF vscale x 2
+; CHECK-DBG: LV: Selecting VF: vscale x 2.
 ; CHECK-LABEL: @test4
 ; CHECK: <vscale x 2 x i32>
 define void @test4(i32* %a, i32* %b) {
 entry:
   br label %loop
 
 loop:
   %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop ]

 ; }
 ;
 ; Max fixed VF=128, Max scalable VF=8, unsafe to vectorize. Should clamp to 8.
 
 ; CHECK-DBG-LABEL: LV: Checking a loop in "test6"
 ; CHECK-DBG: LV: The max safe VF is: vscale x 8.
 ; CHECK-DBG: LV: User VF=vscale x 16 is unsafe, clamping to max safe VF=vscale x 8.
 ; CHECK-DBG: remark: <unknown>:0:0: User-specified vectorization factor vscale x 16 is unsafe, clamping to maximum safe vectorization factor vscale x 8
-; CHECK-DBG: LV: Using max VF vscale x 8
+; CHECK-DBG: LV: Selecting VF: vscale x 8.
 ; CHECK-LABEL: @test6
 ; CHECK: <vscale x 8 x i32>
 define void @test6(i32* %a, i32* %b) {
 entry:
   br label %loop
 
 loop:
   %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop ]