!fixup address comments, thanks!

Author: fhahn

llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp

@@ -463,7 +463,7 @@ static void addCanonicalIVRecipes(VPlan &Plan, VPBasicBlock *HeaderVPBB,
 void VPlanTransforms::prepareForVectorization(
     VPlan &Plan, Type *InductionTy, PredicatedScalarEvolution &PSE,
     bool RequiresScalarEpilogueCheck, bool TailFolded, Loop *TheLoop,
-    DebugLoc IVDL, bool HandleUncountableExit, VFRange &Range) {
+    DebugLoc IVDL, bool HasUncountableEarlyExit, VFRange &Range) {
   VPDominatorTree VPDT;
   VPDT.recalculate(Plan);
 
@@ -490,23 +490,35 @@ void VPlanTransforms::prepareForVectorization(
   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);
+  [[maybe_unused]] bool HandledUncountableEarlyExit = false;
+  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");
+
   // Create SCEV and VPValue for the trip count.
   // We use the symbolic max backedge-taken-count, which works also when
   // vectorizing loops with uncountable early exits.

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

@@ -2458,68 +2458,42 @@ void VPlanTransforms::convertToConcreteRecipes(VPlan &Plan,
     R->eraseFromParent();
 }
 
-void VPlanTransforms::handleUncountableEarlyExit(VPlan &Plan,
+void VPlanTransforms::handleUncountableEarlyExit(VPBasicBlock *EarlyExitingVPBB,
+                                                 VPBasicBlock *EarlyExitVPBB,
+
+                                                 VPlan &Plan,
                                                  VPBasicBlock *HeaderVPBB,
                                                  VPBasicBlock *LatchVPBB,
                                                  VFRange &Range) {
-  // First find the uncountable early exiting block by looking at the
-  // predecessors of the exit blocks.
-  VPBlockBase *MiddleVPBB = LatchVPBB->getSuccessors()[0];
-  VPBasicBlock *EarlyExitingVPBB = nullptr;
-  VPIRBasicBlock *EarlyExitVPBB = nullptr;
-  for (auto *EB : Plan.getExitBlocks()) {
-    for (VPBlockBase *Pred : EB->getPredecessors()) {
-      if (Pred != MiddleVPBB) {
-        EarlyExitingVPBB = cast<VPBasicBlock>(Pred);
-        EarlyExitVPBB = EB;
-        break;
-      }
-    }
-  }
-  assert(EarlyExitVPBB && "Must have a early exiting block.");
-  assert(all_of(Plan.getExitBlocks(),
-                [EarlyExitingVPBB, MiddleVPBB](VPIRBasicBlock *EB) {
-                  return all_of(
-                      EB->getPredecessors(),
-                      [EarlyExitingVPBB, MiddleVPBB](VPBlockBase *Pred) {
-                        return Pred == EarlyExitingVPBB || Pred == MiddleVPBB;
-                      });
-                }) &&
-         "All exit blocks must only have EarlyExitingVPBB or MiddleVPBB as "
-         "predecessors.");
-
-  VPBuilder Builder(LatchVPBB->getTerminator());
-  VPBlockBase *TrueSucc = EarlyExitingVPBB->getSuccessors()[0];
-  VPValue *EarlyExitCond = EarlyExitingVPBB->getTerminator()->getOperand(0);
-  auto *EarlyExitTakenCond = TrueSucc == EarlyExitVPBB
-                                 ? EarlyExitCond
-                                 : Builder.createNot(EarlyExitCond);
+  using namespace llvm::VPlanPatternMatch;
 
+  VPBlockBase *MiddleVPBB = LatchVPBB->getSuccessors()[0];
   if (!EarlyExitVPBB->getSinglePredecessor() &&
       EarlyExitVPBB->getPredecessors()[0] != MiddleVPBB) {
-    for (VPRecipeBase &R : EarlyExitVPBB->phis()) {
-      // Early exit operand should always be last, i.e., 0 if EarlyExitVPBB has
-      // a single predecessor and 1 if it has two.
-      // If EarlyExitVPBB has two predecessors, they are already ordered such
-      // that early exit is second (and latch exit is first), by construction.
-      // But its underlying IRBB (EarlyExitIRBB) may have its predecessors
-      // ordered the other way around, and it is the order of the latter which
-      // corresponds to the order of operands of EarlyExitVPBB's phi recipes.
-      // Therefore, if early exit (UncountableExitingBlock) is the first
-      // predecessor of EarlyExitIRBB, we swap the operands of phi recipes,
-      // thereby bringing them to match EarlyExitVPBB's predecessor order,
-      // with early exit being last (second). Otherwise they already match.
+    // 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.
+    for (VPRecipeBase &R : EarlyExitVPBB->phis())
       cast<VPIRPhi>(&R)->swapOperands();
-    }
   }
 
+  VPBuilder Builder(LatchVPBB->getTerminator());
+  VPBlockBase *TrueSucc = EarlyExitingVPBB->getSuccessors()[0];
+  assert(
+      match(EarlyExitingVPBB->getTerminator(), m_BranchOnCond(m_VPValue())) &&
+      "Terminator must be be BranchOnCond");
+  VPValue *CondOfEarlyExitingVPBB =
+      EarlyExitingVPBB->getTerminator()->getOperand(0);
+  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.
   VPValue *IsEarlyExitTaken =
-      Builder.createNaryOp(VPInstruction::AnyOf, {EarlyExitTakenCond});
+      Builder.createNaryOp(VPInstruction::AnyOf, {CondToEarlyExit});
   VPBasicBlock *NewMiddle = Plan.createVPBasicBlock("middle.split");
   VPBasicBlock *VectorEarlyExitVPBB =
       Plan.createVPBasicBlock("vector.early.exit");
@@ -2537,7 +2511,7 @@ void VPlanTransforms::handleUncountableEarlyExit(VPlan &Plan,
     // Early exit operand should always be last, i.e., 0 if EarlyExitVPBB has
     // a single predecessor and 1 if it has two.
     unsigned EarlyExitIdx = ExitIRI->getNumOperands() - 1;
-    if (!EarlyExitVPBB->getSinglePredecessor()) {
+    if (ExitIRI->getNumOperands() != 1) {
       // The first of two operands corresponds to the latch exit, via MiddleVPBB
       // predecessor. Extract its last lane.
       ExitIRI->extractLastLaneOfFirstOperand(MiddleBuilder);
@@ -2553,7 +2527,7 @@ void VPlanTransforms::handleUncountableEarlyExit(VPlan &Plan,
         LoopVectorizationPlanner::getDecisionAndClampRange(IsVector, Range)) {
       // Update the incoming value from the early exit.
       VPValue *FirstActiveLane = EarlyExitB.createNaryOp(
-          VPInstruction::FirstActiveLane, {EarlyExitTakenCond}, nullptr,
+          VPInstruction::FirstActiveLane, {CondToEarlyExit}, nullptr,
           "first.active.lane");
       IncomingFromEarlyExit = EarlyExitB.createNaryOp(
           Instruction::ExtractElement, {IncomingFromEarlyExit, FirstActiveLane},

llvm/lib/Transforms/Vectorize/VPlanTransforms.h

@@ -174,13 +174,15 @@ struct VPlanTransforms {
   /// Remove dead recipes from \p Plan.
   static void removeDeadRecipes(VPlan &Plan);
 
-  /// Update \p Plan to account for the uncountable early exit block in \p
-  /// UncountableExitingBlock by
+  /// 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
   ///  * splitting the original middle block to branch to the early exit block
   ///    if taken.
-  static void handleUncountableEarlyExit(VPlan &Plan, VPBasicBlock *HeaderVPBB,
+  static void handleUncountableEarlyExit(VPBasicBlock *EarlyExitingVPBB,
+                                         VPBasicBlock *EarlyExitVPBB,
+                                         VPlan &Plan, VPBasicBlock *HeaderVPBB,
                                          VPBasicBlock *LatchVPBB,
                                          VFRange &Range);