Skip to content

[LV] Support argmin/argmax with strict predicates. #170223

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.

Sign up for GitHub

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
fhahn wants to merge 2 commits into
base: main
Choose a base branch
from
Open

[LV] Support argmin/argmax with strict predicates. #170223

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a9b3c0f
[LV] Support argmin/argmax with strict predicates.
fhahn Dec 1, 2025
c6a0bc0
!fixup generalize, create wide IV if needed.
fhahn Dec 1, 2025
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
6 changes: 6 additions & 0 deletions llvm/include/llvm/Analysis/IVDescriptors.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ class Loop;
class PredicatedScalarEvolution;
class ScalarEvolution;
class SCEV;
class SCEVAddRecExpr;
class StoreInst;

/// These are the kinds of recurrences that we support.
Expand Down Expand Up @@ -310,6 +311,11 @@ class RecurrenceDescriptor {
isFindLastIVRecurrenceKind(Kind);
}

/// Returns true if \p AR's range is valid for either FindFirstIV or
/// FindLastIV reductions i.e. if the sentinel value is outside \p AR's range.
Comment on lines +314 to +315
Copy link
Collaborator

@ayalz ayalz Dec 2, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The sentinel value should be outside AR's range, when FindFirstIV or FindLastIV work with a condition that may be false for all values inside AR's range, as in that case the out-of-bound sentinel value helps to convey that answer. But for argmin/argmax with strict/FirstIV or non-strict/LastIV, the condition is guaranteed to be true at-least once inside AR's range. So having a sentinel inside AR's range should be fine?

Follow-up thought: even if the condition may be false for all values inside AR's range, but the final value to return in that case coincides with the natural sentinel inside AR's range, it could potentially also be optimized as above. E.g., suppose FindLastIV work with an arbitrary condition a[i]>0 in the range [0,100], where if all a's are zeroes the value to return is 0.

Another thought: argmin/argmax is interesting not only because it's condition is guaranteed to hold at-least once in the range, but also because it supports FindFirstIV naturally (w/o negating the IV - which conceptually still means it's a FindLastIV but on a decreasing derived IV) - by relying on the running min/max recurrence whose strict comparison precludes repeated occurrences. This relates to a discussion we had on some earlier patch, where adding another, dependent (boolean "found") recurrence was found to have significant negative impact.

Copy link
Contributor Author

@fhahn fhahn Dec 3, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am not sure if currently the condition must hold at least once unless I am missing something, because there are no restrictions on the start values e.g. the minimum reduction does not have to start at the max value, so the IV may never get selected in the loop, and the minimuxm index reduction may start at a value that's different to the IV start value.

But IV start == min.idx start and condition known to be true at least once are good to optimize separately. Note that for the initial FindLastIV, the range is checked during the IVDescriptor analysis.

static bool isValidIVRangeForFindIV(const SCEVAddRecExpr *AR, bool IsSigned,
bool IsFindFirstIV, ScalarEvolution &SE);

/// Returns the type of the recurrence. This type can be narrower than the
/// actual type of the Phi if the recurrence has been type-promoted.
Type *getRecurrenceType() const { return RecurrenceType; }
Expand Down
73 changes: 39 additions & 34 deletions llvm/lib/Analysis/IVDescriptors.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -715,6 +715,36 @@ RecurrenceDescriptor::isAnyOfPattern(Loop *Loop, PHINode *OrigPhi,
return InstDesc(I, RecurKind::AnyOf);
}

bool RecurrenceDescriptor::isValidIVRangeForFindIV(const SCEVAddRecExpr *AR,
bool IsSigned,
bool IsFindFirstIV,
ScalarEvolution &SE) {
const ConstantRange IVRange =
IsSigned ? SE.getSignedRange(AR) : SE.getUnsignedRange(AR);
unsigned NumBits = AR->getType()->getIntegerBitWidth();
ConstantRange ValidRange = ConstantRange::getEmpty(NumBits);

if (IsFindFirstIV) {
if (IsSigned)
ValidRange =
ConstantRange::getNonEmpty(APInt::getSignedMinValue(NumBits),
APInt::getSignedMaxValue(NumBits) - 1);
else
ValidRange = ConstantRange::getNonEmpty(APInt::getMinValue(NumBits),
APInt::getMaxValue(NumBits) - 1);
} else {
APInt Sentinel = IsSigned ? APInt::getSignedMinValue(NumBits)
: APInt::getMinValue(NumBits);
ValidRange = ConstantRange::getNonEmpty(Sentinel + 1, Sentinel);
}

LLVM_DEBUG(dbgs() << "LV: " << (IsFindFirstIV ? "FindFirstIV" : "FindLastIV")
<< " valid range is " << ValidRange << ", and the range of "
<< *AR << " is " << IVRange << "\n");

return ValidRange.contains(IVRange);
}

// We are looking for loops that do something like this:
// int r = 0;
// for (int i = 0; i < n; i++) {
Expand Down Expand Up @@ -792,49 +822,24 @@ RecurrenceDescriptor::isFindIVPattern(RecurKind Kind, Loop *TheLoop,
// [Signed|Unsigned]Max(<recurrence type>) for FindFirstIV.
// TODO: This range restriction can be lifted by adding an additional
// virtual OR reduction.
auto CheckRange = [&](bool IsSigned) {
const ConstantRange IVRange =
IsSigned ? SE.getSignedRange(AR) : SE.getUnsignedRange(AR);
unsigned NumBits = Ty->getIntegerBitWidth();
ConstantRange ValidRange = ConstantRange::getEmpty(NumBits);
if (isFindLastIVRecurrenceKind(Kind)) {
APInt Sentinel = IsSigned ? APInt::getSignedMinValue(NumBits)
: APInt::getMinValue(NumBits);
ValidRange = ConstantRange::getNonEmpty(Sentinel + 1, Sentinel);
} else {
if (IsSigned)
ValidRange =
ConstantRange::getNonEmpty(APInt::getSignedMinValue(NumBits),
APInt::getSignedMaxValue(NumBits) - 1);
else
ValidRange = ConstantRange::getNonEmpty(
APInt::getMinValue(NumBits), APInt::getMaxValue(NumBits) - 1);
}

LLVM_DEBUG(dbgs() << "LV: "
<< (isFindLastIVRecurrenceKind(Kind) ? "FindLastIV"
: "FindFirstIV")
<< " valid range is " << ValidRange
<< ", and the range of " << *AR << " is " << IVRange
<< "\n");

// Ensure the induction variable does not wrap around by verifying that
// its range is fully contained within the valid range.
return ValidRange.contains(IVRange);
};
bool IsFindFirstIV = isFindFirstIVRecurrenceKind(Kind);
if (isFindLastIVRecurrenceKind(Kind)) {
if (CheckRange(true))
if (RecurrenceDescriptor::isValidIVRangeForFindIV(
cast<SCEVAddRecExpr>(AR), /*IsSigned=*/true, IsFindFirstIV, SE))
return RecurKind::FindLastIVSMax;
if (CheckRange(false))
if (RecurrenceDescriptor::isValidIVRangeForFindIV(
cast<SCEVAddRecExpr>(AR), /*IsSigned=*/false, IsFindFirstIV, SE))
return RecurKind::FindLastIVUMax;
return std::nullopt;
}
assert(isFindFirstIVRecurrenceKind(Kind) &&
"Kind must either be a FindLastIV or FindFirstIV");

if (CheckRange(true))
if (RecurrenceDescriptor::isValidIVRangeForFindIV(
cast<SCEVAddRecExpr>(AR), /*IsSigned=*/true, IsFindFirstIV, SE))
return RecurKind::FindFirstIVSMin;
if (CheckRange(false))
if (RecurrenceDescriptor::isValidIVRangeForFindIV(
cast<SCEVAddRecExpr>(AR), /*IsSigned=*/false, IsFindFirstIV, SE))
return RecurKind::FindFirstIVUMin;
return std::nullopt;
};
Expand Down
3 changes: 1 addition & 2 deletions llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8587,8 +8587,7 @@ VPlanPtr LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(

// Apply mandatory transformation to handle reductions with multiple in-loop
// uses if possible, bail out otherwise.
if (!VPlanTransforms::runPass(VPlanTransforms::handleMultiUseReductions,
*Plan))
if (!VPlanTransforms::handleMultiUseReductions(*Plan))
return nullptr;
// Apply mandatory transformation to handle FP maxnum/minnum reduction with
// NaNs if possible, bail out otherwise.
Expand Down
4 changes: 3 additions & 1 deletion llvm/lib/Transforms/Vectorize/VPlan.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2155,7 +2155,9 @@ class VPWidenInductionRecipe : public VPHeaderPHIRecipe {
/// incoming value, its start value.
unsigned getNumIncoming() const override { return 1; }

PHINode *getPHINode() const { return cast<PHINode>(getUnderlyingValue()); }
PHINode *getPHINode() const {
return cast_if_present<PHINode>(getUnderlyingValue());
}

/// Returns the induction descriptor for the recipe.
const InductionDescriptor &getInductionDescriptor() const { return IndDesc; }
Expand Down
162 changes: 151 additions & 11 deletions llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,13 +13,16 @@

#include "LoopVectorizationPlanner.h"
#include "VPlan.h"
#include "VPlanAnalysis.h"
#include "VPlanCFG.h"
#include "VPlanDominatorTree.h"
#include "VPlanPatternMatch.h"
#include "VPlanTransforms.h"
#include "VPlanUtils.h"
#include "llvm/Analysis/LoopInfo.h"
#include "llvm/Analysis/LoopIterator.h"
#include "llvm/Analysis/ScalarEvolution.h"
#include "llvm/Analysis/ScalarEvolutionExpressions.h"
#include "llvm/IR/InstrTypes.h"
#include "llvm/IR/MDBuilder.h"
#include "llvm/Transforms/Utils/LoopUtils.h"
Expand Down Expand Up @@ -997,6 +1000,134 @@ bool VPlanTransforms::handleMaxMinNumReductions(VPlan &Plan) {
return true;
}

/// For argmin/argmax reductions with strict predicates, convert the existing
/// FindLastIV reduction to a new UMin reduction of a wide canonical IV. If the
/// original IV was not canonical, a new canonical wide IV is added, and the
/// final result is scaled back to the original IV.
static bool handleStrictArgMinArgMax(VPlan &Plan,
VPReductionPHIRecipe *MinMaxPhiR,
VPReductionPHIRecipe *FindIVPhiR,
VPWidenIntOrFpInductionRecipe *WideIV,
VPInstruction *MinMaxResult) {
Type *Ty = Plan.getVectorLoopRegion()->getCanonicalIVType();
if (Ty != VPTypeAnalysis(Plan).inferScalarType(FindIVPhiR))
return false;

// If the original wide IV is not canonical, create a new one. The wide IV is
// guaranteed to not wrap for all lanes that are active in the vector loop.
if (!WideIV->isCanonical()) {
VPValue *Zero = Plan.getOrAddLiveIn(ConstantInt::get(Ty, 0));
VPValue *One = Plan.getOrAddLiveIn(ConstantInt::get(Ty, 1));
auto *WidenCanIV = new VPWidenIntOrFpInductionRecipe(
nullptr, Zero, One, WideIV->getVFValue(),
WideIV->getInductionDescriptor(), VPIRFlags(), WideIV->getDebugLoc());
WidenCanIV->insertBefore(WideIV);

// Update the select to use the wide canonical IV.
auto *SelectRecipe = cast<VPSingleDefRecipe>(
FindIVPhiR->getBackedgeValue()->getDefiningRecipe());
if (SelectRecipe->getOperand(1) == WideIV)
SelectRecipe->setOperand(1, WidenCanIV);
else if (SelectRecipe->getOperand(2) == WideIV)
SelectRecipe->setOperand(2, WidenCanIV);
}

// Create the new UMin reduction recipe to track the minimum index.
assert(!FindIVPhiR->isInLoop() && !FindIVPhiR->isOrdered() &&
"inloop and ordered reductions not supported");
VPValue *MaxInt =
Plan.getConstantInt(APInt::getMaxValue(Ty->getIntegerBitWidth()));
ReductionStyle Style = RdxUnordered{FindIVPhiR->getVFScaleFactor()};
auto *MinIdxPhiR = new VPReductionPHIRecipe(
dyn_cast_or_null<PHINode>(FindIVPhiR->getUnderlyingValue()),
RecurKind::UMin, *MaxInt, Style,
FindIVPhiR->hasUsesOutsideReductionChain());
MinIdxPhiR->addOperand(FindIVPhiR->getBackedgeValue());
MinIdxPhiR->insertBefore(FindIVPhiR);

VPInstruction *FindLastIVResult =
findUserOf<VPInstruction::ComputeFindIVResult>(FindIVPhiR);
MinMaxResult->moveBefore(*FindLastIVResult->getParent(),
FindLastIVResult->getIterator());

// The reduction using MinMaxPhiR needs adjusting to compute the correct
// result:
// 1. We need to find the first canonical IV for which the condition based
// on the min/max recurrence is true,
// 2. Compare the partial min/max reduction result to its final value and,
// 3. Select the lanes of the partial UMin reduction of the canonical wide
// IV which correspond to the lanes matching the min/max reduction result.
// 4. Scale the final select canonical IV back to the original IV using
// VPDerivedIVRecipe.
// 5. If the minimum value matches the start value, the condition in the
// loop was never true, return the start value in that case.
//
// The original reductions need adjusting:
// For example, this transforms
// vp<%min.result> = compute-reduction-result ir<%min.val>,
// ir<%min.val.next>
// vp<%find.iv.result = compute-find-iv-result ir<%min.idx>, ir<0>,
// SENTINEL, vp<%min.idx.next>
//
// into:
// vp<%min.result> = compute-reduction-result ir<%min.val>, ir<%min.val.next>
// vp<%final.min.cmp> = icmp eq ir<%min.val.next>, vp<%min.result>
// vp<%final.min.iv> = select vp<%final.min.cmp>, ir<%min.idx.next>, ir<-1>
// vp<%13> = compute-reduction-result ir<%min.idx>, vp<%final.min.iv>
// vp<%scaled.result.iv> = DERIVED-IV ir<20> + vp<%13> * ir<1>
// vp<%threshold.cmp> = icmp slt vp<%min.result>, ir<0>
// vp<%final.result> = select vp<%threshold.cmp>, vp<%scaled.result.iv>,
// ir<%original.start>

VPBuilder Builder(FindLastIVResult);
VPValue *MinMaxExiting = MinMaxResult->getOperand(1);
auto *FinalMinMaxCmp =
Builder.createICmp(CmpInst::ICMP_EQ, MinMaxExiting, MinMaxResult);
VPValue *LastIVExiting = FindLastIVResult->getOperand(3);
auto *FinalIVSelect =
Builder.createSelect(FinalMinMaxCmp, LastIVExiting, MaxInt);
VPSingleDefRecipe *FinalResult = Builder.createNaryOp(
VPInstruction::ComputeReductionResult, {MinIdxPhiR, FinalIVSelect}, {},
FindLastIVResult->getDebugLoc());

// If we used a new wide canonical IV convert the reduction result back to the
// original IV scale before the final select.
if (!WideIV->isCanonical()) {
auto *DerivedIVRecipe =
new VPDerivedIVRecipe(InductionDescriptor::IK_IntInduction,
nullptr, // No FPBinOp for integer induction
WideIV->getStartValue(), FinalResult,
WideIV->getStepValue(), "derived.iv.result");
DerivedIVRecipe->insertBefore(&*Builder.getInsertPoint());
FinalResult = DerivedIVRecipe;
}

auto GetPred = [&MinMaxPhiR]() {
switch (MinMaxPhiR->getRecurrenceKind()) {
case RecurKind::UMin:
return CmpInst::ICMP_ULT;
case RecurKind::SMin:
return CmpInst::ICMP_SLT;
case RecurKind::UMax:
return CmpInst::ICMP_UGT;
case RecurKind::SMax:
return CmpInst::ICMP_SGT;
default:
llvm_unreachable("must be an integer min/max recurrence kind");
}
};
// If the final min/max value matches the start value, the condition in the
// loop was always false, i.e. no induction value has been selected. If that's
// the case, use the original start value.
VPValue *MinMaxLT =
Builder.createICmp(GetPred(), MinMaxResult, MinMaxPhiR->getStartValue());
VPValue *Res = Builder.createSelect(MinMaxLT, FinalResult,
FindLastIVResult->getOperand(1));
FindIVPhiR->replaceAllUsesWith(MinIdxPhiR);
FindLastIVResult->replaceAllUsesWith(Res);
return true;
}

bool VPlanTransforms::handleMultiUseReductions(VPlan &Plan) {
for (auto &PhiR : make_early_inc_range(
Plan.getVectorLoopRegion()->getEntryBasicBlock()->phis())) {
Expand All @@ -1008,7 +1139,7 @@ bool VPlanTransforms::handleMultiUseReductions(VPlan &Plan) {
// MinMaxPhiR has users outside the reduction cycle in the loop. Check if
// the only other user is a FindLastIV reduction. MinMaxPhiR must have
// exactly 3 users: 1) the min/max operation, the compare of a FindLastIV
// reduction and ComputeReductionResult. The comparisom must compare
// reduction and ComputeReductionResult. The comparison must compare
// MinMaxPhiR against the min/max operand used for the min/max reduction
// and only be used by the select of the FindLastIV reduction.
RecurKind RdxKind = MinMaxPhiR->getRecurrenceKind();
Expand Down Expand Up @@ -1080,33 +1211,42 @@ bool VPlanTransforms::handleMultiUseReductions(VPlan &Plan) {
FindIVPhiR->getRecurrenceKind()))
return false;

assert(!FindIVPhiR->isInLoop() && !FindIVPhiR->isOrdered() &&
"cannot handle inloop/ordered reductions yet");

// TODO: Support cases where IVOp is the IV increment.
if (!match(IVOp, m_TruncOrSelf(m_VPValue(IVOp))) ||
!isa<VPWidenIntOrFpInductionRecipe>(IVOp))
return false;

CmpInst::Predicate RdxPredicate = [RdxKind]() {
// Check if the predicate is compatible with the reduction kind.
bool IsValidPredicate = [RdxKind, Pred]() {
switch (RdxKind) {
case RecurKind::UMin:
return CmpInst::ICMP_UGE;
return Pred == CmpInst::ICMP_UGE || Pred == CmpInst::ICMP_UGT;
case RecurKind::UMax:
return CmpInst::ICMP_ULE;
return Pred == CmpInst::ICMP_ULE || Pred == CmpInst::ICMP_ULT;
case RecurKind::SMax:
return CmpInst::ICMP_SLE;
return Pred == CmpInst::ICMP_SLE || Pred == CmpInst::ICMP_SLT;
case RecurKind::SMin:
return CmpInst::ICMP_SGE;
return Pred == CmpInst::ICMP_SGE || Pred == CmpInst::ICMP_SGT;
default:
llvm_unreachable("unhandled recurrence kind");
}
}();

// TODO: Strict predicates need to find the first IV value for which the
// predicate holds, not the last.
if (Pred != RdxPredicate)
if (!IsValidPredicate)
return false;

assert(!FindIVPhiR->isInLoop() && !FindIVPhiR->isOrdered() &&
"cannot handle inloop/ordered reductions yet");
// For strict predicates, use a UMin reduction to find the minimum index.
// Canonical IVs (0, 1, 2, ...) are guaranteed not to wrap in the vector
// loop, so UMin can always be used.
bool IsStrictPredicate = ICmpInst::isLT(Pred) || ICmpInst::isGT(Pred);
if (IsStrictPredicate) {
return handleStrictArgMinArgMax(Plan, MinMaxPhiR, FindIVPhiR,
cast<VPWidenIntOrFpInductionRecipe>(IVOp),
MinMaxResult);
}

// The reduction using MinMaxPhiR needs adjusting to compute the correct
// result:
Expand Down
2 changes: 2 additions & 0 deletions llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -163,6 +163,7 @@ bool VPRecipeBase::mayHaveSideEffects() const {
return cast<VPExpressionRecipe>(this)->mayHaveSideEffects();
case VPDerivedIVSC:
case VPFirstOrderRecurrencePHISC:
case VPReductionPHISC:
case VPPredInstPHISC:
case VPVectorEndPointerSC:
return false;
Expand Down Expand Up @@ -1189,6 +1190,7 @@ bool VPInstruction::opcodeMayReadOrWriteFromMemory() const {
case VPInstruction::BuildVector:
case VPInstruction::CalculateTripCountMinusVF:
case VPInstruction::CanonicalIVIncrementForPart:
case VPInstruction::ComputeFindIVResult:
case VPInstruction::ExtractLane:
case VPInstruction::ExtractLastLane:
case VPInstruction::ExtractLastPart:
Expand Down
4 changes: 2 additions & 2 deletions llvm/lib/Transforms/Vectorize/VPlanTransforms.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -146,8 +146,8 @@ struct VPlanTransforms {
const TargetLibraryInfo &TLI);

/// Try to legalize reductions with multiple in-loop uses. Currently only
/// min/max reductions used by FindLastIV reductions are supported. Otherwise
/// return false.
/// min/max reductions used by FindLastIV and FindFirstIV reductions are
/// supported. Otherwise return false.
static bool handleMultiUseReductions(VPlan &Plan);

/// Try to have all users of fixed-order recurrences appear after the recipe
Expand Down
Loading
Loading