!fixup address more comments.

Author: fhahn

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -8839,11 +8839,7 @@ static SetVector<VPIRInstruction *> collectUsersInExitBlocks(
     Loop *OrigLoop, VPRecipeBuilder &Builder, VPlan &Plan,
     const MapVector<PHINode *, InductionDescriptor> &Inductions) {
   SetVector<VPIRInstruction *> ExitUsersToFix;
-  for (VPBlockBase *VPB : vp_depth_first_shallow(
-           Plan.getVectorLoopRegion()->getSingleSuccessor())) {
-    if (VPB->getNumSuccessors() != 0 || VPB == Plan.getScalarHeader())
-      continue;
-    auto *ExitVPBB = cast<VPIRBasicBlock>(VPB);
+  for (VPIRBasicBlock *ExitVPBB : Plan.getExitBlocks()) {
     BasicBlock *ExitBB = ExitVPBB->getIRBasicBlock();
     for (VPRecipeBase &R : *ExitVPBB) {
       auto *ExitIRI = dyn_cast<VPIRInstruction>(&R);
@@ -9178,14 +9174,31 @@ LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(VFRange &Range) {
   RecipeBuilder.fixHeaderPhis();
 
   if (Legal->hasUncountableEarlyExit()) {
-    VPlanTransforms::handleUncountableEarlyExit(*Plan, *PSE.getSE(), OrigLoop,
-                                                RecipeBuilder);
+    VPlanTransforms::handleUncountableEarlyExit(
+        *Plan, *PSE.getSE(), OrigLoop, Legal->getUncountableExitingBlocks(),
+        RecipeBuilder);
   }
   addScalarResumePhis(RecipeBuilder, *Plan);
   SetVector<VPIRInstruction *> ExitUsersToFix = collectUsersInExitBlocks(
       OrigLoop, RecipeBuilder, *Plan, Legal->getInductionVars());
   addExitUsersForFirstOrderRecurrences(*Plan, ExitUsersToFix);
   addUsersInExitBlocks(*Plan, ExitUsersToFix);
+
+  // Currently only live-ins can be used by exit values. We also bail out if any
+  // exit value isn't handled in VPlan yet, i.e. a VPIRInstruction in the exit
+  // without any operands.
+  if (Legal->hasUncountableEarlyExit()) {
+    if (any_of(Plan->getExitBlocks(), [](VPIRBasicBlock *ExitBB) {
+          return any_of(*ExitBB, [](VPRecipeBase &R) {
+            auto VPIRI = cast<VPIRInstruction>(&R);
+            return VPIRI->getNumOperands() == 0 ||
+                   any_of(VPIRI->operands(),
+                          [](VPValue *Op) { return !Op->isLiveIn(); });
+          });
+        }))
+      return nullptr;
+  }
+
   // ---------------------------------------------------------------------------
   // Transform initial VPlan: Apply previously taken decisions, in order, to
   // bring the VPlan to its final state.

llvm/lib/Transforms/Vectorize/VPlan.h

@@ -3832,6 +3832,10 @@ class VPlan {
     return cast<VPBasicBlock>(getScalarPreheader()->getSinglePredecessor());
   }
 
+  /// Return the exit blocks of the VPlan, that is leaf nodes except the scalar
+  /// header.
+  auto getExitBlocks();
+
   /// The trip count of the original loop.
   VPValue *getTripCount() const {
     assert(TripCount && "trip count needs to be set before accessing it");

llvm/lib/Transforms/Vectorize/VPlanCFG.h

@@ -306,6 +306,15 @@ template <> struct GraphTraits<VPlan *> {
   }
 };
 
+inline auto VPlan::getExitBlocks() {
+  VPBlockBase *ScalarHeader = getScalarHeader();
+  return make_filter_range(
+      VPBlockUtils::blocksOnly<VPIRBasicBlock>(
+          vp_depth_first_shallow(getVectorLoopRegion()->getSingleSuccessor())),
+      [ScalarHeader](VPIRBasicBlock *VPIRBB) {
+        return VPIRBB != ScalarHeader && VPIRBB->getNumSuccessors() == 0;
+      });
+}
 } // namespace llvm
 
 #endif // LLVM_TRANSFORMS_VECTORIZE_VPLANCFG_H

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

@@ -872,7 +872,6 @@ void VPIRInstruction::print(raw_ostream &O, const Twine &Indent,
   O << Indent << "IR " << I;
 
   if (getNumOperands() != 0) {
-    // assert(getNumOperands() == 1 && "can have at most 1 operand");
     O << " (extra operand: ";
     printOperands(O, SlotTracker);
     O << ")";

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

@@ -1793,31 +1793,26 @@ void VPlanTransforms::createInterleaveGroups(
 
 void VPlanTransforms::handleUncountableEarlyExit(
     VPlan &Plan, ScalarEvolution &SE, Loop *OrigLoop,
+    ArrayRef<BasicBlock *> UncountableExitingBlocks,
     VPRecipeBuilder &RecipeBuilder) {
   auto *LatchVPBB =
       cast<VPBasicBlock>(Plan.getVectorLoopRegion()->getExiting());
   VPBuilder Builder(LatchVPBB->getTerminator());
   auto *MiddleVPBB = Plan.getMiddleBlock();
   VPRegionBlock *LoopRegion = Plan.getVectorLoopRegion();
-
   VPValue *EarlyExitTaken = nullptr;
-  SmallVector<BasicBlock *> ExitingBBs;
-  OrigLoop->getExitingBlocks(ExitingBBs);
 
   // Process all uncountable exiting blocks. For each exiting block, update the
   // EarlyExitTaken, which tracks if any uncountable early exit has been taken.
   // Also split the middle block and branch to the exit block for the early exit
   // if it has been taken.
-  for (BasicBlock *Exiting : ExitingBBs) {
-    if (Exiting == OrigLoop->getLoopLatch())
-      continue;
-
+  for (BasicBlock *Exiting : UncountableExitingBlocks) {
     auto *ExitingTerm = cast<BranchInst>(Exiting->getTerminator());
     BasicBlock *TrueSucc = ExitingTerm->getSuccessor(0);
     BasicBlock *FalseSucc = ExitingTerm->getSuccessor(1);
-    VPBasicBlock *VPExitBlock;
+    VPIRBasicBlock *VPExitBlock;
     if (OrigLoop->getUniqueExitBlock()) {
-      VPExitBlock = cast<VPBasicBlock>(MiddleVPBB->getSuccessors()[0]);
+      VPExitBlock = cast<VPIRBasicBlock>(MiddleVPBB->getSuccessors()[0]);
     } else {
       VPExitBlock = VPIRBasicBlock::fromBasicBlock(
           !OrigLoop->contains(TrueSucc) ? TrueSucc : FalseSucc);

llvm/lib/Transforms/Vectorize/VPlanTransforms.h

@@ -124,9 +124,16 @@ struct VPlanTransforms {
   /// Remove dead recipes from \p Plan.
   static void removeDeadRecipes(VPlan &Plan);
 
-  static void handleUncountableEarlyExit(VPlan &Plan, ScalarEvolution &SE,
-                                         Loop *OrigLoop,
-                                         VPRecipeBuilder &RecipeBuilder);
+  /// Update \p Plan to account for uncountable exit blocks in \p
+  /// UncountableExitingBlocks by
+  ///  * updating the condition to exit the vector loop to include the early
+  ///  exit conditions
+  ///  * splitting the original middle block to branch to the early exit blocks
+  ///  if taken. Returns false if the transformation wasn't successful.
+  static void
+  handleUncountableEarlyExit(VPlan &Plan, ScalarEvolution &SE, Loop *OrigLoop,
+                             ArrayRef<BasicBlock *> UncountableExitingBlocks,
+                             VPRecipeBuilder &RecipeBuilder);
 };
 
 } // namespace llvm