[llvm] [VPlan] Handle early exit before forming regions. (NFC) (PR #138393)

[Florian Hahn via llvm-commits]

================
@@ -491,19 +490,38 @@ void VPlanTransforms::prepareForVectorization(VPlan &Plan, Type *InductionTy,
   addCanonicalIVRecipes(Plan, cast<VPBasicBlock>(HeaderVPB),
                         cast<VPBasicBlock>(LatchVPB), InductionTy, IVDL);
 
-  // Disconnect all edges to exit blocks other than from the middle block.
-  // TODO: VPlans with early exits should be explicitly converted to a form
-  // exiting only via the latch here, including adjusting the exit condition,
-  // instead of simply disconnecting the edges and adjusting the VPlan later.
-  for (VPBlockBase *EB : Plan.getExitBlocks()) {
+  [[maybe_unused]] bool HandledUncountableEarlyExit = false;
+  // Handle the remaining early exits, either by converting the plan to one only
+  // exiting via the latch or by disconnecting all early exiting edges and
+  // requiring a scalar epilogue.
+  for (VPIRBasicBlock *EB : Plan.getExitBlocks()) {
     for (VPBlockBase *Pred : to_vector(EB->getPredecessors())) {
       if (Pred == MiddleVPBB)
         continue;
+
+      if (HasUncountableEarlyExit) {
+        assert(!HandledUncountableEarlyExit &&
+               "can handle exactly one uncountable early exit");
+        // Convert VPlans with early exits to a form exiting only via the latch
+        // here, including adjusting the exit condition of the latch.
+        handleUncountableEarlyExit(cast<VPBasicBlock>(Pred), EB, Plan,
+                                   cast<VPBasicBlock>(HeaderVPB),
+                                   cast<VPBasicBlock>(LatchVPB), Range);
+        HandledUncountableEarlyExit = true;
+        continue;
+      }
+
+      // Otherwise all early exits must be countable and we require at least one
+      // iteration in the scalar epilogue. Disconnect all edges to exit blocks
+      // other than from the middle block.
       cast<VPBasicBlock>(Pred)->getTerminator()->eraseFromParent();
       VPBlockUtils::disconnectBlocks(Pred, EB);
     }
   }
 
+  assert((!HasUncountableEarlyExit || HandledUncountableEarlyExit) &&
+         "did not handle uncountable early exit");
----------------
fhahn wrote:

Tried to make this clearer in the message: `missed an uncountable exit that must be handled`

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