[VPlan] Create initial skeleton before creating regions. (NFC)

Move out the logic to prepare for vectorization to a separate transform,
before creating loop regions. This was discussed as follow-up
in llvm/llvm-project#136455.

This just moves the existing code around slightly  and will simplify
follow-up patches to include the exiting edges during initial VPlan
construction.

Author: fhahn

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -9457,9 +9457,10 @@ LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(VFRange &Range) {
           Range);
   DenseMap<VPBlockBase *, BasicBlock *> VPB2IRBB;
   auto Plan = VPlanTransforms::buildPlainCFG(OrigLoop, *LI, VPB2IRBB);
-  VPlanTransforms::createLoopRegions(*Plan, Legal->getWidestInductionType(),
-                                     PSE, RequiresScalarEpilogueCheck,
-                                     CM.foldTailByMasking(), OrigLoop);
+  VPlanTransforms::prepareForVectorization(
+      *Plan, Legal->getWidestInductionType(), PSE, RequiresScalarEpilogueCheck,
+      CM.foldTailByMasking(), OrigLoop);
+  VPlanTransforms::createLoopRegions(*Plan);
 
   // Don't use getDecisionAndClampRange here, because we don't know the UF
   // so this function is better to be conservative, rather than to split
@@ -9749,8 +9750,9 @@ VPlanPtr LoopVectorizationPlanner::tryToBuildVPlan(VFRange &Range) {
 
   DenseMap<VPBlockBase *, BasicBlock *> VPB2IRBB;
   auto Plan = VPlanTransforms::buildPlainCFG(OrigLoop, *LI, VPB2IRBB);
-  VPlanTransforms::createLoopRegions(*Plan, Legal->getWidestInductionType(),
-                                     PSE, true, false, OrigLoop);
+  VPlanTransforms::prepareForVectorization(
+      *Plan, Legal->getWidestInductionType(), PSE, true, false, OrigLoop);
+  VPlanTransforms::createLoopRegions(*Plan);
 
   for (ElementCount VF : Range)
     Plan->addVF(VF);

llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp

@@ -461,19 +461,23 @@ static void createLoopRegion(VPlan &Plan, VPBlockBase *HeaderVPB) {
     VPBlockUtils::connectBlocks(R, Succ);
 }
 
-void VPlanTransforms::createLoopRegions(VPlan &Plan, Type *InductionTy,
-                                        PredicatedScalarEvolution &PSE,
-                                        bool RequiresScalarEpilogueCheck,
-                                        bool TailFolded, Loop *TheLoop) {
+void VPlanTransforms::prepareForVectorization(VPlan &Plan, Type *InductionTy,
+                                              PredicatedScalarEvolution &PSE,
+                                              bool RequiresScalarEpilogueCheck,
+                                              bool TailFolded, Loop *TheLoop) {
   VPDominatorTree VPDT;
   VPDT.recalculate(Plan);
-  for (VPBlockBase *HeaderVPB : vp_depth_first_shallow(Plan.getEntry()))
-    if (canonicalHeaderAndLatch(HeaderVPB, VPDT))
-      createLoopRegion(Plan, HeaderVPB);
 
-  VPRegionBlock *TopRegion = Plan.getVectorLoopRegion();
-  TopRegion->setName("vector loop");
-  TopRegion->getEntryBasicBlock()->setName("vector.body");
+  VPBlockBase *HeaderVPB = Plan.getEntry()->getSingleSuccessor();
+  canonicalHeaderAndLatch(HeaderVPB, VPDT);
+  VPBlockBase *LatchVPB = HeaderVPB->getPredecessors()[1];
+
+  VPBasicBlock *VecPreheader = Plan.createVPBasicBlock("vector.ph");
+  VPBlockUtils::insertBlockAfter(VecPreheader, Plan.getEntry());
+
+  VPBasicBlock *MiddleVPBB = Plan.createVPBasicBlock("middle.block");
+  VPBlockUtils::connectBlocks(LatchVPB, MiddleVPBB);
+  LatchVPB->swapSuccessors();
 
   // Create SCEV and VPValue for the trip count.
   // We use the symbolic max backedge-taken-count, which works also when
@@ -487,11 +491,6 @@ void VPlanTransforms::createLoopRegions(VPlan &Plan, Type *InductionTy,
   Plan.setTripCount(
       vputils::getOrCreateVPValueForSCEVExpr(Plan, TripCount, SE));
 
-  VPBasicBlock *VecPreheader = Plan.createVPBasicBlock("vector.ph");
-  VPBlockUtils::insertBlockAfter(VecPreheader, Plan.getEntry());
-  VPBasicBlock *MiddleVPBB = Plan.createVPBasicBlock("middle.block");
-  VPBlockUtils::insertBlockAfter(MiddleVPBB, TopRegion);
-
   VPBasicBlock *ScalarPH = Plan.createVPBasicBlock("scalar.ph");
   VPBlockUtils::connectBlocks(ScalarPH, Plan.getScalarHeader());
 
@@ -516,10 +515,10 @@ void VPlanTransforms::createLoopRegions(VPlan &Plan, Type *InductionTy,
     return;
   }
 
+  // The connection order corresponds to the operands of the conditional branch.
   BasicBlock *IRExitBlock = TheLoop->getUniqueLatchExitBlock();
   auto *VPExitBlock = Plan.getExitBlock(IRExitBlock);
-  // The connection order corresponds to the operands of the conditional branch.
-  VPBlockUtils::insertBlockAfter(VPExitBlock, MiddleVPBB);
+  VPBlockUtils::connectBlocks(MiddleVPBB, VPExitBlock);
   VPBlockUtils::connectBlocks(MiddleVPBB, ScalarPH);
 
   auto *ScalarLatchTerm = TheLoop->getLoopLatch()->getTerminator();
@@ -538,3 +537,15 @@ void VPlanTransforms::createLoopRegions(VPlan &Plan, Type *InductionTy,
   Builder.createNaryOp(VPInstruction::BranchOnCond, {Cmp},
                        ScalarLatchTerm->getDebugLoc());
 }
+
+void VPlanTransforms::createLoopRegions(VPlan &Plan) {
+  VPDominatorTree VPDT;
+  VPDT.recalculate(Plan);
+  for (VPBlockBase *HeaderVPB : vp_depth_first_shallow(Plan.getEntry()))
+    if (canonicalHeaderAndLatch(HeaderVPB, VPDT))
+      createLoopRegion(Plan, HeaderVPB);
+
+  VPRegionBlock *TopRegion = Plan.getVectorLoopRegion();
+  TopRegion->setName("vector loop");
+  TopRegion->getEntryBasicBlock()->setName("vector.body");
+}

llvm/lib/Transforms/Vectorize/VPlanTransforms.h

@@ -57,19 +57,22 @@ struct VPlanTransforms {
   buildPlainCFG(Loop *TheLoop, LoopInfo &LI,
                 DenseMap<VPBlockBase *, BasicBlock *> &VPB2IRBB);
 
-  /// Replace loops in \p Plan's flat CFG with VPRegionBlocks, turing \p Plan's
-  /// flat CFG into a hierarchical CFG. It also creates a VPValue expression for
-  /// the original trip count. It will also introduce a dedicated VPBasicBlock
-  /// for the vector pre-header as well a VPBasicBlock as exit block of the
-  /// region (middle.block). If a check is needed to guard executing the scalar
-  /// epilogue loop, it will be added to the middle block, together with
-  /// VPBasicBlocks for the scalar preheader and exit blocks. \p InductionTy is
-  /// the type of the canonical induction and used for related values, like the
-  /// trip count expression.
-  static void createLoopRegions(VPlan &Plan, Type *InductionTy,
-                                PredicatedScalarEvolution &PSE,
-                                bool RequiresScalarEpilogueCheck,
-                                bool TailFolded, Loop *TheLoop);
+  /// Prepare the plan for vectorization. It will introduce a dedicated
+  /// VPBasicBlock for the vector pre-header as well as a VPBasicBlock as exit
+  /// block of the main vector loop (middle.block). If a check is needed to
+  /// guard executing the scalar epilogue loop, it will be added to the middle
+  /// block, together with VPBasicBlocks for the scalar preheader and exit
+  /// blocks. \p InductionTy is the type of the canonical induction and used for
+  /// related values, like the trip count expression.  It also creates a VPValue
+  /// expression for the original trip count.
+  static void prepareForVectorization(VPlan &Plan, Type *InductionTy,
+                                      PredicatedScalarEvolution &PSE,
+                                      bool RequiresScalarEpilogueCheck,
+                                      bool TailFolded, Loop *TheLoop);
+
+  /// Replace loops in \p Plan's flat CFG with VPRegionBlocks, turning \p Plan's
+  /// flat CFG into a hierarchical CFG.
+  static void createLoopRegions(VPlan &Plan);
 
   /// Replaces the VPInstructions in \p Plan with corresponding
   /// widen recipes. Returns false if any VPInstructions could not be converted

llvm/unittests/Transforms/Vectorize/VPlanTestBase.h

@@ -72,8 +72,9 @@ class VPlanTestIRBase : public testing::Test {
     PredicatedScalarEvolution PSE(*SE, *L);
     DenseMap<VPBlockBase *, BasicBlock *> VPB2IRBB;
     auto Plan = VPlanTransforms::buildPlainCFG(L, *LI, VPB2IRBB);
-    VPlanTransforms::createLoopRegions(*Plan, IntegerType::get(*Ctx, 64), PSE,
-                                       true, false, L);
+    VPlanTransforms::prepareForVectorization(*Plan, IntegerType::get(*Ctx, 64),
+                                             PSE, true, false, L);
+    VPlanTransforms::createLoopRegions(*Plan);
     return Plan;
   }
 };