Skip to content

[LoopUnroll] Introduce parallel reduction phis when unrolling. #149470

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 6 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions llvm/include/llvm/Analysis/TargetTransformInfo.h
Original file line number Diff line number Diff line change
Expand Up @@ -638,6 +638,8 @@ class TargetTransformInfo {
/// Fall back to the generic logic to determine whether multi-exit unrolling
/// is profitable if set to false.
bool RuntimeUnrollMultiExit;

DenseMap<PHINode *, RecurrenceDescriptor> ParallelizeReductions;
};

/// Get target-customized preferences for the generic loop unrolling
Expand Down
3 changes: 3 additions & 0 deletions llvm/include/llvm/Transforms/Utils/UnrollLoop.h
Original file line number Diff line number Diff line change
Expand Up @@ -163,6 +163,9 @@ LLVM_ABI bool computeUnrollCount(
TargetTransformInfo::UnrollingPreferences &UP,
TargetTransformInfo::PeelingPreferences &PP, bool &UseUpperBound);

LLVM_ABI std::optional<RecurrenceDescriptor>
canParallelizeReductionWhenUnrolling(PHINode &Phi, Loop *L,
ScalarEvolution *SE);
} // end namespace llvm

#endif // LLVM_TRANSFORMS_UTILS_UNROLLLOOP_H
137 changes: 137 additions & 0 deletions llvm/lib/Transforms/Utils/LoopUnroll.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@
#include "llvm/IR/DiagnosticInfo.h"
#include "llvm/IR/Dominators.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/Instruction.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/IntrinsicInst.h"
Expand Down Expand Up @@ -108,6 +109,9 @@ UnrollVerifyLoopInfo("unroll-verify-loopinfo", cl::Hidden,
#endif
);

static cl::opt<bool> UnrollAddParallelReductions(
"unroll-add-parallel-reductions", cl::init(false), cl::Hidden,
cl::desc("Allow unrolling to add parallel reduction phis."));

/// Check if unrolling created a situation where we need to insert phi nodes to
/// preserve LCSSA form.
Expand Down Expand Up @@ -660,6 +664,39 @@ llvm::UnrollLoop(Loop *L, UnrollLoopOptions ULO, LoopInfo *LI,
OrigPHINode.push_back(cast<PHINode>(I));
}

// Collect phi nodes for reductions for which we can introduce multiple
// parallel reduction phis and compute the final reduction result after the
// loop. This requires a single exit block after unrolling. This is ensured by
// restricting to single-block loops where the unrolled iterations are known
// to not exit.
DenseMap<PHINode *, RecurrenceDescriptor> Reductions;
bool CanAddAdditionalAccumulators =
UnrollAddParallelReductions && !CompletelyUnroll &&
L->getNumBlocks() == 1 &&
(ULO.Runtime ||
(ExitInfos.contains(Header) && ((ExitInfos[Header].TripCount != 0 &&
ExitInfos[Header].BreakoutTrip == 0))));

// Limit parallelizing reductions to unroll counts of 4 or less for now.
// TODO: The number of parallel reductions should depend on the number of
// execution units. We also don't have to add a parallel reduction phi per
// unrolled iteration, but could for example add a parallel phi for every 2
// unrolled iterations.
if (CanAddAdditionalAccumulators && ULO.Count <= 4) {
for (PHINode &Phi : Header->phis()) {
auto RdxDesc = canParallelizeReductionWhenUnrolling(Phi, L, SE);
if (!RdxDesc)
continue;

// Only handle duplicate phis for a single reduction for now.
// TODO: Handle any number of reductions
if (!Reductions.empty())
continue;

Reductions[&Phi] = *RdxDesc;
}
}

std::vector<BasicBlock *> Headers;
std::vector<BasicBlock *> Latches;
Headers.push_back(Header);
Expand Down Expand Up @@ -710,6 +747,7 @@ llvm::UnrollLoop(Loop *L, UnrollLoopOptions ULO, LoopInfo *LI,
// latch. This is a reasonable default placement if we don't have block
// frequencies, and if we do, well the layout will be adjusted later.
auto BlockInsertPt = std::next(LatchBlock->getIterator());
SmallVector<Instruction *> PartialReductions;
for (unsigned It = 1; It != ULO.Count; ++It) {
SmallVector<BasicBlock *, 8> NewBlocks;
SmallDenseMap<const Loop *, Loop *, 4> NewLoops;
Expand All @@ -733,6 +771,31 @@ llvm::UnrollLoop(Loop *L, UnrollLoopOptions ULO, LoopInfo *LI,
for (PHINode *OrigPHI : OrigPHINode) {
PHINode *NewPHI = cast<PHINode>(VMap[OrigPHI]);
Value *InVal = NewPHI->getIncomingValueForBlock(LatchBlock);

// Use cloned phis as parallel phis for partial reductions, which will
// get combined to the final reduction result after the loop.
if (Reductions.contains(OrigPHI)) {
// Collect partial reduction results.
if (PartialReductions.empty())
PartialReductions.push_back(cast<Instruction>(InVal));
PartialReductions.push_back(cast<Instruction>(VMap[InVal]));

// Update the start value for the cloned phis to use the identity
// value for the reduction.
const RecurrenceDescriptor &RdxDesc = Reductions[OrigPHI];
NewPHI->setIncomingValueForBlock(
L->getLoopPreheader(),
getRecurrenceIdentity(RdxDesc.getRecurrenceKind(),
OrigPHI->getType(),
RdxDesc.getFastMathFlags()));

// Update NewPHI to use the cloned value for the iteration and move
// to header.
NewPHI->replaceUsesOfWith(InVal, VMap[InVal]);
NewPHI->moveBefore(OrigPHI->getIterator());
continue;
}

if (Instruction *InValI = dyn_cast<Instruction>(InVal))
if (It > 1 && L->contains(InValI))
InVal = LastValueMap[InValI];
Expand Down Expand Up @@ -832,7 +895,11 @@ llvm::UnrollLoop(Loop *L, UnrollLoopOptions ULO, LoopInfo *LI,
PN->replaceAllUsesWith(PN->getIncomingValueForBlock(Preheader));
PN->eraseFromParent();
} else if (ULO.Count > 1) {
if (Reductions.contains(PN))
continue;

Value *InVal = PN->removeIncomingValue(LatchBlock, false);

// If this value was defined in the loop, take the value defined by the
// last iteration of the loop.
if (Instruction *InValI = dyn_cast<Instruction>(InVal)) {
Expand Down Expand Up @@ -1010,6 +1077,38 @@ llvm::UnrollLoop(Loop *L, UnrollLoopOptions ULO, LoopInfo *LI,
}
}

// If there are partial reductions, create code in the exit block to compute
// the final result and update users of the final result.
if (!PartialReductions.empty()) {
BasicBlock *ExitBlock = L->getExitBlock();
assert(ExitBlock &&
"Can only introduce parallel reduction phis with single exit block");
assert(Reductions.size() == 1 &&
"currently only a single reduction is supported");
Value *FinalRdxValue = PartialReductions.back();
Value *RdxResult = nullptr;
for (PHINode &Phi : ExitBlock->phis()) {
if (Phi.getIncomingValueForBlock(L->getLoopLatch()) != FinalRdxValue)
continue;
if (!RdxResult) {
RdxResult = PartialReductions.front();
IRBuilder Builder(ExitBlock, ExitBlock->getFirstNonPHIIt());
RecurKind RK = Reductions.begin()->second.getRecurrenceKind();
for (Instruction *RdxPart : drop_begin(PartialReductions)) {
RdxResult = Builder.CreateBinOp(
(Instruction::BinaryOps)RecurrenceDescriptor::getOpcode(RK),
RdxPart, RdxResult, "bin.rdx");
}
NeedToFixLCSSA = true;
for (Instruction *RdxPart : PartialReductions)
RdxPart->dropPoisonGeneratingFlags();
}

Phi.replaceAllUsesWith(RdxResult);
continue;
}
}

if (DTUToUse) {
// Apply updates to the DomTree.
DT = &DTU.getDomTree();
Expand Down Expand Up @@ -1111,3 +1210,41 @@ MDNode *llvm::GetUnrollMetadata(MDNode *LoopID, StringRef Name) {
}
return nullptr;
}

std::optional<RecurrenceDescriptor>
llvm::canParallelizeReductionWhenUnrolling(PHINode &Phi, Loop *L,
ScalarEvolution *SE) {
RecurrenceDescriptor RdxDesc;
if (!RecurrenceDescriptor::isReductionPHI(&Phi, L, RdxDesc,
/*DemandedBits=*/nullptr,
/*AC=*/nullptr, /*DT=*/nullptr, SE))
return std::nullopt;
RecurKind RK = RdxDesc.getRecurrenceKind();
// Skip unsupported reductions.
// TODO: Handle additional reductions, including FP and min-max
// reductions.
if (!RecurrenceDescriptor::isIntegerRecurrenceKind(RK) ||
RecurrenceDescriptor::isAnyOfRecurrenceKind(RK) ||
RecurrenceDescriptor::isFindIVRecurrenceKind(RK) ||
RecurrenceDescriptor::isMinMaxRecurrenceKind(RK))
return std::nullopt;

if (RdxDesc.IntermediateStore)
return std::nullopt;

// Don't unroll reductions with constant ops; those can be folded to a
// single induction update.
if (any_of(cast<Instruction>(Phi.getIncomingValueForBlock(L->getLoopLatch()))
->operands(),
IsaPred<Constant>))
return std::nullopt;

BasicBlock *Latch = L->getLoopLatch();
if (!Latch ||
!is_contained(
cast<Instruction>(Phi.getIncomingValueForBlock(Latch))->operands(),
&Phi))
return std::nullopt;

return RdxDesc;
}
Loading
Loading