[LoopVectorize] Enable vectorisation of early exit loops with live-outs

This work feeds part of PR llvm#88385, and adds support for vectorising
loops with uncountable early exits and outside users of loop-defined
variables. When calculating the final value from an uncountable early
exit we need to calculate the vector lane that triggered the exit,
and hence determine the value at the point we exited.

All code for calculating the last value when exiting the loop early
now lives in a new vector.early.exit block, which sits between the
middle.split block and the original exit block. Doing this required
two fixes:

1. The vplan verifier incorrectly assumed that the block containing
a definition always dominates the block of the user. That's not true
if you can arrive at the use block from multiple incoming blocks.
This is possible for early exit loops where both the early exit and
the latch jump to the same block.

I've added a new ExtractFirstActive VPInstruction that extracts the
first active lane of a vector, i.e. the lane of the vector predicate
that triggered the exit.

NOTE: The IR generated for dealing with live-outs from early exit
loops is unoptimised, as opposed to normal loops. This inevitably
leads to poor quality code, but this can be fixed up later.

Author: david-arm

llvm/docs/Vectorizers.rst

@@ -405,9 +405,11 @@ Early Exit Vectorization
 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. Otherwise
-``middle.block`` selects between the exit block from the latch or the scalar
-remainder loop.
+(``middle.split`` below) branches to the early exit block via an intermediate
+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/docs/vplan-early-exit.dot

@@ -19,23 +19,27 @@ compound=true
     "middle.split"
   ]
   N4 -> N5 [ label=""]
-  N4 -> N6 [ label=""]
+  N4 -> N7 [ label=""]
   N5 [label =
-    "early.exit"
+    "vector.early.exit"
   ]
+  N5 -> N6 [ label=""]
   N6 [label =
-    "middle.block"
+    "early.exit"
   ]
-  N6 -> N9 [ label=""]
-  N6 -> N7 [ label=""]
   N7 [label =
-    "scalar.ph"
+    "middle.block"
   ]
+  N7 -> N10 [ label=""]
   N7 -> N8 [ label=""]
   N8 [label =
-    "loop.header"
+    "scalar.ph"
   ]
+  N8 -> N9 [ label=""]
   N9 [label =
+    "loop.header"
+  ]
+  N10 [label =
     "latch.exit"
   ]
 }

llvm/docs/vplan-early-exit.png

@@ -407,6 +407,11 @@ class LoopVectorizationLegality {
 
   /// Returns the destination of an uncountable early exiting block.
   BasicBlock *getUncountableEarlyExitBlock() const {
+    if (!HasUncountableEarlyExit) {
+      assert(getUncountableExitBlocks().empty() &&
+             "Expected no uncountable exiting blocks");
+      return nullptr;
+    }
     assert(getUncountableExitBlocks().size() == 1 &&
            "Expected only a single uncountable exit block");
     return getUncountableExitBlocks()[0];

llvm/include/llvm/Transforms/Vectorize/LoopVectorizationLegality.h

@@ -9067,15 +9067,21 @@ collectUsersInExitBlocks(Loop *OrigLoop, VPRecipeBuilder &Builder,
 // Add exit values to \p Plan. Extracts are added for each entry in \p
 // ExitUsersToFix if needed and their operands are updated. Returns true if all
 // exit users can be handled, otherwise return false.
-static bool
+static void
 addUsersInExitBlocks(VPlan &Plan,
                      const SetVector<VPIRInstruction *> &ExitUsersToFix) {
   if (ExitUsersToFix.empty())
-    return true;
+    return;
 
   auto *MiddleVPBB = Plan.getMiddleBlock();
   VPBuilder B(MiddleVPBB, MiddleVPBB->getFirstNonPhi());
-  VPTypeAnalysis TypeInfo(Plan.getCanonicalIV()->getScalarType());
+  VPBuilder MiddleB(MiddleVPBB, MiddleVPBB->getFirstNonPhi());
+  VPBuilder EarlyExitB;
+  VPBasicBlock *VectorEarlyExitVPBB = Plan.getEarlyExit();
+  VPValue *EarlyExitMask = nullptr;
+  if (VectorEarlyExitVPBB)
+    EarlyExitB.setInsertPoint(VectorEarlyExitVPBB,
+                              VectorEarlyExitVPBB->getFirstNonPhi());
 
   // Introduce extract for exiting values and update the VPIRInstructions
   // modeling the corresponding LCSSA phis.
@@ -9086,19 +9092,38 @@ addUsersInExitBlocks(VPlan &Plan,
       if (Op->isLiveIn())
         continue;
 
-      // Currently only live-ins can be used by exit values from blocks not
-      // exiting via the vector latch through to the middle block.
-      if (ExitIRI->getParent()->getSinglePredecessor() != MiddleVPBB)
-        return false;
-
       LLVMContext &Ctx = ExitIRI->getInstruction().getContext();
-      VPValue *Ext = B.createNaryOp(VPInstruction::ExtractFromEnd,
-                                    {Op, Plan.getOrAddLiveIn(ConstantInt::get(
-                                             IntegerType::get(Ctx, 32), 1))});
+      VPValue *Ext;
+      VPBasicBlock *PredVPBB =
+          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 {
+        Ext = MiddleB.createNaryOp(VPInstruction::ExtractFromEnd,
+                                   {Op, Plan.getOrAddLiveIn(ConstantInt::get(
+                                            IntegerType::get(Ctx, 32), 1))});
+      }
       ExitIRI->setOperand(Idx, Ext);
     }
   }
-  return true;
 }
 
 /// Handle users in the exit block for first order reductions in the original
@@ -9402,12 +9427,7 @@ LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(VFRange &Range) {
   SetVector<VPIRInstruction *> ExitUsersToFix =
       collectUsersInExitBlocks(OrigLoop, RecipeBuilder, *Plan);
   addExitUsersForFirstOrderRecurrences(*Plan, ExitUsersToFix);
-  if (!addUsersInExitBlocks(*Plan, ExitUsersToFix)) {
-    reportVectorizationFailure(
-        "Some exit values in loop with uncountable exit not supported yet",
-        "UncountableEarlyExitLoopsUnsupportedExitValue", ORE, OrigLoop);
-    return nullptr;
-  }
+  addUsersInExitBlocks(*Plan, ExitUsersToFix);
 
   // ---------------------------------------------------------------------------
   // Transform initial VPlan: Apply previously taken decisions, in order, to

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -501,8 +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 (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.cpp

@@ -1228,6 +1228,9 @@ class VPInstruction : public VPRecipeWithIRFlags,
     // Returns a scalar boolean value, which is true if any lane of its (only
     // boolean) vector operand is true.
     AnyOf,
+    // Extracts the first active lane of a vector, where the first operand is
+    // the predicate, and the second operand is the vector to extract.
+    ExtractFirstActive,
   };
 
 private:
@@ -3938,6 +3941,22 @@ class VPlan {
   VPRegionBlock *getVectorLoopRegion();
   const VPRegionBlock *getVectorLoopRegion() const;
 
+  /// Get the vector early exit block
+  VPBasicBlock *getEarlyExit() {
+    auto LoopRegion = getVectorLoopRegion();
+    if (!LoopRegion)
+      return nullptr;
+
+    auto *SuccessorVPBB = LoopRegion->getSingleSuccessor();
+    auto *MiddleVPBB = getMiddleBlock();
+    if (SuccessorVPBB == MiddleVPBB)
+      return nullptr;
+
+    assert(SuccessorVPBB->getSuccessors()[1] == MiddleVPBB &&
+           "Expected second successor to be the middle block");
+    return cast<VPBasicBlock>(SuccessorVPBB->getSuccessors()[0]);
+  }
+
   /// Returns the 'middle' block of the plan, that is the block that selects
   /// whether to execute the scalar tail loop or the exit block from the loop
   /// latch.
@@ -3962,6 +3981,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/VPlan.h

@@ -78,6 +78,7 @@ Type *VPTypeAnalysis::inferScalarTypeForRecipe(const VPInstruction *R) {
   case VPInstruction::CanonicalIVIncrementForPart:
   case VPInstruction::AnyOf:
     return SetResultTyFromOp();
+  case VPInstruction::ExtractFirstActive:
   case VPInstruction::ExtractFromEnd: {
     Type *BaseTy = inferScalarType(R->getOperand(0));
     if (auto *VecTy = dyn_cast<VectorType>(BaseTy))

llvm/lib/Transforms/Vectorize/VPlanAnalysis.cpp

@@ -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(

llvm/lib/Transforms/Vectorize/VPlanCFG.h

@@ -696,14 +696,21 @@ Value *VPInstruction::generate(VPTransformState &State) {
     Value *A = State.get(getOperand(0));
     return Builder.CreateOrReduce(A);
   }
-
+  case VPInstruction::ExtractFirstActive: {
+    Value *Vec = State.get(getOperand(0));
+    Value *Mask = State.get(getOperand(1));
+    Value *Ctz =
+        Builder.CreateCountTrailingZeroElems(Builder.getInt64Ty(), Mask);
+    return Builder.CreateExtractElement(Vec, Ctz);
+  }
   default:
     llvm_unreachable("Unsupported opcode for instruction");
   }
 }
 
 bool VPInstruction::isVectorToScalar() const {
   return getOpcode() == VPInstruction::ExtractFromEnd ||
+         getOpcode() == VPInstruction::ExtractFirstActive ||
          getOpcode() == VPInstruction::ComputeReductionResult ||
          getOpcode() == VPInstruction::AnyOf;
 }
@@ -768,6 +775,7 @@ bool VPInstruction::opcodeMayReadOrWriteFromMemory() const {
   case VPInstruction::CalculateTripCountMinusVF:
   case VPInstruction::CanonicalIVIncrementForPart:
   case VPInstruction::ExtractFromEnd:
+  case VPInstruction::ExtractFirstActive:
   case VPInstruction::FirstOrderRecurrenceSplice:
   case VPInstruction::LogicalAnd:
   case VPInstruction::Not:
@@ -887,6 +895,9 @@ void VPInstruction::print(raw_ostream &O, const Twine &Indent,
   case VPInstruction::AnyOf:
     O << "any-of";
     break;
+  case VPInstruction::ExtractFirstActive:
+    O << "extract-first-active";
+    break;
   default:
     O << Instruction::getOpcodeName(getOpcode());
   }

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

@@ -2099,10 +2099,13 @@ void VPlanTransforms::handleUncountableEarlyExit(
       Builder.createNaryOp(VPInstruction::AnyOf, {EarlyExitTakenCond});
 
   VPBasicBlock *NewMiddle = Plan.createVPBasicBlock("middle.split");
+  VPBasicBlock *EarlyExitVPBB = Plan.createVPBasicBlock("vector.early.exit");
   VPBlockUtils::insertOnEdge(LoopRegion, MiddleVPBB, NewMiddle);
-  VPBlockUtils::connectBlocks(NewMiddle, VPEarlyExitBlock);
+  VPBlockUtils::connectBlocks(NewMiddle, EarlyExitVPBB);
   NewMiddle->swapSuccessors();
 
+  VPBlockUtils::connectBlocks(EarlyExitVPBB, VPEarlyExitBlock);
+
   VPBuilder MiddleBuilder(NewMiddle);
   MiddleBuilder.createNaryOp(VPInstruction::BranchOnCond, {IsEarlyExitTaken});
 

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

@@ -222,7 +222,11 @@ bool VPlanVerifier::verifyVPBasicBlock(const VPBasicBlock *VPBB) {
           continue;
         }
 
-        if (!VPDT.dominates(VPBB, UI->getParent())) {
+        // Now that we support vectorising loops with uncountable early exits
+        // we can end up in situations where VPBB does not dominate the exit
+        // block. Only do the check if the user is not in a VPIRBasicBlock.
+        if (!isa<VPIRBasicBlock>(UI->getParent()) &&
+            !VPDT.dominates(VPBB, UI->getParent())) {
           errs() << "Use before def!\n";
           return false;
         }