1c20c1960SStella Laurenzo# MLIR Python Bindings 2c20c1960SStella Laurenzo 3bacb0cacSAlex Zinenko**Current status**: Under development and not enabled by default 4bacb0cacSAlex Zinenko 5bacb0cacSAlex Zinenko[TOC] 6c20c1960SStella Laurenzo 7c20c1960SStella Laurenzo## Building 8c20c1960SStella Laurenzo 9c20c1960SStella Laurenzo### Pre-requisites 10c20c1960SStella Laurenzo 11c20c1960SStella Laurenzo* A relatively recent Python3 installation 124ca39dadSStella Laurenzo* Installation of python dependencies as specified in 139f3f6d7bSStella Laurenzo `mlir/python/requirements.txt` 14c20c1960SStella Laurenzo 15c20c1960SStella Laurenzo### CMake variables 16c20c1960SStella Laurenzo 179c21ddb7SUday Bondhugula* **`MLIR_ENABLE_BINDINGS_PYTHON`**`:BOOL` 18c20c1960SStella Laurenzo 19c20c1960SStella Laurenzo Enables building the Python bindings. Defaults to `OFF`. 20c20c1960SStella Laurenzo 21417f6137SStella Laurenzo* **`Python3_EXECUTABLE`**:`STRING` 22417f6137SStella Laurenzo 23417f6137SStella Laurenzo Specifies the `python` executable used for the LLVM build, including for 24417f6137SStella Laurenzo determining header/link flags for the Python bindings. On systems with 25417f6137SStella Laurenzo multiple Python implementations, setting this explicitly to the preferred 26417f6137SStella Laurenzo `python3` executable is strongly recommended. 27417f6137SStella Laurenzo 28417f6137SStella Laurenzo### Recommended development practices 29c20c1960SStella Laurenzo 30417f6137SStella LaurenzoIt is recommended to use a python virtual environment. Many ways exist for this, 31417f6137SStella Laurenzobut the following is the simplest: 32417f6137SStella Laurenzo 33417f6137SStella Laurenzo```shell 34417f6137SStella Laurenzo# Make sure your 'python' is what you expect. Note that on multi-python 35417f6137SStella Laurenzo# systems, this may have a version suffix, and on many Linuxes and MacOS where 36417f6137SStella Laurenzo# python2 and python3 co-exist, you may also want to use `python3`. 37417f6137SStella Laurenzowhich python 38417f6137SStella Laurenzopython -m venv ~/.venv/mlirdev 39417f6137SStella Laurenzosource ~/.venv/mlirdev/bin/activate 40417f6137SStella Laurenzo 414ca39dadSStella Laurenzo# Note that many LTS distros will bundle a version of pip itself that is too 424ca39dadSStella Laurenzo# old to download all of the latest binaries for certain platforms. 434ca39dadSStella Laurenzo# The pip version can be obtained with `python -m pip --version`, and for 444ca39dadSStella Laurenzo# Linux specifically, this should be cross checked with minimum versions 454ca39dadSStella Laurenzo# here: https://github.com/pypa/manylinux 464ca39dadSStella Laurenzo# It is recommended to upgrade pip: 474ca39dadSStella Laurenzopython -m pip install --upgrade pip 484ca39dadSStella Laurenzo 494ca39dadSStella Laurenzo 50417f6137SStella Laurenzo# Now the `python` command will resolve to your virtual environment and 51417f6137SStella Laurenzo# packages will be installed there. 529f3f6d7bSStella Laurenzopython -m pip install -r mlir/python/requirements.txt 53417f6137SStella Laurenzo 54417f6137SStella Laurenzo# Now run `cmake`, `ninja`, et al. 55417f6137SStella Laurenzo``` 56417f6137SStella Laurenzo 57310c9496SStella LaurenzoFor interactive use, it is sufficient to add the 58310c9496SStella Laurenzo`tools/mlir/python_packages/mlir_core/` directory in your `build/` directory to 59310c9496SStella Laurenzothe `PYTHONPATH`. Typically: 60417f6137SStella Laurenzo 61417f6137SStella Laurenzo```shell 62310c9496SStella Laurenzoexport PYTHONPATH=$(cd build && pwd)/tools/mlir/python_packages/mlir_core 63417f6137SStella Laurenzo``` 64c20c1960SStella Laurenzo 65a54f4eaeSMogballNote that if you have installed (i.e. via `ninja install`, et al), then python 66a54f4eaeSMogballpackages for all enabled projects will be in your install tree under 67310c9496SStella Laurenzo`python_packages/` (i.e. `python_packages/mlir_core`). Official distributions 68310c9496SStella Laurenzoare built with a more specialized setup. 69310c9496SStella Laurenzo 70c20c1960SStella Laurenzo## Design 71c20c1960SStella Laurenzo 72c20c1960SStella Laurenzo### Use cases 73c20c1960SStella Laurenzo 74c20c1960SStella LaurenzoThere are likely two primary use cases for the MLIR python bindings: 75c20c1960SStella Laurenzo 76c20c1960SStella Laurenzo1. Support users who expect that an installed version of LLVM/MLIR will yield 77c20c1960SStella Laurenzo the ability to `import mlir` and use the API in a pure way out of the box. 78c20c1960SStella Laurenzo 79a54f4eaeSMogball1. Downstream integrations will likely want to include parts of the API in 80a54f4eaeSMogball their private namespace or specially built libraries, probably mixing it 81a54f4eaeSMogball with other python native bits. 82c20c1960SStella Laurenzo 83c20c1960SStella Laurenzo### Composable modules 84c20c1960SStella Laurenzo 85895ae487SStella LaurenzoIn order to support use case \#2, the Python bindings are organized into 86c20c1960SStella Laurenzocomposable modules that downstream integrators can include and re-export into 87c20c1960SStella Laurenzotheir own namespace if desired. This forces several design points: 88c20c1960SStella Laurenzo 89a54f4eaeSMogball* Separate the construction/populating of a `py::module` from 90a54f4eaeSMogball `PYBIND11_MODULE` global constructor. 91c20c1960SStella Laurenzo 92c20c1960SStella Laurenzo* Introduce headers for C++-only wrapper classes as other related C++ modules 93c20c1960SStella Laurenzo will need to interop with it. 94c20c1960SStella Laurenzo 95c20c1960SStella Laurenzo* Separate any initialization routines that depend on optional components into 96c20c1960SStella Laurenzo its own module/dependency (currently, things like `registerAllDialects` fall 97c20c1960SStella Laurenzo into this category). 98c20c1960SStella Laurenzo 99c20c1960SStella LaurenzoThere are a lot of co-related issues of shared library linkage, distribution 100c20c1960SStella Laurenzoconcerns, etc that affect such things. Organizing the code into composable 101c20c1960SStella Laurenzomodules (versus a monolithic `cpp` file) allows the flexibility to address many 102c20c1960SStella Laurenzoof these as needed over time. Also, compilation time for all of the template 103c20c1960SStella Laurenzometa-programming in pybind scales with the number of things you define in a 104c20c1960SStella Laurenzotranslation unit. Breaking into multiple translation units can significantly aid 105c20c1960SStella Laurenzocompile times for APIs with a large surface area. 106c20c1960SStella Laurenzo 107c20c1960SStella Laurenzo### Submodules 108c20c1960SStella Laurenzo 109c20c1960SStella LaurenzoGenerally, the C++ codebase namespaces most things into the `mlir` namespace. 110c20c1960SStella LaurenzoHowever, in order to modularize and make the Python bindings easier to 111c20c1960SStella Laurenzounderstand, sub-packages are defined that map roughly to the directory structure 112c20c1960SStella Laurenzoof functional units in MLIR. 113c20c1960SStella Laurenzo 114c20c1960SStella LaurenzoExamples: 115c20c1960SStella Laurenzo 116c20c1960SStella Laurenzo* `mlir.ir` 117c20c1960SStella Laurenzo* `mlir.passes` (`pass` is a reserved word :( ) 118c20c1960SStella Laurenzo* `mlir.dialect` 119c20c1960SStella Laurenzo* `mlir.execution_engine` (aside from namespacing, it is important that 120c20c1960SStella Laurenzo "bulky"/optional parts like this are isolated) 121c20c1960SStella Laurenzo 122a54f4eaeSMogballIn addition, initialization functions that imply optional dependencies should be 123a54f4eaeSMogballin underscored (notionally private) modules such as `_init` and linked 124c20c1960SStella Laurenzoseparately. This allows downstream integrators to completely customize what is 125a54f4eaeSMogballincluded "in the box" and covers things like dialect registration, pass 126a54f4eaeSMogballregistration, etc. 127c20c1960SStella Laurenzo 128c20c1960SStella Laurenzo### Loader 129c20c1960SStella Laurenzo 130c20c1960SStella LaurenzoLLVM/MLIR is a non-trivial python-native project that is likely to co-exist with 131c20c1960SStella Laurenzoother non-trivial native extensions. As such, the native extension (i.e. the 132c20c1960SStella Laurenzo`.so`/`.pyd`/`.dylib`) is exported as a notionally private top-level symbol 133e31c77b1SStella Laurenzo(`_mlir`), while a small set of Python code is provided in 134a54f4eaeSMogball`mlir/_cext_loader.py` and siblings which loads and re-exports it. This split 135a54f4eaeSMogballprovides a place to stage code that needs to prepare the environment *before* 136a54f4eaeSMogballthe shared library is loaded into the Python runtime, and also provides a place 137a54f4eaeSMogballthat one-time initialization code can be invoked apart from module constructors. 138c20c1960SStella Laurenzo 139e31c77b1SStella LaurenzoIt is recommended to avoid using `__init__.py` files to the extent possible, 140a54f4eaeSMogballuntil reaching a leaf package that represents a discrete component. The rule to 141a54f4eaeSMogballkeep in mind is that the presence of an `__init__.py` file prevents the ability 142a54f4eaeSMogballto split anything at that level or below in the namespace into different 143a54f4eaeSMogballdirectories, deployment packages, wheels, etc. 144c20c1960SStella Laurenzo 145e31c77b1SStella LaurenzoSee the documentation for more information and advice: 146e31c77b1SStella Laurenzohttps://packaging.python.org/guides/packaging-namespace-packages/ 147c20c1960SStella Laurenzo 148c20c1960SStella Laurenzo### Use the C-API 149c20c1960SStella Laurenzo 150c20c1960SStella LaurenzoThe Python APIs should seek to layer on top of the C-API to the degree possible. 151c20c1960SStella LaurenzoEspecially for the core, dialect-independent parts, such a binding enables 152c20c1960SStella Laurenzopackaging decisions that would be difficult or impossible if spanning a C++ ABI 153c20c1960SStella Laurenzoboundary. In addition, factoring in this way side-steps some very difficult 154c20c1960SStella Laurenzoissues that arise when combining RTTI-based modules (which pybind derived things 155c20c1960SStella Laurenzoare) with non-RTTI polymorphic C++ code (the default compilation mode of LLVM). 156c20c1960SStella Laurenzo 1577abb0ff7SStella Laurenzo### Ownership in the Core IR 1587abb0ff7SStella Laurenzo 159a54f4eaeSMogballThere are several top-level types in the core IR that are strongly owned by 160a54f4eaeSMogballtheir python-side reference: 1617abb0ff7SStella Laurenzo 1627abb0ff7SStella Laurenzo* `PyContext` (`mlir.ir.Context`) 1637abb0ff7SStella Laurenzo* `PyModule` (`mlir.ir.Module`) 1647abb0ff7SStella Laurenzo* `PyOperation` (`mlir.ir.Operation`) - but with caveats 1657abb0ff7SStella Laurenzo 166895ae487SStella LaurenzoAll other objects are dependent. All objects maintain a back-reference 167895ae487SStella Laurenzo(keep-alive) to their closest containing top-level object. Further, dependent 168895ae487SStella Laurenzoobjects fall into two categories: a) uniqued (which live for the life-time of 169895ae487SStella Laurenzothe context) and b) mutable. Mutable objects need additional machinery for 170895ae487SStella Laurenzokeeping track of when the C++ instance that backs their Python object is no 171895ae487SStella Laurenzolonger valid (typically due to some specific mutation of the IR, deletion, or 172895ae487SStella Laurenzobulk operation). 1737abb0ff7SStella Laurenzo 174af66cd17SStella Laurenzo### Optionality and argument ordering in the Core IR 175af66cd17SStella Laurenzo 176a54f4eaeSMogballThe following types support being bound to the current thread as a context 177a54f4eaeSMogballmanager: 178af66cd17SStella Laurenzo 179af66cd17SStella Laurenzo* `PyLocation` (`loc: mlir.ir.Location = None`) 180af66cd17SStella Laurenzo* `PyInsertionPoint` (`ip: mlir.ir.InsertionPoint = None`) 181af66cd17SStella Laurenzo* `PyMlirContext` (`context: mlir.ir.Context = None`) 182af66cd17SStella Laurenzo 183895ae487SStella LaurenzoIn order to support composability of function arguments, when these types appear 184895ae487SStella Laurenzoas arguments, they should always be the last and appear in the above order and 185895ae487SStella Laurenzowith the given names (which is generally the order in which they are expected to 186895ae487SStella Laurenzoneed to be expressed explicitly in special cases) as necessary. Each should 187895ae487SStella Laurenzocarry a default value of `py::none()` and use either a manual or automatic 188895ae487SStella Laurenzoconversion for resolving either with the explicit value or a value from the 189895ae487SStella Laurenzothread context manager (i.e. `DefaultingPyMlirContext` or 190895ae487SStella Laurenzo`DefaultingPyLocation`). 191af66cd17SStella Laurenzo 192895ae487SStella LaurenzoThe rationale for this is that in Python, trailing keyword arguments to the 193895ae487SStella Laurenzo*right* are the most composable, enabling a variety of strategies such as kwarg 194895ae487SStella Laurenzopassthrough, default values, etc. Keeping function signatures composable 195895ae487SStella Laurenzoincreases the chances that interesting DSLs and higher level APIs can be 196895ae487SStella Laurenzoconstructed without a lot of exotic boilerplate. 197af66cd17SStella Laurenzo 198895ae487SStella LaurenzoUsed consistently, this enables a style of IR construction that rarely needs to 199895ae487SStella Laurenzouse explicit contexts, locations, or insertion points but is free to do so when 200895ae487SStella Laurenzoextra control is needed. 201af66cd17SStella Laurenzo 2027abb0ff7SStella Laurenzo#### Operation hierarchy 2037abb0ff7SStella Laurenzo 204895ae487SStella LaurenzoAs mentioned above, `PyOperation` is special because it can exist in either a 205895ae487SStella Laurenzotop-level or dependent state. The life-cycle is unidirectional: operations can 206895ae487SStella Laurenzobe created detached (top-level) and once added to another operation, they are 207895ae487SStella Laurenzothen dependent for the remainder of their lifetime. The situation is more 208895ae487SStella Laurenzocomplicated when considering construction scenarios where an operation is added 209895ae487SStella Laurenzoto a transitive parent that is still detached, necessitating further accounting 210895ae487SStella Laurenzoat such transition points (i.e. all such added children are initially added to 211895ae487SStella Laurenzothe IR with a parent of their outer-most detached operation, but then once it is 212895ae487SStella Laurenzoadded to an attached operation, they need to be re-parented to the containing 213895ae487SStella Laurenzomodule). 2147abb0ff7SStella Laurenzo 215895ae487SStella LaurenzoDue to the validity and parenting accounting needs, `PyOperation` is the owner 216895ae487SStella Laurenzofor regions and blocks and needs to be a top-level type that we can count on not 217895ae487SStella Laurenzoaliasing. This let's us do things like selectively invalidating instances when 218895ae487SStella Laurenzomutations occur without worrying that there is some alias to the same operation 219895ae487SStella Laurenzoin the hierarchy. Operations are also the only entity that are allowed to be in 220895ae487SStella Laurenzoa detached state, and they are interned at the context level so that there is 221895ae487SStella Laurenzonever more than one Python `mlir.ir.Operation` object for a unique 222895ae487SStella Laurenzo`MlirOperation`, regardless of how it is obtained. 2237abb0ff7SStella Laurenzo 224895ae487SStella LaurenzoThe C/C++ API allows for Region/Block to also be detached, but it simplifies the 225895ae487SStella Laurenzoownership model a lot to eliminate that possibility in this API, allowing the 226895ae487SStella LaurenzoRegion/Block to be completely dependent on its owning operation for accounting. 227895ae487SStella LaurenzoThe aliasing of Python `Region`/`Block` instances to underlying 228895ae487SStella Laurenzo`MlirRegion`/`MlirBlock` is considered benign and these objects are not interned 229895ae487SStella Laurenzoin the context (unlike operations). 2307abb0ff7SStella Laurenzo 231895ae487SStella LaurenzoIf we ever want to re-introduce detached regions/blocks, we could do so with new 232895ae487SStella Laurenzo"DetachedRegion" class or similar and also avoid the complexity of accounting. 233895ae487SStella LaurenzoWith the way it is now, we can avoid having a global live list for regions and 234895ae487SStella Laurenzoblocks. We may end up needing an op-local one at some point TBD, depending on 235895ae487SStella Laurenzohow hard it is to guarantee how mutations interact with their Python peer 236895ae487SStella Laurenzoobjects. We can cross that bridge easily when we get there. 2377abb0ff7SStella Laurenzo 238895ae487SStella LaurenzoModule, when used purely from the Python API, can't alias anyway, so we can use 239895ae487SStella Laurenzoit as a top-level ref type without a live-list for interning. If the API ever 240895ae487SStella Laurenzochanges such that this cannot be guaranteed (i.e. by letting you marshal a 241895ae487SStella Laurenzonative-defined Module in), then there would need to be a live table for it too. 2427abb0ff7SStella Laurenzo 243bacb0cacSAlex Zinenko## User-level API 244bacb0cacSAlex Zinenko 245bacb0cacSAlex Zinenko### Context Management 246bacb0cacSAlex Zinenko 247bacb0cacSAlex ZinenkoThe bindings rely on Python 248bacb0cacSAlex Zinenko[context managers](https://docs.python.org/3/reference/datamodel.html#context-managers) 249bacb0cacSAlex Zinenko(`with` statements) to simplify creation and handling of IR objects by omitting 250bacb0cacSAlex Zinenkorepeated arguments such as MLIR contexts, operation insertion points and 251bacb0cacSAlex Zinenkolocations. A context manager sets up the default object to be used by all 252bacb0cacSAlex Zinenkobinding calls within the following context and in the same thread. This default 253bacb0cacSAlex Zinenkocan be overridden by specific calls through the dedicated keyword arguments. 254bacb0cacSAlex Zinenko 255bacb0cacSAlex Zinenko#### MLIR Context 256bacb0cacSAlex Zinenko 257bacb0cacSAlex ZinenkoAn MLIR context is a top-level entity that owns attributes and types and is 258bacb0cacSAlex Zinenkoreferenced from virtually all IR constructs. Contexts also provide thread safety 259bacb0cacSAlex Zinenkoat the C++ level. In Python bindings, the MLIR context is also a Python context 260bacb0cacSAlex Zinenkomanager, one can write: 261bacb0cacSAlex Zinenko 262bacb0cacSAlex Zinenko```python 263bacb0cacSAlex Zinenkofrom mlir.ir import Context, Module 264bacb0cacSAlex Zinenko 265bacb0cacSAlex Zinenkowith Context() as ctx: 266bacb0cacSAlex Zinenko # IR construction using `ctx` as context. 267bacb0cacSAlex Zinenko 268bacb0cacSAlex Zinenko # For example, parsing an MLIR module from string requires the context. 269bacb0cacSAlex Zinenko Module.parse("builtin.module {}") 270bacb0cacSAlex Zinenko``` 271bacb0cacSAlex Zinenko 272bacb0cacSAlex ZinenkoIR objects referencing a context usually provide access to it through the 273bacb0cacSAlex Zinenko`.context` property. Most IR-constructing functions expect the context to be 274bacb0cacSAlex Zinenkoprovided in some form. In case of attributes and types, the context may be 275bacb0cacSAlex Zinenkoextracted from the contained attribute or type. In case of operations, the 276bacb0cacSAlex Zinenkocontext is systematically extracted from Locations (see below). When the context 277bacb0cacSAlex Zinenkocannot be extracted from any argument, the bindings API expects the (keyword) 278bacb0cacSAlex Zinenkoargument `context`. If it is not provided or set to `None` (default), it will be 279bacb0cacSAlex Zinenkolooked up from an implicit stack of contexts maintained by the bindings in the 280bacb0cacSAlex Zinenkocurrent thread and updated by context managers. If there is no surrounding 281bacb0cacSAlex Zinenkocontext, an error will be raised. 282bacb0cacSAlex Zinenko 283bacb0cacSAlex ZinenkoNote that it is possible to manually specify the MLIR context both inside and 284bacb0cacSAlex Zinenkooutside of the `with` statement: 285bacb0cacSAlex Zinenko 286bacb0cacSAlex Zinenko```python 287bacb0cacSAlex Zinenkofrom mlir.ir import Context, Module 288bacb0cacSAlex Zinenko 289bacb0cacSAlex Zinenkostandalone_ctx = Context() 290bacb0cacSAlex Zinenkowith Context() as managed_ctx: 291bacb0cacSAlex Zinenko # Parse a module in managed_ctx. 292bacb0cacSAlex Zinenko Module.parse("...") 293bacb0cacSAlex Zinenko 294bacb0cacSAlex Zinenko # Parse a module in standalone_ctx (override the context manager). 295bacb0cacSAlex Zinenko Module.parse("...", context=standalone_ctx) 296bacb0cacSAlex Zinenko 297bacb0cacSAlex Zinenko# Parse a module without using context managers. 298bacb0cacSAlex ZinenkoModule.parse("...", context=standalone_ctx) 299bacb0cacSAlex Zinenko``` 300bacb0cacSAlex Zinenko 301bacb0cacSAlex ZinenkoThe context object remains live as long as there are IR objects referencing it. 302bacb0cacSAlex Zinenko 303bacb0cacSAlex Zinenko#### Insertion Points and Locations 304bacb0cacSAlex Zinenko 305bacb0cacSAlex ZinenkoWhen constructing an MLIR operation, two pieces of information are required: 306bacb0cacSAlex Zinenko 307bacb0cacSAlex Zinenko- an *insertion point* that indicates where the operation is to be created in 308bacb0cacSAlex Zinenko the IR region/block/operation structure (usually before or after another 309bacb0cacSAlex Zinenko operation, or at the end of some block); it may be missing, at which point 310bacb0cacSAlex Zinenko the operation is created in the *detached* state; 311bacb0cacSAlex Zinenko- a *location* that contains user-understandable information about the source 312bacb0cacSAlex Zinenko of the operation (for example, file/line/column information), which must 313bacb0cacSAlex Zinenko always be provided as it carries a reference to the MLIR context. 314bacb0cacSAlex Zinenko 315bacb0cacSAlex ZinenkoBoth can be provided using context managers or explicitly as keyword arguments 316bacb0cacSAlex Zinenkoin the operation constructor. They can be also provided as keyword arguments 317bacb0cacSAlex Zinenko`ip` and `loc` both within and outside of the context manager. 318bacb0cacSAlex Zinenko 319bacb0cacSAlex Zinenko```python 320bacb0cacSAlex Zinenkofrom mlir.ir import Context, InsertionPoint, Location, Module, Operation 321bacb0cacSAlex Zinenko 322bacb0cacSAlex Zinenkowith Context() as ctx: 323bacb0cacSAlex Zinenko module = Module.create() 324bacb0cacSAlex Zinenko 325bacb0cacSAlex Zinenko # Prepare for inserting operations into the body of the module and indicate 326bacb0cacSAlex Zinenko # that these operations originate in the "f.mlir" file at the given line and 327bacb0cacSAlex Zinenko # column. 328bacb0cacSAlex Zinenko with InsertionPoint(module.body), Location.file("f.mlir", line=42, col=1): 329bacb0cacSAlex Zinenko # This operation will be inserted at the end of the module body and will 330bacb0cacSAlex Zinenko # have the location set up by the context manager. 331bacb0cacSAlex Zinenko Operation(<...>) 332bacb0cacSAlex Zinenko 333bacb0cacSAlex Zinenko # This operation will be inserted at the end of the module (and after the 334bacb0cacSAlex Zinenko # previously constructed operation) and will have the location provided as 335bacb0cacSAlex Zinenko # the keyword argument. 336bacb0cacSAlex Zinenko Operation(<...>, loc=Location.file("g.mlir", line=1, col=10)) 337bacb0cacSAlex Zinenko 338bacb0cacSAlex Zinenko # This operation will be inserted at the *beginning* of the block rather 339bacb0cacSAlex Zinenko # than at its end. 340bacb0cacSAlex Zinenko Operation(<...>, ip=InsertionPoint.at_block_begin(module.body)) 341bacb0cacSAlex Zinenko``` 342bacb0cacSAlex Zinenko 343bacb0cacSAlex ZinenkoNote that `Location` needs an MLIR context to be constructed. It can take the 344bacb0cacSAlex Zinenkocontext set up in the current thread by some surrounding context manager, or 345bacb0cacSAlex Zinenkoaccept it as an explicit argument: 346bacb0cacSAlex Zinenko 347bacb0cacSAlex Zinenko```python 348bacb0cacSAlex Zinenkofrom mlir.ir import Context, Location 349bacb0cacSAlex Zinenko 350bacb0cacSAlex Zinenko# Create a context and a location in this context in the same `with` statement. 351bacb0cacSAlex Zinenkowith Context() as ctx, Location.file("f.mlir", line=42, col=1, context=ctx): 352bacb0cacSAlex Zinenko pass 353bacb0cacSAlex Zinenko``` 354bacb0cacSAlex Zinenko 35590a6c3c2SAlex ZinenkoLocations are owned by the context and live as long as they are (transitively) 35690a6c3c2SAlex Zinenkoreferenced from somewhere in Python code. 357bacb0cacSAlex Zinenko 358bacb0cacSAlex ZinenkoUnlike locations, the insertion point may be left unspecified (or, equivalently, 359bacb0cacSAlex Zinenkoset to `None` or `False`) during operation construction. In this case, the 360bacb0cacSAlex Zinenkooperation is created in the *detached* state, that is, it is not added into the 361bacb0cacSAlex Zinenkoregion of another operation and is owned by the caller. This is usually the case 362bacb0cacSAlex Zinenkofor top-level operations that contain the IR, such as modules. Regions, blocks 363bacb0cacSAlex Zinenkoand values contained in an operation point back to it and maintain it live. 364bacb0cacSAlex Zinenko 365bacb0cacSAlex Zinenko### Inspecting IR Objects 366bacb0cacSAlex Zinenko 367bacb0cacSAlex ZinenkoInspecting the IR is one of the primary tasks the Python bindings are designed 368bacb0cacSAlex Zinenkofor. One can traverse the IR operation/region/block structure and inspect their 369bacb0cacSAlex Zinenkoaspects such as operation attributes and value types. 370bacb0cacSAlex Zinenko 371bacb0cacSAlex Zinenko#### Operations, Regions and Blocks 372bacb0cacSAlex Zinenko 373bacb0cacSAlex ZinenkoOperations are represented as either: 374bacb0cacSAlex Zinenko 375bacb0cacSAlex Zinenko- the generic `Operation` class, useful in particular for generic processing 376bacb0cacSAlex Zinenko of unregistered operations; or 377bacb0cacSAlex Zinenko- a specific subclass of `OpView` that provides more semantically-loaded 378bacb0cacSAlex Zinenko accessors to operation properties. 379bacb0cacSAlex Zinenko 380bacb0cacSAlex ZinenkoGiven an `OpView` subclass, one can obtain an `Operation` using its `.operation` 381bacb0cacSAlex Zinenkoproperty. Given an `Operation`, one can obtain the corresponding `OpView` using 382bacb0cacSAlex Zinenkoits `.opview` property *as long as* the corresponding class has been set up. 383bacb0cacSAlex ZinenkoThis typically means that the Python module of its dialect has been loaded. By 384bacb0cacSAlex Zinenkodefault, the `OpView` version is produced when navigating the IR tree. 385bacb0cacSAlex Zinenko 386bacb0cacSAlex ZinenkoOne can check if an operation has a specific type by means of Python's 387bacb0cacSAlex Zinenko`isinstance` function: 388bacb0cacSAlex Zinenko 389bacb0cacSAlex Zinenko```python 390bacb0cacSAlex Zinenkooperation = <...> 391bacb0cacSAlex Zinenkoopview = <...> 392bacb0cacSAlex Zinenkoif isinstance(operation.opview, mydialect.MyOp): 393bacb0cacSAlex Zinenko pass 394bacb0cacSAlex Zinenkoif isinstance(opview, mydialect.MyOp): 395bacb0cacSAlex Zinenko pass 396bacb0cacSAlex Zinenko``` 397bacb0cacSAlex Zinenko 398bacb0cacSAlex ZinenkoThe components of an operation can be inspected using its properties. 399bacb0cacSAlex Zinenko 400bacb0cacSAlex Zinenko- `attributes` is a collection of operation attributes . It can be subscripted 401bacb0cacSAlex Zinenko as both dictionary and sequence, e.g., both `operation.attributes["value"]` 402bacb0cacSAlex Zinenko and `operation.attributes[0]` will work. There is no guarantee on the order 403bacb0cacSAlex Zinenko in which the attributes are traversed when iterating over the `attributes` 404bacb0cacSAlex Zinenko property as sequence. 405bacb0cacSAlex Zinenko- `operands` is a sequence collection of operation operands. 406bacb0cacSAlex Zinenko- `results` is a sequence collection of operation results. 407bacb0cacSAlex Zinenko- `regions` is a sequence collection of regions attached to the operation. 408bacb0cacSAlex Zinenko 409bacb0cacSAlex ZinenkoThe objects produced by `operands` and `results` have a `.types` property that 410bacb0cacSAlex Zinenkocontains a sequence collection of types of the corresponding values. 411bacb0cacSAlex Zinenko 412bacb0cacSAlex Zinenko```python 413bacb0cacSAlex Zinenkofrom mlir.ir import Operation 414bacb0cacSAlex Zinenko 415bacb0cacSAlex Zinenkooperation1 = <...> 416bacb0cacSAlex Zinenkooperation2 = <...> 417bacb0cacSAlex Zinenkoif operation1.results.types == operation2.operand.types: 418bacb0cacSAlex Zinenko pass 419bacb0cacSAlex Zinenko``` 420bacb0cacSAlex Zinenko 421bacb0cacSAlex Zinenko`OpView` subclasses for specific operations may provide leaner accessors to 422286a7a40SMarkus Böckproperties of an operation. For example, named attributes, operand and results 423bacb0cacSAlex Zinenkoare usually accessible as properties of the `OpView` subclass with the same 424bacb0cacSAlex Zinenkoname, such as `operation.const_value` instead of 425bacb0cacSAlex Zinenko`operation.attributes["const_value"]`. If this name is a reserved Python 426bacb0cacSAlex Zinenkokeyword, it is suffixed with an underscore. 427bacb0cacSAlex Zinenko 428bacb0cacSAlex ZinenkoThe operation itself is iterable, which provides access to the attached regions 429bacb0cacSAlex Zinenkoin order: 430bacb0cacSAlex Zinenko 431bacb0cacSAlex Zinenko```python 432bacb0cacSAlex Zinenkofrom mlir.ir import Operation 433bacb0cacSAlex Zinenko 434bacb0cacSAlex Zinenkooperation = <...> 435bacb0cacSAlex Zinenkofor region in operation: 436bacb0cacSAlex Zinenko do_something_with_region(region) 437bacb0cacSAlex Zinenko``` 438bacb0cacSAlex Zinenko 439bacb0cacSAlex ZinenkoA region is conceptually a sequence of blocks. Objects of the `Region` class are 440bacb0cacSAlex Zinenkothus iterable, which provides access to the blocks. One can also use the 441bacb0cacSAlex Zinenko`.blocks` property. 442bacb0cacSAlex Zinenko 443bacb0cacSAlex Zinenko```python 444286a7a40SMarkus Böck# Regions are directly iterable and give access to blocks. 445bacb0cacSAlex Zinenkofor block1, block2 in zip(operation.regions[0], operation.regions[0].blocks) 446bacb0cacSAlex Zinenko assert block1 == block2 447bacb0cacSAlex Zinenko``` 448bacb0cacSAlex Zinenko 449bacb0cacSAlex ZinenkoA block contains a sequence of operations, and has several additional 450bacb0cacSAlex Zinenkoproperties. Objects of the `Block` class are iterable and provide access to the 451bacb0cacSAlex Zinenkooperations contained in the block. So does the `.operations` property. Blocks 452bacb0cacSAlex Zinenkoalso have a list of arguments available as a sequence collection using the 453bacb0cacSAlex Zinenko`.arguments` property. 454bacb0cacSAlex Zinenko 455bacb0cacSAlex ZinenkoBlock and region belong to the parent operation in Python bindings and keep it 456bacb0cacSAlex Zinenkoalive. This operation can be accessed using the `.owner` property. 457bacb0cacSAlex Zinenko 458bacb0cacSAlex Zinenko#### Attributes and Types 459bacb0cacSAlex Zinenko 460bacb0cacSAlex ZinenkoAttributes and types are (mostly) immutable context-owned objects. They are 461bacb0cacSAlex Zinenkorepresented as either: 462bacb0cacSAlex Zinenko 463286a7a40SMarkus Böck- an opaque `Attribute` or `Type` object supporting printing and comparison; 464bacb0cacSAlex Zinenko or 465bacb0cacSAlex Zinenko- a concrete subclass thereof with access to properties of the attribute or 466bacb0cacSAlex Zinenko type. 467bacb0cacSAlex Zinenko 468bacb0cacSAlex ZinenkoGiven an `Attribute` or `Type` object, one can obtain a concrete subclass using 469bacb0cacSAlex Zinenkothe constructor of the subclass. This may raise a `ValueError` if the attribute 47090a6c3c2SAlex Zinenkoor type is not of the expected subclass: 471bacb0cacSAlex Zinenko 472bacb0cacSAlex Zinenko```python 473bacb0cacSAlex Zinenkofrom mlir.ir import Attribute, Type 474bacb0cacSAlex Zinenkofrom mlir.<dialect> import ConcreteAttr, ConcreteType 475bacb0cacSAlex Zinenko 476bacb0cacSAlex Zinenkoattribute = <...> 477bacb0cacSAlex Zinenkotype = <...> 478bacb0cacSAlex Zinenkotry: 479bacb0cacSAlex Zinenko concrete_attr = ConcreteAttr(attribute) 480bacb0cacSAlex Zinenko concrete_type = ConcreteType(type) 481bacb0cacSAlex Zinenkoexcept ValueError as e: 482bacb0cacSAlex Zinenko # Handle incorrect subclass. 483bacb0cacSAlex Zinenko``` 484bacb0cacSAlex Zinenko 485bacb0cacSAlex ZinenkoIn addition, concrete attribute and type classes provide a static `isinstance` 486bacb0cacSAlex Zinenkomethod to check whether an object of the opaque `Attribute` or `Type` type can 487bacb0cacSAlex Zinenkobe downcasted: 488bacb0cacSAlex Zinenko 489bacb0cacSAlex Zinenko```python 490bacb0cacSAlex Zinenkofrom mlir.ir import Attribute, Type 491bacb0cacSAlex Zinenkofrom mlir.<dialect> import ConcreteAttr, ConcreteType 492bacb0cacSAlex Zinenko 493bacb0cacSAlex Zinenkoattribute = <...> 494bacb0cacSAlex Zinenkotype = <...> 495bacb0cacSAlex Zinenko 496bacb0cacSAlex Zinenko# No need to handle errors here. 497bacb0cacSAlex Zinenkoif ConcreteAttr.isinstance(attribute): 498bacb0cacSAlex Zinenko concrete_attr = ConcreteAttr(attribute) 499bacb0cacSAlex Zinenkoif ConcreteType.isinstance(type): 500bacb0cacSAlex Zinenko concrete_type = ConcreteType(type) 501bacb0cacSAlex Zinenko``` 502bacb0cacSAlex Zinenko 503bacb0cacSAlex ZinenkoBy default, and unlike operations, attributes and types are returned from IR 504bacb0cacSAlex Zinenkotraversals using the opaque `Attribute` or `Type` that needs to be downcasted. 505bacb0cacSAlex Zinenko 506bacb0cacSAlex ZinenkoConcrete attribute and type classes usually expose their properties as Python 507bacb0cacSAlex Zinenkoreadonly properties. For example, the elemental type of a tensor type can be 508bacb0cacSAlex Zinenkoaccessed using the `.element_type` property. 509bacb0cacSAlex Zinenko 510bacb0cacSAlex Zinenko#### Values 511bacb0cacSAlex Zinenko 512bacb0cacSAlex ZinenkoMLIR has two kinds of values based on their defining object: block arguments and 513bacb0cacSAlex Zinenkooperation results. Values are handled similarly to attributes and types. They 514bacb0cacSAlex Zinenkoare represented as either: 515bacb0cacSAlex Zinenko 516bacb0cacSAlex Zinenko- a generic `Value` object; or 517bacb0cacSAlex Zinenko- a concrete `BlockArgument` or `OpResult` object. 518bacb0cacSAlex Zinenko 519bacb0cacSAlex ZinenkoThe former provides all the generic functionality such as comparison, type 520bacb0cacSAlex Zinenkoaccess and printing. The latter provide access to the defining block or 521bacb0cacSAlex Zinenkooperation and the position of the value within it. By default, the generic 522bacb0cacSAlex Zinenko`Value` objects are returned from IR traversals. Downcasting is implemented 523bacb0cacSAlex Zinenkothrough concrete subclass constructors, similarly to attribtues and types: 524bacb0cacSAlex Zinenko 525bacb0cacSAlex Zinenko```python 526bacb0cacSAlex Zinenkofrom mlir.ir import BlockArgument, OpResult, Value 527bacb0cacSAlex Zinenko 528bacb0cacSAlex Zinenkovalue = ... 529bacb0cacSAlex Zinenko 530bacb0cacSAlex Zinenko# Set `concrete` to the specific value subclass. 531bacb0cacSAlex Zinenkotry: 532bacb0cacSAlex Zinenko concrete = BlockArgument(value) 533bacb0cacSAlex Zinenkoexcept ValueError: 534bacb0cacSAlex Zinenko # This must not raise another ValueError as values are either block arguments 535bacb0cacSAlex Zinenko # or op results. 536bacb0cacSAlex Zinenko concrete = OpResult(value) 537bacb0cacSAlex Zinenko``` 538bacb0cacSAlex Zinenko 53914c92070SAlex Zinenko#### Interfaces 54014c92070SAlex Zinenko 54114c92070SAlex ZinenkoMLIR interfaces are a mechanism to interact with the IR without needing to know 54214c92070SAlex Zinenkospecific types of operations but only some of their aspects. Operation 54314c92070SAlex Zinenkointerfaces are available as Python classes with the same name as their C++ 54414c92070SAlex Zinenkocounterparts. Objects of these classes can be constructed from either: 54514c92070SAlex Zinenko 54614c92070SAlex Zinenko- an object of the `Operation` class or of any `OpView` subclass; in this 54714c92070SAlex Zinenko case, all interface methods are available; 54814c92070SAlex Zinenko- a subclass of `OpView` and a context; in this case, only the *static* 54914c92070SAlex Zinenko interface methods are available as there is no associated operation. 55014c92070SAlex Zinenko 55114c92070SAlex ZinenkoIn both cases, construction of the interface raises a `ValueError` if the 55214c92070SAlex Zinenkooperation class does not implement the interface in the given context (or, for 55314c92070SAlex Zinenkooperations, in the context that the operation is defined in). Similarly to 55414c92070SAlex Zinenkoattributes and types, the MLIR context may be set up by a surrounding context 55514c92070SAlex Zinenkomanager. 55614c92070SAlex Zinenko 55714c92070SAlex Zinenko```python 55814c92070SAlex Zinenkofrom mlir.ir import Context, InferTypeOpInterface 55914c92070SAlex Zinenko 56014c92070SAlex Zinenkowith Context(): 56114c92070SAlex Zinenko op = <...> 56214c92070SAlex Zinenko 56314c92070SAlex Zinenko # Attempt to cast the operation into an interface. 56414c92070SAlex Zinenko try: 56514c92070SAlex Zinenko iface = InferTypeOpInterface(op) 56614c92070SAlex Zinenko except ValueError: 56714c92070SAlex Zinenko print("Operation does not implement InferTypeOpInterface.") 56814c92070SAlex Zinenko raise 56914c92070SAlex Zinenko 57014c92070SAlex Zinenko # All methods are available on interface objects constructed from an Operation 57114c92070SAlex Zinenko # or an OpView. 57214c92070SAlex Zinenko iface.someInstanceMethod() 57314c92070SAlex Zinenko 57414c92070SAlex Zinenko # An interface object can also be constructed given an OpView subclass. It 57514c92070SAlex Zinenko # also needs a context in which the interface will be looked up. The context 57614c92070SAlex Zinenko # can be provided explicitly or set up by the surrounding context manager. 57714c92070SAlex Zinenko try: 57814c92070SAlex Zinenko iface = InferTypeOpInterface(some_dialect.SomeOp) 57914c92070SAlex Zinenko except ValueError: 58014c92070SAlex Zinenko print("SomeOp does not implement InferTypeOpInterface.") 58114c92070SAlex Zinenko raise 58214c92070SAlex Zinenko 58314c92070SAlex Zinenko # Calling an instance method on an interface object constructed from a class 58414c92070SAlex Zinenko # will raise TypeError. 58514c92070SAlex Zinenko try: 58614c92070SAlex Zinenko iface.someInstanceMethod() 58714c92070SAlex Zinenko except TypeError: 58814c92070SAlex Zinenko pass 58914c92070SAlex Zinenko 59014c92070SAlex Zinenko # One can still call static interface methods though. 59114c92070SAlex Zinenko iface.inferOpReturnTypes(<...>) 59214c92070SAlex Zinenko``` 59314c92070SAlex Zinenko 59414c92070SAlex ZinenkoIf an interface object was constructed from an `Operation` or an `OpView`, they 59514c92070SAlex Zinenkoare available as `.operation` and `.opview` properties of the interface object, 59614c92070SAlex Zinenkorespectively. 59714c92070SAlex Zinenko 59814c92070SAlex ZinenkoOnly a subset of operation interfaces are currently provided in Python bindings. 59914c92070SAlex ZinenkoAttribute and type interfaces are not yet available in Python bindings. 60014c92070SAlex Zinenko 601bacb0cacSAlex Zinenko### Creating IR Objects 602bacb0cacSAlex Zinenko 603bacb0cacSAlex ZinenkoPython bindings also support IR creation and manipulation. 604bacb0cacSAlex Zinenko 605bacb0cacSAlex Zinenko#### Operations, Regions and Blocks 606bacb0cacSAlex Zinenko 607bacb0cacSAlex ZinenkoOperations can be created given a `Location` and an optional `InsertionPoint`. 608bacb0cacSAlex ZinenkoIt is often easier to user context managers to specify locations and insertion 609286a7a40SMarkus Böckpoints for several operations created in a row as described above. 610bacb0cacSAlex Zinenko 611bacb0cacSAlex ZinenkoConcrete operations can be created by using constructors of the corresponding 612bacb0cacSAlex Zinenko`OpView` subclasses. The generic, default form of the constructor accepts: 613bacb0cacSAlex Zinenko 614bacb0cacSAlex Zinenko- an optional sequence of types for operation results (`results`); 615bacb0cacSAlex Zinenko- an optional sequence of values for operation operands, or another operation 616bacb0cacSAlex Zinenko producing those values (`operands`); 617bacb0cacSAlex Zinenko- an optional dictionary of operation attributes (`attributes`); 618bacb0cacSAlex Zinenko- an optional sequence of successor blocks (`successors`); 619bacb0cacSAlex Zinenko- the number of regions to attach to the operation (`regions`, default `0`); 620bacb0cacSAlex Zinenko- the `loc` keyword argument containing the `Location` of this operation; if 621bacb0cacSAlex Zinenko `None`, the location created by the closest context manager is used or an 622bacb0cacSAlex Zinenko exception will be raised if there is no context manager; 623bacb0cacSAlex Zinenko- the `ip` keyword argument indicating where the operation will be inserted in 624bacb0cacSAlex Zinenko the IR; if `None`, the insertion point created by the closest context 625bacb0cacSAlex Zinenko manager is used; if there is no surrounding context manager, the operation 626bacb0cacSAlex Zinenko is created in the detached state. 627bacb0cacSAlex Zinenko 628bacb0cacSAlex ZinenkoMost operations will customize the constructor to accept a reduced list of 629bacb0cacSAlex Zinenkoarguments that are relevant for the operation. For example, zero-result 630bacb0cacSAlex Zinenkooperations may omit the `results` argument, so can the operations where the 631bacb0cacSAlex Zinenkoresult types can be derived from operand types unambiguously. As a concrete 632bacb0cacSAlex Zinenkoexample, built-in function operations can be constructed by providing a function 633bacb0cacSAlex Zinenkoname as string and its argument and result types as a tuple of sequences: 634bacb0cacSAlex Zinenko 635bacb0cacSAlex Zinenko```python 636bacb0cacSAlex Zinenkofrom mlir.ir import Context, Module 637bacb0cacSAlex Zinenkofrom mlir.dialects import builtin 638bacb0cacSAlex Zinenko 639bacb0cacSAlex Zinenkowith Context(): 640bacb0cacSAlex Zinenko module = Module.create() 641bacb0cacSAlex Zinenko with InsertionPoint(module.body), Location.unknown(): 64236550692SRiver Riddle func = func.FuncOp("main", ([], [])) 643bacb0cacSAlex Zinenko``` 644bacb0cacSAlex Zinenko 645bacb0cacSAlex ZinenkoAlso see below for constructors generated from ODS. 646bacb0cacSAlex Zinenko 647bacb0cacSAlex ZinenkoOperations can also be constructed using the generic class and based on the 648bacb0cacSAlex Zinenkocanonical string name of the operation using `Operation.create`. It accepts the 649bacb0cacSAlex Zinenkooperation name as string, which must exactly match the canonical name of the 650bacb0cacSAlex Zinenkooperation in C++ or ODS, followed by the same argument list as the default 651bacb0cacSAlex Zinenkoconstructor for `OpView`. *This form is discouraged* from use and is intended 652bacb0cacSAlex Zinenkofor generic operation processing. 653bacb0cacSAlex Zinenko 654bacb0cacSAlex Zinenko```python 655bacb0cacSAlex Zinenkofrom mlir.ir import Context, Module 656bacb0cacSAlex Zinenkofrom mlir.dialects import builtin 657bacb0cacSAlex Zinenko 658bacb0cacSAlex Zinenkowith Context(): 659bacb0cacSAlex Zinenko module = Module.create() 660bacb0cacSAlex Zinenko with InsertionPoint(module.body), Location.unknown(): 661bacb0cacSAlex Zinenko # Operations can be created in a generic way. 662bacb0cacSAlex Zinenko func = Operation.create( 66336550692SRiver Riddle "func.func", results=[], operands=[], 6644a3460a7SRiver Riddle attributes={"function_type":TypeAttr.get(FunctionType.get([], []))}, 665bacb0cacSAlex Zinenko successors=None, regions=1) 666bacb0cacSAlex Zinenko # The result will be downcasted to the concrete `OpView` subclass if 667bacb0cacSAlex Zinenko # available. 66836550692SRiver Riddle assert isinstance(func, func.FuncOp) 669bacb0cacSAlex Zinenko``` 670bacb0cacSAlex Zinenko 671bacb0cacSAlex ZinenkoRegions are created for an operation when constructing it on the C++ side. They 672bacb0cacSAlex Zinenkoare not constructible in Python and are not expected to exist outside of 673bacb0cacSAlex Zinenkooperations (unlike in C++ that supports detached regions). 674bacb0cacSAlex Zinenko 675bacb0cacSAlex ZinenkoBlocks can be created within a given region and inserted before or after another 676bacb0cacSAlex Zinenkoblock of the same region using `create_before()`, `create_after()` methods of 67778f2dae0SAlex Zinenkothe `Block` class, or the `create_at_start()` static method of the same class. 67878f2dae0SAlex ZinenkoThey are not expected to exist outside of regions (unlike in C++ that supports 67978f2dae0SAlex Zinenkodetached blocks). 68078f2dae0SAlex Zinenko 68178f2dae0SAlex Zinenko```python 68278f2dae0SAlex Zinenkofrom mlir.ir import Block, Context, Operation 68378f2dae0SAlex Zinenko 68478f2dae0SAlex Zinenkowith Context(): 68578f2dae0SAlex Zinenko op = Operation.create("generic.op", regions=1) 68678f2dae0SAlex Zinenko 68778f2dae0SAlex Zinenko # Create the first block in the region. 68878f2dae0SAlex Zinenko entry_block = Block.create_at_start(op.regions[0]) 68978f2dae0SAlex Zinenko 69078f2dae0SAlex Zinenko # Create further blocks. 69178f2dae0SAlex Zinenko other_block = entry_block.create_after() 69278f2dae0SAlex Zinenko``` 693bacb0cacSAlex Zinenko 694bacb0cacSAlex ZinenkoBlocks can be used to create `InsertionPoint`s, which can point to the beginning 695bacb0cacSAlex Zinenkoor the end of the block, or just before its terminator. It is common for 696bacb0cacSAlex Zinenko`OpView` subclasses to provide a `.body` property that can be used to construct 697bacb0cacSAlex Zinenkoan `InsertionPoint`. For example, builtin `Module` and `FuncOp` provide a 698bacb0cacSAlex Zinenko`.body` and `.add_entry_blocK()`, respectively. 699bacb0cacSAlex Zinenko 700bacb0cacSAlex Zinenko#### Attributes and Types 701bacb0cacSAlex Zinenko 702bacb0cacSAlex ZinenkoAttributes and types can be created given a `Context` or another attribute or 703bacb0cacSAlex Zinenkotype object that already references the context. To indicate that they are owned 704bacb0cacSAlex Zinenkoby the context, they are obtained by calling the static `get` method on the 705bacb0cacSAlex Zinenkoconcrete attribute or type class. These method take as arguments the data 706bacb0cacSAlex Zinenkonecessary to construct the attribute or type and a the keyword `context` 707bacb0cacSAlex Zinenkoargument when the context cannot be derived from other arguments. 708bacb0cacSAlex Zinenko 709bacb0cacSAlex Zinenko```python 710bacb0cacSAlex Zinenkofrom mlir.ir import Context, F32Type, FloatAttr 711bacb0cacSAlex Zinenko 712bacb0cacSAlex Zinenko# Attribute and types require access to an MLIR context, either directly or 713bacb0cacSAlex Zinenko# through another context-owned object. 714bacb0cacSAlex Zinenkoctx = Context() 715bacb0cacSAlex Zinenkof32 = F32Type.get(context=ctx) 716bacb0cacSAlex Zinenkopi = FloatAttr.get(f32, 3.14) 717bacb0cacSAlex Zinenko 718bacb0cacSAlex Zinenko# They may use the context defined by the surrounding context manager. 719bacb0cacSAlex Zinenkowith Context(): 720bacb0cacSAlex Zinenko f32 = F32Type.get() 721bacb0cacSAlex Zinenko pi = FloatAttr.get(f32, 3.14) 722bacb0cacSAlex Zinenko``` 723bacb0cacSAlex Zinenko 724bacb0cacSAlex ZinenkoSome attributes provide additional construction methods for clarity. 725bacb0cacSAlex Zinenko 726bacb0cacSAlex Zinenko```python 727bacb0cacSAlex Zinenkofrom mlir.ir import Context, IntegerAttr, IntegerType 728bacb0cacSAlex Zinenko 729bacb0cacSAlex Zinenkowith Context(): 730bacb0cacSAlex Zinenko i8 = IntegerType.get_signless(8) 731bacb0cacSAlex Zinenko IntegerAttr.get(i8, 42) 732bacb0cacSAlex Zinenko``` 733bacb0cacSAlex Zinenko 734bacb0cacSAlex ZinenkoBuiltin attribute can often be constructed from Python types with similar 735bacb0cacSAlex Zinenkostructure. For example, `ArrayAttr` can be constructed from a sequence 736bacb0cacSAlex Zinenkocollection of attributes, and a `DictAttr` can be constructed from a dictionary: 737bacb0cacSAlex Zinenko 738bacb0cacSAlex Zinenko```python 739bacb0cacSAlex Zinenkofrom mlir.ir import ArrayAttr, Context, DictAttr, UnitAttr 740bacb0cacSAlex Zinenko 741bacb0cacSAlex Zinenkowith Context(): 742bacb0cacSAlex Zinenko array = ArrayAttr.get([UnitAttr.get(), UnitAttr.get()]) 743bacb0cacSAlex Zinenko dictionary = DictAttr.get({"array": array, "unit": UnitAttr.get()}) 744bacb0cacSAlex Zinenko``` 745bacb0cacSAlex Zinenko 746b57acb9aSJacques PienaarCustom builders for Attributes to be used during Operation creation can be 747b57acb9aSJacques Pienaarregistered by way of the `register_attribute_builder`. In particular the 748b57acb9aSJacques Pienaarfollowing is how a custom builder is registered for `I32Attr`: 749b57acb9aSJacques Pienaar 750b57acb9aSJacques Pienaar```python 751b57acb9aSJacques Pienaar@register_attribute_builder("I32Attr") 752b57acb9aSJacques Pienaardef _i32Attr(x: int, context: Context): 753b57acb9aSJacques Pienaar return IntegerAttr.get( 754b57acb9aSJacques Pienaar IntegerType.get_signless(32, context=context), x) 755b57acb9aSJacques Pienaar``` 756b57acb9aSJacques Pienaar 757b57acb9aSJacques PienaarThis allows to invoke op creation of an op with a `I32Attr` with 758b57acb9aSJacques Pienaar 759b57acb9aSJacques Pienaar```python 760b57acb9aSJacques Pienaarfoo.Op(30) 761b57acb9aSJacques Pienaar``` 762b57acb9aSJacques Pienaar 763b57acb9aSJacques PienaarThe registration is based on the ODS name but registry is via pure python 764b57acb9aSJacques Pienaarmethod. Only single custom builder is allowed to be registered per ODS attribute 765b57acb9aSJacques Pienaartype (e.g., I32Attr can have only one, which can correspond to multiple of the 766b57acb9aSJacques Pienaarunderlying IntegerAttr type). 767b57acb9aSJacques Pienaar 768b57acb9aSJacques Pienaarinstead of 769b57acb9aSJacques Pienaar 770b57acb9aSJacques Pienaar```python 771b57acb9aSJacques Pienaarfoo.Op(IntegerAttr.get(IndexType.get_signless(32, context=context), 30)) 772b57acb9aSJacques Pienaar``` 773b57acb9aSJacques Pienaar 774c20c1960SStella Laurenzo## Style 775c20c1960SStella Laurenzo 776c20c1960SStella LaurenzoIn general, for the core parts of MLIR, the Python bindings should be largely 777c20c1960SStella Laurenzoisomorphic with the underlying C++ structures. However, concessions are made 778c20c1960SStella Laurenzoeither for practicality or to give the resulting library an appropriately 779c20c1960SStella Laurenzo"Pythonic" flavor. 780c20c1960SStella Laurenzo 781895ae487SStella Laurenzo### Properties vs get\*() methods 782c20c1960SStella Laurenzo 783c20c1960SStella LaurenzoGenerally favor converting trivial methods like `getContext()`, `getName()`, 784c20c1960SStella Laurenzo`isEntryBlock()`, etc to read-only Python properties (i.e. `context`). It is 785c20c1960SStella Laurenzoprimarily a matter of calling `def_property_readonly` vs `def` in binding code, 786c20c1960SStella Laurenzoand makes things feel much nicer to the Python side. 787c20c1960SStella Laurenzo 788c20c1960SStella LaurenzoFor example, prefer: 789c20c1960SStella Laurenzo 790c20c1960SStella Laurenzo```c++ 791c20c1960SStella Laurenzom.def_property_readonly("context", ...) 792c20c1960SStella Laurenzo``` 793c20c1960SStella Laurenzo 794c20c1960SStella LaurenzoOver: 795c20c1960SStella Laurenzo 796c20c1960SStella Laurenzo```c++ 797c20c1960SStella Laurenzom.def("getContext", ...) 798c20c1960SStella Laurenzo``` 799c20c1960SStella Laurenzo 800a54f4eaeSMogball### **repr** methods 801c20c1960SStella Laurenzo 802c20c1960SStella LaurenzoThings that have nice printed representations are really great :) If there is a 803c20c1960SStella Laurenzoreasonable printed form, it can be a significant productivity boost to wire that 804c20c1960SStella Laurenzoto the `__repr__` method (and verify it with a [doctest](#sample-doctest)). 805c20c1960SStella Laurenzo 806895ae487SStella Laurenzo### CamelCase vs snake\_case 807c20c1960SStella Laurenzo 808c20c1960SStella LaurenzoName functions/methods/properties in `snake_case` and classes in `CamelCase`. As 809c20c1960SStella Laurenzoa mechanical concession to Python style, this can go a long way to making the 810c20c1960SStella LaurenzoAPI feel like it fits in with its peers in the Python landscape. 811c20c1960SStella Laurenzo 812c20c1960SStella LaurenzoIf in doubt, choose names that will flow properly with other 813c20c1960SStella Laurenzo[PEP 8 style names](https://pep8.org/#descriptive-naming-styles). 814c20c1960SStella Laurenzo 815c20c1960SStella Laurenzo### Prefer pseudo-containers 816c20c1960SStella Laurenzo 817c20c1960SStella LaurenzoMany core IR constructs provide methods directly on the instance to query count 818c20c1960SStella Laurenzoand begin/end iterators. Prefer hoisting these to dedicated pseudo containers. 819c20c1960SStella Laurenzo 820c20c1960SStella LaurenzoFor example, a direct mapping of blocks within regions could be done this way: 821c20c1960SStella Laurenzo 822c20c1960SStella Laurenzo```python 823c20c1960SStella Laurenzoregion = ... 824c20c1960SStella Laurenzo 825c20c1960SStella Laurenzofor block in region: 826c20c1960SStella Laurenzo 827c20c1960SStella Laurenzo pass 828c20c1960SStella Laurenzo``` 829c20c1960SStella Laurenzo 830c20c1960SStella LaurenzoHowever, this way is preferred: 831c20c1960SStella Laurenzo 832c20c1960SStella Laurenzo```python 833c20c1960SStella Laurenzoregion = ... 834c20c1960SStella Laurenzo 835c20c1960SStella Laurenzofor block in region.blocks: 836c20c1960SStella Laurenzo 837c20c1960SStella Laurenzo pass 838c20c1960SStella Laurenzo 839c20c1960SStella Laurenzoprint(len(region.blocks)) 840c20c1960SStella Laurenzoprint(region.blocks[0]) 841c20c1960SStella Laurenzoprint(region.blocks[-1]) 842c20c1960SStella Laurenzo``` 843c20c1960SStella Laurenzo 844c20c1960SStella LaurenzoInstead of leaking STL-derived identifiers (`front`, `back`, etc), translate 845c20c1960SStella Laurenzothem to appropriate `__dunder__` methods and iterator wrappers in the bindings. 846c20c1960SStella Laurenzo 847c20c1960SStella LaurenzoNote that this can be taken too far, so use good judgment. For example, block 848c20c1960SStella Laurenzoarguments may appear container-like but have defined methods for lookup and 849c20c1960SStella Laurenzomutation that would be hard to model properly without making semantics 850c20c1960SStella Laurenzocomplicated. If running into these, just mirror the C/C++ API. 851c20c1960SStella Laurenzo 852c20c1960SStella Laurenzo### Provide one stop helpers for common things 853c20c1960SStella Laurenzo 854c20c1960SStella LaurenzoOne stop helpers that aggregate over multiple low level entities can be 855c20c1960SStella Laurenzoincredibly helpful and are encouraged within reason. For example, making 856c20c1960SStella Laurenzo`Context` have a `parse_asm` or equivalent that avoids needing to explicitly 857c20c1960SStella Laurenzoconstruct a SourceMgr can be quite nice. One stop helpers do not have to be 858c20c1960SStella Laurenzomutually exclusive with a more complete mapping of the backing constructs. 859c20c1960SStella Laurenzo 860c20c1960SStella Laurenzo## Testing 861c20c1960SStella Laurenzo 862c20c1960SStella LaurenzoTests should be added in the `test/Bindings/Python` directory and should 863c20c1960SStella Laurenzotypically be `.py` files that have a lit run line. 864c20c1960SStella Laurenzo 865417f6137SStella LaurenzoWe use `lit` and `FileCheck` based tests: 866c20c1960SStella Laurenzo 867c20c1960SStella Laurenzo* For generative tests (those that produce IR), define a Python module that 868c20c1960SStella Laurenzo constructs/prints the IR and pipe it through `FileCheck`. 869c20c1960SStella Laurenzo* Parsing should be kept self-contained within the module under test by use of 870c20c1960SStella Laurenzo raw constants and an appropriate `parse_asm` call. 871c20c1960SStella Laurenzo* Any file I/O code should be staged through a tempfile vs relying on file 872c20c1960SStella Laurenzo artifacts/paths outside of the test module. 873417f6137SStella Laurenzo* For convenience, we also test non-generative API interactions with the same 874417f6137SStella Laurenzo mechanisms, printing and `CHECK`ing as needed. 875c20c1960SStella Laurenzo 876c20c1960SStella Laurenzo### Sample FileCheck test 877c20c1960SStella Laurenzo 878c20c1960SStella Laurenzo```python 879c20c1960SStella Laurenzo# RUN: %PYTHON %s | mlir-opt -split-input-file | FileCheck 880c20c1960SStella Laurenzo 881c20c1960SStella Laurenzo# TODO: Move to a test utility class once any of this actually exists. 882c20c1960SStella Laurenzodef print_module(f): 883c20c1960SStella Laurenzo m = f() 884c20c1960SStella Laurenzo print("// -----") 885c20c1960SStella Laurenzo print("// TEST_FUNCTION:", f.__name__) 886c20c1960SStella Laurenzo print(m.to_asm()) 887c20c1960SStella Laurenzo return f 888c20c1960SStella Laurenzo 889c20c1960SStella Laurenzo# CHECK-LABEL: TEST_FUNCTION: create_my_op 890c20c1960SStella Laurenzo@print_module 891c20c1960SStella Laurenzodef create_my_op(): 892c20c1960SStella Laurenzo m = mlir.ir.Module() 893c20c1960SStella Laurenzo builder = m.new_op_builder() 894c20c1960SStella Laurenzo # CHECK: mydialect.my_operation ... 895c20c1960SStella Laurenzo builder.my_op() 896c20c1960SStella Laurenzo return m 897c20c1960SStella Laurenzo``` 898d9b6e4d5SStella Laurenzo 899d9b6e4d5SStella Laurenzo## Integration with ODS 900d9b6e4d5SStella Laurenzo 901d9b6e4d5SStella LaurenzoThe MLIR Python bindings integrate with the tablegen-based ODS system for 902a54f4eaeSMogballproviding user-friendly wrappers around MLIR dialects and operations. There are 903a54f4eaeSMogballmultiple parts to this integration, outlined below. Most details have been 904a54f4eaeSMogballelided: refer to the build rules and python sources under `mlir.dialects` for 905a54f4eaeSMogballthe canonical way to use this facility. 906d9b6e4d5SStella Laurenzo 907a54f4eaeSMogballUsers are responsible for providing a `{DIALECT_NAMESPACE}.py` (or an equivalent 908a54f4eaeSMogballdirectory with `__init__.py` file) as the entrypoint. 909e31c77b1SStella Laurenzo 910e31c77b1SStella Laurenzo### Generating `_{DIALECT_NAMESPACE}_ops_gen.py` wrapper modules 911d9b6e4d5SStella Laurenzo 912d9b6e4d5SStella LaurenzoEach dialect with a mapping to python requires that an appropriate 913e31c77b1SStella Laurenzo`_{DIALECT_NAMESPACE}_ops_gen.py` wrapper module is created. This is done by 914e31c77b1SStella Laurenzoinvoking `mlir-tblgen` on a python-bindings specific tablegen wrapper that 915e31c77b1SStella Laurenzoincludes the boilerplate and actual dialect specific `td` file. An example, for 91623aa5a74SRiver Riddlethe `Func` (which is assigned the namespace `func` as a special case): 917d9b6e4d5SStella Laurenzo 918d9b6e4d5SStella Laurenzo```tablegen 91923aa5a74SRiver Riddle#ifndef PYTHON_BINDINGS_FUNC_OPS 92023aa5a74SRiver Riddle#define PYTHON_BINDINGS_FUNC_OPS 921d9b6e4d5SStella Laurenzo 92223aa5a74SRiver Riddleinclude "mlir/Dialect/Func/IR/FuncOps.td" 923d9b6e4d5SStella Laurenzo 92423aa5a74SRiver Riddle#endif // PYTHON_BINDINGS_FUNC_OPS 925d9b6e4d5SStella Laurenzo``` 926d9b6e4d5SStella Laurenzo 927d9b6e4d5SStella LaurenzoIn the main repository, building the wrapper is done via the CMake function 928edcac733SDenys Shabalin`declare_mlir_dialect_python_bindings`, which invokes: 929d9b6e4d5SStella Laurenzo 930d9b6e4d5SStella Laurenzo``` 93171b6b010SStella Laurenzomlir-tblgen -gen-python-op-bindings -bind-dialect={DIALECT_NAMESPACE} \ 932d9b6e4d5SStella Laurenzo {PYTHON_BINDING_TD_FILE} 933d9b6e4d5SStella Laurenzo``` 934d9b6e4d5SStella Laurenzo 935e31c77b1SStella LaurenzoThe generates op classes must be included in the `{DIALECT_NAMESPACE}.py` file 936e31c77b1SStella Laurenzoin a similar way that generated headers are included for C++ generated code: 937e31c77b1SStella Laurenzo 938e31c77b1SStella Laurenzo```python 939e31c77b1SStella Laurenzofrom ._my_dialect_ops_gen import * 940e31c77b1SStella Laurenzo``` 941e31c77b1SStella Laurenzo 942d9b6e4d5SStella Laurenzo### Extending the search path for wrapper modules 943d9b6e4d5SStella Laurenzo 944d9b6e4d5SStella LaurenzoWhen the python bindings need to locate a wrapper module, they consult the 945a54f4eaeSMogball`dialect_search_path` and use it to find an appropriately named module. For the 946a54f4eaeSMogballmain repository, this search path is hard-coded to include the `mlir.dialects` 947286a7a40SMarkus Böckmodule, which is where wrappers are emitted by the above build rule. Out of tree 9485192e299SMaksim Leventaldialects can add their modules to the search path by calling: 949d9b6e4d5SStella Laurenzo 950d9b6e4d5SStella Laurenzo```python 9515192e299SMaksim Leventalfrom mlir.dialects._ods_common import _cext 9525192e299SMaksim Levental_cext.globals.append_dialect_search_prefix("myproject.mlir.dialects") 953d9b6e4d5SStella Laurenzo``` 954d9b6e4d5SStella Laurenzo 955d9b6e4d5SStella Laurenzo### Wrapper module code organization 956d9b6e4d5SStella Laurenzo 957d9b6e4d5SStella LaurenzoThe wrapper module tablegen emitter outputs: 958d9b6e4d5SStella Laurenzo 959d9b6e4d5SStella Laurenzo* A `_Dialect` class (extending `mlir.ir.Dialect`) with a `DIALECT_NAMESPACE` 960d9b6e4d5SStella Laurenzo attribute. 961d9b6e4d5SStella Laurenzo* An `{OpName}` class for each operation (extending `mlir.ir.OpView`). 962d9b6e4d5SStella Laurenzo* Decorators for each of the above to register with the system. 963d9b6e4d5SStella Laurenzo 964d9b6e4d5SStella LaurenzoNote: In order to avoid naming conflicts, all internal names used by the wrapper 965d9b6e4d5SStella Laurenzomodule are prefixed by `_ods_`. 966d9b6e4d5SStella Laurenzo 96771b6b010SStella LaurenzoEach concrete `OpView` subclass further defines several public-intended 96871b6b010SStella Laurenzoattributes: 969d9b6e4d5SStella Laurenzo 970d9b6e4d5SStella Laurenzo* `OPERATION_NAME` attribute with the `str` fully qualified operation name 97100f7096dSJeff Niu (i.e. `math.absf`). 972d9b6e4d5SStella Laurenzo* An `__init__` method for the *default builder* if one is defined or inferred 973d9b6e4d5SStella Laurenzo for the operation. 974d9b6e4d5SStella Laurenzo* `@property` getter for each operand or result (using an auto-generated name 975d9b6e4d5SStella Laurenzo for unnamed of each). 976d9b6e4d5SStella Laurenzo* `@property` getter, setter and deleter for each declared attribute. 977d9b6e4d5SStella Laurenzo 97871b6b010SStella LaurenzoIt further emits additional private-intended attributes meant for subclassing 979a54f4eaeSMogballand customization (default cases omit these attributes in favor of the defaults 980a54f4eaeSMogballon `OpView`): 98171b6b010SStella Laurenzo 98271b6b010SStella Laurenzo* `_ODS_REGIONS`: A specification on the number and types of regions. 98371b6b010SStella Laurenzo Currently a tuple of (min_region_count, has_no_variadic_regions). Note that 98471b6b010SStella Laurenzo the API does some light validation on this but the primary purpose is to 98571b6b010SStella Laurenzo capture sufficient information to perform other default building and region 98671b6b010SStella Laurenzo accessor generation. 98771b6b010SStella Laurenzo* `_ODS_OPERAND_SEGMENTS` and `_ODS_RESULT_SEGMENTS`: Black-box value which 98871b6b010SStella Laurenzo indicates the structure of either the operand or results with respect to 98971b6b010SStella Laurenzo variadics. Used by `OpView._ods_build_default` to decode operand and result 99071b6b010SStella Laurenzo lists that contain lists. 99171b6b010SStella Laurenzo 992594e0ba9SStella Laurenzo#### Default Builder 993d9b6e4d5SStella Laurenzo 994d9b6e4d5SStella LaurenzoPresently, only a single, default builder is mapped to the `__init__` method. 995a54f4eaeSMogballThe intent is that this `__init__` method represents the *most specific* of the 996a54f4eaeSMogballbuilders typically generated for C++; however currently it is just the generic 997a54f4eaeSMogballform below. 998d9b6e4d5SStella Laurenzo 999d9b6e4d5SStella Laurenzo* One argument for each declared result: 1000d9b6e4d5SStella Laurenzo * For single-valued results: Each will accept an `mlir.ir.Type`. 1001d9b6e4d5SStella Laurenzo * For variadic results: Each will accept a `List[mlir.ir.Type]`. 1002d9b6e4d5SStella Laurenzo* One argument for each declared operand or attribute: 1003d9b6e4d5SStella Laurenzo * For single-valued operands: Each will accept an `mlir.ir.Value`. 1004d9b6e4d5SStella Laurenzo * For variadic operands: Each will accept a `List[mlir.ir.Value]`. 1005d9b6e4d5SStella Laurenzo * For attributes, it will accept an `mlir.ir.Attribute`. 1006d9b6e4d5SStella Laurenzo* Trailing usage-specific, optional keyword arguments: 1007d9b6e4d5SStella Laurenzo * `loc`: An explicit `mlir.ir.Location` to use. Defaults to the location 1008a54f4eaeSMogball bound to the thread (i.e. `with Location.unknown():`) or an error if 1009a54f4eaeSMogball none is bound nor specified. 1010a54f4eaeSMogball * `ip`: An explicit `mlir.ir.InsertionPoint` to use. Default to the 1011a54f4eaeSMogball insertion point bound to the thread (i.e. `with InsertionPoint(...):`). 1012fd226c9bSStella Laurenzo 1013fd226c9bSStella LaurenzoIn addition, each `OpView` inherits a `build_generic` method which allows 1014fd226c9bSStella Laurenzoconstruction via a (nested in the case of variadic) sequence of `results` and 1015fd226c9bSStella Laurenzo`operands`. This can be used to get some default construction semantics for 1016a54f4eaeSMogballoperations that are otherwise unsupported in Python, at the expense of having a 1017a54f4eaeSMogballvery generic signature. 1018594e0ba9SStella Laurenzo 1019594e0ba9SStella Laurenzo#### Extending Generated Op Classes 1020594e0ba9SStella Laurenzo 1021594e0ba9SStella LaurenzoAs mentioned above, the build system generates Python sources like 1022a54f4eaeSMogball`_{DIALECT_NAMESPACE}_ops_gen.py` for each dialect with Python bindings. It is 1023a2288a89SMaksim Leventaloften desirable to use these generated classes as a starting point for 1024a2288a89SMaksim Leventalfurther customization, so an extension mechanism is provided to make this easy. 1025a2288a89SMaksim LeventalThis mechanism uses conventional inheritance combined with `OpView` registration. 1026a2288a89SMaksim LeventalFor example, the default builder for `arith.constant` 1027594e0ba9SStella Laurenzo 1028594e0ba9SStella Laurenzo```python 1029a2288a89SMaksim Leventalclass ConstantOp(_ods_ir.OpView): 1030a2288a89SMaksim Levental OPERATION_NAME = "arith.constant" 1031a2288a89SMaksim Levental 1032a2288a89SMaksim Levental _ODS_REGIONS = (0, True) 1033a2288a89SMaksim Levental 1034a2288a89SMaksim Levental def __init__(self, value, *, loc=None, ip=None): 1035a2288a89SMaksim Levental ... 1036594e0ba9SStella Laurenzo``` 1037594e0ba9SStella Laurenzo 1038a2288a89SMaksim Leventalexpects `value` to be a `TypedAttr` (e.g., `IntegerAttr` or `FloatAttr`). 1039a2288a89SMaksim LeventalThus, a natural extension is a builder that accepts a MLIR type and a Python value and instantiates the appropriate `TypedAttr`: 1040594e0ba9SStella Laurenzo 1041594e0ba9SStella Laurenzo```python 1042a2288a89SMaksim Leventalfrom typing import Union 1043a2288a89SMaksim Levental 1044a2288a89SMaksim Leventalfrom mlir.ir import Type, IntegerAttr, FloatAttr 1045a2288a89SMaksim Leventalfrom mlir.dialects._arith_ops_gen import _Dialect, ConstantOp 1046a2288a89SMaksim Leventalfrom mlir.dialects._ods_common import _cext 1047a2288a89SMaksim Levental 1048a2288a89SMaksim Levental@_cext.register_operation(_Dialect, replace=True) 1049a2288a89SMaksim Leventalclass ConstantOpExt(ConstantOp): 1050a2288a89SMaksim Levental def __init__( 1051a2288a89SMaksim Levental self, result: Type, value: Union[int, float], *, loc=None, ip=None 1052a2288a89SMaksim Levental ): 1053a2288a89SMaksim Levental if isinstance(value, int): 1054a2288a89SMaksim Levental super().__init__(IntegerAttr.get(result, value), loc=loc, ip=ip) 1055a2288a89SMaksim Levental elif isinstance(value, float): 1056a2288a89SMaksim Levental super().__init__(FloatAttr.get(result, value), loc=loc, ip=ip) 1057a2288a89SMaksim Levental else: 1058a2288a89SMaksim Levental raise NotImplementedError(f"Building `arith.constant` not supported for {result=} {value=}") 1059594e0ba9SStella Laurenzo``` 1060594e0ba9SStella Laurenzo 1061a2288a89SMaksim Leventalwhich enables building an instance of `arith.constant` like so: 1062594e0ba9SStella Laurenzo 1063594e0ba9SStella Laurenzo```python 1064a2288a89SMaksim Leventalfrom mlir.ir import F32Type 1065a2288a89SMaksim Levental 1066a2288a89SMaksim Leventala = ConstantOpExt(F32Type.get(), 42.42) 1067a2288a89SMaksim Leventalb = ConstantOpExt(IntegerType.get_signless(32), 42) 1068594e0ba9SStella Laurenzo``` 1069594e0ba9SStella Laurenzo 1070a2288a89SMaksim LeventalNote, three key aspects of the extension mechanism in this example: 1071a2288a89SMaksim Levental 1072a2288a89SMaksim Levental1. `ConstantOpExt` directly inherits from the generated `ConstantOp`; 1073a2288a89SMaksim Levental2. in this, simplest, case all that's required is a call to the super class' initializer, i.e., `super().__init__(...)`; 1074a2288a89SMaksim Levental3. in order to register `ConstantOpExt` as the preferred `OpView` that is returned by `mlir.ir.Operation.opview` (see [Operations, Regions and Blocks](#operations-regions-and-blocks)) 1075a2288a89SMaksim Levental we decorate the class with `@_cext.register_operation(_Dialect, replace=True)`, **where the `replace=True` must be used**. 1076a2288a89SMaksim Levental 1077a2288a89SMaksim LeventalIn some more complex cases it might be necessary to explicitly build the `OpView` through `OpView.build_generic` (see [Default Builder](#default-builder)), just as is performed by the generated builders. 1078a2288a89SMaksim LeventalI.e., we must call `OpView.build_generic` **and pass the result to `OpView.__init__`**, where the small issue becomes that the latter is already overridden by the generated builder. 1079a2288a89SMaksim LeventalThus, we must call a method of a super class' super class (the "grandparent"); for example: 1080a2288a89SMaksim Levental 1081a2288a89SMaksim Levental```python 1082a2288a89SMaksim Leventalfrom mlir.dialects._scf_ops_gen import _Dialect, ForOp 1083a2288a89SMaksim Leventalfrom mlir.dialects._ods_common import _cext 1084a2288a89SMaksim Levental 1085a2288a89SMaksim Levental@_cext.register_operation(_Dialect, replace=True) 1086a2288a89SMaksim Leventalclass ForOpExt(ForOp): 1087a2288a89SMaksim Levental def __init__(self, lower_bound, upper_bound, step, iter_args, *, loc=None, ip=None): 1088a2288a89SMaksim Levental ... 1089a2288a89SMaksim Levental super(ForOp, self).__init__(self.build_generic(...)) 1090a2288a89SMaksim Levental``` 1091a2288a89SMaksim Levental 1092a2288a89SMaksim Leventalwhere `OpView.__init__` is called via `super(ForOp, self).__init__`. 1093a2288a89SMaksim LeventalNote, there are alternatives ways to implement this (e.g., explicitly writing `OpView.__init__`); see any discussion on Python inheritance. 109451460675SAlex Zinenko 109551460675SAlex Zinenko## Providing Python bindings for a dialect 109651460675SAlex Zinenko 109751460675SAlex ZinenkoPython bindings are designed to support MLIR’s open dialect ecosystem. A dialect 109851460675SAlex Zinenkocan be exposed to Python as a submodule of `mlir.dialects` and interoperate with 109951460675SAlex Zinenkothe rest of the bindings. For dialects containing only operations, it is 110051460675SAlex Zinenkosufficient to provide Python APIs for those operations. Note that the majority 110151460675SAlex Zinenkoof boilerplate APIs can be generated from ODS. For dialects containing 110251460675SAlex Zinenkoattributes and types, it is necessary to thread those through the C API since 110351460675SAlex Zinenkothere is no generic mechanism to create attributes and types. Passes need to be 110451460675SAlex Zinenkoregistered with the context in order to be usable in a text-specified pass 110551460675SAlex Zinenkomanager, which may be done at Python module load time. Other functionality can 110651460675SAlex Zinenkobe provided, similar to attributes and types, by exposing the relevant C API and 110751460675SAlex Zinenkobuilding Python API on top. 110851460675SAlex Zinenko 110951460675SAlex Zinenko 111051460675SAlex Zinenko### Operations 111151460675SAlex Zinenko 111251460675SAlex ZinenkoDialect operations are provided in Python by wrapping the generic 111351460675SAlex Zinenko`mlir.ir.Operation` class with operation-specific builder functions and 111451460675SAlex Zinenkoproperties. Therefore, there is no need to implement a separate C API for them. 111551460675SAlex ZinenkoFor operations defined in ODS, `mlir-tblgen -gen-python-op-bindings 111651460675SAlex Zinenko-bind-dialect=<dialect-namespace>` generates the Python API from the declarative 111767a910bbSRahul Kayaithdescription. 111867a910bbSRahul KayaithIt is sufficient to create a new `.td` file that includes the original ODS 111967a910bbSRahul Kayaithdefinition and use it as source for the `mlir-tblgen` call. 112067a910bbSRahul KayaithSuch `.td` files reside in 112151460675SAlex Zinenko[`python/mlir/dialects/`](https://github.com/llvm/llvm-project/tree/main/mlir/python/mlir/dialects). 112251460675SAlex ZinenkoThe results of `mlir-tblgen` are expected to produce a file named 112351460675SAlex Zinenko`_<dialect-namespace>_ops_gen.py` by convention. The generated operation classes 112451460675SAlex Zinenkocan be extended as described above. MLIR provides [CMake 112551460675SAlex Zinenkofunctions](https://github.com/llvm/llvm-project/blob/main/mlir/cmake/modules/AddMLIRPython.cmake) 112651460675SAlex Zinenkoto automate the production of such files. Finally, a 112751460675SAlex Zinenko`python/mlir/dialects/<dialect-namespace>.py` or a 112851460675SAlex Zinenko`python/mlir/dialects/<dialect-namespace>/__init__.py` file must be created and 112951460675SAlex Zinenkofilled with `import`s from the generated files to enable `import 113051460675SAlex Zinenkomlir.dialects.<dialect-namespace>` in Python. 113151460675SAlex Zinenko 113251460675SAlex Zinenko 113351460675SAlex Zinenko### Attributes and Types 113451460675SAlex Zinenko 113551460675SAlex ZinenkoDialect attributes and types are provided in Python as subclasses of the 113651460675SAlex Zinenko`mlir.ir.Attribute` and `mlir.ir.Type` classes, respectively. Python APIs for 113751460675SAlex Zinenkoattributes and types must connect to the relevant C APIs for building and 113851460675SAlex Zinenkoinspection, which must be provided first. Bindings for `Attribute` and `Type` 113951460675SAlex Zinenkosubclasses can be defined using 114051460675SAlex Zinenko[`include/mlir/Bindings/Python/PybindAdaptors.h`](https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Bindings/Python/PybindAdaptors.h) 1141392622d0SMaksim Leventalor 1142392622d0SMaksim Levental[`include/mlir/Bindings/Python/NanobindAdaptors.h`](https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h) 1143392622d0SMaksim Leventalutilities that mimic pybind11/nanobind API for defining functions and 1144392622d0SMaksim Leventalproperties. These bindings are to be included in a separate module. The 1145392622d0SMaksim Leventalutilities also provide automatic casting between C API handles `MlirAttribute` 1146392622d0SMaksim Leventaland `MlirType` and their Python counterparts so that the C API handles can be 1147392622d0SMaksim Leventalused directly in binding implementations. The methods and properties provided by 1148392622d0SMaksim Leventalthe bindings should follow the principles discussed above. 114951460675SAlex Zinenko 115051460675SAlex ZinenkoThe attribute and type bindings for a dialect can be located in 115151460675SAlex Zinenko`lib/Bindings/Python/Dialect<Name>.cpp` and should be compiled into a separate 115251460675SAlex Zinenko“Python extension” library placed in `python/mlir/_mlir_libs` that will be 115351460675SAlex Zinenkoloaded by Python at runtime. MLIR provides [CMake 115451460675SAlex Zinenkofunctions](https://github.com/llvm/llvm-project/blob/main/mlir/cmake/modules/AddMLIRPython.cmake) 115551460675SAlex Zinenkoto automate the production of such libraries. This library should be `import`ed 115651460675SAlex Zinenkofrom the main dialect file, i.e. `python/mlir/dialects/<dialect-namespace>.py` 115751460675SAlex Zinenkoor `python/mlir/dialects/<dialect-namespace>/__init__.py`, to ensure the types 115851460675SAlex Zinenkoare available when the dialect is loaded from Python. 115951460675SAlex Zinenko 116051460675SAlex Zinenko 116151460675SAlex Zinenko### Passes 116251460675SAlex Zinenko 116351460675SAlex ZinenkoDialect-specific passes can be made available to the pass manager in Python by 116451460675SAlex Zinenkoregistering them with the context and relying on the API for pass pipeline 116551460675SAlex Zinenkoparsing from string descriptions. This can be achieved by creating a new 116651460675SAlex Zinenkopybind11 module, defined in `lib/Bindings/Python/<Dialect>Passes.cpp`, that 116751460675SAlex Zinenkocalls the registration C API, which must be provided first. For passes defined 116851460675SAlex Zinenkodeclaratively using Tablegen, `mlir-tblgen -gen-pass-capi-header` and 116951460675SAlex Zinenko`-mlir-tblgen -gen-pass-capi-impl` automate the generation of C API. The 117051460675SAlex Zinenkopybind11 module must be compiled into a separate “Python extension” library, 117151460675SAlex Zinenkowhich can be `import`ed from the main dialect file, i.e. 117251460675SAlex Zinenko`python/mlir/dialects/<dialect-namespace>.py` or 117351460675SAlex Zinenko`python/mlir/dialects/<dialect-namespace>/__init__.py`, or from a separate 117451460675SAlex Zinenko`passes` submodule to be put in 117551460675SAlex Zinenko`python/mlir/dialects/<dialect-namespace>/passes.py` if it is undesirable to 117651460675SAlex Zinenkomake the passes available along with the dialect. 117751460675SAlex Zinenko 117851460675SAlex Zinenko 117951460675SAlex Zinenko### Other functionality 118051460675SAlex Zinenko 118151460675SAlex ZinenkoDialect functionality other than IR objects or passes, such as helper functions, 118251460675SAlex Zinenkocan be exposed to Python similarly to attributes and types. C API is expected to 118351460675SAlex Zinenkoexist for this functionality, which can then be wrapped using pybind11 and 118496f8cfe4Svfdev[`include/mlir/Bindings/Python/PybindAdaptors.h`](https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Bindings/Python/PybindAdaptors.h), 1185392622d0SMaksim Leventalor nanobind and 118696f8cfe4Svfdev[`include/mlir/Bindings/Python/NanobindAdaptors.h`](https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h) 118751460675SAlex Zinenkoutilities to connect to the rest of Python API. The bindings can be located in a 1188392622d0SMaksim Leventalseparate module or in the same module as attributes and types, and 118951460675SAlex Zinenkoloaded along with the dialect. 1190*f136c800Svfdev 1191*f136c800Svfdev## Free-threading (No-GIL) support 1192*f136c800Svfdev 1193*f136c800SvfdevFree-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/). 1194*f136c800Svfdev 1195*f136c800SvfdevMLIR Python bindings are free-threading compatible with exceptions (discussed below) in the following sense: it is safe to work in multiple threads with **independent** contexts. Below we show an example code of safe usage: 1196*f136c800Svfdev 1197*f136c800Svfdev```python 1198*f136c800Svfdev# python3.13t example.py 1199*f136c800Svfdevimport concurrent.futures 1200*f136c800Svfdev 1201*f136c800Svfdevimport mlir.dialects.arith as arith 1202*f136c800Svfdevfrom mlir.ir import Context, Location, Module, IntegerType, InsertionPoint 1203*f136c800Svfdev 1204*f136c800Svfdev 1205*f136c800Svfdevdef func(py_value): 1206*f136c800Svfdev with Context() as ctx: 1207*f136c800Svfdev module = Module.create(loc=Location.file("foo.txt", 0, 0)) 1208*f136c800Svfdev 1209*f136c800Svfdev dtype = IntegerType.get_signless(64) 1210*f136c800Svfdev with InsertionPoint(module.body), Location.name("a"): 1211*f136c800Svfdev arith.constant(dtype, py_value) 1212*f136c800Svfdev 1213*f136c800Svfdev return module 1214*f136c800Svfdev 1215*f136c800Svfdev 1216*f136c800Svfdevnum_workers = 8 1217*f136c800Svfdevwith concurrent.futures.ThreadPoolExecutor(max_workers=num_workers) as executor: 1218*f136c800Svfdev futures = [] 1219*f136c800Svfdev for i in range(num_workers): 1220*f136c800Svfdev futures.append(executor.submit(func, i)) 1221*f136c800Svfdev assert len(list(f.result() for f in futures)) == num_workers 1222*f136c800Svfdev``` 1223*f136c800Svfdev 1224*f136c800SvfdevThe exceptions to the free-threading compatibility: 1225*f136c800Svfdev- IR printing is unsafe, e.g. when using `PassManager` with `PassManager.enable_ir_printing()` which calls thread-unsafe `llvm::raw_ostream`. 1226*f136c800Svfdev- Usage of `Location.emit_error` is unsafe (due to thread-unsafe `llvm::raw_ostream`). 1227*f136c800Svfdev- Usage of `Module.dump` is unsafe (due to thread-unsafe `llvm::raw_ostream`). 1228*f136c800Svfdev- Usage of `mlir.dialects.transform.interpreter` is unsafe. 1229*f136c800Svfdev- Usage of `mlir.dialects.gpu` and `gpu-module-to-binary` is unsafe.