[VPlan] Simplify branch to scalar ph in VPlan transform. (NFC)

Simplify branch on false, starting with the branch from the middle block
to the scalar preheader.

Depends on #140405.

Author: fhahn

llvm/lib/Transforms/Vectorize/VPlan.h

@@ -1155,6 +1155,8 @@ class VPPhiAccessors {
     return getAsRecipe()->getNumOperands();
   }
 
+  void removeIncomingValue(VPBlockBase *VPB) const;
+
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
   /// Print the recipe.
   void printPhiOperands(raw_ostream &O, VPSlotTracker &SlotTracker) const;
@@ -3736,6 +3738,15 @@ VPPhiAccessors::getIncomingBlock(unsigned Idx) const {
   return getAsRecipe()->getParent()->getCFGPredecessor(Idx);
 }
 
+inline void VPPhiAccessors::removeIncomingValue(VPBlockBase *VPB) const {
+  VPRecipeBase *R = const_cast<VPRecipeBase *>(getAsRecipe());
+  const VPBasicBlock *Parent = R->getParent();
+  assert(R->getNumOperands() == Parent->getNumPredecessors());
+  auto I = find(Parent->getPredecessors(), VPB);
+  R->getOperand(I - Parent->getPredecessors().begin())->removeUser(*R);
+  R->removeOperand(I - Parent->getPredecessors().begin());
+}
+
 /// A special type of VPBasicBlock that wraps an existing IR basic block.
 /// Recipes of the block get added before the first non-phi instruction in the
 /// wrapped block.

llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp

@@ -507,8 +507,12 @@ void VPlanTransforms::prepareForVectorization(
                                    cast<VPBasicBlock>(HeaderVPB),
                                    cast<VPBasicBlock>(LatchVPB), Range);
         HandledUncountableEarlyExit = true;
+      } else {
+        for (VPRecipeBase &R : cast<VPIRBasicBlock>(EB)->phis()) {
+          if (auto *PhiR = dyn_cast<VPIRPhi>(&R))
+            PhiR->removeIncomingValue(Pred);
+        }
       }
-
       cast<VPBasicBlock>(Pred)->getTerminator()->eraseFromParent();
       VPBlockUtils::disconnectBlocks(Pred, EB);
     }
@@ -541,43 +545,31 @@ void VPlanTransforms::prepareForVectorization(
   //    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) {
-    if (auto *LatchExitVPB = MiddleVPBB->getSingleSuccessor())
-      VPBlockUtils::disconnectBlocks(MiddleVPBB, LatchExitVPB);
-    VPBlockUtils::connectBlocks(MiddleVPBB, ScalarPH);
-    VPBlockUtils::connectBlocks(Plan.getEntry(), ScalarPH);
-    Plan.getEntry()->swapSuccessors();
-
-    // 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;
-  }
-
-  // The connection order corresponds to the operands of the conditional branch,
-  // with the middle block already connected to the exit block.
   VPBlockUtils::connectBlocks(MiddleVPBB, ScalarPH);
   // Also connect the entry block to the scalar preheader.
   VPBlockUtils::connectBlocks(Plan.getEntry(), ScalarPH);
   Plan.getEntry()->swapSuccessors();
 
+  if (MiddleVPBB->getNumSuccessors() != 2)
+    return;
+
   auto *ScalarLatchTerm = TheLoop->getLoopLatch()->getTerminator();
   // Here we use the same DebugLoc as the scalar loop latch terminator instead
   // of the corresponding compare because they may have ended up with
   // different line numbers and we want to avoid awkward line stepping while
   // debugging. Eg. if the compare has got a line number inside the loop.
   VPBuilder Builder(MiddleVPBB);
-  VPValue *Cmp =
-      TailFolded
-          ? Plan.getOrAddLiveIn(ConstantInt::getTrue(
-                IntegerType::getInt1Ty(TripCount->getType()->getContext())))
-          : Builder.createICmp(CmpInst::ICMP_EQ, Plan.getTripCount(),
-                               &Plan.getVectorTripCount(),
-                               ScalarLatchTerm->getDebugLoc(), "cmp.n");
+  VPValue *Cmp;
+  if (TailFolded)
+    Cmp = Plan.getOrAddLiveIn(ConstantInt::getTrue(
+        IntegerType::getInt1Ty(TripCount->getType()->getContext())));
+  else if (!RequiresScalarEpilogueCheck)
+    Cmp = Plan.getOrAddLiveIn(ConstantInt::getFalse(
+        IntegerType::getInt1Ty(TripCount->getType()->getContext())));
+  else
+    Cmp = Builder.createICmp(CmpInst::ICMP_EQ, Plan.getTripCount(),
+                             &Plan.getVectorTripCount(),
+                             ScalarLatchTerm->getDebugLoc(), "cmp.n");
   Builder.createNaryOp(VPInstruction::BranchOnCond, {Cmp},
                        ScalarLatchTerm->getDebugLoc());
 }

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

@@ -1732,11 +1732,21 @@ static void removeBranchOnCondTrue(VPlan &Plan) {
   using namespace llvm::VPlanPatternMatch;
   for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly<VPBasicBlock>(
            vp_depth_first_shallow(Plan.getEntry()))) {
-    if (VPBB->getNumSuccessors() != 2 || isa<VPIRBasicBlock>(VPBB) ||
-        !match(&VPBB->back(), m_BranchOnCond(m_True())))
+    VPValue *Cond;
+    if (VPBB->getNumSuccessors() != 2 || VPBB->empty() ||
+        !match(&VPBB->back(), m_BranchOnCond(m_VPValue(Cond))))
       continue;
 
-    VPBasicBlock *RemovedSucc = cast<VPBasicBlock>(VPBB->getSuccessors()[1]);
+    unsigned RemovedIdx;
+    if (match(Cond, m_True()))
+      RemovedIdx = 1;
+    else if (match(Cond, m_False()))
+      RemovedIdx = 0;
+    else
+      continue;
+
+    VPBasicBlock *RemovedSucc =
+        cast<VPBasicBlock>(VPBB->getSuccessors()[RemovedIdx]);
     const auto &Preds = RemovedSucc->getPredecessors();
     assert(count(Preds, VPBB) == 1 &&
            "There must be a single edge between VPBB and its successor");
@@ -1745,10 +1755,12 @@ static void removeBranchOnCondTrue(VPlan &Plan) {
     // Values coming from VPBB into ResumePhi recipes of RemoveSucc are removed
     // from these recipes.
     for (VPRecipeBase &R : make_early_inc_range(*RemovedSucc)) {
-      assert((!isa<VPIRInstruction>(&R) ||
-              !isa<PHINode>(cast<VPIRInstruction>(&R)->getInstruction())) &&
-             !isa<VPHeaderPHIRecipe>(&R) &&
-             "Cannot update VPIRInstructions wrapping phis or header phis yet");
+      if (isa<VPIRPhi>(&R)) {
+        assert(RemovedSucc->getNumPredecessors() == 1);
+        cast<VPIRPhi>(&R)->removeIncomingValue(VPBB);
+        continue;
+      }
+
       auto *VPI = dyn_cast<VPPhi>(&R);
       if (!VPI)
         break;

llvm/lib/Transforms/Vectorize/VPlanValue.h

@@ -38,6 +38,7 @@ class VPSlotTracker;
 class VPUser;
 class VPRecipeBase;
 class VPInterleaveRecipe;
+class VPPhiAccessors;
 
 // This is the base class of the VPlan Def/Use graph, used for modeling the data
 // flow into, within and out of the VPlan. VPValues can stand for live-ins
@@ -199,8 +200,12 @@ raw_ostream &operator<<(raw_ostream &OS, const VPRecipeBase &R);
 /// This class augments VPValue with operands which provide the inverse def-use
 /// edges from VPValue's users to their defs.
 class VPUser {
+  friend class VPPhiAccessors;
+
   SmallVector<VPValue *, 2> Operands;
 
+  void removeOperand(unsigned Idx) { Operands.erase(Operands.begin() + Idx); }
+
 protected:
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
   /// Print the operands to \p O.