[llvm] [LV] Support argmin/argmax with strict predicates. (PR #170223)

[Florian Hahn via llvm-commits]

================
@@ -1120,6 +1122,133 @@ bool VPlanTransforms::handleMaxMinNumReductions(VPlan &Plan) {
   return true;
 }
 
+/// For argmin/argmax reductions with strict predicates, convert the existing
+/// FindLastIV reduction to a new UMin reduction of a wide canonical IV. If the
+/// original IV was not canonical, a new canonical wide IV is added, and the
+/// final result is scaled back to the original IV.
+static bool handleStrictArgMinArgMax(VPlan &Plan,
+                                     VPReductionPHIRecipe *MinMaxPhiR,
+                                     VPReductionPHIRecipe *FindIVPhiR,
+                                     VPWidenIntOrFpInductionRecipe *WideIV,
+                                     VPInstruction *MinMaxResult) {
+  Type *Ty = Plan.getVectorLoopRegion()->getCanonicalIVType();
+  if (Ty != VPTypeAnalysis(Plan).inferScalarType(FindIVPhiR))
+    return false;
+
+  // If the original wide IV is not canonical, create a new one. The wide IV is
+  // guaranteed to not wrap for all lanes that are active in the vector loop.
+  if (!WideIV->isCanonical()) {
+    VPValue *Zero = Plan.getOrAddLiveIn(ConstantInt::get(Ty, 0));
+    VPValue *One = Plan.getOrAddLiveIn(ConstantInt::get(Ty, 1));
+    auto *WidenCanIV = new VPWidenIntOrFpInductionRecipe(
+        nullptr, Zero, One, WideIV->getVFValue(),
+        WideIV->getInductionDescriptor(), VPIRFlags(), WideIV->getDebugLoc());
+    WidenCanIV->insertBefore(WideIV);
+
+    // Update the select to use the wide canonical IV.
+    auto *SelectRecipe = cast<VPSingleDefRecipe>(
+        FindIVPhiR->getBackedgeValue()->getDefiningRecipe());
+    if (SelectRecipe->getOperand(1) == WideIV)
+      SelectRecipe->setOperand(1, WidenCanIV);
+    else if (SelectRecipe->getOperand(2) == WideIV)
+      SelectRecipe->setOperand(2, WidenCanIV);
+  }
+
+  // Create the new UMin reduction recipe to track the minimum index.
+  assert(!FindIVPhiR->isInLoop() && !FindIVPhiR->isOrdered() &&
+         "inloop and ordered reductions not supported");
+  VPValue *MaxInt =
+      Plan.getConstantInt(APInt::getMaxValue(Ty->getIntegerBitWidth()));
+  ReductionStyle Style = RdxUnordered{FindIVPhiR->getVFScaleFactor()};
+  auto *MinIdxPhiR = new VPReductionPHIRecipe(
+      dyn_cast_or_null<PHINode>(FindIVPhiR->getUnderlyingValue()),
+      RecurKind::UMin, *MaxInt, *FindIVPhiR->getBackedgeValue(), Style,
+      FindIVPhiR->hasUsesOutsideReductionChain());
+  MinIdxPhiR->insertBefore(FindIVPhiR);
+
+  VPInstruction *FindLastIVResult =
+      findUserOf<VPInstruction::ComputeFindIVResult>(FindIVPhiR);
+  MinMaxResult->moveBefore(*FindLastIVResult->getParent(),
+                           FindLastIVResult->getIterator());
+
+  // The reduction using MinMaxPhiR needs adjusting to compute the correct
+  // result:
+  //  1. We need to find the first canonical 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 UMin reduction of the canonical wide
+  //     IV which correspond to the lanes matching the min/max reduction result.
+  //  4. Scale the final select canonical IV back to the original IV using
+  //     VPDerivedIVRecipe.
+  //  5. If the minimum value matches the start value, the condition in the
+  //     loop was never true, return the start value in that case.
+  //
+  // The original reductions need adjusting:
+  // For example, this transforms
+  // vp<%min.result> = compute-reduction-result ir<%min.val>,
+  //                                            ir<%min.val.next>
+  // vp<%find.iv.result = compute-find-iv-result ir<%min.idx>, ir<0>,
+  //                                             SENTINEL, vp<%min.idx.next>
+  //
+  // into:
+  //  vp<%min.result> = compute-reduction-result ir<%min.val>, ir<%min.val.next>
+  //  vp<%final.min.cmp> = icmp eq ir<%min.val.next>, vp<%min.result>
+  //  vp<%final.min.iv> = select vp<%final.min.cmp>, ir<%min.idx.next>, ir<-1>
+  //  vp<%13> = compute-reduction-result ir<%min.idx>, vp<%final.min.iv>
----------------
fhahn wrote:

The compute-reduction-result will do that, but also take care of handling multiple parts when unrolling

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