!fixup address latest comments, thanks

Author: fhahn

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -9324,6 +9324,8 @@ LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(VFRange &Range) {
   auto Plan = std::make_unique<VPlan>(OrigLoop);
   // Build hierarchical CFG.
   VPlanHCFGBuilder HCFGBuilder(OrigLoop, LI, *Plan);
+  // TODO: Convert to VPlan-transform and consoliate all transforms for VPlan
+  // creation.
   HCFGBuilder.buildHierarchicalCFG();
 
   VPlanTransforms::introduceTopLevelVectorLoopRegion(

llvm/lib/Transforms/Vectorize/VPlan.h

@@ -3602,11 +3602,17 @@ class VPlan {
   /// the original trip count have been replaced.
   void resetTripCount(VPValue *NewTripCount) {
     assert(TripCount && NewTripCount && TripCount->getNumUsers() == 0 &&
-           "TripCount always must be set");
+           "TripCount must be set when resetting");
     TripCount = NewTripCount;
   }
 
-  void setTripCount(VPValue *NewTripCount) { TripCount = NewTripCount; }
+  // Set the trip count assuming it is currently null; if it is not - use
+  // resetTripCount().
+  void setTripCount(VPValue *NewTripCount) {
+    assert(!TripCount && NewTripCount && "TripCount should not be set yet.");
+
+    TripCount = NewTripCount;
+  }
 
   /// The backedge taken count of the original loop.
   VPValue *getOrCreateBackedgeTakenCount() {

llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.cpp

@@ -178,7 +178,8 @@ VPBasicBlock *PlainCFGBuilder::getOrCreateVPBB(BasicBlock *BB) {
   VPBasicBlock *VPBB = Plan.createVPBasicBlock(Name);
   BB2VPBB[BB] = VPBB;
 
-  // Get or create a region for the loop containing BB.
+  // Get or create a region for the loop containing BB, except for the top
+  // region of TheLoop which is created later.
   Loop *LoopOfBB = LI->getLoopFor(BB);
   if (!LoopOfBB || LoopOfBB == TheLoop || !doesContainLoop(LoopOfBB, TheLoop))
     return VPBB;
@@ -194,12 +195,8 @@ VPBasicBlock *PlainCFGBuilder::getOrCreateVPBB(BasicBlock *BB) {
   assert(!RegionOfVPBB &&
          "First visit of a header basic block expects to register its region.");
   // Handle a header - take care of its Region.
-  if (LoopOfBB == TheLoop) {
-    RegionOfVPBB = Plan.getVectorLoopRegion();
-  } else {
-    RegionOfVPBB = Plan.createVPRegionBlock(Name.str(), false /*isReplicator*/);
-    RegionOfVPBB->setParent(Loop2Region[LoopOfBB->getParentLoop()]);
-  }
+  RegionOfVPBB = Plan.createVPRegionBlock(Name.str(), false /*isReplicator*/);
+  RegionOfVPBB->setParent(Loop2Region[LoopOfBB->getParentLoop()]);
   RegionOfVPBB->setEntry(VPBB);
   Loop2Region[LoopOfBB] = RegionOfVPBB;
   return VPBB;
@@ -383,11 +380,10 @@ void PlainCFGBuilder::buildPlainCFG(
     // Set VPBB predecessors in the same order as they are in the incoming BB.
     if (!isHeaderBB(BB, LoopForBB)) {
       setVPBBPredsFromBB(VPBB, BB);
-    } else {
-      // BB is a loop header, set the predecessor for the region, except for the
-      // top region, whose predecessor was set when creating VPlan's skeleton.
-      if (LoopForBB != TheLoop)
-        setRegionPredsFromBB(Region, BB);
+    } else if (Region) {
+      // BB is a loop header and there's a corresponding region , set the
+      // predecessor for it.
+      setRegionPredsFromBB(Region, BB);
     }
 
     // Create VPInstructions for BB.
@@ -427,22 +423,15 @@ void PlainCFGBuilder::buildPlainCFG(
     VPBasicBlock *Successor0 = getOrCreateVPBB(IRSucc0);
     VPBasicBlock *Successor1 = getOrCreateVPBB(IRSucc1);
     if (BB == LoopForBB->getLoopLatch()) {
-      // For a latch we need to set the successor of the region rather
-      // than that
-      // of VPBB and it should be set to the exit, i.e., non-header
-      // successor,
-      // except for the top region, whose successor was set when creating
-      // VPlan's skeleton.
+      // For a latch we need to set the successor of the region rather than that
+      // of VPBB and it should be set to the exit, i.e., non-header successor,
+      // except for the top region, which is handled elsewhere.
       assert(LoopForBB != TheLoop &&
              "Latch of the top region should have been handled earlier");
       Region->setOneSuccessor(isHeaderVPBB(Successor0) ? Successor1
                                                        : Successor0);
       Region->setExiting(VPBB);
       continue;
-
-      VPBasicBlock *HeaderVPBB = getOrCreateVPBB(LoopForBB->getHeader());
-      VPBlockUtils::connectBlocks(VPBB, HeaderVPBB);
-      continue;
     }
 
     // Don't connect any blocks outside the current loop except the latch for

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

@@ -35,6 +35,7 @@ using namespace llvm;
 void VPlanTransforms::introduceTopLevelVectorLoopRegion(
     VPlan &Plan, Type *InductionTy, PredicatedScalarEvolution &PSE,
     bool RequiresScalarEpilogueCheck, bool TailFolded, Loop *TheLoop) {
+  // TODO: Generalize to introduce all loop regions.
   auto *HeaderVPBB = cast<VPBasicBlock>(Plan.getEntry()->getSingleSuccessor());
   VPBlockUtils::disconnectBlocks(Plan.getEntry(), HeaderVPBB);
 
@@ -43,6 +44,7 @@ void VPlanTransforms::introduceTopLevelVectorLoopRegion(
   VPBlockUtils::disconnectBlocks(OriginalLatch, HeaderVPBB);
   VPBasicBlock *VecPreheader = Plan.createVPBasicBlock("vector.ph");
   VPBlockUtils::connectBlocks(Plan.getEntry(), VecPreheader);
+  assert(OriginalLatch->getNumSuccessors() == 0 && "expected no predecessors");
 
   // Create SCEV and VPValue for the trip count.
   // We use the symbolic max backedge-taken-count, which works also when
@@ -56,15 +58,14 @@ void VPlanTransforms::introduceTopLevelVectorLoopRegion(
   Plan.setTripCount(
       vputils::getOrCreateVPValueForSCEVExpr(Plan, TripCount, SE));
 
-  // Create VPRegionBlock, with empty header and latch blocks, to be filled
-  // during processing later.
+  // Create VPRegionBlock, with existing header and new empty latch block, to be
+  // filled
   VPBasicBlock *LatchVPBB = Plan.createVPBasicBlock("vector.latch");
   VPBlockUtils::insertBlockAfter(LatchVPBB, OriginalLatch);
   auto *TopRegion = Plan.createVPRegionBlock(
       HeaderVPBB, LatchVPBB, "vector loop", false /*isReplicator*/);
-  for (VPBlockBase *VPBB : vp_depth_first_shallow(HeaderVPBB)) {
+  for (VPBlockBase *VPBB : vp_depth_first_shallow(HeaderVPBB))
     VPBB->setParent(TopRegion);
-  }
 
   VPBlockUtils::insertBlockAfter(TopRegion, VecPreheader);
   VPBasicBlock *MiddleVPBB = Plan.createVPBasicBlock("middle.block");

llvm/lib/Transforms/Vectorize/VPlanTransforms.h

@@ -53,9 +53,9 @@ struct VPlanTransforms {
   }
 
   /// Introduce the top-level VPRegionBlock for the main loop in \p Plan. Coming
-  /// in this function, \p Plan's top-level loop is modeled using a plain CFG.
-  /// This transforms replaces the plain CFG with a VPRegionBlock wrapping the
-  /// top-level loop and creates a VPValue expressions for the original trip
+  /// into this function, \p Plan's top-level loop is modeled using a plain CFG.
+  /// This transform wraps the plain CFG of the top-level loop within a
+  /// VPRegionBlock and creates a VPValue expressions 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