Skip to content

Commit 594e0ba

Browse files
[mlir][python] Add docs for op class extension mechanism.
Differential Revision: https://reviews.llvm.org/D99387
1 parent 4f3ea27 commit 594e0ba

File tree

1 file changed

+88
-1
lines changed

1 file changed

+88
-1
lines changed

mlir/docs/Bindings/Python.md

Lines changed: 88 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -449,7 +449,7 @@ defaults on `OpView`):
449449
variadics. Used by `OpView._ods_build_default` to decode operand and result
450450
lists that contain lists.
451451

452-
#### Builders
452+
#### Default Builder
453453

454454
Presently, only a single, default builder is mapped to the `__init__` method.
455455
The intent is that this `__init__` method represents the *most specific* of
@@ -475,3 +475,90 @@ construction via a (nested in the case of variadic) sequence of `results` and
475475
`operands`. This can be used to get some default construction semantics for
476476
operations that are otherwise unsupported in Python, at the expense of having
477477
a very generic signature.
478+
479+
#### Extending Generated Op Classes
480+
481+
Note that this is a rather complex mechanism and this section errs on the side
482+
of explicitness. Users are encouraged to find an example and duplicate it if
483+
they don't feel the need to understand the subtlety. The `builtin` dialect
484+
provides some relatively simple examples.
485+
486+
As mentioned above, the build system generates Python sources like
487+
`_{DIALECT_NAMESPACE}_ops_gen.py` for each dialect with Python bindings. It
488+
is often desirable to to use these generated classes as a starting point for
489+
further customization, so an extension mechanism is provided to make this
490+
easy (you are always free to do ad-hoc patching in your `{DIALECT_NAMESPACE}.py`
491+
file but we prefer a more standard mechanism that is applied uniformly).
492+
493+
To provide extensions, add a `_{DIALECT_NAMESPACE}_ops_ext.py` file to the
494+
`dialects` module (i.e. adjacent to your `{DIALECT_NAMESPACE}.py` top-level
495+
and the `*_ops_gen.py` file). Using the `builtin` dialect and `FuncOp` as an
496+
example, the generated code will include an import like this:
497+
498+
```python
499+
try:
500+
from . import _builtin_ops_ext as _ods_ext_module
501+
except ImportError:
502+
_ods_ext_module = None
503+
```
504+
505+
Then for each generated concrete `OpView` subclass, it will apply a decorator
506+
like:
507+
508+
```python
509+
@_ods_cext.register_operation(_Dialect)
510+
@_ods_extend_opview_class(_ods_ext_module)
511+
class FuncOp(_ods_ir.OpView):
512+
```
513+
514+
See the `_ods_common.py` `extend_opview_class` function for details of the
515+
mechanism. At a high level:
516+
517+
* If the extension module exists, locate an extension class for the op (in
518+
this example, `FuncOp`):
519+
* First by looking for an attribute with the exact name in the extension
520+
module.
521+
* Falling back to calling a `select_opview_mixin(parent_opview_cls)`
522+
function defined in the extension module.
523+
* If a mixin class is found, a new subclass is dynamically created that multiply
524+
inherits from `({_builtin_ops_ext.FuncOp}, _builtin_ops_gen.FuncOp)`.
525+
526+
The mixin class should not inherit from anything (i.e. directly extends
527+
`object` only). The facility is typically used to define custom `__init__`
528+
methods, properties, instance methods and static methods. Due to the
529+
inheritance ordering, the mixin class can act as though it extends the
530+
generated `OpView` subclass in most contexts (i.e.
531+
`issubclass(_builtin_ops_ext.FuncOp, OpView)` will return `False` but usage
532+
generally allows you treat it as duck typed as an `OpView`).
533+
534+
There are a couple of recommendations, given how the class hierarchy is
535+
defined:
536+
537+
* For static methods that need to instantiate the actual "leaf" op (which
538+
is dynamically generated and would result in circular dependencies to try
539+
to reference by name), prefer to use `@classmethod` and the concrete
540+
subclass will be provided as your first `cls` argument. See
541+
`_builtin_ops_ext.FuncOp.from_py_func` as an example.
542+
* If seeking to replace the generated `__init__` method entirely, you may
543+
actually want to invoke the super-super-class `mlir.ir.OpView` constructor
544+
directly, as it takes an `mlir.ir.Operation`, which is likely what you
545+
are constructing (i.e. the generated `__init__` method likely adds more
546+
API constraints than you want to expose in a custom builder).
547+
548+
A pattern that comes up frequently is wanting to provide a sugared `__init__`
549+
method which has optional or type-polymorphism/implicit conversions but to
550+
otherwise want to invoke the default op building logic. For such cases,
551+
it is recommended to use an idiom such as:
552+
553+
```python
554+
def __init__(self, sugar, spice, *, loc=None, ip=None):
555+
... massage into result_type, operands, attributes ...
556+
OpView.__init__(self, self.build_generic(
557+
results=[result_type],
558+
operands=operands,
559+
attributes=attributes,
560+
loc=loc,
561+
ip=ip))
562+
```
563+
564+
Refer to the documentation for `build_generic` for more information.

0 commit comments

Comments
 (0)