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

llvm/include/llvm/Analysis/IVDescriptors.h

@@ -42,9 +42,13 @@ enum class RecurKind {
   And,      ///< Bitwise or logical AND of integers.
   Xor,      ///< Bitwise or logical XOR of integers.
   SMin,     ///< Signed integer min implemented in terms of select(cmp()).
+  SMinMultiUse,     ///< Signed integer min implemented in terms of select(cmp()).
   SMax,     ///< Signed integer max implemented in terms of select(cmp()).
+  SMaxMultiUse,     ///< Signed integer max implemented in terms of select(cmp()).
   UMin,     ///< Unsigned integer min implemented in terms of select(cmp()).
+  UMinMultiUse,     ///< Unsigned integer min implemented in terms of select(cmp()).
   UMax,     ///< Unsigned integer max implemented in terms of select(cmp()).
+  UMaxMultiUse,     ///< Unsigned integer max implemented in terms of select(cmp()).
   FAdd,     ///< Sum of floats.
   FMul,     ///< Product of floats.
   FMin,     ///< FP min implemented in terms of select(cmp()).
@@ -247,8 +251,26 @@ class RecurrenceDescriptor {
 
   /// Returns true if the recurrence kind is an integer min/max kind.
   static bool isIntMinMaxRecurrenceKind(RecurKind Kind) {
-    return Kind == RecurKind::UMin || Kind == RecurKind::UMax ||
-           Kind == RecurKind::SMin || Kind == RecurKind::SMax;
+    return Kind == RecurKind::UMin || Kind == RecurKind::UMinMultiUse ||
+           Kind == RecurKind::UMax || Kind == RecurKind::UMaxMultiUse ||
+           Kind == RecurKind::SMin || Kind == RecurKind::SMinMultiUse ||
+           Kind == RecurKind::SMax || Kind == RecurKind::SMaxMultiUse;
+  }
+
+  static RecurKind convertFromMultiUseKind(RecurKind Kind) {
+    switch (Kind) {
+    case RecurKind::UMaxMultiUse:
+      return RecurKind::UMax;
+    case RecurKind::UMinMultiUse:
+      return RecurKind::UMin;
+    case RecurKind::SMinMultiUse:
+      return RecurKind::SMin;
+    case RecurKind::SMaxMultiUse:
+      return RecurKind::SMax;
+    default:
+      return Kind;
+    }
+    llvm_unreachable("all cases must be handled above");
   }
 
   /// Returns true if the recurrence kind is a floating-point minnum/maxnum

llvm/lib/Analysis/IVDescriptors.cpp

@@ -214,6 +214,41 @@ 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::UMinMultiUse;
+  else if (match(Inc, m_OneUse(m_UMax(m_Value(A), m_Value(B)))))
+    RK = RecurKind::UMaxMultiUse;
+  else if (match(Inc, m_OneUse(m_SMax(m_Value(A), m_Value(B)))))
+    RK = RecurKind::SMaxMultiUse;
+  else if (match(Inc, m_OneUse(m_SMin(m_Value(A), m_Value(B)))))
+    RK = RecurKind::SMinMultiUse;
+  else
+    return std::nullopt;
+
+  if (RecurrenceDescriptor::convertFromMultiUseKind(RK) != Kind || 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);
+  return {RD};
+}
+
 bool RecurrenceDescriptor::AddReductionVar(
     PHINode *Phi, RecurKind Kind, Loop *TheLoop, FastMathFlags FuncFMF,
     RecurrenceDescriptor &RedDes, DemandedBits *DB, AssumptionCache *AC,
@@ -225,6 +260,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/Vectorize/LoopVectorize.cpp

@@ -7160,6 +7160,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,
@@ -7170,7 +7173,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
@@ -7882,6 +7889,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);
   }
@@ -8046,9 +8056,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);
@@ -8061,12 +8068,17 @@ VPRecipeBase *VPRecipeBuilder::tryToCreateWidenRecipe(VPSingleDefRecipe *R,
       PhiRecipe = new VPReductionPHIRecipe(
           Phi, RdxDesc.getRecurrenceKind(), *StartV, CM.isInLoopReduction(Phi),
           CM.useOrderedReductions(RdxDesc), ScaleFactor);
-    } else {
+    } 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]);
@@ -8347,8 +8359,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)) {
@@ -8410,6 +8425,10 @@ VPlanPtr LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(
   // Adjust the recipes for any inloop reductions.
   adjustRecipesForReductions(Plan, RecipeBuilder, Range.Start);
 
+  // Try to convert remaining VPWidenPHIRecipes to reduction recipes.
+  if (!VPlanTransforms::runPass(VPlanTransforms::legalizeUnclassifiedPhis,
+                                *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/SLPVectorizer.cpp

@@ -25002,6 +25002,10 @@ class HorizontalReduction {
         case RecurKind::FMinNum:
         case RecurKind::FMaximumNum:
         case RecurKind::FMinimumNum:
+        case RecurKind::UMaxMultiUse:
+        case RecurKind::UMinMultiUse:
+        case RecurKind::SMinMultiUse:
+        case RecurKind::SMaxMultiUse:
         case RecurKind::None:
           llvm_unreachable("Unexpected reduction kind for repeated scalar.");
         }
@@ -25143,6 +25147,10 @@ class HorizontalReduction {
     case RecurKind::FMinNum:
     case RecurKind::FMaximumNum:
     case RecurKind::FMinimumNum:
+    case RecurKind::UMaxMultiUse:
+    case RecurKind::UMinMultiUse:
+    case RecurKind::SMinMultiUse:
+    case RecurKind::SMaxMultiUse:
     case RecurKind::None:
       llvm_unreachable("Unexpected reduction kind for repeated scalar.");
     }
@@ -25249,6 +25257,10 @@ class HorizontalReduction {
     case RecurKind::FMinNum:
     case RecurKind::FMaximumNum:
     case RecurKind::FMinimumNum:
+    case RecurKind::UMaxMultiUse:
+    case RecurKind::UMinMultiUse:
+    case RecurKind::SMinMultiUse:
+    case RecurKind::SMaxMultiUse:
     case RecurKind::None:
       llvm_unreachable("Unexpected reduction kind for reused scalars.");
     }

llvm/lib/Transforms/Vectorize/VPlan.h

@@ -1988,7 +1988,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;
   }
@@ -1997,6 +1998,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;
@@ -2061,7 +2066,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));
   }
 

llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp

@@ -950,3 +950,105 @@ bool VPlanTransforms::handleMaxMinNumReductions(VPlan &Plan) {
   MiddleTerm->setOperand(0, NewCond);
   return true;
 }
+
+bool VPlanTransforms::legalizeUnclassifiedPhis(VPlan &Plan) {
+  using namespace VPlanPatternMatch;
+  for (auto &PhiR : make_early_inc_range(
+           Plan.getVectorLoopRegion()->getEntryBasicBlock()->phis())) {
+    auto *MinMaxPhiR = dyn_cast<VPReductionPHIRecipe>(&PhiR);
+    if (!MinMaxPhiR || !RecurrenceDescriptor::isIntMinMaxRecurrenceKind(
+                           MinMaxPhiR->getRecurrenceKind()))
+      continue;
+
+    RecurKind RdxKind = RecurrenceDescriptor::convertFromMultiUseKind(
+        MinMaxPhiR->getRecurrenceKind());
+    if (RdxKind == MinMaxPhiR->getRecurrenceKind())
+      continue;
+
+    // One user of MinMaxPhiR is MinMaxOp, the other users must be a compare
+    // that's part of a FindLastIV chain.
+    auto *MinMaxOp = cast<VPRecipeWithIRFlags>(MinMaxPhiR->getBackedgeValue());
+    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");
+
+    auto NewPhiR = new VPReductionPHIRecipe(
+        cast<PHINode>(MinMaxPhiR->getUnderlyingInstr()), RdxKind,
+        *MinMaxPhiR->getOperand(0), false, false, 1);
+    NewPhiR->insertBefore(MinMaxPhiR);
+    MinMaxPhiR->replaceAllUsesWith(NewPhiR);
+    NewPhiR->addOperand(MinMaxPhiR->getOperand(1));
+    MinMaxPhiR->eraseFromParent();
+
+    // 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 = cast<VPInstruction>(
+        *(Sel->user_begin() + (*Sel->user_begin() == FindIVPhiR ? 1 : 0)));
+    VPBuilder B(FindIVResult);
+    VPInstruction *MinMaxResult =
+        B.createNaryOp(VPInstruction::ComputeReductionResult,
+                       {NewPhiR, NewPhiR->getBackedgeValue()}, VPIRFlags(), {});
+    NewPhiR->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;
+}

llvm/lib/Transforms/Vectorize/VPlanTransforms.h

@@ -143,6 +143,11 @@ struct VPlanTransforms {
           GetIntOrFpInductionDescriptor,
       const TargetLibraryInfo &TLI);
 
+  /// Try to legalize unclassified phis by converting VPWidenPHIRecipes to
+  /// min-max reductions used by FindLastIV reductions if possible. Returns
+  /// false if the VPlan contains VPWidenPHIRecipes that cannot be legalized.
+  static bool legalizeUnclassifiedPhis(VPlan &Plan);
+
   /// Try to have all users of fixed-order recurrences appear after the recipe
   /// defining their previous value, by either sinking users or hoisting recipes
   /// defining their previous value (and its operands). Then introduce