Skip to content

Commit 36f26d4

Browse files
authored
[MLIR][Docs] Add a section for Python-defined dialects, passes and rewrite patterns in bindings (#163123)
The MLIR Python bindings now support defining new passes, new rewrite patterns (through either `RewritePatternSet` or `PDLModule`), as well as new dialects using the IRDL bindings. Adding a dedicated section to document these features would make it easier for users to discover and understand the full capabilities of the Python bindings.
1 parent 2118721 commit 36f26d4

File tree

1 file changed

+53
-0
lines changed

1 file changed

+53
-0
lines changed

mlir/docs/Bindings/Python.md

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1188,6 +1188,21 @@ which can be `import`ed from the main dialect file, i.e.
11881188
`python/mlir/dialects/<dialect-namespace>/passes.py` if it is undesirable to
11891189
make the passes available along with the dialect.
11901190

1191+
## Extending MLIR in Python
1192+
1193+
The MLIR Python bindings provide support for defining custom components in Python,
1194+
mainly including dialects, passes, and rewrite patterns.
1195+
The following sections outline how each of these can be implemented.
1196+
1197+
### Dialects
1198+
1199+
Dialects can be defined through the IRDL dialect bindings in Python.
1200+
The IRDL bindings offer a `load_dialects` function that
1201+
converts an MLIR module containing `irdl.dialect` ops into MLIR dialects.
1202+
For further details, see the documentation of [the IRDL dialect](../Dialects/IRDL.md).
1203+
1204+
### Passes
1205+
11911206
Passes can be defined as Python callables via the `PassManager.add` API.
11921207
In such case, the callable is wrapped as an `mlir::Pass` internally and
11931208
executed as part of the pass pipeline when `PassManager.run` is invoked.
@@ -1209,6 +1224,44 @@ pm.add('some-cpp-defined-passes')
12091224
pm.run(some_op)
12101225
```
12111226

1227+
### Rewrite Patterns
1228+
1229+
Rewrite patterns can be registered via the `add` method
1230+
of `mlir.rewrite.RewritePatternSet` in Python.
1231+
This method takes the operation type to be rewritten
1232+
and a Python callable that defines the *match and rewrite* logic.
1233+
Note that the Python callable should be defined so that
1234+
the rewrite is applied if and only if the match succeeds,
1235+
which corresponds to the return value being castable to `False`.
1236+
1237+
The `RewritePatternSet` can be converted into
1238+
a `FrozenRewritePatternSet` using the `freeze` method,
1239+
which can be applied to an operation through
1240+
the greedy pattern driver using `apply_patterns_and_fold_greedily`.
1241+
The following example demonstrates the typical usage:
1242+
1243+
```python
1244+
def to_muli(op, rewriter):
1245+
with rewriter.ip:
1246+
new_op = arith.muli(op.lhs, op.rhs, loc=op.location)
1247+
rewriter.replace_op(op, new_op)
1248+
1249+
patterns = RewritePatternSet()
1250+
patterns.add(arith.AddIOp, to_muli) # Rewrite arith.addi into arith.muli
1251+
patterns.add(...)
1252+
frozen = patterns.freeze()
1253+
1254+
module = ...
1255+
apply_patterns_and_fold_greedily(module, frozen)
1256+
```
1257+
1258+
The PDL dialect bindings also enable defining and generating rewrite patterns in Python.
1259+
The `mlir.rewrite.PDLModule` class accepts a module containing `pdl.pattern` ops,
1260+
which can be transformed into a `FrozenRewritePatternSet` using the `freeze` method.
1261+
This frozen set can then be applied to an operation
1262+
using the greedy rewrite pattern driver via `apply_patterns_and_fold_greedily`.
1263+
For further information, see [the PDL dialect documentation](/docs/Dialects/PDLOps/).
1264+
12121265
### Other functionality
12131266

12141267
Dialect functionality other than IR objects or passes, such as helper functions,

0 commit comments

Comments
 (0)