Skip to content

Commit

Permalink
[LV] Retrieve reduction resume values directly for epilogue vec. (NFC)
Browse files Browse the repository at this point in the history
Use the reduction resume values from the phis in the scalar header,
instead of collecting them in a map. This removes some complexity from
the general executePlan code paths and pushes it to only the epilogue
vectorization part.
  • Loading branch information
fhahn committed Sep 29, 2024
1 parent 1714b11 commit 1edd220
Show file tree
Hide file tree
Showing 2 changed files with 12 additions and 23 deletions.
9 changes: 4 additions & 5 deletions llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h
Original file line number Diff line number Diff line change
Expand Up @@ -435,11 +435,10 @@ class LoopVectorizationPlanner {
/// \p ExpandedSCEVs is passed during execution of the plan for epilogue loop
/// to re-use expansion results generated during main plan execution.
///
/// Returns a mapping of SCEVs to their expanded IR values and a mapping for
/// the reduction resume values. Note that this is a temporary workaround
/// needed due to the current epilogue handling.
std::pair<DenseMap<const SCEV *, Value *>,
DenseMap<const RecurrenceDescriptor *, Value *>>
/// Returns a mapping of SCEVs to their expanded IR values.
/// Note that this is a temporary workaround needed due to the current
/// epilogue handling.
DenseMap<const SCEV *, Value *>
executePlan(ElementCount VF, unsigned UF, VPlan &BestPlan,
InnerLoopVectorizer &LB, DominatorTree *DT,
bool IsEpilogueVectorization,
Expand Down
26 changes: 8 additions & 18 deletions llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -676,11 +676,6 @@ class InnerLoopVectorizer {
/// Structure to hold information about generated runtime checks, responsible
/// for cleaning the checks, if vectorization turns out unprofitable.
GeneratedRTChecks &RTChecks;

// Holds the resume values for reductions in the loops, used to set the
// correct start value of reduction PHIs when vectorizing the epilogue.
SmallMapVector<const RecurrenceDescriptor *, PHINode *, 4>
ReductionResumeValues;
};

/// Encapsulate information regarding vectorization of a loop and its epilogue.
Expand Down Expand Up @@ -7426,10 +7421,9 @@ static void addRuntimeUnrollDisableMetaData(Loop *L) {
}

// Check if \p RedResult is a ComputeReductionResult instruction, and if it is
// create a merge phi node for it and add it to \p ReductionResumeValues.
// create a merge phi node for it.
static void createAndCollectMergePhiForReduction(
VPInstruction *RedResult,
DenseMap<const RecurrenceDescriptor *, Value *> &ReductionResumeValues,
VPTransformState &State, Loop *OrigLoop, BasicBlock *LoopMiddleBlock,
bool VectorizingEpilogue) {
if (!RedResult ||
Expand Down Expand Up @@ -7487,13 +7481,9 @@ static void createAndCollectMergePhiForReduction(
OrigPhi->setIncomingValue(SelfEdgeBlockIdx, BCBlockPhi);
Instruction *LoopExitInst = RdxDesc.getLoopExitInstr();
OrigPhi->setIncomingValue(IncomingEdgeBlockIdx, LoopExitInst);

ReductionResumeValues[&RdxDesc] = BCBlockPhi;
}

std::pair<DenseMap<const SCEV *, Value *>,
DenseMap<const RecurrenceDescriptor *, Value *>>
LoopVectorizationPlanner::executePlan(
DenseMap<const SCEV *, Value *> LoopVectorizationPlanner::executePlan(
ElementCount BestVF, unsigned BestUF, VPlan &BestVPlan,
InnerLoopVectorizer &ILV, DominatorTree *DT, bool IsEpilogueVectorization,
const DenseMap<const SCEV *, Value *> *ExpandedSCEVs) {
Expand Down Expand Up @@ -7579,12 +7569,11 @@ LoopVectorizationPlanner::executePlan(
BestVPlan.execute(&State);

// 2.5 Collect reduction resume values.
DenseMap<const RecurrenceDescriptor *, Value *> ReductionResumeValues;
auto *ExitVPBB =
cast<VPBasicBlock>(BestVPlan.getVectorLoopRegion()->getSingleSuccessor());
for (VPRecipeBase &R : *ExitVPBB) {
createAndCollectMergePhiForReduction(
dyn_cast<VPInstruction>(&R), ReductionResumeValues, State, OrigLoop,
dyn_cast<VPInstruction>(&R), State, OrigLoop,
State.CFG.VPBB2IRBB[ExitVPBB], ExpandedSCEVs);
}

Expand Down Expand Up @@ -7634,7 +7623,7 @@ LoopVectorizationPlanner::executePlan(
setBranchWeights(*MiddleTerm, Weights, /*IsExpected=*/false);
}

return {State.ExpandedSCEVs, ReductionResumeValues};
return State.ExpandedSCEVs;
}

//===--------------------------------------------------------------------===//
Expand Down Expand Up @@ -10121,8 +10110,8 @@ bool LoopVectorizePass::processLoop(Loop *L) {
EPI, &LVL, &CM, BFI, PSI, Checks);

std::unique_ptr<VPlan> BestMainPlan(BestPlan.duplicate());
const auto &[ExpandedSCEVs, ReductionResumeValues] = LVP.executePlan(
EPI.MainLoopVF, EPI.MainLoopUF, *BestMainPlan, MainILV, DT, true);
auto ExpandedSCEVs = LVP.executePlan(EPI.MainLoopVF, EPI.MainLoopUF,
*BestMainPlan, MainILV, DT, true);
++LoopsVectorized;

// Second pass vectorizes the epilogue and adjusts the control flow
Expand Down Expand Up @@ -10167,10 +10156,11 @@ bool LoopVectorizePass::processLoop(Loop *L) {
Value *ResumeV = nullptr;
// TODO: Move setting of resume values to prepareToExecute.
if (auto *ReductionPhi = dyn_cast<VPReductionPHIRecipe>(&R)) {
ResumeV = cast<PHINode>(ReductionPhi->getUnderlyingInstr())
->getIncomingValueForBlock(L->getLoopPreheader());
const RecurrenceDescriptor &RdxDesc =
ReductionPhi->getRecurrenceDescriptor();
RecurKind RK = RdxDesc.getRecurrenceKind();
ResumeV = ReductionResumeValues.find(&RdxDesc)->second;
if (RecurrenceDescriptor::isAnyOfRecurrenceKind(RK)) {
// VPReductionPHIRecipes for AnyOf reductions expect a boolean as
// start value; compare the final value from the main vector loop
Expand Down

0 comments on commit 1edd220

Please sign in to comment.