[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:

We can, but then we want to print it before calling the cost function otherwise it would look weird with debug output like:

Cost of ...
Calculating cost of work in vector early exit block

So the only way to do this is to duplicate the loop below and check for an early exit. I can collect all the blocks in a vector and check for non-zero size I suppose. This debug printing issue is actually the reason I used to guard calling this function by querying Legal->hasUncountableEarlyExit(), but now that's been removed we unconditionally call this function.

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