Skip to content

[LoopVectorize] Vectorize the reduction pattern of integer min/max with index. (1/2) #141467

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.

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
wants to merge 1 commit into
base: main
Choose a base branch
from

Conversation

Mel-Chen
Copy link
Contributor

Consider the following loop:

int idx = idx_start;
int max = max_start;
for (int i = 0; i < n; ++i) {
  int x = a[i];
  if (max < x) {
    max = x;
    idx = i;
  }
}

There are two recurrences involve this idiom -- min/max recurrence and index reduction. This patch performs two-stage reduction detection for this idiom. First, min/max recurrences and FindLastIV reductions are identified separately. A new RecurrenceDescriptor::isMinMaxRecurrencePHI function is introduced to detect min/max recurrences, instead of reusing the general reduction detection mechanism. This decouples min/max recurrences from reduction constraints—like, external loop users of recurrence are no longer required—allowing cases where the min/max reduction result is not live out the loop. In the second stage, a check ensures that all min/max recurrences are involved and exclusively participate in min/max with index patterns. Based on this analysis, the index reduction is classified as either min/max with first index or min/max with last index. New recurrence kinds, MinMaxFirstIdx and MinMaxLastIdx for this pattern. Those kinds are not directly generated by function AddReductionVar, but converted from FindLastIV recurrence by function RecurrenceDescriptor::isMinMaxIdxReduction.

TODOs:

  • Support the min/max recurrence in select(cmp()). Refer to test case smax_idx_select_cmp.
  • Support FP min/max recurrence.
  • Support intermediate store.
  • Support type-promoted recurrence.
  • Support min/max with 2-D indexes.

…th index. (1/2)

Consider the following loop:
```
int idx = idx_start;
int max = max_start;
for (int i = 0; i < n; ++i) {
  int x = a[i];
  if (max < x) {
    max = x;
    idx = i;
  }
}
```

There are two recurrences involve this idiom -- min/max recurrence and index reduction.
This patch performs two-stage reduction detection for this idiom.
First, min/max recurrences and FindLastIV reductions are identified separately. A new RecurrenceDescriptor::isMinMaxRecurrencePHI function is introduced to detect min/max recurrences, instead of reusing the general reduction detection mechanism. This decouples min/max recurrences from reduction constraints—like, external loop users of recurrence are no longer required—allowing cases where the min/max reduction result is not live out the loop.
In the second stage, a check ensures that all min/max recurrences are involved and exclusively participate in min/max with index patterns. Based on this analysis, the index reduction is classified as either min/max with first index or min/max with last index.
New recurrence kinds, MinMaxFirstIdx and MinMaxLastIdx for this pattern. Those kinds are not directly generated by function AddReductionVar, but converted from FindLastIV recurrence by function RecurrenceDescriptor::isMinMaxIdxReduction.

TODOs:
* Support the min/max recurrence in select(cmp()). Refer to test case smax_idx_select_cmp.
* Support FP min/max recurrence.
* Support intermediate store.
* Support type-promoted recurrence.
* Support min/max with 2-D indexes.
@llvmbot
Copy link
Member

llvmbot commented May 26, 2025

@llvm/pr-subscribers-llvm-transforms
@llvm/pr-subscribers-llvm-analysis

@llvm/pr-subscribers-vectorizers

Author: Mel Chen (Mel-Chen)

Changes

Consider the following loop:

int idx = idx_start;
int max = max_start;
for (int i = 0; i &lt; n; ++i) {
  int x = a[i];
  if (max &lt; x) {
    max = x;
    idx = i;
  }
}

There are two recurrences involve this idiom -- min/max recurrence and index reduction. This patch performs two-stage reduction detection for this idiom. First, min/max recurrences and FindLastIV reductions are identified separately. A new RecurrenceDescriptor::isMinMaxRecurrencePHI function is introduced to detect min/max recurrences, instead of reusing the general reduction detection mechanism. This decouples min/max recurrences from reduction constraints—like, external loop users of recurrence are no longer required—allowing cases where the min/max reduction result is not live out the loop. In the second stage, a check ensures that all min/max recurrences are involved and exclusively participate in min/max with index patterns. Based on this analysis, the index reduction is classified as either min/max with first index or min/max with last index. New recurrence kinds, MinMaxFirstIdx and MinMaxLastIdx for this pattern. Those kinds are not directly generated by function AddReductionVar, but converted from FindLastIV recurrence by function RecurrenceDescriptor::isMinMaxIdxReduction.

TODOs:

  • Support the min/max recurrence in select(cmp()). Refer to test case smax_idx_select_cmp.
  • Support FP min/max recurrence.
  • Support intermediate store.
  • Support type-promoted recurrence.
  • Support min/max with 2-D indexes.

Patch is 23.18 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/141467.diff

6 Files Affected:

  • (modified) llvm/include/llvm/Analysis/IVDescriptors.h (+36)
  • (modified) llvm/include/llvm/Transforms/Vectorize/LoopVectorizationLegality.h (+14)
  • (modified) llvm/lib/Analysis/IVDescriptors.cpp (+222)
  • (modified) llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp (+85)
  • (modified) llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp (+6)
  • (modified) llvm/test/Transforms/LoopVectorize/smax-idx.ll (+35-8)
diff --git a/llvm/include/llvm/Analysis/IVDescriptors.h b/llvm/include/llvm/Analysis/IVDescriptors.h
index d94ffa7287db3..971f7564a3430 100644
--- a/llvm/include/llvm/Analysis/IVDescriptors.h
+++ b/llvm/include/llvm/Analysis/IVDescriptors.h
@@ -56,6 +56,8 @@ enum class RecurKind {
   FindLastIV, ///< FindLast reduction with select(cmp(),x,y) where one of
               ///< (x,y) is increasing loop induction, and both x and y are
               ///< integer type.
+  MinMaxFirstIdx, ///< Integer Min/Max with first index
+  MinMaxLastIdx,  ///< Integer Min/Max with last index
   // clang-format on
   // TODO: Any_of and FindLast reduction need not be restricted to integer type
   // only.
@@ -206,6 +208,26 @@ class RecurrenceDescriptor {
   static bool isFixedOrderRecurrence(PHINode *Phi, Loop *TheLoop,
                                      DominatorTree *DT);
 
+  /// Returns the recurrence chain if \p Phi is an integer min/max recurrence in
+  /// \p TheLoop. The RecurrenceDescriptor is returned in \p RecurDes.
+  static SmallVector<Instruction *, 2>
+  tryToGetMinMaxRecurrenceChain(PHINode *Phi, Loop *TheLoop,
+                                RecurrenceDescriptor &RecurDes);
+
+  /// Returns true if the recurrence is a min/max with index pattern, and
+  /// updates the recurrence kind to RecurKind::MinMaxFirstIdx or
+  /// RecurKind::MinMaxLastIdx.
+  ///
+  /// \param IdxPhi         The phi representing the index recurrence.
+  /// \param MinMaxPhi      The phi representing the min/max recurrence involved
+  ///                       in the min/max with index pattern.
+  /// \param MinMaxDesc     The descriptor of the min/max recurrence.
+  /// \param MinMaxChain    The chain of instructions involved in the min/max
+  ///                       recurrence.
+  bool isMinMaxIdxReduction(PHINode *IdxPhi, PHINode *MinMaxPhi,
+                            const RecurrenceDescriptor &MinMaxDesc,
+                            ArrayRef<Instruction *> MinMaxChain);
+
   RecurKind getRecurrenceKind() const { return Kind; }
 
   unsigned getOpcode() const { return getOpcode(getRecurrenceKind()); }
@@ -259,6 +281,20 @@ class RecurrenceDescriptor {
     return Kind == RecurKind::FindLastIV;
   }
 
+  /// Returns true if the recurrence kind is of the form:
+  ///   select(icmp(a,b),x,y)
+  /// where one of (x,y) is an increasing loop induction variable, and icmp(a,b)
+  /// depends on a min/max recurrence.
+  static bool isMinMaxIdxRecurrenceKind(RecurKind Kind) {
+    return Kind == RecurKind::MinMaxFirstIdx ||
+           Kind == RecurKind::MinMaxLastIdx;
+  }
+
+  /// Returns true if the recurrence kind is an integer max kind.
+  static bool isIntMaxRecurrenceKind(RecurKind Kind) {
+    return Kind == RecurKind::UMax || Kind == RecurKind::SMax;
+  }
+
   /// 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; }
diff --git a/llvm/include/llvm/Transforms/Vectorize/LoopVectorizationLegality.h b/llvm/include/llvm/Transforms/Vectorize/LoopVectorizationLegality.h
index d654ac3ec9273..3f82a81dd99c0 100644
--- a/llvm/include/llvm/Transforms/Vectorize/LoopVectorizationLegality.h
+++ b/llvm/include/llvm/Transforms/Vectorize/LoopVectorizationLegality.h
@@ -345,6 +345,9 @@ class LoopVectorizationLegality {
   /// Returns True if Phi is a fixed-order recurrence in this loop.
   bool isFixedOrderRecurrence(const PHINode *Phi) const;
 
+  /// Returns True if \p Phi is a min/max recurrence in this loop.
+  bool isMinMaxRecurrence(const PHINode *Phi) const;
+
   /// Return true if the block BB needs to be predicated in order for the loop
   /// to be vectorized.
   bool blockNeedsPredication(BasicBlock *BB) const;
@@ -519,6 +522,14 @@ class LoopVectorizationLegality {
   /// specific checks for outer loop vectorization.
   bool canVectorizeOuterLoop();
 
+  // Min/max recurrences can only be vectorized when involved in a min/max with
+  // index reduction pattern. This function checks whether the \p Phi, which
+  // represents the min/max recurrence, can be vectorized based on the given \p
+  // Chain, which is the recurrence chain for the min/max recurrence. Returns
+  // true if the min/max recurrence can be vectorized.
+  bool canVectorizeMinMaxRecurrence(PHINode *Phi,
+                                    ArrayRef<Instruction *> Chain);
+
   /// Returns true if this is an early exit loop that can be vectorized.
   /// Currently, a loop with an uncountable early exit is considered
   /// vectorizable if:
@@ -606,6 +617,9 @@ class LoopVectorizationLegality {
   /// Holds the phi nodes that are fixed-order recurrences.
   RecurrenceSet FixedOrderRecurrences;
 
+  /// Holds the min/max recurrences variables.
+  RecurrenceSet MinMaxRecurrences;
+
   /// Holds the widest induction type encountered.
   IntegerType *WidestIndTy = nullptr;
 
diff --git a/llvm/lib/Analysis/IVDescriptors.cpp b/llvm/lib/Analysis/IVDescriptors.cpp
index b7c7bcab168cc..43948c1de30c1 100644
--- a/llvm/lib/Analysis/IVDescriptors.cpp
+++ b/llvm/lib/Analysis/IVDescriptors.cpp
@@ -51,6 +51,8 @@ bool RecurrenceDescriptor::isIntegerRecurrenceKind(RecurKind Kind) {
   case RecurKind::UMin:
   case RecurKind::AnyOf:
   case RecurKind::FindLastIV:
+  case RecurKind::MinMaxFirstIdx:
+  case RecurKind::MinMaxLastIdx:
     return true;
   }
   return false;
@@ -1130,6 +1132,226 @@ bool RecurrenceDescriptor::isFixedOrderRecurrence(PHINode *Phi, Loop *TheLoop,
   return true;
 }
 
+/// Return the recurrence kind if \p I is matched by the min/max operation
+/// pattern. Otherwise, return RecurKind::None.
+static RecurKind isMinMaxRecurOp(const Instruction *I) {
+  if (match(I, m_UMin(m_Value(), m_Value())))
+    return RecurKind::UMin;
+  if (match(I, m_UMax(m_Value(), m_Value())))
+    return RecurKind::UMax;
+  if (match(I, m_SMax(m_Value(), m_Value())))
+    return RecurKind::SMax;
+  if (match(I, m_SMin(m_Value(), m_Value())))
+    return RecurKind::SMin;
+  // TODO: support fp-min/max
+  return RecurKind::None;
+}
+
+SmallVector<Instruction *, 2>
+RecurrenceDescriptor::tryToGetMinMaxRecurrenceChain(
+    PHINode *Phi, Loop *TheLoop, RecurrenceDescriptor &RecurDes) {
+  SmallVector<Instruction *, 2> Chain;
+  // Check the phi is in the loop header and has two incoming values.
+  if (Phi->getParent() != TheLoop->getHeader() ||
+      Phi->getNumIncomingValues() != 2)
+    return {};
+
+  // Ensure the loop has a preheader and a latch block.
+  auto *Preheader = TheLoop->getLoopPreheader();
+  auto *Latch = TheLoop->getLoopLatch();
+  if (!Preheader || !Latch)
+    return {};
+
+  // Ensure that one of the incoming values of the PHI node is from the
+  // preheader, and the other one is from the loop latch.
+  if (Phi->getBasicBlockIndex(Preheader) < 0 ||
+      Phi->getBasicBlockIndex(Latch) < 0)
+    return {};
+
+  Value *StartValue = Phi->getIncomingValueForBlock(Preheader);
+  auto *BEValue = dyn_cast<Instruction>(Phi->getIncomingValueForBlock(Latch));
+  if (!BEValue || BEValue == Phi)
+    return {};
+
+  auto HasLoopExternalUse = [TheLoop](const Instruction *I) {
+    return any_of(I->users(), [TheLoop](auto *U) {
+      return !TheLoop->contains(cast<Instruction>(U));
+    });
+  };
+
+  // Ensure the recurrence phi has no users outside the loop, as such cases
+  // cannot be vectorized.
+  if (HasLoopExternalUse(Phi))
+    return {};
+
+  // Ensure the backedge value of the phi is only used internally by the phi;
+  // all other users must be outside the loop.
+  // TODO: support intermediate store.
+  if (any_of(BEValue->users(), [&](auto *U) {
+        auto *UI = cast<Instruction>(U);
+        return TheLoop->contains(UI) && UI != Phi;
+      }))
+    return {};
+
+  // Ensure the backedge value of the phi matches the min/max operation pattern.
+  RecurKind TargetKind = isMinMaxRecurOp(BEValue);
+  if (TargetKind == RecurKind::None)
+    return {};
+
+  // TODO: type-promoted recurrence
+  SmallPtrSet<Instruction *, 4> CastInsts;
+
+  // Trace the use-def chain from the backedge value to the phi, ensuring a
+  // unique in-loop path where all operations match the expected recurrence
+  // kind.
+  bool FoundRecurPhi = false;
+  SmallVector<Instruction *, 8> Worklist(1, BEValue);
+  SmallDenseMap<Instruction *, Instruction *, 4> VisitedFrom;
+
+  VisitedFrom.try_emplace(BEValue);
+
+  while (!Worklist.empty()) {
+    Instruction *Cur = Worklist.pop_back_val();
+    if (Cur == Phi) {
+      if (FoundRecurPhi)
+        return {};
+      FoundRecurPhi = true;
+      continue;
+    }
+
+    if (!TheLoop->contains(Cur))
+      continue;
+
+    // TODO: support the min/max recurrence in cmp-select pattern.
+    if (!isa<CallInst>(Cur) || isMinMaxRecurOp(Cur) != TargetKind)
+      continue;
+
+    for (Use &Op : Cur->operands()) {
+      if (auto *OpInst = dyn_cast<Instruction>(Op)) {
+        if (!VisitedFrom.try_emplace(OpInst, Cur).second)
+          return {};
+        Worklist.push_back(OpInst);
+      }
+    }
+  }
+
+  if (!FoundRecurPhi)
+    return {};
+
+  Instruction *ExitInstruction = nullptr;
+  // Get the recurrence chain by visited trace.
+  Instruction *VisitedInst = VisitedFrom.at(Phi);
+  while (VisitedInst) {
+    // Ensure that no instruction in the recurrence chain is used outside the
+    // loop, except for the backedge value, which is permitted.
+    if (HasLoopExternalUse(VisitedInst)) {
+      if (VisitedInst != BEValue)
+        return {};
+      ExitInstruction = BEValue;
+    }
+    Chain.push_back(VisitedInst);
+    VisitedInst = VisitedFrom.at(VisitedInst);
+  }
+
+  RecurDes = RecurrenceDescriptor(
+      StartValue, ExitInstruction, /*IntermediateStore=*/nullptr, TargetKind,
+      FastMathFlags(), /*ExactFPMathInst=*/nullptr, Phi->getType(),
+      /*IsSigned=*/false, /*IsOrdered=*/false, CastInsts,
+      /*MinWidthCastToRecurTy=*/-1U);
+
+  LLVM_DEBUG(dbgs() << "Found a min/max recurrence PHI: " << *Phi << "\n");
+
+  return Chain;
+}
+
+bool RecurrenceDescriptor::isMinMaxIdxReduction(
+    PHINode *IdxPhi, PHINode *MinMaxPhi, const RecurrenceDescriptor &MinMaxDesc,
+    ArrayRef<Instruction *> MinMaxChain) {
+  // Return early if the recurrence kind is already known to be min/max with
+  // index.
+  if (isMinMaxIdxRecurrenceKind(Kind))
+    return true;
+
+  if (!isFindLastIVRecurrenceKind(Kind))
+    return false;
+
+  // Ensure index reduction phi and min/max recurrence phi are in the same basic
+  // block.
+  if (IdxPhi->getParent() != MinMaxPhi->getParent())
+    return false;
+
+  RecurKind MinMaxRK = MinMaxDesc.getRecurrenceKind();
+  // TODO: support floating-point min/max with index.
+  if (!isIntMinMaxRecurrenceKind(MinMaxRK))
+    return false;
+
+  // FindLastIV only supports a single select operation in the recurrence chain
+  // so far. Therefore, do not consider min/max recurrences with more than one
+  // operation in the recurrence chain.
+  // TODO: support FindLastIV with multiple operations in the recurrence chain.
+  if (MinMaxChain.size() != 1)
+    return false;
+
+  Instruction *MinMaxChainCur = MinMaxPhi;
+  Instruction *MinMaxChainNext = MinMaxChain.front();
+  Value *OutOfChain;
+  bool IsMinMaxOperation = match(
+      MinMaxChainNext,
+      m_CombineOr(m_MaxOrMin(m_Specific(MinMaxChainCur), m_Value(OutOfChain)),
+                  m_MaxOrMin(m_Value(OutOfChain), m_Specific(MinMaxChainCur))));
+  assert(IsMinMaxOperation && "Unexpected operation in the recurrence chain");
+
+  auto *IdxExit = cast<SelectInst>(LoopExitInstr);
+  Value *IdxCond = IdxExit->getCondition();
+  // Check if the operands used by cmp instruction of index select is the same
+  // as the operands used by min/max recurrence.
+  bool IsMatchLHSInMinMaxChain =
+      match(IdxCond, m_Cmp(m_Specific(MinMaxChainCur), m_Specific(OutOfChain)));
+  bool IsMatchRHSInMinMaxChain =
+      match(IdxCond, m_Cmp(m_Specific(OutOfChain), m_Specific(MinMaxChainCur)));
+  if (!IsMatchLHSInMinMaxChain && !IsMatchRHSInMinMaxChain)
+    return false;
+
+  CmpInst::Predicate IdxPred = cast<CmpInst>(IdxCond)->getPredicate();
+  // The predicate of cmp instruction must be relational in min/max with index.
+  if (CmpInst::isEquality(IdxPred))
+    return false;
+
+  // Normalize predicate from
+  //   m_Cmp(pred, out_of_chain, in_chain)
+  // to
+  //   m_Cmp(swapped_pred, in_chain, out_of_chain).
+  if (IsMatchRHSInMinMaxChain)
+    IdxPred = CmpInst::getSwappedPredicate(IdxPred);
+
+  // Verify that the select operation is updated on the correct side based on
+  // the min/max kind.
+  bool IsTrueUpdateIdx = IdxExit->getFalseValue() == IdxPhi;
+  bool IsMaxRK = isIntMaxRecurrenceKind(MinMaxRK);
+  bool IsLess = ICmpInst::isLT(IdxPred) || ICmpInst::isLE(IdxPred);
+  bool IsExpectedTrueUpdateIdx = IsMaxRK == IsLess;
+  if (IsTrueUpdateIdx != IsExpectedTrueUpdateIdx)
+    return false;
+
+  RecurKind NewIdxRK;
+  // The index recurrence kind is the same for both the predicate and its
+  // inverse.
+  if (!IsLess)
+    IdxPred = CmpInst::getInversePredicate(IdxPred);
+  // For max recurrence, a strict less-than predicate indicates that the first
+  // matching index will be selected. For min recurrence, the opposite holds.
+  NewIdxRK = IsMaxRK != ICmpInst::isLE(IdxPred) ? RecurKind::MinMaxFirstIdx
+                                                : RecurKind::MinMaxLastIdx;
+
+  // Update the kind of index recurrence.
+  Kind = NewIdxRK;
+  LLVM_DEBUG(
+      dbgs() << "Found a min/max with "
+             << (NewIdxRK == RecurKind::MinMaxFirstIdx ? "first" : "last")
+             << " index reduction PHI." << *IdxPhi << "\n");
+  return true;
+}
+
 unsigned RecurrenceDescriptor::getOpcode(RecurKind Kind) {
   switch (Kind) {
   case RecurKind::Add:
diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp
index 8e09e6f8d4935..25894ac14df33 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp
@@ -795,6 +795,10 @@ static bool canWidenCallReturnType(Type *Ty) {
 bool LoopVectorizationLegality::canVectorizeInstrs() {
   BasicBlock *Header = TheLoop->getHeader();
 
+  // Tracks the operation chain for each min/max recurrence phi that is
+  // considered vectorizable.
+  SmallDenseMap<PHINode *, SmallVector<Instruction *>> MinMaxRecurrenceChains;
+
   // For each block in the loop.
   for (BasicBlock *BB : TheLoop->blocks()) {
     // Scan the instructions in the block and look for hazards.
@@ -840,6 +844,18 @@ bool LoopVectorizationLegality::canVectorizeInstrs() {
           continue;
         }
 
+        RecurrenceDescriptor MinMaxRecurDes;
+        if (auto Chain = RecurrenceDescriptor::tryToGetMinMaxRecurrenceChain(
+                Phi, TheLoop, MinMaxRecurDes);
+            !Chain.empty()) {
+          if (MinMaxRecurDes.getLoopExitInstr())
+            AllowedExit.insert(MinMaxRecurDes.getLoopExitInstr());
+          Reductions[Phi] = MinMaxRecurDes;
+          MinMaxRecurrences.insert(Phi);
+          MinMaxRecurrenceChains[Phi] = std::move(Chain);
+          continue;
+        }
+
         // We prevent matching non-constant strided pointer IVS to preserve
         // historical vectorizer behavior after a generalization of the
         // IVDescriptor code.  The intent is to remove this check, but we
@@ -1069,9 +1085,74 @@ bool LoopVectorizationLegality::canVectorizeInstrs() {
   if (PrimaryInduction && WidestIndTy != PrimaryInduction->getType())
     PrimaryInduction = nullptr;
 
+  // The second stage check for reduction. Confirm if the min/max with index
+  // reduction, involving two PHIs, is legal to vectorize.
+  for (auto &Entry : MinMaxRecurrenceChains) {
+    PHINode *Phi = Entry.first;
+    ArrayRef<Instruction *> Chain = Entry.second;
+    if (!canVectorizeMinMaxRecurrence(Phi, Chain))
+      return false;
+  }
+  // FIXME: Remove this after the IR generation of min/max with index is
+  // supported.
+  if (!MinMaxRecurrences.empty())
+    return false;
+
   return true;
 }
 
+bool LoopVectorizationLegality::canVectorizeMinMaxRecurrence(
+    PHINode *Phi, ArrayRef<Instruction *> Chain) {
+  assert(!Chain.empty() && "Unexpected empty recurrence chain");
+  assert(isMinMaxRecurrence(Phi) && "The PHI is not a min/max recurrence phi");
+
+  auto IsMinMaxIdxReductionPhi = [this, Phi, &Chain](Value *Candidate) -> bool {
+    auto *IdxPhi = dyn_cast<PHINode>(Candidate);
+    if (!IdxPhi || !isReductionVariable(IdxPhi))
+      return false;
+
+    RecurrenceDescriptor &IdxRdxDesc = Reductions.find(IdxPhi)->second;
+    const RecurrenceDescriptor &MinMaxDesc = Reductions.find(Phi)->second;
+    return IdxRdxDesc.isMinMaxIdxReduction(IdxPhi, Phi, MinMaxDesc, Chain);
+  };
+
+  // Find the potential index recurrence chain head.
+  // Note: Only one chain head can be found since 2-D indexes are not yet
+  // supported.
+  SelectInst *IdxChainHead = nullptr;
+  // TODO: support min/max with 2-D indexes.
+  if (!Phi->hasNUses(2))
+    return false;
+
+  for (User *U : Phi->users()) {
+    if (auto *Cmp = dyn_cast<CmpInst>(U)) {
+      if (!Cmp->hasOneUse())
+        return false;
+      if (!match(Cmp->user_back(),
+                 m_Select(m_Specific(Cmp), m_Value(), m_Value())))
+        return false;
+      assert(!IdxChainHead &&
+             "Unexpected multiple index recurrence chain head");
+      IdxChainHead = cast<SelectInst>(Cmp->user_back());
+      continue;
+    }
+
+    // Skip the user in the min/max recurrence chain
+    if (llvm::is_contained(Chain, cast<Instruction>(U)))
+      continue;
+
+    // Unexpected user
+    return false;
+  }
+
+  if (!IdxChainHead)
+    return false;
+
+  auto *TrueVal = IdxChainHead->getTrueValue();
+  auto *FalseVal = IdxChainHead->getFalseValue();
+  return IsMinMaxIdxReductionPhi(TrueVal) || IsMinMaxIdxReductionPhi(FalseVal);
+}
+
 /// Find histogram operations that match high-level code in loops:
 /// \code
 /// buckets[indices[i]]+=step;
@@ -1394,6 +1475,10 @@ bool LoopVectorizationLegality::isFixedOrderRecurrence(
   return FixedOrderRecurrences.count(Phi);
 }
 
+bool LoopVectorizationLegality::isMinMaxRecurrence(const PHINode *Phi) const {
+  return MinMaxRecurrences.contains(Phi);
+}
+
 bool LoopVectorizationLegality::blockNeedsPredication(BasicBlock *BB) const {
   // When vectorizing early exits, create predicates for the latch block only.
   // The early exiting block must be a direct predecessor of the latch at the
diff --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
index 831703b375d9b..e15562996fb73 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -23092,6 +23092,8 @@ class HorizontalReduction {
         case RecurKind::FMulAdd:
         case RecurKind::AnyOf:
         case RecurKind::FindLastIV:
+        case RecurKind::MinMaxFirstIdx:
+        case RecurKind::MinMaxLastIdx:
         case RecurKind::FMaximumNum:
         case RecurKind::FMinimumNum:
         case RecurKind::None:
@@ -23226,6 +23228,8 @@ class HorizontalReduction {
     case RecurKind::FMulAdd:
     case RecurKind::AnyOf:
     case RecurKind::FindLastIV:
+    case RecurKind::MinMaxFirstIdx:
+    case RecurKind::MinMaxLastIdx:
     case RecurKind::FMaximumNum:
     case RecurKind::FMinimumNum:
     case RecurKind::None:
@@ -23325,6 +23329,8 @@ class HorizontalReduction {
     case RecurKind::FMulAdd:
     case RecurKind::AnyOf:
     case RecurKind::FindLastIV:
+    case RecurKind::MinMaxFirstIdx:
+    case RecurKind::MinMaxLastIdx:
     case RecurKind::FMaximumNum:
     case RecurKind::FMinimumNum:
     case RecurKind::None:
diff --git a/llvm/test/Transforms/LoopVectorize/smax-idx.ll b/llvm/test/Transforms/LoopVectorize/smax-idx.ll
index 37dcd7fc7e39f..ce29818d05913 100644
--- a/llvm/test/Transforms/LoopVectorize/smax-idx.ll
+++ b/llvm/test/Transforms/LoopVectorize/smax-idx.ll
@@ -1,6 +1,38 @@
-; RUN: opt -passes=loop-vectorize -force-vector-width=4 -force-vector-interleave=1 -S < %s | FileCheck %s --check-prefix=CHECK
-; RUN: opt -passes=loop-vectorize -force-vector-width=4 -force-vector-interleave=4 -S < %s | FileCheck %s --check-prefix=CHECK
-; RUN: opt -passes=loop-vectorize -force-vector-width=1 -force-vector-interleave=4 -S < %s | FileCheck %s --check-prefix=CHECK
...
[truncated]

@Mel-Chen
Copy link
Contributor Author

Mel-Chen commented Jun 6, 2025

Ping

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants