[llvm] [LoopVectorizer] Add support for chaining partial reductions (PR #120272)

[via llvm-commits]

github-actions[bot] wrote:

<!--LLVM CODE FORMAT COMMENT: {clang-format}-->


:warning: C/C++ code formatter, clang-format found issues in your code. :warning:

<details>
<summary>
You can test this locally with the following command:
</summary>

``````````bash
git-clang-format --diff e33f456ae591559883e89a1f18b2dec21225e90f 34d5f25a026e6bdb337a1ba8e1a2cf7a8a4291d5 --extensions h,cpp -- llvm/lib/Target/AArch64/AArch64TargetTransformInfo.h llvm/lib/Transforms/Vectorize/LoopVectorize.cpp llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h llvm/lib/Transforms/Vectorize/VPlan.h llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
``````````

</details>

<details>
<summary>
View the diff from clang-format here.
</summary>

``````````diff
diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index 79be3e1559..a18cf3c9be 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -8824,11 +8824,10 @@ void VPRecipeBuilder::collectScaledReductions(VFRange &Range) {
 }
 
 std::optional<SmallVector<std::pair<PartialReductionChain, unsigned>>>
-VPRecipeBuilder::getScaledReduction(Instruction *PHI,
-                                    Instruction *RdxExitInstr,
+VPRecipeBuilder::getScaledReduction(Instruction *PHI, Instruction *RdxExitInstr,
                                     VFRange &Range) {
 
-  if(!CM.TheLoop->contains(RdxExitInstr))
+  if (!CM.TheLoop->contains(RdxExitInstr))
     return std::nullopt;
 
   // TODO: Allow scaling reductions when predicating. The select at
@@ -8850,7 +8849,7 @@ VPRecipeBuilder::getScaledReduction(Instruction *PHI,
   SmallVector<std::pair<PartialReductionChain, unsigned>> Chains;
 
   if (auto *OpInst = dyn_cast<Instruction>(Op)) {
-    if(auto SR0 = getScaledReduction(PHI, OpInst, Range)) {
+    if (auto SR0 = getScaledReduction(PHI, OpInst, Range)) {
       Chains.append(*SR0);
       PHI = SR0->rbegin()->first.Reduction;
 
@@ -9000,7 +8999,8 @@ VPRecipeBuilder::tryToCreatePartialReduction(Instruction *Reduction,
   VPValue *BinOp = Operands[0];
   VPValue *Phi = Operands[1];
   VPRecipeBase *BinOpRecipe = BinOp->getDefiningRecipe();
-  if (isa<VPReductionPHIRecipe>(BinOpRecipe) || isa<VPPartialReductionRecipe>(BinOpRecipe))
+  if (isa<VPReductionPHIRecipe>(BinOpRecipe) ||
+      isa<VPPartialReductionRecipe>(BinOpRecipe))
     std::swap(BinOp, Phi);
 
   return new VPPartialReductionRecipe(Reduction->getOpcode(), BinOp, Phi,
diff --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h
index b2d3d3944c..9e09bfc621 100644
--- a/llvm/lib/Transforms/Vectorize/VPlan.h
+++ b/llvm/lib/Transforms/Vectorize/VPlan.h
@@ -2454,13 +2454,15 @@ public:
                           ArrayRef<VPValue *>({Op0, Op1}), ReductionInst),
         Opcode(Opcode) {
     auto *DefiningRecipe = getOperand(1)->getDefiningRecipe();
-    assert((isa<VPReductionPHIRecipe>(DefiningRecipe) || isa<VPPartialReductionRecipe>(DefiningRecipe)) &&
+    assert((isa<VPReductionPHIRecipe>(DefiningRecipe) ||
+            isa<VPPartialReductionRecipe>(DefiningRecipe)) &&
            "Unexpected operand order for partial reduction recipe");
   }
   ~VPPartialReductionRecipe() override = default;
 
   VPPartialReductionRecipe *clone() override {
-    return new VPPartialReductionRecipe(Opcode, getOperand(0), getOperand(1), getUnderlyingInstr());
+    return new VPPartialReductionRecipe(Opcode, getOperand(0), getOperand(1),
+                                        getUnderlyingInstr());
   }
 
   VP_CLASSOF_IMPL(VPDef::VPPartialReductionSC)
diff --git a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
index 4fb32c1439..668c317033 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
@@ -325,8 +325,9 @@ void VPPartialReductionRecipe::execute(VPTransformState &State) {
 
   // Currently we don't have a partial_reduce_sub intrinsic,
   // so mimic the behaviour by negating the second operand
-  if(Opcode == Instruction::Sub) {
-    BinOpVal = Builder.CreateSub(Constant::getNullValue(BinOpVal->getType()), BinOpVal);
+  if (Opcode == Instruction::Sub) {
+    BinOpVal = Builder.CreateSub(Constant::getNullValue(BinOpVal->getType()),
+                                 BinOpVal);
     Opcode = Instruction::Add;
   }
 

``````````

</details>


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