!fixup address latest comments, thanks

Author: fhahn

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -9317,9 +9317,9 @@ LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(VFRange &Range) {
   VPlanHCFGBuilder HCFGBuilder(OrigLoop, LI, *Plan);
   HCFGBuilder.buildPlainCFG();
 
-  VPlanTransforms::introduceRegions(*Plan, Legal->getWidestInductionType(), PSE,
-                                    RequiresScalarEpilogueCheck,
-                                    CM.foldTailByMasking(), OrigLoop);
+  VPlanTransforms::createLoopRegions(*Plan, Legal->getWidestInductionType(),
+                                     PSE, RequiresScalarEpilogueCheck,
+                                     CM.foldTailByMasking(), OrigLoop);
 
   // Don't use getDecisionAndClampRange here, because we don't know the UF
   // so this function is better to be conservative, rather than to split
@@ -9621,8 +9621,8 @@ VPlanPtr LoopVectorizationPlanner::tryToBuildVPlan(VFRange &Range) {
   VPlanHCFGBuilder HCFGBuilder(OrigLoop, LI, *Plan);
   HCFGBuilder.buildPlainCFG();
 
-  VPlanTransforms::introduceRegions(*Plan, Legal->getWidestInductionType(), PSE,
-                                    true, false, OrigLoop);
+  VPlanTransforms::createLoopRegions(*Plan, Legal->getWidestInductionType(),
+                                     PSE, true, false, OrigLoop);
 
   for (ElementCount VF : Range)
     Plan->addVF(VF);

llvm/lib/Transforms/Vectorize/VPlan.cpp

@@ -652,22 +652,6 @@ bool VPBasicBlock::isExiting() const {
   return getParent() && getParent()->getExitingBasicBlock() == this;
 }
 
-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)
 void VPBlockBase::print(raw_ostream &O) const {
   VPSlotTracker SlotTracker(getPlan());

llvm/lib/Transforms/Vectorize/VPlan.h

@@ -3255,14 +3255,6 @@ class VPBasicBlock : public VPBlockBase {
   /// Returns true if the block is exiting it's parent region.
   bool isExiting() 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.
   VPBasicBlock *clone() override;

llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp

@@ -21,47 +21,79 @@
 
 using namespace llvm;
 
-/// Create and return a new VPRegionBlock for loop starting at \p HeaderVPBB and
-/// return it.
-static VPRegionBlock *introduceRegion(VPlan &Plan, VPBasicBlock *PreheaderVPBB,
-                                      VPBasicBlock *HeaderVPBB,
-                                      VPBasicBlock *LatchVPBB) {
-  VPBlockUtils::disconnectBlocks(PreheaderVPBB, HeaderVPBB);
-  VPBlockUtils::disconnectBlocks(LatchVPBB, HeaderVPBB);
+/// Checks if \p HeaderVPB 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.
+static std::optional<std::pair<VPBasicBlock *, VPBasicBlock *>>
+getPreheaderAndLatch(VPBlockBase *HeaderVPB, const VPDominatorTree &VPDT) {
+  ArrayRef<VPBlockBase *> Preds = HeaderVPB->getPredecessors();
+  if (Preds.size() != 2)
+    return std::nullopt;
+
+  auto *PreheaderVPBB = cast<VPBasicBlock>(Preds[0]);
+  auto *LatchVPBB = cast<VPBasicBlock>(Preds[1]);
+  if (VPDT.dominates(PreheaderVPBB, HeaderVPB) &&
+      VPDT.dominates(HeaderVPB, LatchVPBB))
+    return {std::make_pair(PreheaderVPBB, LatchVPBB)};
+
+  std::swap(PreheaderVPBB, LatchVPBB);
+  if (VPDT.dominates(PreheaderVPBB, HeaderVPB) &&
+      VPDT.dominates(HeaderVPB, LatchVPBB))
+    return {std::make_pair(PreheaderVPBB, LatchVPBB)};
+
+  return std::nullopt;
+}
+
+/// Create a new VPRegionBlock if there is a loop starting at \p HeaderVPB.
+static void createLoopRegion(VPlan &Plan, VPBlockBase *HeaderVPB,
+                             VPDominatorTree &VPDT) {
+  auto Res = getPreheaderAndLatch(HeaderVPB, VPDT);
+  if (!Res)
+    return;
+
+  const auto &[PreheaderVPBB, LatchVPBB] = *Res;
+
+  // Swap the operands of header phis if needed. After creating the region, the
+  // incoming value from the preheader must be the first operand and the one
+  // from the latch must be the second operand.
+  if (HeaderVPB->getPredecessors()[0] != PreheaderVPBB) {
+    for (VPRecipeBase &R : cast<VPBasicBlock>(HeaderVPB)->phis()) {
+      VPValue *Inc0 = R.getOperand(0);
+      R.setOperand(0, R.getOperand(1));
+      R.setOperand(1, Inc0);
+    }
+  }
+  VPBlockUtils::disconnectBlocks(PreheaderVPBB, HeaderVPB);
+  VPBlockUtils::disconnectBlocks(LatchVPBB, HeaderVPB);
   VPBlockBase *Succ = LatchVPBB->getSingleSuccessor();
   assert(LatchVPBB->getNumSuccessors() <= 1 &&
          "Latch has more than one successor");
   if (Succ)
     VPBlockUtils::disconnectBlocks(LatchVPBB, Succ);
 
-  auto *R = Plan.createVPRegionBlock(HeaderVPBB, LatchVPBB, "",
+  auto *R = Plan.createVPRegionBlock(HeaderVPB, LatchVPBB, "",
                                      false /*isReplicator*/);
-  R->setParent(HeaderVPBB->getParent());
-  // All VPBB's reachable shallowly from HeaderVPBB belong to top level loop,
+  R->setParent(HeaderVPB->getParent());
+  // All VPBB's reachable shallowly from HeaderVPB belong to top level loop,
   // because VPlan is expected to end at top level latch disconnected above.
-  for (VPBlockBase *VPBB : vp_depth_first_shallow(HeaderVPBB))
+  for (VPBlockBase *VPBB : vp_depth_first_shallow(HeaderVPB))
     VPBB->setParent(R);
 
   VPBlockUtils::insertBlockAfter(R, PreheaderVPBB);
   if (Succ)
     VPBlockUtils::connectBlocks(R, Succ);
-  return R;
 }
 
-void VPlanTransforms::introduceRegions(VPlan &Plan, Type *InductionTy,
-                                       PredicatedScalarEvolution &PSE,
-                                       bool RequiresScalarEpilogueCheck,
-                                       bool TailFolded, Loop *TheLoop) {
+void VPlanTransforms::createLoopRegions(VPlan &Plan, Type *InductionTy,
+                                        PredicatedScalarEvolution &PSE,
+                                        bool RequiresScalarEpilogueCheck,
+                                        bool TailFolded, Loop *TheLoop) {
   VPDominatorTree VPDT;
   VPDT.recalculate(Plan);
-  for (VPBasicBlock *HeaderVPBB : VPBlockUtils::blocksOnly<VPBasicBlock>(
-           vp_depth_first_shallow(Plan.getEntry()))) {
-    auto Res = HeaderVPBB->isHeader(VPDT);
-    if (!Res)
-      continue;
-    const auto &[PreheaderVPBB, LatchVPBB] = *Res;
-    introduceRegion(Plan, PreheaderVPBB, HeaderVPBB, LatchVPBB);
-  }
+  for (VPBlockBase *HeaderVPB : vp_depth_first_shallow(Plan.getEntry()))
+    createLoopRegion(Plan, HeaderVPB, VPDT);
 
   VPRegionBlock *TopRegion = Plan.getVectorLoopRegion();
   auto *OrigExiting = TopRegion->getExiting();

llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.cpp

@@ -98,19 +98,14 @@ void PlainCFGBuilder::fixHeaderPhis() {
     auto *VPPhi = cast<VPWidenPHIRecipe>(VPVal);
     assert(VPPhi->getNumOperands() == 0 &&
            "Expected VPInstruction with no operands.");
-
-    Loop *L = LI->getLoopFor(Phi->getParent());
-    assert(isHeaderBB(Phi->getParent(), L));
+    assert(isHeaderBB(Phi->getParent(), LI->getLoopFor(Phi->getParent())));
     // For header phis, make sure the incoming value from the loop
     // predecessor is the first operand of the recipe.
     assert(Phi->getNumOperands() == 2 &&
            "header phi must have exactly 2 operands");
-    BasicBlock *LoopPred = L->getLoopPredecessor();
-    VPPhi->addOperand(
-        getOrCreateVPOperand(Phi->getIncomingValueForBlock(LoopPred)));
-    BasicBlock *LoopLatch = L->getLoopLatch();
-    VPPhi->addOperand(
-        getOrCreateVPOperand(Phi->getIncomingValueForBlock(LoopLatch)));
+    for (BasicBlock *Pred : predecessors(Phi->getParent()))
+      VPPhi->addOperand(
+          getOrCreateVPOperand(Phi->getIncomingValueForBlock(Pred)));
   }
 }
 
@@ -344,8 +339,8 @@ void PlainCFGBuilder::buildPlainCFG(
     VPBasicBlock *Successor0 = getOrCreateVPBB(IRSucc0);
     VPBasicBlock *Successor1 = getOrCreateVPBB(IRSucc1);
 
-    // Don't connect any blocks outside the current loop except the latch, which
-    // is handled below.
+    // Don't connect any blocks outside the current loop except the latches for
+    // inner loops.
     if (LoopForBB &&
         (LoopForBB == TheLoop || BB != LoopForBB->getLoopLatch())) {
       if (!LoopForBB->contains(IRSucc0)) {

llvm/lib/Transforms/Vectorize/VPlanTransforms.h

@@ -61,10 +61,10 @@ struct VPlanTransforms {
   /// 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 introduceRegions(VPlan &Plan, Type *InductionTy,
-                               PredicatedScalarEvolution &PSE,
-                               bool RequiresScalarEpilogueCheck,
-                               bool TailFolded, Loop *TheLoop);
+  static void createLoopRegions(VPlan &Plan, Type *InductionTy,
+                                PredicatedScalarEvolution &PSE,
+                                bool RequiresScalarEpilogueCheck,
+                                bool TailFolded, Loop *TheLoop);
 
   /// Replaces the VPInstructions in \p Plan with corresponding
   /// widen recipes. Returns false if any VPInstructions could not be converted

llvm/test/Transforms/LoopVectorize/vplan-printing-outer-loop.ll

@@ -13,14 +13,14 @@ define void @foo(i64 %n) {
 ; CHECK-NEXT: Successor(s): vector.body
 ; CHECK-EMPTY:
 ; CHECK-NEXT: vector.body:
-; CHECK-NEXT:   WIDEN-PHI ir<%outer.iv> = phi ir<0>, ir<%outer.iv.next>
+; CHECK-NEXT:   WIDEN-PHI ir<%outer.iv> = phi ir<%outer.iv.next>, ir<0>
 ; CHECK-NEXT:   EMIT ir<%gep.1> = getelementptr ir<@arr2>, ir<0>, ir<%outer.iv>
 ; CHECK-NEXT:   EMIT store ir<%outer.iv>, ir<%gep.1>
 ; CHECK-NEXT:   EMIT ir<%add> = add ir<%outer.iv>, ir<%n>
 ; CHECK-NEXT: Successor(s): inner
 ; CHECK-EMPTY:
 ; CHECK-NEXT: inner:
-; CHECK-NEXT:   WIDEN-PHI ir<%inner.iv> = phi ir<0>, ir<%inner.iv.next>
+; CHECK-NEXT:   WIDEN-PHI ir<%inner.iv> = phi ir<%inner.iv.next>, ir<0>
 ; CHECK-NEXT:   EMIT ir<%gep.2> = getelementptr ir<@arr>, ir<0>, ir<%inner.iv>, ir<%outer.iv>
 ; CHECK-NEXT:   EMIT store ir<%add>, ir<%gep.2>
 ; CHECK-NEXT:   EMIT ir<%inner.iv.next> = add ir<%inner.iv>, ir<1>

llvm/unittests/Transforms/Vectorize/VPlanTestBase.h

@@ -74,8 +74,8 @@ class VPlanTestIRBase : public testing::Test {
     auto Plan = std::make_unique<VPlan>(L);
     VPlanHCFGBuilder HCFGBuilder(L, LI.get(), *Plan);
     HCFGBuilder.buildPlainCFG();
-    VPlanTransforms::introduceRegions(*Plan, IntegerType::get(*Ctx, 64), PSE,
-                                      true, false, L);
+    VPlanTransforms::createLoopRegions(*Plan, IntegerType::get(*Ctx, 64), PSE,
+                                       true, false, L);
     return Plan;
   }
 };