!fixup, Address comments and fix VPReductionRecipe::computeCost

Note that we should use std::nullopt when quering cost of non-FPMathOperator instructions.

Author: ElvisWang123

llvm/lib/Transforms/Vectorize/VPlan.h

@@ -2437,14 +2437,6 @@ class VPExtendedReductionRecipe : public VPReductionRecipe {
   /// Opcode of the extend recipe will be lowered to.
   Instruction::CastOps ExtOp;
 
-public:
-  VPExtendedReductionRecipe(VPReductionRecipe *R, VPWidenCastRecipe *Ext)
-      : VPReductionRecipe(VPDef::VPExtendedReductionSC,
-                          R->getRecurrenceDescriptor(),
-                          {R->getChainOp(), Ext->getOperand(0)}, R->getCondOp(),
-                          R->isOrdered(), Ext->isNonNeg(), Ext->getDebugLoc()),
-        ExtOp(Ext->getOpcode()) {}
-
   /// For cloning VPExtendedReductionRecipe.
   VPExtendedReductionRecipe(VPExtendedReductionRecipe *ExtRed)
       : VPReductionRecipe(
@@ -2453,6 +2445,14 @@ class VPExtendedReductionRecipe : public VPReductionRecipe {
             ExtRed->isOrdered(), ExtRed->isNonNeg(), ExtRed->getDebugLoc()),
         ExtOp(ExtRed->getExtOpcode()) {}
 
+public:
+  VPExtendedReductionRecipe(VPReductionRecipe *R, VPWidenCastRecipe *Ext)
+      : VPReductionRecipe(VPDef::VPExtendedReductionSC,
+                          R->getRecurrenceDescriptor(),
+                          {R->getChainOp(), Ext->getOperand(0)}, R->getCondOp(),
+                          R->isOrdered(), Ext->isNonNeg(), Ext->getDebugLoc()),
+        ExtOp(Ext->getOpcode()) {}
+
   ~VPExtendedReductionRecipe() override = default;
 
   VPExtendedReductionRecipe *clone() override {
@@ -2500,8 +2500,16 @@ class VPMulAccumulateReductionRecipe : public VPReductionRecipe {
   /// Non-neg flag of the extend recipe.
   bool IsNonNeg = false;
 
-  /// Is this multiply-accumulate-reduction recipe contains extend?
-  bool IsExtended = false;
+  /// For cloning VPMulAccumulateReductionRecipe.
+  VPMulAccumulateReductionRecipe(VPMulAccumulateReductionRecipe *MulAcc)
+      : VPReductionRecipe(
+            VPDef::VPMulAccumulateReductionSC,
+            MulAcc->getRecurrenceDescriptor(),
+            {MulAcc->getChainOp(), MulAcc->getVecOp0(), MulAcc->getVecOp1()},
+            MulAcc->getCondOp(), MulAcc->isOrdered(),
+            MulAcc->hasNoUnsignedWrap(), MulAcc->hasNoSignedWrap(),
+            MulAcc->getDebugLoc()),
+        ExtOp(MulAcc->getExtOpcode()), IsNonNeg(MulAcc->isNonNeg()) {}
 
 public:
   VPMulAccumulateReductionRecipe(VPReductionRecipe *R, VPWidenRecipe *Mul,
@@ -2516,32 +2524,20 @@ class VPMulAccumulateReductionRecipe : public VPReductionRecipe {
     assert(getRecurrenceDescriptor().getOpcode() == Instruction::Add &&
            "The reduction instruction in MulAccumulateteReductionRecipe must "
            "be Add");
-    IsExtended = true;
   }
 
   VPMulAccumulateReductionRecipe(VPReductionRecipe *R, VPWidenRecipe *Mul)
       : VPReductionRecipe(
             VPDef::VPMulAccumulateReductionSC, R->getRecurrenceDescriptor(),
             {R->getChainOp(), Mul->getOperand(0), Mul->getOperand(1)},
             R->getCondOp(), R->isOrdered(), Mul->hasNoUnsignedWrap(),
-            Mul->hasNoSignedWrap(), R->getDebugLoc()) {
+            Mul->hasNoSignedWrap(), R->getDebugLoc()),
+        ExtOp(Instruction::CastOps::CastOpsEnd) {
     assert(getRecurrenceDescriptor().getOpcode() == Instruction::Add &&
            "The reduction instruction in MulAccumulateReductionRecipe must be "
            "Add");
   }
 
-  /// For cloning VPMulAccumulateReductionRecipe.
-  VPMulAccumulateReductionRecipe(VPMulAccumulateReductionRecipe *MulAcc)
-      : VPReductionRecipe(
-            VPDef::VPMulAccumulateReductionSC,
-            MulAcc->getRecurrenceDescriptor(),
-            {MulAcc->getChainOp(), MulAcc->getVecOp0(), MulAcc->getVecOp1()},
-            MulAcc->getCondOp(), MulAcc->isOrdered(),
-            MulAcc->hasNoUnsignedWrap(), MulAcc->hasNoSignedWrap(),
-            MulAcc->getDebugLoc()),
-        ExtOp(MulAcc->getExtOpcode()), IsNonNeg(MulAcc->isNonNeg()),
-        IsExtended(MulAcc->isExtended()) {}
-
   ~VPMulAccumulateReductionRecipe() override = default;
 
   VPMulAccumulateReductionRecipe *clone() override {
@@ -2571,7 +2567,7 @@ class VPMulAccumulateReductionRecipe : public VPReductionRecipe {
   VPValue *getVecOp1() const { return getOperand(2); }
 
   /// Return if this MulAcc recipe contains extend instructions.
-  bool isExtended() const { return IsExtended; }
+  bool isExtended() const { return ExtOp != Instruction::CastOps::CastOpsEnd; }
 
   /// Return if the operands of mul instruction come from same extend.
   bool isSameExtend() const { return getVecOp0() == getVecOp1(); }
@@ -2580,11 +2576,7 @@ class VPMulAccumulateReductionRecipe : public VPReductionRecipe {
   Instruction::CastOps getExtOpcode() const { return ExtOp; }
 
   /// Return if the extend opcode is ZExt.
-  bool isZExt() const {
-    if (!isExtended())
-      return true;
-    return ExtOp == Instruction::CastOps::ZExt;
-  }
+  bool isZExt() const { return ExtOp == Instruction::CastOps::ZExt; }
 
   /// Return the non negative flag of the ext recipe.
   bool isNonNeg() const { return IsNonNeg; }

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

@@ -2399,8 +2399,13 @@ InstructionCost VPReductionRecipe::computeCost(ElementCount VF,
            Ctx.TTI.getMinMaxReductionCost(Id, VectorTy, FMFs, Ctx.CostKind);
   }
 
-  return Ctx.TTI.getArithmeticReductionCost(Opcode, VectorTy, FMFs,
-                                                   Ctx.CostKind);
+  if (ElementTy->isFloatingPointTy())
+    return Ctx.TTI.getArithmeticReductionCost(Opcode, VectorTy, FMFs,
+                                              Ctx.CostKind);
+  // Cannot get correct cost when quering TTI with FMFs not contains `reassoc`
+  // for non-FP reductions.
+  return Ctx.TTI.getArithmeticReductionCost(Opcode, VectorTy, std::nullopt,
+                                            Ctx.CostKind);
 }
 
 InstructionCost

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

@@ -2146,7 +2146,7 @@ expandVPMulAccumulateReduction(VPMulAccumulateReductionRecipe *MulAcc) {
   }
 
   // Generate VPWidenRecipe.
-  std::array<VPValue *, 2> MulOps = {Op0, Op1};
+  ArrayRef<VPValue *> MulOps = {Op0, Op1};
   auto *Mul = new VPWidenRecipe(
       Instruction::Mul, make_range(MulOps.begin(), MulOps.end()),
       MulAcc->hasNoUnsignedWrap(), MulAcc->hasNoSignedWrap(),