[LoopVectorize] NFCI: BuildVPlansWithVPRecipes to include ScalableVFs.

[sdesmalen]

Summary

This patch makes sure that the loopvectorizer generates VPlans for
all power of two VFs upto the provided maximum VFs (can be both fixed
and/or scalable). This patch is intended to be a non-functional change,
because selectVectorizationFactor does not automatically pick any of
the scalable VPlans yet.

This patch is a preparatory patch with the ultimate goal of making
computeMaxVF() return both a max fixed VF and a max scalable VF,
so that selectVectorizationFactor() can pick the most cost-effective
vectorization factor.

[sdesmalen]

Rebased

[sdesmalen]

Simplified genFeasibleVFCandidates

[nasherm]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

1979

There is a code path that's being travelled down after your genFeasibleVFCandidates function that causes a runtime crash where VF elements
. The issue is caused by the fact that this code path assumes that the supplied MinFactor and MaxFactor, of type ElementCount, will be in the Scalars map. For illustration here's the piece of code that triggers this error code path

Optional<VectorizationFactor>
7595 LoopVectorizationPlanner::plan(ElementCount UserVF, unsigned UserIC) {
...
7640   genFeasibleVFCandidates(CM, VFCandidates, MinFactors, MaxFactors);
7641   for (auto &VF : VFCandidates) {
7642     // Collect Uniform and Scalar instructions after vectorization with VF.
7643     CM.collectUniformsAndScalars(VF);
7644 
7645     // Collect the instructions (and their associated costs) that will be more
7646     // profitable to scalarize.
7647     if (VF.isVector())
7648       CM.collectInstsToScalarize(VF);
7649   }
7650 
7651   CM.collectInLoopReductions();
7652 
7653   buildVPlansWithVPRecipes(MinFactors, MaxFactors); // <-- code path starts from this call
7654  ....

The buildVPlansWithVPRecipes function will traverse MinFactors->MaxFactors, similar to the genFeasibleVFCandidates candidates, but it will not perform the same !CM.canVectorizeReductions check which allows for the early breakout. As a consequence it will eventually hit this function

8374 bool VPRecipeBuilder::shouldWiden(Instruction *I, VFRange &Range) const {
8375   assert(!isa<BranchInst>(I) && !isa<PHINode>(I) && !isa<LoadInst>(I) &&
8376          !isa<StoreInst>(I) && "Instruction should have been handled earlier");
8377   // Instruction should be widened, unless it is scalar after vectorization,
8378   // scalarization is profitable or it is predicated.
8379   auto WillScalarize = [this, I](ElementCount VF) -> bool {
8380     return CM.isScalarAfterVectorization(I, VF) ||
8381            CM.isProfitableToScalarize(I, VF) ||
8382            CM.isScalarWithPredication(I, VF);
8383   };
8384   return !LoopVectorizationPlanner::getDecisionAndClampRange(WillScalarize,
8385                                                              Range);
8386 }

The call to CM.isScalarVectorization is the source of the error where the assert call is causing the crash. A simple, but perhaps inelegant solution, would be to add the !CM.canVectorizeReductions to the list of checks.

[sdesmalen]

• Rebased patch.
• Removed call to canVectorizeReductions in GenerateRange, as this is now handled in computeFeasibleMaxVF.

[sdesmalen]

Abandoning in favour of a new patch series (see D98509)

llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h

private:
   /// Legal. This method is only used for the legacy inner loop vectorizer.
   VPlanPtr buildVPlanWithVPRecipes(
       VFRange &Range, SmallPtrSetImpl<Instruction *> &DeadInstructions,
       const DenseMap<Instruction *, Instruction *> &SinkAfter);
 
   /// Build VPlans for power-of-2 VF's between \p MinVF and \p MaxVF inclusive,
   /// according to the information gathered by Legal when it checked if it is
   /// legal to vectorize the loop. This method creates VPlans using VPRecipes.
-  void buildVPlansWithVPRecipes(ElementCount MinVF, ElementCount MaxVF);
+  /// When the Maximum VF is set, the corresponding Minimum VF must also be set.
+  void buildVPlansWithVPRecipes(const MinMaxVFCandidates &MinVFs,
+                                const MinMaxVFCandidates &MaxVFs);
 
   /// Adjust the recipes for any inloop reductions. The chain of instructions
   /// leading from the loop exit instr to the phi need to be converted to
   /// reductions, with one operand being vector and the other being the scalar
   /// reduction chain.
   void adjustRecipesForInLoopReductions(VPlanPtr &Plan,
                                         VPRecipeBuilder &RecipeBuilder);
 };
 
 } // namespace llvm
 
 #endif // LLVM_TRANSFORMS_VECTORIZE_LOOPVECTORIZATIONPLANNER_H

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

public:
   /// \return True if runtime checks are required for vectorization, and false
   /// 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(ElementCount MaxVF);
+  VectorizationFactor
+  selectVectorizationFactor(const MinMaxVFCandidates &MaxVF);
   VectorizationFactor
   selectEpilogueVectorizationFactor(const ElementCount MaxVF,
                                     const LoopVectorizationPlanner &LVP);
 
   /// Setup cost-based decisions for user vectorization factor.
   void selectUserVectorizationFactor(ElementCount UserVF) {
     collectUniformsAndScalars(UserVF);
     collectInstsToScalarize(UserVF);
bool isLegalGatherOrScatter(Value *V) {
     auto *Ty = getMemInstValueType(V);
     Align Align = getLoadStoreAlignment(V);
     return (LI && isLegalMaskedGather(Ty, Align)) ||
            (SI && isLegalMaskedScatter(Ty, Align));
   }
 
   /// Returns true if the target machine supports all of the reduction
   /// variables found for the given VF.
-  bool canVectorizeReductions(ElementCount VF) {
+  bool canVectorizeReductions(ElementCount VF) const {
     return (all_of(Legal->getReductionVars(), [&](auto &Reduction) -> bool {
       RecurrenceDescriptor RdxDesc = Reduction.second;
       return TTI.isLegalToVectorizeReduction(RdxDesc, VF);
     }));
   }
 
   /// Returns true if \p I is an instruction that will be scalarized with
   /// predication. Such instructions include conditional stores and
if (!containsIrreducibleCFG<const BasicBlock *>(RPOT, *LI)) {
       // that.
       return;
     }
   }
   for (Loop *InnerL : L)
     collectSupportedLoops(*InnerL, LI, ORE, V);
 }
 
+/// Generate power-of-2 VFs from MinFactors up to and including MaxFactors. It
+/// will generate fixed- and/or scalable VFs depending on which values are set.
+/// For scalable VFs, it is only the minimum known value that is a power of 2,
+/// the eventual VF may not be if the runtime value for vscale is not also a
+/// power of two.
+void genFeasibleVFCandidates(const LoopVectorizationCostModel &CM,
+                             SmallVectorImpl<ElementCount> &List,
+                             const MinMaxVFCandidates &MinFactors,
+                             const MinMaxVFCandidates &MaxFactors) {
+  auto AddToListIfVectorizable = [&](const ElementCount &VF) {
+    if (!CM.canVectorizeReductions(VF))
+      return;
+
+    List.push_back(VF);
+  };
+
+  if (MaxFactors.hasFixedVF())
+    for (ElementCount VF = *MinFactors.getFixedVF(),
+                      Max = *MaxFactors.getFixedVF();
+         ElementCount::isKnownLE(VF, Max); VF *= 2)
+      AddToListIfVectorizable(VF);
+
+  if (MaxFactors.hasScalableVF())
+    for (ElementCount VF = *MinFactors.getScalableVF(),
+                      Max = *MaxFactors.getScalableVF();
+         ElementCount::isKnownLE(VF, Max); VF *= 2)
+      AddToListIfVectorizable(VF);
+}
+
 namespace {
 
 /// The LoopVectorize Pass.
 struct LoopVectorize : public FunctionPass {
   /// Pass identification, replacement for typeid
   static char ID;
 
   LoopVectorizePass Impl;
if (ElementCount MinVF =
                           << ") with target's minimum: " << MinVF << '\n');
         MaxVF = MinVF;
       }
     }
   }
   return MaxVF;
 }
 
-VectorizationFactor
-LoopVectorizationCostModel::selectVectorizationFactor(ElementCount MaxVF) {
+VectorizationFactor LoopVectorizationCostModel::selectVectorizationFactor(
+    const MinMaxVFCandidates &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");
LLVM_DEBUG(
            "which requires masked-interleaved support.\n");
     if (CM.InterleaveInfo.invalidateGroups())
       // Invalidating interleave groups also requires invalidating all decisions
       // based on them, which includes widening decisions and uniform and scalar
       // values.
       CM.invalidateCostModelingDecisions();
   }
 
-  ElementCount MaxVF =
+  ElementCount MaxUserVF =
       (UserVF && UserVF.isScalable() && MaxFactors.hasScalableVF())
           ? *MaxFactors.getScalableVF()
           : *MaxFactors.getFixedVF();
-  assert(MaxVF.isNonZero() && "MaxVF is zero.");
+  assert(MaxUserVF.isNonZero() && "MaxUserVF is zero.");
 
-  bool UserVFIsLegal = ElementCount::isKnownLE(UserVF, MaxVF);
+  bool UserVFIsLegal = ElementCount::isKnownLE(UserVF, MaxUserVF);
   if (!UserVF.isZero() &&
-      (UserVFIsLegal || (UserVF.isScalable() && MaxVF.isScalable()))) {
-    // FIXME: MaxVF is temporarily used inplace of UserVF for illegal
+      (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 : MaxVF;
+    ElementCount VF = UserVFIsLegal ? UserVF : MaxUserVF;
     LLVM_DEBUG(dbgs() << "LV: Using " << (UserVFIsLegal ? "user" : "max")
                       << " VF " << VF << ".\n");
     assert(isPowerOf2_32(VF.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.collectInLoopReductions();
-    buildVPlansWithVPRecipes(VF, VF);
+    buildVPlansWithVPRecipes({VF}, {VF});
     LLVM_DEBUG(printPlans(dbgs()));
     return {{VF, 0}};
   }
 
-  assert(!MaxVF.isScalable() &&
-         "Scalable vectors not yet supported beyond this point");
+  MinMaxVFCandidates MinFactors(ElementCount::getFixed(1),
+                                ElementCount::getScalable(1));
 
-  for (ElementCount VF = ElementCount::getFixed(1);
-       ElementCount::isKnownLE(VF, MaxVF); VF *= 2) {
+  SmallVector<ElementCount> VFCandidates;
+  genFeasibleVFCandidates(CM, VFCandidates, MinFactors, MaxFactors);
+  for (auto &VF : VFCandidates) {
     // Collect Uniform and Scalar instructions after vectorization with VF.
     CM.collectUniformsAndScalars(VF);
 
     // Collect the instructions (and their associated costs) that will be more
     // profitable to scalarize.
     if (VF.isVector())
       CM.collectInstsToScalarize(VF);
   }
 
   CM.collectInLoopReductions();
 
-  buildVPlansWithVPRecipes(ElementCount::getFixed(1), MaxVF);
+  buildVPlansWithVPRecipes(MinFactors, MaxFactors);
+
   LLVM_DEBUG(printPlans(dbgs()));
-  if (MaxVF.isScalar())
+  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(MaxVF);
+  return CM.selectVectorizationFactor(MaxFactors);
 }
 
 void LoopVectorizationPlanner::setBestPlan(ElementCount VF, unsigned UF) {
   LLVM_DEBUG(dbgs() << "Setting best plan to VF=" << VF << ", UF=" << UF
                     << '\n');
   BestVF = VF;
   BestUF = UF;
 
bool InvariantCond =
         PSE.getSE()->isLoopInvariant(PSE.getSCEV(SI->getOperand(0)), OrigLoop);
     return new VPWidenSelectRecipe(*SI, Plan->mapToVPValues(SI->operands()),
                                    InvariantCond);
   }
 
   return tryToWiden(Instr, *Plan);
 }
 
-void LoopVectorizationPlanner::buildVPlansWithVPRecipes(ElementCount MinVF,
-                                                        ElementCount MaxVF) {
+void LoopVectorizationPlanner::buildVPlansWithVPRecipes(
+    const MinMaxVFCandidates &MinVFs, const MinMaxVFCandidates &MaxVFs) {
   assert(OrigLoop->isInnermost() && "Inner loop expected.");
 
   // Collect instructions from the original loop that will become trivially dead
   // in the vectorized loop. We don't need to vectorize these instructions. For
   // example, original induction update instructions can become dead because we
   // separately emit induction "steps" when generating code for the new loop.
   // Similarly, we create a new latch condition when setting up the structure
   // of the new loop, so the old one can become dead.
   SmallPtrSet<Instruction *, 4> DeadInstructions;
   collectTriviallyDeadInstructions(DeadInstructions);
 
   // Add assume instructions we need to drop to DeadInstructions, to prevent
   // them from being added to the VPlan.
   // TODO: We only need to drop assumes in blocks that get flattend. If the
   // control flow is preserved, we should keep them.
   auto &ConditionalAssumes = Legal->getConditionalAssumes();
   DeadInstructions.insert(ConditionalAssumes.begin(), ConditionalAssumes.end());
 
   DenseMap<Instruction *, Instruction *> &SinkAfter = Legal->getSinkAfter();
   // Dead instructions do not need sinking. Remove them from SinkAfter.
   for (Instruction *I : DeadInstructions)
     SinkAfter.erase(I);
 
+  auto CollectVPlans = [&](const ElementCount &MinVF, const ElementCount &MaxVF) {
     auto MaxVFPlusOne = MaxVF.getWithIncrement(1);
     for (ElementCount VF = MinVF; ElementCount::isKnownLT(VF, MaxVFPlusOne);) {
       VFRange SubRange = {VF, MaxVFPlusOne};
       VPlans.push_back(
           buildVPlanWithVPRecipes(SubRange, DeadInstructions, SinkAfter));
       VF = SubRange.End;
     }
+  };
+
+  if (MaxVFs.hasFixedVF()) {
+    assert(MinVFs.hasFixedVF() && "Expected MinVFs.FixedVF to be set");
+    CollectVPlans(*MinVFs.getFixedVF(), *MaxVFs.getFixedVF());
+  }
+
+  if (MaxVFs.hasScalableVF()) {
+    assert(MinVFs.hasScalableVF() && "Expected MinVFs.ScalableVF to be set");
+    CollectVPlans(*MinVFs.getScalableVF(), *MaxVFs.getScalableVF());
+  }
 }
 
 VPlanPtr LoopVectorizationPlanner::buildVPlanWithVPRecipes(
     VFRange &Range, SmallPtrSetImpl<Instruction *> &DeadInstructions,
     const DenseMap<Instruction *, Instruction *> &SinkAfter) {
 
   // Hold a mapping from predicated instructions to their recipes, in order to
   // fix their AlsoPack behavior if a user is determined to replicate and use a