[VPlan] Add exit opernds early.

Author: fhahn

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -9141,11 +9141,7 @@ 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);
+      VPValue *V = ExitIRI->getOperand(0);
       if (V->isLiveIn())
         continue;
       assert(V->getDefiningRecipe()->getParent()->getEnclosingLoopRegion() &&

llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp

@@ -375,6 +375,24 @@ void PlainCFGBuilder::buildPlainCFG(
   VPBlockUtils::connectBlocks(Plan.getEntry(),
                               getOrCreateVPBB(TheLoop->getHeader()));
 
+  for (auto *EB : Plan.getExitBlocks()) {
+    BasicBlock *IRBB = EB->getIRBasicBlock();
+    for (VPRecipeBase &R : *EB) {
+      auto *PhiR = cast<VPIRInstruction>(&R);
+      auto *Phi = dyn_cast<PHINode>(&PhiR->getInstruction());
+      if (!Phi)
+        break;
+      for (Value *Inc : Phi->incoming_values())
+        PhiR->addOperand(getOrCreateVPOperand(Inc));
+      if (R.getNumOperands() > 1 &&
+          Phi->getIncomingBlock(0) != TheLoop->getLoopLatch()) {
+        VPValue *Tmp = R.getOperand(0);
+        R.setOperand(0, R.getOperand(1));
+        R.setOperand(1, Tmp);
+      }
+    }
+  }
+
   for (const auto &[IRBB, VPB] : BB2VPBB)
     VPB2IRBB[VPB] = IRBB;
 
@@ -466,6 +484,11 @@ void VPlanTransforms::introduceTopLevelVectorLoopRegion(
   VPBlockUtils::connectBlocks(ScalarPH, Plan.getScalarHeader());
   if (!RequiresScalarEpilogueCheck) {
     VPBlockUtils::connectBlocks(MiddleVPBB, ScalarPH);
+    for (auto *EB : Plan.getExitBlocks()) {
+      for (VPRecipeBase &R : *EB)
+        for (unsigned Idx = 0; Idx != R.getNumOperands(); ++Idx)
+          R.setOperand(Idx, Plan.getOrAddLiveIn(PoisonValue::get(InductionTy)));
+    }
     return;
   }
 

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

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

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

@@ -2317,29 +2317,27 @@ void VPlanTransforms::handleUncountableEarlyExit(
     if (!ExitIRI)
       break;
 
-    PHINode &ExitPhi = ExitIRI->getIRPhi();
-    VPValue *IncomingFromEarlyExit = RecipeBuilder.getVPValueOrAddLiveIn(
-        ExitPhi.getIncomingValueForBlock(UncountableExitingBlock));
+    unsigned EarlyExitIdx = 0;
 
     if (OrigLoop->getUniqueExitBlock()) {
+      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);
     }
+    VPValue *IncomingFromEarlyExit = ExitIRI->getOperand(EarlyExitIdx);
     // Add the incoming value from the early exit.
     if (!IncomingFromEarlyExit->isLiveIn() && !Plan.hasScalarVFOnly()) {
       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,
+          EarlyExitB.createNaryOp(Instruction::ExtractElement,
+                                  {IncomingFromEarlyExit, FirstActiveLane},
+                                  nullptr, "early.exit.value"));
     }
-    ExitIRI->addOperand(IncomingFromEarlyExit);
   }
   MiddleBuilder.createNaryOp(VPInstruction::BranchOnCond, {IsEarlyExitTaken});