!fixup address comments, thanks

Author: fhahn

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -2380,10 +2380,14 @@ void InnerLoopVectorizer::introduceCheckBlockInVPlan(BasicBlock *CheckIRBB) {
 
   // We just connected a new block to the scalar preheader. Update all
   // VPPhis by adding an incoming value for it, replicating the last value.
+  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 &&
+           "must have incoming values for all operands");
   }
 }
 

llvm/lib/Transforms/Vectorize/VPlan.h

@@ -1122,7 +1122,9 @@ class VPPhiAccessors {
     return getAsRecipe()->getNumOperands();
   }
 
-  void removeIncomingValue(VPBlockBase *VPB) const;
+  /// Removes the incoming value corresponding to \p IncomingBlock, which must
+  /// be a predecessor.
+  void removeIncomingValue(VPBlockBase *IncomingBlock) const;
 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
   /// Print the recipe.
@@ -3717,15 +3719,6 @@ 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

@@ -502,10 +502,8 @@ void VPlanTransforms::prepareForVectorization(
                                    cast<VPBasicBlock>(LatchVPB), Range);
         HandledUncountableEarlyExit = true;
       } else {
-        for (VPRecipeBase &R : cast<VPIRBasicBlock>(EB)->phis()) {
-          if (auto *PhiR = dyn_cast<VPIRPhi>(&R))
-            PhiR->removeIncomingValue(Pred);
-        }
+        for (VPRecipeBase &R : EB->phis())
+          cast<VPIRPhi>(&R)->removeIncomingValue(Pred);
       }
       cast<VPBasicBlock>(Pred)->getTerminator()->eraseFromParent();
       VPBlockUtils::disconnectBlocks(Pred, EB);
@@ -530,44 +528,51 @@ void VPlanTransforms::prepareForVectorization(
   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.
   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->getNumSuccessors() != 2)
+  // 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->getNumSuccessors() == 1) {
+    assert(MiddleVPBB->getSingleSuccessor() == ScalarPH &&
+           "must have ScalarPH as single successor");
     return;
+  }
+
+  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.
+  // Three cases:
+  // 1) If we require a scalar epilogue, the scalar ph must execute. Set the
+  //    condition to false.
+  // 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.
 
-  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.
+  // 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.
+  // E.g., if the compare has got a line number inside the loop.
+  DebugLoc LatchDL = TheLoop->getLoopLatch()->getTerminator()->getDebugLoc();
   VPBuilder Builder(MiddleVPBB);
   VPValue *Cmp;
-  if (TailFolded)
-    Cmp = Plan.getOrAddLiveIn(ConstantInt::getTrue(
-        IntegerType::getInt1Ty(TripCount->getType()->getContext())));
-  else if (!RequiresScalarEpilogueCheck)
+  if (!RequiresScalarEpilogueCheck)
     Cmp = Plan.getOrAddLiveIn(ConstantInt::getFalse(
         IntegerType::getInt1Ty(TripCount->getType()->getContext())));
+  else if (TailFolded)
+    Cmp = Plan.getOrAddLiveIn(ConstantInt::getTrue(
+        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());
+                             &Plan.getVectorTripCount(), LatchDL, "cmp.n");
+  Builder.createNaryOp(VPInstruction::BranchOnCond, {Cmp}, LatchDL);
 }
 
 void VPlanTransforms::createLoopRegions(VPlan &Plan) {

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

@@ -1185,6 +1185,16 @@ void VPIRPhi::execute(VPTransformState &State) {
   State.Builder.SetInsertPoint(Phi->getParent(), std::next(Phi->getIterator()));
 }
 
+void VPPhiAccessors::removeIncomingValue(VPBlockBase *IncomingBlock) const {
+  VPRecipeBase *R = const_cast<VPRecipeBase *>(getAsRecipe());
+  auto &Preds = R->getParent()->getPredecessors();
+  assert(R->getNumOperands() == Preds.size() &&
+         "Number of phi operands must match number of predecessors");
+  unsigned Position = std::distance(Preds.begin(), find(Preds, IncomingBlock));
+  R->getOperand(Position)->removeUser(*R);
+  R->removeOperand(Position);
+}
+
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
 void VPPhiAccessors::printPhiOperands(raw_ostream &O,
                                       VPSlotTracker &SlotTracker) const {

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

@@ -1841,9 +1841,9 @@ void VPlanTransforms::truncateToMinimalBitwidths(
   }
 }
 
-/// Remove BranchOnCond recipes with true conditions together with removing
-/// dead edges to their successors.
-static void removeBranchOnCondTrue(VPlan &Plan) {
+/// Remove BranchOnCond recipes with true or false conditions together with
+/// removing dead edges to their successors.
+static void removeBranchOnConst(VPlan &Plan) {
   using namespace llvm::VPlanPatternMatch;
   for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly<VPBasicBlock>(
            vp_depth_first_shallow(Plan.getEntry()))) {
@@ -1871,7 +1871,8 @@ static void removeBranchOnCondTrue(VPlan &Plan) {
     // from these recipes.
     for (VPRecipeBase &R : make_early_inc_range(*RemovedSucc)) {
       if (isa<VPIRPhi>(&R)) {
-        assert(RemovedSucc->getNumPredecessors() == 1);
+        assert(RemovedSucc->getNumPredecessors() == 1 &&
+               "VPIRPhis must have a single predecessor");
         cast<VPIRPhi>(&R)->removeIncomingValue(VPBB);
         continue;
       }
@@ -1909,7 +1910,7 @@ void VPlanTransforms::optimize(VPlan &Plan) {
   runPass(legalizeAndOptimizeInductions, Plan);
   runPass(removeRedundantExpandSCEVRecipes, Plan);
   runPass(simplifyRecipes, Plan, *Plan.getCanonicalIV()->getScalarType());
-  runPass(removeBranchOnCondTrue, Plan);
+  runPass(removeBranchOnConst, Plan);
   runPass(removeDeadRecipes, Plan);
 
   runPass(createAndOptimizeReplicateRegions, Plan);