Skip to content

Commit fc9324f

Browse files
david-armMDevereau
authored and
github-actions[bot]
committed
Automerge: [LV] Add support for partial reductions without a binary op (#133922)
Consider IR such as this: for.body: %iv = phi i64 [ 0, %entry ], [ %iv.next, %for.body ] %accum = phi i32 [ 0, %entry ], [ %add, %for.body ] %gep.a = getelementptr i8, ptr %a, i64 %iv %load.a = load i8, ptr %gep.a, align 1 %ext.a = zext i8 %load.a to i32 %add = add i32 %ext.a, %accum %iv.next = add i64 %iv, 1 %exitcond.not = icmp eq i64 %iv.next, 1025 br i1 %exitcond.not, label %for.exit, label %for.body Conceptually we can vectorise this using partial reductions too, although the current loop vectoriser implementation requires the accumulation of a multiply. For AArch64 this is easily done with a udot or sdot with an identity operand, i.e. a vector of (i16 1). In order to do this I had to teach getScaledReductions that the accumulated value may come from a unary op, hence there is only one extension to consider. Similarly, I updated the vplan and AArch64 TTI cost model to understand the possible unary op. --------- Co-authored-by: Matt Devereau <[email protected]>
2 parents 89b525b + f575b18 commit fc9324f

File tree

10 files changed

+1579
-640
lines changed

10 files changed

+1579
-640
lines changed

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

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1325,9 +1325,21 @@ class TargetTransformInfo {
13251325
/// \return The cost of a partial reduction, which is a reduction from a
13261326
/// vector to another vector with fewer elements of larger size. They are
13271327
/// represented by the llvm.experimental.partial.reduce.add intrinsic, which
1328-
/// takes an accumulator and a binary operation operand that itself is fed by
1329-
/// two extends. An example of an operation that uses a partial reduction is a
1330-
/// dot product, which reduces two vectors to another of 4 times fewer and 4
1328+
/// takes an accumulator of type \p AccumType and a second vector operand to
1329+
/// be accumulated, whose element count is specified by \p VF. The type of
1330+
/// reduction is specified by \p Opcode. The second operand passed to the
1331+
/// intrinsic could be the result of an extend, such as sext or zext. In
1332+
/// this case \p BinOp is nullopt, \p InputTypeA represents the type being
1333+
/// extended and \p OpAExtend the operation, i.e. sign- or zero-extend.
1334+
/// Also, \p InputTypeB should be nullptr and OpBExtend should be None.
1335+
/// Alternatively, the second operand could be the result of a binary
1336+
/// operation performed on two extends, i.e.
1337+
/// mul(zext i8 %a -> i32, zext i8 %b -> i32).
1338+
/// In this case \p BinOp may specify the opcode of the binary operation,
1339+
/// \p InputTypeA and \p InputTypeB the types being extended, and
1340+
/// \p OpAExtend, \p OpBExtend the form of extensions. An example of an
1341+
/// operation that uses a partial reduction is a dot product, which reduces
1342+
/// two vectors in binary mul operation to another of 4 times fewer and 4
13311343
/// times larger elements.
13321344
LLVM_ABI InstructionCost getPartialReductionCost(
13331345
unsigned Opcode, Type *InputTypeA, Type *InputTypeB, Type *AccumType,

‎llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp‎

Lines changed: 12 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -5405,11 +5405,21 @@ InstructionCost AArch64TTIImpl::getPartialReductionCost(
54055405

54065406
// Sub opcodes currently only occur in chained cases.
54075407
// Independent partial reduction subtractions are still costed as an add
5408-
if (Opcode != Instruction::Add && Opcode != Instruction::Sub)
5408+
if ((Opcode != Instruction::Add && Opcode != Instruction::Sub) ||
5409+
OpAExtend == TTI::PR_None)
54095410
return Invalid;
54105411

5411-
if (InputTypeA != InputTypeB)
5412+
// We only support multiply binary operations for now, and for muls we
5413+
// require the types being extended to be the same.
5414+
// NOTE: For muls AArch64 supports lowering mixed extensions to a usdot but
5415+
// only if the i8mm or sve/streaming features are available.
5416+
if (BinOp && (*BinOp != Instruction::Mul || InputTypeA != InputTypeB ||
5417+
OpBExtend == TTI::PR_None ||
5418+
(OpAExtend != OpBExtend && !ST->hasMatMulInt8() &&
5419+
!ST->isSVEorStreamingSVEAvailable())))
54125420
return Invalid;
5421+
assert((BinOp || (OpBExtend == TTI::PR_None && !InputTypeB)) &&
5422+
"Unexpected values for OpBExtend or InputTypeB");
54135423

54145424
EVT InputEVT = EVT::getEVT(InputTypeA);
54155425
EVT AccumEVT = EVT::getEVT(AccumType);
@@ -5456,15 +5466,6 @@ InstructionCost AArch64TTIImpl::getPartialReductionCost(
54565466
} else
54575467
return Invalid;
54585468

5459-
// AArch64 supports lowering mixed fixed-width extensions to a usdot but only
5460-
// if the i8mm feature is available.
5461-
if (OpAExtend == TTI::PR_None || OpBExtend == TTI::PR_None ||
5462-
(OpAExtend != OpBExtend && !ST->hasMatMulInt8()))
5463-
return Invalid;
5464-
5465-
if (!BinOp || *BinOp != Instruction::Mul)
5466-
return Invalid;
5467-
54685469
return Cost;
54695470
}
54705471

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

Lines changed: 56 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -8149,15 +8149,15 @@ void VPRecipeBuilder::collectScaledReductions(VFRange &Range) {
81498149
// something that isn't another partial reduction. This is because the
81508150
// extends are intended to be lowered along with the reduction itself.
81518151

8152-
// Build up a set of partial reduction bin ops for efficient use checking.
8153-
SmallSet<User *, 4> PartialReductionBinOps;
8152+
// Build up a set of partial reduction ops for efficient use checking.
8153+
SmallSet<User *, 4> PartialReductionOps;
81548154
for (const auto &[PartialRdx, _] : PartialReductionChains)
8155-
PartialReductionBinOps.insert(PartialRdx.BinOp);
8155+
PartialReductionOps.insert(PartialRdx.ExtendUser);
81568156

81578157
auto ExtendIsOnlyUsedByPartialReductions =
8158-
[&PartialReductionBinOps](Instruction *Extend) {
8158+
[&PartialReductionOps](Instruction *Extend) {
81598159
return all_of(Extend->users(), [&](const User *U) {
8160-
return PartialReductionBinOps.contains(U);
8160+
return PartialReductionOps.contains(U);
81618161
});
81628162
};
81638163

@@ -8166,15 +8166,14 @@ void VPRecipeBuilder::collectScaledReductions(VFRange &Range) {
81668166
for (auto Pair : PartialReductionChains) {
81678167
PartialReductionChain Chain = Pair.first;
81688168
if (ExtendIsOnlyUsedByPartialReductions(Chain.ExtendA) &&
8169-
ExtendIsOnlyUsedByPartialReductions(Chain.ExtendB))
8169+
(!Chain.ExtendB || ExtendIsOnlyUsedByPartialReductions(Chain.ExtendB)))
81708170
ScaledReductionMap.try_emplace(Chain.Reduction, Pair.second);
81718171
}
81728172
}
81738173

81748174
bool VPRecipeBuilder::getScaledReductions(
81758175
Instruction *PHI, Instruction *RdxExitInstr, VFRange &Range,
81768176
SmallVectorImpl<std::pair<PartialReductionChain, unsigned>> &Chains) {
8177-
81788177
if (!CM.TheLoop->contains(RdxExitInstr))
81798178
return false;
81808179

@@ -8203,43 +8202,71 @@ bool VPRecipeBuilder::getScaledReductions(
82038202
if (PhiOp != PHI)
82048203
return false;
82058204

8206-
auto *BinOp = dyn_cast<BinaryOperator>(Op);
8207-
if (!BinOp || !BinOp->hasOneUse())
8208-
return false;
8209-
82108205
using namespace llvm::PatternMatch;
8211-
// Use the side-effect of match to replace BinOp only if the pattern is
8212-
// matched, we don't care at this point whether it actually matched.
8213-
match(BinOp, m_Neg(m_BinOp(BinOp)));
82148206

8215-
Value *A, *B;
8216-
if (!match(BinOp->getOperand(0), m_ZExtOrSExt(m_Value(A))) ||
8217-
!match(BinOp->getOperand(1), m_ZExtOrSExt(m_Value(B))))
8218-
return false;
8207+
// If the update is a binary operator, check both of its operands to see if
8208+
// they are extends. Otherwise, see if the update comes directly from an
8209+
// extend.
8210+
Instruction *Exts[2] = {nullptr};
8211+
BinaryOperator *ExtendUser = dyn_cast<BinaryOperator>(Op);
8212+
std::optional<unsigned> BinOpc;
8213+
Type *ExtOpTypes[2] = {nullptr};
8214+
8215+
auto CollectExtInfo = [&Exts,
8216+
&ExtOpTypes](SmallVectorImpl<Value *> &Ops) -> bool {
8217+
unsigned I = 0;
8218+
for (Value *OpI : Ops) {
8219+
Value *ExtOp;
8220+
if (!match(OpI, m_ZExtOrSExt(m_Value(ExtOp))))
8221+
return false;
8222+
Exts[I] = cast<Instruction>(OpI);
8223+
ExtOpTypes[I] = ExtOp->getType();
8224+
I++;
8225+
}
8226+
return true;
8227+
};
8228+
8229+
if (ExtendUser) {
8230+
if (!ExtendUser->hasOneUse())
8231+
return false;
82198232

8220-
Instruction *ExtA = cast<Instruction>(BinOp->getOperand(0));
8221-
Instruction *ExtB = cast<Instruction>(BinOp->getOperand(1));
8233+
// Use the side-effect of match to replace BinOp only if the pattern is
8234+
// matched, we don't care at this point whether it actually matched.
8235+
match(ExtendUser, m_Neg(m_BinOp(ExtendUser)));
8236+
8237+
SmallVector<Value *> Ops(ExtendUser->operands());
8238+
if (!CollectExtInfo(Ops))
8239+
return false;
8240+
8241+
BinOpc = std::make_optional(ExtendUser->getOpcode());
8242+
} else if (match(Update, m_Add(m_Value(), m_Value()))) {
8243+
// We already know the operands for Update are Op and PhiOp.
8244+
SmallVector<Value *> Ops({Op});
8245+
if (!CollectExtInfo(Ops))
8246+
return false;
8247+
8248+
ExtendUser = Update;
8249+
BinOpc = std::nullopt;
8250+
} else
8251+
return false;
82228252

82238253
TTI::PartialReductionExtendKind OpAExtend =
8224-
TargetTransformInfo::getPartialReductionExtendKind(ExtA);
8254+
TTI::getPartialReductionExtendKind(Exts[0]);
82258255
TTI::PartialReductionExtendKind OpBExtend =
8226-
TargetTransformInfo::getPartialReductionExtendKind(ExtB);
8227-
8228-
PartialReductionChain Chain(RdxExitInstr, ExtA, ExtB, BinOp);
8256+
Exts[1] ? TTI::getPartialReductionExtendKind(Exts[1]) : TTI::PR_None;
8257+
PartialReductionChain Chain(RdxExitInstr, Exts[0], Exts[1], ExtendUser);
82298258

82308259
TypeSize PHISize = PHI->getType()->getPrimitiveSizeInBits();
8231-
TypeSize ASize = A->getType()->getPrimitiveSizeInBits();
8232-
8260+
TypeSize ASize = ExtOpTypes[0]->getPrimitiveSizeInBits();
82338261
if (!PHISize.hasKnownScalarFactor(ASize))
82348262
return false;
8235-
82368263
unsigned TargetScaleFactor = PHISize.getKnownScalarFactor(ASize);
82378264

82388265
if (LoopVectorizationPlanner::getDecisionAndClampRange(
82398266
[&](ElementCount VF) {
82408267
InstructionCost Cost = TTI->getPartialReductionCost(
8241-
Update->getOpcode(), A->getType(), B->getType(), PHI->getType(),
8242-
VF, OpAExtend, OpBExtend, BinOp->getOpcode(), CM.CostKind);
8268+
Update->getOpcode(), ExtOpTypes[0], ExtOpTypes[1],
8269+
PHI->getType(), VF, OpAExtend, OpBExtend, BinOpc, CM.CostKind);
82438270
return Cost.isValid();
82448271
},
82458272
Range)) {

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

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -24,22 +24,23 @@ struct HistogramInfo;
2424
struct VFRange;
2525

2626
/// A chain of instructions that form a partial reduction.
27-
/// Designed to match: reduction_bin_op (bin_op (extend (A), (extend (B))),
28-
/// accumulator).
27+
/// Designed to match either:
28+
/// reduction_bin_op (extend (A), accumulator), or
29+
/// reduction_bin_op (bin_op (extend (A), (extend (B))), accumulator).
2930
struct PartialReductionChain {
3031
PartialReductionChain(Instruction *Reduction, Instruction *ExtendA,
31-
Instruction *ExtendB, Instruction *BinOp)
32-
: Reduction(Reduction), ExtendA(ExtendA), ExtendB(ExtendB), BinOp(BinOp) {
33-
}
32+
Instruction *ExtendB, Instruction *ExtendUser)
33+
: Reduction(Reduction), ExtendA(ExtendA), ExtendB(ExtendB),
34+
ExtendUser(ExtendUser) {}
3435
/// The top-level binary operation that forms the reduction to a scalar
3536
/// after the loop body.
3637
Instruction *Reduction;
3738
/// The extension of each of the inner binary operation's operands.
3839
Instruction *ExtendA;
3940
Instruction *ExtendB;
4041

41-
/// The binary operation using the extends that is then reduced.
42-
Instruction *BinOp;
42+
/// The user of the extends that is then reduced.
43+
Instruction *ExtendUser;
4344
};
4445

4546
/// Helper class to create VPRecipies from IR instructions.

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

Lines changed: 52 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -296,49 +296,70 @@ bool VPRecipeBase::isScalarCast() const {
296296
InstructionCost
297297
VPPartialReductionRecipe::computeCost(ElementCount VF,
298298
VPCostContext &Ctx) const {
299-
std::optional<unsigned> Opcode = std::nullopt;
300-
VPValue *BinOp = getOperand(1);
299+
std::optional<unsigned> Opcode;
300+
VPValue *Op = getOperand(0);
301+
VPRecipeBase *OpR = Op->getDefiningRecipe();
301302

302-
// If the partial reduction is predicated, a select will be operand 0 rather
303-
// than the binary op
303+
// If the partial reduction is predicated, a select will be operand 0
304304
using namespace llvm::VPlanPatternMatch;
305-
if (match(getOperand(1), m_Select(m_VPValue(), m_VPValue(), m_VPValue())))
306-
BinOp = BinOp->getDefiningRecipe()->getOperand(1);
307-
308-
// If BinOp is a negation, use the side effect of match to assign the actual
309-
// binary operation to BinOp
310-
match(BinOp, m_Binary<Instruction::Sub>(m_SpecificInt(0), m_VPValue(BinOp)));
311-
VPRecipeBase *BinOpR = BinOp->getDefiningRecipe();
312-
313-
if (auto *WidenR = dyn_cast<VPWidenRecipe>(BinOpR))
314-
Opcode = std::make_optional(WidenR->getOpcode());
315-
316-
VPRecipeBase *ExtAR = BinOpR->getOperand(0)->getDefiningRecipe();
317-
VPRecipeBase *ExtBR = BinOpR->getOperand(1)->getDefiningRecipe();
305+
if (match(getOperand(1), m_Select(m_VPValue(), m_VPValue(Op), m_VPValue()))) {
306+
OpR = Op->getDefiningRecipe();
307+
}
318308

319-
auto *PhiType = Ctx.Types.inferScalarType(getOperand(1));
320-
auto *InputTypeA = Ctx.Types.inferScalarType(ExtAR ? ExtAR->getOperand(0)
321-
: BinOpR->getOperand(0));
322-
auto *InputTypeB = Ctx.Types.inferScalarType(ExtBR ? ExtBR->getOperand(0)
323-
: BinOpR->getOperand(1));
309+
Type *InputTypeA = nullptr, *InputTypeB = nullptr;
310+
TTI::PartialReductionExtendKind ExtAType = TTI::PR_None,
311+
ExtBType = TTI::PR_None;
324312

325313
auto GetExtendKind = [](VPRecipeBase *R) {
326-
// The extend could come from outside the plan.
327314
if (!R)
328-
return TargetTransformInfo::PR_None;
315+
return TTI::PR_None;
329316
auto *WidenCastR = dyn_cast<VPWidenCastRecipe>(R);
330317
if (!WidenCastR)
331-
return TargetTransformInfo::PR_None;
318+
return TTI::PR_None;
332319
if (WidenCastR->getOpcode() == Instruction::CastOps::ZExt)
333-
return TargetTransformInfo::PR_ZeroExtend;
320+
return TTI::PR_ZeroExtend;
334321
if (WidenCastR->getOpcode() == Instruction::CastOps::SExt)
335-
return TargetTransformInfo::PR_SignExtend;
336-
return TargetTransformInfo::PR_None;
322+
return TTI::PR_SignExtend;
323+
return TTI::PR_None;
337324
};
338325

339-
return Ctx.TTI.getPartialReductionCost(
340-
getOpcode(), InputTypeA, InputTypeB, PhiType, VF, GetExtendKind(ExtAR),
341-
GetExtendKind(ExtBR), Opcode, Ctx.CostKind);
326+
// Pick out opcode, type/ext information and use sub side effects from a widen
327+
// recipe.
328+
auto HandleWiden = [&](VPWidenRecipe *Widen) {
329+
if (match(Widen,
330+
m_Binary<Instruction::Sub>(m_SpecificInt(0), m_VPValue(Op)))) {
331+
Widen = dyn_cast<VPWidenRecipe>(Op->getDefiningRecipe());
332+
}
333+
Opcode = Widen->getOpcode();
334+
VPRecipeBase *ExtAR = Widen->getOperand(0)->getDefiningRecipe();
335+
VPRecipeBase *ExtBR = Widen->getOperand(1)->getDefiningRecipe();
336+
InputTypeA = Ctx.Types.inferScalarType(ExtAR ? ExtAR->getOperand(0)
337+
: Widen->getOperand(0));
338+
InputTypeB = Ctx.Types.inferScalarType(ExtBR ? ExtBR->getOperand(0)
339+
: Widen->getOperand(1));
340+
ExtAType = GetExtendKind(ExtAR);
341+
ExtBType = GetExtendKind(ExtBR);
342+
};
343+
344+
if (isa<VPWidenCastRecipe>(OpR)) {
345+
InputTypeA = Ctx.Types.inferScalarType(OpR->getOperand(0));
346+
ExtAType = GetExtendKind(OpR);
347+
} else if (isa<VPReductionPHIRecipe>(OpR)) {
348+
auto RedPhiOp1R = getOperand(1)->getDefiningRecipe();
349+
if (isa<VPWidenCastRecipe>(RedPhiOp1R)) {
350+
InputTypeA = Ctx.Types.inferScalarType(RedPhiOp1R->getOperand(0));
351+
ExtAType = GetExtendKind(RedPhiOp1R);
352+
} else if (auto Widen = dyn_cast<VPWidenRecipe>(RedPhiOp1R))
353+
HandleWiden(Widen);
354+
} else if (auto Widen = dyn_cast<VPWidenRecipe>(OpR)) {
355+
HandleWiden(Widen);
356+
} else if (auto Reduction = dyn_cast<VPPartialReductionRecipe>(OpR)) {
357+
return Reduction->computeCost(VF, Ctx);
358+
}
359+
auto *PhiType = Ctx.Types.inferScalarType(getOperand(1));
360+
return Ctx.TTI.getPartialReductionCost(getOpcode(), InputTypeA, InputTypeB,
361+
PhiType, VF, ExtAType, ExtBType,
362+
Opcode, Ctx.CostKind);
342363
}
343364

344365
void VPPartialReductionRecipe::execute(VPTransformState &State) {

‎llvm/test/CodeGen/AArch64/neon-partial-reduce-dot-product.ll‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1255,4 +1255,3 @@ entry:
12551255
%partial.reduce = tail call <2 x i64> @llvm.experimental.vector.partial.reduce.add(<2 x i64> %acc, <8 x i64> %input.wide)
12561256
ret <2 x i64> %partial.reduce
12571257
}
1258-

0 commit comments

Comments
 (0)