From 852ad201df95eee7ea8abcdb4e6ebaff4ea2638f Mon Sep 17 00:00:00 2001 From: Florian Hahn Date: Tue, 5 Nov 2024 15:07:50 +0000 Subject: [PATCH 1/4] [VPlan] Generalize collectUsersInExitBlocks for multiple exit bbs. Generalize collectUsersInExitBlock to collecting exit users in multiple exit blocks. Exit blocks are leaf nodes in the VPlan (without successors) except the scalar header. Split off from https://github.com/llvm/llvm-project/pull/112138 --- .../Transforms/Vectorize/LoopVectorize.cpp | 76 ++++++++++--------- 1 file changed, 40 insertions(+), 36 deletions(-) diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp index 35113c204b3c1..944896a75cc33 100644 --- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp +++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp @@ -8819,47 +8819,51 @@ static void addScalarResumePhis(VPRecipeBuilder &Builder, VPlan &Plan) { } } -// Collect VPIRInstructions for phis in the original exit block that are modeled +// Collect VPIRInstructions for phis in the exit blocks that are modeled // in VPlan and add the exiting VPValue as operand. Some exiting values are not // modeled explicitly yet and won't be included. Those are un-truncated // VPWidenIntOrFpInductionRecipe, VPWidenPointerInductionRecipe and induction // increments. -static SetVector collectUsersInExitBlock( +static SetVector collectUsersInExitBlocks( Loop *OrigLoop, VPRecipeBuilder &Builder, VPlan &Plan, const MapVector &Inductions) { - auto *MiddleVPBB = Plan.getMiddleBlock(); - // No edge from the middle block to the unique exit block has been inserted - // and there is nothing to fix from vector loop; phis should have incoming - // from scalar loop only. - if (MiddleVPBB->getNumSuccessors() != 2) - return {}; SetVector ExitUsersToFix; - VPBasicBlock *ExitVPBB = cast(MiddleVPBB->getSuccessors()[0]); - BasicBlock *ExitingBB = OrigLoop->getExitingBlock(); - for (VPRecipeBase &R : *ExitVPBB) { - auto *ExitIRI = dyn_cast(&R); - if (!ExitIRI) - continue; - auto *ExitPhi = dyn_cast(&ExitIRI->getInstruction()); - if (!ExitPhi) - break; - Value *IncomingValue = ExitPhi->getIncomingValueForBlock(ExitingBB); - VPValue *V = Builder.getVPValueOrAddLiveIn(IncomingValue); - // Exit values for inductions are computed and updated outside of VPlan and - // independent of induction recipes. - // TODO: Compute induction exit values in VPlan. - if ((isa(V) && - !cast(V)->getTruncInst()) || - isa(V) || - (isa(IncomingValue) && - OrigLoop->contains(cast(IncomingValue)) && - any_of(IncomingValue->users(), [&Inductions](User *U) { - auto *P = dyn_cast(U); - return P && Inductions.contains(P); - }))) + for (VPBlockBase *VPB : vp_depth_first_shallow( + Plan.getVectorLoopRegion()->getSingleSuccessor())) { + if (VPB->getNumSuccessors() != 0 || VPB == Plan.getScalarHeader()) continue; - ExitUsersToFix.insert(ExitIRI); - ExitIRI->addOperand(V); + auto *ExitVPBB = cast(VPB); + BasicBlock *ExitBB = ExitVPBB->getIRBasicBlock(); + BasicBlock *ExitingBB = find_singleton( + to_vector(predecessors(ExitBB)), + [OrigLoop](BasicBlock *Pred, bool AllowRepeats) { + return OrigLoop->contains(Pred) ? Pred : nullptr; + }); + for (VPRecipeBase &R : *ExitVPBB) { + auto *ExitIRI = dyn_cast(&R); + if (!ExitIRI) + continue; + auto *ExitPhi = dyn_cast(&ExitIRI->getInstruction()); + if (!ExitPhi) + break; + Value *IncomingValue = ExitPhi->getIncomingValueForBlock(ExitingBB); + VPValue *V = Builder.getVPValueOrAddLiveIn(IncomingValue); + // Exit values for inductions are computed and updated outside of VPlan + // and independent of induction recipes. + // TODO: Compute induction exit values in VPlan. + if ((isa(V) && + !cast(V)->getTruncInst()) || + isa(V) || + (isa(IncomingValue) && + OrigLoop->contains(cast(IncomingValue)) && + any_of(IncomingValue->users(), [&Inductions](User *U) { + auto *P = dyn_cast(U); + return P && Inductions.contains(P); + }))) + continue; + ExitUsersToFix.insert(ExitIRI); + ExitIRI->addOperand(V); + } } return ExitUsersToFix; } @@ -8867,7 +8871,7 @@ static SetVector collectUsersInExitBlock( // Add exit values to \p Plan. Extracts are added for each entry in \p // ExitUsersToFix if needed and their operands are updated. static void -addUsersInExitBlock(VPlan &Plan, +addUsersInExitBlocks(VPlan &Plan, const SetVector &ExitUsersToFix) { if (ExitUsersToFix.empty()) return; @@ -9161,10 +9165,10 @@ LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(VFRange &Range) { RecipeBuilder.fixHeaderPhis(); addScalarResumePhis(RecipeBuilder, *Plan); - SetVector ExitUsersToFix = collectUsersInExitBlock( + SetVector ExitUsersToFix = collectUsersInExitBlocks( OrigLoop, RecipeBuilder, *Plan, Legal->getInductionVars()); addExitUsersForFirstOrderRecurrences(*Plan, ExitUsersToFix); - addUsersInExitBlock(*Plan, ExitUsersToFix); + addUsersInExitBlocks(*Plan, ExitUsersToFix); // --------------------------------------------------------------------------- // Transform initial VPlan: Apply previously taken decisions, in order, to // bring the VPlan to its final state. From dda14f7937b4bf4c48c31355ee09c78d9a4976bc Mon Sep 17 00:00:00 2001 From: Florian Hahn Date: Sun, 17 Nov 2024 18:11:14 +0000 Subject: [PATCH 2/4] !fixup clarify type of getExitBlocks in comment. --- llvm/lib/Transforms/Vectorize/VPlan.h | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h index 08341caf41b1b..d3f295ee7c0ad 100644 --- a/llvm/lib/Transforms/Vectorize/VPlan.h +++ b/llvm/lib/Transforms/Vectorize/VPlan.h @@ -3837,8 +3837,10 @@ class VPlan { return cast(getVectorLoopRegion()->getSingleSuccessor()); } - /// Return the exit blocks of the VPlan, that is leaf nodes except the scalar - /// header. + /// Return an iterator range over the VPIRBasicBlock wrapping the exit blocks + /// of the VPlan, that is leaf nodes except the scalar header. Defined in + /// VPlanHCFG, as the definition of the type needs access to the definitions + /// of VPBlockShallowTraversalWrapper. auto getExitBlocks(); /// The trip count of the original loop. From 7c70e37f8d9140285ccc9e6cc7e59d01e3388891 Mon Sep 17 00:00:00 2001 From: Florian Hahn Date: Sun, 17 Nov 2024 19:31:02 +0000 Subject: [PATCH 3/4] !fixup add assertion guarding against unsupported exit values. --- llvm/lib/Transforms/Vectorize/LoopVectorize.cpp | 2 ++ 1 file changed, 2 insertions(+) diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp index 3c9befd24e3c9..f970cf6988add 100644 --- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp +++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp @@ -8892,6 +8892,8 @@ addUsersInExitBlocks(VPlan &Plan, if (V->isLiveIn()) continue; + assert(ExitIRI->getParent()->getSinglePredecessor() == MiddleVPBB && + "Exit value not handled yet for this edge."); LLVMContext &Ctx = ExitIRI->getInstruction().getContext(); VPValue *Ext = B.createNaryOp(VPInstruction::ExtractFromEnd, {V, Plan.getOrAddLiveIn(ConstantInt::get( From db4e3a0d64d8a425c9362b4a7025c21063cbfaaa Mon Sep 17 00:00:00 2001 From: Florian Hahn Date: Sun, 17 Nov 2024 19:38:45 +0000 Subject: [PATCH 4/4] !fixup fix formatting --- llvm/lib/Transforms/Vectorize/LoopVectorize.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp index f970cf6988add..6a518150ac472 100644 --- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp +++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp @@ -8876,7 +8876,7 @@ static SetVector collectUsersInExitBlocks( // ExitUsersToFix if needed and their operands are updated. static void addUsersInExitBlocks(VPlan &Plan, - const SetVector &ExitUsersToFix) { + const SetVector &ExitUsersToFix) { if (ExitUsersToFix.empty()) return;