Address review comments

Author: david-arm

llvm/docs/Vectorizers.rst

@@ -406,8 +406,10 @@ When vectorizing a loop with a single early exit, the loop blocks following the
 early exit are predicated and the vector loop will always exit via the latch.
 If the early exit has been taken, the vector loop's successor block
 (``middle.split`` below) branches to the early exit block via an intermediate
-block (``vector.early.exit`` below). Otherwise ``middle.block`` selects between
-the exit block from the latch or the scalar remainder loop.
+block (``vector.early.exit`` below). This intermediate block is responsible for
+calculating any exit values of loop-defined variables that are used in the
+early exit block. Otherwise, ``middle.block`` selects between the exit block
+from the latch or the scalar remainder loop.
 
 .. image:: vplan-early-exit.png
 

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -7818,9 +7818,6 @@ DenseMap<const SCEV *, Value *> LoopVectorizationPlanner::executePlan(
     State.LVer->prepareNoAliasMetadata();
   }
 
-  // Set the uncountable early exit block in the VPTransformState.
-  State.CFG.UncountableEarlyExitBB = ILV.Legal->getUncountableEarlyExitBlock();
-
   ILV.printDebugTracesAtStart();
 
   //===------------------------------------------------===//
@@ -9278,7 +9275,7 @@ addUsersInExitBlocks(VPlan &Plan,
       if (PredVPBB != MiddleVPBB) {
         assert(ExitIRI->getParent()->getNumPredecessors() <= 2);
 
-        // Cache the early exit mask
+        // Lookup and cache the early exit mask.
         if (!EarlyExitMask) {
           VPBasicBlock *MiddleSplitVPBB =
               cast<VPBasicBlock>(VectorEarlyExitVPBB->getSinglePredecessor());

llvm/lib/Transforms/Vectorize/VPlan.cpp

@@ -501,16 +501,15 @@ void VPBasicBlock::execute(VPTransformState *State) {
     UnreachableInst *Terminator = State->Builder.CreateUnreachable();
     // Register NewBB in its loop. In innermost loops its the same for all
     // BB's.
-    if (this == State->Plan->getEarlyExit()) {
-      // If this is the vector early exit block then it has a single successor,
-      // which is the uncountable early exit block of the original loop. The
-      // parent loop for the exit block may not be the same as the parent loop
-      // of the vectorised loop, so we have to treat this differently.
-      Loop *EEL = State->LI->getLoopFor(State->CFG.UncountableEarlyExitBB);
-      if (EEL)
-        EEL->addBasicBlockToLoop(NewBB, *State->LI);
-    } else if (State->CurrentParentLoop)
-      State->CurrentParentLoop->addBasicBlockToLoop(NewBB, *State->LI);
+    Loop *ParentLoop = State->CurrentParentLoop;
+    // If this block has a sole successor that is an exit block then it needs
+    // adding to the same parent loop as the exit block.
+    VPBlockBase *SuccVPBB = getSingleSuccessor();
+    if (SuccVPBB && State->Plan->isExitBlock(SuccVPBB))
+      ParentLoop = State->LI->getLoopFor(
+          cast<VPIRBasicBlock>(SuccVPBB)->getIRBasicBlock());
+    if (ParentLoop)
+      ParentLoop->addBasicBlockToLoop(NewBB, *State->LI);
     State->Builder.SetInsertPoint(Terminator);
 
     State->CFG.PrevBB = NewBB;

llvm/lib/Transforms/Vectorize/VPlan.h

@@ -347,9 +347,6 @@ struct VPTransformState {
     /// vector loop.
     BasicBlock *ExitBB = nullptr;
 
-    /// The uncountable early exit block in the original scalar loop.
-    BasicBlock *UncountableEarlyExitBB = nullptr;
-
     /// A mapping of each VPBasicBlock to the corresponding BasicBlock. In case
     /// of replication, maps the BasicBlock of the last replica created.
     SmallDenseMap<VPBasicBlock *, BasicBlock *> VPBB2IRBB;
@@ -3976,6 +3973,9 @@ class VPlan {
   /// of VPBlockShallowTraversalWrapper.
   auto getExitBlocks();
 
+  /// Returns true if \p VPBB is an exit block.
+  bool isExitBlock(VPBlockBase *VPBB);
+
   /// 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,24 @@ template <> struct GraphTraits<VPlan *> {
   }
 };
 
+inline bool VPlan::isExitBlock(VPBlockBase *VPBB) {
+  if (!isa<VPIRBasicBlock>(VPBB) || VPBB->getNumSuccessors() ||
+      VPBB == getScalarHeader())
+    return false;
+
+  VPRegionBlock *RegionBlock = getVectorLoopRegion();
+  if (!RegionBlock)
+    return false;
+
+  // The block must be a successor of the region block.
+  for (auto *OtherVPBB :
+       vp_depth_first_shallow(RegionBlock->getSingleSuccessor()))
+    if (OtherVPBB == VPBB)
+      return true;
+
+  return false;
+}
+
 inline auto VPlan::getExitBlocks() {
   VPBlockBase *ScalarHeader = getScalarHeader();
   return make_filter_range(