!fixup address latest comments, thanks

Author: fhahn

llvm/lib/Transforms/Vectorize/VPlanPatternMatch.h

@@ -251,6 +251,23 @@ m_VPInstruction(const Op0_t &Op0, const Op1_t &Op1) {
   return BinaryVPInstruction_match<Op0_t, Op1_t, Opcode>(Op0, Op1);
 }
 
+template <typename Op0_t, typename Op1_t, typename Op2_t, unsigned Opcode,
+          bool Commutative, typename... RecipeTys>
+using TernaryRecipe_match = Recipe_match<std::tuple<Op0_t, Op1_t, Op2_t>,
+                                         Opcode, Commutative, RecipeTys...>;
+
+template <typename Op0_t, typename Op1_t, typename Op2_t, unsigned Opcode>
+using TernaryVPInstruction_match =
+    TernaryRecipe_match<Op0_t, Op1_t, Op2_t, Opcode, /*Commutative*/ false,
+                        VPInstruction>;
+
+template <unsigned Opcode, typename Op0_t, typename Op1_t, typename Op2_t>
+inline TernaryVPInstruction_match<Op0_t, Op1_t, Op2_t, Opcode>
+m_VPInstruction(const Op0_t &Op0, const Op1_t &Op1, const Op2_t &Op2) {
+  return TernaryVPInstruction_match<Op0_t, Op1_t, Op2_t, Opcode>(
+      {Op0, Op1, Op2});
+}
+
 template <typename Op0_t>
 inline UnaryVPInstruction_match<Op0_t, VPInstruction::Not>
 m_Not(const Op0_t &Op0) {

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

@@ -773,6 +773,7 @@ bool VPInstruction::opcodeMayReadOrWriteFromMemory() const {
   case VPInstruction::LogicalAnd:
   case VPInstruction::Not:
   case VPInstruction::PtrAdd:
+  case VPInstruction::WideIVStep:
     return false;
   default:
     return true;

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

@@ -968,6 +968,16 @@ static void simplifyRecipe(VPRecipeBase &R, VPTypeAnalysis &TypeInfo) {
       TypeInfo.inferScalarType(R.getOperand(1)) ==
           TypeInfo.inferScalarType(R.getVPSingleValue()))
     return R.getVPSingleValue()->replaceAllUsesWith(R.getOperand(1));
+
+  if (match(&R, m_VPInstruction<VPInstruction::WideIVStep>(
+                    m_VPValue(X), m_SpecificInt(1), m_VPValue(Y)))) {
+    if (TypeInfo.inferScalarType(X) != TypeInfo.inferScalarType(Y)) {
+      X = new VPWidenCastRecipe(Instruction::Trunc, X,
+                                TypeInfo.inferScalarType(Y));
+      X->getDefiningRecipe()->insertBefore(&R);
+    }
+    R.getVPSingleValue()->replaceAllUsesWith(X);
+  }
 }
 
 /// Try to simplify the recipes in \p Plan. Use \p CanonicalIVTy as type for all
@@ -2050,9 +2060,10 @@ void VPlanTransforms::createInterleaveGroups(
   }
 }
 
-void VPlanTransforms::convertToConcreteRecipes(VPlan &Plan) {
-  Type *CanonicalIVType = Plan.getCanonicalIV()->getScalarType();
-  VPTypeAnalysis TypeInfo(CanonicalIVType);
+void VPlanTransforms::convertToConcreteRecipes(VPlan &Plan,
+                                               Type *CanonicalIVTy) {
+  using namespace llvm::VPlanPatternMatch;
+  VPTypeAnalysis TypeInfo(CanonicalIVTy);
 
   for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly<VPBasicBlock>(
            vp_depth_first_deep(Plan.getEntry()))) {
@@ -2070,42 +2081,44 @@ void VPlanTransforms::convertToConcreteRecipes(VPlan &Plan) {
         continue;
       }
 
-      auto *VPI = dyn_cast<VPInstruction>(&R);
-      if (VPI && VPI->getOpcode() == VPInstruction::WideIVStep) {
-        VPBuilder Builder(VPI->getParent(), VPI->getIterator());
-        VPValue *VectorStep = VPI->getOperand(0);
-        Type *IVTy = TypeInfo.inferScalarType(VPI->getOperand(2));
-        if (TypeInfo.inferScalarType(VectorStep) != IVTy) {
-          Instruction::CastOps CastOp = IVTy->isFloatingPointTy()
-                                            ? Instruction::UIToFP
-                                            : Instruction::Trunc;
-          VectorStep = Builder.createWidenCast(CastOp, VectorStep, IVTy);
-        }
-
-        VPValue *ScalarStep = VPI->getOperand(1);
-        auto *ConstStep =
-            ScalarStep->isLiveIn()
-                ? dyn_cast<ConstantInt>(ScalarStep->getLiveInIRValue())
-                : nullptr;
-        if (!ConstStep || ConstStep->getValue() != 1) {
-          if (TypeInfo.inferScalarType(ScalarStep) != IVTy) {
-            ScalarStep =
-                Builder.createWidenCast(Instruction::Trunc, ScalarStep, IVTy);
-          }
-
-          std::optional<FastMathFlags> FMFs;
-          if (IVTy->isFloatingPointTy())
-            FMFs = VPI->getFastMathFlags();
+      VPValue *VectorStep;
+      VPValue *ScalarStep;
+      VPValue *IVTyOp;
+      if (!match(&R, m_VPInstruction<VPInstruction::WideIVStep>(
+                         m_VPValue(VectorStep), m_VPValue(ScalarStep),
+                         m_VPValue(IVTyOp))))
+        continue;
+      auto *VPI = cast<VPInstruction>(&R);
+      VPBuilder Builder(VPI->getParent(), VPI->getIterator());
+      Type *IVTy = TypeInfo.inferScalarType(IVTyOp);
+      if (TypeInfo.inferScalarType(VectorStep) != IVTy) {
+        Instruction::CastOps CastOp = IVTy->isFloatingPointTy()
+                                          ? Instruction::UIToFP
+                                          : Instruction::Trunc;
+        VectorStep = Builder.createWidenCast(CastOp, VectorStep, IVTy);
+      }
 
-          unsigned MulOpc =
-              IVTy->isFloatingPointTy() ? Instruction::FMul : Instruction::Mul;
-          VPInstruction *Mul = Builder.createNaryOp(
-              MulOpc, {VectorStep, ScalarStep}, FMFs, R.getDebugLoc());
-          VectorStep = Mul;
-        }
-        VPI->replaceAllUsesWith(VectorStep);
-        VPI->eraseFromParent();
+      auto *ConstStep =
+          ScalarStep->isLiveIn()
+              ? dyn_cast<ConstantInt>(ScalarStep->getLiveInIRValue())
+              : nullptr;
+      assert(!ConstStep || ConstStep->getValue() != 1);
+      if (TypeInfo.inferScalarType(ScalarStep) != IVTy) {
+        ScalarStep =
+            Builder.createWidenCast(Instruction::Trunc, ScalarStep, IVTy);
       }
+
+      std::optional<FastMathFlags> FMFs;
+      if (IVTy->isFloatingPointTy())
+        FMFs = VPI->getFastMathFlags();
+
+      unsigned MulOpc =
+          IVTy->isFloatingPointTy() ? Instruction::FMul : Instruction::Mul;
+      VPInstruction *Mul = Builder.createNaryOp(
+          MulOpc, {VectorStep, ScalarStep}, FMFs, R.getDebugLoc());
+      VectorStep = Mul;
+      VPI->replaceAllUsesWith(VectorStep);
+      VPI->eraseFromParent();
     }
   }
 }

llvm/lib/Transforms/Vectorize/VPlanTransforms.h

@@ -136,9 +136,9 @@ struct VPlanTransforms {
                                          VPRecipeBuilder &RecipeBuilder);
 
   /// Lower abstract recipes to concrete ones, that can be codegen'd.
-  static void convertToConcreteRecipes(VPlan &Plan);
+  static void convertToConcreteRecipes(VPlan &Plan, Type *CanonicalIVTy);
 
-static void simplifyRecipes(VPlan &Plan, Type *CanonicalIVTy);
+  static void simplifyRecipes(VPlan &Plan, Type *CanonicalIVTy);
 
   /// If there's a single exit block, optimize its phi recipes that use exiting
   /// IV values by feeding them precomputed end values instead, possibly taken