[llvm] [LV] Vectorize selecting last IV of min/max element. (PR #141431)

[via llvm-commits]

================
@@ -992,3 +992,92 @@ bool VPlanTransforms::handleMaxMinNumReductions(VPlan &Plan) {
   MiddleTerm->setOperand(0, NewCond);
   return true;
 }
+
+bool VPlanTransforms::handleMultiUseReductions(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 (!MinMaxPhiR->hasLoopUsesOutsideReductionChain())
+      continue;
+
+    assert(
+        RecurrenceDescriptor::isIntMinMaxRecurrenceKind(RdxKind) &&
+        "only min/max recurrences support users outside the reduction chain");
+
+    // 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;
+
+    CmpInst::Predicate Pred = Cmp->getPredicate();
+    // TODO: Strict predicates need to find the first IV value for which the
+    // predicate holds, not the last.
+    if (Pred == CmpInst::ICMP_EQ || Pred == CmpInst::ICMP_NE ||
+        ICmpInst::isLT(Pred) || ICmpInst::isGT(Pred))
+      return false;
+
+    // Normalize the predicate so MinMaxPhiR is on the right side.
+    if (CmpOpA == MinMaxPhiR)
+      Pred = CmpInst::getSwappedPredicate(Pred);
+
+    // 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)
+      return false;
+
+    auto *FindIVPhiR = dyn_cast<VPReductionPHIRecipe>(FindIV);
+    if (!FindIVPhiR || !RecurrenceDescriptor::isFindLastIVRecurrenceKind(
+                           FindIVPhiR->getRecurrenceKind()))
+      return false;
+
+    assert(isa<VPWidenIntOrFpInductionRecipe>(IVOp) &&
+           "IVOp must be a wide induction");
+    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)));
----------------
ayalz wrote:

Is there a clearer and consistent way to "get the other user" among two users of a VPValue, as done above when setting Cmp above to the user of MinMaxPhiR other than MinMaxOp (using to_vector etc.) and here setting FindIVResult to the user of Sel other than FindIVPhiR?

https://github.com/llvm/llvm-project/pull/141431