pyMLIR: Python Interface for the Multi-Level Intermediate Representation
pyMLIR is a full Python interface to parse, process, and output MLIR files according to the syntax described in the MLIR documentation. pyMLIR supports the basic dialects and can be extended with other dialects. It uses Lark to parse the MLIR syntax, and mirrors the classes into Python classes. Custom dialects can also be implemented with a Python string-format-like syntax, or via direct parsing.
Note that the tool does not depend on LLVM or MLIR. It can be installed and invoked directly from Python.
Instructions
How to install: pip install pymlir
Requirements: Python 3.6 or newer, and the requirements in setup.py or requirements.txt. To manually install the
requirements, use pip install -r requirements.txt
Problem parsing MLIR files? Run the file through LLVM's mlir-opt --mlir-print-op-generic to get the generic form
of the IR (instructions on how to build/install MLIR can be found here):
$ mlir-opt file.mlir --mlir-print-op-generic > output.mlir
Found other problems parsing files? Not all dialects and modes are supported. Feel free to send us an issue or create a pull request! This is a community project and we welcome any contribution.
Usage examples
Parsing MLIR files into Python
import mlir
# Read a file path, file handle (stream), or a string
ast1 = mlir.parse_path('/path/to/file.mlir')
ast2 = mlir.parse_file(open('/path/to/file.mlir', 'r'))
ast3 = mlir.parse_string('''
module {
func @toy_func(%tensor: tensor<2x3xf64>) -> tensor<3x2xf64> {
%t_tensor = "toy.transpose"(%tensor) { inplace = true } : (tensor<2x3xf64>) -> tensor<3x2xf64>
return %t_tensor : tensor<3x2xf64>
}
}
''')Inspecting MLIR files in Python
MLIR files can be inspected by dumping their contents (which will print standard MLIR code), or by using the same tools as you would with Python's ast module.
import mlir
# Dump valid MLIR files
m = mlir.parse_path('/path/to/file.mlir')
print(m.dump())
print('---')
# Dump the AST directly
print(m.dump_ast())
print('---')
# Or visit each node type by implementing visitor functions
class MyVisitor(mlir.NodeVisitor):
def visit_Function(self, node: mlir.astnodes.Function):
print('Function detected:', node.name.value)
MyVisitor().visit(m)Transforming MLIR files
MLIR files can also be transformed with a Python-like NodeTransformer object.
import mlir
m = mlir.parse_path('/path/to/file.mlir')
# Simple node transformer that removes all operations with a result
class RemoveAllResultOps(mlir.NodeTransformer):
def visit_Operation(self, node: mlir.astnodes.Operation):
# There are one or more outputs, return None to remove from AST
if len(node.result_list) > 0:
return None
# No outputs, no need to do anything
return self.generic_visit(node)
m = RemoveAllResultOps().visit(m)
# Write back to file
with open('output.mlir', 'w') as fp:
fp.write(m.dump())Using custom dialects
Custom dialects can be written and loaded as part of the pyMLIR parser. See full tutorial here.
import mlir
from lark import UnexpectedCharacters
from .mydialect import dialect
# Try to parse as-is
try:
m = mlir.parse_path('/path/to/matrixfile.mlir')
except UnexpectedCharacters: # MyMatrix dialect not recognized
pass
# Add dialect to the parser
m = mlir.parse_path('/path/to/matrixfile.mlir',
dialects=[dialect])
# Print output back
print(m.dump_ast())MLIR from scratch with the builder API
pyMLIR has a Builder API that can create MLIR ASTs on the fly within Python code.
import mlir.builder
builder = mlir.builder.IRBuilder()
mlirfile = builder.make_mlir_file()
module = mlirfile.default_module
with builder.goto_block(builder.make_block(module.region)):
hello = builder.function("hello_world")
block = builder.make_block(hello.region)
builder.position_at_entry(block)
x, y = builder.add_function_args(hello, [builder.F64, builder.F64], ['a', 'b'])
adder = builder.addf(x, y, builder.F64)
builder.ret([adder], [builder.F64])
print(mlirfile.dump())prints:
module {
func @hello_world(%a: f64, %b: f64) {
%_pymlir_ssa = addf %a , %b : f64
return %_pymlir_ssa : f64
}
}See also saxpy for a full example that registers and uses a dialect in the builder.
Built-in dialect implementations and more examples
All dialect implementations can be found in the dialects subfolder. Additional uses of the library, including a custom dialect implementation, can be found in the tests subfolder.