!fixup address latest comments, thanks

Author: fhahn

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -9370,8 +9370,8 @@ static void addScalarResumePhis(VPRecipeBuilder &Builder, VPlan &Plan,
 
 // Collect VPIRInstructions for phis in the exit block from the latch only.
 static SetVector<VPIRInstruction *>
-collectUsersInExitBlocks(Loop *OrigLoop, VPRecipeBuilder &Builder,
-                         VPlan &Plan) {
+collectUsersInLatchExitBlock(Loop *OrigLoop, VPRecipeBuilder &Builder,
+                             VPlan &Plan) {
   SetVector<VPIRInstruction *> ExitUsersToFix;
   for (VPIRBasicBlock *ExitVPBB : Plan.getExitBlocks()) {
     // Nothing to do for unreachable exit blocks.
@@ -9758,7 +9758,7 @@ LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(VFRange &Range) {
   DenseMap<VPValue *, VPValue *> IVEndValues;
   addScalarResumePhis(RecipeBuilder, *Plan, IVEndValues);
   SetVector<VPIRInstruction *> ExitUsersToFix =
-      collectUsersInExitBlocks(OrigLoop, RecipeBuilder, *Plan);
+      collectUsersInLatchExitBlock(OrigLoop, RecipeBuilder, *Plan);
   addExitUsersForFirstOrderRecurrences(*Plan, ExitUsersToFix);
   addUsersInExitBlocks(*Plan, ExitUsersToFix);
 

llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp

@@ -352,8 +352,10 @@ 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
+  // Fix VPlan loop-closed-ssa exit phi's by adding incoming operands to the
   // VPIRInstructions wrapping them.
+  // Note that the operand order may need adjusting when predecessors are added,
+  // if an exit block has multiple predecessor.
   for (auto *EB : Plan->getExitBlocks()) {
     for (VPRecipeBase &R : EB->phis()) {
       auto *PhiR = cast<VPIRPhi>(&R);
@@ -476,6 +478,17 @@ void VPlanTransforms::createLoopRegions(VPlan &Plan, Type *InductionTy,
 
   VPBasicBlock *ScalarPH = Plan.createVPBasicBlock("scalar.ph");
   VPBlockUtils::connectBlocks(ScalarPH, Plan.getScalarHeader());
+
+  // If needed, add a check in the middle block to see if we have completed
+  // all of the iterations in the first vector loop.  Three cases:
+  // 1) If we require a scalar epilogue, there is no conditional branch as
+  //    we unconditionally branch to the scalar preheader.  Remove the recipes
+  //    from the exit blocks.
+  // 2) If (N - N%VF) == N, then we *don't* need to run the remainder.
+  //    Thus if tail is to be folded, we know we don't need to run the
+  //    remainder and we can set the condition to true.
+  // 3) Otherwise, construct a runtime check.
+
   if (!RequiresScalarEpilogueCheck) {
     VPBlockUtils::connectBlocks(MiddleVPBB, ScalarPH);
     // The exit blocks are unreachable, remove their recipes to make sure no
@@ -487,14 +500,6 @@ void VPlanTransforms::createLoopRegions(VPlan &Plan, Type *InductionTy,
     return;
   }
 
-  // If needed, add a check in the middle block to see if we have completed
-  // all of the iterations in the first vector loop.  Three cases:
-  // 1) If (N - N%VF) == N, then we *don't* need to run the remainder.
-  //    Thus if tail is to be folded, we know we don't need to run the
-  //    remainder and we can set the condition to true.
-  // 2) If we require a scalar epilogue, there is no conditional branch as
-  //    we unconditionally branch to the scalar preheader.  Do nothing.
-  // 3) Otherwise, construct a runtime check.
   BasicBlock *IRExitBlock = TheLoop->getUniqueLatchExitBlock();
   auto *VPExitBlock = Plan.getExitBlock(IRExitBlock);
   // The connection order corresponds to the operands of the conditional branch.

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

@@ -2505,23 +2505,25 @@ void VPlanTransforms::handleUncountableEarlyExit(
   VPBuilder EarlyExitB(VectorEarlyExitVPBB);
   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;
+    // Early exit operand should always be last, i.e., 0 if VPEarlyExitBlock has
+    // a single predecessor and 1 if it has two.
+    unsigned EarlyExitIdx = ExitIRI->getNumOperands() - 1;
     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())
+      // If VPEarlyExitBlock 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 VPEarlyExitBlock'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 VPEarlyExitBlock's predecessor order,
+      // with early exit being last (second). Otherwise they already match.
+      if (*pred_begin(VPEarlyExitBlock->getIRBasicBlock()) ==
+          UncountableExitingBlock)
         ExitIRI->swapOperands();
 
-      EarlyExitIdx = 1;
-      // If there's a unique exit block, VPEarlyExitBlock has 2 predecessors
-      // (MiddleVPBB and NewMiddle). Extract the last lane of the incoming value
-      // from MiddleVPBB which is coming from the original latch.
+      // The first of two operands corresponds to the latch exit, via MiddleVPBB
+      // predecessor. Extract its last lane.
       ExitIRI->extractLastLaneOfFirstOperand(MiddleBuilder);
     }