MLIR Python Bindings

Current status: Under development and not enabled by default

Building ¶

Pre-requisites ¶

• A relatively recent Python3 installation
• Installation of python dependencies as specified in mlir/python/requirements.txt

CMake variables ¶

• MLIR_ENABLE_BINDINGS_PYTHON:BOOL

  Enables building the Python bindings. Defaults to OFF.

• Python3_EXECUTABLE:STRING

  Specifies the python executable used for the LLVM build, including for determining header/link flags for the Python bindings. On systems with multiple Python implementations, setting this explicitly to the preferred python3 executable is strongly recommended.

Recommended development practices ¶

It is recommended to use a python virtual environment. Many ways exist for this, but the following is the simplest:

# Make sure your 'python' is what you expect. Note that on multi-python
# systems, this may have a version suffix, and on many Linuxes and MacOS where
# python2 and python3 co-exist, you may also want to use `python3`.
which python
python -m venv ~/.venv/mlirdev
source ~/.venv/mlirdev/bin/activate

# Note that many LTS distros will bundle a version of pip itself that is too
# old to download all of the latest binaries for certain platforms.
# The pip version can be obtained with `python -m pip --version`, and for
# Linux specifically, this should be cross checked with minimum versions
# here: https://github.com/pypa/manylinux
# It is recommended to upgrade pip:
python -m pip install --upgrade pip


# Now the `python` command will resolve to your virtual environment and
# packages will be installed there.
python -m pip install -r mlir/python/requirements.txt

# Now run `cmake`, `ninja`, et al.

For interactive use, it is sufficient to add the tools/mlir/python_packages/mlir_core/ directory in your build/ directory to the PYTHONPATH. Typically:

export PYTHONPATH=$(cd build && pwd)/tools/mlir/python_packages/mlir_core

Note that if you have installed (i.e. via ninja install, et al), then python packages for all enabled projects will be in your install tree under python_packages/ (i.e. python_packages/mlir_core). Official distributions are built with a more specialized setup.

Design ¶

Use cases ¶

There are likely two primary use cases for the MLIR python bindings:

1. Support users who expect that an installed version of LLVM/MLIR will yield the ability to import mlir and use the API in a pure way out of the box.

2. Downstream integrations will likely want to include parts of the API in their private namespace or specially built libraries, probably mixing it with other python native bits.

Composable modules ¶

In order to support use case #2, the Python bindings are organized into composable modules that downstream integrators can include and re-export into their own namespace if desired. This forces several design points:

• Separate the construction/populating of a py::module from PYBIND11_MODULE global constructor.

• Introduce headers for C++-only wrapper classes as other related C++ modules will need to interop with it.

• Separate any initialization routines that depend on optional components into its own module/dependency (currently, things like registerAllDialects fall into this category).

There are a lot of co-related issues of shared library linkage, distribution concerns, etc that affect such things. Organizing the code into composable modules (versus a monolithic cpp file) allows the flexibility to address many of these as needed over time. Also, compilation time for all of the template meta-programming in pybind scales with the number of things you define in a translation unit. Breaking into multiple translation units can significantly aid compile times for APIs with a large surface area.

Submodules ¶

Generally, the C++ codebase namespaces most things into the mlir namespace. However, in order to modularize and make the Python bindings easier to understand, sub-packages are defined that map roughly to the directory structure of functional units in MLIR.

Examples:

• mlir.ir
• mlir.passes (pass is a reserved word :( )
• mlir.dialect
• mlir.execution_engine (aside from namespacing, it is important that “bulky”/optional parts like this are isolated)

In addition, initialization functions that imply optional dependencies should be in underscored (notionally private) modules such as _init and linked separately. This allows downstream integrators to completely customize what is included “in the box” and covers things like dialect registration, pass registration, etc.

Loader ¶

LLVM/MLIR is a non-trivial python-native project that is likely to co-exist with other non-trivial native extensions. As such, the native extension (i.e. the .so/.pyd/.dylib) is exported as a notionally private top-level symbol (_mlir), while a small set of Python code is provided in mlir/_cext_loader.py and siblings which loads and re-exports it. This split provides a place to stage code that needs to prepare the environment before the shared library is loaded into the Python runtime, and also provides a place that one-time initialization code can be invoked apart from module constructors.

It is recommended to avoid using __init__.py files to the extent possible, until reaching a leaf package that represents a discrete component. The rule to keep in mind is that the presence of an __init__.py file prevents the ability to split anything at that level or below in the namespace into different directories, deployment packages, wheels, etc.

See the documentation for more information and advice: https://packaging.python.org/guides/packaging-namespace-packages/

Use the C-API ¶

The Python APIs should seek to layer on top of the C-API to the degree possible. Especially for the core, dialect-independent parts, such a binding enables packaging decisions that would be difficult or impossible if spanning a C++ ABI boundary. In addition, factoring in this way side-steps some very difficult issues that arise when combining RTTI-based modules (which pybind derived things are) with non-RTTI polymorphic C++ code (the default compilation mode of LLVM).

Ownership in the Core IR ¶

There are several top-level types in the core IR that are strongly owned by their python-side reference:

• PyContext (mlir.ir.Context)
• PyModule (mlir.ir.Module)
• PyOperation (mlir.ir.Operation) - but with caveats

All other objects are dependent. All objects maintain a back-reference (keep-alive) to their closest containing top-level object. Further, dependent objects fall into two categories: a) uniqued (which live for the life-time of the context) and b) mutable. Mutable objects need additional machinery for keeping track of when the C++ instance that backs their Python object is no longer valid (typically due to some specific mutation of the IR, deletion, or bulk operation).

Optionality and argument ordering in the Core IR ¶

The following types support being bound to the current thread as a context manager:

• PyLocation (loc: mlir.ir.Location = None)
• PyInsertionPoint (ip: mlir.ir.InsertionPoint = None)
• PyMlirContext (context: mlir.ir.Context = None)

In order to support composability of function arguments, when these types appear as arguments, they should always be the last and appear in the above order and with the given names (which is generally the order in which they are expected to need to be expressed explicitly in special cases) as necessary. Each should carry a default value of py::none() and use either a manual or automatic conversion for resolving either with the explicit value or a value from the thread context manager (i.e. DefaultingPyMlirContext or DefaultingPyLocation).

The rationale for this is that in Python, trailing keyword arguments to the right are the most composable, enabling a variety of strategies such as kwarg passthrough, default values, etc. Keeping function signatures composable increases the chances that interesting DSLs and higher level APIs can be constructed without a lot of exotic boilerplate.

Used consistently, this enables a style of IR construction that rarely needs to use explicit contexts, locations, or insertion points but is free to do so when extra control is needed.

Operation hierarchy ¶

As mentioned above, PyOperation is special because it can exist in either a top-level or dependent state. The life-cycle is unidirectional: operations can be created detached (top-level) and once added to another operation, they are then dependent for the remainder of their lifetime. The situation is more complicated when considering construction scenarios where an operation is added to a transitive parent that is still detached, necessitating further accounting at such transition points (i.e. all such added children are initially added to the IR with a parent of their outer-most detached operation, but then once it is added to an attached operation, they need to be re-parented to the containing module).

Due to the validity and parenting accounting needs, PyOperation is the owner for regions and blocks and needs to be a top-level type that we can count on not aliasing. This let’s us do things like selectively invalidating instances when mutations occur without worrying that there is some alias to the same operation in the hierarchy. Operations are also the only entity that are allowed to be in a detached state, and they are interned at the context level so that there is never more than one Python mlir.ir.Operation object for a unique MlirOperation, regardless of how it is obtained.

The C/C++ API allows for Region/Block to also be detached, but it simplifies the ownership model a lot to eliminate that possibility in this API, allowing the Region/Block to be completely dependent on its owning operation for accounting. The aliasing of Python Region/Block instances to underlying MlirRegion/MlirBlock is considered benign and these objects are not interned in the context (unlike operations).

If we ever want to re-introduce detached regions/blocks, we could do so with new “DetachedRegion” class or similar and also avoid the complexity of accounting. With the way it is now, we can avoid having a global live list for regions and blocks. We may end up needing an op-local one at some point TBD, depending on how hard it is to guarantee how mutations interact with their Python peer objects. We can cross that bridge easily when we get there.

Module, when used purely from the Python API, can’t alias anyway, so we can use it as a top-level ref type without a live-list for interning. If the API ever changes such that this cannot be guaranteed (i.e. by letting you marshal a native-defined Module in), then there would need to be a live table for it too.

User-level API ¶

Context Management ¶

The bindings rely on Python context managers (with statements) to simplify creation and handling of IR objects by omitting repeated arguments such as MLIR contexts, operation insertion points and locations. A context manager sets up the default object to be used by all binding calls within the following context and in the same thread. This default can be overridden by specific calls through the dedicated keyword arguments.

MLIR Context ¶

An MLIR context is a top-level entity that owns attributes and types and is referenced from virtually all IR constructs. Contexts also provide thread safety at the C++ level. In Python bindings, the MLIR context is also a Python context manager, one can write:

from mlir.ir import Context, Module

with Context() as ctx:
  # IR construction using `ctx` as context.

  # For example, parsing an MLIR module from string requires the context.
  Module.parse("builtin.module {}")

IR objects referencing a context usually provide access to it through the .context property. Most IR-constructing functions expect the context to be provided in some form. In case of attributes and types, the context may be extracted from the contained attribute or type. In case of operations, the context is systematically extracted from Locations (see below). When the context cannot be extracted from any argument, the bindings API expects the (keyword) argument context. If it is not provided or set to None (default), it will be looked up from an implicit stack of contexts maintained by the bindings in the current thread and updated by context managers. If there is no surrounding context, an error will be raised.

Note that it is possible to manually specify the MLIR context both inside and outside of the with statement:

from mlir.ir import Context, Module

standalone_ctx = Context()
with Context() as managed_ctx:
  # Parse a module in managed_ctx.
  Module.parse("...")

  # Parse a module in standalone_ctx (override the context manager).
  Module.parse("...", context=standalone_ctx)

# Parse a module without using context managers.
Module.parse("...", context=standalone_ctx)

The context object remains live as long as there are IR objects referencing it.

Insertion Points and Locations ¶

When constructing an MLIR operation, two pieces of information are required:

• an insertion point that indicates where the operation is to be created in the IR region/block/operation structure (usually before or after another operation, or at the end of some block); it may be missing, at which point the operation is created in the detached state;
• a location that contains user-understandable information about the source of the operation (for example, file/line/column information), which must always be provided as it carries a reference to the MLIR context.

Both can be provided using context managers or explicitly as keyword arguments in the operation constructor. They can be also provided as keyword arguments ip and loc both within and outside of the context manager.

from mlir.ir import Context, InsertionPoint, Location, Module, Operation

with Context() as ctx:
  module = Module.create()

  # Prepare for inserting operations into the body of the module and indicate
  # that these operations originate in the "f.mlir" file at the given line and
  # column.
  with InsertionPoint(module.body), Location.file("f.mlir", line=42, col=1):
    # This operation will be inserted at the end of the module body and will
    # have the location set up by the context manager.
    Operation(<...>)

    # This operation will be inserted at the end of the module (and after the
    # previously constructed operation) and will have the location provided as
    # the keyword argument.
    Operation(<...>, loc=Location.file("g.mlir", line=1, col=10))

    # This operation will be inserted at the *beginning* of the block rather
    # than at its end.
    Operation(<...>, ip=InsertionPoint.at_block_begin(module.body))

Note that Location needs an MLIR context to be constructed. It can take the context set up in the current thread by some surrounding context manager, or accept it as an explicit argument:

from mlir.ir import Context, Location

# Create a context and a location in this context in the same `with` statement.
with Context() as ctx, Location.file("f.mlir", line=42, col=1, context=ctx):
  pass

Locations are owned by the context and live as long as they are (transitively) referenced from somewhere in Python code.

Unlike locations, the insertion point may be left unspecified (or, equivalently, set to None or False) during operation construction. In this case, the operation is created in the detached state, that is, it is not added into the region of another operation and is owned by the caller. This is usually the case for top-level operations that contain the IR, such as modules. Regions, blocks and values contained in an operation point back to it and maintain it live.

Inspecting IR Objects ¶

Inspecting the IR is one of the primary tasks the Python bindings are designed for. One can traverse the IR operation/region/block structure and inspect their aspects such as operation attributes and value types.

Operations, Regions and Blocks ¶

Operations are represented as either:

• the generic Operation class, useful in particular for generic processing of unregistered operations; or
• a specific subclass of OpView that provides more semantically-loaded accessors to operation properties.

Given an OpView subclass, one can obtain an Operation using its .operation property. Given an Operation, one can obtain the corresponding OpView using its .opview property as long as the corresponding class has been set up. This typically means that the Python module of its dialect has been loaded. By default, the OpView version is produced when navigating the IR tree.

One can check if an operation has a specific type by means of Python’s isinstance function:

operation = <...>
opview = <...>
if isinstance(operation.opview, mydialect.MyOp):
  pass
if isinstance(opview, mydialect.MyOp):
  pass

The components of an operation can be inspected using its properties.

• attributes is a collection of operation attributes . It can be subscripted as both dictionary and sequence, e.g., both operation.attributes["value"] and operation.attributes[0] will work. There is no guarantee on the order in which the attributes are traversed when iterating over the attributes property as sequence.
• operands is a sequence collection of operation operands.
• results is a sequence collection of operation results.
• regions is a sequence collection of regions attached to the operation.

The objects produced by operands and results have a .types property that contains a sequence collection of types of the corresponding values.

from mlir.ir import Operation

operation1 = <...>
operation2 = <...>
if operation1.results.types == operation2.operand.types:
  pass

OpView subclasses for specific operations may provide leaner accessors to properties of an operation. For example, named attributes, operand and results are usually accessible as properties of the OpView subclass with the same name, such as operation.const_value instead of operation.attributes["const_value"]. If this name is a reserved Python keyword, it is suffixed with an underscore.

The operation itself is iterable, which provides access to the attached regions in order:

from mlir.ir import Operation

operation = <...>
for region in operation:
  do_something_with_region(region)

A region is conceptually a sequence of blocks. Objects of the Region class are thus iterable, which provides access to the blocks. One can also use the .blocks property.

# Regions are directly iterable and give access to blocks.
for block1, block2 in zip(operation.regions[0], operation.regions[0].blocks)
  assert block1 == block2

A block contains a sequence of operations, and has several additional properties. Objects of the Block class are iterable and provide access to the operations contained in the block. So does the .operations property. Blocks also have a list of arguments available as a sequence collection using the .arguments property.

Block and region belong to the parent operation in Python bindings and keep it alive. This operation can be accessed using the .owner property.

Attributes and Types ¶

Attributes and types are (mostly) immutable context-owned objects. They are represented as either:

• an opaque Attribute or Type object supporting printing and comparison; or
• a concrete subclass thereof with access to properties of the attribute or type.

Given an Attribute or Type object, one can obtain a concrete subclass using the constructor of the subclass. This may raise a ValueError if the attribute or type is not of the expected subclass:

from mlir.ir import Attribute, Type
from mlir.<dialect> import ConcreteAttr, ConcreteType

attribute = <...>
type = <...>
try:
  concrete_attr = ConcreteAttr(attribute)
  concrete_type = ConcreteType(type)
except ValueError as e:
  # Handle incorrect subclass.

In addition, concrete attribute and type classes provide a static isinstance method to check whether an object of the opaque Attribute or Type type can be downcasted:

from mlir.ir import Attribute, Type
from mlir.<dialect> import ConcreteAttr, ConcreteType

attribute = <...>
type = <...>

# No need to handle errors here.
if ConcreteAttr.isinstance(attribute):
  concrete_attr = ConcreteAttr(attribute)
if ConcreteType.isinstance(type):
  concrete_type = ConcreteType(type)

By default, and unlike operations, attributes and types are returned from IR traversals using the opaque Attribute or Type that needs to be downcasted.

Concrete attribute and type classes usually expose their properties as Python readonly properties. For example, the elemental type of a tensor type can be accessed using the .element_type property.

Values ¶

MLIR has two kinds of values based on their defining object: block arguments and operation results. Values are handled similarly to attributes and types. They are represented as either:

• a generic Value object; or
• a concrete BlockArgument or OpResult object.

The former provides all the generic functionality such as comparison, type access and printing. The latter provide access to the defining block or operation and the position of the value within it. By default, the generic Value objects are returned from IR traversals. Downcasting is implemented through concrete subclass constructors, similarly to attribtues and types:

from mlir.ir import BlockArgument, OpResult, Value

value = ...

# Set `concrete` to the specific value subclass.
try:
  concrete = BlockArgument(value)
except ValueError:
  # This must not raise another ValueError as values are either block arguments
  # or op results.
  concrete = OpResult(value)

Interfaces ¶

MLIR interfaces are a mechanism to interact with the IR without needing to know specific types of operations but only some of their aspects. Operation interfaces are available as Python classes with the same name as their C++ counterparts. Objects of these classes can be constructed from either:

• an object of the Operation class or of any OpView subclass; in this case, all interface methods are available;
• a subclass of OpView and a context; in this case, only the static interface methods are available as there is no associated operation.

In both cases, construction of the interface raises a ValueError if the operation class does not implement the interface in the given context (or, for operations, in the context that the operation is defined in). Similarly to attributes and types, the MLIR context may be set up by a surrounding context manager.

from mlir.ir import Context, InferTypeOpInterface

with Context():
  op = <...>

  # Attempt to cast the operation into an interface.
  try:
    iface = InferTypeOpInterface(op)
  except ValueError:
    print("Operation does not implement InferTypeOpInterface.")
    raise

  # All methods are available on interface objects constructed from an Operation
  # or an OpView.
  iface.someInstanceMethod()

  # An interface object can also be constructed given an OpView subclass. It
  # also needs a context in which the interface will be looked up. The context
  # can be provided explicitly or set up by the surrounding context manager.
  try:
    iface = InferTypeOpInterface(some_dialect.SomeOp)
  except ValueError:
    print("SomeOp does not implement InferTypeOpInterface.")
    raise

  # Calling an instance method on an interface object constructed from a class
  # will raise TypeError.
  try:
    iface.someInstanceMethod()
  except TypeError:
    pass

  # One can still call static interface methods though.
  iface.inferOpReturnTypes(<...>)

If an interface object was constructed from an Operation or an OpView, they are available as .operation and .opview properties of the interface object, respectively.

Only a subset of operation interfaces are currently provided in Python bindings. Attribute and type interfaces are not yet available in Python bindings.

Creating IR Objects ¶

Python bindings also support IR creation and manipulation.

Operations, Regions and Blocks ¶

Operations can be created given a Location and an optional InsertionPoint. It is often easier to user context managers to specify locations and insertion points for several operations created in a row as described above.

Concrete operations can be created by using constructors of the corresponding OpView subclasses. The generic, default form of the constructor accepts:

• an optional sequence of types for operation results (results);
• an optional sequence of values for operation operands, or another operation producing those values (operands);
• an optional dictionary of operation attributes (attributes);
• an optional sequence of successor blocks (successors);
• the number of regions to attach to the operation (regions, default 0);
• the loc keyword argument containing the Location of this operation; if None, the location created by the closest context manager is used or an exception will be raised if there is no context manager;
• the ip keyword argument indicating where the operation will be inserted in the IR; if None, the insertion point created by the closest context manager is used; if there is no surrounding context manager, the operation is created in the detached state.

Most operations will customize the constructor to accept a reduced list of arguments that are relevant for the operation. For example, zero-result operations may omit the results argument, so can the operations where the result types can be derived from operand types unambiguously. As a concrete example, built-in function operations can be constructed by providing a function name as string and its argument and result types as a tuple of sequences:

from mlir.ir import Context, Module
from mlir.dialects import builtin

with Context():
  module = Module.create()
  with InsertionPoint(module.body), Location.unknown():
    func = func.FuncOp("main", ([], []))

Also see below for constructors generated from ODS.

Operations can also be constructed using the generic class and based on the canonical string name of the operation using Operation.create. It accepts the operation name as string, which must exactly match the canonical name of the operation in C++ or ODS, followed by the same argument list as the default constructor for OpView. This form is discouraged from use and is intended for generic operation processing.

from mlir.ir import Context, Module
from mlir.dialects import builtin

with Context():
  module = Module.create()
  with InsertionPoint(module.body), Location.unknown():
    # Operations can be created in a generic way.
    func = Operation.create(
        "func.func", results=[], operands=[],
        attributes={"function_type":TypeAttr.get(FunctionType.get([], []))},
        successors=None, regions=1)
    # The result will be downcasted to the concrete `OpView` subclass if
    # available.
    assert isinstance(func, func.FuncOp)

Regions are created for an operation when constructing it on the C++ side. They are not constructible in Python and are not expected to exist outside of operations (unlike in C++ that supports detached regions).

Blocks can be created within a given region and inserted before or after another block of the same region using create_before(), create_after() methods of the Block class, or the create_at_start() static method of the same class. They are not expected to exist outside of regions (unlike in C++ that supports detached blocks).

from mlir.ir import Block, Context, Operation

with Context():
  op = Operation.create("generic.op", regions=1)

  # Create the first block in the region.
  entry_block = Block.create_at_start(op.regions[0])

  # Create further blocks.
  other_block = entry_block.create_after()

Blocks can be used to create InsertionPoints, which can point to the beginning or the end of the block, or just before its terminator. It is common for OpView subclasses to provide a .body property that can be used to construct an InsertionPoint. For example, builtin Module and FuncOp provide a .body and .add_entry_blocK(), respectively.

Attributes and Types ¶

Attributes and types can be created given a Context or another attribute or type object that already references the context. To indicate that they are owned by the context, they are obtained by calling the static get method on the concrete attribute or type class. These method take as arguments the data necessary to construct the attribute or type and a the keyword context argument when the context cannot be derived from other arguments.

from mlir.ir import Context, F32Type, FloatAttr

# Attribute and types require access to an MLIR context, either directly or
# through another context-owned object.
ctx = Context()
f32 = F32Type.get(context=ctx)
pi = FloatAttr.get(f32, 3.14)

# They may use the context defined by the surrounding context manager.
with Context():
  f32 = F32Type.get()
  pi = FloatAttr.get(f32, 3.14)

Some attributes provide additional construction methods for clarity.

from mlir.ir import Context, IntegerAttr, IntegerType

with Context():
  i8 = IntegerType.get_signless(8)
  IntegerAttr.get(i8, 42)

Builtin attribute can often be constructed from Python types with similar structure. For example, ArrayAttr can be constructed from a sequence collection of attributes, and a DictAttr can be constructed from a dictionary:

from mlir.ir import ArrayAttr, Context, DictAttr, UnitAttr

with Context():
  array = ArrayAttr.get([UnitAttr.get(), UnitAttr.get()])
  dictionary = DictAttr.get({"array": array, "unit": UnitAttr.get()})

Custom builders for Attributes to be used during Operation creation can be registered by way of the register_attribute_builder. In particular the following is how a custom builder is registered for I32Attr:

@register_attribute_builder("I32Attr")
def _i32Attr(x: int, context: Context):
  return IntegerAttr.get(
        IntegerType.get_signless(32, context=context), x)

This allows to invoke op creation of an op with a I32Attr with

foo.Op(30)

The registration is based on the ODS name but registry is via pure python method. Only single custom builder is allowed to be registered per ODS attribute type (e.g., I32Attr can have only one, which can correspond to multiple of the underlying IntegerAttr type).

instead of

foo.Op(IntegerAttr.get(IndexType.get_signless(32, context=context), 30))

Style ¶

In general, for the core parts of MLIR, the Python bindings should be largely isomorphic with the underlying C++ structures. However, concessions are made either for practicality or to give the resulting library an appropriately “Pythonic” flavor.

Properties vs get*() methods ¶

Generally favor converting trivial methods like getContext(), getName(), isEntryBlock(), etc to read-only Python properties (i.e. context). It is primarily a matter of calling def_property_readonly vs def in binding code, and makes things feel much nicer to the Python side.

For example, prefer:

m.def_property_readonly("context", ...)

Over:

m.def("getContext", ...)

repr methods ¶

Things that have nice printed representations are really great :) If there is a reasonable printed form, it can be a significant productivity boost to wire that to the __repr__ method (and verify it with a doctest).

CamelCase vs snake_case ¶

Name functions/methods/properties in snake_case and classes in CamelCase. As a mechanical concession to Python style, this can go a long way to making the API feel like it fits in with its peers in the Python landscape.

If in doubt, choose names that will flow properly with other PEP 8 style names.

Prefer pseudo-containers ¶

Many core IR constructs provide methods directly on the instance to query count and begin/end iterators. Prefer hoisting these to dedicated pseudo containers.

For example, a direct mapping of blocks within regions could be done this way:

region = ...

for block in region:

  pass

However, this way is preferred:

region = ...

for block in region.blocks:

  pass

print(len(region.blocks))
print(region.blocks[0])
print(region.blocks[-1])

Instead of leaking STL-derived identifiers (front, back, etc), translate them to appropriate __dunder__ methods and iterator wrappers in the bindings.

Note that this can be taken too far, so use good judgment. For example, block arguments may appear container-like but have defined methods for lookup and mutation that would be hard to model properly without making semantics complicated. If running into these, just mirror the C/C++ API.

Provide one stop helpers for common things ¶

One stop helpers that aggregate over multiple low level entities can be incredibly helpful and are encouraged within reason. For example, making Context have a parse_asm or equivalent that avoids needing to explicitly construct a SourceMgr can be quite nice. One stop helpers do not have to be mutually exclusive with a more complete mapping of the backing constructs.

Testing ¶

Tests should be added in the test/Bindings/Python directory and should typically be .py files that have a lit run line.

We use lit and FileCheck based tests:

• For generative tests (those that produce IR), define a Python module that constructs/prints the IR and pipe it through FileCheck.
• Parsing should be kept self-contained within the module under test by use of raw constants and an appropriate parse_asm call.
• Any file I/O code should be staged through a tempfile vs relying on file artifacts/paths outside of the test module.
• For convenience, we also test non-generative API interactions with the same mechanisms, printing and CHECKing as needed.

Sample FileCheck test ¶

# RUN: %PYTHON %s | mlir-opt -split-input-file | FileCheck

# TODO: Move to a test utility class once any of this actually exists.
def print_module(f):
  m = f()
  print("// -----")
  print("// TEST_FUNCTION:", f.__name__)
  print(m.to_asm())
  return f

# CHECK-LABEL: TEST_FUNCTION: create_my_op
@print_module
def create_my_op():
  m = mlir.ir.Module()
  builder = m.new_op_builder()
  # CHECK: mydialect.my_operation ...
  builder.my_op()
  return m

Integration with ODS ¶

The MLIR Python bindings integrate with the tablegen-based ODS system for providing user-friendly wrappers around MLIR dialects and operations. There are multiple parts to this integration, outlined below. Most details have been elided: refer to the build rules and python sources under mlir.dialects for the canonical way to use this facility.

Users are responsible for providing a {DIALECT_NAMESPACE}.py (or an equivalent directory with __init__.py file) as the entrypoint.

Generating _{DIALECT_NAMESPACE}_ops_gen.py wrapper modules ¶

Each dialect with a mapping to python requires that an appropriate _{DIALECT_NAMESPACE}_ops_gen.py wrapper module is created. This is done by invoking mlir-tblgen on a python-bindings specific tablegen wrapper that includes the boilerplate and actual dialect specific td file. An example, for the Func (which is assigned the namespace func as a special case):

#ifndef PYTHON_BINDINGS_FUNC_OPS
#define PYTHON_BINDINGS_FUNC_OPS

include "mlir/Dialect/Func/IR/FuncOps.td"

#endif // PYTHON_BINDINGS_FUNC_OPS

In the main repository, building the wrapper is done via the CMake function declare_mlir_dialect_python_bindings, which invokes:

mlir-tblgen -gen-python-op-bindings -bind-dialect={DIALECT_NAMESPACE} \
    {PYTHON_BINDING_TD_FILE}

The generates op classes must be included in the {DIALECT_NAMESPACE}.py file in a similar way that generated headers are included for C++ generated code:

from ._my_dialect_ops_gen import *

Extending the search path for wrapper modules ¶

When the python bindings need to locate a wrapper module, they consult the dialect_search_path and use it to find an appropriately named module. For the main repository, this search path is hard-coded to include the mlir.dialects module, which is where wrappers are emitted by the above build rule. Out of tree dialects can add their modules to the search path by calling:

from mlir.dialects._ods_common import _cext
_cext.globals.append_dialect_search_prefix("myproject.mlir.dialects")

Wrapper module code organization ¶

The wrapper module tablegen emitter outputs:

• A _Dialect class (extending mlir.ir.Dialect) with a DIALECT_NAMESPACE attribute.
• An {OpName} class for each operation (extending mlir.ir.OpView).
• Decorators for each of the above to register with the system.

Note: In order to avoid naming conflicts, all internal names used by the wrapper module are prefixed by _ods_.

Each concrete OpView subclass further defines several public-intended attributes:

• OPERATION_NAME attribute with the str fully qualified operation name (i.e. math.absf).
• An __init__ method for the default builder if one is defined or inferred for the operation.
• @property getter for each operand or result (using an auto-generated name for unnamed of each).
• @property getter, setter and deleter for each declared attribute.

It further emits additional private-intended attributes meant for subclassing and customization (default cases omit these attributes in favor of the defaults on OpView):

• _ODS_REGIONS: A specification on the number and types of regions. Currently a tuple of (min_region_count, has_no_variadic_regions). Note that the API does some light validation on this but the primary purpose is to capture sufficient information to perform other default building and region accessor generation.
• _ODS_OPERAND_SEGMENTS and _ODS_RESULT_SEGMENTS: Black-box value which indicates the structure of either the operand or results with respect to variadics. Used by OpView._ods_build_default to decode operand and result lists that contain lists.

Default Builder ¶

Presently, only a single, default builder is mapped to the __init__ method. The intent is that this __init__ method represents the most specific of the builders typically generated for C++; however currently it is just the generic form below.

• One argument for each declared result:
  • For single-valued results: Each will accept an mlir.ir.Type.
  • For variadic results: Each will accept a List[mlir.ir.Type].
• One argument for each declared operand or attribute:
  • For single-valued operands: Each will accept an mlir.ir.Value.
  • For variadic operands: Each will accept a List[mlir.ir.Value].
  • For attributes, it will accept an mlir.ir.Attribute.
• Trailing usage-specific, optional keyword arguments:
  • loc: An explicit mlir.ir.Location to use. Defaults to the location bound to the thread (i.e. with Location.unknown():) or an error if none is bound nor specified.
  • ip: An explicit mlir.ir.InsertionPoint to use. Default to the insertion point bound to the thread (i.e. with InsertionPoint(...):).

In addition, each OpView inherits a build_generic method which allows construction via a (nested in the case of variadic) sequence of results and operands. This can be used to get some default construction semantics for operations that are otherwise unsupported in Python, at the expense of having a very generic signature.

Extending Generated Op Classes ¶

As mentioned above, the build system generates Python sources like _{DIALECT_NAMESPACE}_ops_gen.py for each dialect with Python bindings. It is often desirable to use these generated classes as a starting point for further customization, so an extension mechanism is provided to make this easy. This mechanism uses conventional inheritance combined with OpView registration. For example, the default builder for arith.constant

class ConstantOp(_ods_ir.OpView):
  OPERATION_NAME = "arith.constant"

  _ODS_REGIONS = (0, True)

  def __init__(self, value, *, loc=None, ip=None):
    ...

expects value to be a TypedAttr (e.g., IntegerAttr or FloatAttr). Thus, a natural extension is a builder that accepts a MLIR type and a Python value and instantiates the appropriate TypedAttr:

from typing import Union

from mlir.ir import Type, IntegerAttr, FloatAttr
from mlir.dialects._arith_ops_gen import _Dialect, ConstantOp
from mlir.dialects._ods_common import _cext

@_cext.register_operation(_Dialect, replace=True)
class ConstantOpExt(ConstantOp):
    def __init__(
        self, result: Type, value: Union[int, float], *, loc=None, ip=None
    ):
        if isinstance(value, int):
            super().__init__(IntegerAttr.get(result, value), loc=loc, ip=ip)
        elif isinstance(value, float):
            super().__init__(FloatAttr.get(result, value), loc=loc, ip=ip)
        else:
            raise NotImplementedError(f"Building `arith.constant` not supported for {result=} {value=}")

which enables building an instance of arith.constant like so:

from mlir.ir import F32Type

a = ConstantOpExt(F32Type.get(), 42.42)
b = ConstantOpExt(IntegerType.get_signless(32), 42)

Note, three key aspects of the extension mechanism in this example:

1. ConstantOpExt directly inherits from the generated ConstantOp;
2. in this, simplest, case all that’s required is a call to the super class’ initializer, i.e., super().__init__(...);
3. in order to register ConstantOpExt as the preferred OpView that is returned by mlir.ir.Operation.opview (see Operations, Regions and Blocks) we decorate the class with @_cext.register_operation(_Dialect, replace=True), where the replace=True must be used.

In some more complex cases it might be necessary to explicitly build the OpView through OpView.build_generic (see Default Builder), just as is performed by the generated builders. I.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. Thus, we must call a method of a super class’ super class (the “grandparent”); for example:

from mlir.dialects._scf_ops_gen import _Dialect, ForOp
from mlir.dialects._ods_common import _cext

@_cext.register_operation(_Dialect, replace=True)
class ForOpExt(ForOp):
    def __init__(self, lower_bound, upper_bound, step, iter_args, *, loc=None, ip=None):
        ...
        super(ForOp, self).__init__(self.build_generic(...))

where OpView.__init__ is called via super(ForOp, self).__init__. Note, there are alternatives ways to implement this (e.g., explicitly writing OpView.__init__); see any discussion on Python inheritance.

Providing Python bindings for a dialect ¶

Python bindings are designed to support MLIR’s open dialect ecosystem. A dialect can be exposed to Python as a submodule of mlir.dialects and interoperate with the rest of the bindings. For dialects containing only operations, it is sufficient to provide Python APIs for those operations. Note that the majority of boilerplate APIs can be generated from ODS. For dialects containing attributes and types, it is necessary to thread those through the C API since there is no generic mechanism to create attributes and types. Passes need to be registered with the context in order to be usable in a text-specified pass manager, which may be done at Python module load time. Other functionality can be provided, similar to attributes and types, by exposing the relevant C API and building Python API on top.

Operations ¶

Dialect operations are provided in Python by wrapping the generic mlir.ir.Operation class with operation-specific builder functions and properties. Therefore, there is no need to implement a separate C API for them. For operations defined in ODS, mlir-tblgen -gen-python-op-bindings -bind-dialect=<dialect-namespace> generates the Python API from the declarative description. It is sufficient to create a new .td file that includes the original ODS definition and use it as source for the mlir-tblgen call. Such .td files reside in python/mlir/dialects/. The results of mlir-tblgen are expected to produce a file named _<dialect-namespace>_ops_gen.py by convention. The generated operation classes can be extended as described above. MLIR provides CMake functions to automate the production of such files. Finally, a python/mlir/dialects/<dialect-namespace>.py or a python/mlir/dialects/<dialect-namespace>/__init__.py file must be created and filled with imports from the generated files to enable import mlir.dialects.<dialect-namespace> in Python.

Attributes and Types ¶

Dialect attributes and types are provided in Python as subclasses of the mlir.ir.Attribute and mlir.ir.Type classes, respectively. Python APIs for attributes and types must connect to the relevant C APIs for building and inspection, which must be provided first. Bindings for Attribute and Type subclasses can be defined using include/mlir/Bindings/Python/PybindAdaptors.h utilities that mimic pybind11 API for defining functions and properties. These bindings are to be included in a separate pybind11 module. The utilities also provide automatic casting between C API handles MlirAttribute and MlirType and their Python counterparts so that the C API handles can be used directly in binding implementations. The methods and properties provided by the bindings should follow the principles discussed above.

The attribute and type bindings for a dialect can be located in lib/Bindings/Python/Dialect<Name>.cpp and should be compiled into a separate “Python extension” library placed in python/mlir/_mlir_libs that will be loaded by Python at runtime. MLIR provides CMake functions to automate the production of such libraries. This library should be imported from the main dialect file, i.e. python/mlir/dialects/<dialect-namespace>.py or python/mlir/dialects/<dialect-namespace>/__init__.py, to ensure the types are available when the dialect is loaded from Python.

Passes ¶

Dialect-specific passes can be made available to the pass manager in Python by registering them with the context and relying on the API for pass pipeline parsing from string descriptions. This can be achieved by creating a new pybind11 module, defined in lib/Bindings/Python/<Dialect>Passes.cpp, that calls the registration C API, which must be provided first. For passes defined declaratively using Tablegen, mlir-tblgen -gen-pass-capi-header and -mlir-tblgen -gen-pass-capi-impl automate the generation of C API. The pybind11 module must be compiled into a separate “Python extension” library, which can be imported from the main dialect file, i.e. python/mlir/dialects/<dialect-namespace>.py or python/mlir/dialects/<dialect-namespace>/__init__.py, or from a separate passes submodule to be put in 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, can be exposed to Python similarly to attributes and types. C API is expected to exist for this functionality, which can then be wrapped using pybind11 and [include/mlir/Bindings/Python/PybindAdaptors.h](https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Bindings/Python/PybindAdaptors.h) utilities to connect to the rest of Python API. The bindings can be located in a separate pybind11 module or in the same module as attributes and types, and loaded along with the dialect.