[mlir][IR] Move match and rewrite functions into separate class (llvm#129861)

The vast majority of rewrite / conversion patterns uses a combined
`matchAndRewrite` instead of separate `match` and `rewrite` functions.

This PR optimizes the code base for the most common case where users
implement a combined `matchAndRewrite`. There are no longer any `match`
and `rewrite` functions in `RewritePattern`, `ConversionPattern` and
their derived classes. Instead, there is a `SplitMatchAndRewriteImpl`
class that implements `matchAndRewrite` in terms of `match` and
`rewrite`.

Details:
* The `RewritePattern` and `ConversionPattern` classes are simpler
(fewer functions). Especially the `ConversionPattern` class, which now
has 5 fewer functions. (There were various `rewrite` overloads to
account for 1:1 / 1:N patterns.)
* There is a new class `SplitMatchAndRewriteImpl` that derives from
`RewritePattern` / `OpRewritePatern` / ..., along with a type alias
`RewritePattern::SplitMatchAndRewrite` for convenience.
* Fewer `llvm_unreachable` are needed throughout the code base. Instead,
we can use pure virtual functions. (In cases where users previously had
to implement `rewrite` or `matchAndRewrite`, etc.)
* This PR may also improve the number of [`-Woverload-virtual`
warnings](https://discourse.llvm.org/t/matchandrewrite-hiding-virtual-functions/84933)
that are produced by GCC. (To be confirmed...)

Note for LLVM integration: Patterns with separate `match` / `rewrite`
implementations, must derive from `X::SplitMatchAndRewrite` instead of
`X`.

---------

Co-authored-by: River Riddle <riddleriver@gmail.com>

Author: matthias-springer

flang/include/flang/Optimizer/CodeGen/FIROpPatterns.h

@@ -187,7 +187,6 @@ class ConvertFIRToLLVMPattern : public mlir::ConvertToLLVMPattern {
 
   const fir::FIRToLLVMPassOptions &options;
 
-  using ConvertToLLVMPattern::match;
   using ConvertToLLVMPattern::matchAndRewrite;
 };
 
@@ -206,20 +205,6 @@ class FIROpConversion : public ConvertFIRToLLVMPattern {
                                 options, benefit) {}
 
   /// Wrappers around the RewritePattern methods that pass the derived op type.
-  void rewrite(mlir::Operation *op, mlir::ArrayRef<mlir::Value> operands,
-               mlir::ConversionPatternRewriter &rewriter) const final {
-    rewrite(mlir::cast<SourceOp>(op),
-            OpAdaptor(operands, mlir::cast<SourceOp>(op)), rewriter);
-  }
-  void rewrite(mlir::Operation *op, mlir::ArrayRef<mlir::ValueRange> operands,
-               mlir::ConversionPatternRewriter &rewriter) const final {
-    auto sourceOp = llvm::cast<SourceOp>(op);
-    rewrite(llvm::cast<SourceOp>(op), OneToNOpAdaptor(operands, sourceOp),
-            rewriter);
-  }
-  llvm::LogicalResult match(mlir::Operation *op) const final {
-    return match(mlir::cast<SourceOp>(op));
-  }
   llvm::LogicalResult
   matchAndRewrite(mlir::Operation *op, mlir::ArrayRef<mlir::Value> operands,
                   mlir::ConversionPatternRewriter &rewriter) const final {
@@ -235,28 +220,12 @@ class FIROpConversion : public ConvertFIRToLLVMPattern {
     return matchAndRewrite(sourceOp, OneToNOpAdaptor(operands, sourceOp),
                            rewriter);
   }
-  /// Rewrite and Match methods that operate on the SourceOp type. These must be
+  /// Methods that operate on the SourceOp type. These must be
   /// overridden by the derived pattern class.
-  virtual llvm::LogicalResult match(SourceOp op) const {
-    llvm_unreachable("must override match or matchAndRewrite");
-  }
-  virtual void rewrite(SourceOp op, OpAdaptor adaptor,
-                       mlir::ConversionPatternRewriter &rewriter) const {
-    llvm_unreachable("must override rewrite or matchAndRewrite");
-  }
-  virtual void rewrite(SourceOp op, OneToNOpAdaptor adaptor,
-                       mlir::ConversionPatternRewriter &rewriter) const {
-    llvm::SmallVector<mlir::Value> oneToOneOperands =
-        getOneToOneAdaptorOperands(adaptor.getOperands());
-    rewrite(op, OpAdaptor(oneToOneOperands, adaptor), rewriter);
-  }
   virtual llvm::LogicalResult
   matchAndRewrite(SourceOp op, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const {
-    if (mlir::failed(match(op)))
-      return mlir::failure();
-    rewrite(op, adaptor, rewriter);
-    return mlir::success();
+    llvm_unreachable("matchAndRewrite is not implemented");
   }
   virtual llvm::LogicalResult
   matchAndRewrite(SourceOp op, OneToNOpAdaptor adaptor,
@@ -268,7 +237,6 @@ class FIROpConversion : public ConvertFIRToLLVMPattern {
 
 private:
   using ConvertFIRToLLVMPattern::matchAndRewrite;
-  using ConvertToLLVMPattern::match;
 };
 
 /// FIR conversion pattern template

mlir/docs/PatternRewriter.md

@@ -38,22 +38,23 @@ possible cost and use the predicate to guard the match.
 ### Root Operation Name (Optional)
 
 The name of the root operation that this pattern matches against. If specified,
-only operations with the given root name will be provided to the `match` and
-`rewrite` implementation. If not specified, any operation type may be provided.
-The root operation name should be provided whenever possible, because it
-simplifies the analysis of patterns when applying a cost model. To match any
+only operations with the given root name will be provided to the
+`matchAndRewrite` implementation. If not specified, any operation type may be
+provided. The root operation name should be provided whenever possible, because
+it simplifies the analysis of patterns when applying a cost model. To match any
 operation type, a special tag must be provided to make the intent explicit:
 `MatchAnyOpTypeTag`.
 
-### `match` and `rewrite` implementation
+### `matchAndRewrite` implementation
 
 This is the chunk of code that matches a given root `Operation` and performs a
 rewrite of the IR. A `RewritePattern` can specify this implementation either via
-separate `match` and `rewrite` methods, or via a combined `matchAndRewrite`
-method. When using the combined `matchAndRewrite` method, no IR mutation should
-take place before the match is deemed successful. The combined `matchAndRewrite`
-is useful when non-trivially recomputable information is required by the
-matching and rewriting phase. See below for examples:
+the `matchAndRewrite` method or via separate `match` and `rewrite` methods when
+deriving from `RewritePattern::SplitMatchAndRewrite`. When using the combined
+`matchAndRewrite` method, no IR mutation should take place before the match is
+deemed successful. The combined `matchAndRewrite` is useful when non-trivially
+recomputable information is required by the matching and rewriting phase. See
+below for examples:
 
 ```c++
 class MyPattern : public RewritePattern {
@@ -105,6 +106,10 @@ Within the `rewrite` section of a pattern, the following constraints apply:
     `eraseOp`) should be used instead.
 *   The root operation is required to either be: updated in-place, replaced, or
     erased.
+*   `matchAndRewrite` must return "success" if and only if the IR was modified.
+    `match` must return "success" if and only if the IR is going to be modified
+    during `rewrite`.
+
 
 ### Application Recursion
 

mlir/include/mlir/Conversion/LLVMCommon/Pattern.h

@@ -40,6 +40,9 @@ LogicalResult oneToOneRewrite(
 /// during the entire pattern lifetime.
 class ConvertToLLVMPattern : public ConversionPattern {
 public:
+  using SplitMatchAndRewrite =
+      detail::ConversionSplitMatchAndRewriteImpl<ConvertToLLVMPattern>;
+
   ConvertToLLVMPattern(StringRef rootOpName, MLIRContext *context,
                        const LLVMTypeConverter &typeConverter,
                        PatternBenefit benefit = 1);
@@ -142,9 +145,12 @@ class ConvertToLLVMPattern : public ConversionPattern {
 template <typename SourceOp>
 class ConvertOpToLLVMPattern : public ConvertToLLVMPattern {
 public:
+  using OperationT = SourceOp;
   using OpAdaptor = typename SourceOp::Adaptor;
   using OneToNOpAdaptor =
       typename SourceOp::template GenericAdaptor<ArrayRef<ValueRange>>;
+  using SplitMatchAndRewrite = detail::ConversionSplitMatchAndRewriteImpl<
+      ConvertOpToLLVMPattern<SourceOp>>;
 
   explicit ConvertOpToLLVMPattern(const LLVMTypeConverter &typeConverter,
                                   PatternBenefit benefit = 1)
@@ -153,19 +159,6 @@ class ConvertOpToLLVMPattern : public ConvertToLLVMPattern {
                              benefit) {}
 
   /// Wrappers around the RewritePattern methods that pass the derived op type.
-  void rewrite(Operation *op, ArrayRef<Value> operands,
-               ConversionPatternRewriter &rewriter) const final {
-    auto sourceOp = cast<SourceOp>(op);
-    rewrite(sourceOp, OpAdaptor(operands, sourceOp), rewriter);
-  }
-  void rewrite(Operation *op, ArrayRef<ValueRange> operands,
-               ConversionPatternRewriter &rewriter) const final {
-    auto sourceOp = cast<SourceOp>(op);
-    rewrite(sourceOp, OneToNOpAdaptor(operands, sourceOp), rewriter);
-  }
-  LogicalResult match(Operation *op) const final {
-    return match(cast<SourceOp>(op));
-  }
   LogicalResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const final {
@@ -180,28 +173,12 @@ class ConvertOpToLLVMPattern : public ConvertToLLVMPattern {
                            rewriter);
   }
 
-  /// Rewrite and Match methods that operate on the SourceOp type. These must be
+  /// Methods that operate on the SourceOp type. One of these must be
   /// overridden by the derived pattern class.
-  virtual LogicalResult match(SourceOp op) const {
-    llvm_unreachable("must override match or matchAndRewrite");
-  }
-  virtual void rewrite(SourceOp op, OpAdaptor adaptor,
-                       ConversionPatternRewriter &rewriter) const {
-    llvm_unreachable("must override rewrite or matchAndRewrite");
-  }
-  virtual void rewrite(SourceOp op, OneToNOpAdaptor adaptor,
-                       ConversionPatternRewriter &rewriter) const {
-    SmallVector<Value> oneToOneOperands =
-        getOneToOneAdaptorOperands(adaptor.getOperands());
-    rewrite(op, OpAdaptor(oneToOneOperands, adaptor), rewriter);
-  }
   virtual LogicalResult
   matchAndRewrite(SourceOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const {
-    if (failed(match(op)))
-      return failure();
-    rewrite(op, adaptor, rewriter);
-    return success();
+    llvm_unreachable("matchAndRewrite is not implemented");
   }
   virtual LogicalResult
   matchAndRewrite(SourceOp op, OneToNOpAdaptor adaptor,
@@ -212,7 +189,6 @@ class ConvertOpToLLVMPattern : public ConvertToLLVMPattern {
   }
 
 private:
-  using ConvertToLLVMPattern::match;
   using ConvertToLLVMPattern::matchAndRewrite;
 };
 

mlir/include/mlir/IR/PatternMatch.h

@@ -234,41 +234,52 @@ class Pattern {
 // RewritePattern
 //===----------------------------------------------------------------------===//
 
-/// RewritePattern is the common base class for all DAG to DAG replacements.
-/// There are two possible usages of this class:
-///   * Multi-step RewritePattern with "match" and "rewrite"
-///     - By overloading the "match" and "rewrite" functions, the user can
-///       separate the concerns of matching and rewriting.
-///   * Single-step RewritePattern with "matchAndRewrite"
-///     - By overloading the "matchAndRewrite" function, the user can perform
-///       the rewrite in the same call as the match.
-///
-class RewritePattern : public Pattern {
-public:
-  virtual ~RewritePattern() = default;
+namespace detail {
+/// Helper class that derives from a RewritePattern class and provides separate
+/// `match` and `rewrite` entry points instead of a combined `matchAndRewrite`.
+template <typename PatternT>
+class SplitMatchAndRewriteImpl : public PatternT {
+  using PatternT::PatternT;
+
+  /// Attempt to match against IR rooted at the specified operation, which is
+  /// the same operation kind as getRootKind().
+  ///
+  /// Note: This function must not modify the IR.
+  virtual LogicalResult match(typename PatternT::OperationT op) const = 0;
 
   /// Rewrite the IR rooted at the specified operation with the result of
   /// this pattern, generating any new operations with the specified
-  /// builder.  If an unexpected error is encountered (an internal
-  /// compiler error), it is emitted through the normal MLIR diagnostic
-  /// hooks and the IR is left in a valid state.
-  virtual void rewrite(Operation *op, PatternRewriter &rewriter) const;
-
-  /// Attempt to match against code rooted at the specified operation,
-  /// which is the same operation code as getRootKind().
-  virtual LogicalResult match(Operation *op) const;
+  /// rewriter.
+  virtual void rewrite(typename PatternT::OperationT op,
+                       PatternRewriter &rewriter) const = 0;
 
-  /// Attempt to match against code rooted at the specified operation,
-  /// which is the same operation code as getRootKind(). If successful, this
-  /// function will automatically perform the rewrite.
-  virtual LogicalResult matchAndRewrite(Operation *op,
-                                        PatternRewriter &rewriter) const {
+  LogicalResult matchAndRewrite(typename PatternT::OperationT op,
+                                PatternRewriter &rewriter) const final {
     if (succeeded(match(op))) {
       rewrite(op, rewriter);
       return success();
     }
     return failure();
   }
+};
+} // namespace detail
+
+/// RewritePattern is the common base class for all DAG to DAG replacements.
+class RewritePattern : public Pattern {
+public:
+  using OperationT = Operation *;
+  using SplitMatchAndRewrite = detail::SplitMatchAndRewriteImpl<RewritePattern>;
+
+  virtual ~RewritePattern() = default;
+
+  /// Attempt to match against code rooted at the specified operation,
+  /// which is the same operation code as getRootKind(). If successful, perform
+  /// the rewrite.
+  ///
+  /// Note: Implementations must modify the IR if and only if the function
+  /// returns "success".
+  virtual LogicalResult matchAndRewrite(Operation *op,
+                                        PatternRewriter &rewriter) const = 0;
 
   /// This method provides a convenient interface for creating and initializing
   /// derived rewrite patterns of the given type `T`.
@@ -317,36 +328,19 @@ namespace detail {
 /// class or Interface.
 template <typename SourceOp>
 struct OpOrInterfaceRewritePatternBase : public RewritePattern {
+  using OperationT = SourceOp;
   using RewritePattern::RewritePattern;
 
-  /// Wrappers around the RewritePattern methods that pass the derived op type.
-  void rewrite(Operation *op, PatternRewriter &rewriter) const final {
-    rewrite(cast<SourceOp>(op), rewriter);
-  }
-  LogicalResult match(Operation *op) const final {
-    return match(cast<SourceOp>(op));
-  }
+  /// Wrapper around the RewritePattern method that passes the derived op type.
   LogicalResult matchAndRewrite(Operation *op,
                                 PatternRewriter &rewriter) const final {
     return matchAndRewrite(cast<SourceOp>(op), rewriter);
   }
 
-  /// Rewrite and Match methods that operate on the SourceOp type. These must be
-  /// overridden by the derived pattern class.
-  virtual void rewrite(SourceOp op, PatternRewriter &rewriter) const {
-    llvm_unreachable("must override rewrite or matchAndRewrite");
-  }
-  virtual LogicalResult match(SourceOp op) const {
-    llvm_unreachable("must override match or matchAndRewrite");
-  }
+  /// Method that operates on the SourceOp type. Must be overridden by the
+  /// derived pattern class.
   virtual LogicalResult matchAndRewrite(SourceOp op,
-                                        PatternRewriter &rewriter) const {
-    if (succeeded(match(op))) {
-      rewrite(op, rewriter);
-      return success();
-    }
-    return failure();
-  }
+                                        PatternRewriter &rewriter) const = 0;
 };
 } // namespace detail
 
@@ -356,6 +350,9 @@ struct OpOrInterfaceRewritePatternBase : public RewritePattern {
 template <typename SourceOp>
 struct OpRewritePattern
     : public detail::OpOrInterfaceRewritePatternBase<SourceOp> {
+  using SplitMatchAndRewrite =
+      detail::SplitMatchAndRewriteImpl<OpRewritePattern<SourceOp>>;
+
   /// Patterns must specify the root operation name they match against, and can
   /// also specify the benefit of the pattern matching and a list of generated
   /// ops.
@@ -371,6 +368,9 @@ struct OpRewritePattern
 template <typename SourceOp>
 struct OpInterfaceRewritePattern
     : public detail::OpOrInterfaceRewritePatternBase<SourceOp> {
+  using SplitMatchAndRewrite =
+      detail::SplitMatchAndRewriteImpl<OpInterfaceRewritePattern<SourceOp>>;
+
   OpInterfaceRewritePattern(MLIRContext *context, PatternBenefit benefit = 1)
       : detail::OpOrInterfaceRewritePatternBase<SourceOp>(
             Pattern::MatchInterfaceOpTypeTag(), SourceOp::getInterfaceID(),