!fixup don't special case header/latch predecessor order.

Author: fhahn

llvm/lib/Transforms/Vectorize/VPlan.cpp

@@ -648,15 +648,20 @@ bool VPBasicBlock::isExiting() const {
   return getParent() && getParent()->getExitingBasicBlock() == this;
 }
 
-bool VPBasicBlock::isHeader(const VPDominatorTree &VPDT) const {
-  if (getNumPredecessors() != 2)
-    return false;
-  VPBlockBase *LatchVPBB = getPredecessors()[1];
-  if (!VPDT.dominates(this, LatchVPBB))
-    return false;
-  assert(VPDT.dominates(getPredecessors()[0], this) &&
-         "preheader must dominate header");
-  return true;
+std::optional<std::pair<VPBasicBlock *, VPBasicBlock *>>
+VPBasicBlock::isHeader(const VPDominatorTree &VPDT) const {
+  ArrayRef<VPBlockBase *> Preds = getPredecessors();
+  if (Preds.size() != 2)
+    return std::nullopt;
+
+  for (unsigned Idx : {0, 1}) {
+    auto *PreheaderVPBB = cast<VPBasicBlock>(Preds[Idx]);
+    auto *LatchVPBB = cast<VPBasicBlock>(Preds[1 - Idx]);
+    if (VPDT.dominates(PreheaderVPBB, this) && VPDT.dominates(this, LatchVPBB))
+      return {std::make_pair(PreheaderVPBB, LatchVPBB)};
+  }
+
+  return std::nullopt;
 }
 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)

llvm/lib/Transforms/Vectorize/VPlan.h

@@ -3255,10 +3255,13 @@ class VPBasicBlock : public VPBlockBase {
   /// Returns true if the block is exiting it's parent region.
   bool isExiting() const;
 
-  /// Returns true if the block is a loop header block in the plain CFG; that
-  /// is, it has exactly 2 predecessors (preheader and latch), where the block
-  /// dominates the latch and the preheader dominates the block.
-  bool isHeader(const VPDominatorTree &VPDT) const;
+  /// Checks if the block is a loop header block in the plain CFG; that is, it
+  /// has exactly 2 predecessors (preheader and latch), where the block
+  /// dominates the latch and the preheader dominates the block. If it is a
+  /// header block, returns a pair with the corresponding preheader and latch
+  /// blocks. Otherwise return std::nullopt.
+  std::optional<std::pair<VPBasicBlock *, VPBasicBlock *>>
+  isHeader(const VPDominatorTree &VPDT) const;
 
   /// Clone the current block and it's recipes, without updating the operands of
   /// the cloned recipes.

llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp

@@ -23,9 +23,9 @@ using namespace llvm;
 
 /// Create and return a new VPRegionBlock for loop starting at \p HeaderVPBB and
 /// return it.
-static VPRegionBlock *introduceRegion(VPlan &Plan, VPBlockBase *HeaderVPBB) {
-  VPBlockBase *PreheaderVPBB = HeaderVPBB->getPredecessors()[0];
-  VPBlockBase *LatchVPBB = HeaderVPBB->getPredecessors()[1];
+static VPRegionBlock *introduceRegion(VPlan &Plan, VPBasicBlock *PreheaderVPBB,
+                                      VPBasicBlock *HeaderVPBB,
+                                      VPBasicBlock *LatchVPBB) {
   VPBlockUtils::disconnectBlocks(PreheaderVPBB, HeaderVPBB);
   VPBlockUtils::disconnectBlocks(LatchVPBB, HeaderVPBB);
   VPBlockBase *Succ = LatchVPBB->getSingleSuccessor();
@@ -56,9 +56,11 @@ void VPlanTransforms::introduceRegions(VPlan &Plan, Type *InductionTy,
   VPDT.recalculate(Plan);
   for (VPBasicBlock *HeaderVPBB : VPBlockUtils::blocksOnly<VPBasicBlock>(
            vp_depth_first_shallow(Plan.getEntry()))) {
-    if (!HeaderVPBB->isHeader(VPDT))
+    auto Res = HeaderVPBB->isHeader(VPDT);
+    if (!Res)
       continue;
-    introduceRegion(Plan, HeaderVPBB);
+    const auto &[PreheaderVPBB, LatchVPBB] = *Res;
+    introduceRegion(Plan, PreheaderVPBB, HeaderVPBB, LatchVPBB);
   }
 
   VPRegionBlock *TopRegion = Plan.getVectorLoopRegion();

llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.cpp

@@ -313,22 +313,11 @@ void PlainCFGBuilder::buildPlainCFG(
     VPBasicBlock *VPBB = getOrCreateVPBB(BB);
     Loop *LoopForBB = LI->getLoopFor(BB);
     // Set VPBB predecessors in the same order as they are in the incoming BB.
-    if (!isHeaderBB(BB, LoopForBB)) {
-      setVPBBPredsFromBB(VPBB, BB);
-    } else {
-      VPBB->setPredecessors({getOrCreateVPBB(LoopForBB->getLoopPredecessor()),
-                             getOrCreateVPBB(LoopForBB->getLoopLatch())});
-    }
+    setVPBBPredsFromBB(VPBB, BB);
 
     // Create VPInstructions for BB.
     createVPInstructionsForVPBB(VPBB, BB);
 
-    if (BB == TheLoop->getLoopLatch()) {
-      VPBasicBlock *HeaderVPBB = getOrCreateVPBB(LoopForBB->getHeader());
-      VPBB->setOneSuccessor(HeaderVPBB);
-      continue;
-    }
-
     // Set VPBB successors. We create empty VPBBs for successors if they don't
     // exist already. Recipes will be created when the successor is visited
     // during the RPO traversal.