[LV] Vectorize select min/max index.

fhahn · fhahn · commit 0554871550bc · 2025-05-25T22:48:19.000+01:00
Add support for vectorizing loops that select the index of the minimum or
maximum element. The patch implements vectorizing those patterns by
combining Min/Max and FindFirstIV reductions.

It extends matching Min/Max reductions to allow in-loop users that are
FindLastIV reductions. It records a flag indicating that the Min/Max
reduction is used by another reduction.

When creating reduction recipes, we process any reduction that has other
reduction users. The reduction using the min/max reduction needs adjusting
to compute the correct result:
 1. We need to find the first IV for which the condition based on the
    min/max reduction 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.
diff --git a/llvm/include/llvm/Analysis/IVDescriptors.h b/llvm/include/llvm/Analysis/IVDescriptors.h
@@ -296,6 +296,8 @@ class RecurrenceDescriptor {
     }
   }
 
+  void setKind(RecurKind NewKind) { Kind = NewKind; }
+
   /// Returns a reference to the instructions used for type-promoting the
   /// recurrence.
   const SmallPtrSet<Instruction *, 8> &getCastInsts() const { return CastInsts; }
@@ -327,6 +329,10 @@ class RecurrenceDescriptor {
   /// AddReductionVar method, this field will be assigned the last met store.
   StoreInst *IntermediateStore = nullptr;
 
+  /// True if this recurrence is used by another recurrence in the loop. Users
+  /// need to ensure that the final code-gen accounts for the use in the loop.
+  bool IsUsedByOtherRecurrence = false;
+
 private:
   // The starting value of the recurrence.
   // It does not have to be zero!
diff --git a/llvm/lib/Analysis/IVDescriptors.cpp b/llvm/lib/Analysis/IVDescriptors.cpp
@@ -255,7 +255,7 @@ bool RecurrenceDescriptor::AddReductionVar(
   // Data used for determining if the recurrence has been type-promoted.
   Type *RecurrenceType = Phi->getType();
   SmallPtrSet<Instruction *, 4> CastInsts;
-  unsigned MinWidthCastToRecurrenceType;
+  unsigned MinWidthCastToRecurrenceType = -1ull;
   Instruction *Start = Phi;
   bool IsSigned = false;
 
@@ -310,6 +310,7 @@ bool RecurrenceDescriptor::AddReductionVar(
   //    This is either:
   //      * An instruction type other than PHI or the reduction operation.
   //      * A PHI in the header other than the initial PHI.
+  bool IsUsedByOtherRecurrence = false;
   while (!Worklist.empty()) {
     Instruction *Cur = Worklist.pop_back_val();
 
@@ -371,15 +372,37 @@ bool RecurrenceDescriptor::AddReductionVar(
 
     // Any reduction instruction must be of one of the allowed kinds. We ignore
     // the starting value (the Phi or an AND instruction if the Phi has been
-    // type-promoted).
+    // type-promoted) and other in-loop users, if they form a FindLastIV
+    // reduction. In the latter case, the user of the IVDescriptors must account
+    // for that during codegen.
     if (Cur != Start) {
       ReduxDesc =
           isRecurrenceInstr(TheLoop, Phi, Cur, Kind, ReduxDesc, FuncFMF, SE);
       ExactFPMathInst = ExactFPMathInst == nullptr
                             ? ReduxDesc.getExactFPMathInst()
                             : ExactFPMathInst;
-      if (!ReduxDesc.isRecurrence())
+      if (!ReduxDesc.isRecurrence()) {
+        if (isMinMaxRecurrenceKind(Kind)) {
+          // If the current recurrence is Min/Max, check if the current user is
+          // a select that is a FindLastIV reduction. During codegen, this
+          // recurrence needs to be turned into one that finds the first IV, as
+          // the value to compare against is a Min/Max recurrence.
+          auto *Sel = dyn_cast<SelectInst>(Cur);
+          if (!Sel || !Sel->getType()->isIntegerTy())
+            return false;
+          auto *OtherPhi = dyn_cast<PHINode>(Sel->getOperand(2));
+          if (!OtherPhi)
+            return false;
+          auto NewReduxDesc =
+              isRecurrenceInstr(TheLoop, OtherPhi, Cur, RecurKind::FindLastIV,
+                                ReduxDesc, FuncFMF, SE);
+          if (NewReduxDesc.isRecurrence()) {
+            IsUsedByOtherRecurrence = true;
+            continue;
+          }
+        }
         return false;
+      }
       // FIXME: FMF is allowed on phi, but propagation is not handled correctly.
       if (isa<FPMathOperator>(ReduxDesc.getPatternInst()) && !IsAPhi) {
         FastMathFlags CurFMF = ReduxDesc.getPatternInst()->getFastMathFlags();
@@ -503,7 +526,7 @@ bool RecurrenceDescriptor::AddReductionVar(
   // pattern or more than just a select and cmp. Zero implies that we saw a
   // llvm.min/max intrinsic, which is always OK.
   if (isMinMaxRecurrenceKind(Kind) && NumCmpSelectPatternInst != 2 &&
-      NumCmpSelectPatternInst != 0)
+      NumCmpSelectPatternInst != 0 && !IsUsedByOtherRecurrence)
     return false;
 
   if (isAnyOfRecurrenceKind(Kind) && NumCmpSelectPatternInst != 1)
@@ -535,7 +558,13 @@ bool RecurrenceDescriptor::AddReductionVar(
       ExitInstruction = cast<Instruction>(IntermediateStore->getValueOperand());
   }
 
-  if (!FoundStartPHI || !FoundReduxOp || !ExitInstruction)
+  if (!FoundStartPHI || !FoundReduxOp)
+    return false;
+
+  if (IsUsedByOtherRecurrence) {
+    if (ExitInstruction)
+      return false;
+  } else if (!ExitInstruction)
     return false;
 
   const bool IsOrdered =
@@ -586,8 +615,9 @@ bool RecurrenceDescriptor::AddReductionVar(
   //       without needing a white list of instructions to ignore.
   //       This may also be useful for the inloop reductions, if it can be
   //       kept simple enough.
-  collectCastInstrs(TheLoop, ExitInstruction, RecurrenceType, CastInsts,
-                    MinWidthCastToRecurrenceType);
+  if (ExitInstruction)
+    collectCastInstrs(TheLoop, ExitInstruction, RecurrenceType, CastInsts,
+                      MinWidthCastToRecurrenceType);
 
   // We found a reduction var if we have reached the original phi node and we
   // only have a single instruction with out-of-loop users.
@@ -600,7 +630,7 @@ bool RecurrenceDescriptor::AddReductionVar(
                           FMF, ExactFPMathInst, RecurrenceType, IsSigned,
                           IsOrdered, CastInsts, MinWidthCastToRecurrenceType);
   RedDes = RD;
-
+  RedDes.IsUsedByOtherRecurrence = IsUsedByOtherRecurrence;
   return true;
 }
 
diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -8499,7 +8499,7 @@ bool VPRecipeBuilder::getScaledReductions(
     Instruction *PHI, Instruction *RdxExitInstr, VFRange &Range,
     SmallVectorImpl<std::pair<PartialReductionChain, unsigned>> &Chains) {
 
-  if (!CM.TheLoop->contains(RdxExitInstr))
+  if (!RdxExitInstr || !CM.TheLoop->contains(RdxExitInstr))
     return false;
 
   auto *Update = dyn_cast<BinaryOperator>(RdxExitInstr);
@@ -9621,6 +9621,75 @@ void LoopVectorizationPlanner::adjustRecipesForReductions(
       PhiR->setOperand(0, Plan->getOrAddLiveIn(RdxDesc.getSentinelValue()));
     }
   }
+
+  // Check if any reduction is used by another reduction, by starting from
+  // ComputeReductionResults and checking if the recurrence descriptor is marked
+  // has having another recurrence user.
+  for (auto &R : *MiddleVPBB) {
+    auto *MinMaxResult = dyn_cast<VPInstruction>(&R);
+    if (!MinMaxResult ||
+        MinMaxResult->getOpcode() != VPInstruction::ComputeReductionResult)
+      continue;
+    auto *MinMaxPhiR = cast<VPReductionPHIRecipe>(MinMaxResult->getOperand(0));
+    auto &MinMaxRdxDesc = MinMaxPhiR->getRecurrenceDescriptor();
+    if (!MinMaxRdxDesc.IsUsedByOtherRecurrence)
+      continue;
+
+    assert(RecurrenceDescriptor::isMinMaxRecurrenceKind(
+               MinMaxRdxDesc.getRecurrenceKind()) &&
+           "only min/max reductions can be used by other recurrences");
+
+    SmallVector<VPUser *> Worklist;
+    append_range(Worklist, MinMaxPhiR->users());
+    VPReductionPHIRecipe *FindIVPhiR = nullptr;
+    // Starting from MinMaxPhiR's users, find the other reduction phi using
+    // MinMaxPhiR.
+    while (!Worklist.empty()) {
+      VPUser *Cur = Worklist.pop_back_val();
+      if (isa<VPHeaderPHIRecipe>(Cur)) {
+        if (Cur != MinMaxPhiR) {
+          assert(!FindIVPhiR &&
+                 "Only the starting MinMaxPhiR or another reduction "
+                 "phi must be reachable");
+          FindIVPhiR = cast<VPReductionPHIRecipe>(Cur);
+        }
+        continue;
+      }
+      // Skip recipes outside any region.
+      if (!cast<VPSingleDefRecipe>(Cur)->getParent()->getParent())
+        continue;
+      append_range(Worklist, cast<VPSingleDefRecipe>(Cur)->users());
+    }
+
+    // Find the  recipe computing the result of the other reduction.
+    VPInstruction *FindIVResult = nullptr;
+    for (auto *U : FindIVPhiR->users()) {
+      auto *VPI = dyn_cast<VPInstruction>(U);
+      if (!VPI || VPI->getOpcode() != VPInstruction::ComputeFindIVResult)
+        continue;
+      FindIVResult = VPI;
+    }
+    assert(FindIVResult && "must find a matching ComputeFindIVResult");
+
+    // The reduction using MinMaxPhiR needs adjusting to compute the correct
+    // result:
+    //  1. We need to find the first 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.
+    FindIVResult->moveAfter(MinMaxResult);
+    VPBuilder B(FindIVResult);
+    const_cast<RecurrenceDescriptor &>(FindIVPhiR->getRecurrenceDescriptor())
+        .setKind(RecurKind::FindFirstIVUMax);
+    auto *Cmp = B.createICmp(CmpInst::ICMP_EQ, MinMaxResult->getOperand(1),
+                             MinMaxResult);
+    auto *Sel = B.createSelect(
+        Cmp, FindIVResult->getOperand(2),
+        Plan->getOrAddLiveIn(
+            FindIVPhiR->getRecurrenceDescriptor().getSentinelValue()));
+    FindIVResult->setOperand(2, Sel);
+  }
   for (VPRecipeBase *R : ToDelete)
     R->eraseFromParent();
 
diff --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h
@@ -1804,7 +1804,8 @@ class VPHeaderPHIRecipe : public VPSingleDefRecipe, public VPPhiAccessors {
   ~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;
   }
@@ -1813,6 +1814,10 @@ class VPHeaderPHIRecipe : public VPSingleDefRecipe, public VPPhiAccessors {
     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;
diff --git a/llvm/test/Transforms/LoopVectorize/select-min-index.ll b/llvm/test/Transforms/LoopVectorize/select-min-index.ll
