[llvm] 43e6f46 - [VPlan] Pre-compute cost for all instrs only feeding exit conditions.

[Florian Hahn via llvm-commits]

Author: Florian Hahn
Date: 2024-06-14T15:12:25+01:00
New Revision: 43e6f46936e177e47de6627a74b047ba27561b44

URL: https://github.com/llvm/llvm-project/commit/43e6f46936e177e47de6627a74b047ba27561b44
DIFF: https://github.com/llvm/llvm-project/commit/43e6f46936e177e47de6627a74b047ba27561b44.diff

LOG: [VPlan] Pre-compute cost for all instrs only feeding exit conditions.

This fixes the following buildbot failures after 90fd99c07957:
    https://lab.llvm.org/buildbot/#/builders/17/builds/47
    https://lab.llvm.org/buildbot/#/builders/168/builds/37

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index 5b652068a7ba9..9fc068a068926 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -7357,16 +7357,30 @@ InstructionCost LoopVectorizationPlanner::cost(VPlan &Plan,
   /// be a single condition to control the vector loop.
   SmallVector<BasicBlock *> Exiting;
   CM.TheLoop->getExitingBlocks(Exiting);
-  // Add the cost of all exit conditions.
+  SetVector<Instruction *> ExitInstrs;
+  // Collect all exit conditions.
   for (BasicBlock *EB : Exiting) {
     auto *Term = dyn_cast<BranchInst>(EB->getTerminator());
     if (!Term)
       continue;
     if (auto *CondI = dyn_cast<Instruction>(Term->getOperand(0))) {
-      assert(!CostCtx.SkipCostComputation.contains(CondI) &&
-             "Condition already skipped?");
-      CostCtx.SkipCostComputation.insert(CondI);
-      Cost += CostCtx.getLegacyCost(CondI, VF);
+      ExitInstrs.insert(CondI);
+    }
+  }
+  // Compute the cost of all instructions only feeding the exit conditions.
+  for (unsigned I = 0; I != ExitInstrs.size(); ++I) {
+    Instruction *CondI = ExitInstrs[I];
+    if (!OrigLoop->contains(CondI) ||
+        !CostCtx.SkipCostComputation.insert(CondI).second)
+      continue;
+    Cost += CostCtx.getLegacyCost(CondI, VF);
+    for (Value *Op : CondI->operands()) {
+      auto *OpI = dyn_cast<Instruction>(Op);
+      if (!OpI || any_of(OpI->users(), [&ExitInstrs](User *U) {
+            return !ExitInstrs.contains(cast<Instruction>(U));
+          }))
+        continue;
+      ExitInstrs.insert(OpI);
     }
   }