[llvm] [VPlan] Introduce child regions as VPlan transform. (PR #129402)

[Florian Hahn via llvm-commits]

================
@@ -14,26 +14,95 @@
 #include "LoopVectorizationPlanner.h"
 #include "VPlan.h"
 #include "VPlanCFG.h"
+#include "VPlanDominatorTree.h"
 #include "VPlanTransforms.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/ScalarEvolution.h"
 
 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);
+/// 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;
 
-  VPBasicBlock *OriginalLatch =
-      cast<VPBasicBlock>(HeaderVPBB->getSinglePredecessor());
-  VPBlockUtils::disconnectBlocks(OriginalLatch, HeaderVPBB);
-  VPBasicBlock *VecPreheader = Plan.createVPBasicBlock("vector.ph");
-  VPBlockUtils::connectBlocks(Plan.getEntry(), VecPreheader);
-  assert(OriginalLatch->getNumSuccessors() == 0 &&
-         "Plan should end at top level latch");
+  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;
+}
+
+/// Try to create a new VPRegionBlock if there is a loop starting at \p
+/// HeaderVPB.
+static void tryToCreateLoopRegion(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(HeaderVPB, LatchVPBB, "",
+                                     false /*isReplicator*/);
+  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(HeaderVPB))
+    VPBB->setParent(R);
+
+  VPBlockUtils::insertBlockAfter(R, PreheaderVPBB);
+  if (Succ)
+    VPBlockUtils::connectBlocks(R, Succ);
+}
+
+void VPlanTransforms::createLoopRegions(VPlan &Plan, Type *InductionTy,
+                                        PredicatedScalarEvolution &PSE,
+                                        bool RequiresScalarEpilogueCheck,
+                                        bool TailFolded, Loop *TheLoop) {
+  VPDominatorTree VPDT;
+  VPDT.recalculate(Plan);
+  for (VPBlockBase *HeaderVPB : vp_depth_first_shallow(Plan.getEntry()))
+    tryToCreateLoopRegion(Plan, HeaderVPB, VPDT);
----------------
fhahn wrote:

done thanks!

https://github.com/llvm/llvm-project/pull/129402