[LV] Strengthen calls to collectInstsToScalarize (NFC)

[artagnon]

Avoid the pattern of always calling collectInstsToScalarize after collectUniformsAndScalars, and call it in collectUniformsAndScalars instead. Also strengthen checks for early exits in the function.

[llvmbot]

@llvm/pr-subscribers-vectorizers

@llvm/pr-subscribers-llvm-transforms

Author: Ramkumar Ramachandra (artagnon)

Changes

Avoid the pattern of always calling collectInstsToScalarize after collectUniformsAndScalars, and call it in collectUniformsAndScalars instead. Also strengthen checks for early exits in the function.

---

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

1 Files Affected:

• (modified) llvm/lib/Transforms/Vectorize/LoopVectorize.cpp (+10-15)

diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index b987863127994..a4a9f20e5bd71 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -1011,7 +1011,6 @@ class LoopVectorizationCostModel {
   /// \return true if the UserVF is a feasible VF to be chosen.
   bool selectUserVectorizationFactor(ElementCount UserVF) {
     collectUniformsAndScalars(UserVF);
-    collectInstsToScalarize(UserVF);
     return expectedCost(UserVF).isValid();
   }
 
@@ -1281,6 +1280,7 @@ class LoopVectorizationCostModel {
     collectLoopUniforms(VF);
     setVectorizedCallDecision(VF);
     collectLoopScalars(VF);
+    collectInstsToScalarize(VF);
   }
 
   /// Returns true if the target machine supports masked store operation
@@ -5393,11 +5393,14 @@ bool LoopVectorizationCostModel::useEmulatedMaskMemRefHack(Instruction *I,
 }
 
 void LoopVectorizationCostModel::collectInstsToScalarize(ElementCount VF) {
-  // If we aren't vectorizing the loop, or if we've already collected the
-  // instructions to scalarize, there's nothing to do. Collection may already
-  // have occurred if we have a user-selected VF and are now computing the
-  // expected cost for interleaving.
-  if (VF.isScalar() || VF.isZero() || InstsToScalarize.contains(VF))
+  assert(VF.isVector() && "Expected VF >= 2");
+
+  // If we've already collected the instructions to scalarize or the predicated
+  // BBs after vectorization, there's nothing to do. Collection may already have
+  // occurred if we have a user-selected VF and are now computing the expected
+  // cost for interleaving.
+  if (InstsToScalarize.contains(VF) ||
+      PredicatedBBsAfterVectorization.contains(VF))
     return;
 
   // Initialize a mapping for VF in InstsToScalalarize. If we find that it's
@@ -5405,8 +5408,6 @@ void LoopVectorizationCostModel::collectInstsToScalarize(ElementCount VF) {
   // map will indicate that we've analyzed it already.
   ScalarCostsTy &ScalarCostsVF = InstsToScalarize[VF];
 
-  PredicatedBBsAfterVectorization[VF].clear();
-
   // Find all the instructions that are scalar with predication in the loop and
   // determine if it would be better to not if-convert the blocks they are in.
   // If so, we also record the instructions to scalarize.
@@ -7193,16 +7194,10 @@ void LoopVectorizationPlanner::plan(ElementCount UserVF, unsigned UserIC) {
     VFCandidates.push_back(VF);
 
   CM.collectInLoopReductions();
-  for (const auto &VF : VFCandidates) {
+  for (const 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);
-  }
-
   buildVPlansWithVPRecipes(ElementCount::getFixed(1), MaxFactors.FixedVF);
   buildVPlansWithVPRecipes(ElementCount::getScalable(1), MaxFactors.ScalableVF);
 

[david-arm]

LGTM!

[fhahn]

LGTM, thanks