[VPlan] Retain exit conditions early

fhahn · fhahn · commit 317d975da01c · 2025-04-05T14:16:02.000+01:00
diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -9350,6 +9350,24 @@ LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(VFRange &Range) {
   bool HasNUW = !IVUpdateMayOverflow || Style == TailFoldingStyle::None;
   addCanonicalIVRecipes(*Plan, Legal->getWidestInductionType(), HasNUW, DL);
 
+  if (auto *UncountableExitingBlock =
+          Legal->getUncountableEarlyExitingBlock()) {
+    VPlanTransforms::runPass(VPlanTransforms::handleUncountableEarlyExit, *Plan,
+                             *PSE.getSE(), OrigLoop, UncountableExitingBlock);
+  } else {
+    SmallPtrSet<VPBlockBase *, 2> ExitBlocks(Plan->getExitBlocks().begin(),
+                                             Plan->getExitBlocks().end());
+    for (VPBlockBase *VPBB : to_vector(
+             vp_depth_first_shallow(Plan->getVectorLoopRegion()->getEntry()))) {
+      for (VPBlockBase *EB : ExitBlocks) {
+        if (is_contained(VPBB->getSuccessors(), EB)) {
+          cast<VPBasicBlock>(VPBB)->getTerminator()->eraseFromParent();
+          VPBlockUtils::disconnectBlocks(VPBB, EB);
+        }
+      }
+    }
+  }
+
   VPRecipeBuilder RecipeBuilder(*Plan, OrigLoop, TLI, &TTI, Legal, CM, PSE,
                                 Builder);
 
@@ -9528,12 +9546,6 @@ LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(VFRange &Range) {
     R->setOperand(1, WideIV->getStepValue());
   }
 
-  if (auto *UncountableExitingBlock =
-          Legal->getUncountableEarlyExitingBlock()) {
-    VPlanTransforms::runPass(VPlanTransforms::handleUncountableEarlyExit, *Plan,
-                             *PSE.getSE(), OrigLoop, UncountableExitingBlock,
-                             RecipeBuilder);
-  }
   DenseMap<VPValue *, VPValue *> IVEndValues;
   addScalarResumePhis(RecipeBuilder, *Plan, IVEndValues);
   SetVector<VPIRInstruction *> ExitUsersToFix =
@@ -9631,6 +9643,17 @@ VPlanPtr LoopVectorizationPlanner::tryToBuildVPlan(VFRange &Range) {
   auto Plan = VPlanTransforms::buildPlainCFG(OrigLoop, *LI, VPB2IRBB);
   VPlanTransforms::introduceRegions(*Plan, Legal->getWidestInductionType(), PSE,
                                     true, false, OrigLoop);
+  SmallPtrSet<VPBlockBase *, 2> ExitBlocks(Plan->getExitBlocks().begin(),
+                                           Plan->getExitBlocks().end());
+  for (VPBlockBase *VPBB : to_vector(
+           vp_depth_first_shallow(Plan->getVectorLoopRegion()->getEntry()))) {
+    for (VPBlockBase *EB : ExitBlocks) {
+      if (is_contained(VPBB->getSuccessors(), EB)) {
+        cast<VPBasicBlock>(VPBB)->getTerminator()->eraseFromParent();
+        VPBlockUtils::disconnectBlocks(VPBB, EB);
+      }
+    }
+  }
 
   for (ElementCount VF : Range)
     Plan->addVF(VF);
diff --git a/llvm/lib/Transforms/Vectorize/VPlan.cpp b/llvm/lib/Transforms/Vectorize/VPlan.cpp
@@ -655,9 +655,10 @@ VPBasicBlock::isHeader(const VPDominatorTree &VPDT) const {
     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))
+    auto *PreheaderVPBB = dyn_cast<VPBasicBlock>(Preds[Idx]);
+    auto *LatchVPBB = dyn_cast<VPBasicBlock>(Preds[1 - Idx]);
+    if (PreheaderVPBB && LatchVPBB && VPDT.dominates(PreheaderVPBB, this) &&
+        VPDT.dominates(this, LatchVPBB))
       return {std::make_pair(PreheaderVPBB, LatchVPBB)};
   }
 
@@ -871,7 +872,7 @@ VPlan::VPlan(Loop *L) {
   ScalarHeader = createVPIRBasicBlock(L->getHeader());
 
   SmallVector<BasicBlock *> IRExitBlocks;
-  L->getExitBlocks(IRExitBlocks);
+  L->getUniqueExitBlocks(IRExitBlocks);
   for (BasicBlock *EB : IRExitBlocks)
     ExitBlocks.push_back(createVPIRBasicBlock(EB));
 }
diff --git a/llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp b/llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp
@@ -122,6 +122,9 @@ VPBasicBlock *PlainCFGBuilder::getOrCreateVPBB(BasicBlock *BB) {
     return VPBB;
   }
 
+  if (!TheLoop->contains(BB))
+    return Plan.getExitBlock(BB);
+
   // Create new VPBB.
   StringRef Name = isHeaderBB(BB, TheLoop) ? "vector.body" : BB->getName();
   LLVM_DEBUG(dbgs() << "Creating VPBasicBlock for " << Name << "\n");
@@ -155,14 +158,6 @@ bool PlainCFGBuilder::isExternalDef(Value *Val) {
     // Instruction definition is in outermost loop PH.
     return false;
 
-  // Check whether Instruction definition is in a loop exit.
-  SmallVector<BasicBlock *> ExitBlocks;
-  TheLoop->getExitBlocks(ExitBlocks);
-  if (is_contained(ExitBlocks, InstParent)) {
-    // Instruction definition is in outermost loop exit.
-    return false;
-  }
-
   // Check whether Instruction definition is in loop body.
   return !TheLoop->contains(Inst);
 }
@@ -211,11 +206,8 @@ void PlainCFGBuilder::createVPInstructionsForVPBB(VPBasicBlock *VPBB,
            "Instruction shouldn't have been visited.");
 
     if (auto *Br = dyn_cast<BranchInst>(Inst)) {
-      if (TheLoop->getLoopLatch() == BB ||
-          any_of(successors(BB),
-                 [this](BasicBlock *Succ) { return !TheLoop->contains(Succ); }))
+      if (TheLoop->getLoopLatch() == BB)
         continue;
-
       // Conditional branch instruction are represented using BranchOnCond
       // recipes.
       if (Br->isConditional()) {
@@ -305,7 +297,6 @@ void PlainCFGBuilder::buildPlainCFG(
   for (BasicBlock *BB : RPO) {
     // Create or retrieve the VPBasicBlock for this BB.
     VPBasicBlock *VPBB = getOrCreateVPBB(BB);
-    Loop *LoopForBB = LI.getLoopFor(BB);
     // Set VPBB predecessors in the same order as they are in the incoming BB.
     setVPBBPredsFromBB(VPBB, BB);
 
@@ -339,24 +330,12 @@ void PlainCFGBuilder::buildPlainCFG(
     BasicBlock *IRSucc1 = BI->getSuccessor(1);
     VPBasicBlock *Successor0 = getOrCreateVPBB(IRSucc0);
     VPBasicBlock *Successor1 = getOrCreateVPBB(IRSucc1);
-
-    // Don't connect any blocks outside the current loop except the latch, which
-    // is handled below.
-    if (LoopForBB &&
-        (LoopForBB == TheLoop || BB != LoopForBB->getLoopLatch())) {
-      if (!LoopForBB->contains(IRSucc0)) {
-        VPBB->setOneSuccessor(Successor1);
-        continue;
-      }
-      if (!LoopForBB->contains(IRSucc1)) {
-        VPBB->setOneSuccessor(Successor0);
-        continue;
-      }
-    }
-
     VPBB->setTwoSuccessors(Successor0, Successor1);
   }
 
+  for (auto *EB : Plan.getExitBlocks()) {
+    setVPBBPredsFromBB(EB, EB->getIRBasicBlock());
+  }
   // 2. The whole CFG has been built at this point so all the input Values must
   // have a VPlan counterpart. Fix VPlan header phi by adding their
   // corresponding VPlan operands.
@@ -413,10 +392,15 @@ static VPRegionBlock *introduceRegion(VPlan &Plan, VPBasicBlock *PreheaderVPBB,
   auto *R = Plan.createVPRegionBlock(HeaderVPBB, LatchVPBB, "",
                                      false /*isReplicator*/);
   R->setParent(HeaderVPBB->getParent());
+
   // All VPBB's reachable shallowly from HeaderVPBB 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))
-    VPBB->setParent(R);
+  SmallPtrSet<VPBlockBase *, 2> ExitBlocks(Plan.getExitBlocks().begin(),
+                                           Plan.getExitBlocks().end());
+  for (VPBlockBase *VPBB : vp_depth_first_shallow(HeaderVPBB)) {
+    if (!ExitBlocks.contains(VPBB))
+      VPBB->setParent(R);
+  }
 
   VPBlockUtils::insertBlockAfter(R, PreheaderVPBB);
   if (Succ)
@@ -466,7 +450,11 @@ void VPlanTransforms::introduceRegions(VPlan &Plan, Type *InductionTy,
 
   VPBasicBlock *ScalarPH = Plan.createVPBasicBlock("scalar.ph");
   VPBlockUtils::connectBlocks(ScalarPH, Plan.getScalarHeader());
+  BasicBlock *IRExitBlock = TheLoop->getUniqueLatchExitBlock();
+  auto *VPExitBlock = IRExitBlock ? Plan.getExitBlock(IRExitBlock) : nullptr;
   if (!RequiresScalarEpilogueCheck) {
+    if (VPExitBlock)
+      VPBlockUtils::disconnectBlocks(MiddleVPBB, VPExitBlock);
     VPBlockUtils::connectBlocks(MiddleVPBB, ScalarPH);
     for (auto *EB : Plan.getExitBlocks()) {
       for (VPRecipeBase &R : *EB)
@@ -484,10 +472,7 @@ void VPlanTransforms::introduceRegions(VPlan &Plan, Type *InductionTy,
   // 2) If we require a scalar epilogue, there is no conditional branch as
   //    we unconditionally branch to the scalar preheader.  Do nothing.
   // 3) Otherwise, construct a runtime check.
-  BasicBlock *IRExitBlock = TheLoop->getUniqueLatchExitBlock();
-  auto *VPExitBlock = Plan.getExitBlock(IRExitBlock);
   // The connection order corresponds to the operands of the conditional branch.
-  VPBlockUtils::insertBlockAfter(VPExitBlock, MiddleVPBB);
   VPBlockUtils::connectBlocks(MiddleVPBB, ScalarPH);
 
   auto *ScalarLatchTerm = TheLoop->getLoopLatch()->getTerminator();
diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
@@ -2368,7 +2368,7 @@ void VPlanTransforms::convertToConcreteRecipes(VPlan &Plan) {
 
 void VPlanTransforms::handleUncountableEarlyExit(
     VPlan &Plan, ScalarEvolution &SE, Loop *OrigLoop,
-    BasicBlock *UncountableExitingBlock, VPRecipeBuilder &RecipeBuilder) {
+    BasicBlock *UncountableExitingBlock) {
   VPRegionBlock *LoopRegion = Plan.getVectorLoopRegion();
   auto *LatchVPBB = cast<VPBasicBlock>(LoopRegion->getExiting());
   VPBuilder Builder(LatchVPBB->getTerminator());
@@ -2379,17 +2379,29 @@ void VPlanTransforms::handleUncountableEarlyExit(
   // tracks if the uncountable early exit has been taken. Also split the middle
   // block and have it conditionally branch to the early exit block if
   // EarlyExitTaken.
-  auto *EarlyExitingBranch =
-      cast<BranchInst>(UncountableExitingBlock->getTerminator());
-  BasicBlock *TrueSucc = EarlyExitingBranch->getSuccessor(0);
-  BasicBlock *FalseSucc = EarlyExitingBranch->getSuccessor(1);
-  BasicBlock *EarlyExitIRBB =
-      !OrigLoop->contains(TrueSucc) ? TrueSucc : FalseSucc;
-  VPIRBasicBlock *VPEarlyExitBlock = Plan.getExitBlock(EarlyExitIRBB);
-
-  VPValue *EarlyExitNotTakenCond = RecipeBuilder.getBlockInMask(
-      OrigLoop->contains(TrueSucc) ? TrueSucc : FalseSucc);
-  auto *EarlyExitTakenCond = Builder.createNot(EarlyExitNotTakenCond);
+  VPBasicBlock *EEB = nullptr;
+  for (auto *EB : Plan.getExitBlocks()) {
+    for (VPBlockBase *Pred : EB->getPredecessors()) {
+      if (Pred != MiddleVPBB) {
+        EEB = cast<VPBasicBlock>(Pred);
+        break;
+      }
+    }
+  }
+
+  VPBlockBase *TrueSucc = EEB->getSuccessors()[0];
+  VPBlockBase *FalseSucc = EEB->getSuccessors()[1];
+  auto *VPEarlyExitBlock =
+      cast<VPIRBasicBlock>(TrueSucc->getParent() ? FalseSucc : TrueSucc);
+
+  VPValue *EarlyExitCond = EEB->getTerminator()->getOperand(0);
+  auto *EarlyExitTakenCond = TrueSucc == VPEarlyExitBlock
+                                 ? EarlyExitCond
+                                 : Builder.createNot(EarlyExitCond);
+
+  EEB->getTerminator()->eraseFromParent();
+  VPBlockUtils::disconnectBlocks(EEB, VPEarlyExitBlock);
+
   IsEarlyExitTaken =
       Builder.createNaryOp(VPInstruction::AnyOf, {EarlyExitTakenCond});
 
diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.h b/llvm/lib/Transforms/Vectorize/VPlanTransforms.h
@@ -176,8 +176,7 @@ struct VPlanTransforms {
   ///    if taken.
   static void handleUncountableEarlyExit(VPlan &Plan, ScalarEvolution &SE,
                                          Loop *OrigLoop,
-                                         BasicBlock *UncountableExitingBlock,
-                                         VPRecipeBuilder &RecipeBuilder);
+                                         BasicBlock *UncountableExitingBlock);
 
   /// Lower abstract recipes to concrete ones, that can be codegen'd.
   static void convertToConcreteRecipes(VPlan &Plan);
diff --git a/llvm/test/Transforms/LoopVectorize/vplan-printing-outer-loop.ll b/llvm/test/Transforms/LoopVectorize/vplan-printing-outer-loop.ll
@@ -31,7 +31,10 @@ define void @foo(i64 %n) {
 ; CHECK-NEXT: outer.latch:
 ; CHECK-NEXT:   EMIT ir<%outer.iv.next> = add ir<%outer.iv>, ir<1>
 ; CHECK-NEXT:   EMIT ir<%outer.ec> = icmp ir<%outer.iv.next>, ir<8>
-; CHECK-NEXT: Successor(s): vector.body
+; CHECK-NEXT: Successor(s): ir-bb<exit>, vector.body
+; CHECK-EMPTY:
+; CHECK-NEXT: ir-bb<exit>:
+; CHECK-NEXT: No successors
 ; CHECK-NEXT: }
 entry:
   br label %outer.header
