!fixup address latest comments, thanks

Author: fhahn

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -2383,11 +2383,10 @@ void InnerLoopVectorizer::introduceCheckBlockInVPlan(BasicBlock *CheckIRBB) {
   unsigned NumPredecessors = ScalarPH->getNumPredecessors();
   (void)NumPredecessors;
   for (VPRecipeBase &R : cast<VPBasicBlock>(ScalarPH)->phis()) {
-    auto *ResumePhi = cast<VPPhi>(&R);
-    ResumePhi->addOperand(
-        ResumePhi->getOperand(ResumePhi->getNumOperands() - 1));
-    assert(ResumePhi->getNumIncoming() == NumPredecessors &&
+    assert(isa<VPPhi>(&R) && "Phi expected to be VPPhi");
+    assert(cast<VPPhi>(&R)->getNumIncoming() == NumPredecessors - 1 &&
            "must have incoming values for all operands");
+    R.addOperand(R.getOperand(NumPredecessors - 2));
   }
 }
 

llvm/lib/Transforms/Vectorize/VPlan.h

@@ -1122,9 +1122,9 @@ class VPPhiAccessors {
     return getAsRecipe()->getNumOperands();
   }
 
-  /// Removes the incoming value corresponding to \p IncomingBlock, which must
-  /// be a predecessor.
-  void removeIncomingValue(VPBlockBase *IncomingBlock) const;
+  /// Removes the incoming value for \p IncomingBlock, which must be a
+  /// predecessor.
+  void removeIncomingValueFor(VPBlockBase *IncomingBlock) const;
 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
   /// Print the recipe.
@@ -3530,22 +3530,22 @@ class VPScalarIVStepsRecipe : public VPRecipeWithIRFlags,
 
 /// Casting from VPRecipeBase -> VPPhiAccessors is supported for all recipe
 /// types implementing VPPhiAccessors. Used by isa<> & co.
-template <> struct CastIsPossible<VPPhiAccessors, const VPRecipeBase *> {
-  static inline bool isPossible(const VPRecipeBase *f) {
+template <typename SrcTy> struct CastIsPossible<VPPhiAccessors, SrcTy> {
+  static inline bool isPossible(SrcTy f) {
     // TODO: include VPPredInstPHIRecipe too, once it implements VPPhiAccessors.
     return isa<VPIRPhi, VPHeaderPHIRecipe, VPWidenPHIRecipe, VPPhi>(f);
   }
 };
 /// Support casting from VPRecipeBase -> VPPhiAccessors, by down-casting to the
 /// recipe types implementing VPPhiAccessors. Used by cast<>, dyn_cast<> & co.
-template <>
-struct CastInfo<VPPhiAccessors, const VPRecipeBase *>
-    : public CastIsPossible<VPPhiAccessors, const VPRecipeBase *> {
+template <typename SrcTy>
+struct CastInfo<VPPhiAccessors, SrcTy>
+    : public CastIsPossible<VPPhiAccessors, SrcTy> {
 
-  using Self = CastInfo<VPPhiAccessors, const VPRecipeBase *>;
+  using Self = CastInfo<VPPhiAccessors, SrcTy>;
 
   /// doCast is used by cast<>.
-  static inline VPPhiAccessors *doCast(const VPRecipeBase *R) {
+  static inline VPPhiAccessors *doCast(SrcTy R) {
     return const_cast<VPPhiAccessors *>([R]() -> const VPPhiAccessors * {
       switch (R->getVPDefID()) {
       case VPDef::VPInstructionSC:
@@ -3561,7 +3561,7 @@ struct CastInfo<VPPhiAccessors, const VPRecipeBase *>
   }
 
   /// doCastIfPossible is used by dyn_cast<>.
-  static inline VPPhiAccessors *doCastIfPossible(const VPRecipeBase *f) {
+  static inline VPPhiAccessors *doCastIfPossible(SrcTy f) {
     if (!Self::isPossible(f))
       return nullptr;
     return doCast(f);

llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp

@@ -503,7 +503,7 @@ void VPlanTransforms::prepareForVectorization(
         HandledUncountableEarlyExit = true;
       } else {
         for (VPRecipeBase &R : EB->phis())
-          cast<VPIRPhi>(&R)->removeIncomingValue(Pred);
+          cast<VPIRPhi>(&R)->removeIncomingValueFor(Pred);
       }
       cast<VPBasicBlock>(Pred)->getTerminator()->eraseFromParent();
       VPBlockUtils::disconnectBlocks(Pred, EB);
@@ -528,16 +528,18 @@ void VPlanTransforms::prepareForVectorization(
   VPBasicBlock *ScalarPH = Plan.createVPBasicBlock("scalar.ph");
   VPBlockUtils::connectBlocks(ScalarPH, Plan.getScalarHeader());
 
+  // 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.
   // TODO: Also introduce a branch recipe together with the minimum trip count
   // check.
   VPBlockUtils::connectBlocks(Plan.getEntry(), ScalarPH);
   Plan.getEntry()->swapSuccessors();
 
-  // If MiddleVPBB has a single successor the original loop exits via the latch
-  // and the single successor must be the scalar preheader. There's no need to
-  // add a runtime check to MiddleVPBB.
+  // If MiddleVPBB has a single successor then the original loop does not exit
+  // via the latch and the single successor must be the scalar preheader.
+  // There's no need to add a runtime check to MiddleVPBB.
   if (MiddleVPBB->getNumSuccessors() == 1) {
     assert(MiddleVPBB->getSingleSuccessor() == ScalarPH &&
            "must have ScalarPH as single successor");
@@ -546,8 +548,9 @@ void VPlanTransforms::prepareForVectorization(
 
   assert(MiddleVPBB->getNumSuccessors() == 2 && "must have 2 successors");
 
-  // If needed, add a check in the middle block to see if we have completed
-  // all of the iterations in the first vector loop.
+  // 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, the scalar ph must execute. Set the
   //    condition to false.

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

@@ -1185,7 +1185,7 @@ void VPIRPhi::execute(VPTransformState &State) {
   State.Builder.SetInsertPoint(Phi->getParent(), std::next(Phi->getIterator()));
 }
 
-void VPPhiAccessors::removeIncomingValue(VPBlockBase *IncomingBlock) const {
+void VPPhiAccessors::removeIncomingValueFor(VPBlockBase *IncomingBlock) const {
   VPRecipeBase *R = const_cast<VPRecipeBase *>(getAsRecipe());
   auto &Preds = R->getParent()->getPredecessors();
   assert(R->getNumOperands() == Preds.size() &&

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

@@ -1865,32 +1865,15 @@ static void removeBranchOnConst(VPlan &Plan) {
     const auto &Preds = RemovedSucc->getPredecessors();
     assert(count(Preds, VPBB) == 1 &&
            "There must be a single edge between VPBB and its successor");
-    unsigned DeadIdx = std::distance(Preds.begin(), find(Preds, VPBB));
-
-    // Values coming from VPBB into ResumePhi recipes of RemoveSucc are removed
-    // from these recipes.
+    // Values coming from VPBB into phi recipes of RemoveSucc are removed from
+    // these recipes.
     for (VPRecipeBase &R : make_early_inc_range(*RemovedSucc)) {
-      if (isa<VPIRPhi>(&R)) {
-        assert(RemovedSucc->getNumPredecessors() == 1 &&
-               "VPIRPhis must have a single predecessor");
-        cast<VPIRPhi>(&R)->removeIncomingValue(VPBB);
-        continue;
-      }
-
-      auto *VPI = dyn_cast<VPPhi>(&R);
-      if (!VPI)
+      auto *Phi = dyn_cast<VPPhiAccessors>(&R);
+      if (!Phi)
         break;
-      VPBuilder B(VPI);
-      SmallVector<VPValue *> NewOperands;
-      // Create new operand list, with the dead incoming value filtered out.
-      for (const auto &[Idx, Op] : enumerate(VPI->operands())) {
-        if (Idx == DeadIdx)
-          continue;
-        NewOperands.push_back(Op);
-      }
-      VPI->replaceAllUsesWith(
-          B.createScalarPhi(NewOperands, VPI->getDebugLoc(), VPI->getName()));
-      VPI->eraseFromParent();
+      assert((!isa<VPIRPhi>(&R) || RemovedSucc->getNumPredecessors() == 1) &&
+             "VPIRPhis must have a single predecessor");
+      Phi->removeIncomingValueFor(VPBB);
     }
     // Disconnect blocks and remove the terminator. RemovedSucc will be deleted
     // automatically on VPlan destruction if it becomes unreachable.

llvm/lib/Transforms/Vectorize/VPlanValue.h

@@ -200,6 +200,7 @@ 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 {
+  /// Grant access to removeOperand for VPPhiAccessors, the only supported user.
   friend class VPPhiAccessors;
 
   SmallVector<VPValue *, 2> Operands;