[mlir] Allow specifying benefit for C func ptr style patterns.

[jpienaar]

Could you add test coverage?

[wecing]

Add test and rebase to main.

[rriddle]

Nice! Thanks for adding test coverage as well.

[rriddle]

mlir/unittests/IR/PatternMatchTest.cpp

29–49

[wecing]

Group unit tests in anonymous namespaces.

[wecing]

Thank you for your feedbacks!

Since I don't have write access, could someone land this change for me? Thanks!

[jpienaar]

Sure I'll help submit

mlir/include/mlir/IR/PatternMatch.h

public:
   /// `this` for chaining insertions.
   RewritePatternSet &add(PDLPatternModule &&pattern) {
     pdlPatterns.mergeIn(std::move(pattern));
     return *this;
   }
 
   // Add a matchAndRewrite style pattern represented as a C function pointer.
   template <typename OpType>
-  RewritePatternSet &add(LogicalResult (*implFn)(OpType,
-                                                 PatternRewriter &rewriter)) {
+  RewritePatternSet &
+  add(LogicalResult (*implFn)(OpType, PatternRewriter &rewriter),
+      PatternBenefit benefit = 1, ArrayRef<StringRef> generatedNames = {}) {
     struct FnPattern final : public OpRewritePattern<OpType> {
       FnPattern(LogicalResult (*implFn)(OpType, PatternRewriter &rewriter),
-                MLIRContext *context)
-          : OpRewritePattern<OpType>(context), implFn(implFn) {}
+                MLIRContext *context, PatternBenefit benefit,
+                ArrayRef<StringRef> generatedNames)
+          : OpRewritePattern<OpType>(context, benefit, generatedNames),
+            implFn(implFn) {}
 
       LogicalResult matchAndRewrite(OpType op,
                                     PatternRewriter &rewriter) const override {
         return implFn(op, rewriter);
       }
 
     private:
       LogicalResult (*implFn)(OpType, PatternRewriter &rewriter);
     };
-    add(std::make_unique<FnPattern>(std::move(implFn), getContext()));
+    add(std::make_unique<FnPattern>(std::move(implFn), getContext(), benefit,
+                                    generatedNames));
     return *this;
   }
 
   //===--------------------------------------------------------------------===//
   // Pattern Insertion
   //===--------------------------------------------------------------------===//
 
   // TODO: These are soft deprecated in favor of the 'add' methods above.

mlir/unittests/IR/PatternMatchTest.cpp

 };
 TEST(OpRewritePatternTest, GetGeneratedNames) {
   MLIRContext context;
   AnOpRewritePattern pattern(&context);
   ArrayRef<OperationName> ops = pattern.getGeneratedOps();
 
   ASSERT_EQ(ops.size(), 1u);
   ASSERT_EQ(ops.front().getStringRef(), test::OpB::getOperationName());
 }
 } // end anonymous namespace
+
+namespace {
+LogicalResult anOpRewritePatternFunc(test::OpA op, PatternRewriter &rewriter) {
+  return failure();
+}
+TEST(AnOpRewritePatternTest, PatternFuncAttributes) {
+  MLIRContext context;
+  RewritePatternSet patterns(&context);
+
+  patterns.add(anOpRewritePatternFunc, /*benefit=*/3,
+               /*generatedNames=*/{test::OpB::getOperationName()});
+  ASSERT_EQ(patterns.getNativePatterns().size(), 1);
+  auto &pattern = patterns.getNativePatterns().front();
+  ASSERT_EQ(pattern->getBenefit(), 3);
+  ASSERT_EQ(pattern->getGeneratedOps().size(), 1);
+  ASSERT_EQ(pattern->getGeneratedOps().front().getStringRef(),
+            test::OpB::getOperationName());
+}
+} // end anonymous namespace