[LoopVectorize] Return both fixed and scalable Max VF from computeMaxVF.

[sdesmalen]

Summary

This patch introduces a new class, MinMaxVFCandidates, that holds
the possible maximum vectorization factors that have been computed. The
algorithm guarantees there is always at least one maximum VF (in worst
case ElementCount::getFixed(1) for scalarization), but may have both a
fixed MaxVF and a scalable MaxVF.

This patch is intended to be a non-functional change for now,
because LoopVectorizationPlanner::plan doesn't actually use the
scalable MaxVF yet to automatically pick a scalable VPlan
(at least not beyond the UserVF if that's given).

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.

[rogfer01]

Hi Sander, one question OptionalVFCandidates::addMaxVF combines using the maximum but AFAICT it is not used in a loop. Maybe you plan to "update" the corresponding VF values (scalable and/or fixed) in other places (not shown in the current patch, yet)?

[sdesmalen]

In D96025#2565165, @rogfer01 wrote:

Hi Sander, one question OptionalVFCandidates::addMaxVF combines using the maximum but AFAICT it is not used in a loop. Maybe you plan to "update" the corresponding VF values (scalable and/or fixed) in other places (not shown in the current patch, yet)?

Hi Roger, thanks for having a look! This patch is part of a series that allows the vectorizer to reason about a suitable VF (fixed or scalable):

• D96021, D96022, D96023, D96025 are mostly NFC changes that restructure the code to make it easier to reason about both fixed/scalable VFs (and adds the OptionalVFCandidates class).
• D96546 then adds a function genAllPowerOfTwoFactorsUptoMaxFactors, which uses two loops to generate all power-of-two VFs upto the max-VFs. It is still more or less NFC, because it doesn't yet consider the scalable VF for vectorization.
• D96547 let's the cost-model reason about those scalable VFs, so that it can pick a scalable VF as the most profitable VF.

[sdesmalen]

Rebased

[sdesmalen]

Removed unrelated change.

[sdesmalen]

Rebased patch.

[sdesmalen]

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

llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h

bool operator==(const VectorizationFactor &rhs) const {
     return Width == rhs.Width && Cost == rhs.Cost;
   }
 
   bool operator!=(const VectorizationFactor &rhs) const {
     return !(*this == rhs);
   }
 };
 
+/// A class that can hold upto two possible vectorization factors, a fixed VF
+/// and a scalable VF. This can be used by the vectorizer to choose from a range
+/// of fixed and/or scalable VFs in order to find the most cost-effective VF to
+/// vectorize with.
+class MinMaxVFCandidates {
+  Optional<ElementCount> FixedVF;
+  Optional<ElementCount> ScalableVF;
+
+public:
+  /// Initializes both FixedVF/ScalableVF with None
+  MinMaxVFCandidates() : FixedVF({}), ScalableVF({}) {}
+
+  /// Initializes either FixedVF or ScalableVF with VF, the other with None.
+  MinMaxVFCandidates(const ElementCount &VF) : MinMaxVFCandidates() {
+    Optional<ElementCount> *LHSPtr = VF.isScalable() ? &ScalableVF : &FixedVF;
+    *LHSPtr = VF;
+  }
+
+  /// Initializes FixedVF and ScalableVF
+  MinMaxVFCandidates(const ElementCount &FixedVF,
+                     const ElementCount &ScalableVF)
+      : FixedVF(FixedVF), ScalableVF(ScalableVF) {}
+
+  static MinMaxVFCandidates None() { return MinMaxVFCandidates(); }
+
+  /// Assign VF to either FixedVF or ScalableVF, depending on whether VF
+  /// is fixed/scalable. If the VF value is already set, the value is
+  /// overwritten with the largest VF.
+  void addMaxVF(const ElementCount &VF) {
+    Optional<ElementCount> *LHSPtr =
+        VF.isScalable() ? &ScalableVF : &FixedVF;
+    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(); }
+  Optional<ElementCount> getFixedVF() const { return FixedVF; }
+  Optional<ElementCount> getScalableVF() const { return ScalableVF; }
+};
+
 /// Planner drives the vectorization process after having passed
 /// Legality checks.
 class LoopVectorizationPlanner {
   /// The loop that we evaluate.
   Loop *OrigLoop;
 
   /// Loop Info analysis.
   LoopInfo *LI;

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

LoopVectorizationCostModel(ScalarEpilogueLowering SEL, Loop *L,
                              const TargetLibraryInfo *TLI, DemandedBits *DB,
                              AssumptionCache *AC,
                              OptimizationRemarkEmitter *ORE, const Function *F,
                              const LoopVectorizeHints *Hints,
                              InterleavedAccessInfo &IAI)
       : ScalarEpilogueStatus(SEL), TheLoop(L), PSE(PSE), LI(LI), Legal(Legal),
         TTI(TTI), TLI(TLI), DB(DB), AC(AC), ORE(ORE), TheFunction(F),
         Hints(Hints), InterleaveInfo(IAI) {}
-
   /// \return An upper bound for the vectorization factor, or None if
   /// vectorization and interleaving should be avoided up front.
-  Optional<ElementCount> computeMaxVF(ElementCount UserVF, unsigned UserIC);
+  MinMaxVFCandidates computeMaxVF(ElementCount UserVF, unsigned UserIC);
 
   /// \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
reportVectorizationFailure("Runtime stride check for small trip count",
         "this loop without such check by compiling with -Os/-Oz",
         "CantVersionLoopWithOptForSize", ORE, TheLoop);
     return true;
   }
 
   return false;
 }
 
-Optional<ElementCount>
-LoopVectorizationCostModel::computeMaxVF(ElementCount UserVF, unsigned UserIC) {
+MinMaxVFCandidates LoopVectorizationCostModel::computeMaxVF(ElementCount UserVF,
+                                                            unsigned UserIC) {
   if (Legal->getRuntimePointerChecking()->Need && TTI.hasBranchDivergence()) {
     // TODO: It may by useful to do since it's still likely to be dynamically
     // uniform if the target can skip.
     reportVectorizationFailure(
         "Not inserting runtime ptr check for divergent target",
         "runtime pointer checks needed. Not enabled for divergent target",
         "CantVersionLoopWithDivergentTarget", ORE, TheLoop);
-    return None;
+    return MinMaxVFCandidates::None();
   }
 
   unsigned TC = PSE.getSE()->getSmallConstantTripCount(TheLoop);
   LLVM_DEBUG(dbgs() << "LV: Found trip count: " << TC << '\n');
   if (TC == 1) {
     reportVectorizationFailure("Single iteration (non) loop",
         "loop trip count is one, irrelevant for vectorization",
         "SingleIterationLoop", ORE, TheLoop);
-    return None;
+    return MinMaxVFCandidates::None();
   }
 
-  auto GetFeasibleMaxVF = [&]() -> ElementCount {
+  auto GetFeasibleMaxVF = [&]() -> MinMaxVFCandidates {
     MinBWs = computeMinimumValueSizes(TheLoop->getBlocks(), *DB, &TTI);
     unsigned SmallestType, WidestType;
     std::tie(SmallestType, WidestType) = getSmallestAndWidestTypes();
 
     // Get the maximum safe dependence distance in bits computed by LAA.
     // It is computed by MaxVF * sizeOf(type) * 8, where type is taken from
     // the memory accesses that is most restrictive (involved in the smallest
     // dependence distance).
     auto MaxSafeElements = ElementCount::getFixed(
         PowerOf2Floor(Legal->getMaxSafeVectorWidthInBits() / WidestType));
 
-    // Try to find a max scalable VF.
+    MinMaxVFCandidates Result;
+
+    // Always make sure to add at least a max fixed-VF of 1 (scalar) as a
+    // fallback when vectorization is not possible.
+    Result.addMaxVF(ElementCount::getFixed(1));
+
+    // First analyze the UserVF, fall back if the UserVF should be ignored.
+    if (auto MaybeMaxVF = getFeasibleUserVF(UserVF, MaxSafeElements))
+      Result.addMaxVF(*MaybeMaxVF);
+    else {
+      // If we don't force the use of scalable vectors or if the UserVF was
+      // ignored for some reason, we should consider auto-selecting a
+      // fixed-width VF.
+      if (!Hints->isForcedScalable() || UserVF.isNonZero()) {
+        auto MaxSafeVF = clampFeasibleMaxVF(ElementCount::getFixed(1 << 16),
+                                            MaxSafeElements);
+        Result.addMaxVF(
+            computeFeasibleMaxVF(TC, MaxSafeVF, SmallestType, WidestType));
+      }
+
+      // Try to find a max scalable VF as well if the target supports it.
       if (Hints->allowScalable() && TTI.supportsScalableVectors()) {
         auto MaxSafeVF = clampFeasibleMaxVF(ElementCount::getScalable(1 << 16),
                                             MaxSafeElements);
         if (MaxSafeVF.isScalableVector()) {
           ElementCount MaxScalableVF =
               computeFeasibleMaxVF(TC, MaxSafeVF, SmallestType, WidestType);
           if (MaxScalableVF.isScalable())
             LLVM_DEBUG(dbgs() << "LV: Found feasible scalable VF = "
                               << MaxScalableVF << "\n");
           else
             LLVM_DEBUG(dbgs() << "LV: No feasible scalable VF found.\n");
+          Result.addMaxVF(MaxScalableVF);
+        }
       }
     }
 
-    // First analyze the UserVF, fall back if the UserVF should be ignored.
-    if (auto MaybeMaxVF = getFeasibleUserVF(UserVF, MaxSafeElements))
-      return MaybeMaxVF.getValue();
-
-    // Try to automatically determine a suitable maximum VF.
-    auto MaxSafeVF =
-        clampFeasibleMaxVF(ElementCount::getFixed(1 << 16), MaxSafeElements);
-    return computeFeasibleMaxVF(TC, MaxSafeVF, SmallestType, WidestType);
+    return Result;
   };
 
   switch (ScalarEpilogueStatus) {
   case CM_ScalarEpilogueAllowed:
     return GetFeasibleMaxVF();
   case CM_ScalarEpilogueNotAllowedUsePredicate:
     LLVM_FALLTHROUGH;
   case CM_ScalarEpilogueNotNeededUsePredicate:
if (ScalarEpilogueStatus == CM_ScalarEpilogueNotAllowedOptSize)
           dbgs() << "LV: Not allowing scalar epilogue due to -Os/-Oz.\n");
     else
       LLVM_DEBUG(dbgs() << "LV: Not allowing scalar epilogue due to low trip "
                         << "count.\n");
 
     // Bail if runtime checks are required, which are not good when optimising
     // for size.
     if (runtimeChecksRequired())
-      return None;
+      return MinMaxVFCandidates::None();
 
     break;
   }
 
   // The only loops we can vectorize without a scalar epilogue, are loops with
   // a bottom-test and a single exiting block. We'd have to handle the fact
   // that not every instruction executes on the last iteration.  This will
   // require a lane mask which varies through the vector loop body.  (TODO)
   if (TheLoop->getExitingBlock() != TheLoop->getLoopLatch()) {
     // If there was a tail-folding hint/switch, but we can't fold the tail by
     // masking, fallback to a vectorization with a scalar epilogue.
     if (ScalarEpilogueStatus == CM_ScalarEpilogueNotNeededUsePredicate) {
       LLVM_DEBUG(dbgs() << "LV: Cannot fold tail by masking: vectorize with a "
                            "scalar epilogue instead.\n");
       ScalarEpilogueStatus = CM_ScalarEpilogueAllowed;
       return GetFeasibleMaxVF();
     }
-    return None;
+    return MinMaxVFCandidates::None();
   }
 
   // Now try the tail folding
 
   // Invalidate interleave groups that require an epilogue if we can't mask
   // the interleave-group.
   if (!useMaskedInterleavedAccesses(TTI)) {
     assert(WideningDecisions.empty() && Uniforms.empty() && Scalars.empty() &&
            "No decisions should have been taken at this point");
     // Note: There is no need to invalidate any cost modeling decisions here, as
     // non where taken so far.
     InterleaveInfo.invalidateGroupsRequiringScalarEpilogue();
   }
 
-  ElementCount MaxVF = GetFeasibleMaxVF();
-  assert(!MaxVF.isScalable() &&
-         "Scalable vectors do not yet support tail folding");
-  assert((UserVF.isNonZero() || isPowerOf2_32(MaxVF.getFixedValue())) &&
-         "MaxVF must be a power of 2");
-  unsigned MaxVFtimesIC =
-      UserIC ? MaxVF.getFixedValue() * UserIC : MaxVF.getFixedValue();
-  // Avoid tail folding if the trip count is known to be a multiple of any VF we
-  // chose.
+  MinMaxVFCandidates MaxFactors = GetFeasibleMaxVF();
+  if (MaxFactors.hasFixedVF() && !MaxFactors.hasScalableVF()) {
+    ElementCount MaxFixedVF = *MaxFactors.getFixedVF();
+    assert((UserVF.isNonZero() || isPowerOf2_32(MaxFixedVF.getFixedValue())) &&
+           "MaxFixedVF must be a power of 2");
+    unsigned MaxVFtimesIC = UserIC ? MaxFixedVF.getFixedValue() * UserIC
+                                   : MaxFixedVF.getFixedValue();
+    // Avoid tail folding if the trip count is known to be a multiple of any VF
+    // we chose.
     ScalarEvolution *SE = PSE.getSE();
     const SCEV *BackedgeTakenCount = PSE.getBackedgeTakenCount();
     const SCEV *ExitCount = SE->getAddExpr(
         BackedgeTakenCount, SE->getOne(BackedgeTakenCount->getType()));
     const SCEV *Rem = SE->getURemExpr(
         SE->applyLoopGuards(ExitCount, TheLoop),
         SE->getConstant(BackedgeTakenCount->getType(), MaxVFtimesIC));
     if (Rem->isZero()) {
-    // Accept MaxVF if we do not have a tail.
+      // Accept MaxFixedVF if we do not have a tail.
       LLVM_DEBUG(dbgs() << "LV: No tail will remain for any chosen VF.\n");
-    return MaxVF;
+      return MaxFactors;
+    }
   }
 
   // If we don't know the precise trip count, or if the trip count that we
   // found modulo the vectorization factor is not zero, try to fold the tail
   // by masking.
   // FIXME: look for a smaller MaxVF that does divide TC rather than masking.
   if (Legal->prepareToFoldTailByMasking()) {
     FoldTailByMasking = true;
-    return MaxVF;
+    return MaxFactors;
   }
 
   // If there was a tail-folding hint/switch, but we can't fold the tail by
   // masking, fallback to a vectorization with a scalar epilogue.
   if (ScalarEpilogueStatus == CM_ScalarEpilogueNotNeededUsePredicate) {
     LLVM_DEBUG(dbgs() << "LV: Cannot fold tail by masking: vectorize with a "
                          "scalar epilogue instead.\n");
     ScalarEpilogueStatus = CM_ScalarEpilogueAllowed;
-    return MaxVF;
+    return MaxFactors;
   }
 
   if (ScalarEpilogueStatus == CM_ScalarEpilogueNotAllowedUsePredicate) {
     LLVM_DEBUG(dbgs() << "LV: Can't fold tail by masking: don't vectorize\n");
-    return None;
+    return MinMaxVFCandidates::None();
   }
 
   if (TC == 0) {
     reportVectorizationFailure(
         "Unable to calculate the loop count due to complex control flow",
         "unable to calculate the loop count due to complex control flow",
         "UnknownLoopCountComplexCFG", ORE, TheLoop);
-    return None;
+    return MinMaxVFCandidates::None();
   }
 
   reportVectorizationFailure(
       "Cannot optimize for size and vectorize at the same time.",
       "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",
       "NoTailLoopWithOptForSize", ORE, TheLoop);
-  return None;
+  return MinMaxVFCandidates::None();
 }
 
 ElementCount LoopVectorizationCostModel::clampFeasibleMaxVF(
     ElementCount SuggestedVF, ElementCount ClampValue,
     Optional<ElementCount *> OutputMaxValidVF) {
   assert((SuggestedVF.isScalable() || !ClampValue.isScalable()) &&
          "Cannot clamp a fixed-width VF to a scalable VF");
 
LLVM_DEBUG(
       dbgs() << "LV: Not vectorizing. Inner loops aren't supported in the "
                 "VPlan-native path.\n");
   return VectorizationFactor::Disabled();
 }
 
 Optional<VectorizationFactor>
 LoopVectorizationPlanner::plan(ElementCount UserVF, unsigned UserIC) {
   assert(OrigLoop->isInnermost() && "Inner loop expected.");
-  Optional<ElementCount> MaybeMaxVF = CM.computeMaxVF(UserVF, UserIC);
-  if (!MaybeMaxVF) // Cases that should not to be vectorized nor interleaved.
+  MinMaxVFCandidates MaxFactors = CM.computeMaxVF(UserVF, UserIC);
+  if (!MaxFactors) // Cases that should not to be vectorized nor interleaved.
     return None;
 
   // Invalidate interleave groups if all blocks of loop will be predicated.
   if (CM.blockNeedsPredication(OrigLoop->getHeader()) &&
       !useMaskedInterleavedAccesses(*TTI)) {
     LLVM_DEBUG(
         dbgs()
         << "LV: Invalidate all interleaved groups due to fold-tail by masking "
            "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 = MaybeMaxVF.getValue();
+  ElementCount MaxVF =
+      (UserVF && UserVF.isScalable() && MaxFactors.hasScalableVF())
+          ? *MaxFactors.getScalableVF()
+          : *MaxFactors.getFixedVF();
   assert(MaxVF.isNonZero() && "MaxVF is zero.");
 
   bool UserVFIsLegal = ElementCount::isKnownLE(UserVF, MaxVF);
   if (!UserVF.isZero() &&
       (UserVFIsLegal || (UserVF.isScalable() && MaxVF.isScalable()))) {
     // FIXME: MaxVF 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.