[llvm] [LoopVectorize] Further improve cost model for early exit loops (PR #126235)

[David Sherwood via llvm-commits]

================
@@ -10175,19 +10175,57 @@ static void checkMixedPrecision(Loop *L, OptimizationRemarkEmitter *ORE) {
   }
 }
 
-static bool areRuntimeChecksProfitable(GeneratedRTChecks &Checks,
-                                       VectorizationFactor &VF, Loop *L,
-                                       PredicatedScalarEvolution &PSE,
-                                       ScalarEpilogueLowering SEL,
-                                       std::optional<unsigned> VScale) {
-  InstructionCost CheckCost = Checks.getCost();
-  if (!CheckCost.isValid())
+/// For loops with uncountable early exits, find the cost of doing work when
+/// exiting the loop early, such as calculating the final exit values of
+/// variables used outside the loop.
+/// TODO: This is currently overly pessimistic because the loop may not take
+/// the early exit, but better to keep this conservative for now. In future,
+/// it might be possible to relax this by using branch probabilities.
+static InstructionCost calculateEarlyExitCost(LoopVectorizationCostModel &CM,
+                                              VPlan &Plan, ElementCount VF) {
+  InstructionCost Cost = 0;
+  VPCostContext CostCtx(CM.TTI, *CM.TLI, CM.Legal->getWidestInductionType(), CM,
+                        CM.CostKind);
+  LLVM_DEBUG(
+      dbgs() << "Calculating cost of work in vector early exit block:\n");
----------------
david-arm wrote:

Yep, good suggestion. Thanks!

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