[LV] Vectorize selecting last IV of min/max element.

llvm/include/llvm/Analysis/IVDescriptors.h

@@ -95,12 +95,15 @@ class RecurrenceDescriptor {
                        RecurKind K, FastMathFlags FMF, Instruction *ExactFP,
                        Type *RT, bool Signed, bool Ordered,
                        SmallPtrSetImpl<Instruction *> &CI,
-                       unsigned MinWidthCastToRecurTy)
+                       unsigned MinWidthCastToRecurTy, bool PhiMultiUse = false)
       : IntermediateStore(Store), StartValue(Start), LoopExitInstr(Exit),
         Kind(K), FMF(FMF), ExactFPMathInst(ExactFP), RecurrenceType(RT),
-        IsSigned(Signed), IsOrdered(Ordered),
+        IsSigned(Signed), IsOrdered(Ordered), IsPhiMultiUse(PhiMultiUse),
         MinWidthCastToRecurrenceType(MinWidthCastToRecurTy) {
     CastInsts.insert_range(CI);
+    assert(
+        (!PhiMultiUse || isMinMaxRecurrenceKind(K)) &&
+        "Only min/max recurrences are allowed to have multiple uses currently");
   }
 
   /// This POD struct holds information about a potential recurrence operation.
@@ -339,6 +342,10 @@ class RecurrenceDescriptor {
   /// Expose an ordered FP reduction to the instance users.
   bool isOrdered() const { return IsOrdered; }
 
+  /// Returns true if the reduction PHI has multiple in-loop users. This is
+  /// relevant for min/max reductions that are part of a FindLastIV pattern.
+  bool isPhiMultiUse() const { return IsPhiMultiUse; }
+
   /// Attempts to find a chain of operations from Phi to LoopExitInst that can
   /// be treated as a set of reductions instructions for in-loop reductions.
   LLVM_ABI SmallVector<Instruction *, 4> getReductionOpChain(PHINode *Phi,
@@ -376,6 +383,9 @@ class RecurrenceDescriptor {
   // Currently only a non-reassociative FAdd can be considered in-order,
   // if it is also the only FAdd in the PHI's use chain.
   bool IsOrdered = false;
+  // True if the reduction PHI has multiple in-loop users. This is relevant
+  // for min/max reductions that are part of a FindLastIV pattern.
+  bool IsPhiMultiUse = false;
   // Instructions used for type-promoting the recurrence.
   SmallPtrSet<Instruction *, 8> CastInsts;
   // The minimum width used by the recurrence.

llvm/lib/Analysis/IVDescriptors.cpp

@@ -214,6 +214,40 @@ static bool checkOrderedReduction(RecurKind Kind, Instruction *ExactFPMathInst,
   return true;
 }
 
+static std::optional<RecurrenceDescriptor>
+getMultiUseMinMax(PHINode *Phi, RecurKind Kind, Loop *TheLoop) {
+  BasicBlock *Latch = TheLoop->getLoopLatch();
+  if (!Latch)
+    return std::nullopt;
+  Value *Inc = Phi->getIncomingValueForBlock(Latch);
+  RecurKind RK;
+  if (Phi->hasOneUse() ||
+      !RecurrenceDescriptor::isIntMinMaxRecurrenceKind(Kind))
+    return std::nullopt;
+
+  Value *A, *B;
+  if (match(Inc, m_OneUse(m_UMin(m_Value(A), m_Value(B)))))
+    RK = RecurKind::UMin;
+  else if (match(Inc, m_OneUse(m_UMax(m_Value(A), m_Value(B)))))
+    RK = RecurKind::UMax;
+  else if (match(Inc, m_OneUse(m_SMax(m_Value(A), m_Value(B)))))
+    RK = RecurKind::SMax;
+  else if (match(Inc, m_OneUse(m_SMin(m_Value(A), m_Value(B)))))
+    RK = RecurKind::SMin;
+  else
+    return std::nullopt;
+
+  if (A == B || (A != Phi && B != Phi))
+    return std::nullopt;
+
+  SmallPtrSet<Instruction *, 4> CastInsts;
+  Value *RdxStart = Phi->getIncomingValueForBlock(TheLoop->getLoopPreheader());
+  RecurrenceDescriptor RD(RdxStart, nullptr, nullptr, RK, FastMathFlags(),
+                          nullptr, Phi->getType(), false, false, CastInsts, -1U,
+                          true);
+  return {RD};
+}
+
 bool RecurrenceDescriptor::AddReductionVar(
     PHINode *Phi, RecurKind Kind, Loop *TheLoop, FastMathFlags FuncFMF,
     RecurrenceDescriptor &RedDes, DemandedBits *DB, AssumptionCache *AC,
@@ -225,6 +259,10 @@ bool RecurrenceDescriptor::AddReductionVar(
   if (Phi->getParent() != TheLoop->getHeader())
     return false;
 
+  if (auto RD = getMultiUseMinMax(Phi, Kind, TheLoop)) {
+    RedDes = *RD;
+    return true;
+  }
   // Obtain the reduction start value from the value that comes from the loop
   // preheader.
   Value *RdxStart = Phi->getIncomingValueForBlock(TheLoop->getLoopPreheader());

llvm/lib/Transforms/Utils/LoopUnroll.cpp

@@ -1252,7 +1252,9 @@ llvm::canParallelizeReductionWhenUnrolling(PHINode &Phi, Loop *L,
   RecurrenceDescriptor RdxDesc;
   if (!RecurrenceDescriptor::isReductionPHI(&Phi, L, RdxDesc,
                                             /*DemandedBits=*/nullptr,
-                                            /*AC=*/nullptr, /*DT=*/nullptr, SE))
+                                            /*AC=*/nullptr, /*DT=*/nullptr,
+                                            SE) ||
+      RdxDesc.isPhiMultiUse())
     return std::nullopt;
   RecurKind RK = RdxDesc.getRecurrenceKind();
   // Skip unsupported reductions.

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

@@ -878,6 +878,11 @@ bool LoopVectorizationLegality::canVectorizeInstr(Instruction &I) {
       Requirements->addExactFPMathInst(RedDes.getExactFPMathInst());
       AllowedExit.insert(RedDes.getLoopExitInstr());
       Reductions[Phi] = RedDes;
+      assert((!RedDes.isPhiMultiUse() ||
+              RecurrenceDescriptor::isMinMaxRecurrenceKind(
+                  RedDes.getRecurrenceKind())) &&
+             "Only min/max recurrences are allowed to have multiple uses "
+             "currently");
       return true;
     }
 

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -6533,6 +6533,11 @@ void LoopVectorizationCostModel::collectInLoopReductions() {
     PHINode *Phi = Reduction.first;
     const RecurrenceDescriptor &RdxDesc = Reduction.second;
 
+    // Multi-use reductions (e.g., used in FindLastIV patterns) are handled
+    // separately and should not be considered for in-loop reductions.
+    if (RdxDesc.isPhiMultiUse())
+      continue;
+
     // We don't collect reductions that are type promoted (yet).
     if (RdxDesc.getRecurrenceType() != Phi->getType())
       continue;
@@ -7184,6 +7189,9 @@ static void fixReductionScalarResumeWhenVectorizingEpilog(
     Value *StartV = getStartValueFromReductionResult(EpiRedResult);
     Value *SentinelV = EpiRedResult->getOperand(2)->getLiveInIRValue();
     using namespace llvm::PatternMatch;
+    MainResumeValue = cast<VPInstruction>(EpiRedHeaderPhi->getStartValue())
+                          ->getOperand(0)
+                          ->getUnderlyingValue();
     Value *Cmp, *OrigResumeV, *CmpOp;
     [[maybe_unused]] bool IsExpectedPattern =
         match(MainResumeValue,
@@ -7194,7 +7202,11 @@ static void fixReductionScalarResumeWhenVectorizingEpilog(
          ((CmpOp == StartV && isGuaranteedNotToBeUndefOrPoison(CmpOp))));
     assert(IsExpectedPattern && "Unexpected reduction resume pattern");
     MainResumeValue = OrigResumeV;
+  } else if (auto *VPI =
+                 dyn_cast<VPInstruction>(EpiRedHeaderPhi->getStartValue())) {
+    MainResumeValue = VPI->getOperand(0)->getUnderlyingValue();
   }
+
   PHINode *MainResumePhi = cast<PHINode>(MainResumeValue);
 
   // When fixing reductions in the epilogue loop we should already have
@@ -7906,6 +7918,9 @@ void VPRecipeBuilder::collectScaledReductions(VFRange &Range) {
   SmallVector<std::pair<PartialReductionChain, unsigned>>
       PartialReductionChains;
   for (const auto &[Phi, RdxDesc] : Legal->getReductionVars()) {
+    if (RecurrenceDescriptor::isMinMaxRecurrenceKind(
+            RdxDesc.getRecurrenceKind()))
+      continue;
     getScaledReductions(Phi, RdxDesc.getLoopExitInstr(), Range,
                         PartialReductionChains);
   }
@@ -8070,9 +8085,6 @@ VPRecipeBase *VPRecipeBuilder::tryToCreateWidenRecipe(VPSingleDefRecipe *R,
       return Recipe;
 
     VPHeaderPHIRecipe *PhiRecipe = nullptr;
-    assert((Legal->isReductionVariable(Phi) ||
-            Legal->isFixedOrderRecurrence(Phi)) &&
-           "can only widen reductions and fixed-order recurrences here");
     VPValue *StartV = Operands[0];
     if (Legal->isReductionVariable(Phi)) {
       const RecurrenceDescriptor &RdxDesc = Legal->getRecurrenceDescriptor(Phi);
@@ -8084,13 +8096,19 @@ VPRecipeBase *VPRecipeBuilder::tryToCreateWidenRecipe(VPSingleDefRecipe *R,
           getScalingForReduction(RdxDesc.getLoopExitInstr()).value_or(1);
       PhiRecipe = new VPReductionPHIRecipe(
           Phi, RdxDesc.getRecurrenceKind(), *StartV, CM.isInLoopReduction(Phi),
-          CM.useOrderedReductions(RdxDesc), ScaleFactor);
-    } else {
+          CM.useOrderedReductions(RdxDesc), ScaleFactor,
+          RdxDesc.isPhiMultiUse());
+    } else if (Legal->isFixedOrderRecurrence(Phi)) {
       // TODO: Currently fixed-order recurrences are modeled as chains of
       // first-order recurrences. If there are no users of the intermediate
       // recurrences in the chain, the fixed order recurrence should be modeled
       // directly, enabling more efficient codegen.
       PhiRecipe = new VPFirstOrderRecurrencePHIRecipe(Phi, *StartV);
+    } else {
+      // Failed to identify phi as reduction or fixed-order recurrence. Keep the
+      // original VPWidenPHIRecipe for now, to be legalized later if possible.
+      setRecipe(Phi, R);
+      return nullptr;
     }
     // Add backedge value.
     PhiRecipe->addOperand(Operands[1]);
@@ -8365,8 +8383,11 @@ VPlanPtr LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(
 
       VPRecipeBase *Recipe =
           RecipeBuilder.tryToCreateWidenRecipe(SingleDef, Range);
-      if (!Recipe)
+      if (!Recipe) {
+        if (isa<VPPhi>(SingleDef))
+          continue;
         Recipe = RecipeBuilder.handleReplication(Instr, R.operands(), Range);
+      }
 
       RecipeBuilder.setRecipe(Instr, Recipe);
       if (isa<VPWidenIntOrFpInductionRecipe>(Recipe) && isa<TruncInst>(Instr)) {
@@ -8428,6 +8449,10 @@ VPlanPtr LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(
   // Adjust the recipes for any inloop reductions.
   adjustRecipesForReductions(Plan, RecipeBuilder, Range.Start);
 
+  // Try to legalize reductions with multiple in-loop uses.
+  if (!VPlanTransforms::runPass(VPlanTransforms::legalizeMultiUseReductions,
+                                *Plan))
+    return nullptr;
   // Apply mandatory transformation to handle FP maxnum/minnum reduction with
   // NaNs if possible, bail out otherwise.
   if (!VPlanTransforms::runPass(VPlanTransforms::handleMaxMinNumReductions,

llvm/lib/Transforms/Vectorize/VPlan.h

@@ -1994,7 +1994,8 @@ class LLVM_ABI_FOR_TEST VPHeaderPHIRecipe : public VPSingleDefRecipe,
   ~VPHeaderPHIRecipe() override = default;
 
   /// Method to support type inquiry through isa, cast, and dyn_cast.
-  static inline bool classof(const VPRecipeBase *B) {
+  static inline bool classof(const VPUser *U) {
+    auto *B = cast<VPRecipeBase>(U);
     return B->getVPDefID() >= VPDef::VPFirstHeaderPHISC &&
            B->getVPDefID() <= VPDef::VPLastHeaderPHISC;
   }
@@ -2003,6 +2004,10 @@ class LLVM_ABI_FOR_TEST VPHeaderPHIRecipe : public VPSingleDefRecipe,
     return B && B->getVPDefID() >= VPRecipeBase::VPFirstHeaderPHISC &&
            B->getVPDefID() <= VPRecipeBase::VPLastHeaderPHISC;
   }
+  static inline bool classof(const VPSingleDefRecipe *B) {
+    return B->getVPDefID() >= VPDef::VPFirstHeaderPHISC &&
+           B->getVPDefID() <= VPDef::VPLastHeaderPHISC;
+  }
 
   /// Generate the phi nodes.
   void execute(VPTransformState &State) override = 0;
@@ -2067,7 +2072,7 @@ class VPWidenInductionRecipe : public VPHeaderPHIRecipe {
     return R && classof(R);
   }
 
-  static inline bool classof(const VPHeaderPHIRecipe *R) {
+  static inline bool classof(const VPSingleDefRecipe *R) {
     return classof(static_cast<const VPRecipeBase *>(R));
   }
 
@@ -2344,6 +2349,10 @@ class VPReductionPHIRecipe : public VPHeaderPHIRecipe,
   /// The phi is part of an ordered reduction. Requires IsInLoop to be true.
   bool IsOrdered;
 
+  /// The phi is part of a multi-use reduction (e.g., used in FindLastIV
+  /// patterns).
+  bool IsPhiMultiUse;
+
   /// When expanding the reduction PHI, the plan's VF element count is divided
   /// by this factor to form the reduction phi's VF.
   unsigned VFScaleFactor = 1;
@@ -2352,9 +2361,10 @@ class VPReductionPHIRecipe : public VPHeaderPHIRecipe,
   /// Create a new VPReductionPHIRecipe for the reduction \p Phi.
   VPReductionPHIRecipe(PHINode *Phi, RecurKind Kind, VPValue &Start,
                        bool IsInLoop = false, bool IsOrdered = false,
-                       unsigned VFScaleFactor = 1)
+                       unsigned VFScaleFactor = 1, bool IsPhiMultiUse = false)
       : VPHeaderPHIRecipe(VPDef::VPReductionPHISC, Phi, &Start), Kind(Kind),
-        IsInLoop(IsInLoop), IsOrdered(IsOrdered), VFScaleFactor(VFScaleFactor) {
+        IsInLoop(IsInLoop), IsOrdered(IsOrdered), IsPhiMultiUse(IsPhiMultiUse),
+        VFScaleFactor(VFScaleFactor) {
     assert((!IsOrdered || IsInLoop) && "IsOrdered requires IsInLoop");
   }
 
@@ -2363,7 +2373,7 @@ class VPReductionPHIRecipe : public VPHeaderPHIRecipe,
   VPReductionPHIRecipe *clone() override {
     auto *R = new VPReductionPHIRecipe(
         dyn_cast_or_null<PHINode>(getUnderlyingValue()), getRecurrenceKind(),
-        *getOperand(0), IsInLoop, IsOrdered, VFScaleFactor);
+        *getOperand(0), IsInLoop, IsOrdered, VFScaleFactor, IsPhiMultiUse);
     R->addOperand(getBackedgeValue());
     return R;
   }
@@ -2396,6 +2406,9 @@ class VPReductionPHIRecipe : public VPHeaderPHIRecipe,
   /// Returns true, if the phi is part of an in-loop reduction.
   bool isInLoop() const { return IsInLoop; }
 
+  /// Returns true, if the phi is part of a multi-use reduction.
+  bool isPhiMultiUse() const { return IsPhiMultiUse; }
+
   /// Returns true if the recipe only uses the first lane of operand \p Op.
   bool onlyFirstLaneUsed(const VPValue *Op) const override {
     assert(is_contained(operands(), Op) &&

llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp

@@ -948,3 +948,101 @@ bool VPlanTransforms::handleMaxMinNumReductions(VPlan &Plan) {
   MiddleTerm->setOperand(0, NewCond);
   return true;
 }
+
+bool VPlanTransforms::legalizeMultiUseReductions(VPlan &Plan) {
+  for (auto &PhiR : make_early_inc_range(
+           Plan.getVectorLoopRegion()->getEntryBasicBlock()->phis())) {
+    auto *MinMaxPhiR = dyn_cast<VPReductionPHIRecipe>(&PhiR);
+    if (!MinMaxPhiR)
+      continue;
+
+    RecurKind RdxKind = MinMaxPhiR->getRecurrenceKind();
+    // TODO: check for multi-uses in VPlan directly.
+    if (!RecurrenceDescriptor::isIntMinMaxRecurrenceKind(RdxKind) ||
+        !MinMaxPhiR->isPhiMultiUse())
+      continue;
+
+    // One user of MinMaxPhiR is MinMaxOp, the other users must be a compare
+    // that's part of a FindLastIV chain.
+    auto *MinMaxOp =
+        dyn_cast<VPRecipeWithIRFlags>(MinMaxPhiR->getBackedgeValue());
+    if (!MinMaxOp || MinMaxOp->getNumUsers() != 2)
+      return false;
+    auto MinMaxUsers = to_vector(MinMaxPhiR->users());
+    auto *Cmp = dyn_cast<VPRecipeWithIRFlags>(
+        MinMaxUsers[0] == MinMaxOp ? MinMaxUsers[1] : MinMaxUsers[0]);
+    VPValue *CmpOpA;
+    VPValue *CmpOpB;
+    if (!Cmp || Cmp->getNumUsers() != 1 ||
+        !match(Cmp, m_Binary<Instruction::ICmp>(m_VPValue(CmpOpA),
+                                                m_VPValue(CmpOpB))))
+      return false;
+
+    // Normalize the predicate so MinMaxPhiR is on the right side.
+    CmpInst::Predicate Pred = Cmp->getPredicate();
+    if (CmpOpA == MinMaxPhiR)
+      Pred = CmpInst::getSwappedPredicate(Pred);
+
+    // Determine if the predicate is not strict.
+    bool IsNonStrictPred = ICmpInst::isLE(Pred) || ICmpInst::isGE(Pred);
+    // Account for a mis-match between RdxKind and the predicate.
+    switch (RdxKind) {
+    case RecurKind::UMin:
+    case RecurKind::SMin:
+      IsNonStrictPred |= ICmpInst::isGT(Pred);
+      break;
+    case RecurKind::UMax:
+    case RecurKind::SMax:
+      IsNonStrictPred |= ICmpInst::isLT(Pred);
+      break;
+    default:
+      llvm_unreachable("unsupported kind");
+    }
+
+    // TODO: Strict predicates need to find the first IV value for which the
+    // predicate holds, not the last.
+    if (Pred == CmpInst::ICMP_NE || !IsNonStrictPred)
+      return false;
+
+    // Cmp must be used by the select of a FindLastIV chain.
+    VPValue *Sel = dyn_cast<VPSingleDefRecipe>(*Cmp->user_begin());
+    VPValue *IVOp, *FindIV;
+    if (!Sel ||
+        !match(Sel,
+               m_Select(m_Specific(Cmp), m_VPValue(IVOp), m_VPValue(FindIV))) ||
+        Sel->getNumUsers() != 2 || !isa<VPWidenIntOrFpInductionRecipe>(IVOp))
+      return false;
+    auto *FindIVPhiR = dyn_cast<VPReductionPHIRecipe>(FindIV);
+    if (!FindIVPhiR || !RecurrenceDescriptor::isFindLastIVRecurrenceKind(
+                           FindIVPhiR->getRecurrenceKind()))
+      return false;
+
+    assert(!FindIVPhiR->isInLoop() && !FindIVPhiR->isOrdered() &&
+           "cannot handle inloop/ordered reductions yet");
+
+    // The reduction using MinMaxPhiR needs adjusting to compute the correct
+    // result:
+    //  1. We need to find the last 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 FindLastIV reductions which
+    //  correspond to the lanes matching the min/max reduction result.
+    VPInstruction *FindIVResult = dyn_cast<VPInstruction>(
+        *(Sel->user_begin() + (*Sel->user_begin() == FindIVPhiR ? 1 : 0)));
+    if (!FindIVResult)
+      return false;
+    VPBuilder B(FindIVResult);
+    VPInstruction *MinMaxResult = B.createNaryOp(
+        VPInstruction::ComputeReductionResult,
+        {MinMaxPhiR, MinMaxPhiR->getBackedgeValue()}, VPIRFlags(), {});
+    MinMaxPhiR->getBackedgeValue()->replaceUsesWithIf(
+        MinMaxResult, [](VPUser &U, unsigned) { return isa<VPPhi>(&U); });
+    auto *FinalMinMaxCmp = B.createICmp(
+        CmpInst::ICMP_EQ, MinMaxResult->getOperand(1), MinMaxResult);
+    auto *FinalIVSelect =
+        B.createSelect(FinalMinMaxCmp, FindIVResult->getOperand(3),
+                       FindIVResult->getOperand(2));
+    FindIVResult->setOperand(3, FinalIVSelect);
+  }
+  return true;
+}