!fixup address comments, thanks!

Author: fhahn

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -2428,30 +2428,24 @@ InnerLoopVectorizer::getOrCreateVectorTripCount(BasicBlock *InsertBlock) {
   return VectorTripCount;
 }
 
-/// Helper to connect both the vector and scalar preheaders to the Plan's
-/// entry. This is used when adjusting \p Plan during skeleton
-/// creation, i.e. adjusting the plan after introducing an initial runtime
-/// check.
-/// TODO: Connect scalar preheader during initial VPlan construction.
-static void connectScalarPreheaderAsBypassInVPlan(VPlan &Plan) {
-  VPBlockBase *VectorPH = Plan.getVectorPreheader();
-  VPBlockBase *ScalarPH = Plan.getScalarPreheader();
-  VPBlockBase *Entry = Plan.getEntry();
-  VPBlockUtils::connectBlocks(Entry, ScalarPH);
-  std::swap(Entry->getSuccessors()[0], Entry->getSuccessors()[1]);
-}
-
 /// Introduces a new VPIRBasicBlock for \p CheckIRBB to \p Plan between the
-/// vector preheader and its predecessor, also connecting to the scalar
-/// preheader.
+/// vector preheader and its predecessor, also connecting the new block to the
+/// scalar preheader.
 static void introduceCheckBlockInVPlan(VPlan &Plan, BasicBlock *CheckIRBB) {
+
   VPBlockBase *ScalarPH = Plan.getScalarPreheader();
   VPBlockBase *VectorPH = Plan.getVectorPreheader();
   VPBlockBase *PreVectorPH = VectorPH->getSinglePredecessor();
-  assert(PreVectorPH->getSuccessors()[0] == ScalarPH && "Unexpected successor");
-  VPIRBasicBlock *CheckVPIRBB = VPIRBasicBlock::fromBasicBlock(CheckIRBB);
-  VPBlockUtils::connectBlocks(CheckVPIRBB, ScalarPH);
-  VPBlockUtils::insertOnEdge(PreVectorPH, VectorPH, CheckVPIRBB);
+  if (PreVectorPH->getNumSuccessors() != 1) {
+    assert(PreVectorPH->getNumSuccessors() == 2 && "Expected 2 successors");
+    assert(PreVectorPH->getSuccessors()[0] == ScalarPH &&
+           "Unexpected successor");
+    VPIRBasicBlock *CheckVPIRBB = VPIRBasicBlock::fromBasicBlock(CheckIRBB);
+    VPBlockUtils::insertOnEdge(PreVectorPH, VectorPH, CheckVPIRBB);
+    PreVectorPH = CheckVPIRBB;
+  }
+  VPBlockUtils::connectBlocks(PreVectorPH, ScalarPH);
+  PreVectorPH->swapSuccessors();
 }
 
 void InnerLoopVectorizer::emitIterationCountCheck(BasicBlock *Bypass) {
@@ -2536,7 +2530,7 @@ void InnerLoopVectorizer::emitIterationCountCheck(BasicBlock *Bypass) {
   LoopBypassBlocks.push_back(TCCheckBlock);
 
   // TODO: Wrap LoopVectorPreHeader in VPIRBasicBlock here.
-  connectScalarPreheaderAsBypassInVPlan(Plan);
+  introduceCheckBlockInVPlan(Plan, nullptr);
 }
 
 BasicBlock *InnerLoopVectorizer::emitSCEVChecks(BasicBlock *Bypass) {
@@ -7919,13 +7913,8 @@ EpilogueVectorizerMainLoop::emitIterationCountCheck(BasicBlock *Bypass,
     setBranchWeights(BI, MinItersBypassWeights, /*IsExpected=*/false);
   ReplaceInstWithInst(TCCheckBlock->getTerminator(), &BI);
 
-  // Connect the vector preheader to the scalar preheader when creating the trip
-  // count for the epilogue loop. Otherwise there is already a runtime check and
-  // connection to the scalar preheader, so we introduce it as new check block.
-  if (ForEpilogue)
-    connectScalarPreheaderAsBypassInVPlan(Plan);
-  else
-    introduceCheckBlockInVPlan(Plan, TCCheckBlock);
+  // Connect TCCheckblock to the VPlan.
+  introduceCheckBlockInVPlan(Plan, TCCheckBlock);
   return TCCheckBlock;
 }
 
@@ -8071,14 +8060,15 @@ EpilogueVectorizerEpilogueLoop::emitMinimumVectorEpilogueIterCountCheck(
   ReplaceInstWithInst(Insert->getTerminator(), &BI);
   LoopBypassBlocks.push_back(Insert);
 
-  // A new entry block has been created for the epilogue VPlan. Hook it in.
+  // A new entry block has been created for the epilogue VPlan. Hook it in, as
+  // otherwise we would try to modify the entry to the main vector loop.
   VPIRBasicBlock *NewEntry = VPIRBasicBlock::fromBasicBlock(Insert);
   VPBasicBlock *OldEntry = Plan.getEntry();
   VPBlockUtils::reassociateBlocks(OldEntry, NewEntry);
   Plan.setEntry(NewEntry);
   delete OldEntry;
 
-  connectScalarPreheaderAsBypassInVPlan(Plan);
+  introduceCheckBlockInVPlan(Plan, nullptr);
   return Insert;
 }
 

llvm/lib/Transforms/Vectorize/VPlan.cpp

@@ -821,16 +821,16 @@ void VPRegionBlock::print(raw_ostream &O, const Twine &Indent,
 }
 #endif
 
-VPlan::VPlan(VPBasicBlock *Preheader, VPValue *TC, VPBasicBlock *Entry,
-             VPIRBasicBlock *ScalarHeader)
-    : VPlan(Preheader, TC, ScalarHeader) {
-  VPBlockUtils::connectBlocks(Preheader, Entry);
+VPlan::VPlan(VPBasicBlock *OriginalPreheader, VPValue *TC,
+             VPBasicBlock *EntryVectorPreHeader, VPIRBasicBlock *ScalarHeader)
+    : VPlan(OriginalPreheader, TC, ScalarHeader) {
+  VPBlockUtils::connectBlocks(OriginalPreheader, EntryVectorPreHeader);
 }
 
-VPlan::VPlan(VPBasicBlock *Preheader, VPBasicBlock *Entry,
-             VPIRBasicBlock *ScalarHeader)
-    : VPlan(Preheader, ScalarHeader) {
-  VPBlockUtils::connectBlocks(Preheader, Entry);
+VPlan::VPlan(VPBasicBlock *OriginalPreheader,
+             VPBasicBlock *EntryVectorPreHeader, VPIRBasicBlock *ScalarHeader)
+    : VPlan(OriginalPreheader, ScalarHeader) {
+  VPBlockUtils::connectBlocks(OriginalPreheader, EntryVectorPreHeader);
 }
 
 VPlan::~VPlan() {
@@ -1187,6 +1187,7 @@ std::string VPlan::getName() const {
 }
 
 VPRegionBlock *VPlan::getVectorLoopRegion() {
+  // TODO: Cache if possible.
   for (VPBlockBase *B : vp_depth_first_shallow(getEntry()))
     if (auto *R = dyn_cast<VPRegionBlock>(B))
       return R;
@@ -1200,6 +1201,16 @@ const VPRegionBlock *VPlan::getVectorLoopRegion() const {
   return nullptr;
 }
 
+VPBasicBlock *VPlan::getScalarPreheader() const {
+  auto *MiddleVPBB =
+      cast<VPBasicBlock>(getVectorLoopRegion()->getSingleSuccessor());
+  auto *LastSucc = MiddleVPBB->getSuccessors().back();
+  // If scalar preheader is connected to VPlan, it is the last successor of
+  // MiddleVPBB. If this last successor is a VPIRBasicBlock, it is the Exit
+  // block rather than the scalar preheader.
+  return isa<VPIRBasicBlock>(LastSucc) ? nullptr : cast<VPBasicBlock>(LastSucc);
+}
+
 LLVM_DUMP_METHOD
 void VPlan::printDOT(raw_ostream &O) const {
   VPlanPrinter Printer(O, *this);
@@ -1292,15 +1303,6 @@ VPlan *VPlan::duplicate() {
   return NewPlan;
 }
 
-VPBasicBlock *VPlan::getScalarPreheader() {
-  auto *MiddleVPBB =
-      cast<VPBasicBlock>(getVectorLoopRegion()->getSingleSuccessor());
-  auto *LastSucc = MiddleVPBB->getSuccessors().back();
-  // Order is strict: if the last successor is VPIRBasicBlock, it must be the
-  // single exit.
-  return isa<VPIRBasicBlock>(LastSucc) ? nullptr : cast<VPBasicBlock>(LastSucc);
-}
-
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
 
 Twine VPlanPrinter::getUID(const VPBlockBase *Block) {

llvm/lib/Transforms/Vectorize/VPlan.h

@@ -621,6 +621,12 @@ class VPBlockBase {
   /// Remove all the successors of this block.
   void clearSuccessors() { Successors.clear(); }
 
+  /// Swap successors of the block. The block must have exactly 2 successors.
+  void swapSuccessors() {
+    assert(Successors.size() == 2 && "must have 2 successors to swap");
+    std::swap(Successors[0], Successors[1]);
+  }
+
   /// The method which generates the output IR that correspond to this
   /// VPBlockBase, thereby "executing" the VPlan.
   virtual void execute(VPTransformState *State) = 0;
@@ -3814,14 +3820,15 @@ class VPlan {
   }
 
   /// Constructor variants that take disconnected preheader and entry blocks,
-  /// connecting them as part of construction. Only used to reduce the need of
-  /// code changes during transition.
-  VPlan(VPBasicBlock *Preheader, VPValue *TC, VPBasicBlock *Entry,
-        VPIRBasicBlock *ScalarHeader);
-  VPlan(VPBasicBlock *Preheader, VPBasicBlock *Entry,
+  /// connecting them as part of construction.
+  /// FIXME: Only used to reduce the need of code changes during transition.
+  VPlan(VPBasicBlock *OriginalPreheader, VPValue *TC,
+        VPBasicBlock *EntryVectorPreHeader, VPIRBasicBlock *ScalarHeader);
+  VPlan(VPBasicBlock *OriginalPreheader, VPBasicBlock *EntryVectorPreHeader,
         VPIRBasicBlock *ScalarHeader);
 
-  /// Construct a VPlan with \p Entry to the plan.
+  /// Construct a VPlan with \p Entry to the plan and with \p ScalarHeader
+  /// wrapping the original header of the scalar loop.
   VPlan(VPBasicBlock *Entry, VPIRBasicBlock *ScalarHeader)
       : Entry(Entry), ScalarHeader(ScalarHeader) {
     Entry->setPlan(this);
@@ -3883,9 +3890,7 @@ class VPlan {
   }
 
   /// Return the VPBasicBlock for the preheader of the scalar loop.
-  VPBasicBlock *getScalarPreheader() const {
-    return cast<VPBasicBlock>(ScalarHeader->getSinglePredecessor());
-  }
+  VPBasicBlock *getScalarPreheader() const;
 
   /// Return the VPIRBasicBlock wrapping the header of the scalar loop.
   VPIRBasicBlock *getScalarHeader() const { return ScalarHeader; }
@@ -4020,14 +4025,14 @@ class VPlan {
   }
 
   /// \return The block corresponding to the original preheader.
+  /// FIXME: There's no separate preheader any longer and Entry now serves the
+  /// same purpose as the original preheader. Remove after transition.
   VPBasicBlock *getPreheader() { return Entry; }
   const VPBasicBlock *getPreheader() const { return Entry; }
 
   /// Clone the current VPlan, update all VPValues of the new VPlan and cloned
   /// recipes to refer to the clones, and return it.
   VPlan *duplicate();
-
-  VPBasicBlock *getScalarPreheader();
 };
 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)