[VPlan] Match more GEP-like in m_GetElementPtr

llvm/lib/Transforms/Vectorize/VPlan.h

@@ -1799,6 +1799,9 @@ class LLVM_ABI_FOR_TEST VPWidenGEPRecipe : public VPRecipeWithIRFlags {
 
   VP_CLASSOF_IMPL(VPDef::VPWidenGEPSC)
 
+  /// This recipe generates a GEP instruction.
+  unsigned getOpcode() const { return Instruction::GetElementPtr; }
+
   /// Generate the gep nodes.
   void execute(VPTransformState &State) override;
 

llvm/lib/Transforms/Vectorize/VPlanPatternMatch.h

@@ -252,10 +252,9 @@ struct Recipe_match {
   static bool matchRecipeAndOpcode(const VPRecipeBase *R) {
     auto *DefR = dyn_cast<RecipeTy>(R);
     // Check for recipes that do not have opcodes.
-    if constexpr (std::is_same<RecipeTy, VPScalarIVStepsRecipe>::value ||
-                  std::is_same<RecipeTy, VPCanonicalIVPHIRecipe>::value ||
-                  std::is_same<RecipeTy, VPDerivedIVRecipe>::value ||
-                  std::is_same<RecipeTy, VPWidenGEPRecipe>::value)
+    if constexpr (std::is_same_v<RecipeTy, VPScalarIVStepsRecipe> ||
+                  std::is_same_v<RecipeTy, VPCanonicalIVPHIRecipe> ||
+                  std::is_same_v<RecipeTy, VPDerivedIVRecipe>)
       return DefR;
     else
       return DefR && DefR->getOpcode() == Opcode;
@@ -524,15 +523,24 @@ m_SpecificCmp(CmpPredicate MatchPred, const Op0_t &Op0, const Op1_t &Op1) {
 }
 
 template <typename Op0_t, typename Op1_t>
-using GEPLikeRecipe_match =
+using GEPLikeRecipe_match = match_combine_or<
     Recipe_match<std::tuple<Op0_t, Op1_t>, Instruction::GetElementPtr,
-                 /*Commutative*/ false, VPWidenRecipe, VPReplicateRecipe,
-                 VPWidenGEPRecipe, VPInstruction>;
+                 /*Commutative*/ false, VPReplicateRecipe, VPWidenGEPRecipe>,
+    match_combine_or<
+        VPInstruction_match<VPInstruction::PtrAdd, Op0_t, Op1_t>,
+        VPInstruction_match<VPInstruction::WidePtrAdd, Op0_t, Op1_t>>>;
 
 template <typename Op0_t, typename Op1_t>
 inline GEPLikeRecipe_match<Op0_t, Op1_t> m_GetElementPtr(const Op0_t &Op0,
                                                          const Op1_t &Op1) {
-  return GEPLikeRecipe_match<Op0_t, Op1_t>(Op0, Op1);
+  return m_CombineOr(
+      Recipe_match<std::tuple<Op0_t, Op1_t>, Instruction::GetElementPtr,
+                   /*Commutative*/ false, VPReplicateRecipe, VPWidenGEPRecipe>(
+          Op0, Op1),
+      m_CombineOr(
+          VPInstruction_match<VPInstruction::PtrAdd, Op0_t, Op1_t>(Op0, Op1),
+          VPInstruction_match<VPInstruction::WidePtrAdd, Op0_t, Op1_t>(Op0,
+                                                                       Op1)));
 }
 
 template <typename Op0_t, typename Op1_t, typename Op2_t>

llvm/unittests/Transforms/Vectorize/VPlanPatternMatchTest.cpp

@@ -51,5 +51,29 @@ TEST_F(VPPatternMatchTest, ScalarIVSteps) {
                                             m_SpecificInt(2), m_Specific(VF))));
 }
 
+TEST_F(VPPatternMatchTest, GetElementPtr) {
+  VPlan &Plan = getPlan();
+  VPBasicBlock *VPBB = Plan.createVPBasicBlock("entry");
+  VPBuilder Builder(VPBB);
+
+  IntegerType *I64Ty = IntegerType::get(C, 64);
+  VPValue *One = Plan.getOrAddLiveIn(ConstantInt::get(I64Ty, 1));
+  VPValue *Two = Plan.getOrAddLiveIn(ConstantInt::get(I64Ty, 2));
+  VPValue *Ptr =
+      Plan.getOrAddLiveIn(Constant::getNullValue(PointerType::get(C, 0)));
+
+  VPInstruction *PtrAdd = Builder.createPtrAdd(Ptr, One);
+  VPInstruction *WidePtrAdd = Builder.createWidePtrAdd(Ptr, Two);
+
+  using namespace VPlanPatternMatch;
+  ASSERT_TRUE(
+      match(PtrAdd, m_GetElementPtr(m_Specific(Ptr), m_SpecificInt(1))));
+  ASSERT_FALSE(
+      match(PtrAdd, m_GetElementPtr(m_Specific(Ptr), m_SpecificInt(2))));
+  ASSERT_TRUE(
+      match(WidePtrAdd, m_GetElementPtr(m_Specific(Ptr), m_SpecificInt(2))));
+  ASSERT_FALSE(
+      match(WidePtrAdd, m_GetElementPtr(m_Specific(Ptr), m_SpecificInt(1))));
+}
 } // namespace
 } // namespace llvm