[llvm] [VPlan] Compute cost of some VPWidenIntOrFpInductions in VPlan (PR #202232)

[David Sherwood via llvm-commits]

================
@@ -5639,7 +5639,36 @@ LoopVectorizationPlanner::precomputeCosts(VPlan &Plan, ElementCount VF,
   // former case) later on, if they exist, to avoid counting them twice.
   // Similarly we pre-compute the cost of any optimized truncates.
   // TODO: Switch to more accurate costing based on VPlan.
+
+  // If the vector loop gets executed exactly once with the given VF, ignore the
+  // costs of comparison and induction instructions, as they'll get simplified
+  // away.
+  // TODO: Remove this code after stepping away from the legacy cost model and
+  // adding code to simplify VPlans before calculating their costs.
+  auto TC = getSmallConstantTripCount(PSE.getSE(), OrigLoop);
+  bool IsFullyUnrolled = TC == VF && !CM.foldTailByMasking();
+  SmallPtrSet<const Value *, 4> WidenedIVs;
+  bool HasTruncatedIV = false;
+  if (IsFullyUnrolled) {
+    addFullyUnrolledInstructionsToIgnore(OrigLoop, Legal->getInductionVars(),
+                                         CostCtx.SkipCostComputation);
+  } else {
+    // Inductions represented by a VPWidenIntOrFpInductionRecipe have their cost
+    // computed by the recipe, so collect their phis to skip the legacy
+    // increment cost below. If any induction is truncated the VPlan-based cost
+    // will diverge. Still fall back to the legacy cost model for now.
+    VPRegionBlock *LoopRegion = Plan.getVectorLoopRegion();
+    for (VPRecipeBase &R : *LoopRegion->getEntryBasicBlock())
+      if (auto *WideIV = dyn_cast<VPWidenIntOrFpInductionRecipe>(&R)) {
+        HasTruncatedIV |= WideIV->getTruncInst() != nullptr;
+        if (PHINode *IVPhi = WideIV->getPHINode())
+          WidenedIVs.insert(IVPhi);
+      }
+  }
+
   for (const auto &[IV, IndDesc] : Legal->getInductionVars()) {
+    if (!HasTruncatedIV && WidenedIVs.contains(IV))
----------------
david-arm wrote:

It looks like if there are both truncated and non-truncated IVs in the loop we will also bail out and fall back on the legacy cost model even for the non-truncated IVs. I guess you've written it this way just to make it easier to write the code and ensure forward progress?

https://github.com/llvm/llvm-project/pull/202232