[llvm] [VPlan] Reject partial reductions with invalid costs in getScaledReds. (PR #180438)

[Sander de Smalen via llvm-commits]

================
@@ -5542,137 +5542,120 @@ struct VPPartialReductionChain {
   unsigned ScaleFactor;
 };
 
-// Helper to transform a partial reduction chain into a partial reduction
-// recipe. Returns true if transformation succeeded. Checks profitability and
-// clamps VF range.
-static bool transformToPartialReduction(const VPPartialReductionChain &Chain,
-                                        VFRange &Range, VPCostContext &CostCtx,
-                                        VPlan &Plan) {
+// Transform a partial reduction chain into a partial reduction recipe.
+// Assumes profitability has been checked. Performs RAUW so chained reductions'
+// accumulators are automatically updated. Updates the PHI's scale factor when
+// the accumulator is the reduction PHI.
+static void transformToPartialReduction(const VPPartialReductionChain &Chain,
+                                        VPTypeAnalysis &TypeInfo, VPlan &Plan) {
   VPWidenRecipe *WidenRecipe = Chain.ReductionBinOp;
-  unsigned ScaleFactor = Chain.ScaleFactor;
-  assert(WidenRecipe->getNumOperands() == 2 && "Expected binary operation");
 
   VPValue *BinOp = WidenRecipe->getOperand(0);
   VPValue *Accumulator = WidenRecipe->getOperand(1);
-
-  // Swap if needed to ensure Accumulator is the PHI or partial reduction.
-  if (isa_and_present<VPReductionPHIRecipe, VPReductionRecipe>(BinOp))
+  if (isa_and_present<VPReductionPHIRecipe, VPReductionRecipe>(
+          BinOp->getDefiningRecipe()))
     std::swap(BinOp, Accumulator);
 
-  // For chained reductions, only transform if accumulator is already a PHI or
-  // partial reduction. Otherwise, it needs to be transformed first.
-  auto *AccumRecipe = Accumulator->getDefiningRecipe();
-  if (!isa_and_present<VPReductionPHIRecipe, VPReductionRecipe>(AccumRecipe))
-    return false;
+  Type *PhiType = TypeInfo.inferScalarType(Accumulator);
 
-  // Check if the partial reduction is profitable for the VF range.
-  Type *PhiType = CostCtx.Types.inferScalarType(Accumulator);
+  // Update PHI and start vector scale factors when accumulator is the PHI.
+  if (auto *RdxPhi =
+          dyn_cast<VPReductionPHIRecipe>(Accumulator->getDefiningRecipe())) {
+    assert(RdxPhi->getVFScaleFactor() == 1 && "scale factor must not be set");
+    RdxPhi->setVFScaleFactor(Chain.ScaleFactor);
+    auto *StartInst = cast<VPInstruction>(RdxPhi->getOperand(0));
+    assert(StartInst->getOpcode() == VPInstruction::ReductionStartVector &&
+           "start value must be ReductionStartVector");
+    StartInst->setOperand(2, Plan.getConstantInt(32, Chain.ScaleFactor));
+  }
 
-  // Derive extend info from the stored extends.
+  // Handle SUB by negating the operand.
+  if (WidenRecipe->getOpcode() == Instruction::Sub) {
+    VPBuilder Builder(WidenRecipe);
+    auto *Zero = Plan.getConstantInt(TypeInfo.inferScalarType(BinOp), 0);
+    VPIRFlags Flags = WidenRecipe->getUnderlyingInstr()
+                          ? VPIRFlags(*WidenRecipe->getUnderlyingInstr())
+                          : VPIRFlags();
+    auto *NegRecipe = new VPWidenRecipe(Instruction::Sub, {Zero, BinOp}, Flags,
+                                        VPIRMetadata(), DebugLoc::getUnknown());
+    Builder.insert(NegRecipe);
+    BinOp = NegRecipe;
+  }
+
+  // Look for predicated reduction's condition and replace select if found.
+  VPValue *Cond = nullptr;
+  VPReductionPHIRecipe *RdxPhi;
+  auto *Sel = cast_or_null<VPInstruction>(vputils::findUserOf(
+      WidenRecipe, m_Select(m_VPValue(Cond), m_Specific(WidenRecipe),
+                            m_ReductionPhi(RdxPhi))));
+
+  RecurKind RdxKind =
+      PhiType->isFloatingPointTy() ? RecurKind::FAdd : RecurKind::Add;
+  auto *PartialRed = new VPReductionRecipe(
+      RdxKind,
+      RdxKind == RecurKind::FAdd ? WidenRecipe->getFastMathFlags()
+                                 : FastMathFlags(),
+      WidenRecipe->getUnderlyingInstr(), Accumulator, BinOp, Cond,
+      RdxUnordered{/*VFScaleFactor=*/Chain.ScaleFactor});
+  PartialRed->insertBefore(WidenRecipe);
+
+  if (Sel)
+    Sel->replaceAllUsesWith(PartialRed);
+  WidenRecipe->replaceAllUsesWith(PartialRed);
+}
+
+/// Check if a partial reduction chain is 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 profitable for any VF.
+static bool isValidPartialReduction(const VPPartialReductionChain &Chain,
+                                    Type *PhiType, VPCostContext &CostCtx,
+                                    VFRange &Range) {
   auto GetExtInfo = [&CostCtx](VPWidenCastRecipe *Ext)
       -> std::pair<Type *, TargetTransformInfo::PartialReductionExtendKind> {
     if (!Ext)
       return {nullptr, TargetTransformInfo::PR_None};
     Type *ExtOpType = CostCtx.Types.inferScalarType(Ext->getOperand(0));
-    auto ExtKind = TargetTransformInfo::getPartialReductionExtendKind(
-        static_cast<Instruction::CastOps>(Ext->getOpcode()));
-    return {ExtOpType, ExtKind};
+    return {ExtOpType,
+            TargetTransformInfo::getPartialReductionExtendKind(
+                static_cast<Instruction::CastOps>(Ext->getOpcode()))};
   };
   auto ExtInfoA = GetExtInfo(Chain.ExtendA);
   auto ExtInfoB = GetExtInfo(Chain.ExtendB);
   Type *ExtOpTypeA = ExtInfoA.first;
   Type *ExtOpTypeB = ExtInfoB.first;
   auto ExtKindA = ExtInfoA.second;
   auto ExtKindB = ExtInfoB.second;
+
   // If ExtendB is nullptr but there's a separate BinOp, the second operand
   // was a constant that can use the same extend kind as the first.
   if (!Chain.ExtendB && Chain.BinOp && Chain.BinOp != Chain.ReductionBinOp) {
-    // Validate that the constant can be extended to the narrow type.
     const APInt *Const = nullptr;
-    for (VPValue *Op : Chain.BinOp->operands()) {
+    for (VPValue *Op : Chain.BinOp->operands())
----------------
sdesmalen-arm wrote:

unnecessary change

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