!fixup address latest comments, thanks

Author: fhahn

llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp

@@ -491,36 +491,31 @@ void VPlanTransforms::prepareForVectorization(
                         cast<VPBasicBlock>(LatchVPB), InductionTy, IVDL);
 
   [[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.
+  // Disconnect all early exits from the loop leaving it with a single exit from
+  // the latch. Early exits that are countable are left for a scalar epilog. The
+  // condition of uncountable early exits (currently at most one is supported)
+  // is fused into the latch exit, and used to branch from middle block to the
+  // early exit destination.
   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");
+         "missed an uncountable exit that must be handled");
 
   // Create SCEV and VPValue for the trip count.
   // We use the symbolic max backedge-taken-count, which works also when

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

@@ -2462,7 +2462,6 @@ void VPlanTransforms::convertToConcreteRecipes(VPlan &Plan,
 
 void VPlanTransforms::handleUncountableEarlyExit(VPBasicBlock *EarlyExitingVPBB,
                                                  VPBasicBlock *EarlyExitVPBB,
-
                                                  VPlan &Plan,
                                                  VPBasicBlock *HeaderVPBB,
                                                  VPBasicBlock *LatchVPBB,
@@ -2471,13 +2470,13 @@ void VPlanTransforms::handleUncountableEarlyExit(VPBasicBlock *EarlyExitingVPBB,
 
   VPBlockBase *MiddleVPBB = LatchVPBB->getSuccessors()[0];
   if (!EarlyExitVPBB->getSinglePredecessor() &&
-      EarlyExitVPBB->getPredecessors()[0] != MiddleVPBB) {
+      EarlyExitVPBB->getPredecessors()[1] == MiddleVPBB) {
     assert(EarlyExitVPBB->getNumPredecessors() == 2 &&
-           EarlyExitVPBB->getPredecessors()[1] == MiddleVPBB &&
-           "unsupported earl exit VPBB");
+           EarlyExitVPBB->getPredecessors()[0] == EarlyExitingVPBB &&
+           "unsupported early exit VPBB");
     // Early exit operand should always be last phi operand. If EarlyExitVPBB
-    // has two predecessors and MiddleVPBB isn't the first, swap the operands of
-    // the phis.
+    // has two predecessors and EarlyExitingVPBB is the first, swap the operands
+    // of the phis.
     for (VPRecipeBase &R : EarlyExitVPBB->phis())
       cast<VPIRPhi>(&R)->swapOperands();
   }
@@ -2492,11 +2491,9 @@ void VPlanTransforms::handleUncountableEarlyExit(VPBasicBlock *EarlyExitingVPBB,
   auto *CondToEarlyExit = TrueSucc == EarlyExitVPBB
                               ? CondOfEarlyExitingVPBB
                               : Builder.createNot(CondOfEarlyExitingVPBB);
-  EarlyExitingVPBB->getTerminator()->eraseFromParent();
-  VPBlockUtils::disconnectBlocks(EarlyExitingVPBB, EarlyExitVPBB);
 
   // Split the middle block and have it conditionally branch to the early exit
-  // block if EarlyExitTaken.
+  // block if CondToEarlyExit.
   VPValue *IsEarlyExitTaken =
       Builder.createNaryOp(VPInstruction::AnyOf, {CondToEarlyExit});
   VPBasicBlock *NewMiddle = Plan.createVPBasicBlock("middle.split");

llvm/lib/Transforms/Vectorize/VPlanTransforms.h

@@ -69,7 +69,7 @@ struct VPlanTransforms {
                                       PredicatedScalarEvolution &PSE,
                                       bool RequiresScalarEpilogueCheck,
                                       bool TailFolded, Loop *TheLoop,
-                                      DebugLoc IVDL, bool HandleUncountableExit,
+                                      DebugLoc IVDL, bool HasUncountableExit,
                                       VFRange &Range);
 
   /// Replace loops in \p Plan's flat CFG with VPRegionBlocks, turning \p Plan's
@@ -177,9 +177,9 @@ struct VPlanTransforms {
   /// Update \p Plan to account for the uncountable early exit from \p
   /// EarlyExitingVPBB to \p EarlyExitVPBB by
   ///  * updating the condition exiting the vector loop to include the early
-  ///    exit conditions
+  ///    exit condition,
   ///  * splitting the original middle block to branch to the early exit block
-  ///    if taken.
+  ///    conditionally - according to the early exit condition.
   static void handleUncountableEarlyExit(VPBasicBlock *EarlyExitingVPBB,
                                          VPBasicBlock *EarlyExitVPBB,
                                          VPlan &Plan, VPBasicBlock *HeaderVPBB,