[llvm-branch-commits] [llvm] [LV] Only create partial reductions when profitable. (PR #181706)

[Benjamin Maxwell via llvm-branch-commits]

================
@@ -6216,20 +6197,30 @@ getScaledReductions(VPReductionPHIRecipe *RedPhiR, VPCostContext &CostCtx,
     if (!PHISize.hasKnownScalarFactor(ExtSrcSize))
       return std::nullopt;
 
-    VPPartialReductionChain Chain(
+    /// Check if a partial reduction chain is supported by the target (i.e.
+    /// does not have an invalid cost) for the given VF range. Clamps the range
+    /// and returns true if feasible for any VF.
+    VPPartialReductionChain Link(
         {UpdateR, *ExtendedOp,
          static_cast<unsigned>(PHISize.getKnownScalarFactor(ExtSrcSize)), RK});
-    if (!isValidPartialReduction(Chain, PhiType, CostCtx, Range))
+    if (!LoopVectorizationPlanner::getDecisionAndClampRange(
+            [&](ElementCount VF) {
+              InstructionCost Cost =
+                  getPartialReductionLinkCost(CostCtx, Link, VF);
+              Link.PartialReductionCost[VF] = Cost;
+              return Cost.isValid();
+            },
+            Range))
----------------
MacDue wrote:

This seems simpler (and fixes the build errors):
```diff
diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
index 226354c2bf48..fb5f3b9c576f 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
@@ -5830,8 +5830,6 @@ struct VPPartialReductionChain {
   /// This allows distinguishing between Sub and AddWithSub recurrences,
   /// when the ReductionBinOp is a Instruction::Sub.
   RecurKind RK;
-  /// The cost of the link in the reduction chain.
-  DenseMap<ElementCount, InstructionCost> PartialReductionCost;
 };
 
 static VPSingleDefRecipe *
@@ -6203,15 +6201,6 @@ getScaledReductions(VPReductionPHIRecipe *RedPhiR, VPCostContext &CostCtx,
     VPPartialReductionChain Link(
         {UpdateR, *ExtendedOp,
          static_cast<unsigned>(PHISize.getKnownScalarFactor(ExtSrcSize)), RK});
-    if (!LoopVectorizationPlanner::getDecisionAndClampRange(
-            [&](ElementCount VF) {
-              InstructionCost Cost =
-                  getPartialReductionLinkCost(CostCtx, Link, VF);
-              Link.PartialReductionCost[VF] = Cost;
-              return Cost.isValid();
-            },
-            Range))
-      return std::nullopt;
 
     Chain.push_back(Link);
     CurrentValue = PrevValue;
@@ -6274,9 +6263,11 @@ void VPlanTransforms::createPartialReductions(VPlan &Plan,
     // chain using regular reductions.
     for (const VPPartialReductionChain &Link : Chain) {
       ExtendedReductionOperand ExtendedOp = Link.ExtendedOp;
-      assert(Link.PartialReductionCost.contains(VF) &&
-             "Expected cost to have been calculated for VF");
-      PartialCost += Link.PartialReductionCost.lookup(VF);
+      InstructionCost LinkCost = getPartialReductionLinkCost(CostCtx, Link, VF);
+      if (!LinkCost.isValid())
+        return false;
+
+      PartialCost += LinkCost;
       RegularCost += Link.ReductionBinOp->computeCost(VF, CostCtx);
       if (ExtendedOp.BinOp && ExtendedOp.BinOp != Link.ReductionBinOp)
         RegularCost += ExtendedOp.BinOp->computeCost(VF, CostCtx);
```

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