[llvm] [LV] Bundle partial reductions inside VPExpressionRecipe (PR #147302)

[Sander de Smalen via llvm-commits]

================
@@ -2841,11 +2851,18 @@ InstructionCost VPExpressionRecipe::computeCost(ElementCount VF,
       cast<VPReductionRecipe>(ExpressionRecipes.back())->getRecurrenceKind());
   switch (ExpressionType) {
   case ExpressionTypes::ExtendedReduction: {
+    unsigned Opcode = RecurrenceDescriptor::getOpcode(
+        cast<VPReductionRecipe>(ExpressionRecipes[1])->getRecurrenceKind());
+    auto *ExtR = cast<VPWidenCastRecipe>(ExpressionRecipes[0]);
+    if (isa<VPPartialReductionRecipe>(ExpressionRecipes.back())) {
+      return Ctx.TTI.getPartialReductionCost(
+          Opcode, Ctx.Types.inferScalarType(getOperand(0)), nullptr, RedTy, VF,
+          TargetTransformInfo::getPartialReductionExtendKind(ExtR->getOpcode()),
+          TargetTransformInfo::PR_None, std::nullopt, Ctx.CostKind);
+    }
     return Ctx.TTI.getExtendedReductionCost(
-        Opcode,
-        cast<VPWidenCastRecipe>(ExpressionRecipes.front())->getOpcode() ==
-            Instruction::ZExt,
-        RedTy, SrcVecTy, std::nullopt, Ctx.CostKind);
+        Opcode, ExtR->getOpcode() == Instruction::ZExt, RedTy, SrcVecTy,
+        std::nullopt, Ctx.CostKind);
----------------
sdesmalen-arm wrote:

nit:
```suggestion
    return isa<VPPartialReductionRecipe>(ExpressionRecipes.back())
               ? Ctx.TTI.getPartialReductionCost(
                     Opcode, Ctx.Types.inferScalarType(getOperand(0)), nullptr,
                     RedTy, VF,
                     TargetTransformInfo::getPartialReductionExtendKind(
                         ExtR->getOpcode()),
                     TargetTransformInfo::PR_None, std::nullopt, Ctx.CostKind)
               : Ctx.TTI.getExtendedReductionCost(
                     Opcode, ExtR->getOpcode() == Instruction::ZExt, RedTy,
                     SrcVecTy, std::nullopt, Ctx.CostKind);
```

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