This document details the design and API of the pattern rewriting infrastructure present in MLIR, a general DAG-to-DAG transformation framework. This framework is widely used throughout MLIR for canonicalization, conversion, and general transformation.
For an introduction to DAG-to-DAG transformation, and the rationale behind this framework please take a look at the Generic DAG Rewriter Rationale.
The pattern rewriting framework can largely be decomposed into two parts: Pattern Definition and Pattern Application.
Patterns are defined by inheriting from the RewritePattern
class. This class represents the base class of all rewrite patterns within MLIR, and is comprised of the following components:
This is the expected benefit of applying a given pattern. This benefit is static upon construction of the pattern, but may be computed dynamically at pattern initialization time, e.g. allowing the benefit to be derived from domain specific information (like the target architecture). This limitation allows for performing pattern fusion and compiling patterns into an efficient state machine, and Thier, Ertl, and Krall have shown that match predicates eliminate the need for dynamically computed costs in almost all cases: you can simply instantiate the same pattern one time for each possible cost and use the predicate to guard the match.
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 operation type, a special tag must be provided to make the intent explicit: MatchAnyOpTypeTag
.
match
and rewrite
implementationThis 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:
class MyPattern : public RewritePattern { public: /// This overload constructs a pattern that only matches operations with the /// root name of `MyOp`. MyPattern(PatternBenefit benefit, MLIRContext *context) : RewritePattern(MyOp::getOperationName(), benefit, context) {} /// This overload constructs a pattern that matches any operation type. MyPattern(PatternBenefit benefit) : RewritePattern(benefit, MatchAnyOpTypeTag()) {} /// In this section, the `match` and `rewrite` implementation is specified /// using the separate hooks. LogicalResult match(Operation *op) const override { // The `match` method returns `success()` if the pattern is a match, failure // otherwise. // ... } void rewrite(Operation *op, PatternRewriter &rewriter) { // The `rewrite` method performs mutations on the IR rooted at `op` using // the provided rewriter. All mutations must go through the provided // rewriter. } /// In this section, the `match` and `rewrite` implementation is specified /// using a single hook. LogicalResult matchAndRewrite(Operation *op, PatternRewriter &rewriter) { // The `matchAndRewrite` method performs both the matching and the mutation. // Note that the match must reach a successful point before IR mutation may // take place. } };
Within the match
section of a pattern, the following constraints apply:
Within the rewrite
section of a pattern, the following constraints apply:
PatternRewriter
. This class provides hooks for performing all of the possible mutations that may take place within a pattern. For example, this means that an operation should not be erased via its erase
method. To erase an operation, the appropriate PatternRewriter
hook (in this case eraseOp
) should be used instead.Recursion is an important topic in the context of pattern rewrites, as a pattern may often be applicable to its own result. For example, imagine a pattern that peels a single iteration from a loop operation. If the loop has multiple peelable iterations, this pattern may apply multiple times during the application process. By looking at the implementation of this pattern, the bound for recursive application may be obvious, e.g. there are no peelable iterations within the loop, but from the perspective of the pattern driver this recursion is potentially dangerous. Often times the recursive application of a pattern indicates a bug in the matching logic. These types of bugs generally do not cause crashes, but create infinite loops within the application process. Given this, the pattern rewriting infrastructure conservatively assumes that no patterns have a proper bounded recursion, and will fail if recursion is detected. A pattern that is known to have proper support for handling recursion can signal this by calling setHasBoundedRewriteRecursion
when initializing the pattern. This will signal to the pattern driver that recursive application of this pattern may happen, and the pattern is equipped to safely handle it.
To aid in debugging, patterns may specify: a debug name (via setDebugName
), which should correspond to an identifier that uniquely identifies the specific pattern; and a set of debug labels (via addDebugLabels
), which correspond to identifiers that uniquely identify groups of patterns. This information is used by various utilities to aid in the debugging of pattern rewrites, e.g. in debug logs, to provide pattern filtering, etc. A simple code example is shown below:
class MyPattern : public RewritePattern { public: /// Inherit constructors from RewritePattern. using RewritePattern::RewritePattern; void initialize() { setDebugName("MyPattern"); addDebugLabels("MyRewritePass"); } // ... }; void populateMyPatterns(RewritePatternSet &patterns, MLIRContext *ctx) { // Debug labels may also be attached to patterns during insertion. This allows // for easily attaching common labels to groups of patterns. patterns.addWithLabel<MyPattern, ...>("MyRewritePatterns", ctx); }
Several pieces of pattern state require explicit initialization by the pattern, for example setting setHasBoundedRewriteRecursion
if a pattern safely handles recursive application. This pattern state can be initialized either in the constructor of the pattern or via the utility initialize
hook. Using the initialize
hook removes the need to redefine pattern constructors just to inject additional pattern state initialization. An example is shown below:
class MyPattern : public RewritePattern { public: /// Inherit the constructors from RewritePattern. using RewritePattern::RewritePattern; /// Initialize the pattern. void initialize() { /// Signal that this pattern safely handles recursive application. setHasBoundedRewriteRecursion(); } // ... };
Constructing a RewritePattern should be performed by using the static RewritePattern::create<T>
utility method. This method ensures that the pattern is properly initialized and prepared for insertion into a RewritePatternSet
.
A PatternRewriter
is a special class that allows for a pattern to communicate with the driver of pattern application. As noted above, all IR mutations, including creations, are required to be performed via the PatternRewriter
class. This is required because the underlying pattern driver may have state that would be invalidated when a mutation takes place. Examples of some of the more prevalent PatternRewriter
API is shown below, please refer to the class documentation for a more up-to-date listing of the available API:
eraseOp
This method erases an operation that either has no results, or whose results are all known to have no uses.
match
failed : notifyMatchFailure
This method allows for providing a diagnostic message within a matchAndRewrite
as to why a pattern failed to match. How this message is displayed back to the user is determined by the specific pattern driver.
replaceOp
/replaceOpWithNewOp
This method replaces an operation's results with a set of provided values, and erases the operation.
(start|cancel|finalize)RootUpdate
This is a collection of methods that provide a transaction-like API for updating the attributes, location, operands, or successors of an operation in-place within a pattern. An in-place update transaction is started with startRootUpdate
, and may either be canceled or finalized with cancelRootUpdate
and finalizeRootUpdate
respectively. A convenience wrapper, updateRootInPlace
, is provided that wraps a start
and finalize
around a callback.
The PatternRewriter
inherits from the OpBuilder
class, and thus provides all of the same functionality present within an OpBuilder
. This includes operation creation, as well as many useful attribute and type construction methods.
After a set of patterns have been defined, they are collected and provided to a specific driver for application. A driver consists of several high levels parts:
RewritePatternSet
The input patterns to a driver are provided in the form of an RewritePatternSet
. This class provides a simplified API for building a list of patterns.
PatternRewriter
To ensure that the driver state does not become invalidated by IR mutations within the pattern rewriters, a driver must provide a PatternRewriter
instance with the necessary hooks overridden. If a driver does not need to hook into certain mutations, a default implementation is provided that will perform the mutation directly.
Each driver is responsible for defining its own operation visitation order as well as pattern cost model, but the final application is performed via a PatternApplicator
class. This class takes as input the RewritePatternSet
and transforms the patterns based upon a provided cost model. This cost model computes a final benefit for a given pattern, using whatever driver specific information necessary. After a cost model has been computed, the driver may begin to match patterns against operations using PatternApplicator::matchAndRewrite
.
An example is shown below:
class MyPattern : public RewritePattern { public: MyPattern(PatternBenefit benefit, MLIRContext *context) : RewritePattern(MyOp::getOperationName(), benefit, context) {} }; /// Populate the pattern list. void collectMyPatterns(RewritePatternSet &patterns, MLIRContext *ctx) { patterns.add<MyPattern>(/*benefit=*/1, ctx); } /// Define a custom PatternRewriter for use by the driver. class MyPatternRewriter : public PatternRewriter { public: MyPatternRewriter(MLIRContext *ctx) : PatternRewriter(ctx) {} /// Override the necessary PatternRewriter hooks here. }; /// Apply the custom driver to `op`. void applyMyPatternDriver(Operation *op, const RewritePatternSet &patterns) { // Initialize the custom PatternRewriter. MyPatternRewriter rewriter(op->getContext()); // Create the applicator and apply our cost model. PatternApplicator applicator(patterns); applicator.applyCostModel([](const Pattern &pattern) { // Apply a default cost model. // Note: This is just for demonstration, if the default cost model is truly // desired `applicator.applyDefaultCostModel()` should be used // instead. return pattern.getBenefit(); }); // Try to match and apply a pattern. LogicalResult result = applicator.matchAndRewrite(op, rewriter); if (failed(result)) { // ... No patterns were applied. } // ... A pattern was successfully applied. }
MLIR provides several common pattern drivers that serve a variety of different use cases.
This driver provides a framework in which to perform operation conversions between, and within dialects using a concept of “legality”. This framework allows for transforming illegal operations to those supported by a provided conversion target, via a set of pattern-based operation rewriting patterns. This framework also provides support for type conversions. More information on this driver can be found here.
This driver walks the provided operations and greedily applies the patterns that locally have the most benefit. The benefit of a pattern is decided solely by the benefit specified on the pattern, and the relative order of the pattern within the pattern list (when two patterns have the same local benefit). Patterns are iteratively applied to operations until a fixed point is reached, at which point the driver finishes. This driver may be used via the following: applyPatternsAndFoldGreedily
and applyOpPatternsAndFold
. The latter of which only applies patterns to the provided operation, and will not traverse the IR.
The driver is configurable and supports two modes: 1) you may opt-in to a “top-down” traversal, which seeds the worklist with each operation top down and in a pre-order over the region tree. This is generally more efficient in compile time. 2) the default is a “bottom up” traversal, which builds the initial worklist with a postorder traversal of the region tree. This may match larger patterns with ambiguous pattern sets.
Note: This driver is the one used by the canonicalization pass in MLIR.
To debug the execution of the greedy pattern rewrite driver, -debug-only=greedy-rewriter
may be used. This command line flag activates LLVM's debug logging infrastructure solely for the greedy pattern rewriter. The output is formatted as a tree structure, mirroring the structure of the pattern application process. This output contains all of the actions performed by the rewriter, how operations get processed and patterns are applied, and why they fail.
Example output is shown below:
//===-------------------------------------------===// Processing operation : 'std.cond_br'(0x60f000001120) { "std.cond_br"(%arg0)[^bb2, ^bb2] {operand_segment_sizes = dense<[1, 0, 0]> : vector<3xi32>} : (i1) -> () * Pattern SimplifyConstCondBranchPred : 'std.cond_br -> ()' { } -> failure : pattern failed to match * Pattern SimplifyCondBranchIdenticalSuccessors : 'std.cond_br -> ()' { ** Insert : 'std.br'(0x60b000003690) ** Replace : 'std.cond_br'(0x60f000001120) } -> success : pattern applied successfully } -> success : pattern matched //===-------------------------------------------===//
This output is describing the processing of a std.cond_br
operation. We first try to apply the SimplifyConstCondBranchPred
, which fails. From there, another pattern (SimplifyCondBranchIdenticalSuccessors
) is applied that matches the std.cond_br
and replaces it with a std.br
.
To simplify test case definition and reduction, the FrozenRewritePatternSet
class provides built-in support for filtering which patterns should be provided to the pattern driver for application. Filtering behavior is specified by providing a disabledPatterns
and enabledPatterns
list when constructing the FrozenRewritePatternSet
. The disabledPatterns
list should contain a set of debug names or labels for patterns that are disabled during pattern application, i.e. which patterns should be filtered out. The enabledPatterns
list should contain a set of debug names or labels for patterns that are enabled during pattern application, patterns that do not satisfy this constraint are filtered out. Note that patterns specified by the disabledPatterns
list will be filtered out even if they match criteria in the enabledPatterns
list. An example is shown below:
void MyPass::initialize(MLIRContext *context) { // No patterns are explicitly disabled. SmallVector<std::string> disabledPatterns; // Enable only patterns with a debug name or label of `MyRewritePatterns`. SmallVector<std::string> enabledPatterns(1, "MyRewritePatterns"); RewritePatternSet rewritePatterns(context); // ... frozenPatterns = FrozenRewritePatternSet(rewritePatterns, disabledPatterns, enabledPatterns); }
Passes that utilize rewrite patterns should aim to provide a common set of options and toggles to simplify the debugging experience when switching between different passes/projects/etc. To aid in this endeavor, MLIR provides a common set of utilities that can be easily included when defining a custom pass. These are defined in mlir/RewritePassUtil.td
; an example usage is shown below:
def MyRewritePass : Pass<"..."> { let summary = "..."; let constructor = "createMyRewritePass()"; // Inherit the common pattern rewrite options from `RewritePassUtils`. let options = RewritePassUtils.options; }
This section documents common pass options that are useful for controlling the behavior of rewrite pattern application.
Two common pattern filtering options are exposed, disable-patterns
and enable-patterns
, matching the behavior of the disabledPatterns
and enabledPatterns
lists described in the Pattern Filtering section above. A snippet of the tablegen definition of these options is shown below:
ListOption<"disabledPatterns", "disable-patterns", "std::string", "Labels of patterns that should be filtered out during application", "llvm::cl::MiscFlags::CommaSeparated">, ListOption<"enabledPatterns", "enable-patterns", "std::string", "Labels of patterns that should be used during application, all " "other patterns are filtered out", "llvm::cl::MiscFlags::CommaSeparated">,
These options may be used to provide filtering behavior when constructing any FrozenRewritePatternSet
s within the pass:
void MyRewritePass::initialize(MLIRContext *context) { RewritePatternSet rewritePatterns(context); // ... // When constructing the `FrozenRewritePatternSet`, we provide the filter // list options. frozenPatterns = FrozenRewritePatternSet(rewritePatterns, disabledPatterns, enabledPatterns); }