How Rewriting works in MLIR
If you only remember one thing from this post: rewriting in MLIR is “find a pattern, make a change, repeat until more changes can’t be made”, with two key components: Greedy pattern application (canonicalization and local clenups), and Dialect conversion (legalize/convert regions with invariants about the legal forms of ops). TL;DR Patterns live in a RewritePatternSet and are driven by either applyPatternsGreedily (for local greedy rewrites) or applyPartial/FullConversion (for dialect conversion with legality constraints). Write patterns by subclassing OpRewritePattern<YourOp> and overriding matchAndRewrite with your logic. Rewrite safely using PatternRewriter methods to create, replace, and erase ops. Canonicalization : MLIR has a single canonicalization pass which applies all registered patterns greedily until no more matches are found. Conversion: MLIR’s conversion framework allows you to define legality constraints and convert ops from one dialect to another while preserving invariants. We do this with ConversionTarget and TypeConverter. Folding: Take series of ops complements rewriting by simplifying constant expressions during pattern application. Part 1: The moving pieces RewritePatternSet and PatternRewriter RewritePatternSet is a container for your rewrite patterns. You populate it with instances of your custom patterns. MLIR runs these patterns for you; you don’t directly loop over operations. In your pattern’s matchAndRewrite, you Inspect the matched op. Optionally create the new IR (using the rewriter’s insertion point). Replace or erase the matched op. Greedy vs. Conversion Greedy (Canonicalization and Local Rewrites) Think “peephole + algebraic simplification”. Use applyPatternsGreedily to apply all patterns in a RewritePatternSet. applyPatternsGreedily(fop, std::move(patterns)); Conversion (Dialect Conversion) Define legality constraints for ops via ConversionTarget. Use TypeConverter to handle type conversions. Use applyPartialConversion or applyFullConversion to convert ops while respecting legality. Part 2: Your first greedy rewrite pattern. Let’s fold away arith.addi %x, 0: i32 into just %x. Yeah, it’s trivial, and MLIR’s canonicalization already does this, but it’s a great starting point. ...