[llvm] [VPlan] Detect and create partial reductions in VPlan. (NFCI) (PR #167851)

[Florian Hahn via llvm-commits]

================
@@ -991,3 +996,381 @@ bool VPlanTransforms::handleMaxMinNumReductions(VPlan &Plan) {
   MiddleTerm->setOperand(0, NewCond);
   return true;
 }
+
+namespace {
+/// A VPlan-based chain of recipes that form a partial reduction.
+/// Designed to match either:
+///   reduction_bin_op (extend (A), accumulator), or
+///   reduction_bin_op (bin_op (extend (A), (extend (B))), accumulator).
+struct VPPartialReductionChain {
+  VPPartialReductionChain(VPWidenRecipe *Reduction, VPWidenCastRecipe *ExtendA,
+                          VPWidenCastRecipe *ExtendB, VPWidenRecipe *ExtendUser)
+      : Reduction(Reduction), ExtendA(ExtendA), ExtendB(ExtendB),
+        ExtendUser(ExtendUser) {}
+  /// The top-level binary operation that forms the reduction to a scalar
+  /// after the loop body.
+  VPWidenRecipe *Reduction;
+  /// The extension of each of the inner binary operation's operands.
+  VPWidenCastRecipe *ExtendA;
+  VPWidenCastRecipe *ExtendB;
+
+  /// The user of the extends that is then reduced.
+  VPWidenRecipe *ExtendUser;
+};
+
+// Helper to transform a single widen recipe into a partial reduction recipe.
+// Returns true if transformation succeeded.
+static bool transformToPartialReduction(VPWidenRecipe *WidenRecipe,
+                                        unsigned ScaleFactor, VPlan &Plan) {
+  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.
+  VPRecipeBase *BinOpRecipe = BinOp->getDefiningRecipe();
+  if (BinOpRecipe && (isa<VPReductionPHIRecipe>(BinOpRecipe) ||
+                      isa<VPPartialReductionRecipe>(BinOpRecipe)))
+    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.
+  VPRecipeBase *AccumRecipe = Accumulator->getDefiningRecipe();
+  if (!AccumRecipe || (!isa<VPReductionPHIRecipe>(AccumRecipe) &&
+                       !isa<VPPartialReductionRecipe>(AccumRecipe)))
+    return false;
+
+  if (auto *RdxPhi = dyn_cast<VPReductionPHIRecipe>(AccumRecipe)) {
+    assert(RdxPhi->getVFScaleFactor() == 1 && "scale factor must not be set");
+    RdxPhi->setVFScaleFactor(ScaleFactor);
+
+    // Update ReductionStartVector instruction scale factor.
+    VPValue *StartValue = RdxPhi->getOperand(0);
+    auto *StartInst = cast<VPInstruction>(StartValue);
+    assert(StartInst->getOpcode() == VPInstruction::ReductionStartVector);
+    auto *NewScaleFactor = Plan.getConstantInt(32, ScaleFactor);
+    StartInst->setOperand(2, NewScaleFactor);
+  }
+
+  // Handle SUB by negating the operand and using ADD for the partial reduction.
+  unsigned ReductionOpcode = WidenRecipe->getOpcode();
+  if (ReductionOpcode == Instruction::Sub) {
+    VPBuilder Builder(WidenRecipe);
+
+    // Infer the scalar type for creating the zero constant
+    Type *ElemTy = VPTypeAnalysis(Plan).inferScalarType(BinOp);
+    auto *Zero = Plan.getConstantInt(ElemTy, 0);
+
+    // Create a negation: 0 - BinOp
+    VPIRFlags Flags;
+    if (auto *I = WidenRecipe->getUnderlyingInstr())
+      Flags = VPIRFlags(*I);
+
+    auto *NegRecipe = new VPWidenRecipe(Instruction::Sub, {Zero, BinOp}, Flags,
+                                        VPIRMetadata(), DebugLoc());
+    Builder.insert(NegRecipe);
+    BinOp = NegRecipe;
+    ReductionOpcode = Instruction::Add;
+  }
+
+  VPValue *Cond = nullptr;
+  VPValue *ExitValue = nullptr;
+  if (auto *RedPhiR = dyn_cast<VPReductionPHIRecipe>(Accumulator)) {
+    ExitValue = findComputeReductionResult(RedPhiR)->getOperand(1);
+    match(ExitValue, m_Select(m_VPValue(Cond), m_VPValue(), m_VPValue()));
----------------
fhahn wrote:

We need to handle the selects here, because they will get added during predication, so here we need to look through them; they will be gone and handled by VPReductionRecipe if they are created.

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