[LV] NFC: Return both fixed and scalable Max VF from computeMaxVF.

[sdesmalen]

Summary

This patch introduces a new class, MaxVFCandidates, that holds the
maximum vectorization factors that have been computed for both scalable
and fixed-width vectors.

This patch is NFC and just propagates the information. A follow-up patch
will use this information to create new VPlans.

[sdesmalen]

Rebased patch.

[david-arm]

LGTM!

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h

202

I’’m not sure what the candidates refers to in the name and I couldn’t spot any reference in any of the comments. AFAICT it’s just a pair of VFs, so perhaps there’s no need to make special reference to the fact that it is used for the maximum VFs?

215

Would be good to address?

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

1228

Comment needs updating

1621–1622

Comment needs updating

5668

Do we need the separate variables or could we just use the VF pair and adjust the scalable/fixed VFs as we go along?

5762

I guess this would pessimize the case where fixed width vectorisation would not require a scalar tail, if a scalable VF is present? Might be good to have a comment/TODO? Also might be good to move the comment that this checks if a tail is required up as well

[sdesmalen]

Addressed review comments.

[sdesmalen]

Thanks for the review/comments!

llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h

202

Yes I wasn't too clear on what name to give this. I've renamed it to FixedScalableVFPair, but feel free to suggest an alternative if you know a better one!

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

5762

Yes, that's correct, that's something to fix in a follow-up patch. I actually have a draft patch that addresses this, but I didn't want to complicate this patch any further. I've added a FIXME as you suggested.

[sdesmalen]

@fhahn are you happy with the rename? I've addressed all other comments as well.

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

1229

This does not return none any more? It returns a pair with both set to 0. Perhaps we can just say Returns an upper bound for the vectorization factors (both fixed and scalable). If both factors are 0, vectorization and interleaving should be avoided up front.

5762

Sounds good. Would it be possible to add a test case for that scenario?

[sdesmalen]

• Added test for trying to eliminate vector tail when scalable vectors are considered.
• Updated comment.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

5762

Good suggestion, I've added a test that is XFAIL'ed for now, so that we can re-enable it when the issue is fixed properly.

[fhahn]

LGTM, thanks!

[fhahn]

(The NFC part should probably be dropped from the title, as it changes behaviour as per fixme?)

[sdesmalen]

In D98721#2765191, @fhahn wrote:

(The NFC part should probably be dropped from the title, as it changes behaviour as per fixme?)

Good point, I'll update the commit message before I land it. Thanks for reviewing the patch!

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 represents two maximum vectorization factors (initialized with
+/// 0 by default). One for fixed-width vectorization and one for scalable
+/// vectorization. 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.
+struct MaxVFCandidates {
+  ElementCount FixedVF;
+  ElementCount ScalableVF;
+
+  MaxVFCandidates()
+      : FixedVF(ElementCount::getFixed(0)),
+        ScalableVF(ElementCount::getScalable(0)) {}
+  MaxVFCandidates(const ElementCount &Max) : MaxVFCandidates() {
+    *(Max.isScalable() ? &ScalableVF : &FixedVF) = Max;
+  }
+  MaxVFCandidates(const ElementCount &FixedVF, const ElementCount &ScalableVF)
+      : FixedVF(FixedVF), ScalableVF(ScalableVF) {}
+
+  static MaxVFCandidates None() { return MaxVFCandidates(); }
+
+  /// \return true if either max fixed- or scalable VF is non-zero.
+  explicit operator bool() const { return FixedVF || ScalableVF; }
+
+  /// \return true if either max fixed- or scalable VF is a valid vector VF.
+  bool hasVectorCandidates() const {
+    return FixedVF.isVector() || ScalableVF.isVector();
+  }
+};
+
 /// 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 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);
+  MaxVFCandidates 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
   /// possible.
public:
   }
 
 private:
   unsigned NumPredStores = 0;
 
   /// \return An upper bound for the vectorization factor, a power-of-2 larger
   /// than zero. One is returned if vectorization should best be avoided due
   /// to cost.
-  ElementCount computeFeasibleMaxVF(unsigned ConstTripCount,
+  MaxVFCandidates computeFeasibleMaxVF(unsigned ConstTripCount,
                                        ElementCount UserVF);
 
   /// \return the maximum vector size based on the target's vector registers,
   /// limited to MaxVF. This is a helper function of computeFeasibleMaxVF.
   ElementCount getMaxVectorSize(unsigned ConstTripCount, unsigned SmallestType,
                                 unsigned WidestType, ElementCount MaxVF);
 
   /// \return the maximum legal scalable VF, based on the safe max number
   /// of elements.
if (!MaxScalableVF)
           "Max legal vector width too small, scalable vectorization "
           "unfeasible.",
           "ScalableVFUnfeasible", ORE, TheLoop);
   }
 
   return MaxScalableVF;
 }
 
-ElementCount
+MaxVFCandidates
 LoopVectorizationCostModel::computeFeasibleMaxVF(unsigned ConstTripCount,
                                                  ElementCount UserVF) {
   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
ORE->emit([&]() {
                 "suitable VF.";
     });
   }
 
   LLVM_DEBUG(dbgs() << "LV: The Smallest and Widest types: " << SmallestType
                     << " / " << WidestType << " bits.\n");
 
   ElementCount MaxFixedVF = ElementCount::getFixed(1);
   if (auto Max = getMaxVectorSize(ConstTripCount, SmallestType, WidestType,
                                   MaxSafeFixedVF))
     MaxFixedVF = Max;
 
+  ElementCount MaxScalableVF = ElementCount::getScalable(0);
   if (auto Max = getMaxVectorSize(ConstTripCount, SmallestType, WidestType,
                                   MaxSafeScalableVF))
-    if (Max.isScalable())
+    if (Max.isScalable()) {
+      MaxScalableVF = Max;
       LLVM_DEBUG(dbgs() << "LV: Found feasible scalable VF = " << Max << "\n");
+    }
 
-  return MaxFixedVF;
+  return {MaxFixedVF, MaxScalableVF};
 }
 
-Optional<ElementCount>
-LoopVectorizationCostModel::computeMaxVF(ElementCount UserVF, unsigned UserIC) {
+MaxVFCandidates 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 MaxVFCandidates::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 MaxVFCandidates::None();
   }
 
   switch (ScalarEpilogueStatus) {
   case CM_ScalarEpilogueAllowed:
     return computeFeasibleMaxVF(TC, UserVF);
   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 MaxVFCandidates::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 computeFeasibleMaxVF(TC, UserVF);
     }
-    return None;
+    return MaxVFCandidates::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 = computeFeasibleMaxVF(TC, UserVF);
-  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.
+  MaxVFCandidates MaxFactors = computeFeasibleMaxVF(TC, UserVF);
+  if (MaxFactors.FixedVF.isVector() && !MaxFactors.ScalableVF) {
+    ElementCount MaxFixedVF = MaxFactors.FixedVF;
+    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 MaxVFCandidates::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 MaxVFCandidates::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 MaxVFCandidates::None();
 }
 
 ElementCount LoopVectorizationCostModel::getMaxVectorSize(
     unsigned ConstTripCount, unsigned SmallestType, unsigned WidestType,
     ElementCount MaxSafeVF) {
   bool ComputeScalableMaxVF = MaxSafeVF.isScalable();
   unsigned WidestRegister = ComputeScalableMaxVF
                                 ? TTI.getMaxScalableBitsPerBlock()
ElementCount LoopVectorizationCostModel::getMaxVectorSize(
   auto MaxVectorSize = ElementCount::get(
       PowerOf2Floor(WidestRegister / WidestType), ComputeScalableMaxVF);
   MaxVectorSize = MinVF(MaxVectorSize, MaxSafeVF);
   LLVM_DEBUG(dbgs() << "LV: The Widest register safe to use is: "
                     << (MaxVectorSize * WidestType) << " bits.\n");
 
   if (MaxVectorSize.getKnownMinValue() == 0) {
     LLVM_DEBUG(dbgs() << "LV: The target has no vector registers.\n");
-    return ElementCount::getFixed(1);
+    return MaxVectorSize;
   }
 
   const auto TripCountEC = ElementCount::getFixed(ConstTripCount);
   if (ConstTripCount && ElementCount::isKnownLE(TripCountEC, MaxVectorSize) &&
       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. If MaxVectorSize
     // is scalable, we only fall back on a fixed VF when the TC is less than or
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.
+  MaxVFCandidates 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();
-  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.
-    ElementCount VF = UserVFIsLegal ? UserVF : MaxVF;
+  ElementCount MaxUserVF =
+      UserVF.isScalable() ? MaxFactors.ScalableVF : MaxFactors.FixedVF;
+  bool UserVFIsLegal = ElementCount::isKnownLE(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}};
   }
 
+  ElementCount MaxVF = MaxFactors.FixedVF;
   assert(!MaxVF.isScalable() &&
          "Scalable vectors not yet supported beyond this point");
 
   for (ElementCount VF = ElementCount::getFixed(1);
        ElementCount::isKnownLE(VF, MaxVF); VF *= 2) {
     // 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);
   LLVM_DEBUG(printPlans(dbgs()));
-  if (MaxVF.isScalar())
+  if (!MaxFactors.hasVectorCandidates())
     return VectorizationFactor::Disabled();
 
   // Select the optimal vectorization factor.
   return CM.selectVectorizationFactor(MaxVF);
 }
 
 void LoopVectorizationPlanner::setBestPlan(ElementCount VF, unsigned UF) {
   LLVM_DEBUG(dbgs() << "Setting best plan to VF=" << VF << ", UF=" << UF