[llvm] 2dbc6d4 - [SLP][NFC]Assert total number of scalar uses not less than number of scalar uses, NFC.

[Alexey Bataev via llvm-commits]

Author: Alexey Bataev
Date: 2024-08-27T09:57:08-07:00
New Revision: 2dbc6d4d4b88c4453acfe51cf20f22847aabeaab

URL: https://github.com/llvm/llvm-project/commit/2dbc6d4d4b88c4453acfe51cf20f22847aabeaab
DIFF: https://github.com/llvm/llvm-project/commit/2dbc6d4d4b88c4453acfe51cf20f22847aabeaab.diff

LOG: [SLP][NFC]Assert total number of scalar uses not less than number of scalar uses, NFC.

Added: 
    

Modified: 
    llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
index be8e2aa99d503d..4613d22fb3acf8 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -10842,17 +10842,22 @@ InstructionCost BoUpSLP::getTreeCost(ArrayRef<Value *> VectorizedVals) {
     // Leave the scalar instructions as is if they are cheaper than extracts.
     if (Entry->Idx != 0 || Entry->getOpcode() == Instruction::GetElementPtr ||
         Entry->getOpcode() == Instruction::Load) {
+      // Checks if the user of the external scalar is phi in loop body.
+      auto IsPhiInLoop = [&](const ExternalUser &U) {
+        if (auto *Phi = dyn_cast_if_present<PHINode>(U.User)) {
+          auto *I = cast<Instruction>(U.Scalar);
+          const Loop *L = LI->getLoopFor(Phi->getParent());
+          return L && (Phi->getParent() == I->getParent() ||
+                       L == LI->getLoopFor(I->getParent()));
+        }
+        return false;
+      };
       if (!ValueToExtUses) {
         ValueToExtUses.emplace();
         for_each(enumerate(ExternalUses), [&](const auto &P) {
           // Ignore phis in loops.
-          if (auto *Phi = dyn_cast_if_present<PHINode>(P.value().User)) {
-            auto *I = cast<Instruction>(P.value().Scalar);
-            const Loop *L = LI->getLoopFor(Phi->getParent());
-            if (L && (Phi->getParent() == I->getParent() ||
-                      L == LI->getLoopFor(I->getParent())))
-              return;
-          }
+          if (IsPhiInLoop(P.value()))
+            return;
 
           ValueToExtUses->try_emplace(P.value().Scalar, P.index());
         });
@@ -10903,8 +10908,12 @@ InstructionCost BoUpSLP::getTreeCost(ArrayRef<Value *> VectorizedVals) {
             return ValueToExtUses->contains(V);
           });
           auto It = ExtractsCount.find(Entry);
-          if (It != ExtractsCount.end())
+          if (It != ExtractsCount.end()) {
+            assert(ScalarUsesCount >= It->getSecond().size() &&
+                   "Expected total number of external uses not less than "
+                   "number of scalar uses.");
             ScalarUsesCount -= It->getSecond().size();
+          }
           // Keep original scalar if number of externally used instructions in
           // the same entry is not power of 2. It may help to do some extra
           // vectorization for now.
@@ -10920,7 +10929,8 @@ InstructionCost BoUpSLP::getTreeCost(ArrayRef<Value *> VectorizedVals) {
             }
           });
           ExtraCost = ScalarCost;
-          ExtractsCount[Entry].insert(Inst);
+          if (!IsPhiInLoop(EU))
+            ExtractsCount[Entry].insert(Inst);
         }
       }
     }