[VPlan] Add exit phi operands during initial construction (NFC).

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -9368,8 +9368,7 @@ static void addScalarResumePhis(VPRecipeBuilder &Builder, VPlan &Plan,
   }
 }
 
-// Collect VPIRInstructions for phis in the exit blocks that are modeled
-// in VPlan and add the exiting VPValue as operand.
+// Collect VPIRInstructions for phis in the exit block from the latch only.
 static SetVector<VPIRInstruction *>
 collectUsersInExitBlocks(Loop *OrigLoop, VPRecipeBuilder &Builder,
                          VPlan &Plan) {
@@ -9388,11 +9387,8 @@ collectUsersInExitBlocks(Loop *OrigLoop, VPRecipeBuilder &Builder,
         continue;
       }
 
-      PHINode &ExitPhi = ExitIRI->getIRPhi();
-      BasicBlock *ExitingBB = OrigLoop->getLoopLatch();
-      Value *IncomingValue = ExitPhi.getIncomingValueForBlock(ExitingBB);
-      VPValue *V = Builder.getVPValueOrAddLiveIn(IncomingValue);
-      ExitIRI->addOperand(V);
+      assert(ExitIRI->getNumOperands() == 1 && "must have a single operand");
+      VPValue *V = ExitIRI->getOperand(0);
       if (V->isLiveIn())
         continue;
       assert(V->getDefiningRecipe()->getParent()->getEnclosingLoopRegion() &&
@@ -9421,7 +9417,7 @@ addUsersInExitBlocks(VPlan &Plan,
            ExitIRI->getParent()->getSinglePredecessor() == MiddleVPBB &&
            "exit values from early exits must be fixed when branch to "
            "early-exit is added");
-    ExitIRI->extractLastLaneOfOperand(B);
+    ExitIRI->extractLastLaneOfFirstOperand(B);
   }
 }
 

llvm/lib/Transforms/Vectorize/VPlan.h

@@ -1161,10 +1161,10 @@ class VPIRInstruction : public VPRecipeBase {
     return true;
   }
 
-  /// Update the recipes single operand to the last lane of the operand using \p
-  /// Builder. Must only be used for single operand VPIRInstructions wrapping a
-  /// PHINode.
-  void extractLastLaneOfOperand(VPBuilder &Builder);
+  /// Update the recipes first operand to the last lane of the operand using \p
+  /// Builder. Must only be used for VPIRInstructions with at least one operand
+  /// wrapping a PHINode.
+  void extractLastLaneOfFirstOperand(VPBuilder &Builder);
 };
 
 /// An overlay for VPIRInstructions wrapping PHI nodes enabling convenient use

llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp

@@ -352,6 +352,20 @@ std::unique_ptr<VPlan> PlainCFGBuilder::buildPlainCFG(
   Plan->getEntry()->setOneSuccessor(getOrCreateVPBB(TheLoop->getHeader()));
   Plan->getEntry()->setPlan(&*Plan);
 
+  // Fix VPlan loop-closed-ssa exit phi's by add incoming operands to the
+  // VPIRInstructions wrapping them.
+  for (auto *EB : Plan->getExitBlocks()) {
+    for (VPRecipeBase &R : EB->phis()) {
+      auto *PhiR = cast<VPIRPhi>(&R);
+      PHINode &Phi = PhiR->getIRPhi();
+      assert(PhiR->getNumOperands() == 0 &&
+             "no phi operands should be added yet");
+      for (BasicBlock *Pred : predecessors(EB->getIRBasicBlock()))
+        PhiR->addOperand(
+            getOrCreateVPOperand(Phi.getIncomingValueForBlock(Pred)));
+    }
+  }
+
   for (const auto &[IRBB, VPB] : BB2VPBB)
     VPB2IRBB[VPB] = IRBB;
 
@@ -464,6 +478,12 @@ void VPlanTransforms::createLoopRegions(VPlan &Plan, Type *InductionTy,
   VPBlockUtils::connectBlocks(ScalarPH, Plan.getScalarHeader());
   if (!RequiresScalarEpilogueCheck) {
     VPBlockUtils::connectBlocks(MiddleVPBB, ScalarPH);
+    // The exit blocks are unreachable, remove their recipes to make sure no
+    // users remain that may pessimize transforms.
+    for (auto *EB : Plan.getExitBlocks()) {
+      for (VPRecipeBase &R : make_early_inc_range(*EB))
+        R.eraseFromParent();
+    }
     return;
   }
 

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

@@ -1137,10 +1137,10 @@ InstructionCost VPIRInstruction::computeCost(ElementCount VF,
   return 0;
 }
 
-void VPIRInstruction::extractLastLaneOfOperand(VPBuilder &Builder) {
+void VPIRInstruction::extractLastLaneOfFirstOperand(VPBuilder &Builder) {
   assert(isa<PHINode>(getInstruction()) &&
-         "can only add exiting operands to phi nodes");
-  assert(getNumOperands() == 1 && "must have a single operand");
+         "can only update exiting operands to phi nodes");
+  assert(getNumOperands() > 0 && "must have at least one operand");
   VPValue *Exiting = getOperand(0);
   if (!Exiting->isLiveIn()) {
     LLVMContext &Ctx = getInstruction().getContext();

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

@@ -2505,35 +2505,43 @@ void VPlanTransforms::handleUncountableEarlyExit(
   VPBuilder EarlyExitB(VectorEarlyExitVPBB);
   for (VPRecipeBase &R : VPEarlyExitBlock->phis()) {
     auto *ExitIRI = cast<VPIRPhi>(&R);
-    PHINode &ExitPhi = ExitIRI->getIRPhi();
-    VPValue *IncomingFromEarlyExit = RecipeBuilder.getVPValueOrAddLiveIn(
-        ExitPhi.getIncomingValueForBlock(UncountableExitingBlock));
-
+    // By default, assume early exit operand is first, e.g., when the two exit
+    // blocks are distinct - VPEarlyExitBlock has a single predecessor.
+    unsigned EarlyExitIdx = 0;
     if (OrigLoop->getUniqueExitBlock()) {
+      // The incoming values currently correspond to the original IR
+      // predecessors. After the transform, the first incoming value is coming
+      // from the original loop latch, while the second operand is from the
+      // early exit. Swap the phi operands, if the first predecessor in the
+      // original IR is not the loop latch.
+      if (*pred_begin(VPEarlyExitBlock->getIRBasicBlock()) !=
+          OrigLoop->getLoopLatch())
+        ExitIRI->swapOperands();
+
+      EarlyExitIdx = 1;
       // If there's a unique exit block, VPEarlyExitBlock has 2 predecessors
-      // (MiddleVPBB and NewMiddle). Add the incoming value from MiddleVPBB
-      // which is coming from the original latch.
-      VPValue *IncomingFromLatch = RecipeBuilder.getVPValueOrAddLiveIn(
-          ExitPhi.getIncomingValueForBlock(OrigLoop->getLoopLatch()));
-      ExitIRI->addOperand(IncomingFromLatch);
-      ExitIRI->extractLastLaneOfOperand(MiddleBuilder);
+      // (MiddleVPBB and NewMiddle). Extract the last lane of the incoming value
+      // from MiddleVPBB which is coming from the original latch.
+      ExitIRI->extractLastLaneOfFirstOperand(MiddleBuilder);
     }
 
+    VPValue *IncomingFromEarlyExit = ExitIRI->getOperand(EarlyExitIdx);
     auto IsVector = [](ElementCount VF) { return VF.isVector(); };
     // When the VFs are vectors, need to add `extract` to get the incoming value
     // from early exit. When the range contains scalar VF, limit the range to
     // scalar VF to prevent mis-compilation for the range containing both scalar
     // and vector VFs.
     if (!IncomingFromEarlyExit->isLiveIn() &&
         LoopVectorizationPlanner::getDecisionAndClampRange(IsVector, Range)) {
+      // Update the incoming value from the early exit.
       VPValue *FirstActiveLane = EarlyExitB.createNaryOp(
           VPInstruction::FirstActiveLane, {EarlyExitTakenCond}, nullptr,
           "first.active.lane");
       IncomingFromEarlyExit = EarlyExitB.createNaryOp(
           Instruction::ExtractElement, {IncomingFromEarlyExit, FirstActiveLane},
           nullptr, "early.exit.value");
+      ExitIRI->setOperand(EarlyExitIdx, IncomingFromEarlyExit);
     }
-    ExitIRI->addOperand(IncomingFromEarlyExit);
   }
   MiddleBuilder.createNaryOp(VPInstruction::BranchOnCond, {IsEarlyExitTaken});