Update to pass assertion comparing two cost models

Author: igogo-x86

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -5559,14 +5559,28 @@ InstructionCost LoopVectorizationCostModel::computePredInstDiscount(
 InstructionCost LoopVectorizationCostModel::expectedCost(ElementCount VF) {
   InstructionCost Cost;
 
+  // If with the given VF loop gets fully unrolled, ignore the costs of
+  // comparison and induction instructions, as they'll get simplified away
+  SmallPtrSet<const Value *, 16> ValuesToIgnoreForVF;
+  auto TC = PSE.getSE()->getSmallConstantTripCount(TheLoop);
+  auto *Cmp = TheLoop->getLatchCmpInst();
+  if (Cmp && TC == VF.getKnownMinValue()) {
+    ValuesToIgnoreForVF.insert(Cmp);
+    for (const auto &[IV, IndDesc] : Legal->getInductionVars()) {
+      Instruction *IVInc = cast<Instruction>(
+          IV->getIncomingValueForBlock(TheLoop->getLoopLatch()));
+      ValuesToIgnoreForVF.insert(IVInc);
+    }
+  }
+
   // For each block.
   for (BasicBlock *BB : TheLoop->blocks()) {
     InstructionCost BlockCost;
 
     // For each instruction in the old loop.
     for (Instruction &I : BB->instructionsWithoutDebug()) {
       // Skip ignored values.
-      if (ValuesToIgnore.count(&I) ||
+      if (ValuesToIgnore.count(&I) || ValuesToIgnoreForVF.count(&I) ||
           (VF.isVector() && VecValuesToIgnore.count(&I)))
         continue;
 
@@ -7249,22 +7263,16 @@ LoopVectorizationPlanner::precomputeCosts(VPlan &Plan, ElementCount VF,
       IVInsts.push_back(CI);
     }
 
-    // If the given VF loop gets fully unrolled, ignore the costs of comparison
-    // and increment instruction, as they'll get simplified away
+    // If with the given VF loop gets fully unrolled, ignore the costs of
+    // comparison and induction instructions, as they'll get simplified away
     auto TC = CM.PSE.getSE()->getSmallConstantTripCount(OrigLoop);
     auto *Cmp = OrigLoop->getLatchCmpInst();
-    if (Cmp && VF.isFixed() && VF.getFixedValue() == TC) {
+    if (Cmp && TC == VF.getKnownMinValue()) {
       CostCtx.SkipCostComputation.insert(Cmp);
-      for (Instruction *IVInst : IVInsts) {
-        bool IsSimplifiedAway = true;
-        for (auto *UIV : IVInst->users()) {
-          if (!Legal->isInductionVariable(UIV) && UIV != Cmp) {
-            IsSimplifiedAway = false;
-            break;
-          }
-        }
-        if (IsSimplifiedAway)
-          CostCtx.SkipCostComputation.insert(IVInst);
+      for (const auto &[IV, IndDesc] : Legal->getInductionVars()) {
+        Instruction *IVInc = cast<Instruction>(
+            IV->getIncomingValueForBlock(OrigLoop->getLoopLatch()));
+        CostCtx.SkipCostComputation.insert(IVInc);
       }
     }