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

[via llvm-commits]

================
@@ -491,12 +484,37 @@ void VPlanTransforms::prepareForVectorization(VPlan &Plan, Type *InductionTy,
   VPBlockUtils::insertBlockAfter(VecPreheader, Plan.getEntry());
 
   VPBasicBlock *MiddleVPBB = Plan.createVPBasicBlock("middle.block");
-  VPBlockUtils::connectBlocks(LatchVPB, MiddleVPBB);
-  LatchVPB->swapSuccessors();
+  // Canonical LatchVPB has header block as last successor. If it has another
+  // successor, the latter is an exit block - insert middle block on its edge.
+  // Otherwise, add middle block as another successor retaining header as last.
+  if (LatchVPB->getNumSuccessors() == 2) {
+    VPBlockBase *LatchExitVPB = LatchVPB->getSuccessors()[0];
+    VPBlockUtils::insertOnEdge(LatchVPB, LatchExitVPB, MiddleVPBB);
+  } else {
+    VPBlockUtils::connectBlocks(LatchVPB, MiddleVPBB);
+    LatchVPB->swapSuccessors();
+  }
 
   addCanonicalIVRecipes(Plan, cast<VPBasicBlock>(HeaderVPB),
                         cast<VPBasicBlock>(LatchVPB), InductionTy, IVDL);
 
+  if (HandleUncountableExit) {
+    // Convert VPlans with early exits to a form only exiting via the latch
+    // here, including adjusting the exit condition.
+    handleUncountableEarlyExit(Plan, cast<VPBasicBlock>(HeaderVPB),
+                               cast<VPBasicBlock>(LatchVPB), Range);
+  } else {
+    // Disconnect all edges to exit blocks other than from the middle block.
+    for (VPBlockBase *EB : to_vector(Plan.getExitBlocks())) {
+      for (VPBlockBase *Pred : to_vector(EB->getPredecessors())) {
+        if (Pred == MiddleVPBB)
+          continue;
+        cast<VPBasicBlock>(Pred)->getTerminator()->eraseFromParent();
+        VPBlockUtils::disconnectBlocks(Pred, EB);
----------------
ayalz wrote:

How many such early exit edges from Pred to EB can there be? One if handled by masking/sinking to middle block via handleUncountableEarlyExit(), but possibly more if handled by a required scalar epilog?
```suggestion
        if (SinkSingleEarlyExitToMiddle)
          return sinkSingleEarlyExitToMiddle(Pred, EB, Plan, HeaderVPB, LatchVPB, Range);
        // Else simply remove the edge - it will be taken care of by the scalar epilog.
        cast<VPBasicBlock>(Pred)->getTerminator()->eraseFromParent();
        VPBlockUtils::disconnectBlocks(Pred, EB);
```
?

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