!fixup address latest comments, thanks

Author: fhahn

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -9417,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,13 +352,14 @@ std::unique_ptr<VPlan> PlainCFGBuilder::buildPlainCFG(
   Plan->getEntry()->setOneSuccessor(getOrCreateVPBB(TheLoop->getHeader()));
   Plan->getEntry()->setPlan(&*Plan);
 
-  // Add incoming operands for the VPIRInstructions wrapping the exit phis.
+  // 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) {
-      auto *PhiR = dyn_cast<VPIRPhi>(&R);
-      if (!PhiR)
-        break;
+    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)));
@@ -477,8 +478,8 @@ void VPlanTransforms::createLoopRegions(VPlan &Plan, Type *InductionTy,
   VPBlockUtils::connectBlocks(ScalarPH, Plan.getScalarHeader());
   if (!RequiresScalarEpilogueCheck) {
     VPBlockUtils::connectBlocks(MiddleVPBB, ScalarPH);
-    // The exit blocks are dead, remove any recipes to make sure no users remain
-    // that may pessimize transforms.
+    // 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();

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

@@ -1137,9 +1137,9 @@ 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");
+         "can only update exiting operands to phi nodes");
   assert(getNumOperands() > 0 && "must have at least one operand");
   VPValue *Exiting = getOperand(0);
   if (!Exiting->isLiveIn()) {

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

@@ -2503,23 +2503,26 @@ void VPlanTransforms::handleUncountableEarlyExit(
   // Update the exit phis in the early exit block.
   VPBuilder MiddleBuilder(NewMiddle);
   VPBuilder EarlyExitB(VectorEarlyExitVPBB);
-  for (VPRecipeBase &R : *VPEarlyExitBlock->phis()) {
+  for (VPRecipeBase &R : VPEarlyExitBlock->phis()) {
     auto *ExitIRI = cast<VPIRPhi>(&R);
+    // 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()) {
-      // After the transform, the first incoming value is coming from the
-      // orignial loop latch, while the second operand is from the early exit.
-      // Sawp the phi operands, if the first predecessor in the original IR is
-      // not the loop latch.
+      // 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.
-      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);
@@ -2530,15 +2533,14 @@ void VPlanTransforms::handleUncountableEarlyExit(
     // and vector VFs.
     if (!IncomingFromEarlyExit->isLiveIn() &&
         LoopVectorizationPlanner::getDecisionAndClampRange(IsVector, Range)) {
-      // Add the incoming value from the early exit.
+      // Update the incoming value from the early exit.
       VPValue *FirstActiveLane = EarlyExitB.createNaryOp(
           VPInstruction::FirstActiveLane, {EarlyExitTakenCond}, nullptr,
           "first.active.lane");
-      ExitIRI->setOperand(
-          EarlyExitIdx,
-          EarlyExitB.createNaryOp(Instruction::ExtractElement,
-                                  {IncomingFromEarlyExit, FirstActiveLane},
-                                  nullptr, "early.exit.value"));
+      IncomingFromEarlyExit = EarlyExitB.createNaryOp(
+          Instruction::ExtractElement, {IncomingFromEarlyExit, FirstActiveLane},
+          nullptr, "early.exit.value");
+      ExitIRI->setOperand(EarlyExitIdx, IncomingFromEarlyExit);
     }
   }
   MiddleBuilder.createNaryOp(VPInstruction::BranchOnCond, {IsEarlyExitTaken});