[LoopVectorize] NFC: Split off clamping from computeFeasibleUserVF into its own function.

[sdesmalen]

Summary

This patch is NFC and cleans up computeFeasibleUserVF() by moving the clamping
into a separate function, that can be reused in multiple places, including
computeFeasibleMaxVF. This helps prepare the code to operate on scalable
VFs.

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]

Rebase

[sdesmalen]

• Simplified clampFeasibleMaxVF (it no longer returns a tuple).
• Reuse calculation for MaxSafeVF by passing it to computeFeasibleMaxVF.

[sdesmalen]

Abandoned in favour of D98509

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

Optional<ElementCount> computeFeasibleUserVF(ElementCount UserVF,
                                                unsigned WidestType);
 
   /// \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,
                                     unsigned SmallestType, unsigned WidestType);
 
+  /// \return A pair with the clamped VF and the maximum safe VF.
+  /// If MaxSafeElements allows it, ClampedVF is the SuggestedVF, otherwise
+  /// it is clamped to MaxSafeElements. If SuggestedVF is scalable, the returned
+  /// VF may be clamped to a non-scalable VF.
+  std::pair<ElementCount, ElementCount>
+  clampFeasibleMaxVF(ElementCount SuggestedVF, unsigned MaxSafeElements);
+
   /// 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).
   using VectorizationCostTy = std::pair<InstructionCost, bool>;
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;
 }
 
+std::pair<ElementCount, ElementCount>
+LoopVectorizationCostModel::clampFeasibleMaxVF(ElementCount SuggestedVF,
+                                               unsigned MaxSafeElements) {
+  ElementCount MaxSafeVF = ElementCount::getFixed(MaxSafeElements);
+  if (SuggestedVF.isScalable()) {
+    Optional<unsigned> MaxVScale = TTI.getMaxVScale();
+
+    // Scale VF by vscale before checking if it's safe.
+    MaxSafeVF = ElementCount::getScalable(
+        MaxVScale ? (MaxSafeElements / MaxVScale.getValue()) : 0);
+
+    // Fall back on fixed-width VF.
+    if (MaxSafeVF.isZero())
+      return clampFeasibleMaxVF(
+          ElementCount::getFixed(SuggestedVF.getKnownMinValue()),
+          MaxSafeElements);
+  }
+
+  if (ElementCount::isKnownLE(SuggestedVF, MaxSafeVF))
+    return {SuggestedVF, MaxSafeVF};
+
+  return {MaxSafeVF, MaxSafeVF};
+}
+
 Optional<ElementCount>
 LoopVectorizationCostModel::computeFeasibleUserVF(ElementCount UserVF,
                                                   unsigned WidestType) {
   if (!UserVF.isNonZero())
     return None;
 
   if (UserVF.isScalable() && !TTI.supportsScalableVectors() &&
       !ForceTargetSupportsScalableVectors) {
LoopVectorizationCostModel::computeFeasibleUserVF(ElementCount UserVF,
   unsigned MaxSafeVectorWidthInBits = Legal->getMaxSafeVectorWidthInBits();
   unsigned MaxSafeElements =
       PowerOf2Floor(MaxSafeVectorWidthInBits / WidestType);
 
   // If UserVF is specified and there are dependencies, check if it's legal.
   if (Legal->isSafeForAnyVectorWidth())
     return UserVF;
 
-  ElementCount MaxSafeVF = ElementCount::getFixed(MaxSafeElements);
-
-  if (UserVF.isScalable()) {
-    Optional<unsigned> MaxVScale = TTI.getMaxVScale();
+  // If the user vectorization factor is legally unsafe, clamp it to a safe
+  // value. Otherwise, return as is.
+  ElementCount NewVF = ElementCount::getNull(), MaxSafeVF = NewVF;
+  std::tie(NewVF, MaxSafeVF) = clampFeasibleMaxVF(UserVF, MaxSafeElements);
 
-    // Scale VF by vscale before checking if it's safe.
-    MaxSafeVF = ElementCount::getScalable(
-        MaxVScale ? (MaxSafeElements / MaxVScale.getValue()) : 0);
-
-    if (MaxSafeVF.isZero()) {
-      // The dependence distance is too small to use scalable vectors,
-      // fallback on fixed.
-      LLVM_DEBUG(
-          dbgs()
-          << "LV: Max legal vector width too small, scalable vectorization "
-             "unfeasible. Using fixed-width vectorization instead.\n");
-      ORE->emit([&]() {
-        return OptimizationRemarkAnalysis(DEBUG_TYPE, "ScalableVFUnfeasible",
-                                          TheLoop->getStartLoc(),
-                                          TheLoop->getHeader())
-               << "Max legal vector width too small, scalable vectorization "
-               << "unfeasible. Using fixed-width vectorization instead.";
-      });
-      return computeFeasibleUserVF(
-          ElementCount::getFixed(UserVF.getKnownMinValue()), WidestType);
-    }
-  }
-
-  LLVM_DEBUG(dbgs() << "LV: The max safe VF is: " << MaxSafeVF << ".\n");
-  if (ElementCount::isKnownLE(UserVF, MaxSafeVF))
-    return UserVF;
-
+  // Emit some useful debug output / opt remarks if the user value is clamped.
+  LLVM_DEBUG(dbgs() << "LV: The max safe VF is: " << MaxSafeVF << ".\n");
+  if (UserVF.isScalable() != NewVF.isScalable()) {
+    auto Diag = "Max legal vector width too small, scalable vectorization "
+                "unfeasible. Using fixed-width vectorization instead.";
+    LLVM_DEBUG(dbgs() << "LV: " << Diag << "\n");
+             << Diag;
+  if (NewVF.getKnownMinValue() != UserVF.getKnownMinValue()) {
     LLVM_DEBUG(dbgs() << "LV: User VF=" << UserVF
-                    << " is unsafe, clamping to max safe VF=" << MaxSafeVF
+                      << " is unsafe, clamping to max safe VF=" << NewVF
                       << ".\n");
     ORE->emit([&]() {
       return OptimizationRemarkAnalysis(DEBUG_TYPE, "VectorizationFactor",
                                         TheLoop->getStartLoc(),
                                         TheLoop->getHeader())
              << "User-specified vectorization factor "
              << ore::NV("UserVectorizationFactor", UserVF)
              << " is unsafe, clamping to maximum safe vectorization factor "
-           << ore::NV("VectorizationFactor", MaxSafeVF);
+             << ore::NV("VectorizationFactor", NewVF);
     });
-  return MaxSafeVF;
+  }
+  return NewVF;
 }
 
 ElementCount LoopVectorizationCostModel::computeFeasibleMaxVF(
     unsigned ConstTripCount, unsigned SmallestType, unsigned WidestType) {
   // 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).
   unsigned MaxSafeVectorWidthInBits = Legal->getMaxSafeVectorWidthInBits();
+  unsigned MaxSafeElements =
+      PowerOf2Floor(MaxSafeVectorWidthInBits / WidestType);
 
   unsigned WidestRegister = TTI.getRegisterBitWidth(true);
-  WidestRegister = std::min(WidestRegister, MaxSafeVectorWidthInBits);
-
-  // Ensure MaxVF is a power of 2; the dependence distance bound may not be.
-  // Note that both WidestRegister and WidestType may not be a powers of 2.
-  auto MaxVectorSize =
-      ElementCount::getFixed(PowerOf2Floor(WidestRegister / WidestType));
 
   LLVM_DEBUG(dbgs() << "LV: The Smallest and Widest types: " << SmallestType
                     << " / " << WidestType << " bits.\n");
   LLVM_DEBUG(dbgs() << "LV: The Widest register safe to use is: "
                     << WidestRegister << " bits.\n");
 
-  assert(MaxVectorSize.getFixedValue() <= WidestRegister &&
-         "Did not expect to pack so many elements"
-         " into one vector!");
-  if (MaxVectorSize.getFixedValue() == 0) {
+  // Ensure MaxVF is a power of 2; the dependence distance bound may not be.
+  // Note that both WidestRegister and WidestType may not be a powers of 2.
+  auto MaxVectorSize =
+      ElementCount::getFixed(PowerOf2Floor(WidestRegister / WidestType));
+  std::tie(MaxVectorSize, std::ignore) =
+      clampFeasibleMaxVF(MaxVectorSize, MaxSafeElements);
+
+  if (MaxVectorSize.getKnownMinValue() == 0) {
     LLVM_DEBUG(dbgs() << "LV: The target has no vector registers.\n");
     return ElementCount::getFixed(1);
-  } else if (ConstTripCount && ConstTripCount < MaxVectorSize.getFixedValue() &&
-             isPowerOf2_32(ConstTripCount)) {
+  } else if (ConstTripCount && 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.
+    ElementCount ClampedVF = MaxVectorSize;
+    std::tie(ClampedVF, std::ignore) =
+        clampFeasibleMaxVF(MaxVectorSize, ConstTripCount);
+    if (ClampedVF != MaxVectorSize) {
       LLVM_DEBUG(dbgs() << "LV: Clamping the MaxVF to the constant trip count: "
                         << ConstTripCount << "\n");
-    return ElementCount::getFixed(ConstTripCount);
+      return ClampedVF;
+    }
   }
 
   ElementCount MaxVF = MaxVectorSize;
   if (TTI.shouldMaximizeVectorBandwidth(!isScalarEpilogueAllowed()) ||
       (MaximizeBandwidth && isScalarEpilogueAllowed())) {
+    auto MaxVectorSizeMaxBW =
+        ElementCount::getFixed(PowerOf2Floor(WidestRegister / SmallestType));
+    std::tie(MaxVectorSizeMaxBW, std::ignore) =
+        clampFeasibleMaxVF(MaxVectorSizeMaxBW, MaxSafeElements);
+
     // Collect all viable vectorization factors larger than the default MaxVF
     // (i.e. MaxVectorSize).
     SmallVector<ElementCount, 8> VFs;
-    auto MaxVectorSizeMaxBW =
-        ElementCount::getFixed(WidestRegister / SmallestType);
     for (ElementCount VS = MaxVectorSize * 2;
          ElementCount::isKnownLE(VS, MaxVectorSizeMaxBW); VS *= 2)
       VFs.push_back(VS);
 
     // For each VF calculate its register usage.
     auto RUs = calculateRegisterUsage(VFs);
 
     // Select the largest VF which doesn't require more registers than existing

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

 ;     a[i + 8] = a[i] + b[i];
 ;   }
 ; }
 ;
 ; For scalable vectorization 'vscale' has to be considered, for this example
 ; unless max(vscale)=2 it's unsafe to vectorize. For SVE max(vscale)=16, check
 ; fixed-width vectorization is used instead.
 
+; CHECK-DBG: LV: The max safe VF is: 8.
 ; CHECK-DBG: LV: Max legal vector width too small, scalable vectorization unfeasible. Using fixed-width vectorization instead.
 ; CHECK-DBG: remark: <unknown>:0:0: Max legal vector width too small, scalable vectorization unfeasible. Using fixed-width vectorization instead.
-; CHECK-DBG: LV: The max safe VF is: 8.
 ; CHECK-DBG: LV: Selecting VF: 4.
 ; CHECK-LABEL: @test1
 ; CHECK: <4 x i32>
 define void @test1(i32* %a, i32* %b) {
 entry:
   br label %loop
 
 loop:

 ;
 ; void test2(int *a, int *b, int N) {
 ;   #pragma clang loop vectorize(enable) vectorize_width(8, scalable)
 ;   for (int i=0; i<N; ++i) {
 ;     a[i + 4] = a[i] + b[i];
 ;   }
 ; }
 
-; CHECK-DBG: LV: Max legal vector width too small, scalable vectorization unfeasible. Using fixed-width vectorization instead.
 ; CHECK-DBG: LV: The max safe VF is: 4.
-; CHECK-DBG: LV: User VF=8 is unsafe, clamping to max safe VF=4.
+; CHECK-DBG: LV: Max legal vector width too small, scalable vectorization unfeasible. Using fixed-width vectorization instead.
+; CHECK-DBG: LV: User VF=vscale x 8 is unsafe, clamping to max safe VF=4.
 ; CHECK-DBG: LV: Selecting VF: 4.
 ; CHECK-LABEL: @test2
 ; CHECK: <4 x i32>
 define void @test2(i32* %a, i32* %b) {
 entry:
   br label %loop
 
 loop:

 !18 = !{!18, !19, !20}
 !19 = !{!"llvm.loop.vectorize.width", i32 4}
 !20 = !{!"llvm.loop.vectorize.scalable.enable", i1 true}
 
 ; Test the LV falls back to fixed-width vectorization if scalable vectors are
 ; supported but max vscale is undefined.
 ;
 ; CHECK-NO-MAX-VSCALE-LABEL: LV: Checking a loop in "test_no_max_vscale"
-; CHECK-NO-MAX-VSCALE: LV: Max legal vector width too small, scalable vectorization unfeasible. Using fixed-width vectorization instead.
 ; CEHCK-NO-MAX-VSCALE: The max safe VF is: 4.
+; CHECK-NO-MAX-VSCALE: LV: Max legal vector width too small, scalable vectorization unfeasible. Using fixed-width vectorization instead.
 ; CHECK-NO-MAX-VSCALE: LV: Selecting VF: 4.
 ; CHECK-NO-MAX-VSCALE: <4 x i32>
 define void @test_no_max_vscale(i32* %a, i32* %b) {
 entry:
   br label %loop
 
 loop:
   %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop ]