[llvm] [LV][AArch64] Provide option to use partial reductions by default (PR #216001)
Graham Hunter via llvm-commits
llvm-commits at lists.llvm.org
Thu Aug 13 03:01:08 PDT 2026
https://github.com/huntergr-arm created https://github.com/llvm/llvm-project/pull/216001
Partial reduction intrinsics are explicitly unordered, which may allow us to perform extra optimizations that may result in a different ordering of lanes within a vector. This would only be done if the target indicated the intrinsic was at least as cheap as the normal accumulation operation.
Off by default to start with.
As an example of an optimization which would be easier/less fragile with this change, see #206047 which currently tries to match the entire reduction. If the in-loop update was a partial reduction intrinsic (even at the same VF), then we could just implement the whole thing as a DAGCombine.
Some follow-up work will be needed before turning on by default; there's 25 AArch64 check-all failures (just need to update test checks) and some build failures + regressions in some spec benchmarks to look at.
>From 36a73b7035d8ec3f8482c3119be09ac435d8a647 Mon Sep 17 00:00:00 2001
From: Graham Hunter <graham.hunter at arm.com>
Date: Wed, 12 Aug 2026 13:12:12 +0000
Subject: [PATCH] [LV][AArch64] Provide option to use partial reductions by
default
Partial reduction intrinsics are explicitly unordered, which may
allow us to perform extra optimizations that may result in a different
ordering of lanes within a vector. This would only be done if
the target indicated the intrinsic was at least as cheap as the
normal accumulation operation.
Off by default to start with.
---
.../AArch64/AArch64TargetTransformInfo.cpp | 18 +++++++-
.../Transforms/Vectorize/LoopVectorize.cpp | 8 +++-
llvm/lib/Transforms/Vectorize/VPlan.h | 4 +-
.../Transforms/Vectorize/VPlanTransforms.cpp | 46 ++++++++++++++++++-
.../Transforms/Vectorize/VPlanTransforms.h | 10 ++--
5 files changed, 77 insertions(+), 9 deletions(-)
diff --git a/llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp b/llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp
index 1f4b1b69c0433..b18f5d2680fbc 100644
--- a/llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp
+++ b/llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp
@@ -6704,10 +6704,24 @@ InstructionCost AArch64TTIImpl::getPartialReductionCost(
return Invalid;
if ((Opcode != Instruction::Add && Opcode != Instruction::Sub &&
- Opcode != Instruction::FAdd && Opcode != Instruction::FSub) ||
- OpAExtend == TTI::PR_None)
+ Opcode != Instruction::FAdd && Opcode != Instruction::FSub))
return Invalid;
+ // If none of the operands are extended and there's no extra BinOp, just
+ // cost this as the equivalent arithmetic instruction.
+ // TODO: Depending on VF and element type, we may be able to improve on this.
+ if (OpAExtend == TTI::PR_None) {
+ if (OpBExtend != TTI::PR_None || BinOp)
+ return Invalid;
+
+ assert(InputTypeA == InputTypeB && InputTypeA == AccumType &&
+ "Type mismatch with no extensions.");
+
+ VectorType *VTy = VectorType::get(AccumType, VF);
+
+ return getArithmeticInstrCost(Opcode, VTy, CostKind);
+ }
+
// Floating-point partial reductions are invalid if `reassoc` and `contract`
// are not allowed.
if (AccumType->isFloatingPointTy()) {
diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index 11a49ba16199c..9f10e06e0b39c 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -413,6 +413,12 @@ static cl::opt<bool> EnableEarlyExitVectorizationWithSideEffects(
cl::desc("Enable vectorization of early exit loops with uncountable exits "
"and side effects"));
+static cl::opt<bool>
+ UsePartialReduceByDefault("use-partial-reduce-by-default", cl::init(false),
+ cl::Hidden,
+ cl::desc("Use partial reduction intrinsics for "
+ "all supported unordered reductions."));
+
// Returns true if the epilogue VF has been set to a non-zero value other than
// VF=1 (scalar).
static bool hasForcedEpilogueVF() {
@@ -6857,7 +6863,7 @@ VPlanPtr LoopVectorizationPlanner::tryToBuildVPlan(VPlanPtr Plan,
// and mulacc-reduction are implemented.
if (!CM.foldTailWithEVL()) {
RUN_VPLAN_PASS(VPlanTransforms::createPartialReductions, *Plan, CostCtx,
- Range);
+ Range, UsePartialReduceByDefault);
RUN_VPLAN_PASS(VPlanTransforms::convertToAbstractRecipes, *Plan, CostCtx,
Range);
}
diff --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h
index b2e87a8f4f52d..f792197a91546 100644
--- a/llvm/lib/Transforms/Vectorize/VPlan.h
+++ b/llvm/lib/Transforms/Vectorize/VPlan.h
@@ -3316,7 +3316,9 @@ class LLVM_ABI_FOR_TEST VPReductionRecipe : public VPRecipeWithIRFlags {
/// Return true if the in-loop reduction is conditional.
bool isConditional() const { return IsConditional; };
/// Returns true if the reduction outputs a vector with a scaled down VF.
- bool isPartialReduction() const { return getVFScaleFactor() > 1; }
+ bool isPartialReduction() const {
+ return std::holds_alternative<RdxUnordered>(Style);
+ }
/// Returns true if the reduction is in-loop.
bool isInLoop() const {
return std::holds_alternative<RdxInLoop>(Style) ||
diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
index eb5958c37043d..00fbf5c18b204 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
@@ -5186,13 +5186,15 @@ getScaledReductions(VPReductionPHIRecipe *RedPhiR) {
void VPlanTransforms::createPartialReductions(VPlan &Plan,
VPCostContext &CostCtx,
- VFRange &Range) {
+ VFRange &Range,
+ bool UsePartialReduceByDefault) {
// Find all possible valid partial reductions, grouping chains by their PHI.
// This grouping allows invalidating the whole chain, if any link is not a
// valid partial reduction.
MapVector<VPReductionPHIRecipe *, SmallVector<VPPartialReductionChain>>
ChainsByPhi;
VPBasicBlock *HeaderVPBB = Plan.getVectorLoopRegion()->getEntryBasicBlock();
+ SmallVector<VPReductionPHIRecipe *, 4> UnorderedReductions;
for (VPRecipeBase &R : HeaderVPBB->phis()) {
auto *RedPhiR = dyn_cast<VPReductionPHIRecipe>(&R);
if (!RedPhiR)
@@ -5200,6 +5202,48 @@ void VPlanTransforms::createPartialReductions(VPlan &Plan,
if (auto Chains = getScaledReductions(RedPhiR))
ChainsByPhi.try_emplace(RedPhiR, std::move(*Chains));
+ else if (UsePartialReduceByDefault &&
+ (RedPhiR->getRecurrenceKind() == RecurKind::Add ||
+ (RedPhiR->getRecurrenceKind() == RecurKind::FAdd &&
+ !RedPhiR->isOrdered() && !RedPhiR->isInLoop())))
+ UnorderedReductions.push_back(RedPhiR);
+ }
+
+ // For general unordered reductions which aren't part of a candidate
+ // chain for a scaled partial reduction, we can potentially still use
+ // the intrinsic to allow for more optimization later on.
+ for (auto *Rdx : UnorderedReductions) {
+ auto *Backedge = dyn_cast<VPWidenRecipe>(Rdx->getBackedgeValue());
+ VPValue *OtherOp;
+ if (!Backedge ||
+ !match(Backedge,
+ m_CombineOr(m_c_FAdd(m_Specific(Rdx), m_VPValue(OtherOp)),
+ m_c_Add(m_Specific(Rdx), m_VPValue(OtherOp)))))
+ continue;
+
+ // If the target indicates that the intrinsic is as cheap as (or cheaper
+ // than) the add, then prefer the intrinsic.
+ if (!LoopVectorizationPlanner::getDecisionAndClampRange(
+ [&CostCtx, Rdx, Backedge](ElementCount VF) {
+ InstructionCost CurrentCost = Backedge->computeCost(VF, CostCtx);
+ Type *ScalarTy = Backedge->getScalarType();
+ InstructionCost PRCost = CostCtx.TTI.getPartialReductionCost(
+ Backedge->getOpcode(), ScalarTy, ScalarTy, ScalarTy, VF,
+ TTI::PR_None, TTI::PR_None, std::nullopt, CostCtx.CostKind,
+ Rdx->getFastMathFlagsOrNone());
+ return PRCost <= CurrentCost;
+ },
+ Range))
+ continue;
+
+ auto *Partial = new VPReductionRecipe(
+ Rdx->getRecurrenceKind(), Rdx->getFastMathFlagsOrNone(),
+ Backedge->getUnderlyingInstr(), Rdx, OtherOp, nullptr,
+ getReductionStyle(/*InLoop=*/false, /*Ordered=*/false,
+ /*ScaleFactor=*/1));
+ Partial->insertBefore(Backedge);
+ Backedge->replaceAllUsesWith(Partial);
+ Backedge->eraseFromParent();
}
if (ChainsByPhi.empty())
diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.h b/llvm/lib/Transforms/Vectorize/VPlanTransforms.h
index 1c9959665ee7b..5ea19fd4ad5cd 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanTransforms.h
+++ b/llvm/lib/Transforms/Vectorize/VPlanTransforms.h
@@ -576,11 +576,13 @@ struct VPlanTransforms {
static void optimizeFindIVReductions(VPlan &Plan,
PredicatedScalarEvolution &PSE, Loop &L);
- /// Detect and create partial reduction recipes for scaled reductions in
- /// \p Plan. Must be called after recipe construction. If partial reductions
- /// are only valid for a subset of VFs in Range, Range.End is updated.
+ /// Detect and create partial reduction recipes for scaled or unordered
+ /// reductions in \p Plan. Must be called after recipe construction. If
+ /// partial reductions are only valid for a subset of VFs in Range, Range.End
+ /// is updated.
static void createPartialReductions(VPlan &Plan, VPCostContext &CostCtx,
- VFRange &Range);
+ VFRange &Range,
+ bool UsePartialReduceByDefault);
/// Convert load/store VPInstructions in \p Plan into widened or replicate
/// recipes. Non load/store input instructions are left unchanged.
More information about the llvm-commits
mailing list