[LV] Support argmin/argmax with strict predicates.

Extend handleMultiUseReductions to support strict predicates (>, <),
matching the first index instead of the last for non-strict predicates.

For strict predicates is detected, the transformation converts the
FindLastIV reduction to FindFirstIV by:
1. Checking the IV range to ensure it does not include the sentinel value (max).
2. Creating a new reduction with the appropriate FindFirstIV kind
   (FindFirstIVSMin or FindFirstIVUMin based on the IV range)
3. Replacing the old reduction recipe with the new one

Author: fhahn

llvm/include/llvm/Analysis/IVDescriptors.h

@@ -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.
@@ -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.
+  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; }

llvm/lib/Analysis/IVDescriptors.cpp

@@ -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++) {
@@ -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;
   };

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -8506,8 +8506,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, *PSE.getSE(), OrigLoop))
     return nullptr;
   // Apply mandatory transformation to handle FP maxnum/minnum reduction with
   // NaNs if possible, bail out otherwise.

llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp

@@ -13,10 +13,12 @@
 
 #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"
@@ -1120,7 +1122,48 @@ bool VPlanTransforms::handleMaxMinNumReductions(VPlan &Plan) {
   return true;
 }
 
-bool VPlanTransforms::handleMultiUseReductions(VPlan &Plan) {
+/// Try to convert FindLastIV to FindFirstIV reduction when using a strict
+/// predicate. Returns the new FindFirstIVPhiR on success, nullptr on failure.
+static VPReductionPHIRecipe *
+tryConvertToFindFirstIV(VPlan &Plan, VPReductionPHIRecipe *FindLastIVPhiR,
+                        VPValue *IVOp, ScalarEvolution &SE, const Loop *L) {
+  Type *Ty = VPTypeAnalysis(Plan).inferScalarType(FindLastIVPhiR);
+  unsigned NumBits = Ty->getIntegerBitWidth();
+
+  // Determine the reduction kind and sentinel based on the IV range.
+  RecurKind NewKind;
+  VPValue *NewSentinel;
+  auto *AR = cast<SCEVAddRecExpr>(vputils::getSCEVExprForVPValue(IVOp, SE, L));
+  if (RecurrenceDescriptor::isValidIVRangeForFindIV(
+          AR, /*IsSigned=*/true, /*IsFindFirstIV=*/true, SE)) {
+    NewKind = RecurKind::FindFirstIVSMin;
+    NewSentinel = Plan.getConstantInt(APInt::getSignedMaxValue(NumBits));
+  } else if (RecurrenceDescriptor::isValidIVRangeForFindIV(
+                 AR, /*IsSigned=*/false, /*IsFindFirstIV=*/true, SE)) {
+    NewKind = RecurKind::FindFirstIVUMin;
+    NewSentinel = Plan.getConstantInt(APInt::getMaxValue(NumBits));
+  } else {
+    return nullptr;
+  }
+
+  // Create the new FindFirstIV reduction recipe.
+  assert(!FindLastIVPhiR->isInLoop() && !FindLastIVPhiR->isOrdered());
+  ReductionStyle Style = RdxUnordered{FindLastIVPhiR->getVFScaleFactor()};
+  auto *FindFirstIVPhiR =
+      new VPReductionPHIRecipe(nullptr, NewKind, *NewSentinel, Style,
+                               FindLastIVPhiR->hasUsesOutsideReductionChain());
+  FindFirstIVPhiR->addOperand(FindLastIVPhiR->getBackedgeValue());
+
+  FindFirstIVPhiR->insertBefore(FindLastIVPhiR);
+  VPInstruction *FindLastIVResult =
+      findUserOf<VPInstruction::ComputeFindIVResult>(FindLastIVPhiR);
+  FindLastIVPhiR->replaceAllUsesWith(FindFirstIVPhiR);
+  FindLastIVResult->setOperand(2, NewSentinel);
+  return FindFirstIVPhiR;
+}
+
+bool VPlanTransforms::handleMultiUseReductions(VPlan &Plan, ScalarEvolution &SE,
+                                               const Loop *L) {
   for (auto &PhiR : make_early_inc_range(
            Plan.getVectorLoopRegion()->getEntryBasicBlock()->phis())) {
     auto *MinMaxPhiR = dyn_cast<VPReductionPHIRecipe>(&PhiR);
@@ -1203,33 +1246,41 @@ 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, transform try to convert FindLastIV to
+    // FindFirstIV.
+    bool IsStrictPredicate = ICmpInst::isLT(Pred) || ICmpInst::isGT(Pred);
+    if (IsStrictPredicate) {
+      FindIVPhiR = tryConvertToFindFirstIV(Plan, FindIVPhiR, IVOp, SE, L);
+      if (!FindIVPhiR)
+        return false;
+    }
 
     // The reduction using MinMaxPhiR needs adjusting to compute the correct
     // result:

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

@@ -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;

llvm/lib/Transforms/Vectorize/VPlanTransforms.h

@@ -157,9 +157,10 @@ 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.
-  static bool handleMultiUseReductions(VPlan &Plan);
+  /// min/max reductions used by FindLastIV and FindFirstIV reductions are
+  /// supported. Otherwise return false.
+  static bool handleMultiUseReductions(VPlan &Plan, ScalarEvolution &SE,
+                                       const Loop *L);
 
   /// Try to have all users of fixed-order recurrences appear after the recipe
   /// defining their previous value, by either sinking users or hoisting recipes