Skip to content

Commit 670d184

Browse files
fhahnfhahn
committed
[VPlan] Handle early exit in preparation for vec.
1 parent cc5b39e commit 670d184

File tree

5 files changed

+75
-59
lines changed

5 files changed

+75
-59
lines changed

‎llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

Lines changed: 5 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -9434,7 +9434,8 @@ LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(VFRange &Range) {
94349434
auto Plan = VPlanTransforms::buildPlainCFG(OrigLoop, *LI, VPB2IRBB);
94359435
VPlanTransforms::prepareForVectorization(
94369436
*Plan, Legal->getWidestInductionType(), PSE, RequiresScalarEpilogueCheck,
9437-
CM.foldTailByMasking(), OrigLoop);
9437+
CM.foldTailByMasking(), OrigLoop, Legal->hasUncountableEarlyExit(),
9438+
Range);
94389439
VPlanTransforms::createLoopRegions(*Plan);
94399440

94409441
// Don't use getDecisionAndClampRange here, because we don't know the UF
@@ -9625,12 +9626,6 @@ LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(VFRange &Range) {
96259626
R->setOperand(1, WideIV->getStepValue());
96269627
}
96279628

9628-
if (auto *UncountableExitingBlock =
9629-
Legal->getUncountableEarlyExitingBlock()) {
9630-
VPlanTransforms::runPass(VPlanTransforms::handleUncountableEarlyExit, *Plan,
9631-
OrigLoop, UncountableExitingBlock, RecipeBuilder,
9632-
Range);
9633-
}
96349629
DenseMap<VPValue *, VPValue *> IVEndValues;
96359630
addScalarResumePhis(RecipeBuilder, *Plan, IVEndValues);
96369631
SetVector<VPIRInstruction *> ExitUsersToFix =
@@ -9726,8 +9721,9 @@ VPlanPtr LoopVectorizationPlanner::tryToBuildVPlan(VFRange &Range) {
97269721

97279722
DenseMap<VPBlockBase *, BasicBlock *> VPB2IRBB;
97289723
auto Plan = VPlanTransforms::buildPlainCFG(OrigLoop, *LI, VPB2IRBB);
9729-
VPlanTransforms::prepareForVectorization(
9730-
*Plan, Legal->getWidestInductionType(), PSE, true, false, OrigLoop);
9724+
VPlanTransforms::prepareForVectorization(*Plan,
9725+
Legal->getWidestInductionType(), PSE,
9726+
true, false, OrigLoop, false, Range);
97319727
VPlanTransforms::createLoopRegions(*Plan);
97329728

97339729
for (ElementCount VF : Range)

‎llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -478,11 +478,12 @@ static void addCanonicalIVRecipes(VPlan &Plan, VPBasicBlock *HeaderVPBB, VPBasic
478478
{CanonicalIVIncrement, &Plan.getVectorTripCount()}, DL);
479479
}
480480

481-
482481
void VPlanTransforms::prepareForVectorization(VPlan &Plan, Type *InductionTy,
483482
PredicatedScalarEvolution &PSE,
484483
bool RequiresScalarEpilogueCheck,
485-
bool TailFolded, Loop *TheLoop) {
484+
bool TailFolded, Loop *TheLoop,
485+
bool HandleUncountableExit,
486+
VFRange &Range) {
486487
VPDominatorTree VPDT;
487488
VPDT.recalculate(Plan);
488489

@@ -506,6 +507,10 @@ void VPlanTransforms::prepareForVectorization(VPlan &Plan, Type *InductionTy,
506507
}
507508

508509
addCanonicalIVRecipes(Plan, cast<VPBasicBlock>(HeaderVPB), cast<VPBasicBlock>(LatchVPB), InductionTy, {});
510+
if (HandleUncountableExit) {
511+
handleUncountableEarlyExit(Plan, cast<VPBasicBlock>(HeaderVPB),
512+
cast<VPBasicBlock>(LatchVPB), Range);
513+
} else {
509514
SmallPtrSet<VPBlockBase *, 2> ExitBlocks(Plan.getExitBlocks().begin(),
510515
Plan.getExitBlocks().end());
511516
for (VPBlockBase *EB : ExitBlocks) {
@@ -516,6 +521,7 @@ addCanonicalIVRecipes(Plan, cast<VPBasicBlock>(HeaderVPB), cast<VPBasicBlock>(La
516521
VPBlockUtils::disconnectBlocks(Pred, EB);
517522
}
518523
}
524+
}
519525

520526
// Create SCEV and VPValue for the trip count.
521527
// We use the symbolic max backedge-taken-count, which works also when

‎llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

Lines changed: 54 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -2487,64 +2487,75 @@ void VPlanTransforms::convertToConcreteRecipes(VPlan &Plan,
24872487
R->eraseFromParent();
24882488
}
24892489

2490-
void VPlanTransforms::handleUncountableEarlyExit(
2491-
VPlan &Plan, Loop *OrigLoop, BasicBlock *UncountableExitingBlock,
2492-
VPRecipeBuilder &RecipeBuilder, VFRange &Range) {
2493-
VPRegionBlock *LoopRegion = Plan.getVectorLoopRegion();
2494-
auto *LatchVPBB = cast<VPBasicBlock>(LoopRegion->getExiting());
2490+
void VPlanTransforms::handleUncountableEarlyExit(VPlan &Plan,
2491+
VPBasicBlock *HeaderVPBB,
2492+
VPBasicBlock *LatchVPBB,
2493+
VFRange &Range) {
2494+
// First find the uncountable early exiting block by looking at the
2495+
// predecessors of the exit blocks.
2496+
VPBlockBase *MiddleVPBB = LatchVPBB->getSuccessors()[0];
2497+
VPBasicBlock *EarlyExitingVPBB = nullptr;
2498+
VPIRBasicBlock *EarlyExitVPBB = nullptr;
2499+
for (auto *EB : Plan.getExitBlocks()) {
2500+
for (VPBlockBase *Pred : EB->getPredecessors()) {
2501+
if (Pred != MiddleVPBB) {
2502+
EarlyExitingVPBB = cast<VPBasicBlock>(Pred);
2503+
EarlyExitVPBB = EB;
2504+
break;
2505+
}
2506+
}
2507+
}
2508+
24952509
VPBuilder Builder(LatchVPBB->getTerminator());
2496-
auto *MiddleVPBB = Plan.getMiddleBlock();
2497-
VPValue *IsEarlyExitTaken = nullptr;
2498-
2499-
// Process the uncountable exiting block. Update IsEarlyExitTaken, which
2500-
// tracks if the uncountable early exit has been taken. Also split the middle
2501-
// block and have it conditionally branch to the early exit block if
2502-
// EarlyExitTaken.
2503-
auto *EarlyExitingBranch =
2504-
cast<BranchInst>(UncountableExitingBlock->getTerminator());
2505-
BasicBlock *TrueSucc = EarlyExitingBranch->getSuccessor(0);
2506-
BasicBlock *FalseSucc = EarlyExitingBranch->getSuccessor(1);
2507-
BasicBlock *EarlyExitIRBB =
2508-
!OrigLoop->contains(TrueSucc) ? TrueSucc : FalseSucc;
2509-
VPIRBasicBlock *VPEarlyExitBlock = Plan.getExitBlock(EarlyExitIRBB);
2510-
2511-
VPValue *EarlyExitNotTakenCond = RecipeBuilder.getBlockInMask(
2512-
OrigLoop->contains(TrueSucc) ? TrueSucc : FalseSucc);
2513-
auto *EarlyExitTakenCond = Builder.createNot(EarlyExitNotTakenCond);
2514-
IsEarlyExitTaken =
2515-
Builder.createNaryOp(VPInstruction::AnyOf, {EarlyExitTakenCond});
2510+
VPBlockBase *TrueSucc = EarlyExitingVPBB->getSuccessors()[0];
2511+
VPValue *EarlyExitCond = EarlyExitingVPBB->getTerminator()->getOperand(0);
2512+
auto *EarlyExitTakenCond = TrueSucc == EarlyExitVPBB
2513+
? EarlyExitCond
2514+
: Builder.createNot(EarlyExitCond);
2515+
2516+
if (!EarlyExitVPBB->getSinglePredecessor() &&
2517+
EarlyExitVPBB->getPredecessors()[0] != MiddleVPBB) {
2518+
for (VPRecipeBase &R : EarlyExitVPBB->phis()) {
2519+
// Early exit operand should always be last, i.e., 0 if EarlyExitVPBB has
2520+
// a single predecessor and 1 if it has two.
2521+
// If EarlyExitVPBB has two predecessors, they are already ordered such
2522+
// that early exit is second (and latch exit is first), by construction.
2523+
// But its underlying IRBB (EarlyExitIRBB) may have its predecessors
2524+
// ordered the other way around, and it is the order of the latter which
2525+
// corresponds to the order of operands of EarlyExitVPBB's phi recipes.
2526+
// Therefore, if early exit (UncountableExitingBlock) is the first
2527+
// predecessor of EarlyExitIRBB, we swap the operands of phi recipes,
2528+
// thereby bringing them to match EarlyExitVPBB's predecessor order,
2529+
// with early exit being last (second). Otherwise they already match.
2530+
cast<VPIRPhi>(&R)->swapOperands();
2531+
}
2532+
}
25162533

2534+
EarlyExitingVPBB->getTerminator()->eraseFromParent();
2535+
VPBlockUtils::disconnectBlocks(EarlyExitingVPBB, EarlyExitVPBB);
2536+
2537+
// Split the middle block and have it conditionally branch to the early exit
2538+
// block if EarlyExitTaken.
2539+
VPValue *IsEarlyExitTaken =
2540+
Builder.createNaryOp(VPInstruction::AnyOf, {EarlyExitTakenCond});
25172541
VPBasicBlock *NewMiddle = Plan.createVPBasicBlock("middle.split");
25182542
VPBasicBlock *VectorEarlyExitVPBB =
25192543
Plan.createVPBasicBlock("vector.early.exit");
2520-
VPBlockUtils::insertOnEdge(LoopRegion, MiddleVPBB, NewMiddle);
2544+
VPBlockUtils::insertOnEdge(LatchVPBB, MiddleVPBB, NewMiddle);
25212545
VPBlockUtils::connectBlocks(NewMiddle, VectorEarlyExitVPBB);
25222546
NewMiddle->swapSuccessors();
25232547

2524-
VPBlockUtils::connectBlocks(VectorEarlyExitVPBB, VPEarlyExitBlock);
2548+
VPBlockUtils::connectBlocks(VectorEarlyExitVPBB, EarlyExitVPBB);
25252549

25262550
// Update the exit phis in the early exit block.
25272551
VPBuilder MiddleBuilder(NewMiddle);
25282552
VPBuilder EarlyExitB(VectorEarlyExitVPBB);
2529-
for (VPRecipeBase &R : VPEarlyExitBlock->phis()) {
2553+
for (VPRecipeBase &R : EarlyExitVPBB->phis()) {
25302554
auto *ExitIRI = cast<VPIRPhi>(&R);
2531-
// Early exit operand should always be last, i.e., 0 if VPEarlyExitBlock has
2555+
// Early exit operand should always be last, i.e., 0 if EarlyExitVPBB has
25322556
// a single predecessor and 1 if it has two.
25332557
unsigned EarlyExitIdx = ExitIRI->getNumOperands() - 1;
2534-
if (!VPEarlyExitBlock->getSinglePredecessor()) {
2535-
// If VPEarlyExitBlock has two predecessors, they are already ordered such
2536-
// that early exit is second (and latch exit is first), by construction.
2537-
// But its underlying IRBB (EarlyExitIRBB) may have its predecessors
2538-
// ordered the other way around, and it is the order of the latter which
2539-
// corresponds to the order of operands of VPEarlyExitBlock's phi recipes.
2540-
// Therefore, if early exit (UncountableExitingBlock) is the first
2541-
// predecessor of EarlyExitIRBB, we swap the operands of phi recipes,
2542-
// thereby bringing them to match VPEarlyExitBlock's predecessor order,
2543-
// with early exit being last (second). Otherwise they already match.
2544-
if (*pred_begin(VPEarlyExitBlock->getIRBasicBlock()) ==
2545-
UncountableExitingBlock)
2546-
ExitIRI->swapOperands();
2547-
2558+
if (!EarlyExitVPBB->getSinglePredecessor()) {
25482559
// The first of two operands corresponds to the latch exit, via MiddleVPBB
25492560
// predecessor. Extract its last lane.
25502561
ExitIRI->extractLastLaneOfFirstOperand(MiddleBuilder);

‎llvm/lib/Transforms/Vectorize/VPlanTransforms.h

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,9 @@ struct VPlanTransforms {
6868
static void prepareForVectorization(VPlan &Plan, Type *InductionTy,
6969
PredicatedScalarEvolution &PSE,
7070
bool RequiresScalarEpilogueCheck,
71-
bool TailFolded, Loop *TheLoop);
71+
bool TailFolded, Loop *TheLoop,
72+
bool HandleUncountableExit,
73+
VFRange &Range);
7274

7375
/// Replace loops in \p Plan's flat CFG with VPRegionBlocks, turning \p Plan's
7476
/// flat CFG into a hierarchical CFG.
@@ -178,9 +180,8 @@ struct VPlanTransforms {
178180
/// exit conditions
179181
/// * splitting the original middle block to branch to the early exit block
180182
/// if taken.
181-
static void handleUncountableEarlyExit(VPlan &Plan, Loop *OrigLoop,
182-
BasicBlock *UncountableExitingBlock,
183-
VPRecipeBuilder &RecipeBuilder,
183+
static void handleUncountableEarlyExit(VPlan &Plan, VPBasicBlock *HeaderVPBB,
184+
VPBasicBlock *LatchVPBB,
184185
VFRange &Range);
185186

186187
/// Lower abstract recipes to concrete ones, that can be codegen'd. Use \p

‎llvm/unittests/Transforms/Vectorize/VPlanTestBase.h

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
#define LLVM_UNITTESTS_TRANSFORMS_VECTORIZE_VPLANTESTBASE_H
1414

1515
#include "../lib/Transforms/Vectorize/VPlan.h"
16+
#include "../lib/Transforms/Vectorize/VPlanHelpers.h"
1617
#include "../lib/Transforms/Vectorize/VPlanTransforms.h"
1718
#include "llvm/Analysis/AssumptionCache.h"
1819
#include "llvm/Analysis/BasicAliasAnalysis.h"
@@ -72,8 +73,9 @@ class VPlanTestIRBase : public testing::Test {
7273
PredicatedScalarEvolution PSE(*SE, *L);
7374
DenseMap<VPBlockBase *, BasicBlock *> VPB2IRBB;
7475
auto Plan = VPlanTransforms::buildPlainCFG(L, *LI, VPB2IRBB);
76+
VFRange R(ElementCount::getFixed(1), ElementCount::getFixed(2));
7577
VPlanTransforms::prepareForVectorization(*Plan, IntegerType::get(*Ctx, 64),
76-
PSE, true, false, L);
78+
PSE, true, false, L, false, R);
7779
VPlanTransforms::createLoopRegions(*Plan);
7880
return Plan;
7981
}

0 commit comments

Comments
 (0)