[llvm] [LV] Vectorize early exit loops with multiple exits. (PR #174864)
Florian Hahn via llvm-commits
llvm-commits at lists.llvm.org
Mon Feb 9 06:43:07 PST 2026
================
@@ -4029,75 +4029,214 @@ void VPlanTransforms::convertToConcreteRecipes(VPlan &Plan) {
R->eraseFromParent();
}
-void VPlanTransforms::handleUncountableEarlyExit(VPBasicBlock *EarlyExitingVPBB,
- VPBasicBlock *EarlyExitVPBB,
- VPlan &Plan,
- VPBasicBlock *HeaderVPBB,
- VPBasicBlock *LatchVPBB) {
- auto *MiddleVPBB = cast<VPBasicBlock>(LatchVPBB->getSuccessors()[0]);
- if (!EarlyExitVPBB->getSinglePredecessor() &&
- EarlyExitVPBB->getPredecessors()[1] == MiddleVPBB) {
- assert(EarlyExitVPBB->getNumPredecessors() == 2 &&
- EarlyExitVPBB->getPredecessors()[0] == EarlyExitingVPBB &&
- "unsupported early exit VPBB");
- // Early exit operand should always be last phi operand. If EarlyExitVPBB
- // has two predecessors and EarlyExitingVPBB is the first, swap the operands
- // of the phis.
- for (VPRecipeBase &R : EarlyExitVPBB->phis())
- cast<VPIRPhi>(&R)->swapOperands();
- }
+void VPlanTransforms::handleUncountableEarlyExits(VPlan &Plan,
+ VPBasicBlock *HeaderVPBB,
+ VPBasicBlock *LatchVPBB,
+ VPBasicBlock *MiddleVPBB) {
+ struct EarlyExitInfo {
+ VPBasicBlock *EarlyExitingVPBB;
+ VPIRBasicBlock *EarlyExitVPBB;
+ VPValue *CondToExit;
+ };
VPBuilder Builder(LatchVPBB->getTerminator());
- VPBlockBase *TrueSucc = EarlyExitingVPBB->getSuccessors()[0];
- assert(match(EarlyExitingVPBB->getTerminator(), m_BranchOnCond()) &&
- "Terminator must be be BranchOnCond");
- VPValue *CondOfEarlyExitingVPBB =
- EarlyExitingVPBB->getTerminator()->getOperand(0);
- auto *CondToEarlyExit = TrueSucc == EarlyExitVPBB
- ? CondOfEarlyExitingVPBB
- : Builder.createNot(CondOfEarlyExitingVPBB);
-
- // Create a BranchOnTwoConds in the latch that branches to:
- // [0] vector.early.exit, [1] middle block, [2] header (continue looping).
- VPValue *IsEarlyExitTaken =
- Builder.createNaryOp(VPInstruction::AnyOf, {CondToEarlyExit});
- VPBasicBlock *VectorEarlyExitVPBB =
- Plan.createVPBasicBlock("vector.early.exit");
- VectorEarlyExitVPBB->setParent(EarlyExitVPBB->getParent());
-
- VPBlockUtils::connectBlocks(VectorEarlyExitVPBB, EarlyExitVPBB);
-
- // Update the exit phis in the early exit block.
+ SmallVector<EarlyExitInfo> Exits;
+ for (VPIRBasicBlock *ExitBlock : Plan.getExitBlocks()) {
+ for (VPBlockBase *Pred : to_vector(ExitBlock->getPredecessors())) {
+ if (Pred == MiddleVPBB)
+ continue;
+ // Collect condition for this early exit.
+ auto *EarlyExitingVPBB = cast<VPBasicBlock>(Pred);
+ VPBlockBase *TrueSucc = EarlyExitingVPBB->getSuccessors()[0];
+ VPValue *CondOfEarlyExitingVPBB;
+ [[maybe_unused]] bool Matched =
+ match(EarlyExitingVPBB->getTerminator(),
+ m_BranchOnCond(m_VPValue(CondOfEarlyExitingVPBB)));
+ assert(Matched && "Terminator must be BranchOnCond");
+ auto *CondToEarlyExit = TrueSucc == ExitBlock
+ ? CondOfEarlyExitingVPBB
+ : Builder.createNot(CondOfEarlyExitingVPBB);
+ Exits.push_back({
+ EarlyExitingVPBB,
+ ExitBlock,
+ CondToEarlyExit,
+ });
+ }
+ }
+
+ // Sort exits by dominance to get the correct program order.
+ VPDominatorTree VPDT(Plan);
+ llvm::sort(Exits, [&VPDT](const EarlyExitInfo &A, const EarlyExitInfo &B) {
+ return VPDT.dominates(A.EarlyExitingVPBB, B.EarlyExitingVPBB);
+ });
+
+ // Helper to check if a VPValue's definition dominates the latch.
+ // Live-in values (with no defining recipe) dominate everything.
+ auto DominatesLatch = [&VPDT, LatchVPBB](VPValue *V) {
+ VPRecipeBase *DefRecipe = V->getDefiningRecipe();
+ if (!DefRecipe)
+ return true;
+ return VPDT.dominates(DefRecipe->getParent(), LatchVPBB);
+ };
+
+ // Build the AnyOf condition for the latch terminator.
+ VPValue *Combined = Exits[0].CondToExit;
+ assert(DominatesLatch(Combined) && "All conditions must dominate the latch");
+ for (const auto &[_, _1, CondToExit] : drop_begin(Exits)) {
+ assert(DominatesLatch(CondToExit) &&
+ "All conditions must dominate the latch");
+ Combined = Builder.createOr(Combined, CondToExit);
+ }
+ VPValue *IsAnyExitTaken =
+ Builder.createNaryOp(VPInstruction::AnyOf, {Combined});
+
+ // Create the vector.early.exit blocks.
+ SmallVector<VPBasicBlock *> VectorEarlyExitVPBBs(Exits.size());
+ for (unsigned Idx = 0; Idx != Exits.size(); ++Idx) {
+ Twine BlockSuffix = Exits.size() == 1 ? "" : Twine(".") + Twine(Idx);
+ VPBasicBlock *VectorEarlyExitVPBB =
+ Plan.createVPBasicBlock("vector.early.exit" + BlockSuffix);
+ VectorEarlyExitVPBBs[Idx] = VectorEarlyExitVPBB;
+ }
+
+ // Create the dispatch block (or reuse the single exit block if only one
+ // exit). The dispatch block computes the first active lane of the combined
+ // condition and, for multiple exits, chains through conditions to determine
+ // which exit to take.
+ VPBasicBlock *DispatchVPBB =
+ Exits.size() == 1 ? VectorEarlyExitVPBBs[0]
+ : Plan.createVPBasicBlock("vector.early.exit.check");
+ VPBuilder DispatchBuilder(DispatchVPBB, DispatchVPBB->begin());
+ VPValue *FirstActiveLane =
+ DispatchBuilder.createNaryOp(VPInstruction::FirstActiveLane, {Combined},
+ DebugLoc::getUnknown(), "first.active.lane");
+
+ // For each early exit, disconnect the original exiting block
+ // (early.exiting.I) from the exit block (ir-bb<exit.I>) and route through a
+ // new vector.early.exit block. Update ir-bb<exit.I>'s phis to extract their
+ // values at the first active lane:
+ //
+ // Input:
+ // early.exiting.I:
+ // ...
+ // EMIT branch-on-cond vp<%cond.I>
+ // Successor(s): in.loop.succ, ir-bb<exit.I>
+ //
+ // ir-bb<exit.I>:
+ // IR %phi = phi [ vp<%incoming.I>, early.exiting.I ], ...
+ //
+ // Output:
+ // early.exiting.I:
+ // ...
+ // Successor(s): in.loop.succ
+ //
+ // vector.early.exit.I:
+ // EMIT vp<%exit.val> = extract-lane vp<%first.lane>, vp<%incoming.I>
+ // Successor(s): ir-bb<exit.I>
+ //
+ // ir-bb<exit.I>:
+ // IR %phi = phi ... (extra operand: vp<%exit.val> from
+ // vector.early.exit.I)
+ //
+ for (auto [Exit, VectorEarlyExitVPBB] : zip(Exits, VectorEarlyExitVPBBs)) {
+ auto &[EarlyExitingVPBB, EarlyExitVPBB, CondToExit] = Exit;
+ // Adjust the phi nodes in EarlyExitVPBB.
+ // 1. remove incoming values from EarlyExitingVPBB,
+ // 2. extract the incoming value at FirstActiveLane
+ // 3. add back the extracts as last operands for the phis
+ // Then adjust the CFG, removing the edge between EarlyExitingVPBB and
+ // EarlyExitVPBB and adding a new edge between VectorEarlyExitVPBB and
+ // EarlyExitVPBB. The extracts at FirstActiveLane are now the incoming
+ // values from VectorEarlyExitVPBB.
+ for (VPRecipeBase &R : EarlyExitVPBB->phis()) {
+ auto *ExitIRI = cast<VPIRPhi>(&R);
+ VPValue *IncomingVal =
+ ExitIRI->getIncomingValueForBlock(EarlyExitingVPBB);
+ VPValue *NewIncoming = IncomingVal;
+ if (!isa<VPIRValue>(IncomingVal)) {
+ VPBuilder EarlyExitBuilder(VectorEarlyExitVPBB);
+ NewIncoming = EarlyExitBuilder.createNaryOp(
+ VPInstruction::ExtractLane, {FirstActiveLane, IncomingVal},
+ DebugLoc::getUnknown(), "early.exit.value");
+ }
+ ExitIRI->removeIncomingValueFor(EarlyExitingVPBB);
+ ExitIRI->addOperand(NewIncoming);
+ }
+
+ EarlyExitingVPBB->getTerminator()->eraseFromParent();
+ VPBlockUtils::disconnectBlocks(EarlyExitingVPBB, EarlyExitVPBB);
+ VPBlockUtils::connectBlocks(VectorEarlyExitVPBB, EarlyExitVPBB);
+ }
+
+ // For exit blocks that also have the middle block as predecessor (latch
+ // exits to the same block as an early exit), extract the last lane of the
+ // first operand for the middle block's incoming value.
VPBuilder MiddleBuilder(MiddleVPBB);
- VPBuilder EarlyExitB(VectorEarlyExitVPBB);
- for (VPRecipeBase &R : EarlyExitVPBB->phis()) {
- auto *ExitIRI = cast<VPIRPhi>(&R);
- // Early exit operand should always be last, i.e., 0 if EarlyExitVPBB has
- // a single predecessor and 1 if it has two.
- unsigned EarlyExitIdx = ExitIRI->getNumOperands() - 1;
- if (ExitIRI->getNumOperands() != 1) {
- // The first of two operands corresponds to the latch exit, via MiddleVPBB
- // predecessor. Extract its final lane.
+ VPBasicBlock *MiddleSuccVPBB =
+ cast<VPIRBasicBlock>(MiddleVPBB->getSuccessors()[0]);
+ if (MiddleSuccVPBB->getNumPredecessors() > 1) {
+ assert(all_of(MiddleSuccVPBB->getPredecessors(),
+ [&](VPBlockBase *Pred) {
+ return Pred == MiddleVPBB ||
+ is_contained(VectorEarlyExitVPBBs, Pred);
+ }) &&
+ "All predecessors must be either the middle block or early exit "
+ "blocks");
+
+ for (VPRecipeBase &R : MiddleSuccVPBB->phis()) {
+ auto *ExitIRI = cast<VPIRPhi>(&R);
+ assert(ExitIRI->getIncomingValueForBlock(MiddleVPBB) ==
+ ExitIRI->getOperand(0) &&
+ "First operand must come from middle block");
ExitIRI->extractLastLaneOfLastPartOfFirstOperand(MiddleBuilder);
}
+ }
- VPValue *IncomingFromEarlyExit = ExitIRI->getOperand(EarlyExitIdx);
- if (!isa<VPIRValue>(IncomingFromEarlyExit)) {
- // Update the incoming value from the early exit.
- VPValue *FirstActiveLane = EarlyExitB.createNaryOp(
- VPInstruction::FirstActiveLane, {CondToEarlyExit},
- DebugLoc::getUnknown(), "first.active.lane");
- IncomingFromEarlyExit = EarlyExitB.createNaryOp(
- VPInstruction::ExtractLane, {FirstActiveLane, IncomingFromEarlyExit},
- DebugLoc::getUnknown(), "early.exit.value");
- ExitIRI->setOperand(EarlyExitIdx, IncomingFromEarlyExit);
+ if (Exits.size() != 1) {
+ // Chain through exits: for each exit, check if its condition is true at
+ // the first active lane. If so, take that exit; otherwise, try the next.
+ // The last exit needs no check since it must be taken if all others fail.
+ //
+ // For 3 exits (cond.0, cond.1, cond.2), this creates:
+ //
+ // vector.early.exit.check:
+ // EMIT vp<%combined> = or vp<%cond.0>, vp<%cond.1>, vp<%cond.2>
+ // EMIT vp<%first.lane> = first-active-lane vp<%combined>
+ // EMIT vp<%at.cond.0> = extract-lane vp<%first.lane>, vp<%cond.0>
+ // EMIT branch-on-cond vp<%at.cond.0>
+ // Successor(s): vector.early.exit.0, vector.early.exit.check.0
+ //
+ // vector.early.exit.check.0:
+ // EMIT vp<%at.cond.1> = extract-lane vp<%first.lane>, vp<%cond.1>
+ // EMIT branch-on-cond vp<%at.cond.1>
+ // Successor(s): vector.early.exit.1, vector.early.exit.2
+ VPBasicBlock *CurrentBB = DispatchVPBB;
+ for (auto [I, Exit] : enumerate(ArrayRef(Exits).drop_back())) {
+ VPValue *LaneVal = DispatchBuilder.createNaryOp(
+ VPInstruction::ExtractLane, {FirstActiveLane, Exit.CondToExit},
+ DebugLoc::getUnknown(), "exit.cond.at.lane");
+
+ // For the last dispatch, branch directly to the last exit on false;
+ // otherwise, create a new check block.
+ bool IsLastDispatch = (I + 2 == Exits.size());
+ VPBasicBlock *FalseBB =
+ IsLastDispatch ? VectorEarlyExitVPBBs.back()
+ : Plan.createVPBasicBlock(
+ Twine("vector.early.exit.check.") + Twine(I));
+ if (!IsLastDispatch)
+ FalseBB->setParent(LatchVPBB->getParent());
----------------
fhahn wrote:
No, removed, thanks
https://github.com/llvm/llvm-project/pull/174864
More information about the llvm-commits
mailing list