!fixup address latest comments, thanks!

Author: fhahn

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -2781,7 +2781,7 @@ void InnerLoopVectorizer::fixVectorizedLoop(VPTransformState &State) {
   // Don't apply optimizations below when no (vector) loop remains, as they all
   // require one at the moment.
   VPBasicBlock *HeaderVPBB =
-      vputils::getTopLevelVectorLoopHeader(*State.Plan, State.VPDT);
+      vputils::getFirstLoopHeader(*State.Plan, State.VPDT);
   if (!HeaderVPBB)
     return;
 
@@ -7801,6 +7801,9 @@ DenseMap<const SCEV *, Value *> LoopVectorizationPlanner::executePlan(
       TTI.getRegisterBitWidth(TargetTransformInfo::RGK_FixedWidthVector));
   VPlanTransforms::removeDeadRecipes(BestVPlan);
 
+  // Retrieve and store the middle block before dissolving regions. Regions are
+  // dissolved after optimizing for VF and UF, which completely removes unneeded
+  // loop regions first.
   VPBasicBlock *MiddleVPBB =
       BestVPlan.getVectorLoopRegion() ? BestVPlan.getMiddleBlock() : nullptr;
   VPlanTransforms::disolveLoopRegions(BestVPlan);
@@ -7899,8 +7902,7 @@ DenseMap<const SCEV *, Value *> LoopVectorizationPlanner::executePlan(
   // 2.6. Maintain Loop Hints
   // Keep all loop hints from the original loop on the vector loop (we'll
   // replace the vectorizer-specific hints below).
-  VPBasicBlock *HeaderVPBB =
-      vputils::getTopLevelVectorLoopHeader(BestVPlan, State.VPDT);
+  VPBasicBlock *HeaderVPBB = vputils::getFirstLoopHeader(BestVPlan, State.VPDT);
   if (HeaderVPBB) {
     MDNode *OrigLoopID = OrigLoop->getLoopID();
 

llvm/lib/Transforms/Vectorize/VPlan.cpp

@@ -212,14 +212,24 @@ bool VPBlockUtils::isHeader(const VPBlockBase *VPB,
   auto *VPBB = dyn_cast<VPBasicBlock>(VPB);
   if (!VPBB)
     return false;
+
+  // If VPBB is in a region R, VPBB is a loop header if R is a loop region with
+  // VPBB as its entry, i.e., free of predecessors.
   if (auto *R = VPBB->getParent())
     return !R->isReplicator() && VPBB->getNumPredecessors() == 0;
 
-  assert(!VPB->getParent() && "checking blocks in regions not implemented yet");
+  // A header dominates its second predecessor (the latch), with the other
+  // predecessor being the preheader
   return VPB->getPredecessors().size() == 2 &&
          VPDT.dominates(VPB, VPB->getPredecessors()[1]);
 }
 
+bool VPBlockUtils::isLatch(const VPBlockBase *VPB,
+                           const VPDominatorTree &VPDT) {
+  return VPB->getNumSuccessors() == 2 &&
+         VPBlockUtils::isHeader(VPB->getSuccessors()[1], VPDT);
+}
+
 VPBasicBlock::iterator VPBasicBlock::getFirstNonPhi() {
   iterator It = begin();
   while (It != end() && It->isPhi())
@@ -435,16 +445,17 @@ void VPBasicBlock::connectToPredecessors(VPTransformState &State) {
 
   auto Preds = to_vector(getHierarchicalPredecessors());
   if (VPBlockUtils::isHeader(this, State.VPDT)) {
-    // There's no block yet for the latch, don't try to connect it yet.
+    // There's no block for the latch yet, connect to the preheader only.
     Preds = {Preds[0]};
   }
 
   // Hook up the new basic block to its predecessors.
   for (VPBlockBase *PredVPBlock : Preds) {
     VPBasicBlock *PredVPBB = PredVPBlock->getExitingBasicBlock();
     auto &PredVPSuccessors = PredVPBB->getHierarchicalSuccessors();
+    assert(CFG.VPBB2IRBB.contains(PredVPBB) &&
+           "Predecessor basic-block not found building successor.");
     BasicBlock *PredBB = CFG.VPBB2IRBB.lookup(PredVPBB);
-    assert(PredBB && "Predecessor basic-block not found building successor.");
     auto *PredBBTerminator = PredBB->getTerminator();
     LLVM_DEBUG(dbgs() << "LV: draw edge from" << PredBB->getName() << '\n');
 
@@ -548,8 +559,7 @@ void VPBasicBlock::execute(VPTransformState *State) {
   executeRecipes(State, NewBB);
 
   // If this block is a latch, update CurrentParentLoop.
-  if (getNumSuccessors() == 2 &&
-      VPBlockUtils::isHeader(getSuccessors()[1], State->VPDT))
+  if (VPBlockUtils::isLatch(this, State->VPDT))
     State->CurrentParentLoop = State->CurrentParentLoop->getParentLoop();
 }
 
@@ -988,8 +998,7 @@ void VPlan::execute(VPTransformState *State) {
   for (VPBlockBase *Block : RPOT)
     Block->execute(State);
 
-  VPBasicBlock *Header =
-      vputils::getTopLevelVectorLoopHeader(*this, State->VPDT);
+  VPBasicBlock *Header = vputils::getFirstLoopHeader(*this, State->VPDT);
   if (!Header)
     return;
 

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

@@ -577,7 +577,7 @@ Value *VPInstruction::generate(VPTransformState &State) {
     Value *Cond = State.get(getOperand(0), VPLane(0));
     // Replace the temporary unreachable terminator with a new conditional
     // branch, hooking it up to backward destination (header) for latch blocks
-    // now to forward destination(s) later when they are created.
+    // now, and to forward destination(s) later when they are created.
     BranchInst *CondBr =
         Builder.CreateCondBr(Cond, Builder.GetInsertBlock(), nullptr);
     CondBr->setSuccessor(0, nullptr);
@@ -601,9 +601,9 @@ Value *VPInstruction::generate(VPTransformState &State) {
 
     // Replace the temporary unreachable terminator with a new conditional
     // branch, hooking it up to backward destination (the header) for latch
-    // blocks now forward destination (the exit/middle block) later when it is
-    // created. Note that CreateCondBr expects a valid BB as first argument, so
-    // we need to set it to nullptr later.
+    // blocks now, and to forward destination (the exit/middle block) later when
+    // it is created. Note that CreateCondBr expects a valid BB as first
+    // argument, so we need to set it to nullptr later.
     BranchInst *CondBr = Builder.CreateCondBr(Cond, Builder.GetInsertBlock(),
                                               State.CFG.VPBB2IRBB[Header]);
     CondBr->setSuccessor(0, nullptr);

llvm/lib/Transforms/Vectorize/VPlanUtils.cpp

@@ -126,8 +126,7 @@ bool vputils::isUniformAcrossVFsAndUFs(VPValue *V) {
       });
 }
 
-VPBasicBlock *vputils::getTopLevelVectorLoopHeader(VPlan &Plan,
-                                                   VPDominatorTree &VPDT) {
+VPBasicBlock *vputils::getFirstLoopHeader(VPlan &Plan, VPDominatorTree &VPDT) {
   auto DepthFirst = vp_depth_first_shallow(Plan.getEntry());
   auto I = find_if(DepthFirst, [&VPDT](VPBlockBase *VPB) {
     return VPBlockUtils::isHeader(VPB, VPDT);

llvm/lib/Transforms/Vectorize/VPlanUtils.h

@@ -94,8 +94,9 @@ bool isHeaderMask(const VPValue *V, VPlan &Plan);
 /// VPDerivedIV or VPCanonicalIVPHI).
 bool isUniformAcrossVFsAndUFs(VPValue *V);
 
-/// Returns the header block of the top-level vector loop, if one exists.
-VPBasicBlock *getTopLevelVectorLoopHeader(VPlan &Plan, VPDominatorTree &VPDT);
+/// Returns the header block of the first, top-level loop, or null if none
+/// exist.
+VPBasicBlock *getFirstLoopHeader(VPlan &Plan, VPDominatorTree &VPDT);
 } // namespace vputils
 
 //===----------------------------------------------------------------------===//
@@ -243,8 +244,12 @@ class VPBlockUtils {
     VPBlockUtils::connectBlocks(BlockPtr, To, PredIx, -1);
   }
 
-  /// Returns true if \p VPB is a loop header.
-  static bool isHeader(const VPBlockBase *VPBB, const VPDominatorTree &VPDT);
+  /// Returns true if \p VPB is a loop header, based on regions or \p VPDT in
+  /// their absence.
+  static bool isHeader(const VPBlockBase *VPB, const VPDominatorTree &VPDT);
+
+  /// Returns true if \p VPB is a loop latch, using isHeader().
+  static bool isLatch(const VPBlockBase *VPB, const VPDominatorTree &VPDT);
 };
 
 } // namespace llvm