[llvm] [VPlan] Move predication to VPlanTransform (NFC) (WIP). (PR #128420)
Florian Hahn via llvm-commits
llvm-commits at lists.llvm.org
Sun Feb 23 06:04:49 PST 2025
https://github.com/fhahn created https://github.com/llvm/llvm-project/pull/128420
This patch moves the logic to predicate and linearize a VPlan to a dedicated VPlan transform.
The main logic to perform predication is ready to review, although there are few things to note that should be improved, either directly in the PR or in the future:
* Edge and block masks are cached in VPRecipeBuilder, so they can be accessed during recipe construction. A better alternative may be to add mask operands to all VPInstructions that need them and use that during recipe construction
* The mask caching in a map also means that this map needs updating each time a new recipe replaces a VPInstruction; this would also be handled by adding mask operands.
Currently this is still WIP due to early-exit loop handling not working due to the exit conditions not being available in the initial VPlans. This will be fixed with https://github.com/llvm/llvm-project/pull/128419 and follow-ups
All tests except early-exit loops are passing
>From 92e45cdd100c2ba196114e6edfa1c9217b8d6640 Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Sat, 22 Feb 2025 19:15:32 +0000
Subject: [PATCH] [VPlan] Move predication to VPlanTransform (NFC) (WIP).
This patch moves the logic to predicate and linearize a VPlan to a
dedicated VPlan transform.
The main logic to perform predication is ready to review, although
there are few things to note that should be improved, either directly in
the PR or in the future:
* Edge and block masks are cached in VPRecipeBuilder, so they can be
accessed during recipe construction. A better alternative may be to
add mask operands to all VPInstructions that need them and use that
during recipe construction
* The mask caching in a map also means that this map needs updating
each time a new recipe replaces a VPInstruction; this would also be
handled by adding mask operands.
Currently this is still WIP due to early-exit loop handling not working
due to the exit conditions not being available in the initial VPlans.
This will be fixed with https://github.com/llvm/llvm-project/pull/128419
and follow-ups
All tests except early-exit loops are passing
---
llvm/lib/Transforms/Vectorize/CMakeLists.txt | 1 +
.../Transforms/Vectorize/LoopVectorize.cpp | 286 ++----------------
.../Transforms/Vectorize/VPRecipeBuilder.h | 45 ++-
.../Transforms/Vectorize/VPlanHCFGBuilder.cpp | 24 +-
.../Transforms/Vectorize/VPlanHCFGBuilder.h | 12 -
.../Transforms/Vectorize/VPlanPredicator.cpp | 274 +++++++++++++++++
.../Transforms/Vectorize/VPlanTransforms.cpp | 5 +-
.../Transforms/Vectorize/VPlanTransforms.h | 3 +
8 files changed, 339 insertions(+), 311 deletions(-)
create mode 100644 llvm/lib/Transforms/Vectorize/VPlanPredicator.cpp
diff --git a/llvm/lib/Transforms/Vectorize/CMakeLists.txt b/llvm/lib/Transforms/Vectorize/CMakeLists.txt
index 38670ba304e53..74ae61440327c 100644
--- a/llvm/lib/Transforms/Vectorize/CMakeLists.txt
+++ b/llvm/lib/Transforms/Vectorize/CMakeLists.txt
@@ -23,6 +23,7 @@ add_llvm_component_library(LLVMVectorize
VPlan.cpp
VPlanAnalysis.cpp
VPlanHCFGBuilder.cpp
+ VPlanPredicator.cpp
VPlanRecipes.cpp
VPlanSLP.cpp
VPlanTransforms.cpp
diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index ced01df7b0d44..a2e20a701d612 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -8115,185 +8115,6 @@ void EpilogueVectorizerEpilogueLoop::printDebugTracesAtEnd() {
});
}
-void VPRecipeBuilder::createSwitchEdgeMasks(SwitchInst *SI) {
- BasicBlock *Src = SI->getParent();
- assert(!OrigLoop->isLoopExiting(Src) &&
- all_of(successors(Src),
- [this](BasicBlock *Succ) {
- return OrigLoop->getHeader() != Succ;
- }) &&
- "unsupported switch either exiting loop or continuing to header");
- // Create masks where the terminator in Src is a switch. We create mask for
- // all edges at the same time. This is more efficient, as we can create and
- // collect compares for all cases once.
- VPValue *Cond = getVPValueOrAddLiveIn(SI->getCondition());
- BasicBlock *DefaultDst = SI->getDefaultDest();
- MapVector<BasicBlock *, SmallVector<VPValue *>> Dst2Compares;
- for (auto &C : SI->cases()) {
- BasicBlock *Dst = C.getCaseSuccessor();
- assert(!EdgeMaskCache.contains({Src, Dst}) && "Edge masks already created");
- // Cases whose destination is the same as default are redundant and can be
- // ignored - they will get there anyhow.
- if (Dst == DefaultDst)
- continue;
- auto &Compares = Dst2Compares[Dst];
- VPValue *V = getVPValueOrAddLiveIn(C.getCaseValue());
- Compares.push_back(Builder.createICmp(CmpInst::ICMP_EQ, Cond, V));
- }
-
- // We need to handle 2 separate cases below for all entries in Dst2Compares,
- // which excludes destinations matching the default destination.
- VPValue *SrcMask = getBlockInMask(Src);
- VPValue *DefaultMask = nullptr;
- for (const auto &[Dst, Conds] : Dst2Compares) {
- // 1. Dst is not the default destination. Dst is reached if any of the cases
- // with destination == Dst are taken. Join the conditions for each case
- // whose destination == Dst using an OR.
- VPValue *Mask = Conds[0];
- for (VPValue *V : ArrayRef<VPValue *>(Conds).drop_front())
- Mask = Builder.createOr(Mask, V);
- if (SrcMask)
- Mask = Builder.createLogicalAnd(SrcMask, Mask);
- EdgeMaskCache[{Src, Dst}] = Mask;
-
- // 2. Create the mask for the default destination, which is reached if none
- // of the cases with destination != default destination are taken. Join the
- // conditions for each case where the destination is != Dst using an OR and
- // negate it.
- DefaultMask = DefaultMask ? Builder.createOr(DefaultMask, Mask) : Mask;
- }
-
- if (DefaultMask) {
- DefaultMask = Builder.createNot(DefaultMask);
- if (SrcMask)
- DefaultMask = Builder.createLogicalAnd(SrcMask, DefaultMask);
- }
- EdgeMaskCache[{Src, DefaultDst}] = DefaultMask;
-}
-
-VPValue *VPRecipeBuilder::createEdgeMask(BasicBlock *Src, BasicBlock *Dst) {
- assert(is_contained(predecessors(Dst), Src) && "Invalid edge");
-
- // Look for cached value.
- std::pair<BasicBlock *, BasicBlock *> Edge(Src, Dst);
- EdgeMaskCacheTy::iterator ECEntryIt = EdgeMaskCache.find(Edge);
- if (ECEntryIt != EdgeMaskCache.end())
- return ECEntryIt->second;
-
- if (auto *SI = dyn_cast<SwitchInst>(Src->getTerminator())) {
- createSwitchEdgeMasks(SI);
- assert(EdgeMaskCache.contains(Edge) && "Mask for Edge not created?");
- return EdgeMaskCache[Edge];
- }
-
- VPValue *SrcMask = getBlockInMask(Src);
-
- // The terminator has to be a branch inst!
- BranchInst *BI = dyn_cast<BranchInst>(Src->getTerminator());
- assert(BI && "Unexpected terminator found");
- if (!BI->isConditional() || BI->getSuccessor(0) == BI->getSuccessor(1))
- return EdgeMaskCache[Edge] = SrcMask;
-
- // If source is an exiting block, we know the exit edge is dynamically dead
- // in the vector loop, and thus we don't need to restrict the mask. Avoid
- // adding uses of an otherwise potentially dead instruction unless we are
- // vectorizing a loop with uncountable exits. In that case, we always
- // materialize the mask.
- if (OrigLoop->isLoopExiting(Src) &&
- Src != Legal->getUncountableEarlyExitingBlock())
- return EdgeMaskCache[Edge] = SrcMask;
-
- VPValue *EdgeMask = getVPValueOrAddLiveIn(BI->getCondition());
- assert(EdgeMask && "No Edge Mask found for condition");
-
- if (BI->getSuccessor(0) != Dst)
- EdgeMask = Builder.createNot(EdgeMask, BI->getDebugLoc());
-
- if (SrcMask) { // Otherwise block in-mask is all-one, no need to AND.
- // The bitwise 'And' of SrcMask and EdgeMask introduces new UB if SrcMask
- // is false and EdgeMask is poison. Avoid that by using 'LogicalAnd'
- // instead which generates 'select i1 SrcMask, i1 EdgeMask, i1 false'.
- EdgeMask = Builder.createLogicalAnd(SrcMask, EdgeMask, BI->getDebugLoc());
- }
-
- return EdgeMaskCache[Edge] = EdgeMask;
-}
-
-VPValue *VPRecipeBuilder::getEdgeMask(BasicBlock *Src, BasicBlock *Dst) const {
- assert(is_contained(predecessors(Dst), Src) && "Invalid edge");
-
- // Look for cached value.
- std::pair<BasicBlock *, BasicBlock *> Edge(Src, Dst);
- EdgeMaskCacheTy::const_iterator ECEntryIt = EdgeMaskCache.find(Edge);
- assert(ECEntryIt != EdgeMaskCache.end() &&
- "looking up mask for edge which has not been created");
- return ECEntryIt->second;
-}
-
-void VPRecipeBuilder::createHeaderMask() {
- BasicBlock *Header = OrigLoop->getHeader();
-
- // When not folding the tail, use nullptr to model all-true mask.
- if (!CM.foldTailByMasking()) {
- BlockMaskCache[Header] = nullptr;
- return;
- }
-
- // Introduce the early-exit compare IV <= BTC to form header block mask.
- // This is used instead of IV < TC because TC may wrap, unlike BTC. Start by
- // constructing the desired canonical IV in the header block as its first
- // non-phi instructions.
-
- VPBasicBlock *HeaderVPBB = Plan.getVectorLoopRegion()->getEntryBasicBlock();
- auto NewInsertionPoint = HeaderVPBB->getFirstNonPhi();
- auto *IV = new VPWidenCanonicalIVRecipe(Plan.getCanonicalIV());
- HeaderVPBB->insert(IV, NewInsertionPoint);
-
- VPBuilder::InsertPointGuard Guard(Builder);
- Builder.setInsertPoint(HeaderVPBB, NewInsertionPoint);
- VPValue *BlockMask = nullptr;
- VPValue *BTC = Plan.getOrCreateBackedgeTakenCount();
- BlockMask = Builder.createICmp(CmpInst::ICMP_ULE, IV, BTC);
- BlockMaskCache[Header] = BlockMask;
-}
-
-VPValue *VPRecipeBuilder::getBlockInMask(BasicBlock *BB) const {
- // Return the cached value.
- BlockMaskCacheTy::const_iterator BCEntryIt = BlockMaskCache.find(BB);
- assert(BCEntryIt != BlockMaskCache.end() &&
- "Trying to access mask for block without one.");
- return BCEntryIt->second;
-}
-
-void VPRecipeBuilder::createBlockInMask(BasicBlock *BB) {
- assert(OrigLoop->contains(BB) && "Block is not a part of a loop");
- assert(BlockMaskCache.count(BB) == 0 && "Mask for block already computed");
- assert(OrigLoop->getHeader() != BB &&
- "Loop header must have cached block mask");
-
- // All-one mask is modelled as no-mask following the convention for masked
- // load/store/gather/scatter. Initialize BlockMask to no-mask.
- VPValue *BlockMask = nullptr;
- // This is the block mask. We OR all unique incoming edges.
- for (auto *Predecessor :
- SetVector<BasicBlock *>(pred_begin(BB), pred_end(BB))) {
- VPValue *EdgeMask = createEdgeMask(Predecessor, BB);
- if (!EdgeMask) { // Mask of predecessor is all-one so mask of block is too.
- BlockMaskCache[BB] = EdgeMask;
- return;
- }
-
- if (!BlockMask) { // BlockMask has its initialized nullptr value.
- BlockMask = EdgeMask;
- continue;
- }
-
- BlockMask = Builder.createOr(BlockMask, EdgeMask, {});
- }
-
- BlockMaskCache[BB] = BlockMask;
-}
-
VPWidenMemoryRecipe *
VPRecipeBuilder::tryToWidenMemory(Instruction *I, ArrayRef<VPValue *> Operands,
VFRange &Range) {
@@ -8318,7 +8139,7 @@ VPRecipeBuilder::tryToWidenMemory(Instruction *I, ArrayRef<VPValue *> Operands,
VPValue *Mask = nullptr;
if (Legal->isMaskRequired(I))
- Mask = getBlockInMask(I->getParent());
+ Mask = getBlockInMask(Builder.getInsertBlock());
// Determine if the pointer operand of the access is either consecutive or
// reverse consecutive.
@@ -8437,38 +8258,6 @@ VPWidenIntOrFpInductionRecipe *VPRecipeBuilder::tryToOptimizeInductionTruncate(
return nullptr;
}
-VPBlendRecipe *VPRecipeBuilder::tryToBlend(PHINode *Phi,
- ArrayRef<VPValue *> Operands) {
- unsigned NumIncoming = Phi->getNumIncomingValues();
-
- // We know that all PHIs in non-header blocks are converted into selects, so
- // we don't have to worry about the insertion order and we can just use the
- // builder. At this point we generate the predication tree. There may be
- // duplications since this is a simple recursive scan, but future
- // optimizations will clean it up.
-
- // Map incoming IR BasicBlocks to incoming VPValues, for lookup below.
- // TODO: Add operands and masks in order from the VPlan predecessors.
- DenseMap<BasicBlock *, VPValue *> VPIncomingValues;
- for (const auto &[Idx, Pred] : enumerate(predecessors(Phi->getParent())))
- VPIncomingValues[Pred] = Operands[Idx];
-
- SmallVector<VPValue *, 2> OperandsWithMask;
- for (unsigned In = 0; In < NumIncoming; In++) {
- BasicBlock *Pred = Phi->getIncomingBlock(In);
- OperandsWithMask.push_back(VPIncomingValues.lookup(Pred));
- VPValue *EdgeMask = getEdgeMask(Pred, Phi->getParent());
- if (!EdgeMask) {
- assert(In == 0 && "Both null and non-null edge masks found");
- assert(all_equal(Operands) &&
- "Distinct incoming values with one having a full mask");
- break;
- }
- OperandsWithMask.push_back(EdgeMask);
- }
- return new VPBlendRecipe(Phi, OperandsWithMask);
-}
-
VPSingleDefRecipe *VPRecipeBuilder::tryToWidenCall(CallInst *CI,
ArrayRef<VPValue *> Operands,
VFRange &Range) {
@@ -8544,7 +8333,7 @@ VPSingleDefRecipe *VPRecipeBuilder::tryToWidenCall(CallInst *CI,
// all-true mask.
VPValue *Mask = nullptr;
if (Legal->isMaskRequired(CI))
- Mask = getBlockInMask(CI->getParent());
+ Mask = getBlockInMask(Builder.getInsertBlock());
else
Mask = Plan.getOrAddLiveIn(
ConstantInt::getTrue(IntegerType::getInt1Ty(CI->getContext())));
@@ -8586,7 +8375,7 @@ VPWidenRecipe *VPRecipeBuilder::tryToWiden(Instruction *I,
// div/rem operation itself. Otherwise fall through to general handling below.
if (CM.isPredicatedInst(I)) {
SmallVector<VPValue *> Ops(Operands);
- VPValue *Mask = getBlockInMask(I->getParent());
+ VPValue *Mask = getBlockInMask(Builder.getInsertBlock());
VPValue *One =
Plan.getOrAddLiveIn(ConstantInt::get(I->getType(), 1u, false));
auto *SafeRHS = Builder.createSelect(Mask, Ops[1], One, I->getDebugLoc());
@@ -8668,7 +8457,7 @@ VPRecipeBuilder::tryToWidenHistogram(const HistogramInfo *HI,
// In case of predicated execution (due to tail-folding, or conditional
// execution, or both), pass the relevant mask.
if (Legal->isMaskRequired(HI->Store))
- HGramOps.push_back(getBlockInMask(HI->Store->getParent()));
+ HGramOps.push_back(getBlockInMask(Builder.getInsertBlock()));
return new VPHistogramRecipe(Opcode,
make_range(HGramOps.begin(), HGramOps.end()),
@@ -8724,7 +8513,7 @@ VPRecipeBuilder::handleReplication(Instruction *I, ArrayRef<VPValue *> Operands,
// added initially. Masked replicate recipes will later be placed under an
// if-then construct to prevent side-effects. Generate recipes to compute
// the block mask for this region.
- BlockInMask = getBlockInMask(I->getParent());
+ BlockInMask = getBlockInMask(Builder.getInsertBlock());
}
// Note that there is some custom logic to mark some intrinsics as uniform
@@ -8857,9 +8646,8 @@ VPRecipeBase *VPRecipeBuilder::tryToCreateWidenRecipe(
// nodes, calls and memory operations.
VPRecipeBase *Recipe;
if (auto *Phi = dyn_cast<PHINode>(Instr)) {
- if (Phi->getParent() != OrigLoop->getHeader())
- return tryToBlend(Phi, Operands);
-
+ assert(Phi->getParent() == OrigLoop->getHeader() &&
+ "Non-header phis should have been handled during predication");
assert(Operands.size() == 2 && "Must have 2 operands for header phis");
if ((Recipe = tryToOptimizeInductionPHI(Phi, Operands, Range)))
return Recipe;
@@ -8964,7 +8752,7 @@ VPRecipeBuilder::tryToCreatePartialReduction(Instruction *Reduction,
ReductionOpcode == Instruction::Sub) &&
"Expected an ADD or SUB operation for predicated partial "
"reductions (because the neutral element in the mask is zero)!");
- VPValue *Mask = getBlockInMask(Reduction->getParent());
+ VPValue *Mask = getBlockInMask(Builder.getInsertBlock());
VPValue *Zero =
Plan.getOrAddLiveIn(ConstantInt::get(Reduction->getType(), 0));
BinOp = Builder.createSelect(Mask, BinOp, Zero, Reduction->getDebugLoc());
@@ -9332,9 +9120,6 @@ LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(VFRange &Range) {
bool HasNUW = !IVUpdateMayOverflow || Style == TailFoldingStyle::None;
addCanonicalIVRecipes(*Plan, Legal->getWidestInductionType(), HasNUW, DL);
- VPRecipeBuilder RecipeBuilder(*Plan, OrigLoop, TLI, &TTI, Legal, CM, PSE,
- Builder);
-
// ---------------------------------------------------------------------------
// Pre-construction: record ingredients whose recipes we'll need to further
// process after constructing the initial VPlan.
@@ -9375,39 +9160,24 @@ LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(VFRange &Range) {
return Legal->blockNeedsPredication(BB) || NeedsBlends;
});
- RecipeBuilder.collectScaledReductions(Range);
auto *MiddleVPBB = Plan->getMiddleBlock();
+ VPRecipeBuilder RecipeBuilder(*Plan, OrigLoop, TLI, &TTI, Legal, CM, PSE,
+ Builder);
+ if (NeedsMasks) {
+ VPlanTransforms::predicateAndLinearize(*Plan, CM.foldTailByMasking(),
+ RecipeBuilder);
+ }
+ RecipeBuilder.collectScaledReductions(Range);
+
// Scan the body of the loop in a topological order to visit each basic block
// after having visited its predecessor basic blocks.
ReversePostOrderTraversal<VPBlockShallowTraversalWrapper<VPBlockBase *>> RPOT(
HeaderVPBB);
VPBasicBlock::iterator MBIP = MiddleVPBB->getFirstNonPhi();
- VPBlockBase *PrevVPBB = nullptr;
for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly<VPBasicBlock>(RPOT)) {
- // Handle VPBBs down to the latch.
- if (VPBB == LoopRegion->getExiting()) {
- assert(!HCFGBuilder.getIRBBForVPB(VPBB) &&
- "the latch block shouldn't have a corresponding IRBB");
- VPBlockUtils::connectBlocks(PrevVPBB, VPBB);
- break;
- }
-
- // Create mask based on the IR BB corresponding to VPBB.
- // TODO: Predicate directly based on VPlan.
- Builder.setInsertPoint(VPBB, VPBB->begin());
- if (VPBB == HeaderVPBB) {
- Builder.setInsertPoint(VPBB, VPBB->getFirstNonPhi());
- RecipeBuilder.createHeaderMask();
- } else if (NeedsMasks) {
- // FIXME: At the moment, masks need to be placed at the beginning of the
- // block, as blends introduced for phi nodes need to use it. The created
- // blends should be sunk after the mask recipes.
- RecipeBuilder.createBlockInMask(HCFGBuilder.getIRBBForVPB(VPBB));
- }
-
// Convert input VPInstructions to widened recipes.
for (VPRecipeBase &R : make_early_inc_range(*VPBB)) {
auto *SingleDef = cast<VPSingleDefRecipe>(&R);
@@ -9417,7 +9187,8 @@ LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(VFRange &Range) {
// latter are added above for masking.
// FIXME: Migrate code relying on the underlying instruction from VPlan0
// to construct recipes below to not use the underlying instruction.
- if (isa<VPCanonicalIVPHIRecipe, VPWidenCanonicalIVRecipe>(&R) ||
+ if (isa<VPCanonicalIVPHIRecipe, VPWidenCanonicalIVRecipe, VPBlendRecipe>(
+ &R) ||
(isa<VPInstruction>(&R) && !UnderlyingValue))
continue;
@@ -9469,22 +9240,18 @@ LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(VFRange &Range) {
} else {
Builder.insert(Recipe);
}
- if (Recipe->getNumDefinedValues() == 1)
+ if (Recipe->getNumDefinedValues() == 1) {
SingleDef->replaceAllUsesWith(Recipe->getVPSingleValue());
- else
+ for (auto &[_, V] : RecipeBuilder.BlockMaskCache) {
+ if (V == SingleDef)
+ V = Recipe->getVPSingleValue();
+ }
+ } else
assert(Recipe->getNumDefinedValues() == 0 &&
"Unexpected multidef recipe");
R.eraseFromParent();
}
- // Flatten the CFG in the loop. Masks for blocks have already been generated
- // and added to recipes as needed. To do so, first disconnect VPBB from its
- // successors. Then connect VPBB to the previously visited VPBB.
- for (auto *Succ : to_vector(VPBB->getSuccessors()))
- VPBlockUtils::disconnectBlocks(VPBB, Succ);
- if (PrevVPBB)
- VPBlockUtils::connectBlocks(PrevVPBB, VPBB);
- PrevVPBB = VPBB;
}
assert(isa<VPRegionBlock>(Plan->getVectorLoopRegion()) &&
@@ -9783,7 +9550,7 @@ void LoopVectorizationPlanner::adjustRecipesForReductions(
BasicBlock *BB = CurrentLinkI->getParent();
VPValue *CondOp = nullptr;
if (CM.blockNeedsPredicationForAnyReason(BB))
- CondOp = RecipeBuilder.getBlockInMask(BB);
+ CondOp = RecipeBuilder.getBlockInMask(CurrentLink->getParent());
auto *RedRecipe = new VPReductionRecipe(
RdxDesc, CurrentLinkI, PreviousLink, VecOp, CondOp,
@@ -9818,7 +9585,8 @@ void LoopVectorizationPlanner::adjustRecipesForReductions(
// different numbers of lanes. Partial reductions mask the input instead.
if (!PhiR->isInLoop() && CM.foldTailByMasking() &&
!isa<VPPartialReductionRecipe>(OrigExitingVPV->getDefiningRecipe())) {
- VPValue *Cond = RecipeBuilder.getBlockInMask(OrigLoop->getHeader());
+ VPValue *Cond =
+ RecipeBuilder.getBlockInMask(VectorLoopRegion->getEntryBasicBlock());
assert(OrigExitingVPV->getDefiningRecipe()->getParent() != LatchVPBB &&
"reduction recipe must be defined before latch");
Type *PhiTy = PhiR->getOperand(0)->getLiveInIRValue()->getType();
diff --git a/llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h b/llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h
index 334cfbad8bd7c..9900c4117c5f6 100644
--- a/llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h
+++ b/llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h
@@ -73,11 +73,14 @@ class VPRecipeBuilder {
/// if-conversion currently takes place during VPlan-construction, so these
/// caches are only used at that stage.
using EdgeMaskCacheTy =
- DenseMap<std::pair<BasicBlock *, BasicBlock *>, VPValue *>;
- using BlockMaskCacheTy = DenseMap<BasicBlock *, VPValue *>;
+ DenseMap<std::pair<VPBasicBlock *, VPBasicBlock *>, VPValue *>;
+ using BlockMaskCacheTy = DenseMap<VPBasicBlock *, VPValue *>;
EdgeMaskCacheTy EdgeMaskCache;
+
+public:
BlockMaskCacheTy BlockMaskCache;
+private:
// VPlan construction support: Hold a mapping from ingredients to
// their recipe.
DenseMap<Instruction *, VPRecipeBase *> Ingredient2Recipe;
@@ -114,11 +117,6 @@ class VPRecipeBuilder {
tryToOptimizeInductionTruncate(TruncInst *I, ArrayRef<VPValue *> Operands,
VFRange &Range);
- /// Handle non-loop phi nodes. Return a new VPBlendRecipe otherwise. Currently
- /// all such phi nodes are turned into a sequence of select instructions as
- /// the vectorizer currently performs full if-conversion.
- VPBlendRecipe *tryToBlend(PHINode *Phi, ArrayRef<VPValue *> Operands);
-
/// Handle call instructions. If \p CI can be widened for \p Range.Start,
/// return a new VPWidenCallRecipe or VPWidenIntrinsicRecipe. Range.End may be
/// decreased to ensure same decision from \p Range.Start to \p Range.End.
@@ -187,27 +185,20 @@ class VPRecipeBuilder {
Ingredient2Recipe[I] = R;
}
- /// Create the mask for the vector loop header block.
- void createHeaderMask();
-
- /// A helper function that computes the predicate of the block BB, assuming
- /// that the header block of the loop is set to True or the loop mask when
- /// tail folding.
- void createBlockInMask(BasicBlock *BB);
-
+ void setBlockInMask(VPBasicBlock *BB, VPValue *Mask) {
+ assert(!BlockMaskCache.contains(BB) && "Mask already set");
+ BlockMaskCache[BB] = Mask;
+ }
/// Returns the *entry* mask for the block \p BB.
- VPValue *getBlockInMask(BasicBlock *BB) const;
-
- /// Create an edge mask for every destination of cases and/or default.
- void createSwitchEdgeMasks(SwitchInst *SI);
-
- /// A helper function that computes the predicate of the edge between SRC
- /// and DST.
- VPValue *createEdgeMask(BasicBlock *Src, BasicBlock *Dst);
-
- /// A helper that returns the previously computed predicate of the edge
- /// between SRC and DST.
- VPValue *getEdgeMask(BasicBlock *Src, BasicBlock *Dst) const;
+ VPValue *getBlockInMask(VPBasicBlock *BB) const {
+ return BlockMaskCache.lookup(BB);
+ }
+ void setEdgeMask(VPBasicBlock *Src, VPBasicBlock *Dst, VPValue *Mask) {
+ EdgeMaskCache[{Src, Dst}] = Mask;
+ }
+ VPValue *getEdgeMask(VPBasicBlock *Src, VPBasicBlock *Dst) const {
+ return EdgeMaskCache.lookup({Src, Dst});
+ }
/// Return the recipe created for given ingredient.
VPRecipeBase *getRecipe(Instruction *I) {
diff --git a/llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.cpp b/llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.cpp
index dcf1057b991ee..aec89ba8c74d5 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.cpp
@@ -75,7 +75,7 @@ class PlainCFGBuilder {
: TheLoop(Lp), LI(LI), Plan(P) {}
/// Build plain CFG for TheLoop and connects it to Plan's entry.
- void buildPlainCFG(DenseMap<VPBlockBase *, BasicBlock *> &VPB2IRBB);
+ void buildPlainCFG();
};
} // anonymous namespace
@@ -340,10 +340,16 @@ void PlainCFGBuilder::createVPInstructionsForVPBB(VPBasicBlock *VPBB,
for (Value *Op : Inst->operands())
VPOperands.push_back(getOrCreateVPOperand(Op));
- // Build VPInstruction for any arbitrary Instruction without specific
- // representation in VPlan.
- NewR = cast<VPInstruction>(
- VPIRBuilder.createNaryOp(Inst->getOpcode(), VPOperands, Inst));
+ if (auto *ICmp = dyn_cast<ICmpInst>(Inst)) {
+ NewR = cast<VPInstruction>(VPIRBuilder.createICmp(
+ ICmp->getPredicate(), VPOperands[0], VPOperands[1]));
+ NewR->setUnderlyingValue(ICmp);
+ } else {
+ // Build VPInstruction for any arbitrary Instruction without specific
+ // representation in VPlan.
+ NewR = cast<VPInstruction>(
+ VPIRBuilder.createNaryOp(Inst->getOpcode(), VPOperands, Inst));
+ }
}
IRDef2VPValue[Inst] = NewR;
@@ -351,8 +357,7 @@ void PlainCFGBuilder::createVPInstructionsForVPBB(VPBasicBlock *VPBB,
}
// Main interface to build the plain CFG.
-void PlainCFGBuilder::buildPlainCFG(
- DenseMap<VPBlockBase *, BasicBlock *> &VPB2IRBB) {
+void PlainCFGBuilder::buildPlainCFG() {
// 0. Reuse the top-level region, vector-preheader and exit VPBBs from the
// skeleton. These were created directly rather than via getOrCreateVPBB(),
// revisit them now to update BB2VPBB. Note that header/entry and
@@ -481,14 +486,11 @@ void PlainCFGBuilder::buildPlainCFG(
// have a VPlan counterpart. Fix VPlan header phi by adding their
// corresponding VPlan operands.
fixHeaderPhis();
-
- for (const auto &[IRBB, VPB] : BB2VPBB)
- VPB2IRBB[VPB] = IRBB;
}
void VPlanHCFGBuilder::buildPlainCFG() {
PlainCFGBuilder PCFGBuilder(TheLoop, LI, Plan);
- PCFGBuilder.buildPlainCFG(VPB2IRBB);
+ PCFGBuilder.buildPlainCFG();
}
// Public interface to build a H-CFG.
diff --git a/llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.h b/llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.h
index bc853bf7a1395..ad6e2ad90a961 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.h
+++ b/llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.h
@@ -53,10 +53,6 @@ class VPlanHCFGBuilder {
// are introduced.
VPDominatorTree VPDomTree;
- /// Map of create VP blocks to their input IR basic blocks, if they have been
- /// created for a input IR basic block.
- DenseMap<VPBlockBase *, BasicBlock *> VPB2IRBB;
-
/// Build plain CFG for TheLoop and connects it to Plan's entry.
void buildPlainCFG();
@@ -66,14 +62,6 @@ class VPlanHCFGBuilder {
/// Build H-CFG for TheLoop and update Plan accordingly.
void buildHierarchicalCFG();
-
- /// Return the input IR BasicBlock corresponding to \p VPB. Returns nullptr if
- /// there is no such corresponding block.
- /// FIXME: This is a temporary workaround to drive the createBlockInMask.
- /// Remove once mask creation is done on VPlan.
- BasicBlock *getIRBBForVPB(const VPBlockBase *VPB) const {
- return VPB2IRBB.lookup(VPB);
- }
};
} // namespace llvm
diff --git a/llvm/lib/Transforms/Vectorize/VPlanPredicator.cpp b/llvm/lib/Transforms/Vectorize/VPlanPredicator.cpp
new file mode 100644
index 0000000000000..1836df6bc9cbe
--- /dev/null
+++ b/llvm/lib/Transforms/Vectorize/VPlanPredicator.cpp
@@ -0,0 +1,274 @@
+//===-- VPlanPredicator.cpp - VPlan predicator ----------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// This file implements predication for VPlans.
+///
+//===----------------------------------------------------------------------===//
+
+#include "VPRecipeBuilder.h"
+#include "VPlan.h"
+#include "VPlanCFG.h"
+#include "VPlanTransforms.h"
+#include "VPlanUtils.h"
+#include "llvm/ADT/PostOrderIterator.h"
+
+using namespace llvm;
+
+struct VPPredicator {
+ /// When we if-convert we need to create edge masks. We have to cache values
+ /// so that we don't end up with exponential recursion/IR. Note that
+ /// if-conversion currently takes place during VPlan-construction, so these
+ /// caches are only used at that stage.
+ using EdgeMaskCacheTy =
+ DenseMap<std::pair<VPBasicBlock *, VPBasicBlock *>, VPValue *>;
+ using BlockMaskCacheTy = DenseMap<VPBasicBlock *, VPValue *>;
+
+ VPPredicator(VPRecipeBuilder &RecipeBuilder) : RecipeBuilder(RecipeBuilder) {}
+
+ VPRecipeBuilder &RecipeBuilder;
+
+ VPBuilder Builder;
+ VPValue *createEdgeMask(VPBasicBlock *Src, VPBasicBlock *Dst) {
+ assert(is_contained(Dst->getPredecessors(), Src) && "Invalid edge");
+
+ // Look for cached value.
+ VPValue *EdgeMask = RecipeBuilder.getEdgeMask(Src, Dst);
+ if (EdgeMask)
+ return EdgeMask;
+
+ VPValue *SrcMask = RecipeBuilder.getBlockInMask(Src);
+
+ // The terminator has to be a branch inst!
+ if (Src->empty() || Src->getNumSuccessors() == 1) {
+ RecipeBuilder.setEdgeMask(Src, Dst, SrcMask);
+ return SrcMask;
+ }
+
+ auto *Term = cast<VPInstruction>(Src->getTerminator());
+ if (Term->getOpcode() == Instruction::Switch) {
+ createSwitchEdgeMasks(Term);
+ return RecipeBuilder.getEdgeMask(Src, Dst);
+ }
+
+ auto *BI = cast<VPInstruction>(Src->getTerminator());
+ assert(BI->getOpcode() == VPInstruction::BranchOnCond);
+ if (Src->getSuccessors()[0] == Src->getSuccessors()[1]) {
+ RecipeBuilder.setEdgeMask(Src, Dst, SrcMask);
+ return SrcMask;
+ }
+
+ EdgeMask = BI->getOperand(0);
+ assert(EdgeMask && "No Edge Mask found for condition");
+
+ if (Src->getSuccessors()[0] != Dst)
+ EdgeMask = Builder.createNot(EdgeMask, BI->getDebugLoc());
+
+ if (SrcMask) { // Otherwise block in-mask is all-one, no need to AND.
+ // The bitwise 'And' of SrcMask and EdgeMask introduces new UB if SrcMask
+ // is false and EdgeMask is poison. Avoid that by using 'LogicalAnd'
+ // instead which generates 'select i1 SrcMask, i1 EdgeMask, i1 false'.
+ EdgeMask = Builder.createLogicalAnd(SrcMask, EdgeMask, BI->getDebugLoc());
+ }
+
+ RecipeBuilder.setEdgeMask(Src, Dst, EdgeMask);
+ return EdgeMask;
+ }
+
+ VPValue *createBlockInMask(VPBasicBlock *VPBB) {
+ Builder.setInsertPoint(VPBB, VPBB->begin());
+ // All-one mask is modelled as no-mask following the convention for masked
+ // load/store/gather/scatter. Initialize BlockMask to no-mask.
+ VPValue *BlockMask = nullptr;
+ // This is the block mask. We OR all unique incoming edges.
+ for (auto *Predecessor : SetVector<VPBlockBase *>(
+ VPBB->getPredecessors().begin(), VPBB->getPredecessors().end())) {
+ VPValue *EdgeMask = createEdgeMask(cast<VPBasicBlock>(Predecessor), VPBB);
+ if (!EdgeMask) { // Mask of predecessor is all-one so mask of block is
+ // too.
+ RecipeBuilder.setBlockInMask(VPBB, EdgeMask);
+ return EdgeMask;
+ }
+
+ if (!BlockMask) { // BlockMask has its initialized nullptr value.
+ BlockMask = EdgeMask;
+ continue;
+ }
+
+ BlockMask = Builder.createOr(BlockMask, EdgeMask, {});
+ }
+
+ RecipeBuilder.setBlockInMask(VPBB, BlockMask);
+ return BlockMask;
+ }
+
+ void createHeaderMask(VPBasicBlock *HeaderVPBB, bool FoldTail) {
+ if (!FoldTail) {
+ RecipeBuilder.setBlockInMask(HeaderVPBB, nullptr);
+ return;
+ }
+
+ // Introduce the early-exit compare IV <= BTC to form header block mask.
+ // This is used instead of IV < TC because TC may wrap, unlike BTC. Start by
+ // constructing the desired canonical IV in the header block as its first
+ // non-phi instructions.
+
+ auto NewInsertionPoint = HeaderVPBB->getFirstNonPhi();
+ auto &Plan = *HeaderVPBB->getPlan();
+ auto *IV = new VPWidenCanonicalIVRecipe(Plan.getCanonicalIV());
+ HeaderVPBB->insert(IV, NewInsertionPoint);
+
+ VPBuilder::InsertPointGuard Guard(Builder);
+ Builder.setInsertPoint(HeaderVPBB, NewInsertionPoint);
+ VPValue *BlockMask = nullptr;
+ VPValue *BTC = Plan.getOrCreateBackedgeTakenCount();
+ BlockMask = Builder.createICmp(CmpInst::ICMP_ULE, IV, BTC);
+ RecipeBuilder.setBlockInMask(HeaderVPBB, BlockMask);
+ }
+
+ void createSwitchEdgeMasks(VPInstruction *SI) {
+ VPBasicBlock *Src = SI->getParent();
+
+ // Create masks where the terminator in Src is a switch. We create mask for
+ // all edges at the same time. This is more efficient, as we can create and
+ // collect compares for all cases once.
+ VPValue *Cond = SI->getOperand(0);
+ VPBasicBlock *DefaultDst = cast<VPBasicBlock>(Src->getSuccessors()[0]);
+ MapVector<VPBasicBlock *, SmallVector<VPValue *>> Dst2Compares;
+ for (const auto &[Idx, Succ] :
+ enumerate(ArrayRef(Src->getSuccessors()).drop_front())) {
+ VPBasicBlock *Dst = cast<VPBasicBlock>(Succ);
+ // assert(!EdgeMaskCache.contains({Src, Dst}) && "Edge masks already
+ // created");
+ // Cases whose destination is the same as default are redundant and can
+ // be ignored - they will get there anyhow.
+ if (Dst == DefaultDst)
+ continue;
+ auto &Compares = Dst2Compares[Dst];
+ VPValue *V = SI->getOperand(Idx + 1);
+ Compares.push_back(Builder.createICmp(CmpInst::ICMP_EQ, Cond, V));
+ }
+
+ // We need to handle 2 separate cases below for all entries in Dst2Compares,
+ // which excludes destinations matching the default destination.
+ VPValue *SrcMask = RecipeBuilder.getBlockInMask(Src);
+ VPValue *DefaultMask = nullptr;
+ for (const auto &[Dst, Conds] : Dst2Compares) {
+ // 1. Dst is not the default destination. Dst is reached if any of the
+ // cases with destination == Dst are taken. Join the conditions for each
+ // case whose destination == Dst using an OR.
+ VPValue *Mask = Conds[0];
+ for (VPValue *V : ArrayRef<VPValue *>(Conds).drop_front())
+ Mask = Builder.createOr(Mask, V);
+ if (SrcMask)
+ Mask = Builder.createLogicalAnd(SrcMask, Mask);
+ RecipeBuilder.setEdgeMask(Src, Dst, Mask);
+
+ // 2. Create the mask for the default destination, which is reached if
+ // none of the cases with destination != default destination are taken.
+ // Join the conditions for each case where the destination is != Dst using
+ // an OR and negate it.
+ DefaultMask = DefaultMask ? Builder.createOr(DefaultMask, Mask) : Mask;
+ }
+
+ if (DefaultMask) {
+ DefaultMask = Builder.createNot(DefaultMask);
+ if (SrcMask)
+ DefaultMask = Builder.createLogicalAnd(SrcMask, DefaultMask);
+ }
+ RecipeBuilder.setEdgeMask(Src, DefaultDst, DefaultMask);
+ }
+};
+
+void VPlanTransforms::predicateAndLinearize(VPlan &Plan, bool FoldTail,
+ VPRecipeBuilder &RecipeBuilder) {
+ VPBlockBase *PrevVPBB = nullptr;
+ VPRegionBlock *LoopRegion = Plan.getVectorLoopRegion();
+ // Scan the body of the loop in a topological order to visit each basic block
+ // after having visited its predecessor basic blocks.
+ VPBasicBlock *Header = LoopRegion->getEntryBasicBlock();
+ ReversePostOrderTraversal<VPBlockShallowTraversalWrapper<VPBlockBase *>> RPOT(
+ Header);
+ VPPredicator Predicator(RecipeBuilder);
+ for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly<VPBasicBlock>(RPOT)) {
+ // Handle VPBBs down to the latch.
+ if (VPBB == LoopRegion->getExiting()) {
+ break;
+ }
+
+ if (VPBB == Header) {
+ Predicator.createHeaderMask(Header, FoldTail);
+ continue;
+ }
+ SmallVector<VPWidenPHIRecipe *> Phis;
+ for (VPRecipeBase &R : VPBB->phis())
+ Phis.push_back(cast<VPWidenPHIRecipe>(&R));
+
+ Predicator.createBlockInMask(VPBB);
+
+ for (VPWidenPHIRecipe *Phi : Phis) {
+ PHINode *IRPhi = cast<PHINode>(Phi->getUnderlyingValue());
+
+ unsigned NumIncoming = IRPhi->getNumIncomingValues();
+
+ // We know that all PHIs in non-header blocks are converted into selects,
+ // so we don't have to worry about the insertion order and we can just use
+ // the builder. At this point we generate the predication tree. There may
+ // be duplications since this is a simple recursive scan, but future
+ // optimizations will clean it up.
+
+ // Map incoming IR BasicBlocks to incoming VPValues, for lookup below.
+ // TODO: Add operands and masks in order from the VPlan predecessors.
+ DenseMap<BasicBlock *, VPValue *> VPIncomingValues;
+ DenseMap<BasicBlock *, VPBasicBlock *> VPIncomingBlocks;
+ for (const auto &[Idx, Pred] :
+ enumerate(predecessors(IRPhi->getParent()))) {
+ VPIncomingValues[Pred] = Phi->getOperand(Idx);
+ VPIncomingBlocks[Pred] =
+ cast<VPBasicBlock>(VPBB->getPredecessors()[Idx]);
+ }
+
+ SmallVector<VPValue *, 2> OperandsWithMask;
+ for (unsigned In = 0; In < NumIncoming; In++) {
+ BasicBlock *Pred = IRPhi->getIncomingBlock(In);
+ OperandsWithMask.push_back(VPIncomingValues.lookup(Pred));
+ VPValue *EdgeMask =
+ RecipeBuilder.getEdgeMask(VPIncomingBlocks.lookup(Pred), VPBB);
+ if (!EdgeMask) {
+ assert(In == 0 && "Both null and non-null edge masks found");
+ assert(all_equal(Phi->operands()) &&
+ "Distinct incoming values with one having a full mask");
+ break;
+ }
+ OperandsWithMask.push_back(EdgeMask);
+ }
+ auto *Blend = new VPBlendRecipe(IRPhi, OperandsWithMask);
+ Blend->insertBefore(Phi);
+ Phi->replaceAllUsesWith(Blend);
+ Phi->eraseFromParent();
+ RecipeBuilder.setRecipe(IRPhi, Blend);
+ }
+ }
+
+ for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly<VPBasicBlock>(RPOT)) {
+ // Handle VPBBs down to the latch.
+ if (VPBB == LoopRegion->getExiting()) {
+ VPBlockUtils::connectBlocks(PrevVPBB, VPBB);
+ break;
+ }
+
+ // Flatten the CFG in the loop. Masks for blocks have already been generated
+ // and added to recipes as needed. To do so, first disconnect VPBB from its
+ // successors. Then connect VPBB to the previously visited VPBB.
+ for (auto *Succ : to_vector(VPBB->getSuccessors()))
+ VPBlockUtils::disconnectBlocks(VPBB, Succ);
+ if (PrevVPBB)
+ VPBlockUtils::connectBlocks(PrevVPBB, VPBB);
+ PrevVPBB = VPBB;
+ }
+}
diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
index ce81b2e147df8..7ca2af4531205 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
@@ -2062,8 +2062,9 @@ void VPlanTransforms::handleUncountableEarlyExit(
!OrigLoop->contains(TrueSucc) ? TrueSucc : FalseSucc);
}
- VPValue *EarlyExitNotTakenCond = RecipeBuilder.getBlockInMask(
- OrigLoop->contains(TrueSucc) ? TrueSucc : FalseSucc);
+ VPValue *EarlyExitNotTakenCond =
+ RecipeBuilder.getBlockInMask(cast<VPBasicBlock>(
+ Plan.getVectorLoopRegion()->getEntry()->getSingleSuccessor()));
auto *EarlyExitTakenCond = Builder.createNot(EarlyExitNotTakenCond);
IsEarlyExitTaken =
Builder.createNaryOp(VPInstruction::AnyOf, {EarlyExitTakenCond});
diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.h b/llvm/lib/Transforms/Vectorize/VPlanTransforms.h
index 3dd476a8526d6..4a8e494645059 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanTransforms.h
+++ b/llvm/lib/Transforms/Vectorize/VPlanTransforms.h
@@ -173,6 +173,9 @@ struct VPlanTransforms {
static void
optimizeInductionExitUsers(VPlan &Plan,
DenseMap<VPValue *, VPValue *> &EndValues);
+
+ static void predicateAndLinearize(VPlan &Plan, bool FoldTail,
+ VPRecipeBuilder &RecipeBuilder);
};
} // namespace llvm
More information about the llvm-commits
mailing list