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

[Florian Hahn via llvm-commits]

================
@@ -14,26 +14,57 @@
 #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;
 
+/// Create and return a new VPRegionBlock for loop starting at \p HeaderVPBB, if
+/// it is a header of a loop.
+static VPRegionBlock *introduceRegion(VPlan &Plan, VPBlockBase *HeaderVPBB,
+                                      VPDominatorTree &VPDT) {
+  if (HeaderVPBB->getNumPredecessors() != 2)
+    return nullptr;
+  VPBlockBase *PreheaderVPBB = HeaderVPBB->getPredecessors()[0];
+  VPBlockBase *LatchVPBB = HeaderVPBB->getPredecessors()[1];
+  if (!VPDT.dominates(HeaderVPBB, LatchVPBB))
+    return nullptr;
+  assert(VPDT.dominates(PreheaderVPBB, HeaderVPBB) &&
+         "preheader must dominate header");
+  VPBlockUtils::disconnectBlocks(PreheaderVPBB, HeaderVPBB);
+  VPBlockUtils::disconnectBlocks(LatchVPBB, HeaderVPBB);
+  VPBlockBase *Succ = LatchVPBB->getSingleSuccessor();
+  if (Succ)
+    VPBlockUtils::disconnectBlocks(LatchVPBB, Succ);
+
+  auto *R = Plan.createVPRegionBlock(HeaderVPBB, LatchVPBB, "",
+                                     false /*isReplicator*/);
+  // All VPBB's reachable shallowly from HeaderVPBB belong to top level loop,
+  // because VPlan is expected to end at top level latch.
+  for (VPBlockBase *VPBB : vp_depth_first_shallow(HeaderVPBB))
+    VPBB->setParent(R);
+
+  VPBlockUtils::insertBlockAfter(R, PreheaderVPBB);
+  if (Succ)
+    VPBlockUtils::connectBlocks(R, Succ);
+  return R;
+}
+
 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);
+  VPDominatorTree VPDT;
+  VPDT.recalculate(Plan);
 
-  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 *HeaderVPBB = cast<VPBasicBlock>(Plan.getEntry()->getSingleSuccessor());
+  VPRegionBlock *TopRegion = introduceRegion(Plan, HeaderVPBB, VPDT);
+  auto *OrigExiting = TopRegion->getExiting();
+  VPBasicBlock *LatchVPBB = Plan.createVPBasicBlock("vector.latch");
+  VPBlockUtils::insertBlockAfter(LatchVPBB, OrigExiting);
+  TopRegion->setExiting(LatchVPBB);
----------------
fhahn wrote:

It is unsed only for convenience, when adjusting reductions it is used to place selects there if needed. I'll look into remove it separartely.

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