[VPlan] Retain exit conditions early

Author: fhahn

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -9324,6 +9324,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);
 
@@ -9502,12 +9520,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 =
@@ -9604,6 +9616,17 @@ VPlanPtr LoopVectorizationPlanner::tryToBuildVPlan(VFRange &Range) {
   auto Plan = VPlanTransforms::buildPlainCFG(OrigLoop, *LI, VPB2IRBB);
   VPlanTransforms::introduceTopLevelVectorLoopRegion(
       *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);

llvm/lib/Transforms/Vectorize/VPlan.cpp

@@ -855,7 +855,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));
 }

llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp

@@ -123,6 +123,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");
@@ -156,14 +159,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);
 }
@@ -212,11 +207,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()) {
@@ -319,6 +311,13 @@ void PlainCFGBuilder::buildPlainCFG(
     if (BB == TheLoop->getLoopLatch()) {
       VPBasicBlock *HeaderVPBB = getOrCreateVPBB(LoopForBB->getHeader());
       VPBlockUtils::connectBlocks(VPBB, HeaderVPBB);
+      assert(isa<BranchInst>(BB->getTerminator()) && "latch must be terminated by branch"
+                            );
+      for (BasicBlock *IRSucc : successors(BB)) {
+        VPBasicBlock *VPSucc = getOrCreateVPBB(IRSucc);
+        if (VPSucc != HeaderVPBB)
+          VPBB->getSuccessors().push_back(VPSucc);
+      }
       continue;
     }
 
@@ -349,24 +348,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.
@@ -448,6 +435,11 @@ void VPlanTransforms::introduceTopLevelVectorLoopRegion(
   VPBasicBlock *OriginalLatch =
       cast<VPBasicBlock>(HeaderVPBB->getSinglePredecessor());
   VPBlockUtils::disconnectBlocks(OriginalLatch, HeaderVPBB);
+  if (auto *RemainingSucc = OriginalLatch->getSingleSuccessor())
+    VPBlockUtils::disconnectBlocks(OriginalLatch,
+                                   RemainingSucc);
+  else
+    assert(OriginalLatch->getSuccessors().empty() && "Unsupported number of successors");
   VPBasicBlock *VecPreheader = Plan.createVPBasicBlock("vector.ph");
   VPBlockUtils::connectBlocks(Plan.getEntry(), VecPreheader);
   assert(OriginalLatch->getNumSuccessors() == 0 &&
@@ -473,8 +465,12 @@ void VPlanTransforms::introduceTopLevelVectorLoopRegion(
       HeaderVPBB, LatchVPBB, "vector loop", 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(TopRegion);
+  SmallPtrSet<VPBlockBase *, 2> ExitBlocks(Plan.getExitBlocks().begin(),
+                                           Plan.getExitBlocks().end());
+  for (VPBlockBase *VPBB : vp_depth_first_shallow(HeaderVPBB)) {
+    if (!ExitBlocks.contains(VPBB))
+      VPBB->setParent(TopRegion);
+  }
 
   VPBlockUtils::insertBlockAfter(TopRegion, VecPreheader);
   VPBasicBlock *MiddleVPBB = Plan.createVPBasicBlock("middle.block");
@@ -503,7 +499,7 @@ void VPlanTransforms::introduceTopLevelVectorLoopRegion(
   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, VPExitBlock);
   VPBlockUtils::connectBlocks(MiddleVPBB, ScalarPH);
 
   auto *ScalarLatchTerm = TheLoop->getLoopLatch()->getTerminator();
@@ -522,5 +518,8 @@ void VPlanTransforms::introduceTopLevelVectorLoopRegion(
   Builder.createNaryOp(VPInstruction::BranchOnCond, {Cmp},
                        ScalarLatchTerm->getDebugLoc());
 
-  introduceInnerLoopRegions(Plan);
+  if (all_of(Plan.getExitBlocks(), [MiddleVPBB](VPBlockBase *EB) {
+        return EB->getSinglePredecessor() == MiddleVPBB;
+      }))
+    introduceInnerLoopRegions(Plan);
 }

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

@@ -2275,7 +2275,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());
@@ -2286,17 +2286,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});
 

llvm/lib/Transforms/Vectorize/VPlanTransforms.h

@@ -177,8 +177,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);

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): vector.body,  ir-bb<exit>
+; CHECK-EMPTY:
+; CHECK-NEXT: ir-bb<exit>:
+; CHECK-NEXT: No successors
 ; CHECK-NEXT: }
 entry:
   br label %outer.header