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

Mel-Chen · Mel-Chen · commit f5c4c1499ba6 · 2025-06-09T19:53:18.000-07:00
diff --git a/llvm/include/llvm/Analysis/IVDescriptors.h b/llvm/include/llvm/Analysis/IVDescriptors.h
@@ -305,7 +305,9 @@ class RecurrenceDescriptor {
   /// Returns the sentinel value for FindLastIV recurrences to replace the start
   /// value.
   Value *getSentinelValue() const {
-    assert(isFindLastIVRecurrenceKind(Kind) && "Unexpected recurrence kind");
+    assert(
+        (isFindLastIVRecurrenceKind(Kind) || isMinMaxIdxRecurrenceKind(Kind)) &&
+        "Unexpected recurrence kind");
     Type *Ty = StartValue->getType();
     return ConstantInt::get(Ty,
                             APInt::getSignedMinValue(Ty->getIntegerBitWidth()));
diff --git a/llvm/include/llvm/Transforms/Utils/LoopUtils.h b/llvm/include/llvm/Transforms/Utils/LoopUtils.h
@@ -423,6 +423,12 @@ Value *createAnyOfReduction(IRBuilderBase &B, Value *Src, Value *InitVal,
 Value *createFindLastIVReduction(IRBuilderBase &B, Value *Src, Value *Start,
                                  Value *Sentinel);
 
+/// Create a reduction of the given vector \p Src for a reduction of the
+/// kind RecurKind::MinMaxFirstIdx or RecurKind::MinMaxLastIdx. The reduction
+/// operation is described by \p Desc.
+Value *createMinMaxIdxReduction(IRBuilderBase &B, Value *Src, Value *Start,
+                                const RecurrenceDescriptor &Desc);
+
 /// Create an ordered reduction intrinsic using the given recurrence
 /// kind \p RdxKind.
 Value *createOrderedReduction(IRBuilderBase &B, RecurKind RdxKind, Value *Src,
diff --git a/llvm/include/llvm/Transforms/Vectorize/LoopVectorizationLegality.h b/llvm/include/llvm/Transforms/Vectorize/LoopVectorizationLegality.h
@@ -307,6 +307,11 @@ class LoopVectorizationLegality {
   /// Return the fixed-order recurrences found in the loop.
   RecurrenceSet &getFixedOrderRecurrences() { return FixedOrderRecurrences; }
 
+  /// Return the min/max recurrences found in the loop.
+  const SmallDenseMap<PHINode *, PHINode *> &getMinMaxRecurrences() {
+    return MinMaxRecurrences;
+  }
+
   /// Returns the widest induction type.
   IntegerType *getWidestInductionType() { return WidestIndTy; }
 
@@ -618,7 +623,7 @@ class LoopVectorizationLegality {
   RecurrenceSet FixedOrderRecurrences;
 
   /// Holds the min/max recurrences variables.
-  RecurrenceSet MinMaxRecurrences;
+  SmallDenseMap<PHINode *, PHINode *> MinMaxRecurrences;
 
   /// Holds the widest induction type encountered.
   IntegerType *WidestIndTy = nullptr;
diff --git a/llvm/lib/Transforms/Utils/LoopUtils.cpp b/llvm/lib/Transforms/Utils/LoopUtils.cpp
@@ -1248,6 +1248,25 @@ Value *llvm::createFindLastIVReduction(IRBuilderBase &Builder, Value *Src,
   return Builder.CreateSelect(Cmp, MaxRdx, Start, "rdx.select");
 }
 
+Value *llvm::createMinMaxIdxReduction(IRBuilderBase &Builder, Value *Src,
+                                      Value *Start,
+                                      const RecurrenceDescriptor &Desc) {
+  RecurKind Kind = Desc.getRecurrenceKind();
+  assert(RecurrenceDescriptor::isMinMaxIdxRecurrenceKind(Kind) &&
+         "Unexpected reduction kind");
+  Value *Sentinel = Desc.getSentinelValue();
+  Value *Rdx = Src;
+  if (Src->getType()->isVectorTy())
+    Rdx = Kind == RecurKind::MinMaxFirstIdx
+              ? Builder.CreateIntMinReduce(Src, true)
+              : Builder.CreateIntMaxReduce(Src, true);
+  // Correct the final reduction result back to the start value if the reduction
+  // result is sentinel value.
+  Value *Cmp =
+      Builder.CreateCmp(CmpInst::ICMP_NE, Rdx, Sentinel, "rdx.select.cmp");
+  return Builder.CreateSelect(Cmp, Rdx, Start, "rdx.select");
+}
+
 Value *llvm::getReductionIdentity(Intrinsic::ID RdxID, Type *Ty,
                                   FastMathFlags Flags) {
   bool Negative = false;
@@ -1336,7 +1355,8 @@ Value *llvm::createSimpleReduction(VectorBuilder &VBuilder, Value *Src,
                                    RecurKind Kind) {
   assert(!RecurrenceDescriptor::isAnyOfRecurrenceKind(Kind) &&
          !RecurrenceDescriptor::isFindLastIVRecurrenceKind(Kind) &&
-         "AnyOf or FindLastIV reductions are not supported.");
+         !RecurrenceDescriptor::isMinMaxIdxRecurrenceKind(Kind) &&
+         "AnyOf, FindLastIV and MinMaxIdx reductions are not supported.");
   Intrinsic::ID Id = getReductionIntrinsicID(Kind);
   auto *SrcTy = cast<VectorType>(Src->getType());
   Type *SrcEltTy = SrcTy->getElementType();
diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp
@@ -851,7 +851,7 @@ bool LoopVectorizationLegality::canVectorizeInstrs() {
           if (MinMaxRecurDes.getLoopExitInstr())
             AllowedExit.insert(MinMaxRecurDes.getLoopExitInstr());
           Reductions[Phi] = MinMaxRecurDes;
-          MinMaxRecurrences.insert(Phi);
+          MinMaxRecurrences.try_emplace(Phi);
           MinMaxRecurrenceChains[Phi] = std::move(Chain);
           continue;
         }
@@ -1093,10 +1093,6 @@ bool LoopVectorizationLegality::canVectorizeInstrs() {
     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;
 }
@@ -1106,6 +1102,10 @@ bool LoopVectorizationLegality::canVectorizeMinMaxRecurrence(
   assert(!Chain.empty() && "Unexpected empty recurrence chain");
   assert(isMinMaxRecurrence(Phi) && "The PHI is not a min/max recurrence phi");
 
+  auto It = MinMaxRecurrences.find(Phi);
+  if (It->second)
+    return true;
+
   auto IsMinMaxIdxReductionPhi = [this, Phi, &Chain](Value *Candidate) -> bool {
     auto *IdxPhi = dyn_cast<PHINode>(Candidate);
     if (!IdxPhi || !isReductionVariable(IdxPhi))
@@ -1150,7 +1150,17 @@ bool LoopVectorizationLegality::canVectorizeMinMaxRecurrence(
 
   auto *TrueVal = IdxChainHead->getTrueValue();
   auto *FalseVal = IdxChainHead->getFalseValue();
-  return IsMinMaxIdxReductionPhi(TrueVal) || IsMinMaxIdxReductionPhi(FalseVal);
+  PHINode *IdxPhi;
+  if (IsMinMaxIdxReductionPhi(TrueVal))
+    IdxPhi = cast<PHINode>(TrueVal);
+  else if (IsMinMaxIdxReductionPhi(FalseVal))
+    IdxPhi = cast<PHINode>(FalseVal);
+  else
+    return false;
+
+  // Record the index reduction phi uses the min/max recurrence.
+  It->second = IdxPhi;
+  return true;
 }
 
 /// Find histogram operations that match high-level code in loops:
@@ -1973,7 +1983,8 @@ bool LoopVectorizationLegality::canFoldTailByMasking() const {
   SmallPtrSet<const Value *, 8> ReductionLiveOuts;
 
   for (const auto &Reduction : getReductionVars())
-    ReductionLiveOuts.insert(Reduction.second.getLoopExitInstr());
+    if (auto *ExitInstr = Reduction.second.getLoopExitInstr())
+      ReductionLiveOuts.insert(ExitInstr);
 
   // TODO: handle non-reduction outside users when tail is folded by masking.
   for (auto *AE : AllowedExit) {
diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -4457,6 +4457,14 @@ bool LoopVectorizationPlanner::isCandidateForEpilogueVectorization(
         return false;
   }
 
+  // TODO: support epilogue vectorization for min/max with index.
+  if (any_of(Legal->getReductionVars(), [](const auto &Reduction) {
+        const RecurrenceDescriptor &RdxDesc = Reduction.second;
+        return RecurrenceDescriptor::isMinMaxIdxRecurrenceKind(
+            RdxDesc.getRecurrenceKind());
+      }))
+    return false;
+
   // Epilogue vectorization code has not been auditted to ensure it handles
   // non-latch exits properly.  It may be fine, but it needs auditted and
   // tested.
@@ -4901,7 +4909,8 @@ LoopVectorizationCostModel::selectInterleaveCount(VPlan &Plan, ElementCount VF,
           const RecurrenceDescriptor &RdxDesc = Reduction.second;
           RecurKind RK = RdxDesc.getRecurrenceKind();
           return RecurrenceDescriptor::isAnyOfRecurrenceKind(RK) ||
-                 RecurrenceDescriptor::isFindLastIVRecurrenceKind(RK);
+                 RecurrenceDescriptor::isFindLastIVRecurrenceKind(RK) ||
+                 RecurrenceDescriptor::isMinMaxIdxRecurrenceKind(RK);
         });
     if (HasSelectCmpReductions) {
       LLVM_DEBUG(dbgs() << "LV: Not interleaving select-cmp reductions.\n");
@@ -6618,6 +6627,10 @@ void LoopVectorizationCostModel::collectInLoopReductions() {
 
   for (const auto &Reduction : Legal->getReductionVars()) {
     PHINode *Phi = Reduction.first;
+    // TODO: support in-loop min/max with index.
+    if (Legal->isMinMaxRecurrence(Phi))
+      continue;
+
     const RecurrenceDescriptor &RdxDesc = Reduction.second;
 
     // We don't collect reductions that are type promoted (yet).
@@ -7233,6 +7246,8 @@ static void fixReductionScalarResumeWhenVectorizingEpilog(
        EpiRedResult->getOpcode() != VPInstruction::ComputeFindLastIVResult))
     return;
 
+  assert(EpiRedResult->getOpcode() != VPInstruction::ComputeMinMaxIdxResult);
+
   auto *EpiRedHeaderPhi =
       cast<VPReductionPHIRecipe>(EpiRedResult->getOperand(0));
   const RecurrenceDescriptor &RdxDesc =
@@ -8140,10 +8155,9 @@ void VPRecipeBuilder::collectScaledReductions(VFRange &Range) {
   // Find all possible partial reductions.
   SmallVector<std::pair<PartialReductionChain, unsigned>>
       PartialReductionChains;
-  for (const auto &[Phi, RdxDesc] : Legal->getReductionVars()) {
-    getScaledReductions(Phi, RdxDesc.getLoopExitInstr(), Range,
-                        PartialReductionChains);
-  }
+  for (const auto &[Phi, RdxDesc] : Legal->getReductionVars())
+    if (auto *ExitInstr = RdxDesc.getLoopExitInstr())
+      getScaledReductions(Phi, ExitInstr, Range, PartialReductionChains);
 
   // A partial reduction is invalid if any of its extends are used by
   // something that isn't another partial reduction. This is because the
@@ -9037,6 +9051,7 @@ void LoopVectorizationPlanner::adjustRecipesForReductions(
     assert(
         !RecurrenceDescriptor::isAnyOfRecurrenceKind(Kind) &&
         !RecurrenceDescriptor::isFindLastIVRecurrenceKind(Kind) &&
+        !RecurrenceDescriptor::isMinMaxIdxRecurrenceKind(Kind) &&
         "AnyOf and FindLast reductions are not allowed for in-loop reductions");
 
     // Collect the chain of "link" recipes for the reduction starting at PhiR.
@@ -9160,15 +9175,32 @@ void LoopVectorizationPlanner::adjustRecipesForReductions(
       PreviousLink = RedRecipe;
     }
   }
+
+  // Collect all VPReductionPHIRecipes in the header block, and sort them based
+  // on the dependency order of the reductions. This ensures that results of
+  // min/max reductions are computed before their corresponding index
+  // reductions, since the index reduction relies on the result of the min/max
+  // reduction to determine which lane produced the min/max.
+  SmallVector<VPReductionPHIRecipe *> VPReductionPHIs;
+  for (VPRecipeBase &R : Header->phis())
+    if (auto *PhiR = dyn_cast<VPReductionPHIRecipe>(&R))
+      VPReductionPHIs.push_back(PhiR);
+
+  stable_sort(VPReductionPHIs, [this](const VPReductionPHIRecipe *R1,
+                                      const VPReductionPHIRecipe *R2) {
+    auto *Phi1 = cast<PHINode>(R1->getUnderlyingInstr());
+    if (!Legal->isMinMaxRecurrence(Phi1))
+      return false;
+
+    auto *Phi2 = cast<PHINode>(R2->getUnderlyingInstr());
+    return Legal->getMinMaxRecurrences().find(Phi1)->second == Phi2;
+  });
+
   VPBasicBlock *LatchVPBB = VectorLoopRegion->getExitingBasicBlock();
   Builder.setInsertPoint(&*std::prev(std::prev(LatchVPBB->end())));
   VPBasicBlock::iterator IP = MiddleVPBB->getFirstNonPhi();
-  for (VPRecipeBase &R :
-       Plan->getVectorLoopRegion()->getEntryBasicBlock()->phis()) {
-    VPReductionPHIRecipe *PhiR = dyn_cast<VPReductionPHIRecipe>(&R);
-    if (!PhiR)
-      continue;
-
+  SmallDenseMap<VPReductionPHIRecipe *, VPValue *> IdxReductionMasks;
+  for (auto *PhiR : VPReductionPHIs) {
     const RecurrenceDescriptor &RdxDesc = PhiR->getRecurrenceDescriptor();
     Type *PhiTy = PhiR->getUnderlyingValue()->getType();
     // If tail is folded by masking, introduce selects between the phi
@@ -9195,7 +9227,9 @@ void LoopVectorizationPlanner::adjustRecipesForReductions(
                 cast<VPInstruction>(&U)->getOpcode() ==
                     VPInstruction::ComputeReductionResult ||
                 cast<VPInstruction>(&U)->getOpcode() ==
-                    VPInstruction::ComputeFindLastIVResult);
+                    VPInstruction::ComputeFindLastIVResult ||
+                cast<VPInstruction>(&U)->getOpcode() ==
+                    VPInstruction::ComputeMinMaxIdxResult);
       });
       if (CM.usePredicatedReductionSelect())
         PhiR->setOperand(1, NewExitingVPV);
@@ -9239,6 +9273,7 @@ void LoopVectorizationPlanner::adjustRecipesForReductions(
     VPInstruction *FinalReductionResult;
     VPBuilder::InsertPointGuard Guard(Builder);
     Builder.setInsertPoint(MiddleVPBB, IP);
+    RecurKind RK = RdxDesc.getRecurrenceKind();
     if (RecurrenceDescriptor::isFindLastIVRecurrenceKind(
             RdxDesc.getRecurrenceKind())) {
       VPValue *Start = PhiR->getStartValue();
@@ -9251,6 +9286,19 @@ void LoopVectorizationPlanner::adjustRecipesForReductions(
       FinalReductionResult =
           Builder.createNaryOp(VPInstruction::ComputeAnyOfResult,
                                {PhiR, Start, NewExitingVPV}, ExitDL);
+    } else if (RecurrenceDescriptor::isMinMaxIdxRecurrenceKind(RK)) {
+      // Mask out lanes that cannot be the index of the min/max value.
+      VPValue *Mask = IdxReductionMasks.at(PhiR);
+      Value *Iden = llvm::getRecurrenceIdentity(
+          RK == RecurKind::MinMaxFirstIdx ? RecurKind::SMin : RecurKind::SMax,
+          PhiTy, RdxDesc.getFastMathFlags());
+      NewExitingVPV = Builder.createSelect(Mask, NewExitingVPV,
+                                           Plan->getOrAddLiveIn(Iden), ExitDL);
+
+      VPValue *Start = PhiR->getStartValue();
+      FinalReductionResult =
+          Builder.createNaryOp(VPInstruction::ComputeMinMaxIdxResult,
+                               {PhiR, Start, NewExitingVPV}, ExitDL);
     } else {
       VPIRFlags Flags = RecurrenceDescriptor::isFloatingPointRecurrenceKind(
                             RdxDesc.getRecurrenceKind())
@@ -9262,11 +9310,25 @@ void LoopVectorizationPlanner::adjustRecipesForReductions(
     }
     // Update all users outside the vector region.
     OrigExitingVPV->replaceUsesWithIf(
-        FinalReductionResult, [FinalReductionResult](VPUser &User, unsigned) {
+        FinalReductionResult,
+        [FinalReductionResult, NewExitingVPV](VPUser &User, unsigned) {
           auto *Parent = cast<VPRecipeBase>(&User)->getParent();
-          return FinalReductionResult != &User && !Parent->getParent();
+          return FinalReductionResult != &User &&
+                 NewExitingVPV->getDefiningRecipe() != &User &&
+                 !Parent->getParent();
         });
 
+    // Generate a mask for the index reduction.
+    auto *Phi = cast<PHINode>(PhiR->getUnderlyingInstr());
+    if (Legal->isMinMaxRecurrence(Phi)) {
+      VPValue *IdxRdxMask = Builder.createICmp(CmpInst::ICMP_EQ, NewExitingVPV,
+                                               FinalReductionResult, ExitDL);
+      PHINode *IdxPhi = Legal->getMinMaxRecurrences().find(Phi)->second;
+      IdxReductionMasks.try_emplace(
+          cast<VPReductionPHIRecipe>(RecipeBuilder.getRecipe(IdxPhi)),
+          IdxRdxMask);
+    }
+
     // Adjust AnyOf reductions; replace the reduction phi for the selected value
     // with a boolean reduction phi node to check if the condition is true in
     // any iteration. The final value is selected by the final
@@ -9301,16 +9363,17 @@ void LoopVectorizationPlanner::adjustRecipesForReductions(
       continue;
     }
 
-    if (RecurrenceDescriptor::isFindLastIVRecurrenceKind(
-            RdxDesc.getRecurrenceKind())) {
-      // Adjust the start value for FindLastIV recurrences to use the sentinel
-      // value after generating the ResumePhi recipe, which uses the original
-      // start value.
+    if (RecurrenceDescriptor::isFindLastIVRecurrenceKind(RK) ||
+        RecurrenceDescriptor::isMinMaxIdxRecurrenceKind(RK)) {
+      // Adjust the start value for FindLastIV/MinMaxIdx recurrences to use the
+      // sentinel value after generating the ResumePhi recipe, which uses the
+      // original start value.
       PhiR->setOperand(0, Plan->getOrAddLiveIn(RdxDesc.getSentinelValue()));
     }
-    RecurKind RK = RdxDesc.getRecurrenceKind();
+
     if ((!RecurrenceDescriptor::isAnyOfRecurrenceKind(RK) &&
          !RecurrenceDescriptor::isFindLastIVRecurrenceKind(RK) &&
+         !RecurrenceDescriptor::isMinMaxIdxRecurrenceKind(RK) &&
          !RecurrenceDescriptor::isMinMaxRecurrenceKind(RK))) {
       VPBuilder PHBuilder(Plan->getVectorPreheader());
       VPValue *Iden = Plan->getOrAddLiveIn(
diff --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h
@@ -909,6 +909,7 @@ class VPInstruction : public VPRecipeWithIRFlags,
     Broadcast,
     ComputeAnyOfResult,
     ComputeFindLastIVResult,
+    ComputeMinMaxIdxResult,
     ComputeReductionResult,
     // Extracts the last lane from its operand if it is a vector, or the last
     // part if scalar. In the latter case, the recipe will be removed during
diff --git a/llvm/lib/Transforms/Vectorize/VPlanAnalysis.cpp b/llvm/lib/Transforms/Vectorize/VPlanAnalysis.cpp
@@ -92,6 +92,7 @@ Type *VPTypeAnalysis::inferScalarTypeForRecipe(const VPInstruction *R) {
     return IntegerType::get(Ctx, 1);
   case VPInstruction::ComputeAnyOfResult:
   case VPInstruction::ComputeFindLastIVResult:
+  case VPInstruction::ComputeMinMaxIdxResult:
   case VPInstruction::ComputeReductionResult: {
     auto *PhiR = cast<VPReductionPHIRecipe>(R->getOperand(0));
     auto *OrigPhi = cast<PHINode>(PhiR->getUnderlyingValue());
diff --git a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
diff --git a/llvm/lib/Transforms/Vectorize/VPlanUnroll.cpp b/llvm/lib/Transforms/Vectorize/VPlanUnroll.cpp
diff --git a/llvm/test/Transforms/LoopVectorize/smax-idx.ll b/llvm/test/Transforms/LoopVectorize/smax-idx.ll
