[LV][NFC]Preselect folding style while chosing the max VF, NFC.

[alexey-bataev]

Selects the tail-folding style while choosing the max vector
factor and storing it in the data member rather than calculating it each
time upon getTailFoldingStyle call.

Part of #76172

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Alexey Bataev (alexey-bataev)

Changes

---

Full diff: https://github.com/llvm/llvm-project/pull/81885.diff

1 Files Affected:

• (modified) llvm/lib/Transforms/Vectorize/LoopVectorize.cpp (+23-12)

diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index 98b177cf5d2d0e..f2435601c8fa31 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -1509,15 +1509,25 @@ class LoopVectorizationCostModel {
   }
 
   /// Returns the TailFoldingStyle that is best for the current loop.
-  TailFoldingStyle
-  getTailFoldingStyle(bool IVUpdateMayOverflow = true) const {
-    if (!CanFoldTailByMasking)
-      return TailFoldingStyle::None;
+  TailFoldingStyle getTailFoldingStyle(bool IVUpdateMayOverflow = true) const {
+    return IVUpdateMayOverflow ? ChosenTailFoldingStyle.first
+                               : ChosenTailFoldingStyle.second;
+  }
+
+  void selectTailFoldinStyle() {
+    if (!Legal->prepareToFoldTailByMasking())
+      return;
 
-    if (ForceTailFoldingStyle.getNumOccurrences())
-      return ForceTailFoldingStyle;
+    if (ForceTailFoldingStyle.getNumOccurrences()) {
+      ChosenTailFoldingStyle.first = ChosenTailFoldingStyle.second =
+          ForceTailFoldingStyle;
+      return;
+    }
 
-    return TTI.getPreferredTailFoldingStyle(IVUpdateMayOverflow);
+    ChosenTailFoldingStyle.first =
+        TTI.getPreferredTailFoldingStyle(/*IVUpdateMayOverflow=*/true);
+    ChosenTailFoldingStyle.second =
+        TTI.getPreferredTailFoldingStyle(/*IVUpdateMayOverflow=*/false);
   }
 
   /// Returns true if all loop blocks should be masked to fold tail loop.
@@ -1674,8 +1684,10 @@ class LoopVectorizationCostModel {
   /// iterations to execute in the scalar loop.
   ScalarEpilogueLowering ScalarEpilogueStatus = CM_ScalarEpilogueAllowed;
 
-  /// All blocks of loop are to be masked to fold tail of scalar iterations.
-  bool CanFoldTailByMasking = false;
+  /// Control finally chosen tail folding style. The first element is used if iv
+  /// update may overflow, the second element - if it may not.
+  std::pair<TailFoldingStyle, TailFoldingStyle> ChosenTailFoldingStyle =
+      std::make_pair(TailFoldingStyle::None, TailFoldingStyle::None);
 
   /// A map holding scalar costs for different vectorization factors. The
   /// presence of a cost for an instruction in the mapping indicates that the
@@ -4632,10 +4644,9 @@ LoopVectorizationCostModel::computeMaxVF(ElementCount UserVF, unsigned UserIC) {
   // 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()) {
-    CanFoldTailByMasking = true;
+  selectTailFoldinStyle();
+  if (foldTailByMasking())
     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.

[alexey-bataev]

Ping!

[fhahn]

Thanks for splitting this off! Could you add a brief description including a reference to the PR that will use this?

It is a bit unfortunate that we need to account for 2 possible TF styles, but I don't think there's really any way around that at the moment, as we have to commit to one (or possibly 2) styles up front

[alexey-bataev]

Ping!

[alexey-bataev]

Ping!

[fhahn]

LGTM, thanks!

In the title perhaps: while chosing the max VF ?
In the description, would that be more accurate?

Does the analysis of the tail folding style before choosing max vector

Selects the tail-folding style while choosing the max vector factor?

[ayalz]

It is a bit unfortunate that we need to account for 2 possible TF styles, but I don't think there's really any way around that at the moment, as we have to commit to one (or possibly 2) styles up front

Adding some complementary thoughts:

foldTailByMasking() checks if the default getTailFoldingStyle() is None - implicitly assuming IVUpdateMayOverflow. Worth commenting that checking for None Style can be done independent of IVUpdateMayOverflow, i.e., both styles are None or both of them are not None (right?).

Indeed, a cyclic dependence seems to currently prevent ordering this properly: which TF style is needed depends on IVUpdateMayOverflow which depends on (all VF's up to, but actually only on) MaxVF which is in turn involved in setting the TF styles.
One feasible way to resolve this may be to compute IVUpdateMayOverflow first, based on computeFeasibleMaxVF(), independent of TF styles - assuming conservatively that tail is folded. Going forward, according to roadmap, IVUpdateMayOverflow would relate to the Strip-mine-Factor (VFxUF), possibly splitting it into ranges where IVUpdate may not overflow.

nit: selectTailFoldinStyle() does slightly more than passively "select and record" the style - it also records all masked operations by calling prepareToFoldTailByMasking(), regardless of IVUpdateMayOverflow. And it actually sets two Styles. Perhaps a more accurate name would be setTailFoldingStyles().

[alexey-bataev]

foldTailByMasking() checks if the default getTailFoldingStyle() is None - implicitly assuming IVUpdateMayOverflow. Worth commenting that checking for None Style can be done independent of IVUpdateMayOverflow, i.e., both styles are None or both of them are not None (right?).

Added TODO comment to check infuture that it is so.

nit: selectTailFoldinStyle() does slightly more than passively "select and record" the style - it also records all masked operations by calling prepareToFoldTailByMasking(), regardless of IVUpdateMayOverflow. And it actually sets two Styles. Perhaps a more accurate name would be setTailFoldingStyles().

Done.