[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,9 @@ 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). 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

@@ -2777,6 +2777,23 @@ BasicBlock *InnerLoopVectorizer::createVectorizedLoopSkeleton(
   return LoopVectorPreHeader;
 }
 
+static bool isValueIncomingFromBlock(BasicBlock *ExitingBB, Value *V,
+                                     Instruction *UI) {
+  PHINode *PHI = dyn_cast<PHINode>(UI);
+  assert(PHI && "Expected LCSSA form");
+
+  // If this loop has an uncountable early exit then there could be
+  // different users of OrigPhi with either:
+  //   1. Multiple users, because each exiting block (countable or
+  //      uncountable) jumps to the same exit block, or ..
+  //   2. A single user with an incoming value from a countable or
+  //      uncountable exiting block.
+  // In both cases there is no guarantee this came from a countable exiting
+  // block, i.e. the latch.
+  int Index = PHI->getBasicBlockIndex(ExitingBB);
+  return Index != -1 && PHI->getIncomingValue(Index) == V;
+}
+
 // Fix up external users of the induction variable. At this point, we are
 // in LCSSA form, with all external PHIs that use the IV having one input value,
 // coming from the remainder loop. We need those PHIs to also have a correct
@@ -2799,12 +2816,13 @@ void InnerLoopVectorizer::fixupIVUsers(PHINode *OrigPhi,
 
   // An external user of the last iteration's value should see the value that
   // the remainder loop uses to initialize its own IV.
-  Value *PostInc = OrigPhi->getIncomingValueForBlock(OrigLoop->getLoopLatch());
+  BasicBlock *OrigLoopLatch = OrigLoop->getLoopLatch();
+  Value *PostInc = OrigPhi->getIncomingValueForBlock(OrigLoopLatch);
   for (User *U : PostInc->users()) {
     Instruction *UI = cast<Instruction>(U);
     if (!OrigLoop->contains(UI)) {
-      assert(isa<PHINode>(UI) && "Expected LCSSA form");
-      MissingVals[UI] = EndValue;
+      if (isValueIncomingFromBlock(OrigLoopLatch, PostInc, UI))
+        MissingVals[cast<PHINode>(UI)] = EndValue;
     }
   }
 
@@ -2814,7 +2832,8 @@ void InnerLoopVectorizer::fixupIVUsers(PHINode *OrigPhi,
   for (User *U : OrigPhi->users()) {
     auto *UI = cast<Instruction>(U);
     if (!OrigLoop->contains(UI)) {
-      assert(isa<PHINode>(UI) && "Expected LCSSA form");
+      if (!isValueIncomingFromBlock(OrigLoopLatch, OrigPhi, UI))
+        continue;
       IRBuilder<> B(MiddleBlock->getTerminator());
 
       // Fast-math-flags propagate from the original induction instruction.
@@ -2844,18 +2863,6 @@ void InnerLoopVectorizer::fixupIVUsers(PHINode *OrigPhi,
     }
   }
 
-  assert((MissingVals.empty() ||
-          all_of(MissingVals,
-                 [MiddleBlock, this](const std::pair<Value *, Value *> &P) {
-                   return all_of(
-                       predecessors(cast<Instruction>(P.first)->getParent()),
-                       [MiddleBlock, this](BasicBlock *Pred) {
-                         return Pred == MiddleBlock ||
-                                Pred == OrigLoop->getLoopLatch();
-                       });
-                 })) &&
-         "Expected escaping values from latch/middle.block only");
-
   for (auto &I : MissingVals) {
     PHINode *PHI = cast<PHINode>(I.first);
     // One corner case we have to handle is two IVs "chasing" each-other,
@@ -7811,6 +7818,9 @@ 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();
 
   //===------------------------------------------------===//
@@ -9237,14 +9247,20 @@ 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());
+  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.
@@ -9255,19 +9271,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);
+
+        // 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
@@ -9570,12 +9605,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,7 +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)
+    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);
     State->Builder.SetInsertPoint(Terminator);
 

llvm/lib/Transforms/Vectorize/VPlan.cpp

@@ -347,6 +347,9 @@ 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;
@@ -1226,6 +1229,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:
@@ -3930,6 +3936,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.

llvm/lib/Transforms/Vectorize/VPlan.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

@@ -1973,10 +1973,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;
         }