!fixup address comments, thanks

Author: fhahn

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -8690,9 +8690,11 @@ VPlanPtr LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(
             return !CM.requiresScalarEpilogue(VF.isVector());
           },
           Range);
-  VPlanTransforms::handleEarlyExitsAndAddMiddleCheck(
-      *Plan, RequiresScalarEpilogueCheck, CM.foldTailByMasking(),
-      Legal->hasUncountableEarlyExit(), Range);
+  VPlanTransforms::handleEarlyExits(*Plan, Legal->hasUncountableEarlyExit(),
+                                    Range);
+  VPlanTransforms::addMiddleCheck(*Plan, RequiresScalarEpilogueCheck,
+                                  CM.foldTailByMasking());
+
   VPlanTransforms::createLoopRegions(*Plan);
   VPlanTransforms::createExtractsForLiveOuts(*Plan);
 
@@ -8981,9 +8983,11 @@ VPlanPtr LoopVectorizationPlanner::tryToBuildVPlan(VFRange &Range) {
   auto Plan = VPlanTransforms::buildVPlan0(
       OrigLoop, *LI, Legal->getWidestInductionType(),
       getDebugLocFromInstOrOperands(Legal->getPrimaryInduction()), PSE);
-  VPlanTransforms::handleEarlyExitsAndAddMiddleCheck(
-      *Plan, /*RequiresScalarEpilogue*/ true, /*TailFolded*/ false,
-      /*HasUncountableExit*/ false, Range);
+  VPlanTransforms::handleEarlyExits(*Plan,
+                                    /*HasUncountableExit*/ false, Range);
+  VPlanTransforms::addMiddleCheck(*Plan, /*RequiresScalarEpilogue*/ true,
+                                  /*TailFolded*/ false);
+
   VPlanTransforms::createLoopRegions(*Plan);
 
   for (ElementCount VF : Range)

llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp

@@ -514,14 +514,13 @@ VPlanTransforms::buildVPlan0(Loop *TheLoop, LoopInfo &LI, Type *InductionTy,
   return VPlan0;
 }
 
-void VPlanTransforms::handleEarlyExitsAndAddMiddleCheck(
-    VPlan &Plan, bool RequiresScalarEpilogueCheck, bool TailFolded,
-    bool HasUncountableEarlyExit, VFRange &Range) {
+void VPlanTransforms::handleEarlyExits(VPlan &Plan,
+                                       bool HasUncountableEarlyExit,
+                                       VFRange &Range) {
   auto *MiddleVPBB = cast<VPBasicBlock>(
       Plan.getScalarHeader()->getSinglePredecessor()->getPredecessors()[0]);
-  VPBlockBase *HeaderVPB =
-      Plan.getEntry()->getSuccessors()[1]->getSingleSuccessor();
-  auto *LatchVPBB = cast<VPBasicBlock>(HeaderVPB->getPredecessors()[1]);
+  auto *LatchVPBB = cast<VPBasicBlock>(MiddleVPBB->getSinglePredecessor());
+  VPBlockBase *HeaderVPB = cast<VPBasicBlock>(LatchVPBB->getSuccessors()[1]);
 
   // Disconnect all early exits from the loop leaving it with a single exit from
   // the latch. Early exits that are countable are left for a scalar epilog. The
@@ -551,7 +550,13 @@ void VPlanTransforms::handleEarlyExitsAndAddMiddleCheck(
 
   assert((!HasUncountableEarlyExit || HandledUncountableEarlyExit) &&
          "missed an uncountable exit that must be handled");
+}
 
+void VPlanTransforms::addMiddleCheck(VPlan &Plan,
+                                     bool RequiresScalarEpilogueCheck,
+                                     bool TailFolded) {
+  auto *MiddleVPBB = cast<VPBasicBlock>(
+      Plan.getScalarHeader()->getSinglePredecessor()->getPredecessors()[0]);
   // 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.
@@ -578,6 +583,7 @@ void VPlanTransforms::handleEarlyExitsAndAddMiddleCheck(
   // 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.
+  auto *LatchVPBB = cast<VPBasicBlock>(MiddleVPBB->getSinglePredecessor());
   DebugLoc LatchDL = LatchVPBB->getTerminator()->getDebugLoc();
   VPBuilder Builder(MiddleVPBB);
   VPValue *Cmp;

llvm/lib/Transforms/Vectorize/VPlanTransforms.h

@@ -59,8 +59,8 @@ struct VPlanTransforms {
   /// \p TheLoop being directly translated to VPBasicBlocks with VPInstruction
   /// corresponding to the input IR.
   ///
-  /// The created loop is wrapped into an initial skeleton to facilitate
-  /// vectorization, consisting of a vector pre-header, a exit block for the
+  /// The created loop is wrapped in an initial skeleton to facilitate
+  /// vectorization, consisting of a vector pre-header, an exit block for the
   /// main vector loop (middle.block) and a new block as preheader of the scalar
   /// loop (scalar.ph). It also adds a canonical IV and its increment, using \p
   /// InductionTy and \p IVDL, and creates a VPValue expression for the original
@@ -69,12 +69,15 @@ struct VPlanTransforms {
   buildVPlan0(Loop *TheLoop, LoopInfo &LI, Type *InductionTy, DebugLoc IVDL,
               PredicatedScalarEvolution &PSE);
 
-  /// Update \p Plan to account for all early exits. If a check is needed to
-  /// guard executing the scalar epilogue loop, it will be added to the middle
-  /// block.
-  LLVM_ABI_FOR_TEST static void handleEarlyExitsAndAddMiddleCheck(
-      VPlan &Plan, bool RequiresScalarEpilogueCheck, bool TailFolded,
-      bool HasUncountableExit, VFRange &Range);
+  /// Update \p Plan to account for all early exits.
+  LLVM_ABI_FOR_TEST static void
+  handleEarlyExits(VPlan &Plan, bool HasUncountableExit, VFRange &Range);
+
+  /// If a check is needed to guard executing the scalar epilogue loop, it will
+  /// be added to the middle block.
+  LLVM_ABI_FOR_TEST static void addMiddleCheck(VPlan &Plan,
+                                               bool RequiresScalarEpilogueCheck,
+                                               bool TailFolded);
 
   /// Replace loops in \p Plan's flat CFG with VPRegionBlocks, turning \p Plan's
   /// flat CFG into a hierarchical CFG.

llvm/unittests/Transforms/Vectorize/VPlanTestBase.h

@@ -76,8 +76,9 @@ class VPlanTestIRBase : public testing::Test {
                                              {}, PSE);
 
     VFRange R(ElementCount::getFixed(1), ElementCount::getFixed(2));
-    VPlanTransforms::handleEarlyExitsAndAddMiddleCheck(*Plan, true, false,
-                                                       false, R);
+    VPlanTransforms::handleEarlyExits(*Plan, false, R);
+    VPlanTransforms::addMiddleCheck(*Plan, true, false);
+
     VPlanTransforms::createLoopRegions(*Plan);
     return Plan;
   }