Skip to content

Commit 595d8ff

Browse files
SamTebbs33SamTebbs33
committed
Bundle partial reductions inside VPMulAccumulateReductionRecipe
1 parent a8edb86 commit 595d8ff

File tree

10 files changed

+674
-783
lines changed

10 files changed

+674
-783
lines changed

‎llvm/include/llvm/Analysis/TargetTransformInfo.h

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -219,6 +219,8 @@ class TargetTransformInfo {
219219
/// Get the kind of extension that an instruction represents.
220220
static PartialReductionExtendKind
221221
getPartialReductionExtendKind(Instruction *I);
222+
static PartialReductionExtendKind
223+
getPartialReductionExtendKind(Instruction::CastOps ExtOpcode);
222224

223225
/// Construct a TTI object using a type implementing the \c Concept
224226
/// API below.

‎llvm/lib/Analysis/TargetTransformInfo.cpp

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -993,6 +993,19 @@ TargetTransformInfo::getPartialReductionExtendKind(Instruction *I) {
993993
return PR_None;
994994
}
995995

996+
TargetTransformInfo::PartialReductionExtendKind
997+
TargetTransformInfo::getPartialReductionExtendKind(
998+
Instruction::CastOps ExtOpcode) {
999+
switch (ExtOpcode) {
1000+
case Instruction::CastOps::ZExt:
1001+
return PR_ZeroExtend;
1002+
case Instruction::CastOps::SExt:
1003+
return PR_SignExtend;
1004+
default:
1005+
return PR_None;
1006+
}
1007+
}
1008+
9961009
TTI::CastContextHint
9971010
TargetTransformInfo::getCastContextHint(const Instruction *I) {
9981011
if (!I)

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

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8886,9 +8886,6 @@ VPRecipeBuilder::tryToCreatePartialReduction(Instruction *Reduction,
88868886
"Expected an ADD or SUB operation for predicated partial "
88878887
"reductions (because the neutral element in the mask is zero)!");
88888888
Cond = getBlockInMask(Reduction->getParent());
8889-
VPValue *Zero =
8890-
Plan.getOrAddLiveIn(ConstantInt::get(Reduction->getType(), 0));
8891-
BinOp = Builder.createSelect(Cond, BinOp, Zero, Reduction->getDebugLoc());
88928889
}
88938890
return new VPPartialReductionRecipe(ReductionOpcode, Accumulator, BinOp, Cond,
88948891
Reduction);

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

Lines changed: 15 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2340,8 +2340,7 @@ class VPPartialReductionRecipe : public VPReductionRecipe {
23402340
: VPPartialReductionRecipe(ReductionInst->getOpcode(), Op0, Op1, Cond,
23412341
ReductionInst) {}
23422342
VPPartialReductionRecipe(unsigned Opcode, VPValue *Op0, VPValue *Op1,
2343-
VPValue *Cond,
2344-
Instruction *ReductionInst = nullptr)
2343+
VPValue *Cond, Instruction *ReductionInst = nullptr)
23452344
: VPReductionRecipe(VPDef::VPPartialReductionSC, RecurKind::Add,
23462345
FastMathFlags(), ReductionInst,
23472346
ArrayRef<VPValue *>({Op0, Op1}), Cond, false, {}),
@@ -2356,8 +2355,7 @@ class VPPartialReductionRecipe : public VPReductionRecipe {
23562355

23572356
VPPartialReductionRecipe *clone() override {
23582357
return new VPPartialReductionRecipe(Opcode, getOperand(0), getOperand(1),
2359-
getCondOp(),
2360-
getUnderlyingInstr());
2358+
getCondOp(), getUnderlyingInstr());
23612359
}
23622360

23632361
VP_CLASSOF_IMPL(VPDef::VPPartialReductionSC)
@@ -2372,6 +2370,9 @@ class VPPartialReductionRecipe : public VPReductionRecipe {
23722370
/// Get the binary op's opcode.
23732371
unsigned getOpcode() const { return Opcode; }
23742372

2373+
/// Get the binary op this reduction is applied to.
2374+
VPValue *getBinOp() const { return getOperand(1); }
2375+
23752376
#if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
23762377
/// Print the recipe.
23772378
void print(raw_ostream &O, const Twine &Indent,
@@ -2499,6 +2500,9 @@ class VPMulAccumulateReductionRecipe : public VPReductionRecipe {
24992500

25002501
Type *ResultTy;
25012502

2503+
/// If the reduction this is based on is a partial reduction.
2504+
bool IsPartialReduction = false;
2505+
25022506
/// For cloning VPMulAccumulateReductionRecipe.
25032507
VPMulAccumulateReductionRecipe(VPMulAccumulateReductionRecipe *MulAcc)
25042508
: VPReductionRecipe(
@@ -2508,7 +2512,8 @@ class VPMulAccumulateReductionRecipe : public VPReductionRecipe {
25082512
WrapFlagsTy(MulAcc->hasNoUnsignedWrap(), MulAcc->hasNoSignedWrap()),
25092513
MulAcc->getDebugLoc()),
25102514
ExtOp(MulAcc->getExtOpcode()), IsNonNeg(MulAcc->isNonNeg()),
2511-
ResultTy(MulAcc->getResultType()) {}
2515+
ResultTy(MulAcc->getResultType()),
2516+
IsPartialReduction(MulAcc->isPartialReduction()) {}
25122517

25132518
public:
25142519
VPMulAccumulateReductionRecipe(VPReductionRecipe *R, VPWidenRecipe *Mul,
@@ -2521,7 +2526,8 @@ class VPMulAccumulateReductionRecipe : public VPReductionRecipe {
25212526
WrapFlagsTy(Mul->hasNoUnsignedWrap(), Mul->hasNoSignedWrap()),
25222527
R->getDebugLoc()),
25232528
ExtOp(Ext0->getOpcode()), IsNonNeg(Ext0->isNonNeg()),
2524-
ResultTy(ResultTy) {
2529+
ResultTy(ResultTy),
2530+
IsPartialReduction(isa<VPPartialReductionRecipe>(R)) {
25252531
assert(RecurrenceDescriptor::getOpcode(getRecurrenceKind()) ==
25262532
Instruction::Add &&
25272533
"The reduction instruction in MulAccumulateteReductionRecipe must "
@@ -2592,6 +2598,9 @@ class VPMulAccumulateReductionRecipe : public VPReductionRecipe {
25922598

25932599
/// Return the non negative flag of the ext recipe.
25942600
bool isNonNeg() const { return IsNonNeg; }
2601+
2602+
/// Return if the underlying reduction recipe is a partial reduction.
2603+
bool isPartialReduction() const { return IsPartialReduction; }
25952604
};
25962605

25972606
/// VPReplicateRecipe replicates a given instruction producing multiple scalar

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

Lines changed: 19 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -287,14 +287,9 @@ InstructionCost
287287
VPPartialReductionRecipe::computeCost(ElementCount VF,
288288
VPCostContext &Ctx) const {
289289
std::optional<unsigned> Opcode = std::nullopt;
290-
VPValue *BinOp = getOperand(1);
290+
VPValue *BinOp = getBinOp();
291291

292-
// If the partial reduction is predicated, a select will be operand 0 rather
293-
// than the binary op
294292
using namespace llvm::VPlanPatternMatch;
295-
if (match(getOperand(1), m_Select(m_VPValue(), m_VPValue(), m_VPValue())))
296-
BinOp = BinOp->getDefiningRecipe()->getOperand(1);
297-
298293
// If BinOp is a negation, use the side effect of match to assign the actual
299294
// binary operation to BinOp
300295
match(BinOp, m_Binary<Instruction::Sub>(m_SpecificInt(0), m_VPValue(BinOp)));
@@ -338,12 +333,18 @@ void VPPartialReductionRecipe::execute(VPTransformState &State) {
338333
assert(getOpcode() == Instruction::Add &&
339334
"Unhandled partial reduction opcode");
340335

341-
Value *BinOpVal = State.get(getOperand(1));
342-
Value *PhiVal = State.get(getOperand(0));
336+
Value *BinOpVal = State.get(getBinOp());
337+
Value *PhiVal = State.get(getChainOp());
343338
assert(PhiVal && BinOpVal && "Phi and Mul must be set");
344339

345340
Type *RetTy = PhiVal->getType();
346341

342+
/// Mask the bin op output.
343+
if (VPValue *Cond = getCondOp()) {
344+
Value *Zero = ConstantInt::get(BinOpVal->getType(), 0);
345+
BinOpVal = Builder.CreateSelect(State.get(Cond), BinOpVal, Zero);
346+
}
347+
347348
CallInst *V = Builder.CreateIntrinsic(
348349
RetTy, Intrinsic::experimental_vector_partial_reduce_add,
349350
{PhiVal, BinOpVal}, nullptr, "partial.reduce");
@@ -2432,6 +2433,14 @@ VPExtendedReductionRecipe::computeCost(ElementCount VF,
24322433
InstructionCost
24332434
VPMulAccumulateReductionRecipe::computeCost(ElementCount VF,
24342435
VPCostContext &Ctx) const {
2436+
if (isPartialReduction()) {
2437+
return Ctx.TTI.getPartialReductionCost(
2438+
Instruction::Add, Ctx.Types.inferScalarType(getVecOp0()),
2439+
Ctx.Types.inferScalarType(getVecOp1()), getResultType(), VF,
2440+
TTI::getPartialReductionExtendKind(getExtOpcode()),
2441+
TTI::getPartialReductionExtendKind(getExtOpcode()), Instruction::Mul);
2442+
}
2443+
24352444
Type *RedTy = Ctx.Types.inferScalarType(this);
24362445
auto *SrcVecTy =
24372446
cast<VectorType>(toVectorTy(Ctx.Types.inferScalarType(getVecOp0()), VF));
@@ -2509,6 +2518,8 @@ void VPMulAccumulateReductionRecipe::print(raw_ostream &O, const Twine &Indent,
25092518
O << " = ";
25102519
getChainOp()->printAsOperand(O, SlotTracker);
25112520
O << " + ";
2521+
if (isPartialReduction())
2522+
O << "partial.";
25122523
O << "reduce."
25132524
<< Instruction::getOpcodeName(
25142525
RecurrenceDescriptor::getOpcode(getRecurrenceKind()))

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

Lines changed: 36 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2158,9 +2158,14 @@ expandVPMulAccumulateReduction(VPMulAccumulateReductionRecipe *MulAcc) {
21582158
Mul->insertBefore(MulAcc);
21592159

21602160
// Generate VPReductionRecipe.
2161-
auto *Red = new VPReductionRecipe(
2162-
MulAcc->getRecurrenceKind(), FastMathFlags(), MulAcc->getChainOp(), Mul,
2163-
MulAcc->getCondOp(), MulAcc->isOrdered(), MulAcc->getDebugLoc());
2161+
VPReductionRecipe *Red = nullptr;
2162+
if (MulAcc->isPartialReduction())
2163+
Red = new VPPartialReductionRecipe(Instruction::Add, MulAcc->getChainOp(),
2164+
Mul, MulAcc->getCondOp());
2165+
else
2166+
Red = new VPReductionRecipe(MulAcc->getRecurrenceKind(), FastMathFlags(),
2167+
MulAcc->getChainOp(), Mul, MulAcc->getCondOp(),
2168+
MulAcc->isOrdered(), MulAcc->getDebugLoc());
21642169
Red->insertBefore(MulAcc);
21652170

21662171
MulAcc->replaceAllUsesWith(Red);
@@ -2432,12 +2437,39 @@ static void tryToCreateAbstractReductionRecipe(VPReductionRecipe *Red,
24322437
Red->replaceAllUsesWith(AbstractR);
24332438
}
24342439

2440+
static void
2441+
tryToCreateAbstractPartialReductionRecipe(VPPartialReductionRecipe *PRed) {
2442+
if (PRed->getOpcode() != Instruction::Add)
2443+
return;
2444+
2445+
VPRecipeBase *BinOpR = PRed->getBinOp()->getDefiningRecipe();
2446+
auto *BinOp = dyn_cast<VPWidenRecipe>(BinOpR);
2447+
if (!BinOp || BinOp->getOpcode() != Instruction::Mul)
2448+
return;
2449+
2450+
auto *Ext0 = dyn_cast<VPWidenCastRecipe>(BinOp->getOperand(0));
2451+
auto *Ext1 = dyn_cast<VPWidenCastRecipe>(BinOp->getOperand(1));
2452+
// TODO: Make work with extends of different signedness
2453+
if (!Ext0 || Ext0->hasMoreThanOneUniqueUser() || !Ext1 ||
2454+
Ext1->hasMoreThanOneUniqueUser() ||
2455+
Ext0->getOpcode() != Ext1->getOpcode())
2456+
return;
2457+
2458+
auto *AbstractR = new VPMulAccumulateReductionRecipe(PRed, BinOp, Ext0, Ext1,
2459+
Ext0->getResultType());
2460+
AbstractR->insertBefore(PRed);
2461+
PRed->replaceAllUsesWith(AbstractR);
2462+
PRed->eraseFromParent();
2463+
}
2464+
24352465
void VPlanTransforms::convertToAbstractRecipes(VPlan &Plan, VPCostContext &Ctx,
24362466
VFRange &Range) {
24372467
for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly<VPBasicBlock>(
24382468
vp_depth_first_deep(Plan.getVectorLoopRegion()))) {
24392469
for (VPRecipeBase &R : make_early_inc_range(*VPBB)) {
2440-
if (auto *Red = dyn_cast<VPReductionRecipe>(&R))
2470+
if (auto *PRed = dyn_cast<VPPartialReductionRecipe>(&R))
2471+
tryToCreateAbstractPartialReductionRecipe(PRed);
2472+
else if (auto *Red = dyn_cast<VPReductionRecipe>(&R))
24412473
tryToCreateAbstractReductionRecipe(Red, Ctx, Range);
24422474
}
24432475
}

0 commit comments

Comments
 (0)