!fixup address latest comments, thanks!

Author: fhahn

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -2486,7 +2486,7 @@ void InnerLoopVectorizer::introduceCheckBlockInVPlan(BasicBlock *CheckIRBB) {
   PreVectorPH->swapSuccessors();
 
   // We just connected a new block to the scalar preheader. Update all
-  // ResumePhis by adding an incoming value for it, replacing the last value.
+  // ResumePhis by adding an incoming value for it, replicating the last value.
   for (VPRecipeBase &R : *cast<VPBasicBlock>(ScalarPH)) {
     auto *ResumePhi = dyn_cast<VPInstruction>(&R);
     if (!ResumePhi || ResumePhi->getOpcode() != VPInstruction::ResumePhi)
@@ -2661,7 +2661,7 @@ void InnerLoopVectorizer::createVectorLoopSkeleton(StringRef Prefix) {
   LoopScalarPreHeader =
       SplitBlock(LoopVectorPreHeader, LoopVectorPreHeader->getTerminator(), DT,
                  LI, nullptr, Twine(Prefix) + "scalar.ph");
-  // NOTE: The Plan's scalar preheader isn't replaced with a VPIRBasicBlock
+  // NOTE: The Plan's scalar preheader VPBB isn't replaced with a VPIRBasicBlock
   // wrapping LoopScalarPreHeader here at the moment, because the Plan's scalar
   // preheader may be unreachable at this point. Instead it is replaced in
   // createVectorizedLoopSkeleton.

llvm/lib/Transforms/Vectorize/VPlan.h

@@ -3529,18 +3529,25 @@ class VPlan {
 
   /// Returns the 'middle' block of the plan, that is the block that selects
   /// whether to execute the scalar tail loop or the exit block from the loop
-  /// latch. If the scalar tail loop or exit block are known to always execute,
-  /// the middle block may branch directly to the block.
+  /// latch. If there is an early exit from the vector loop, the middle block
+  /// conceptully has the early exit block as third successor, split accross 2
+  /// VPBBs. In that case, the second VPBB selects whether to execute the scalar
+  /// tail loop or the exit bock. If the scalar tail loop or exit block are
+  /// known to always execute, the middle block may branch directly to that
+  /// block. If there is no loop region, the middle block cannot be identified
+  /// and nullptr is returned.
   VPBasicBlock *getMiddleBlock() {
     VPRegionBlock *LoopRegion = getVectorLoopRegion();
     if (!LoopRegion)
       return nullptr;
-    auto *RegionSucc = LoopRegion->getSingleSuccessor();
+    auto *RegionSucc = cast<VPBasicBlock>(LoopRegion->getSingleSuccessor());
     if (RegionSucc->getSingleSuccessor() ||
         is_contained(RegionSucc->getSuccessors(), getScalarPreheader()))
-      return cast<VPBasicBlock>(RegionSucc);
+      return RegionSucc;
+    // There is an early exit. The successor of RegionSucc is the middle block.
     return cast<VPBasicBlock>(RegionSucc->getSuccessors()[1]);
   }
+
   const VPBasicBlock *getMiddleBlock() const {
     return const_cast<VPlan *>(this)->getMiddleBlock();
   }

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

@@ -1684,7 +1684,7 @@ void VPlanTransforms::truncateToMinimalBitwidths(
 
 /// Remove BranchOnCond recipes with true conditions together with removing
 /// dead edges to their successors.
-static void simplifyBranchOnCondTrue(VPlan &Plan) {
+static void removeBranchOnCondTrue(VPlan &Plan) {
   using namespace llvm::VPlanPatternMatch;
   for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly<VPBasicBlock>(
            vp_depth_first_shallow(Plan.getEntry()))) {
@@ -1737,7 +1737,7 @@ void VPlanTransforms::optimize(VPlan &Plan) {
   runPass(legalizeAndOptimizeInductions, Plan);
   runPass(removeRedundantExpandSCEVRecipes, Plan);
   runPass(simplifyRecipes, Plan, *Plan.getCanonicalIV()->getScalarType());
-  runPass(simplifyBranchOnCondTrue, Plan);
+  runPass(removeBranchOnCondTrue, Plan);
   runPass(removeDeadRecipes, Plan);
 
   runPass(createAndOptimizeReplicateRegions, Plan);