diff options
Diffstat (limited to 'mlir/docs/Bindings/Python.md')
| -rw-r--r-- | mlir/docs/Bindings/Python.md | 77 |
1 files changed, 76 insertions, 1 deletions
diff --git a/mlir/docs/Bindings/Python.md b/mlir/docs/Bindings/Python.md index 98ac635..6f778b0 100644 --- a/mlir/docs/Bindings/Python.md +++ b/mlir/docs/Bindings/Python.md @@ -1188,7 +1188,6 @@ which can be `import`ed from the main dialect file, i.e. `python/mlir/dialects/<dialect-namespace>/passes.py` if it is undesirable to make the passes available along with the dialect. - ### Other functionality Dialect functionality other than IR objects or passes, such as helper functions, @@ -1201,6 +1200,82 @@ utilities to connect to the rest of Python API. The bindings can be located in a separate module or in the same module as attributes and types, and loaded along with the dialect. + +## Extending MLIR in Python + +The MLIR Python bindings provide support for defining custom components in Python, +mainly including dialects, passes, and rewrite patterns. +The following sections outline how each of these can be implemented. + +### Dialects + +Dialects can be defined through the IRDL dialect bindings in Python. +The IRDL bindings offer a `load_dialects` function that +converts an MLIR module containing `irdl.dialect` ops into MLIR dialects. +For further details, see the documentation of [the IRDL dialect](../Dialects/IRDL.md). + +### Passes + +Passes can be defined as Python callables via the `PassManager.add` API. +In such case, the callable is wrapped as an `mlir::Pass` internally and +executed as part of the pass pipeline when `PassManager.run` is invoked. +In the callable, the `op` parameter represents the current operation being transformed, +while the `pass_` parameter provides access to the current `Pass` object, +allowing actions such as `signalPassFailure()`. +The lifetime of the callable is extended at least until the `PassManager` is destroyed. +The following example code demonstrates how to define Python passes. + +```python +def demo_pass(op, pass_): + # do something with the given op + pass + +pm = PassManager('any') +pm.add(demo_pass) +pm.add('some-cpp-defined-passes') +... +pm.run(some_op) +``` + +### Rewrite Patterns + +Rewrite patterns can be registered via the `add` method +of `mlir.rewrite.RewritePatternSet` in Python. +This method takes the operation type to be rewritten +and a Python callable that defines the *match and rewrite* logic. +Note that the Python callable should be defined so that +the rewrite is applied if and only if the match succeeds, +which corresponds to the return value being castable to `False`. + +The `RewritePatternSet` can be converted into +a `FrozenRewritePatternSet` using the `freeze` method, +which can be applied to an operation through +the greedy pattern driver using `apply_patterns_and_fold_greedily`. +The following example demonstrates the typical usage: + +```python +def to_muli(op, rewriter): + with rewriter.ip: + new_op = arith.muli(op.lhs, op.rhs, loc=op.location) + rewriter.replace_op(op, new_op) + +patterns = RewritePatternSet() +patterns.add(arith.AddIOp, to_muli) # Rewrite arith.addi into arith.muli +patterns.add(...) +frozen = patterns.freeze() + +module = ... +apply_patterns_and_fold_greedily(module, frozen) +``` + +The PDL dialect bindings also enable defining and generating rewrite patterns in Python. +The `mlir.rewrite.PDLModule` class accepts a module containing `pdl.pattern` ops, +which can be transformed into a `FrozenRewritePatternSet` using the `freeze` method. +This frozen set can then be applied to an operation +using the greedy rewrite pattern driver via `apply_patterns_and_fold_greedily`. +For further information, see [the PDL dialect documentation](/docs/Dialects/PDLOps/). + + ## Free-threading (No-GIL) support Free-threading or no-GIL support refers to CPython interpreter (>=3.13) with Global Interpreter Lock made optional. For details on the topic, please check [PEP-703](https://peps.python.org/pep-0703/) and this [Python free-threading guide](https://py-free-threading.github.io/). |