Address review comments

Author: david-arm

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -9070,15 +9070,27 @@ addUsersInExitBlocks(VPlan &Plan,
     return;
 
   auto *MiddleVPBB = Plan.getMiddleBlock();
-  VPBuilder B(MiddleVPBB, MiddleVPBB->getFirstNonPhi());
   VPBuilder MiddleB(MiddleVPBB, MiddleVPBB->getFirstNonPhi());
   VPBuilder EarlyExitB;
   VPBasicBlock *VectorEarlyExitVPBB = Plan.getEarlyExit();
   VPValue *EarlyExitMask = nullptr;
-  if (VectorEarlyExitVPBB)
+  if (VectorEarlyExitVPBB) {
     EarlyExitB.setInsertPoint(VectorEarlyExitVPBB,
                               VectorEarlyExitVPBB->getFirstNonPhi());
 
+    // Lookup and cache the early exit mask.
+    VPBasicBlock *MiddleSplitVPBB =
+        cast<VPBasicBlock>(VectorEarlyExitVPBB->getSinglePredecessor());
+    VPInstruction *PredTerm =
+        cast<VPInstruction>(MiddleSplitVPBB->getTerminator());
+    assert(PredTerm->getOpcode() == VPInstruction::BranchOnCond &&
+           "Unexpected middle split block terminator");
+    VPInstruction *ScalarCond = cast<VPInstruction>(PredTerm->getOperand(0));
+    assert(ScalarCond->getOpcode() == VPInstruction::AnyOf &&
+           "Unexpected condition for middle split block terminator branch");
+    EarlyExitMask = ScalarCond->getOperand(0);
+  }
+
   // Introduce extract for exiting values and update the VPIRInstructions
   // modeling the corresponding LCSSA phis.
   for (VPIRInstruction *ExitIRI : ExitUsersToFix) {
@@ -9094,22 +9106,6 @@ addUsersInExitBlocks(VPlan &Plan,
           cast<VPBasicBlock>(ExitIRI->getParent()->getPredecessors()[Idx]);
       if (PredVPBB != MiddleVPBB) {
         assert(ExitIRI->getParent()->getNumPredecessors() <= 2);
-
-        // Lookup and cache the early exit mask.
-        if (!EarlyExitMask) {
-          VPBasicBlock *MiddleSplitVPBB =
-              cast<VPBasicBlock>(VectorEarlyExitVPBB->getSinglePredecessor());
-          VPInstruction *PredTerm =
-              cast<VPInstruction>(MiddleSplitVPBB->getTerminator());
-          assert(PredTerm->getOpcode() == VPInstruction::BranchOnCond &&
-                 "Unexpected middle split block terminator");
-          VPInstruction *ScalarCond =
-              cast<VPInstruction>(PredTerm->getOperand(0));
-          assert(
-              ScalarCond->getOpcode() == VPInstruction::AnyOf &&
-              "Unexpected condition for middle split block terminator branch");
-          EarlyExitMask = ScalarCond->getOperand(0);
-        }
         Ext = EarlyExitB.createNaryOp(VPInstruction::ExtractFirstActive,
                                       {Op, EarlyExitMask});
       } else {

llvm/lib/Transforms/Vectorize/VPlan.cpp

@@ -957,6 +957,26 @@ void VPlan::prepareToExecute(Value *TripCountV, Value *VectorTripCountV,
   }
 }
 
+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.
+  if (llvm::is_contained(
+          vp_depth_first_shallow(RegionBlock->getSingleSuccessor()), VPBB)) {
+    assert(llvm::is_contained(getExitBlocks(), VPBB) &&
+           "Expected to find VPlan block in list of exit blocks!");
+    return true;
+  }
+
+  return false;
+}
+
 /// Generate the code inside the preheader and body of the vectorized loop.
 /// Assumes a single pre-header basic-block was created for this. Introduce
 /// additional basic-blocks as needed, and fill them all.

llvm/lib/Transforms/Vectorize/VPlanCFG.h

@@ -306,24 +306,6 @@ 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(