Skip to content

Commit 3fdd4ef

Browse files
Address comments. Add RUN line to test files. Change a comment.
Remove Mask as operand in Partial Reduction Recipe. Instead just mask the input when creating the recipe.
1 parent 0e0f8a1 commit 3fdd4ef

File tree

5 files changed

+2135
-23
lines changed

5 files changed

+2135
-23
lines changed

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -8898,8 +8898,13 @@ VPRecipeBuilder::tryToCreatePartialReduction(Instruction *Reduction,
88988898
std::swap(BinOp, Accumulator);
88998899

89008900
VPValue *Mask = getBlockInMask(Reduction->getParent());
8901-
8902-
return new VPPartialReductionRecipe(Reduction->getOpcode(), BinOp, Accumulator, Mask, Reduction);
8901+
if (Mask) {
8902+
VPValue *Zero =
8903+
Plan.getOrAddLiveIn(ConstantInt::get(Reduction->getType(), 0));
8904+
BinOp = Builder.createSelect(Mask, BinOp, Zero, Reduction->getDebugLoc());
8905+
}
8906+
return new VPPartialReductionRecipe(Reduction->getOpcode(), BinOp,
8907+
Accumulator, Reduction);
89038908
}
89048909

89058910
void LoopVectorizationPlanner::buildVPlansWithVPRecipes(ElementCount MinVF,
@@ -9711,11 +9716,9 @@ void LoopVectorizationPlanner::adjustRecipesForReductions(
97119716
// beginning of the dedicated latch block.
97129717
auto *OrigExitingVPV = PhiR->getBackedgeValue();
97139718
auto *NewExitingVPV = PhiR->getBackedgeValue();
9714-
// Don't add selects here for partial reductions because the phi and partial
9715-
// reduction values have less vector elements than Cond. But, each operand
9716-
// in a select instruction needs to have the same number of vector elements,
9717-
// so the compiler would crash. Instead, a select, with the active lane
9718-
// mask, is applied to the inputs to the partial reduction.
9719+
// Don't output selects for partial reductions because they have an output
9720+
// with fewer lanes than the VF. So the operands of the select would have
9721+
// different numbers of lanes. Partial reductions mask the input instead.
97199722
if (!PhiR->isInLoop() && CM.foldTailByMasking() &&
97209723
!isa<VPPartialReductionRecipe>(OrigExitingVPV->getDefiningRecipe())) {
97219724
VPValue *Cond = RecipeBuilder.getBlockInMask(OrigLoop->getHeader());

llvm/lib/Transforms/Vectorize/VPlan.h

Lines changed: 3 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -2455,11 +2455,10 @@ class VPPartialReductionRecipe : public VPSingleDefRecipe {
24552455

24562456
public:
24572457
VPPartialReductionRecipe(Instruction *ReductionInst, VPValue *Op0,
2458-
VPValue *Op1, VPValue *Mask = nullptr)
2459-
: VPPartialReductionRecipe(ReductionInst->getOpcode(), Op0, Op1, Mask,
2458+
VPValue *Op1)
2459+
: VPPartialReductionRecipe(ReductionInst->getOpcode(), Op0, Op1,
24602460
ReductionInst) {}
24612461
VPPartialReductionRecipe(unsigned Opcode, VPValue *Op0, VPValue *Op1,
2462-
VPValue *Mask = nullptr,
24632462
Instruction *ReductionInst = nullptr)
24642463
: VPSingleDefRecipe(VPDef::VPPartialReductionSC,
24652464
ArrayRef<VPValue *>({Op0, Op1}), ReductionInst),
@@ -2469,18 +2468,12 @@ class VPPartialReductionRecipe : public VPSingleDefRecipe {
24692468
assert((isa<VPReductionPHIRecipe>(AccumulatorRecipe) ||
24702469
isa<VPPartialReductionRecipe>(AccumulatorRecipe)) &&
24712470
"Unexpected operand order for partial reduction recipe");
2472-
if (Mask)
2473-
addOperand(Mask);
24742471
}
24752472
~VPPartialReductionRecipe() override = default;
24762473

24772474
VPPartialReductionRecipe *clone() override {
24782475
return new VPPartialReductionRecipe(Opcode, getOperand(0), getOperand(1),
2479-
getMask(), getUnderlyingInstr());
2480-
}
2481-
2482-
VPValue *getMask() const {
2483-
return getNumOperands() == 3 ? getOperand(2) : nullptr;
2476+
getUnderlyingInstr());
24842477
}
24852478

24862479
VP_CLASSOF_IMPL(VPDef::VPPartialReductionSC)

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -283,6 +283,13 @@ VPPartialReductionRecipe::computeCost(ElementCount VF,
283283
VPCostContext &Ctx) const {
284284
std::optional<unsigned> Opcode = std::nullopt;
285285
VPRecipeBase *BinOpR = getOperand(0)->getDefiningRecipe();
286+
287+
// If the partial reduction is predicated, a select will be operand 0 rather
288+
// than the binary op
289+
using namespace llvm::VPlanPatternMatch;
290+
if (match(getOperand(0), m_Select(m_VPValue(), m_VPValue(), m_VPValue())))
291+
BinOpR = BinOpR->getOperand(1)->getDefiningRecipe();
292+
286293
if (auto *WidenR = dyn_cast<VPWidenRecipe>(BinOpR))
287294
Opcode = std::make_optional(WidenR->getOpcode());
288295

@@ -327,12 +334,6 @@ void VPPartialReductionRecipe::execute(VPTransformState &State) {
327334

328335
Type *RetTy = PhiVal->getType();
329336

330-
VPValue *Mask = getMask();
331-
if (Mask) {
332-
Value *MaskVal = State.get(Mask);
333-
Value *Zero = ConstantInt::get(BinOpVal->getType(), 0);
334-
BinOpVal = Builder.CreateSelect(MaskVal, BinOpVal, Zero);
335-
}
336337
CallInst *V = Builder.CreateIntrinsic(
337338
RetTy, Intrinsic::experimental_vector_partial_reduce_add,
338339
{PhiVal, BinOpVal}, nullptr, "partial.reduce");

0 commit comments

Comments
 (0)